Harden Dask error handling payloads and recoverability

Author: Codex-Agent

src/agilab/core/agi-cluster/src/agi_cluster/agi_distributor/api/entrypoint_support.py

@@ -140,6 +140,12 @@ def _connection_error_payload(exc: ConnectionError, *, log: Any = logger) -> Dic
     return {"status": "error", "message": message, "kind": "connection"}
 
 
+def _application_error_payload(exc: BaseException, *, kind: str, log: Any = logger) -> Dict[str, str]:
+    message = str(exc).strip() or "Execution failed."
+    log.error("AGI.run failed (%s): %s", kind, message)
+    return {"status": "error", "message": message, "kind": kind}
+
+
 def _log_unhandled_run_exception(
     exc: Exception,
     *,
@@ -165,13 +171,15 @@ async def _run_main_with_handled_errors(
     try:
         return await agi_cls._main(scheduler)
     except process_error_type as exc:
-        log.error("failed to run \n%s", exc)
-        return None
+        return _application_error_payload(exc, kind="process", log=log)
     except ConnectionError as exc:
         return _connection_error_payload(exc, log=log)
     except ModuleNotFoundError as exc:
-        log.error("failed to load module \n%s", exc)
-        return None
+        return _application_error_payload(
+            exc,
+            kind="module_not_found",
+            log=log,
+        )
     except _AGI_RUN_BOUNDARY_EXCEPTIONS as exc:  # Intentional AGI.run boundary: log and re-raise.
         _log_unhandled_run_exception(
             exc,

src/agilab/core/agi-cluster/src/agi_cluster/agi_distributor/runtime/runtime_distribution_support.py

@@ -548,7 +548,7 @@ async def sync(
                 break
 
             if time_fn() - start_time > timeout:
-                log.error("Timeout waiting for all workers. {remaining} workers missing.")
+                log.error("Timeout waiting for all workers. %s workers missing.", remaining)
                 raise TimeoutError("Timed out waiting for all workers to attach")
             await sleep_fn(_sync_poll_delay(poll_attempt))
             poll_attempt += 1

src/agilab/core/agi-cluster/src/agi_cluster/agi_distributor/service/service_lifecycle_support.py

@@ -28,8 +28,6 @@
     OSError,
     RuntimeError,
     TimeoutError,
-    TypeError,
-    ValueError,
 )
 _SERVICE_BREAK_LOOP_EXCEPTIONS = (
     ConnectionError,

src/agilab/core/test/test_agi_distributor_entrypoint_support.py

@@ -467,8 +467,12 @@ async def _raise_missing(_scheduler):
         traceback_format_exc_fn=lambda: "tb",
         log=process_log,
     )
-    assert process_result is None
-    assert process_log.error_messages == ["failed to run \nprocess failed"]
+    assert process_result == {
+        "status": "error",
+        "message": "process failed",
+        "kind": "process",
+    }
+    assert process_log.error_messages == ["AGI.run failed (process): process failed"]
 
     connection_log = _FakeLogger()
     connection_result = await entrypoint_support._run_main_with_handled_errors(
@@ -496,8 +500,12 @@ async def _raise_missing(_scheduler):
         traceback_format_exc_fn=lambda: "tb",
         log=missing_log,
     )
-    assert missing_result is None
-    assert missing_log.error_messages == ["failed to load module \nmissing module"]
+    assert missing_result == {
+        "status": "error",
+        "message": "missing module",
+        "kind": "module_not_found",
+    }
+    assert missing_log.error_messages == ["AGI.run failed (module_not_found): missing module"]
 
 
 @pytest.mark.asyncio

src/agilab/core/test/test_agi_distributor_run_api.py

@@ -239,7 +239,7 @@ async def _fake_main(_scheduler):
 
 
 @pytest.mark.asyncio
-async def test_agi_run_returns_none_on_process_error(monkeypatch):
+async def test_agi_run_returns_process_error_payload(monkeypatch):
     env = _minimal_app_env()
     env.base_worker_cls = "PandasWorker"
 
@@ -257,7 +257,11 @@ async def _fake_main(_scheduler):
         env,
         request=RunRequest(scheduler="127.0.0.1", workers={"127.0.0.1": 1}, mode=AGI.DASK_MODE),
     )
-    assert result is None
+    assert result == {
+        "status": "error",
+        "message": "process failed",
+        "kind": "process",
+    }
 
 
 @pytest.mark.asyncio
@@ -280,7 +284,7 @@ async def _boom(_scheduler):
 
 
 @pytest.mark.asyncio
-async def test_agi_run_returns_none_on_module_not_found(monkeypatch):
+async def test_agi_run_returns_module_not_found_payload(monkeypatch):
     env = _minimal_app_env()
     env.base_worker_cls = "PandasWorker"
     monkeypatch.setattr(AGI, "_train_capacity", staticmethod(lambda *_args, **_kwargs: None))
@@ -289,7 +293,11 @@ async def _missing(_scheduler):
         raise ModuleNotFoundError("missing module")
 
     monkeypatch.setattr(AGI, "_main", staticmethod(_missing))
-    assert await AGI.run(env, request=RunRequest(workers={"127.0.0.1": 1}, mode=AGI.DASK_MODE)) is None
+    assert await AGI.run(env, request=RunRequest(workers={"127.0.0.1": 1}, mode=AGI.DASK_MODE)) == {
+        "status": "error",
+        "message": "missing module",
+        "kind": "module_not_found",
+    }
 
 
 @pytest.mark.asyncio

src/agilab/core/test/test_agi_distributor_service_lifecycle_support.py

@@ -469,6 +469,31 @@ def _boom(_args):
         await AGI._service_recover(env, allow_stale_cleanup=True)
 
 
+@pytest.mark.asyncio
+async def test_service_recover_propagates_unexpected_value_error(monkeypatch):
+    env = AgiEnv(apps_path=Path("src/agilab/apps/builtin"), app="minimal_app_project", verbose=0)
+    AGI._service_write_state(
+        env,
+        {
+            "schema": "agi.service.state.v1",
+            "target": env.target,
+            "app": env.app,
+            "mode": AGI.DASK_MODE,
+            "run_type": "run --no-sync",
+            "workers": {"127.0.0.1": 1},
+            "args": {"sample": 1},
+        },
+    )
+
+    def _boom(_args):
+        raise ValueError("unexpected service arg mapping")
+
+    monkeypatch.setattr(AGI, "_service_public_args", staticmethod(_boom))
+
+    with pytest.raises(ValueError, match="unexpected service arg mapping"):
+        await AGI._service_recover(env, allow_stale_cleanup=True)
+
+
 @pytest.mark.asyncio
 async def test_service_recover_fails_when_no_workers_attached(monkeypatch):
     env = AgiEnv(apps_path=Path("src/agilab/apps/builtin"), app="minimal_app_project", verbose=0)