From 0f330bf6e7d4c1200eb832dde463756f452d05f3 Mon Sep 17 00:00:00 2001
From: wangzhengzhuo05 <175673456+wangzhengzhuo05@users.noreply.github.com>
Date: Sat, 12 Sep 2026 00:10:38 +0800
Subject: [PATCH] fix(orchestrator): record failed stage in state.json when a
 handler raises

---
 src/cain_agent/orchestrator.py | 23 +++++++++++++++--
 tests/test_orchestrator.py     | 46 +++++++++++++++++++++++++++++++++-
 2 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/src/cain_agent/orchestrator.py b/src/cain_agent/orchestrator.py
index c46be13..d7bae80 100644
--- a/src/cain_agent/orchestrator.py
+++ b/src/cain_agent/orchestrator.py
@@ -165,11 +165,30 @@ def run_stage(self, stage: str) -> StageResult:
             artifacts_dir=self.workspace.stage_dir(stage),
         )
         started_at = _utc_now_iso()
-        result = self.handlers[stage](ctx)
+        history: list[dict[str, Any]] = list(state.get("history", []))
+        try:
+            result = self.handlers[stage](ctx)
+        except Exception as exc:
+            failed_at = _utc_now_iso()
+            history.append({
+                "stage": stage,
+                "started_at": started_at,
+                "finished_at": failed_at,
+                "status": "failed",
+                "error": f"{type(exc).__name__}: {exc}",
+            })
+            self._save_state({
+                "current_stage": stage,
+                "completed_stages": completed,
+                "updated_at": failed_at,
+                "history": history,
+                "failed_stage": stage,
+                "error": f"{type(exc).__name__}: {exc}",
+            })
+            raise
         finished_at = _utc_now_iso()
 
         completed.append(stage)
-        history: list[dict[str, Any]] = list(state.get("history", []))
         history.append({
             "stage": stage,
             "started_at": started_at,
diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py
index be0aaec..97d9907 100644
--- a/tests/test_orchestrator.py
+++ b/tests/test_orchestrator.py
@@ -27,7 +27,7 @@
     StageOrderError,
     StageResult,
 )
-from cain_agent.workspace import Workspace
+from cain_agent.workspace import STATE_FILE, Workspace
 
 
 @pytest.fixture
@@ -170,3 +170,47 @@ def test_artifacts_landed_in_stage_dirs(ws: Workspace) -> None:
     orch.run()
     for stage in STAGES:
         assert (ws.root / stage / f"{stage}-placeholder.json").exists()
+
+
+def _raising_handler(error: Exception) -> StageHandler:
+    def handler(ctx: StageContext) -> StageResult:
+        raise error
+
+    return handler
+
+
+def test_failed_stage_is_recorded_in_state(ws: Workspace) -> None:
+    orch = Orchestrator(
+        SDKExecutor(), ws, handlers={"recon": _raising_handler(RuntimeError("recon blew up"))}
+    )
+    with pytest.raises(RuntimeError, match="recon blew up"):
+        orch.run_stage("recon")
+    state = orch.load_state()
+    assert state["failed_stage"] == "recon"
+    assert "recon blew up" in state["error"]
+    assert state["completed_stages"] == []
+    assert state["current_stage"] == "recon"
+    assert state["history"][-1]["status"] == "failed"
+    assert ws.path(STATE_FILE).exists()
+
+
+def test_run_records_state_when_a_handler_raises(ws: Workspace) -> None:
+    orch = Orchestrator(
+        SDKExecutor(), ws, handlers={"recon": _raising_handler(RuntimeError("recon blew up"))}
+    )
+    with pytest.raises(RuntimeError):
+        orch.run()
+    assert orch.load_state()["failed_stage"] == "recon"
+
+
+def test_failure_marker_is_cleared_after_a_successful_stage(ws: Workspace) -> None:
+    orch = Orchestrator(
+        SDKExecutor(), ws, handlers={"recon": _raising_handler(RuntimeError("recon blew up"))}
+    )
+    with pytest.raises(RuntimeError, match="recon blew up"):
+        orch.run_stage("recon")
+    orch2 = Orchestrator(SDKExecutor(), ws)
+    orch2.run_stage("recon")
+    state = orch2.load_state()
+    assert state["completed_stages"] == ["recon"]
+    assert "failed_stage" not in state