"""SQLite schema, DDL constants, and connection helpers (design §3.3.1, §6.7). This module is the single source of truth for the R720 agent-team durable SQL. It declares: * the ``pending_questions`` human-interaction lifecycle ledger (§3.3.1), * the ``budget_ledger`` shared Claude budget ledger (§6.1, §6.6), * a ``schema_meta`` version row driving :func:`migrate`. The companion ``schema.sql`` mirrors these statements verbatim for tooling. SQL DDL lives ONLY here. The LangGraph ``SqliteSaver`` checkpointer creates its OWN tables against this same connection / database file; the design reserves this DB for it but does not declare its tables. The atomic compare-and-set helpers used by the responder and the deadline timer (§3.3.1) take the write lock up front via ``BEGIN IMMEDIATE`` so concurrent responders are serialized — SQLite's default deferred isolation does not serialize a check-and-set. """ from __future__ import annotations import sqlite3 import time from datetime import datetime, timezone from pathlib import Path from typing import Any __all__ = [ "BUDGET_LEDGER_DDL", "PENDING_QUESTIONS_DDL", "PENDING_QUESTIONS_INDEXES_DDL", "BUDGET_LEDGER_INDEXES_DDL", "SCHEMA_META_DDL", "SCHEMA_VERSION", "QUESTION_STATES", "answer_question", "connect", "expire_question", "init_db", "migrate", "supersede_question", ] # Bump when the DDL below changes; migrate() steps a connection forward. SCHEMA_VERSION: int = 1 # Default SQLite busy timeout (ms) so concurrent writers wait for the write # lock rather than failing immediately. _BUSY_TIMEOUT_MS: int = 5000 # Allowed lifecycle states for a pending question (§3.3.1). Mirrors the DDL # CHECK constraint; exported so leaves can validate without re-listing them. QUESTION_STATES: tuple[str, ...] = ("open", "answered", "expired", "superseded") PENDING_QUESTIONS_DDL: str = """ CREATE TABLE IF NOT EXISTS pending_questions ( question_id TEXT PRIMARY KEY, thread_id TEXT NOT NULL, turn INTEGER NOT NULL, status TEXT NOT NULL CHECK (status IN ('open', 'answered', 'expired', 'superseded')), transport TEXT NOT NULL, channel_ref TEXT, posted_at TEXT, deadline_at TEXT, answer_json TEXT, answered_at TEXT, answered_via TEXT ) """.strip() PENDING_QUESTIONS_INDEXES_DDL: str = """ CREATE INDEX IF NOT EXISTS idx_pending_questions_thread ON pending_questions (thread_id, turn); CREATE INDEX IF NOT EXISTS idx_pending_questions_status ON pending_questions (status); """.strip() BUDGET_LEDGER_DDL: str = """ CREATE TABLE IF NOT EXISTS budget_ledger ( entry_id INTEGER PRIMARY KEY AUTOINCREMENT, thread_id TEXT, stage TEXT, model TEXT NOT NULL, billing_mode TEXT NOT NULL, input_tokens INTEGER NOT NULL DEFAULT 0, output_tokens INTEGER NOT NULL DEFAULT 0, usd_cost REAL NOT NULL DEFAULT 0.0, recorded_at TEXT NOT NULL, day_bucket TEXT NOT NULL ) """.strip() BUDGET_LEDGER_INDEXES_DDL: str = """ CREATE INDEX IF NOT EXISTS idx_budget_ledger_day ON budget_ledger (day_bucket); CREATE INDEX IF NOT EXISTS idx_budget_ledger_thread ON budget_ledger (thread_id); """.strip() SCHEMA_META_DDL: str = """ CREATE TABLE IF NOT EXISTS schema_meta ( id INTEGER PRIMARY KEY CHECK (id = 1), schema_version INTEGER NOT NULL ) """.strip() class _Connection(sqlite3.Connection): """``sqlite3.Connection`` subclass that can carry its backing file path. The base ``Connection`` has no ``__dict__``, so a path cannot be stashed on it. This thin subclass (passed as ``factory=`` to :func:`sqlite3.connect`) lets :func:`connect` record the db file for a thread-safe attribute lookup by the compare-and-set, avoiding a ``PRAGMA`` on a connection shared across threads. """ agent_team_db_path: str = "" def connect(db_path: Path) -> sqlite3.Connection: """Open ``db_path`` with WAL, foreign keys, and a busy timeout. WAL (``journal_mode=WAL``) lets the resume worker read while a responder writes; ``foreign_keys=ON`` enforces referential integrity; the busy timeout makes concurrent writers wait for the write lock instead of failing. ``isolation_level=None`` puts the connection in autocommit mode so the compare-and-set helpers can drive transactions explicitly with ``BEGIN IMMEDIATE`` (§3.3.1). """ db_path = Path(db_path) db_path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect( str(db_path), isolation_level=None, check_same_thread=False, factory=_Connection, ) conn.row_factory = sqlite3.Row # Set the busy timeout FIRST so every subsequent statement — including the # journal-mode pragma below, which briefly needs the write lock — waits for # the lock instead of failing immediately when another connection is mid # -write. (Without this, opening a connection under concurrent writers could # raise "database is locked" before the timeout was ever applied.) conn.execute(f"PRAGMA busy_timeout={_BUSY_TIMEOUT_MS}") conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys=ON") # Record the backing file path so the compare-and-set can derive it via a # thread-safe attribute read instead of running a PRAGMA on a connection that # callers share across threads (a sqlite3.Connection is not safe for # concurrent use — even a read would corrupt its transaction state). Empty # for an in-memory DB (no file to reopen on a second connection). conn.agent_team_db_path = "" if str(db_path) == ":memory:" else str(db_path) return conn def init_db(db_path: Path) -> None: """Create the agent-team tables in ``db_path`` if absent. Creates ``pending_questions`` (+ indexes), the budget ledger (+ indexes), and the ``schema_meta`` version row, and reserves the same DB file for the LangGraph ``SqliteSaver`` checkpointer (which creates its own tables on first use against this connection). Idempotent: safe to call on every startup. """ conn = connect(db_path) try: conn.execute(SCHEMA_META_DDL) conn.execute(PENDING_QUESTIONS_DDL) for stmt in _split_statements(PENDING_QUESTIONS_INDEXES_DDL): conn.execute(stmt) conn.execute(BUDGET_LEDGER_DDL) for stmt in _split_statements(BUDGET_LEDGER_INDEXES_DDL): conn.execute(stmt) # Record the schema version (single-row table). conn.execute( "INSERT INTO schema_meta (id, schema_version) VALUES (1, ?) " "ON CONFLICT(id) DO NOTHING", (SCHEMA_VERSION,), ) finally: conn.close() def migrate(conn: sqlite3.Connection) -> None: """Step ``conn``'s schema forward to :data:`SCHEMA_VERSION`. Reads the recorded version from ``schema_meta`` (treating an empty/absent row as version 0), applies any forward steps, and records the new version. At ``SCHEMA_VERSION == 1`` there are no prior versions to migrate from, so this ensures the base tables exist and stamps the version. Future versions add ordered ``if current < N`` blocks here. """ conn.execute(SCHEMA_META_DDL) row = conn.execute("SELECT schema_version FROM schema_meta WHERE id = 1").fetchone() current = int(row["schema_version"]) if row is not None else 0 if current < 1: # Base schema (v1): ensure all tables/indexes exist. conn.execute(PENDING_QUESTIONS_DDL) for stmt in _split_statements(PENDING_QUESTIONS_INDEXES_DDL): conn.execute(stmt) conn.execute(BUDGET_LEDGER_DDL) for stmt in _split_statements(BUDGET_LEDGER_INDEXES_DDL): conn.execute(stmt) current = 1 # Future steps go here: `if current < 2: ...; current = 2`. conn.execute( "INSERT INTO schema_meta (id, schema_version) VALUES (1, ?) " "ON CONFLICT(id) DO UPDATE SET schema_version = excluded.schema_version", (current,), ) def answer_question( conn: sqlite3.Connection, *, question_id: str, answer_json: str, answered_via: str, answered_at: str | None = None, ) -> bool: """First-answer-wins compare-and-set: flip an ``open`` question to answered. Runs the §3.3.1 atomic statement inside a ``BEGIN IMMEDIATE`` transaction so concurrent responders are serialized (the check-and-set takes the write lock up front). Returns ``True`` when rowcount == 1 (this caller recorded the first valid answer; enqueue a resume job), ``False`` when rowcount == 0 (the question was not ``open`` — already answered/expired/superseded — so the answer is a duplicate or late and must be ignored). """ stamp = answered_at or _utc_now_iso() return _compare_and_set( conn, sql=( "UPDATE pending_questions " "SET status='answered', answer_json=?, answered_via=?, answered_at=? " "WHERE question_id=? AND status='open'" ), params=(answer_json, answered_via, stamp, question_id), ) def expire_question( conn: sqlite3.Connection, *, question_id: str, ) -> bool: """Deadline race: flip an overdue ``open`` question to ``expired``. Same compare-and-set discipline as :func:`answer_question` (§3.3.1): an answer that arrives for an already-expired question loses the race and is ignored. Returns ``True`` if this call expired the question. """ return _compare_and_set( conn, sql=( "UPDATE pending_questions SET status='expired' " "WHERE question_id=? AND status='open'" ), params=(question_id,), ) def supersede_question( conn: sqlite3.Connection, *, question_id: str, ) -> bool: """Mark a stale ``open``/``answered`` question ``superseded``. Used by the turn-guarded resume worker: if the graph already advanced past this turn, the question is superseded and the resume is skipped (§3.3.1). Returns ``True`` if this call superseded the question. """ return _compare_and_set( conn, sql=( "UPDATE pending_questions SET status='superseded' " "WHERE question_id=? AND status IN ('open', 'answered')" ), params=(question_id,), ) # Bounded retry if the write lock is still contended after ``busy_timeout`` # elapses, so transient over-timeout contention does not surface as an error to # the responder / deadline-timer callers. _CAS_RETRY_ATTEMPTS: int = 3 _CAS_RETRY_BACKOFF_S: float = 0.05 def _main_db_file(conn: sqlite3.Connection) -> str | None: """Return the file backing ``conn``'s ``main`` database, or ``None``. ``None`` signals an in-memory database (no file to reopen on a second connection). Prefers the path stashed by :func:`connect` — a thread-safe attribute read, so it is safe even when callers share ``conn`` across threads. Falls back to ``PRAGMA database_list`` (rows of ``(seq, name, file)``, indexed positionally to be ``row_factory``-agnostic) only for a connection not opened via :func:`connect`; such a connection must not be shared across threads. """ stashed = getattr(conn, "agent_team_db_path", None) if stashed is not None: return stashed or None for row in conn.execute("PRAGMA database_list"): if row[1] == "main": return row[2] or None return None def _compare_and_set( conn: sqlite3.Connection, *, sql: str, params: tuple[Any, ...], ) -> bool: """Run a single compare-and-set UPDATE under ``BEGIN IMMEDIATE`` (§3.3.1). Returns ``True`` iff exactly one row changed. The check-and-set takes the write lock up front so concurrent responders cannot both observe ``status='open'`` (SQLite's default deferred isolation would not serialize them). **Concurrency safety.** The write runs on a private, short-lived connection to the same database file — never on the passed ``conn``. A single SQLite connection cannot hold two explicit transactions at once, so if a caller shares one ``conn`` across threads (the responder and resume worker do, and ``connect()`` sets ``check_same_thread=False``), two concurrent ``BEGIN IMMEDIATE`` statements on it would raise "cannot start a transaction within a transaction". Giving each call its own connection makes the compare-and-set safe under that sharing; WAL serializes the writers via the busy handler. A lock that outlasts ``busy_timeout`` is retried a bounded number of times before propagating. ``BEGIN IMMEDIATE`` runs inside the guarded path so its lock error is caught and retried, not raised uncaught. For an in-memory database (no file to reopen) the call falls back to the passed ``conn``; in-memory DBs are single-connection and not the concurrent production path. """ db_file = _main_db_file(conn) if db_file is None: return _cas_once(conn, sql, params) last_err: sqlite3.OperationalError | None = None for attempt in range(_CAS_RETRY_ATTEMPTS): write = connect(Path(db_file)) try: return _cas_once(write, sql, params) except sqlite3.OperationalError as err: if "locked" not in str(err).lower(): raise last_err = err finally: write.close() time.sleep(_CAS_RETRY_BACKOFF_S * (attempt + 1)) assert last_err is not None # loop only exits early via return or raise raise last_err def _cas_once( conn: sqlite3.Connection, sql: str, params: tuple[Any, ...], ) -> bool: """Execute one ``BEGIN IMMEDIATE`` compare-and-set on ``conn``. ``BEGIN IMMEDIATE`` is issued before the try so a lock-acquisition error propagates to the caller's retry loop with no transaction to unwind; once the transaction is open, any failure rolls it back (best-effort) and re-raises. """ conn.execute("BEGIN IMMEDIATE") try: cur = conn.execute(sql, params) changed = cur.rowcount == 1 conn.execute("COMMIT") return changed except BaseException: try: conn.execute("ROLLBACK") except sqlite3.OperationalError: pass raise def _utc_now_iso() -> str: """Return the current UTC time as an ISO-8601 string.""" return datetime.now(timezone.utc).isoformat() def _split_statements(ddl: str) -> list[str]: """Split a multi-statement DDL blob into individual statements.""" return [stmt.strip() for stmt in ddl.split(";") if stmt.strip()]