nisse/nisse-dev/src/nisse_dev/state.py

"""Durable state store (spec A2).

SQLite in WAL mode with ``synchronous=FULL``: every commit is fsynced, so a
power cut mid-write can never produce torn state. Every state transition and
intent-log write is a single immediately-committed statement — in-memory state
that hasn't hit disk is assumed lost.

The intent log records a unit of work *before* it starts, so a restart knows
what was in flight (see ``nisse_dev.reconcile`` for how pending intents are
resolved against external truth on boot).
"""

import json
import sqlite3
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Self

from nisse_dev.canonical import (
    CanonicalEvent,
    StoredEvent,
    event_type,
    from_payload,
    to_payload,
)
from nisse_dev.oplog import OpLog

CURRENT_SCHEMA_VERSION = 3

Migration = Callable[[sqlite3.Connection], None]

_CANONICAL_EVENTS_DDL = """
CREATE TABLE canonical_events (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    ts TEXT NOT NULL,
    type TEXT NOT NULL,
    payload TEXT NOT NULL
)
"""


_EVAL_OUTCOMES_DDL = """
CREATE TABLE eval_outcomes (
    event_id INTEGER PRIMARY KEY,
    verdict TEXT NOT NULL,
    note TEXT NOT NULL DEFAULT '',
    marked_at TEXT NOT NULL
)
"""


def _migrate_v1_add_canonical_events(conn: sqlite3.Connection) -> None:
    conn.execute(_CANONICAL_EVENTS_DDL)


def _migrate_v2_add_eval_outcomes(conn: sqlite3.Connection) -> None:
    conn.execute(_EVAL_OUTCOMES_DDL)


# Maps version N to the migration that takes the schema from N to N+1.
MIGRATIONS: Mapping[int, Migration] = {
    1: _migrate_v1_add_canonical_events,
    2: _migrate_v2_add_eval_outcomes,
}


class StateSchemaError(RuntimeError):
    pass


@dataclass(frozen=True)
class AgentState:
    state: str
    project: str | None
    data: dict[str, Any]
    updated_at: str


@dataclass(frozen=True)
class Intent:
    id: int
    kind: str
    payload: dict[str, Any]
    status: str
    created_at: str
    resolved_at: str | None


_SCHEMA = [
    "CREATE TABLE meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)",
    """
    CREATE TABLE agent_state (
        id INTEGER PRIMARY KEY CHECK (id = 1),
        state TEXT NOT NULL,
        project TEXT,
        data TEXT NOT NULL DEFAULT '{}',
        updated_at TEXT NOT NULL
    )
    """,
    """
    CREATE TABLE intent_log (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        kind TEXT NOT NULL,
        payload TEXT NOT NULL,
        status TEXT NOT NULL DEFAULT 'pending',
        created_at TEXT NOT NULL,
        resolved_at TEXT
    )
    """,
    _CANONICAL_EVENTS_DDL,
    _EVAL_OUTCOMES_DDL,
]


def _now() -> str:
    return datetime.now(UTC).isoformat()


class StateStore:
    def __init__(
        self,
        path: str | Path,
        *,
        schema_version: int = CURRENT_SCHEMA_VERSION,
        migrations: Mapping[int, Migration] = MIGRATIONS,
        oplog: OpLog | None = None,
    ) -> None:
        # Log emission is owned by the store (A9): the same code path that
        # performs a durable write emits its log line, right after commit.
        self.oplog = oplog if oplog is not None else OpLog("nisse-dev")
        self._conn = sqlite3.connect(path, isolation_level=None)
        self._conn.row_factory = sqlite3.Row
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA synchronous=FULL")
        self._conn.execute("PRAGMA busy_timeout=5000")
        self._target_version = schema_version
        self.closed = False
        self._init_or_migrate(migrations)

    def _init_or_migrate(self, migrations: Mapping[int, Migration]) -> None:
        fresh = (
            self._conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='meta'"
            ).fetchone()
            is None
        )
        if fresh:
            self._in_transaction(self._create_schema)
            return
        stored = self.schema_version
        if stored > self._target_version:
            raise StateSchemaError(
                f"state schema version {stored} is newer than this code "
                f"supports ({self._target_version}); refusing to open"
            )
        for version in range(stored, self._target_version):
            migrate = migrations.get(version)
            if migrate is None:
                raise StateSchemaError(
                    f"no migration registered from schema version {version}"
                )
            self._in_transaction(lambda m=migrate, v=version: self._apply(m, v))
            self.oplog.emit(
                "schema.migrate",
                cause="startup",
                from_version=version,
                to_version=version + 1,
            )

    def _create_schema(self) -> None:
        for statement in _SCHEMA:
            self._conn.execute(statement)
        self._conn.execute(
            "INSERT INTO meta (key, value) VALUES ('schema_version', ?)",
            (str(self._target_version),),
        )

    def _apply(self, migrate: Migration, from_version: int) -> None:
        migrate(self._conn)
        self._conn.execute(
            "UPDATE meta SET value = ? WHERE key = 'schema_version'",
            (str(from_version + 1),),
        )

    def _in_transaction(self, work: Callable[[], None]) -> None:
        self._conn.execute("BEGIN IMMEDIATE")
        try:
            work()
        except BaseException:
            self._conn.execute("ROLLBACK")
            raise
        self._conn.execute("COMMIT")

    @property
    def schema_version(self) -> int:
        row = self._conn.execute(
            "SELECT value FROM meta WHERE key = 'schema_version'"
        ).fetchone()
        return int(row["value"])

    def get_state(self) -> AgentState | None:
        row = self._conn.execute("SELECT * FROM agent_state WHERE id = 1").fetchone()
        if row is None:
            return None
        return AgentState(
            state=row["state"],
            project=row["project"],
            data=json.loads(row["data"]),
            updated_at=row["updated_at"],
        )

    def transition(
        self,
        state: str,
        *,
        cause: str,
        project: str | None = None,
        data: dict[str, Any] | None = None,
    ) -> None:
        """Record the new state-machine position; the full position is passed
        every time (no partial updates), committed and fsynced before return."""
        previous = self.get_state()
        self._conn.execute(
            """
            INSERT INTO agent_state (id, state, project, data, updated_at)
            VALUES (1, ?, ?, ?, ?)
            ON CONFLICT(id) DO UPDATE SET
                state = excluded.state,
                project = excluded.project,
                data = excluded.data,
                updated_at = excluded.updated_at
            """,
            (state, project, json.dumps(data or {}), _now()),
        )
        self.oplog.emit(
            "state.transition",
            cause=cause,
            **{
                "from": previous.state if previous else None,
                "to": state,
                "project": project,
            },
        )

    def begin_intent(self, kind: str, payload: dict[str, Any], *, cause: str) -> int:
        """Durably record that a unit of work is about to start."""
        cursor = self._conn.execute(
            "INSERT INTO intent_log (kind, payload, created_at) VALUES (?, ?, ?)",
            (kind, json.dumps(payload), _now()),
        )
        assert cursor.lastrowid is not None
        self.oplog.emit(
            "intent.begin", cause=cause, intent_id=cursor.lastrowid, kind=kind
        )
        return cursor.lastrowid

    def commit_intent(self, intent_id: int, *, cause: str) -> None:
        self._resolve_intent(intent_id, "committed", "intent.commit", cause)

    def abort_intent(self, intent_id: int, *, cause: str) -> None:
        self._resolve_intent(intent_id, "aborted", "intent.abort", cause)

    def _resolve_intent(
        self, intent_id: int, status: str, event: str, cause: str
    ) -> None:
        row = self._conn.execute(
            "UPDATE intent_log SET status = ?, resolved_at = ? "
            "WHERE id = ? AND status = 'pending' RETURNING kind",
            (status, _now(), intent_id),
        ).fetchone()
        if row is None:
            raise KeyError(f"no pending intent with id {intent_id}")
        self.oplog.emit(event, cause=cause, intent_id=intent_id, kind=row["kind"])

    def pending_intents(self) -> list[Intent]:
        rows = self._conn.execute(
            "SELECT * FROM intent_log WHERE status = 'pending' ORDER BY id"
        ).fetchall()
        return [
            Intent(
                id=row["id"],
                kind=row["kind"],
                payload=json.loads(row["payload"]),
                status=row["status"],
                created_at=row["created_at"],
                resolved_at=row["resolved_at"],
            )
            for row in rows
        ]

    def append_event(self, event: CanonicalEvent) -> int:
        """Append to the canonical transcript (A10); committed and fsynced
        before return, like every other durable write."""
        cursor = self._conn.execute(
            "INSERT INTO canonical_events (ts, type, payload) VALUES (?, ?, ?)",
            (_now(), event_type(event), json.dumps(to_payload(event))),
        )
        assert cursor.lastrowid is not None
        return cursor.lastrowid

    def read_events(self, after: int = 0) -> list[StoredEvent]:
        rows = self._conn.execute(
            "SELECT * FROM canonical_events WHERE id > ? ORDER BY id", (after,)
        ).fetchall()
        return [
            StoredEvent(
                id=row["id"],
                ts=row["ts"],
                event=from_payload(row["type"], json.loads(row["payload"])),
            )
            for row in rows
        ]

    def mark_outcome(self, event_id: int, verdict: str, note: str = "") -> None:
        """Record whether a delegation or consult was good (A16).

        Kept out of the canonical log on purpose: the log is an append-only
        record of what happened, and a verdict is a later judgement about it
        that may be revised.
        """
        self._conn.execute(
            "INSERT INTO eval_outcomes (event_id, verdict, note, marked_at) "
            "VALUES (?, ?, ?, ?) ON CONFLICT(event_id) DO UPDATE SET "
            "verdict = excluded.verdict, note = excluded.note, "
            "marked_at = excluded.marked_at",
            (event_id, verdict, note, _now()),
        )

    def outcomes(self) -> dict[int, tuple[str, str]]:
        rows = self._conn.execute(
            "SELECT event_id, verdict, note FROM eval_outcomes"
        ).fetchall()
        return {r["event_id"]: (r["verdict"], r["note"]) for r in rows}

    def close(self) -> None:
        self._conn.close()
        self.closed = True

    def __enter__(self) -> Self:
        return self

    def __exit__(self, *exc_info: object) -> None:
        self.close()