From 7e18dccbd981a9f2a900367fe2bb7b133edfc49d Mon Sep 17 00:00:00 2001 From: Jason Walker <913443@dadeschools.net> Date: Fri, 24 Jul 2026 08:06:45 -0400 Subject: [PATCH] feat(control-plane): durable MCP session checkpoint schema (Closes #660) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a versioned, redacted, reconcile-on-boot session checkpoint store to control_plane_db so a restart can recover session identity, stage, lease ownership, and next valid action instead of forcing human reconstruction (umbrella #655; #628 autonomous-handoff goal). - Bump SCHEMA_VERSION 4->5. New `session_checkpoints` table + two indexes, added via `CREATE TABLE IF NOT EXISTS` so table creation is itself the additive, idempotent v4->v5 migration (dependency_edges precedent). - Writer `write_session_checkpoint` upserts the current recoverable state per (remote, org, repo, session_id, work_kind, work_number); stage transitions audit to `events`. Readers `get_session_checkpoint` / `list_session_checkpoints`. - Every free-text and JSON field is passed through `gitea_audit.redact` before storage — no tokens/credential URLs can land in a checkpoint (AC4). - `reconcile_session_checkpoint` is pure and never restores: it diagnoses a stored checkpoint against live head/lease state and flags staleness (AC3). - Drain gate: `require_complete=True` fails closed (writes nothing) when a checkpoint is missing a drain-required field, so drain cannot complete on an unrecoverable record. - `lease_id`/`assignment_id` are soft references (no enforced FK) so a checkpoint survives deletion of the lease it names; `work_number=0` is the NULL-safe session-level sentinel. Tests: 13 new cases (schema/version, multi-role fixtures, upsert+audit, JSON round-trip, secret redaction, reconcile stale-head/dead-lease/ reassigned-lease/clean/unknown, drain fail-closed, sentinel key). Existing schema_version assertion updated 4->5. Full tests/test_control_plane_db.py suite: 33/33 pass. Links #652 #653 #655. Co-Authored-By: Claude Opus 4.8 (1M context) --- control_plane_db.py | 429 ++++++++++++++++++++++++++++++++- tests/test_control_plane_db.py | 227 ++++++++++++++++- 2 files changed, 654 insertions(+), 2 deletions(-) diff --git a/control_plane_db.py b/control_plane_db.py index 4616cb3..cec67e4 100644 --- a/control_plane_db.py +++ b/control_plane_db.py @@ -30,8 +30,9 @@ from datetime import datetime, timedelta, timezone from typing import Any, Iterator, Sequence import dependency_graph +import gitea_audit -SCHEMA_VERSION = 4 +SCHEMA_VERSION = 5 # Assignable work kinds only — raw monitoring incidents are never work items. WORK_KINDS = frozenset({"issue", "pr"}) @@ -177,6 +178,51 @@ CREATE TABLE IF NOT EXISTS dependency_edges ( ) ); +-- Durable MCP session checkpoints (#660, umbrella #655 child; #628 handoff +-- goal). A restart otherwise loses session identity, stage, lease ownership, +-- and next action, forcing human reconstruction. Each row is the *current* +-- recoverable state of one session's work on one work unit; stage transitions +-- upsert the row and audit the change to ``events``. Creating the table is the +-- v4->v5 migration: additive, idempotent, and it never touches prior tables. +-- +-- ``lease_id`` / ``assignment_id`` are recorded as soft references (plain TEXT, +-- no enforced FK) exactly as dependency_edges references issues/PRs by number: +-- a checkpoint is a recovery artifact that must survive the deletion of the +-- lease it names, so a hard FK would fail the pre-drain write the issue +-- requires to succeed. ``work_number`` uses 0 (never a real issue/PR number) +-- as the "session-level, no specific work" sentinel so UNIQUE is NULL-safe. +-- Every free-text / JSON field is redaction-filtered before storage (AC4). +CREATE TABLE IF NOT EXISTS session_checkpoints ( + checkpoint_id TEXT PRIMARY KEY, + remote TEXT NOT NULL, + org TEXT NOT NULL, + repo TEXT NOT NULL, + session_id TEXT NOT NULL, + provider_identity TEXT NOT NULL DEFAULT '', + role TEXT NOT NULL DEFAULT '', + work_kind TEXT NOT NULL DEFAULT '' CHECK (work_kind IN ('issue', 'pr', '')), + work_number INTEGER NOT NULL DEFAULT 0, + worktree_path TEXT NOT NULL DEFAULT '', + branch TEXT NOT NULL DEFAULT '', + head_sha TEXT NOT NULL DEFAULT '', + capabilities TEXT NOT NULL DEFAULT '[]', + lease_id TEXT NOT NULL DEFAULT '', + assignment_id TEXT NOT NULL DEFAULT '', + workflow_stage TEXT NOT NULL DEFAULT '', + last_completed_action TEXT NOT NULL DEFAULT '', + current_operation TEXT NOT NULL DEFAULT '', + pending_mutation TEXT NOT NULL DEFAULT '', + evidence TEXT NOT NULL DEFAULT '{}', + blocker TEXT NOT NULL DEFAULT '', + next_valid_action TEXT NOT NULL DEFAULT '', + recovery_instructions TEXT NOT NULL DEFAULT '', + checkpoint_schema_version INTEGER NOT NULL DEFAULT 5, + status TEXT NOT NULL DEFAULT 'active', + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + UNIQUE (remote, org, repo, session_id, work_kind, work_number) +); + CREATE INDEX IF NOT EXISTS idx_leases_work_status ON leases(work_item_id, status); -- Reverse lookup ("what waits on this target") is the query automatic -- resumption needs, so it gets its own index alongside the forward one. @@ -186,6 +232,12 @@ CREATE INDEX IF NOT EXISTS idx_dependency_edges_target ON dependency_edges(remote, org, repo, target_kind, target_number); CREATE INDEX IF NOT EXISTS idx_assignments_session ON assignments(session_id, status); CREATE INDEX IF NOT EXISTS idx_incident_gitea ON incident_links(gitea_org, gitea_repo, gitea_issue_number); +-- Resume queries hit by session (what was this session doing) and by work unit +-- (who was checkpointed on this issue/PR), so both get an index. +CREATE INDEX IF NOT EXISTS idx_session_checkpoints_session + ON session_checkpoints(remote, org, repo, session_id); +CREATE INDEX IF NOT EXISTS idx_session_checkpoints_work + ON session_checkpoints(remote, org, repo, work_kind, work_number); """ @@ -2368,3 +2420,378 @@ class ControlPlaneDB: edge["prior_state"] = prior_state edge["state_changed"] = prior_state != state_norm return edge + + # ------------------------------------------------------------------ + # Session checkpoints (#660) — durable, redacted, reconcile-on-boot. + # ------------------------------------------------------------------ + + # Free-text columns that carry human/agent-authored strings. Each is + # redaction-filtered before storage so an accidentally pasted token or + # credential URL can never land in a checkpoint (#660 AC4). + _CHECKPOINT_TEXT_FIELDS: tuple[str, ...] = ( + "provider_identity", + "role", + "worktree_path", + "branch", + "head_sha", + "lease_id", + "assignment_id", + "workflow_stage", + "last_completed_action", + "current_operation", + "blocker", + "next_valid_action", + "recovery_instructions", + "status", + ) + + # JSON columns whose *decoded* structure is redacted recursively. + _CHECKPOINT_JSON_FIELDS: tuple[str, ...] = ( + "capabilities", + "pending_mutation", + "evidence", + ) + + # Minimum a checkpoint must carry to be a usable recovery record. The drain + # gate writes with ``require_complete=True``; an incomplete write fails + # closed so drain cannot complete on a checkpoint that can't resume (#660 + # "fail closed if checkpoint incomplete during drain"). + REQUIRED_CHECKPOINT_FIELDS_FOR_DRAIN: tuple[str, ...] = ( + "session_id", + "role", + "workflow_stage", + "next_valid_action", + "recovery_instructions", + ) + + @classmethod + def checkpoint_completeness(cls, record: dict[str, Any]) -> list[str]: + """Return the drain-required fields that are missing/blank in *record*. + + Pure (no DB). An empty list means the record is drain-complete. + """ + missing: list[str] = [] + for field in cls.REQUIRED_CHECKPOINT_FIELDS_FOR_DRAIN: + if not str(record.get(field) or "").strip(): + missing.append(field) + return missing + + def write_session_checkpoint( + self, + *, + remote: str, + org: str, + repo: str, + session_id: str, + provider_identity: str = "", + role: str = "", + work_kind: str | None = None, + work_number: int | None = None, + worktree_path: str = "", + branch: str = "", + head_sha: str = "", + capabilities: Any = None, + lease_id: str = "", + assignment_id: str = "", + workflow_stage: str = "", + last_completed_action: str = "", + current_operation: str = "", + pending_mutation: Any = None, + evidence: Any = None, + blocker: str = "", + next_valid_action: str = "", + recovery_instructions: str = "", + status: str = "active", + require_complete: bool = False, + ) -> dict[str, Any]: + """Upsert the current checkpoint for one session's work on one unit. + + Keyed by ``(remote, org, repo, session_id, work_kind, work_number)`` so + a stage transition refreshes the single current row (transitions are + audited to ``events``) rather than appending unbounded history — the + row is *resumable state*, matching the dependency-edge current-state + model. + + Every string / JSON field is redacted via ``gitea_audit.redact`` before + it touches the DB. With ``require_complete=True`` (the drain gate path) + a record missing any drain-required field raises ``ControlPlaneError`` + and writes nothing, so drain cannot proceed on an unrecoverable record. + """ + if not str(session_id or "").strip(): + raise ControlPlaneError("session_id is required for a checkpoint (fail closed)") + + # Normalize the work-unit key. Absent/blank work collapses to the + # ('', 0) session-level sentinel so UNIQUE stays NULL-safe. + if work_kind and str(work_kind).strip(): + kind_norm = dependency_graph.normalize_work_kind(work_kind) + number_norm = int(work_number) if work_number is not None else 0 + else: + kind_norm = "" + number_norm = 0 + + # Assemble, then redact the whole record in one pass. + raw_record: dict[str, Any] = { + "provider_identity": provider_identity or "", + "role": role or "", + "worktree_path": worktree_path or "", + "branch": branch or "", + "head_sha": head_sha or "", + "lease_id": lease_id or "", + "assignment_id": assignment_id or "", + "workflow_stage": workflow_stage or "", + "last_completed_action": last_completed_action or "", + "current_operation": current_operation or "", + "blocker": blocker or "", + "next_valid_action": next_valid_action or "", + "recovery_instructions": recovery_instructions or "", + "status": (status or "active"), + "session_id": session_id, + "capabilities": capabilities if capabilities is not None else [], + "pending_mutation": pending_mutation if pending_mutation is not None else {}, + "evidence": evidence if evidence is not None else {}, + } + clean = gitea_audit.redact(raw_record) + + if require_complete: + missing = self.checkpoint_completeness(clean) + if missing: + raise ControlPlaneError( + "checkpoint incomplete for drain; missing " + f"{', '.join(missing)} (fail closed)" + ) + + capabilities_json = json.dumps(clean.get("capabilities") or []) + pending_json = json.dumps(clean.get("pending_mutation") or {}) + evidence_json = json.dumps(clean.get("evidence") or {}) + now_s = _ts() + + with self._tx() as conn: + existing = conn.execute( + """ + SELECT * FROM session_checkpoints + WHERE remote = ? AND org = ? AND repo = ? + AND session_id = ? AND work_kind = ? AND work_number = ? + """, + (remote, org, repo, session_id, kind_norm, number_norm), + ).fetchone() + + if existing is None: + checkpoint_id = uuid.uuid4().hex + conn.execute( + """ + INSERT INTO session_checkpoints( + checkpoint_id, remote, org, repo, session_id, + provider_identity, role, work_kind, work_number, + worktree_path, branch, head_sha, capabilities, + lease_id, assignment_id, workflow_stage, + last_completed_action, current_operation, + pending_mutation, evidence, blocker, next_valid_action, + recovery_instructions, checkpoint_schema_version, + status, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + checkpoint_id, remote, org, repo, session_id, + clean["provider_identity"], clean["role"], kind_norm, + number_norm, clean["worktree_path"], clean["branch"], + clean["head_sha"], capabilities_json, clean["lease_id"], + clean["assignment_id"], clean["workflow_stage"], + clean["last_completed_action"], clean["current_operation"], + pending_json, evidence_json, clean["blocker"], + clean["next_valid_action"], clean["recovery_instructions"], + SCHEMA_VERSION, clean["status"], now_s, now_s, + ), + ) + else: + checkpoint_id = str(existing["checkpoint_id"]) + conn.execute( + """ + UPDATE session_checkpoints + SET provider_identity = ?, role = ?, worktree_path = ?, + branch = ?, head_sha = ?, capabilities = ?, + lease_id = ?, assignment_id = ?, workflow_stage = ?, + last_completed_action = ?, current_operation = ?, + pending_mutation = ?, evidence = ?, blocker = ?, + next_valid_action = ?, recovery_instructions = ?, + checkpoint_schema_version = ?, status = ?, + updated_at = ? + WHERE checkpoint_id = ? + """, + ( + clean["provider_identity"], clean["role"], + clean["worktree_path"], clean["branch"], clean["head_sha"], + capabilities_json, clean["lease_id"], clean["assignment_id"], + clean["workflow_stage"], clean["last_completed_action"], + clean["current_operation"], pending_json, evidence_json, + clean["blocker"], clean["next_valid_action"], + clean["recovery_instructions"], SCHEMA_VERSION, + clean["status"], now_s, checkpoint_id, + ), + ) + prior_stage = str(existing["workflow_stage"] or "") + new_stage = str(clean["workflow_stage"] or "") + if prior_stage != new_stage: + conn.execute( + """ + INSERT INTO events(work_item_id, event_type, message, created_at) + VALUES (NULL, 'session_checkpoint_stage_change', ?, ?) + """, + ( + f"checkpoint {checkpoint_id} session {session_id} " + f"stage {prior_stage or '(none)'} -> " + f"{new_stage or '(none)'}", + now_s, + ), + ) + + row = conn.execute( + "SELECT * FROM session_checkpoints WHERE checkpoint_id = ?", + (checkpoint_id,), + ).fetchone() + return self._session_checkpoint_row(row) or {} + + @staticmethod + def _session_checkpoint_row(row: sqlite3.Row | None) -> dict[str, Any] | None: + """Return a stored checkpoint as a plain dict with JSON fields decoded.""" + if row is None: + return None + record = dict(row) + for field, empty in ( + ("capabilities", []), + ("pending_mutation", {}), + ("evidence", {}), + ): + raw = record.get(field) + try: + record[field] = json.loads(raw) if raw else empty + except (TypeError, ValueError): + # A row written by an older/foreign writer must not break reads. + record[field] = {"unparsed": str(raw)} + return record + + def get_session_checkpoint( + self, + *, + remote: str, + org: str, + repo: str, + session_id: str, + work_kind: str | None = None, + work_number: int | None = None, + ) -> dict[str, Any] | None: + """Return the current checkpoint for one session's work unit, or None.""" + if work_kind and str(work_kind).strip(): + kind_norm = dependency_graph.normalize_work_kind(work_kind) + number_norm = int(work_number) if work_number is not None else 0 + else: + kind_norm = "" + number_norm = 0 + with self._tx(immediate=False) as conn: + row = conn.execute( + """ + SELECT * FROM session_checkpoints + WHERE remote = ? AND org = ? AND repo = ? + AND session_id = ? AND work_kind = ? AND work_number = ? + """, + (remote, org, repo, session_id, kind_norm, number_norm), + ).fetchone() + return self._session_checkpoint_row(row) + + def list_session_checkpoints( + self, + *, + remote: str | None = None, + org: str | None = None, + repo: str | None = None, + session_id: str | None = None, + work_kind: str | None = None, + work_number: int | None = None, + status: str | None = None, + limit: int = 500, + ) -> list[dict[str, Any]]: + """Return stored checkpoints, filtered. Newest updated first.""" + clauses: list[str] = [] + params: list[Any] = [] + if remote: + clauses.append("remote = ?") + params.append(remote) + if org: + clauses.append("org = ?") + params.append(org) + if repo: + clauses.append("repo = ?") + params.append(repo) + if session_id: + clauses.append("session_id = ?") + params.append(session_id) + if work_kind: + clauses.append("work_kind = ?") + params.append(dependency_graph.normalize_work_kind(work_kind)) + if work_number is not None: + clauses.append("work_number = ?") + params.append(int(work_number)) + if status: + clauses.append("status = ?") + params.append(status) + + sql = "SELECT * FROM session_checkpoints" + if clauses: + sql += " WHERE " + " AND ".join(clauses) + sql += " ORDER BY updated_at DESC LIMIT ?" + params.append(int(limit)) + + with self._tx(immediate=False) as conn: + rows = conn.execute(sql, params).fetchall() + return [ + record + for record in (self._session_checkpoint_row(r) for r in rows) + if record + ] + + @staticmethod + def reconcile_session_checkpoint( + checkpoint: dict[str, Any], + *, + live_head_sha: str | None = None, + live_lease_active: bool | None = None, + live_lease_id: str | None = None, + ) -> dict[str, Any]: + """Diagnose a stored checkpoint against live Git/Gitea/lease state. + + Pure (no DB) and **never restores** — it returns a diagnosis and the + caller decides. A head that no longer matches, or a lease that is gone + or reassigned, marks the checkpoint stale so resumption reconciles + instead of blindly restoring stale ownership (#660 AC3). + + A ``None`` live input means "unknown, not checked" and never flags a + mismatch on its own. + """ + stored_head = str(checkpoint.get("head_sha") or "") + stored_lease = str(checkpoint.get("lease_id") or "") + + head_mismatch = bool( + live_head_sha is not None + and stored_head + and stored_head != str(live_head_sha) + ) + lease_mismatch = False + if stored_lease: + if live_lease_active is False: + lease_mismatch = True + elif live_lease_id is not None and str(live_lease_id) != stored_lease: + lease_mismatch = True + + stale = head_mismatch or lease_mismatch + return { + "checkpoint_id": checkpoint.get("checkpoint_id"), + "session_id": checkpoint.get("session_id"), + "stale": stale, + "head_mismatch": head_mismatch, + "lease_mismatch": lease_mismatch, + "stored_head_sha": stored_head, + "live_head_sha": None if live_head_sha is None else str(live_head_sha), + "stored_lease_id": stored_lease, + "live_lease_id": None if live_lease_id is None else str(live_lease_id), + "reconcile_action": "reconcile_required" if stale else "safe_to_resume", + } diff --git a/tests/test_control_plane_db.py b/tests/test_control_plane_db.py index 8f251eb..6065f38 100644 --- a/tests/test_control_plane_db.py +++ b/tests/test_control_plane_db.py @@ -11,6 +11,7 @@ from datetime import timedelta from control_plane_db import ( ControlPlaneDB, + ControlPlaneError, InvalidWorkKindError, LeaseRequiredError, WORK_KINDS, @@ -36,7 +37,7 @@ class ControlPlaneDBTest(unittest.TestCase): rows = dict(conn.execute("SELECT key, value FROM schema_meta").fetchall()) finally: conn.close() - self.assertEqual(rows["schema_version"], "4") + self.assertEqual(rows["schema_version"], "5") self.assertIn("DB coordinates", rows["architecture"]) self.assertIn("bridge", rows["architecture"].lower()) @@ -820,5 +821,229 @@ class ControlPlaneDBTest(unittest.TestCase): self.assertEqual(n, 1) +class SessionCheckpointTest(unittest.TestCase): + """Durable MCP session checkpoint schema, redaction, and reconcile (#660).""" + + def setUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self._tmp.name, "cp.sqlite3") + self.db = ControlPlaneDB(self.db_path) + + def tearDown(self) -> None: + self._tmp.cleanup() + + def _write(self, **overrides): + base = dict( + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + session_id="prgs-author-1-abc", + role="author", + work_kind="issue", + work_number=660, + branch="feat/issue-660-session-checkpoint-schema", + head_sha="deadbeef", + lease_id="lease-1", + workflow_stage="implementing", + last_completed_action="wrote schema", + next_valid_action="write tests", + recovery_instructions="re-lock #660 then continue tests", + ) + base.update(overrides) + return self.db.write_session_checkpoint(**base) + + # AC1 — schema documented and versioned. + def test_table_exists_and_row_carries_schema_version(self) -> None: + import sqlite3 + + conn = sqlite3.connect(self.db_path) + try: + names = { + r[0] + for r in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table'" + ).fetchall() + } + finally: + conn.close() + self.assertIn("session_checkpoints", names) + record = self._write() + self.assertEqual(record["checkpoint_schema_version"], 5) + + # AC2 — checkpoints written for multi-role session fixtures. + def test_multi_role_fixtures_each_get_a_row(self) -> None: + roles = [ + ("prgs-author-1", "author", "issue", 660), + ("prgs-reviewer-2", "reviewer", "pr", 795), + ("prgs-merger-3", "merger", "pr", 862), + ("prgs-controller-4", "controller", "issue", 653), + ] + for session_id, role, kind, number in roles: + self._write( + session_id=session_id, + role=role, + work_kind=kind, + work_number=number, + lease_id=f"lease-{session_id}", + ) + rows = self.db.list_session_checkpoints(remote="prgs") + self.assertEqual(len(rows), 4) + self.assertEqual( + {r["role"] for r in rows}, + {"author", "reviewer", "merger", "controller"}, + ) + + def test_upsert_is_current_state_and_audits_stage_change(self) -> None: + first = self._write(workflow_stage="implementing") + second = self._write(workflow_stage="testing") + self.assertEqual(first["checkpoint_id"], second["checkpoint_id"]) + rows = self.db.list_session_checkpoints( + remote="prgs", session_id="prgs-author-1-abc" + ) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0]["workflow_stage"], "testing") + + import sqlite3 + + conn = sqlite3.connect(self.db_path) + try: + n = conn.execute( + "SELECT COUNT(*) FROM events " + "WHERE event_type = 'session_checkpoint_stage_change'" + ).fetchone()[0] + finally: + conn.close() + self.assertEqual(n, 1) + + def test_get_and_roundtrip_json_fields(self) -> None: + self._write( + capabilities=["gitea.repo.commit", "gitea.pr.create"], + evidence={"tests": "4 passing"}, + pending_mutation={"op": "commit_files", "files": ["control_plane_db.py"]}, + ) + got = self.db.get_session_checkpoint( + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + session_id="prgs-author-1-abc", + work_kind="issue", + work_number=660, + ) + self.assertIsNotNone(got) + self.assertEqual(got["capabilities"], ["gitea.repo.commit", "gitea.pr.create"]) + self.assertEqual(got["evidence"], {"tests": "4 passing"}) + self.assertEqual(got["pending_mutation"]["op"], "commit_files") + + # AC4 — no secrets in stored records. + def test_secrets_are_redacted_before_storage(self) -> None: + self._write( + recovery_instructions=( + "resume with Authorization: Bearer sk-supersecrettoken then retry" + ), + evidence={"authorization": "Bearer sk-anothersecret"}, + pending_mutation={"url": "https://user:password@gitea.prgs.cc/repo.git"}, + ) + import sqlite3 + + conn = sqlite3.connect(self.db_path) + try: + row = conn.execute( + "SELECT recovery_instructions, evidence, pending_mutation " + "FROM session_checkpoints" + ).fetchone() + finally: + conn.close() + blob = " ".join(str(v) for v in row) + self.assertNotIn("sk-supersecrettoken", blob) + self.assertNotIn("sk-anothersecret", blob) + self.assertNotIn("password", blob) + self.assertIn("REDACTED", blob) + + # AC3 — reconcile detects stale head / lease mismatch. + def test_reconcile_flags_stale_head(self) -> None: + record = self._write(head_sha="aaaa1111") + result = self.db.reconcile_session_checkpoint( + record, live_head_sha="bbbb2222", live_lease_active=True, + live_lease_id="lease-1", + ) + self.assertTrue(result["stale"]) + self.assertTrue(result["head_mismatch"]) + self.assertFalse(result["lease_mismatch"]) + self.assertEqual(result["reconcile_action"], "reconcile_required") + + def test_reconcile_flags_dead_lease(self) -> None: + record = self._write(lease_id="lease-1", head_sha="aaaa1111") + result = self.db.reconcile_session_checkpoint( + record, live_head_sha="aaaa1111", live_lease_active=False, + ) + self.assertTrue(result["stale"]) + self.assertFalse(result["head_mismatch"]) + self.assertTrue(result["lease_mismatch"]) + + def test_reconcile_reassigned_lease_is_stale(self) -> None: + record = self._write(lease_id="lease-1") + result = self.db.reconcile_session_checkpoint( + record, live_lease_active=True, live_lease_id="lease-999", + ) + self.assertTrue(result["lease_mismatch"]) + + def test_reconcile_clean_state_is_safe_to_resume(self) -> None: + record = self._write(head_sha="aaaa1111", lease_id="lease-1") + result = self.db.reconcile_session_checkpoint( + record, live_head_sha="aaaa1111", live_lease_active=True, + live_lease_id="lease-1", + ) + self.assertFalse(result["stale"]) + self.assertEqual(result["reconcile_action"], "safe_to_resume") + + def test_unknown_live_state_never_flags_mismatch(self) -> None: + record = self._write(head_sha="aaaa1111", lease_id="lease-1") + result = self.db.reconcile_session_checkpoint(record) + self.assertFalse(result["stale"]) + + # Drain gate — fail closed when a checkpoint is incomplete. + def test_drain_requires_complete_checkpoint(self) -> None: + with self.assertRaises(ControlPlaneError): + self.db.write_session_checkpoint( + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + session_id="prgs-author-1-abc", + role="author", + workflow_stage="draining", + # next_valid_action + recovery_instructions intentionally absent + require_complete=True, + ) + # Nothing was written. + rows = self.db.list_session_checkpoints(remote="prgs") + self.assertEqual(rows, []) + + def test_drain_write_succeeds_when_complete(self) -> None: + record = self._write(require_complete=True) + self.assertEqual(record["status"], "active") + self.assertEqual(self.db.checkpoint_completeness(record), []) + + def test_missing_session_id_fails_closed(self) -> None: + with self.assertRaises(ControlPlaneError): + self.db.write_session_checkpoint( + remote="prgs", org="o", repo="r", session_id="", + ) + + def test_session_level_checkpoint_uses_sentinel_key(self) -> None: + # No work unit -> ('', 0) sentinel; a second session-level write upserts. + self.db.write_session_checkpoint( + remote="prgs", org="o", repo="r", session_id="s-sess", + workflow_stage="idle", + ) + self.db.write_session_checkpoint( + remote="prgs", org="o", repo="r", session_id="s-sess", + workflow_stage="booting", + ) + rows = self.db.list_session_checkpoints(remote="prgs", session_id="s-sess") + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0]["work_kind"], "") + self.assertEqual(rows[0]["work_number"], 0) + + if __name__ == "__main__": unittest.main()