refactor: split run state transitions

Author: rohans0509

src/codealmanac/services/runs/store.py

@@ -3,7 +3,7 @@
 from datetime import UTC, datetime, timedelta
 from pathlib import Path
 
-from codealmanac.core.errors import ConflictError, NotFoundError
+from codealmanac.core.errors import NotFoundError
 from codealmanac.core.paths import normalize_path
 from codealmanac.database.local import connect_local_database
 from codealmanac.services.harnesses.models import HarnessEvent, HarnessTranscriptRef
@@ -23,6 +23,13 @@
     RunStatus,
 )
 from codealmanac.services.runs.tables import RUN_TABLES
+from codealmanac.services.runs.transitions import (
+    attach_harness_transcript,
+    cancel_run,
+    finish_run,
+    queued_message,
+    start_run,
+)
 from codealmanac.services.runs.worker_locks import RunWorkerLockStore
 
 
@@ -188,20 +195,10 @@ def record_event(
 
     def mark_running(self, run_id: str) -> RunRecord:
         record = self.read(run_id)
-        if record.status == RunStatus.RUNNING:
-            return record
-        if record.status != RunStatus.QUEUED:
-            raise ConflictError(
-                f"run {run_id} cannot start from {record.status.value}"
-            )
         now = datetime.now(UTC)
-        running = record.model_copy(
-            update={
-                "status": RunStatus.RUNNING,
-                "updated_at": now,
-                "started_at": now,
-            }
-        )
+        running = start_run(record, now)
+        if running is record:
+            return record
         self.update_record_preserving_spec(running)
         self.event_store.append_status(run_id, now, RunStatus.RUNNING.value)
         return running
@@ -212,12 +209,7 @@ def record_harness_transcript(
         transcript: HarnessTranscriptRef,
     ) -> RunRecord:
         record = self.read(run_id)
-        updated = record.model_copy(
-            update={
-                "harness_transcript": transcript,
-                "updated_at": datetime.now(UTC),
-            }
-        )
+        updated = attach_harness_transcript(record, transcript, datetime.now(UTC))
         self.update_record_preserving_spec(updated)
         return updated
 
@@ -229,36 +221,21 @@ def finish(
         error: str | None,
     ) -> RunRecord:
         record = self.read(run_id)
-        if record.status == RunStatus.CANCELLED:
-            return record
         now = datetime.now(UTC)
-        finished = record.model_copy(
-            update={
-                "status": status,
-                "summary": summary,
-                "error": error,
-                "updated_at": now,
-                "finished_at": now,
-            }
-        )
+        finished = finish_run(record, status, summary, error, now)
+        if finished is record:
+            return record
         self.update_record_preserving_spec(finished)
         self.event_store.append_status(run_id, now, status.value)
         return finished
 
     def cancel(self, run_id: str) -> RunCancelResult:
         record = self.read(run_id)
-        if record.status in TERMINAL_RUN_STATUSES:
-            return RunCancelResult(record=record, changed=False)
         now = datetime.now(UTC)
-        cancelled = record.model_copy(
-            update={
-                "status": RunStatus.CANCELLED,
-                "updated_at": now,
-                "finished_at": now,
-                "summary": record.summary,
-                "error": record.error,
-            }
-        )
+        result = cancel_run(record, now)
+        if not result.changed:
+            return result
+        cancelled = result.record
         self.update_record_preserving_spec(cancelled)
         self.event_store.append_status(run_id, now, RunStatus.CANCELLED.value)
         return RunCancelResult(record=cancelled, changed=True)
@@ -307,7 +284,3 @@ def connect(self):
 
 def run_record_from_json(value: str) -> RunRecord:
     return RunRecord.model_validate_json(value)
-
-
-def queued_message(kind: RunKind) -> str:
-    return f"queued {kind.value}"

src/codealmanac/services/runs/transitions.py

@@ -0,0 +1,79 @@
+from datetime import datetime
+
+from codealmanac.core.errors import ConflictError
+from codealmanac.services.harnesses.models import HarnessTranscriptRef
+from codealmanac.services.runs.models import (
+    TERMINAL_RUN_STATUSES,
+    RunCancelResult,
+    RunKind,
+    RunRecord,
+    RunStatus,
+)
+
+
+def queued_message(kind: RunKind) -> str:
+    return f"queued {kind.value}"
+
+
+def start_run(record: RunRecord, now: datetime) -> RunRecord:
+    if record.status == RunStatus.RUNNING:
+        return record
+    if record.status != RunStatus.QUEUED:
+        raise ConflictError(
+            f"run {record.run_id} cannot start from {record.status.value}"
+        )
+    return record.model_copy(
+        update={
+            "status": RunStatus.RUNNING,
+            "updated_at": now,
+            "started_at": now,
+        }
+    )
+
+
+def attach_harness_transcript(
+    record: RunRecord,
+    transcript: HarnessTranscriptRef,
+    now: datetime,
+) -> RunRecord:
+    return record.model_copy(
+        update={
+            "harness_transcript": transcript,
+            "updated_at": now,
+        }
+    )
+
+
+def finish_run(
+    record: RunRecord,
+    status: RunStatus,
+    summary: str | None,
+    error: str | None,
+    now: datetime,
+) -> RunRecord:
+    if record.status == RunStatus.CANCELLED:
+        return record
+    return record.model_copy(
+        update={
+            "status": status,
+            "summary": summary,
+            "error": error,
+            "updated_at": now,
+            "finished_at": now,
+        }
+    )
+
+
+def cancel_run(record: RunRecord, now: datetime) -> RunCancelResult:
+    if record.status in TERMINAL_RUN_STATUSES:
+        return RunCancelResult(record=record, changed=False)
+    cancelled = record.model_copy(
+        update={
+            "status": RunStatus.CANCELLED,
+            "updated_at": now,
+            "finished_at": now,
+            "summary": record.summary,
+            "error": record.error,
+        }
+    )
+    return RunCancelResult(record=cancelled, changed=True)

tests/test_architecture.py

@@ -1808,6 +1808,7 @@ def test_run_persistence_stays_split_by_responsibility():
     service_text = (runs_root / "service.py").read_text(encoding="utf-8")
     streaming_text = (runs_root / "streaming.py").read_text(encoding="utf-8")
     tables_text = (runs_root / "tables.py").read_text(encoding="utf-8")
+    transitions_text = (runs_root / "transitions.py").read_text(encoding="utf-8")
     worker_locks_text = (runs_root / "worker_locks.py").read_text(encoding="utf-8")
     forbidden_store_fragments = (
         "write_json_atomically",
@@ -1828,6 +1829,7 @@ def test_run_persistence_stays_split_by_responsibility():
         "locks.py",
         "streaming.py",
         "tables.py",
+        "transitions.py",
         "worker_locks.py",
     } <= module_names
     assert len(store_text.splitlines()) <= 330
@@ -1837,6 +1839,10 @@ def test_run_persistence_stays_split_by_responsibility():
     assert "def new_run_record(" in factory_text
     assert "uuid4" in factory_text
     assert "strftime" in factory_text
+    assert "def start_run(" in transitions_text
+    assert "def finish_run(" in transitions_text
+    assert "def cancel_run(" in transitions_text
+    assert "ConflictError" in transitions_text
     assert "class RunEventStore" in events_text
     assert "run_events" in events_text
     assert "RunWorkerLease" in locks_text