fix(engine): session-end durability + raw KeyboardInterrupt handling (#62)

Session-end durability barrier: record completed session before post-session hooks; usage folded in afterward as best-effort metadata. Resume continuation consumes durably-recorded dev/review results into verify/decide instead of restart-rollback. Raw KeyboardInterrupt branch records a controlled stop. Replay preserves attempt/cycle/baseline.

Closes #57

Author: dracic

src/bmad_loop/engine.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import contextlib
+import functools
 import shutil
 import signal
 import sys
@@ -245,6 +246,22 @@ def run(self) -> RunSummary:
                     raise  # nested auto-sweep: let the owner record the stop
                 self.state.stopped = True
                 self.journal.append("run-stop")
+            except KeyboardInterrupt:
+                # Some Windows console/control events can still surface as a raw
+                # KeyboardInterrupt without routing through the installed signal
+                # handler. Persist a controlled stop rather than letting the
+                # engine disappear with stale state.
+                self._stopping = True  # swallow stop signals landing mid-teardown
+                try:
+                    kill_session(self.state.run_id)
+                except (
+                    BaseException
+                ):  # noqa: BLE001  # nosec B110 - best-effort teardown; the stop must still record
+                    pass
+                if self._is_nested:
+                    raise
+                self.state.stopped = True
+                self.journal.append("run-stop", reason="KeyboardInterrupt")
             except Exception as exc:
                 # an unexpected exception escaped the loop (e.g. a transport
                 # hang that leaked past the seam). Don't let it die to the lossy
@@ -1043,6 +1060,36 @@ def _finish_inflight(self) -> None:
                     self._integrate_unit(task, unit)
                 else:
                     self._review_and_commit(task)
+            elif (resumable := self._resumable_session(task)) is not None:
+                # the host died inside the post-session window: the session
+                # itself completed and its recorded result is on disk, so
+                # continue into the normal verify/decide pipeline instead of
+                # rolling the finished work back through resume-restart.
+                role, result = resumable
+                self.journal.append("resume-verify", story_key=task.story_key, role=role)
+                if role == "dev":
+                    # deliberate reset like the restart arm: _dev_phase re-enters
+                    # its loop and consumes the recorded result instead of
+                    # running a session
+                    task.phase = Phase.PENDING
+                    continuation = functools.partial(self._drive_story, task, dev_resume=result)
+                else:
+                    # deliberate reset to the legal pre-review phase
+                    task.phase = Phase.DEV_VERIFY
+                    continuation = functools.partial(
+                        self._review_and_commit, task, resume_result=result
+                    )
+                if isolated:
+                    unit = self._reopen_unit(task)
+                    prev = self.workspace
+                    self.workspace = unit.workspace
+                    try:
+                        continuation()
+                    finally:
+                        self.workspace = prev
+                    self._integrate_unit(task, unit)
+                else:
+                    continuation()
             else:
                 self.journal.append(
                     "resume-restart", story_key=task.story_key, phase=str(task.phase)
@@ -1062,6 +1109,33 @@ def _finish_inflight(self) -> None:
                 self._save()
                 self._run_story(task)
 
+    def _resumable_session(self, task: StoryTask) -> tuple[str, SessionResult] | None:
+        """The in-flight session's durably-recorded result, when complete enough
+        to act on: the task died mid-phase but its current attempt/cycle record
+        is ``completed`` and carries the parsed result. Consumes only evidence
+        the adapter vouched for at session end — no artifact re-scan, no
+        loosening of completion authority. Anything less returns None and the
+        caller falls through to resume-restart."""
+        if task.phase == Phase.DEV_RUNNING:
+            role, seq = "dev", task.attempt
+        elif task.phase == Phase.REVIEW_RUNNING:
+            role, seq = "review", task.review_cycle
+        else:
+            return None
+        task_id = f"{task.story_key}-{role}-{seq}"
+        for record in reversed(task.sessions):
+            if record.task_id != task_id:
+                continue
+            if record.status != "completed" or record.result_json is None:
+                return None
+            return role, SessionResult(
+                status="completed",
+                result_json=record.result_json,
+                session_id=record.session_id,
+                transcript_path=record.transcript_path,
+            )
+        return None
+
     # ------------------------------------------------------------- per story
 
     def _gate_unit(self, task: StoryTask) -> bool:
@@ -1269,8 +1343,8 @@ def _run_story(self, task: StoryTask) -> None:
             self._drive_story(task)
         self._emit("post_story", task)
 
-    def _drive_story(self, task: StoryTask) -> None:
-        if not self._dev_phase(task):
+    def _drive_story(self, task: StoryTask, dev_resume: SessionResult | None = None) -> None:
+        if not self._dev_phase(task, resume_result=dev_resume):
             return
         if gates.pause_after_spec(self.policy):
             gates.notify(
@@ -1286,24 +1360,40 @@ def _drive_story(self, task: StoryTask) -> None:
             )
         self._review_and_commit(task)
 
-    def _dev_phase(self, task: StoryTask) -> bool:
+    def _dev_phase(self, task: StoryTask, resume_result: SessionResult | None = None) -> bool:
         if self._vetoed(self._emit("pre_dev_phase", task), task):
             return False
-        task.baseline_commit = verify.rev_parse_head(self.workspace.root)
-        # snapshot untracked files now so a later rollback removes only what THIS
-        # attempt creates, never files the user already had on disk.
-        task.baseline_untracked = sorted(verify.untracked_files(self.workspace.root))
+        if resume_result is None:
+            task.baseline_commit = verify.rev_parse_head(self.workspace.root)
+            # snapshot untracked files now so a later rollback removes only what
+            # THIS attempt creates, never files the user already had on disk.
+            # A resumed result keeps the persisted baseline: re-capturing here
+            # would shift the rollback/squash reference onto the completed
+            # session's own tree.
+            task.baseline_untracked = sorted(verify.untracked_files(self.workspace.root))
         feedback: Path | None = None
         while True:
-            task.attempt += 1
+            if resume_result is None:
+                # a resumed result replays the attempt it was recorded under, so
+                # the counter (and the session task_id derived from it) must not
+                # advance; a second host death then still finds the record and
+                # re-enters this continuation instead of falling back to restart.
+                task.attempt += 1
             advance(task, Phase.DEV_RUNNING)
             self._save()
-            result = self._run_session(
-                task,
-                role="dev",
-                prompt=self._dev_prompt(task, feedback),
-                seq=task.attempt,
-            )
+            if resume_result is not None:
+                # the session already ran before the host died; its recorded
+                # result re-enters the verify/decide pipeline. Consumed exactly
+                # once — later iterations run sessions normally.
+                result = resume_result
+                resume_result = None
+            else:
+                result = self._run_session(
+                    task,
+                    role="dev",
+                    prompt=self._dev_prompt(task, feedback),
+                    seq=task.attempt,
+                )
             advance(task, Phase.DEV_VERIFY)
             outcome = None
             if result.status == "completed":
@@ -1384,7 +1474,9 @@ def _record_dev_spec(self, task: StoryTask, result_json: dict | None) -> None:
         if spec_path.is_file():
             task.spec_file = str(spec_path)
 
-    def _review_and_commit(self, task: StoryTask) -> None:
+    def _review_and_commit(
+        self, task: StoryTask, resume_result: SessionResult | None = None
+    ) -> None:
         if not self.policy.review.enabled:
             # review.enabled = false: the bmad-dev-auto session's own inline
             # review is the only review; verify the deterministic gates + commit.
@@ -1418,16 +1510,27 @@ def _review_and_commit(self, task: StoryTask) -> None:
         # only state the budget-exhaustion rescue below is allowed to commit.
         refileable_followup = False
         while task.review_cycle < self.policy.limits.max_review_cycles:
-            task.review_cycle += 1
+            if resume_result is None:
+                # a resumed result replays the cycle it was recorded under: the
+                # counter must not advance, or the replay burns a review-budget
+                # slot and mislabels its journal/session ids.
+                task.review_cycle += 1
             refileable_followup = False  # only a completed pass this cycle can set it
             advance(task, Phase.REVIEW_RUNNING)
             self._save()
-            result = self._run_session(
-                task,
-                role="review",
-                prompt=self._review_prompt(task),
-                seq=task.review_cycle,
-            )
+            if resume_result is not None:
+                # the session already ran before the host died; its recorded
+                # result re-enters the decision pipeline. Consumed exactly
+                # once — later cycles run sessions normally.
+                result = resume_result
+                resume_result = None
+            else:
+                result = self._run_session(
+                    task,
+                    role="review",
+                    prompt=self._review_prompt(task),
+                    seq=task.review_cycle,
+                )
             advance(task, Phase.REVIEW_VERIFY)
             self._save()
             self._emit(
@@ -1835,23 +1938,31 @@ def _run_session(
         self.journal.set_active_log(task_id)
         self.journal.append("session-start", task_id=task_id, role=role, prompt=prompt)
         result = adapter.run(spec)
-        usage = adapter.read_usage(result)
         task.record_session(
             SessionRecord(
                 task_id=task_id,
                 role=role,
                 status=result.status,
                 session_id=result.session_id,
                 transcript_path=result.transcript_path,
-                usage=usage,
+                result_json=result.result_json,
             )
         )
+        # Make the completed session durable before the usage read, post-session
+        # hooks, and follow-up verification. If the host kills the process in
+        # that window, the resume path can see the session instead of a stale
+        # dev-running task with no evidence; usage stays best-effort metadata,
+        # not a durability gate.
+        self._save()
+        usage = adapter.read_usage(result)
+        task.attach_session_usage(task_id, usage)
         self.journal.append(
             "session-end",
             task_id=task_id,
             status=result.status,
             tokens=usage.total if usage else None,
         )
+        self._save()
         self._emit(
             "post_session",
             task,

src/bmad_loop/model.py

@@ -91,6 +91,9 @@ class SessionRecord:
     session_id: str | None = None
     transcript_path: str | None = None
     usage: TokenUsage | None = None
+    # the session's parsed result payload, persisted so a durably-saved
+    # completed session is actionable on resume, not just forensics
+    result_json: dict[str, Any] | None = None
 
     def to_dict(self) -> dict[str, Any]:
         return {
@@ -100,6 +103,7 @@ def to_dict(self) -> dict[str, Any]:
             "session_id": self.session_id,
             "transcript_path": self.transcript_path,
             "usage": self.usage.to_dict() if self.usage else None,
+            "result_json": self.result_json,
         }
 
     @classmethod
@@ -112,6 +116,7 @@ def from_dict(cls, d: dict[str, Any]) -> "SessionRecord":
             session_id=d.get("session_id"),
             transcript_path=d.get("transcript_path"),
             usage=TokenUsage.from_dict(usage) if usage else None,
+            result_json=d.get("result_json"),
         )
 
 
@@ -168,6 +173,21 @@ def record_session(self, record: SessionRecord) -> None:
         if record.usage:
             self.tokens.add(record.usage)
 
+    def attach_session_usage(self, task_id: str, usage: TokenUsage | None) -> None:
+        """Fold usage into the most recent session for `task_id`. Usage is
+        best-effort metadata attached after the session itself is saved, so a
+        failed usage read never costs the recorded session."""
+        if usage is None:
+            return
+        for record in reversed(self.sessions):
+            if record.task_id != task_id:
+                continue
+            if record.usage is None:
+                record.usage = usage
+                self.tokens.add(usage)
+            return
+        raise KeyError(task_id)
+
     def to_dict(self) -> dict[str, Any]:
         return {
             "story_key": self.story_key,