1"""Durable state store (SQLite, WAL, explicit transactions, validated transitions)."""2from __future__ import annotations3import json, os, sqlite3, threading4from contextlib import contextmanager5from .contracts import RUN_STATES, RUN_TRANSITIONS, utcnow67SCHEMA = """8CREATE TABLE IF NOT EXISTS runs (9 run_id TEXT PRIMARY KEY, task_id TEXT NOT NULL, split TEXT NOT NULL, treatment TEXT NOT NULL,10 repeat_index INTEGER NOT NULL, block_id TEXT, order_in_block INTEGER, node TEXT,11 status TEXT NOT NULL, terminal_reason TEXT, manifest_json TEXT NOT NULL,12 created_utc TEXT NOT NULL, updated_utc TEXT NOT NULL, frozen INTEGER DEFAULT 013);14CREATE TABLE IF NOT EXISTS run_transitions (15 id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, from_status TEXT, to_status TEXT NOT NULL, utc TEXT NOT NULL, note TEXT16);17CREATE TABLE IF NOT EXISTS sessions (18 session_id TEXT PRIMARY KEY, run_id TEXT NOT NULL, role TEXT NOT NULL, kind TEXT NOT NULL,19 seq INTEGER NOT NULL, command_json TEXT NOT NULL, logged_before_submit_utc TEXT NOT NULL,20 started_utc TEXT, ended_utc TEXT, exit_code INTEGER, result_subtype TEXT, is_error INTEGER,21 total_cost_usd REAL, session_cost_usd REAL, num_turns INTEGER, duration_ms INTEGER,22 usage_json TEXT, model_usage_json TEXT, first_turn_input_tokens INTEGER, model_reported TEXT,23 client_version TEXT, raw_path TEXT, note TEXT24);25CREATE TABLE IF NOT EXISTS events (26 event_id TEXT PRIMARY KEY, run_id TEXT NOT NULL, session_id TEXT, seq INTEGER NOT NULL,27 received_utc TEXT NOT NULL, elapsed_s REAL, event_type TEXT NOT NULL, subtype TEXT, tool_name TEXT,28 status TEXT, summary TEXT, usage_json TEXT, error_class TEXT, payload_sha256 TEXT29);30CREATE TABLE IF NOT EXISTS gate_receipts (31 id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, cycle INTEGER NOT NULL, origin TEXT NOT NULL,32 utc TEXT NOT NULL, passed INTEGER NOT NULL, n_failed INTEGER, receipt_json TEXT NOT NULL, receipt_sha256 TEXT33);34CREATE TABLE IF NOT EXISTS reviews (35 id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, session_id TEXT, utc TEXT NOT NULL,36 verdict TEXT, n_issues INTEGER, review_json TEXT37);38CREATE TABLE IF NOT EXISTS evaluations (39 run_id TEXT PRIMARY KEY, utc TEXT NOT NULL, evaluator_version TEXT, scientific_success INTEGER,40 n_criteria INTEGER, n_criteria_passed INTEGER, n_violations INTEGER, violations_types TEXT,41 reproducibility INTEGER, quality_json TEXT, receipt_json TEXT NOT NULL, receipt_sha256 TEXT, contract_receipt_json TEXT42);43CREATE TABLE IF NOT EXISTS screenshots (44 screenshot_id TEXT PRIMARY KEY, run_id TEXT, task_id TEXT, treatment TEXT, phase TEXT, event_id TEXT,45 capture_timestamp_utc TEXT NOT NULL, capture_mode TEXT NOT NULL, capture_tool TEXT, source_surface TEXT,46 original_path TEXT, original_sha256 TEXT, redacted_public_path TEXT, redacted_sha256 TEXT,47 redaction_description TEXT, selection_rule TEXT, caption TEXT, claim_ids TEXT48);49CREATE TABLE IF NOT EXISTS interventions (50 intervention_id TEXT PRIMARY KEY, timestamp_utc TEXT, actor TEXT, phase_or_run TEXT, category TEXT,51 description TEXT, duration_minutes TEXT, measurement_basis TEXT, scientific_impact TEXT52);53CREATE TABLE IF NOT EXISTS faults (54 id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, fault_type TEXT NOT NULL, injected_utc TEXT NOT NULL,55 trigger TEXT, detail_json TEXT, recovered INTEGER, recovery_json TEXT56);57CREATE TABLE IF NOT EXISTS jobs (58 job_key TEXT PRIMARY KEY, run_id TEXT NOT NULL, kind TEXT NOT NULL, logged_utc TEXT NOT NULL, status TEXT NOT NULL, detail TEXT59);60"""6162class StateStore:63 def __init__(self, path):64 self.path = path65 os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)66 self._lock = threading.RLock()67 self._conn = sqlite3.connect(path, check_same_thread=False, isolation_level=None, timeout=60)68 self._conn.execute("PRAGMA journal_mode=WAL")69 self._conn.execute("PRAGMA synchronous=FULL")70 self._conn.executescript(SCHEMA)7172 @contextmanager73 def tx(self):74 with self._lock:75 self._conn.execute("BEGIN IMMEDIATE")76 try:77 yield self._conn78 self._conn.execute("COMMIT")79 except Exception:80 self._conn.execute("ROLLBACK")81 raise8283 # ---- runs84 def create_run(self, manifest) -> None:85 with self.tx() as c:86 if c.execute("SELECT 1 FROM runs WHERE run_id=?", (manifest.run_id,)).fetchone():87 raise ValueError(f"run {manifest.run_id} already exists (idempotency guard)")88 now = utcnow()89 c.execute("INSERT INTO runs VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,0)",90 (manifest.run_id, manifest.task_id, manifest.split, manifest.treatment, manifest.repeat_index,91 manifest.block_id, manifest.order_in_block, manifest.node, manifest.status, manifest.terminal_reason,92 manifest.to_json(), now, now))93 c.execute("INSERT INTO run_transitions (run_id, from_status, to_status, utc, note) VALUES (?,?,?,?,?)",94 (manifest.run_id, None, manifest.status, now, "created"))9596 def get_run(self, run_id) -> dict:97 with self._lock:98 row = self._conn.execute("SELECT manifest_json, status, frozen FROM runs WHERE run_id=?", (run_id,)).fetchone()99 if not row:100 raise KeyError(run_id)101 m = json.loads(row[0]); m["status"] = row[1]; m["_frozen"] = bool(row[2]); return m102103 def update_manifest(self, run_id, **fields) -> dict:104 with self.tx() as c:105 row = c.execute("SELECT manifest_json, frozen FROM runs WHERE run_id=?", (run_id,)).fetchone()106 if row[1]:107 raise PermissionError(f"run {run_id} is frozen")108 m = json.loads(row[0]); m.update(fields)109 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))110 return m111112 def transition(self, run_id, to_status, note="", **fields) -> None:113 if to_status not in RUN_STATES:114 raise ValueError(to_status)115 with self.tx() as c:116 row = c.execute("SELECT status, manifest_json, frozen FROM runs WHERE run_id=?", (run_id,)).fetchone()117 if row is None:118 raise KeyError(run_id)119 if row[2]:120 raise PermissionError(f"run {run_id} is frozen")121 cur = row[0]122 if to_status not in RUN_TRANSITIONS[cur]:123 raise ValueError(f"illegal transition {cur} -> {to_status} for {run_id}")124 m = json.loads(row[1]); m.update(fields); m["status"] = to_status125 if to_status == "evaluated" and not fields.get("evaluation_receipt") and not m.get("evaluation_receipt"):126 raise ValueError("cannot mark evaluated without an evaluation receipt (invariant 1)")127 c.execute("UPDATE runs SET status=?, terminal_reason=?, manifest_json=?, updated_utc=? WHERE run_id=?",128 (to_status, m.get("terminal_reason"), json.dumps(m, indent=1, sort_keys=True), utcnow(), run_id))129 c.execute("INSERT INTO run_transitions (run_id, from_status, to_status, utc, note) VALUES (?,?,?,?,?)", (run_id, cur, to_status, utcnow(), note))130131 def freeze_split(self, split) -> int:132 with self.tx() as c:133 return c.execute("UPDATE runs SET frozen=1 WHERE split=? AND status='evaluated'", (split,)).rowcount134135 def list_runs(self, split=None, status=None) -> list:136 q = "SELECT manifest_json, status FROM runs WHERE 1=1"; a = []137 if split: q += " AND split=?"; a.append(split)138 if status: q += " AND status=?"; a.append(status)139 q += " ORDER BY block_id, order_in_block, run_id"140 out = []141 with self._lock:142 rows = self._conn.execute(q, a).fetchall()143 for mj, st in rows:144 m = json.loads(mj); m["status"] = st; out.append(m)145 return out146147 def n_error_attempts(self, run_id) -> int:148 with self._lock:149 return self._conn.execute("SELECT COUNT(*) FROM run_transitions WHERE run_id=? AND to_status='error'", (run_id,)).fetchone()[0]150151 # ---- jobs (idempotency ledger: logged before submission)152 def log_job(self, job_key, run_id, kind, detail="") -> bool:153 """Return True if newly logged, False if the job already exists (duplicate submission)."""154 with self.tx() as c:155 if c.execute("SELECT 1 FROM jobs WHERE job_key=?", (job_key,)).fetchone():156 return False157 c.execute("INSERT INTO jobs VALUES (?,?,?,?,?,?)", (job_key, run_id, kind, utcnow(), "submitted", detail))158 return True159160 def set_job(self, job_key, status, detail=None):161 with self.tx() as c:162 if detail is None:163 c.execute("UPDATE jobs SET status=? WHERE job_key=?", (status, job_key))164 else:165 c.execute("UPDATE jobs SET status=?, detail=? WHERE job_key=?", (status, detail, job_key))166167 # ---- sessions / events168 def log_session(self, session_id, run_id, role, kind, seq, command, raw_path, note=""):169 with self.tx() as c:170 c.execute("INSERT INTO sessions (session_id, run_id, role, kind, seq, command_json, logged_before_submit_utc, raw_path, note) VALUES (?,?,?,?,?,?,?,?,?)",171 (session_id, run_id, role, kind, seq, json.dumps(command), utcnow(), raw_path, note))172173 def finish_session(self, session_id, **fields):174 cols = ", ".join(f"{k}=?" for k in fields)175 with self.tx() as c:176 c.execute(f"UPDATE sessions SET {cols} WHERE session_id=?", (*fields.values(), session_id))177178 def add_event(self, ev: dict):179 with self.tx() as c:180 c.execute("INSERT OR IGNORE INTO events VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)",181 (ev["event_id"], ev["run_id"], ev.get("session_id"), ev["seq"], ev["received_utc"], ev.get("elapsed_s"),182 ev["event_type"], ev.get("subtype"), ev.get("tool_name"), ev.get("status"), ev.get("summary"),183 json.dumps(ev.get("usage")) if ev.get("usage") else None, ev.get("error_class"), ev.get("payload_sha256")))184185 def add_gate_receipt(self, run_id, cycle, origin, receipt):186 with self.tx() as c:187 c.execute("INSERT INTO gate_receipts (run_id, cycle, origin, utc, passed, n_failed, receipt_json, receipt_sha256) VALUES (?,?,?,?,?,?,?,?)",188 (run_id, cycle, origin, utcnow(), int(bool(receipt.get("passed"))), receipt.get("n_failed"), json.dumps(receipt), receipt.get("receipt_sha256")))189190 def add_review(self, run_id, session_id, review):191 with self.tx() as c:192 c.execute("INSERT INTO reviews (run_id, session_id, utc, verdict, n_issues, review_json) VALUES (?,?,?,?,?,?)",193 (run_id, session_id, utcnow(), (review or {}).get("verdict"), len((review or {}).get("issues") or []), json.dumps(review)))194195 def add_evaluation(self, run_id, receipt, contract_receipt=None):196 q = receipt.get("quality_components", {})197 with self.tx() as c:198 c.execute("INSERT OR REPLACE INTO evaluations VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",199 (run_id, utcnow(), receipt.get("evaluator_version"), int(bool(receipt.get("scientific_success"))),200 len(receipt.get("required_criteria", [])), sum(1 for x in receipt.get("required_criteria", []) if x.get("passed")),201 len(receipt.get("critical_violations", [])), ",".join(sorted({v.get("type", "?") for v in receipt.get("critical_violations", [])})),202 None if q.get("reproducibility") is None else int(bool(q.get("reproducibility"))), json.dumps(q), json.dumps(receipt),203 receipt.get("receipt_sha256"), json.dumps(contract_receipt) if contract_receipt else None))204205 def add_screenshot(self, rec: dict):206 with self.tx() as c:207 c.execute("INSERT OR REPLACE INTO screenshots VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",208 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")))209210 def add_fault(self, run_id, fault_type, trigger, detail):211 with self.tx() as c:212 c.execute("INSERT INTO faults (run_id, fault_type, injected_utc, trigger, detail_json) VALUES (?,?,?,?,?)", (run_id, fault_type, utcnow(), trigger, json.dumps(detail)))213 return c.execute("SELECT last_insert_rowid()").fetchone()[0]214215 def set_fault_recovery(self, fault_id, recovered, recovery):216 with self.tx() as c:217 c.execute("UPDATE faults SET recovered=?, recovery_json=? WHERE id=?", (int(bool(recovered)), json.dumps(recovery), fault_id))218219 def query(self, sql, args=()):220 with self._lock:221 cur = self._conn.execute(sql, args)222 cols = [d[0] for d in cur.description] if cur.description else []223 return [dict(zip(cols, r)) for r in cur.fetchall()]224