"""Durable state store (SQLite, WAL, explicit transactions, validated transitions).""" from __future__ import annotations import json, os, sqlite3, threading from contextlib import contextmanager from .contracts import RUN_STATES, RUN_TRANSITIONS, utcnow SCHEMA = """ CREATE TABLE IF NOT EXISTS runs ( run_id TEXT PRIMARY KEY, task_id TEXT NOT NULL, split TEXT NOT NULL, treatment TEXT NOT NULL, repeat_index INTEGER NOT NULL, block_id TEXT, order_in_block INTEGER, node TEXT, status TEXT NOT NULL, terminal_reason TEXT, manifest_json TEXT NOT NULL, created_utc TEXT NOT NULL, updated_utc TEXT NOT NULL, frozen INTEGER DEFAULT 0 ); CREATE TABLE IF NOT EXISTS run_transitions ( id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, from_status TEXT, to_status TEXT NOT NULL, utc TEXT NOT NULL, note TEXT ); CREATE TABLE IF NOT EXISTS sessions ( session_id TEXT PRIMARY KEY, run_id TEXT NOT NULL, role TEXT NOT NULL, kind TEXT NOT NULL, seq INTEGER NOT NULL, command_json TEXT NOT NULL, logged_before_submit_utc TEXT NOT NULL, started_utc TEXT, ended_utc TEXT, exit_code INTEGER, result_subtype TEXT, is_error INTEGER, total_cost_usd REAL, session_cost_usd REAL, num_turns INTEGER, duration_ms INTEGER, usage_json TEXT, model_usage_json TEXT, first_turn_input_tokens INTEGER, model_reported TEXT, client_version TEXT, raw_path TEXT, note TEXT ); CREATE TABLE IF NOT EXISTS events ( event_id TEXT PRIMARY KEY, run_id TEXT NOT NULL, session_id TEXT, seq INTEGER NOT NULL, received_utc TEXT NOT NULL, elapsed_s REAL, event_type TEXT NOT NULL, subtype TEXT, tool_name TEXT, status TEXT, summary TEXT, usage_json TEXT, error_class TEXT, payload_sha256 TEXT ); CREATE TABLE IF NOT EXISTS gate_receipts ( id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, cycle INTEGER NOT NULL, origin TEXT NOT NULL, utc TEXT NOT NULL, passed INTEGER NOT NULL, n_failed INTEGER, receipt_json TEXT NOT NULL, receipt_sha256 TEXT ); CREATE TABLE IF NOT EXISTS reviews ( id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, session_id TEXT, utc TEXT NOT NULL, verdict TEXT, n_issues INTEGER, review_json TEXT ); CREATE TABLE IF NOT EXISTS evaluations ( run_id TEXT PRIMARY KEY, utc TEXT NOT NULL, evaluator_version TEXT, scientific_success INTEGER, n_criteria INTEGER, n_criteria_passed INTEGER, n_violations INTEGER, violations_types TEXT, reproducibility INTEGER, quality_json TEXT, receipt_json TEXT NOT NULL, receipt_sha256 TEXT, contract_receipt_json TEXT ); CREATE TABLE IF NOT EXISTS screenshots ( screenshot_id TEXT PRIMARY KEY, run_id TEXT, task_id TEXT, treatment TEXT, phase TEXT, event_id TEXT, capture_timestamp_utc TEXT NOT NULL, capture_mode TEXT NOT NULL, capture_tool TEXT, source_surface TEXT, original_path TEXT, original_sha256 TEXT, redacted_public_path TEXT, redacted_sha256 TEXT, redaction_description TEXT, selection_rule TEXT, caption TEXT, claim_ids TEXT ); CREATE TABLE IF NOT EXISTS interventions ( intervention_id TEXT PRIMARY KEY, timestamp_utc TEXT, actor TEXT, phase_or_run TEXT, category TEXT, description TEXT, duration_minutes TEXT, measurement_basis TEXT, scientific_impact TEXT ); CREATE TABLE IF NOT EXISTS faults ( id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, fault_type TEXT NOT NULL, injected_utc TEXT NOT NULL, trigger TEXT, detail_json TEXT, recovered INTEGER, recovery_json TEXT ); CREATE TABLE IF NOT EXISTS jobs ( job_key TEXT PRIMARY KEY, run_id TEXT NOT NULL, kind TEXT NOT NULL, logged_utc TEXT NOT NULL, status TEXT NOT NULL, detail TEXT ); """ class StateStore: def __init__(self, path): self.path = path os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) self._lock = threading.RLock() self._conn = sqlite3.connect(path, check_same_thread=False, isolation_level=None, timeout=60) self._conn.execute("PRAGMA journal_mode=WAL") self._conn.execute("PRAGMA synchronous=FULL") self._conn.executescript(SCHEMA) @contextmanager def tx(self): with self._lock: self._conn.execute("BEGIN IMMEDIATE") try: yield self._conn self._conn.execute("COMMIT") except Exception: self._conn.execute("ROLLBACK") raise # ---- runs def create_run(self, manifest) -> None: with self.tx() as c: if c.execute("SELECT 1 FROM runs WHERE run_id=?", (manifest.run_id,)).fetchone(): raise ValueError(f"run {manifest.run_id} already exists (idempotency guard)") now = utcnow() c.execute("INSERT INTO runs VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,0)", (manifest.run_id, manifest.task_id, manifest.split, manifest.treatment, manifest.repeat_index, manifest.block_id, manifest.order_in_block, manifest.node, manifest.status, manifest.terminal_reason, manifest.to_json(), now, now)) c.execute("INSERT INTO run_transitions (run_id, from_status, to_status, utc, note) VALUES (?,?,?,?,?)", (manifest.run_id, None, manifest.status, now, "created")) def get_run(self, run_id) -> dict: with self._lock: row = self._conn.execute("SELECT manifest_json, status, frozen FROM runs WHERE run_id=?", (run_id,)).fetchone() if not row: raise KeyError(run_id) m = json.loads(row[0]); m["status"] = row[1]; m["_frozen"] = bool(row[2]); return m def update_manifest(self, run_id, **fields) -> dict: with self.tx() as c: row = c.execute("SELECT manifest_json, frozen FROM runs WHERE run_id=?", (run_id,)).fetchone() if row[1]: raise PermissionError(f"run {run_id} is frozen") m = json.loads(row[0]); m.update(fields) c.execute("UPDATE runs SET manifest_json=?, updated_utc=?, node=? WHERE run_id=?", (json.dumps(m, indent=1, sort_keys=True), utcnow(), m.get("node", ""), run_id)) return m def transition(self, run_id, to_status, note="", **fields) -> None: if to_status not in RUN_STATES: raise ValueError(to_status) with self.tx() as c: row = c.execute("SELECT status, manifest_json, frozen FROM runs WHERE run_id=?", (run_id,)).fetchone() if row is None: raise KeyError(run_id) if row[2]: raise PermissionError(f"run {run_id} is frozen") cur = row[0] if to_status not in RUN_TRANSITIONS[cur]: raise ValueError(f"illegal transition {cur} -> {to_status} for {run_id}") m = json.loads(row[1]); m.update(fields); m["status"] = to_status if to_status == "evaluated" and not fields.get("evaluation_receipt") and not m.get("evaluation_receipt"): raise ValueError("cannot mark evaluated without an evaluation receipt (invariant 1)") c.execute("UPDATE runs SET status=?, terminal_reason=?, manifest_json=?, updated_utc=? WHERE run_id=?", (to_status, m.get("terminal_reason"), json.dumps(m, indent=1, sort_keys=True), utcnow(), run_id)) c.execute("INSERT INTO run_transitions (run_id, from_status, to_status, utc, note) VALUES (?,?,?,?,?)", (run_id, cur, to_status, utcnow(), note)) def freeze_split(self, split) -> int: with self.tx() as c: return c.execute("UPDATE runs SET frozen=1 WHERE split=? AND status='evaluated'", (split,)).rowcount def list_runs(self, split=None, status=None) -> list: q = "SELECT manifest_json, status FROM runs WHERE 1=1"; a = [] if split: q += " AND split=?"; a.append(split) if status: q += " AND status=?"; a.append(status) q += " ORDER BY block_id, order_in_block, run_id" out = [] with self._lock: rows = self._conn.execute(q, a).fetchall() for mj, st in rows: m = json.loads(mj); m["status"] = st; out.append(m) return out def n_error_attempts(self, run_id) -> int: with self._lock: return self._conn.execute("SELECT COUNT(*) FROM run_transitions WHERE run_id=? AND to_status='error'", (run_id,)).fetchone()[0] # ---- jobs (idempotency ledger: logged before submission) def log_job(self, job_key, run_id, kind, detail="") -> bool: """Return True if newly logged, False if the job already exists (duplicate submission).""" with self.tx() as c: if c.execute("SELECT 1 FROM jobs WHERE job_key=?", (job_key,)).fetchone(): return False c.execute("INSERT INTO jobs VALUES (?,?,?,?,?,?)", (job_key, run_id, kind, utcnow(), "submitted", detail)) return True def set_job(self, job_key, status, detail=None): with self.tx() as c: if detail is None: c.execute("UPDATE jobs SET status=? WHERE job_key=?", (status, job_key)) else: c.execute("UPDATE jobs SET status=?, detail=? WHERE job_key=?", (status, detail, job_key)) # ---- sessions / events def log_session(self, session_id, run_id, role, kind, seq, command, raw_path, note=""): with self.tx() as c: c.execute("INSERT INTO sessions (session_id, run_id, role, kind, seq, command_json, logged_before_submit_utc, raw_path, note) VALUES (?,?,?,?,?,?,?,?,?)", (session_id, run_id, role, kind, seq, json.dumps(command), utcnow(), raw_path, note)) def finish_session(self, session_id, **fields): cols = ", ".join(f"{k}=?" for k in fields) with self.tx() as c: c.execute(f"UPDATE sessions SET {cols} WHERE session_id=?", (*fields.values(), session_id)) def add_event(self, ev: dict): with self.tx() as c: c.execute("INSERT OR IGNORE INTO events VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)", (ev["event_id"], ev["run_id"], ev.get("session_id"), ev["seq"], ev["received_utc"], ev.get("elapsed_s"), ev["event_type"], ev.get("subtype"), ev.get("tool_name"), ev.get("status"), ev.get("summary"), json.dumps(ev.get("usage")) if ev.get("usage") else None, ev.get("error_class"), ev.get("payload_sha256"))) def add_gate_receipt(self, run_id, cycle, origin, receipt): with self.tx() as c: c.execute("INSERT INTO gate_receipts (run_id, cycle, origin, utc, passed, n_failed, receipt_json, receipt_sha256) VALUES (?,?,?,?,?,?,?,?)", (run_id, cycle, origin, utcnow(), int(bool(receipt.get("passed"))), receipt.get("n_failed"), json.dumps(receipt), receipt.get("receipt_sha256"))) def add_review(self, run_id, session_id, review): with self.tx() as c: c.execute("INSERT INTO reviews (run_id, session_id, utc, verdict, n_issues, review_json) VALUES (?,?,?,?,?,?)", (run_id, session_id, utcnow(), (review or {}).get("verdict"), len((review or {}).get("issues") or []), json.dumps(review))) def add_evaluation(self, run_id, receipt, contract_receipt=None): q = receipt.get("quality_components", {}) with self.tx() as c: c.execute("INSERT OR REPLACE INTO evaluations VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)", (run_id, utcnow(), receipt.get("evaluator_version"), int(bool(receipt.get("scientific_success"))), len(receipt.get("required_criteria", [])), sum(1 for x in receipt.get("required_criteria", []) if x.get("passed")), len(receipt.get("critical_violations", [])), ",".join(sorted({v.get("type", "?") for v in receipt.get("critical_violations", [])})), None if q.get("reproducibility") is None else int(bool(q.get("reproducibility"))), json.dumps(q), json.dumps(receipt), receipt.get("receipt_sha256"), json.dumps(contract_receipt) if contract_receipt else None)) def add_screenshot(self, rec: dict): with self.tx() as c: c.execute("INSERT OR REPLACE INTO screenshots VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", tuple(rec.get(k) for k in ("screenshot_id", "run_id", "task_id", "treatment", "phase", "event_id", "capture_timestamp_utc", "capture_mode", "capture_tool", "source_surface", "original_path", "original_sha256", "redacted_public_path", "redacted_sha256", "redaction_description", "selection_rule", "caption", "claim_ids"))) def add_fault(self, run_id, fault_type, trigger, detail): with self.tx() as c: c.execute("INSERT INTO faults (run_id, fault_type, injected_utc, trigger, detail_json) VALUES (?,?,?,?,?)", (run_id, fault_type, utcnow(), trigger, json.dumps(detail))) return c.execute("SELECT last_insert_rowid()").fetchone()[0] def set_fault_recovery(self, fault_id, recovered, recovery): with self.tx() as c: c.execute("UPDATE faults SET recovered=?, recovery_json=? WHERE id=?", (int(bool(recovered)), json.dumps(recovery), fault_id)) def query(self, sql, args=()): with self._lock: cur = self._conn.execute(sql, args) cols = [d[0] for d in cur.description] if cur.description else [] return [dict(zip(cols, r)) for r in cur.fetchall()]