Add flyte.is_control_plane_available() and Run.first_failure() for agent edit-and-relaunch loops (#1522)

## Why

The agent-mediated forking example (unionai/unionai-examples#308) — an
agent task that launches a workflow, observes the failure, patches the
workflow's source on disk, reloads it, and `fork()`s the failed run with
the fixed code — had to reach into SDK internals:

1. `isinstance(get_controller(), RemoteController)` to decide whether
launching/forking real runs is possible, or whether to fall back to
running the workflow inline.
2. Hand-rolled `Action.listall(...FAILED...)` iteration to find which
step failed and its error message.
3. `flyte._code_bundle.bundle.build_code_bundle.cache_clear()` to make
the next fork ship the edited working tree — **no longer needed**: #1508
removed the code-bundle memoization entirely, so every launch re-bundles
from disk. This PR originally shipped a
`flyte.refresh_code_bundle_cache()` for that; it was dropped when
merging main, since the cache it managed no longer exists.

## What

**`flyte.is_control_plane_available()`** — True when the process can
submit work to a control plane (launch real runs whose actions can be
awaited and replayed/forked). Inside a task, the orchestration mode
decides (`remote`/`hybrid` → True, `local` → False, even when a client
is configured — `flyte run --local` configures one too); outside a task,
a configured client decides. This replaces the
isinstance-on-internal-controller probe with the `TaskContext.mode` the
runtime already maintains.

**`Run.first_failure()` and `ActionDetails.error_message`** — the
observation half of a repair loop: which step of a run failed, and why.
`first_failure()` returns the `ActionDetails` of the first failed action
in creation order, preferring a failed sub-action over the failed root
(whose error usually just repeats the sub-action's); `error_message` is
the failed action's message or `""`.

With these, the example's loop reduces to:

```python
run = await flyte.run.aio(wf.main, n_records=n_records)
await run.wait.aio(quiet=True)
if failure := await run.first_failure.aio():
    patch_workflow_source(failure.task_name, failure.error_message)
    importlib.reload(sys.modules["workflow"])
    run = await fork.aio(run.name, task_template=wf.main)
```

## Testing

- New unit tests: `is_control_plane_available()` across
uninitialized/client/local/remote/hybrid contexts; `Run.first_failure()`
sub-action preference, root fallback, and no-failure; `error_message`.
- Existing `code_bundle`, `remote`, `deploy`, and `cli/test_run.py`
suites pass (the only failures are pre-existing on `main`: the two
`loaded_modules` discovery tests).
- `ruff`, `mypy`, and `ty` clean via pre-commit hooks.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_01NwQixBcyR5va6BC75jaQx3

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: cosmicBboy

src/flyte/__init__.py

@@ -25,6 +25,7 @@
     init_from_config,
     init_in_cluster,
     init_passthrough,
+    is_control_plane_available,
 )
 from ._interactive_run_context import load_interactive_ctx
 from ._link import Link
@@ -108,6 +109,7 @@ def version() -> str:
     "init_from_config",
     "init_in_cluster",
     "init_passthrough",
+    "is_control_plane_available",
     "latest_checkpoint",
     "load_interactive_ctx",
     "load_plugin_config",

src/flyte/_initialize.py

@@ -801,6 +801,48 @@ def is_initialized() -> bool:
     return _get_init_config() is not None
 
 
+def is_control_plane_available() -> bool:
+    """
+    True when this process can submit work to a Flyte control plane — `flyte.run` launches real
+    remote runs whose actions can be inspected, awaited, and replayed (recovered/forked).
+
+    The answer depends on where the code is executing:
+
+    * Inside a task launched on a Flyte cluster (`flyte.ctx().is_in_cluster()`): True. The
+      in-cluster runtime configures the connection before user code runs.
+    * Inside a task executing locally (`flyte run --local` / `flyte.run(mode="local")`): False,
+      even when a client happens to be configured — the local dev loop is expected to stay
+      local, and a locally-orchestrated run has no control-plane actions to replay.
+    * Outside any task (a driver script, a notebook): True iff a client has been configured via
+      `flyte.init` / `flyte.init_from_config` / `flyte.init_from_api_key`.
+
+    Typical use is a task that adapts to where it runs — e.g. an agent that launches and forks
+    real runs on a cluster, but falls back to invoking the task functions in-process when
+    developed locally:
+
+    ```python
+    @env.task
+    async def agent() -> None:
+        if flyte.is_control_plane_available():
+            run = await flyte.run.aio(my_pipeline, x=1)
+            await run.wait.aio()
+        else:
+            await my_pipeline.func(x=1)
+    ```
+
+    Returns:
+        True when remote submission is available, False otherwise.
+    """
+    from flyte._context import ctx
+
+    tctx = ctx()
+    if tctx:
+        # "remote" and "hybrid" tasks are orchestrated by a control plane; "local" ones are not.
+        return tctx.mode != "local"
+    cfg = _get_init_config()
+    return cfg is not None and cfg.client is not None
+
+
 def is_persistence_enabled() -> bool:
     """Check if local run persistence is enabled."""
     cfg = _get_init_config()

src/flyte/remote/_action.py

@@ -1020,6 +1020,14 @@ def error_info(self) -> run_definition_pb2.ErrorInfo | None:
             return self.pb2.error_info
         return None
 
+    @property
+    def error_message(self) -> str:
+        """
+        The error message of a failed action, or an empty string when the action did not fail
+        (or carries no error details).
+        """
+        return self.pb2.error_info.message if self.pb2.HasField("error_info") else ""
+
     @property
     def abort_info(self) -> run_definition_pb2.AbortInfo | None:
         """

src/flyte/remote/_run.py

@@ -315,6 +315,35 @@ async def details(self) -> RunDetails:
         self._details.action_details._preserve_original_types = self._preserve_original_types
         return self._details
 
+    @syncify
+    async def first_failure(self) -> ActionDetails | None:
+        """
+        Details of the first action that failed in this run, or None when no action failed.
+
+        Failed actions are considered in creation order, and a failed sub-action (a step inside
+        the run) is preferred over the failed root action — the root's error usually just
+        repeats the sub-action's. The root's details are returned only when it is the only
+        failure recorded. Together with `ActionDetails.task_name` and
+        `ActionDetails.error_message`, this answers "which step broke, and why" — the
+        observation an automated repair loop (or a human) needs before patching code and
+        rerunning or forking the run:
+
+        ```python
+        run = flyte.run(my_pipeline, x=1)
+        run.wait()
+        if failure := run.first_failure():
+            print(f"{failure.task_name} failed: {failure.error_message}")
+        ```
+        """
+        fallback: ActionDetails | None = None
+        async for action in Action.listall.aio(for_run_name=self.name, in_phase=(ActionPhase.FAILED,)):
+            details = await action.details()
+            if action.parent_name:
+                return details
+            if fallback is None:
+                fallback = details
+        return fallback
+
     @syncify
     async def inputs(self) -> ActionInputs:
         """

tests/flyte/remote/test_run_first_failure.py

@@ -0,0 +1,131 @@
+"""
+Tests for `Run.first_failure()` and `ActionDetails.error_message`.
+
+These are the observation half of an automated repair loop: after a run fails, an agent (or a
+human script) asks which step broke and why, patches the code, and reruns/forks the run.
+"""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from flyteidl2.common import identifier_pb2, phase_pb2
+from flyteidl2.workflow import run_definition_pb2
+
+from flyte.remote._action import ActionDetails
+from flyte.remote._run import Run
+
+RUN_NAME = "run-abc"
+
+
+def _action_pb2(name: str, parent: str | None = None, task_name: str | None = None) -> run_definition_pb2.Action:
+    action = run_definition_pb2.Action(
+        id=identifier_pb2.ActionIdentifier(
+            run=identifier_pb2.RunIdentifier(org="o", project="p", domain="d", name=RUN_NAME),
+            name=name,
+        ),
+        status=run_definition_pb2.ActionStatus(phase=phase_pb2.ACTION_PHASE_FAILED),
+    )
+    if parent:
+        action.metadata.parent = parent
+    if task_name:
+        action.metadata.task.id.name = task_name
+    return action
+
+
+def _details_pb2(name: str, message: str) -> run_definition_pb2.ActionDetails:
+    return run_definition_pb2.ActionDetails(
+        id=identifier_pb2.ActionIdentifier(
+            run=identifier_pb2.RunIdentifier(org="o", project="p", domain="d", name=RUN_NAME),
+            name=name,
+        ),
+        error_info=run_definition_pb2.ErrorInfo(message=message),
+    )
+
+
+def _run(phase=phase_pb2.ACTION_PHASE_FAILED) -> Run:
+    return Run(
+        run_definition_pb2.Run(
+            action=run_definition_pb2.Action(
+                id=identifier_pb2.ActionIdentifier(
+                    run=identifier_pb2.RunIdentifier(org="o", project="p", domain="d", name=RUN_NAME),
+                    name="a0",
+                ),
+                status=run_definition_pb2.ActionStatus(phase=phase),
+            )
+        )
+    )
+
+
+def _mock_client(failed_actions, details_by_name):
+    client = MagicMock()
+
+    resp = MagicMock()
+    resp.actions = failed_actions
+    resp.token = ""
+    client.run_service.list_actions = AsyncMock(return_value=resp)
+
+    async def get_action_details(request):
+        details_resp = MagicMock()
+        details_resp.details = details_by_name[request.action_id.name]
+        return details_resp
+
+    client.run_service.get_action_details = AsyncMock(side_effect=get_action_details)
+    return client
+
+
+def _first_failure(run, client):
+    cfg = MagicMock()
+    cfg.org, cfg.project, cfg.domain = "o", "p", "d"
+    with (
+        patch("flyte.remote._action.ensure_client"),
+        patch("flyte.remote._action.get_client", return_value=client),
+        patch("flyte.remote._action.get_init_config", return_value=cfg),
+    ):
+        return run.first_failure()
+
+
+class TestRunFirstFailure:
+    def test_prefers_failed_sub_action_over_root(self):
+        # The root action's error just repeats the step's — the step is the useful answer.
+        client = _mock_client(
+            failed_actions=[
+                _action_pb2("a0"),
+                _action_pb2("clean-1", parent="a0", task_name="clean_records"),
+            ],
+            details_by_name={
+                "a0": _details_pb2("a0", "child failed"),
+                "clean-1": _details_pb2("clean-1", "KeyError: 'price'"),
+            },
+        )
+        failure = _first_failure(_run(), client)
+        assert failure is not None
+        assert failure.error_message == "KeyError: 'price'"
+
+    def test_falls_back_to_root_when_only_failure(self):
+        client = _mock_client(
+            failed_actions=[_action_pb2("a0")],
+            details_by_name={"a0": _details_pb2("a0", "OOMKilled")},
+        )
+        failure = _first_failure(_run(), client)
+        assert failure is not None
+        assert failure.error_message == "OOMKilled"
+
+    def test_none_when_no_action_failed(self):
+        client = _mock_client(failed_actions=[], details_by_name={})
+        assert _first_failure(_run(phase=phase_pb2.ACTION_PHASE_SUCCEEDED), client) is None
+
+
+class TestActionDetailsErrorMessage:
+    def test_message_of_failed_action(self):
+        details = ActionDetails(_details_pb2("a0", "boom"))
+        assert details.error_message == "boom"
+
+    def test_empty_when_no_error_info(self):
+        details = ActionDetails(
+            run_definition_pb2.ActionDetails(
+                id=identifier_pb2.ActionIdentifier(
+                    run=identifier_pb2.RunIdentifier(org="o", project="p", domain="d", name=RUN_NAME),
+                    name="a0",
+                ),
+            )
+        )
+        assert details.error_message == ""

tests/flyte/test_is_control_plane_available.py

@@ -0,0 +1,62 @@
+"""
+Tests for `flyte.is_control_plane_available()`.
+
+The predicate answers "can this process submit work to a control plane" — inside a task it is
+decided by how the task is orchestrated (mode), outside a task by whether a client is
+configured. See the function's docstring for the full decision table.
+"""
+
+import pathlib
+
+import pytest
+
+import flyte
+from flyte._context import internal_ctx
+from flyte._initialize import _InitConfig
+from flyte.models import ActionID, RawDataPath, TaskContext
+from flyte.report import Report
+
+
+def _tctx(mode: str) -> TaskContext:
+    return TaskContext(
+        action=ActionID(name="a0"),
+        version="v1",
+        raw_data_path=RawDataPath(path="/tmp/rd"),
+        output_path="/tmp/o",
+        run_base_dir="/tmp",
+        report=Report(name="t"),
+        mode=mode,
+    )
+
+
+def test_false_when_uninitialized(monkeypatch):
+    monkeypatch.setattr("flyte._initialize._init_config", None)
+    assert flyte.is_control_plane_available() is False
+
+
+def test_false_when_initialized_without_client(monkeypatch):
+    monkeypatch.setattr("flyte._initialize._init_config", _InitConfig(root_dir=pathlib.Path.cwd()))
+    assert flyte.is_control_plane_available() is False
+
+
+def test_true_when_client_configured(monkeypatch):
+    monkeypatch.setattr(
+        "flyte._initialize._init_config",
+        _InitConfig(root_dir=pathlib.Path.cwd(), client=object()),
+    )
+    assert flyte.is_control_plane_available() is True
+
+
+@pytest.mark.parametrize(
+    ("mode", "expected"),
+    [("remote", True), ("hybrid", True), ("local", False)],
+)
+def test_task_context_mode_decides(monkeypatch, mode, expected):
+    # Inside a task, the orchestration mode wins over client presence: `flyte run --local`
+    # configures a client too, but a locally-orchestrated run has no control plane behind it.
+    monkeypatch.setattr(
+        "flyte._initialize._init_config",
+        _InitConfig(root_dir=pathlib.Path.cwd(), client=object()),
+    )
+    with internal_ctx().replace_task_context(_tctx(mode)):
+        assert flyte.is_control_plane_available() is expected