orchestrator/agent-team/tests/sim/harness.py

"""P1 exit-criteria simulation harness (design §7.1 P1, demonstrating §3.3.1).

Phase P1 of the Plane-2 pipeline must *demonstrate* the durable
human-in-the-loop suspend/resume contract before anything else is built. The
four exit criteria (§7.1 P1) are:

* (a) kill the box mid-wait and have the task resume after restart;
* (b) submit a duplicate answer and confirm it no-ops;
* (c) submit an answer after the deadline expired and confirm it is rejected
  and the task parked;
* (d) two tasks suspended concurrently resume independently to the correct
  thread.

This module is a *simulation* harness, not the production pipeline. There is no
LangGraph runtime on the box yet (that lands in the P1 build proper), so the
harness stands in a minimal, faithful model of the riskiest mechanic — the
``pending_questions`` ledger and the §3.3.1 first-answer-wins / deadline-race /
turn-guarded-resume compare-and-set — *built on the real committed foundation*:

* :mod:`agent_team.db.schema` — the real ``pending_questions`` ledger DDL and
  the real ``answer_question`` / ``expire_question`` / ``supersede_question``
  ``BEGIN IMMEDIATE`` compare-and-set helpers. The harness never re-implements
  the atomic statements; it drives the committed ones.
* :mod:`agent_team.state_store` — the real atomic, integrity-checked durable
  state store. The "graph checkpoint" each task suspends on is written through
  :func:`agent_team.state_store.write_checked` and read back through
  :func:`agent_team.state_store.read_checked`, so a simulated "kill the box"
  (drop the in-memory harness, reconstruct from disk) exercises real durable
  recovery, not a Python dict.
* :mod:`agent_team.task_model` — the real :class:`TaskRecord` / :class:`Phase`
  / :class:`TaskStatus` model and its JSON serialization.
* :mod:`agent_team.transport.base` — the real :class:`Transport` ABC and
  :class:`QuestionSet` payload; :class:`RecordingTransport` is a faithful
  in-memory adapter subclassing the committed contract (no live Slack).

Nothing here provisions, schedules, or reaches live infrastructure. It is
pre-deployment scaffolding that proves the design's durable mechanic holds.
"""

from __future__ import annotations

import json
import sqlite3
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from agent_team.db.schema import (
    answer_question,
    expire_question,
    init_db,
    supersede_question,
)
from agent_team.state_store import IntegrityError, read_checked, write_checked
from agent_team.task_model import (
    Phase,
    TaskRecord,
    TaskStatus,
    new_thread_id,
    task_from_json,
    task_to_json,
)
from agent_team.transport.base import QuestionSet, Transport

__all__ = [
    "CHECKPOINT_SCHEMA_VERSION",
    "PostFailingTransport",
    "RecordingTransport",
    "ResumeOutcome",
    "SimClock",
    "SimPipeline",
    "SuspendedTask",
]

# Schema version stamped on the harness's durable "graph checkpoint" sidecars.
# Distinct from the SQL ``SCHEMA_VERSION``; this versions the checkpoint blob
# format the harness round-trips through ``state_store``.
CHECKPOINT_SCHEMA_VERSION: int = 1


class SimClock:
    """A monotonically advanceable fake clock for deadline-race tests (§3.3.1).

    Each open question carries a ``deadline_at``. Rather than sleep in tests,
    the harness compares a question's deadline against this clock's "now", and
    tests advance the clock past a deadline to drive the timer loop. Times are
    plain integer ticks (seconds since an arbitrary epoch); the ledger stores
    them as ISO-like sortable strings so the durable column is human-readable.
    """

    def __init__(self, start: int = 0) -> None:
        self._now = int(start)

    def now(self) -> int:
        """Return the current tick."""
        return self._now

    def advance(self, ticks: int) -> int:
        """Advance the clock by ``ticks`` and return the new now."""
        if ticks < 0:
            raise ValueError("cannot advance the clock backwards")
        self._now += int(ticks)
        return self._now

    def stamp(self, tick: int | None = None) -> str:
        """Render ``tick`` (default: now) as a sortable durable timestamp."""
        value = self._now if tick is None else int(tick)
        # Zero-padded so lexical order == numeric order in the ledger column.
        return f"t{value:020d}"


@dataclass
class RecordingTransport(Transport):
    """In-memory :class:`Transport` adapter that records posts (no live Slack).

    A faithful subclass of the committed :class:`agent_team.transport.base.
    Transport` contract: :meth:`post_question` embeds the ``question_id`` in
    the returned ``channel_ref`` (mirroring the real "the post MUST embed the
    question_id" rule) and records the post so recovery/reconcile tests can
    inspect delivery. :meth:`parse_answer` normalizes a ``(question_id,
    answer, via)`` raw payload, mapping it back via the embedded id.
    """

    name: str = "sim"
    posts: list[dict[str, Any]] = field(default_factory=list)

    def post_question(
        self,
        *,
        thread_id: str,
        question_id: str,
        turn: int,
        question_set: QuestionSet,
        deadline: str,
    ) -> str:
        channel_ref = f"{self.name}:{question_id}"
        self.posts.append(
            {
                "thread_id": thread_id,
                "question_id": question_id,
                "turn": turn,
                "questions": list(question_set.questions),
                "deadline": deadline,
                "channel_ref": channel_ref,
            }
        )
        return channel_ref

    def parse_answer(self, raw: Any) -> tuple[str, Any, str]:
        question_id = raw["question_id"]
        answer = raw.get("answer")
        via = raw.get("via", self.name)
        return question_id, answer, via


@dataclass
class PostFailingTransport(RecordingTransport):
    """A transport whose first ``post_question`` raises (lost-post simulation).

    Used to exercise the §3.3.1 "if the post fails, the row stays ``open`` with
    no ref and a reconcile loop retries idempotently" path. The first post
    raises; subsequent posts succeed and record normally.
    """

    fail_times: int = 1
    _attempts: int = 0

    def post_question(
        self,
        *,
        thread_id: str,
        question_id: str,
        turn: int,
        question_set: QuestionSet,
        deadline: str,
    ) -> str:
        self._attempts += 1
        if self._attempts <= self.fail_times:
            raise RuntimeError("simulated transport post failure")
        return super().post_question(
            thread_id=thread_id,
            question_id=question_id,
            turn=turn,
            question_set=question_set,
            deadline=deadline,
        )


@dataclass(frozen=True)
class SuspendedTask:
    """Handle to a task suspended on a question (returned by :meth:`SimPipeline.submit`)."""

    thread_id: str
    question_id: str
    turn: int


@dataclass(frozen=True)
class ResumeOutcome:
    """Result of attempting to resume a task on an answered question.

    ``resumed`` is True iff the turn guard passed and the graph advanced;
    ``superseded`` is True iff a stale/redelivered resume was skipped (§3.3.1
    "a resume can never double-apply").
    """

    thread_id: str
    resumed: bool
    superseded: bool
    new_phase: Phase | None


class SimPipeline:
    """A minimal, durable simulation of the §3.3.1 suspend/resume mechanic.

    The pipeline owns two real durable stores under ``root``:

    * the SQLite ``pending_questions`` ledger (via the committed
      :mod:`agent_team.db.schema`), the single source of truth for the
      question lifecycle, and
    * one integrity-checked "graph checkpoint" file per task (via the
      committed :mod:`agent_team.state_store`), holding the durable
      :class:`agent_team.task_model.TaskRecord`.

    A task is submitted, suspends on a question (status ``WAITING_HUMAN``,
    ledger row ``open``), and later resumes when a first valid answer wins the
    compare-and-set. "Killing the box" is modelled by dropping the in-memory
    pipeline and calling :meth:`reopen`, which reconstructs purely from the two
    durable stores — proving recovery has no in-memory-only state.
    """

    def __init__(self, root: Path, *, clock: SimClock, transport: Transport) -> None:
        self._root = Path(root)
        self._db_path = self._root / "agent_team.sqlite"
        self._checkpoints = self._root / "checkpoints"
        self._clock = clock
        self._transport = transport
        self._root.mkdir(parents=True, exist_ok=True)
        self._checkpoints.mkdir(parents=True, exist_ok=True)
        init_db(self._db_path)

    # -- durable checkpoint helpers (real state_store) --------------------

    def _checkpoint_path(self, thread_id: str) -> Path:
        return self._checkpoints / f"{thread_id}.json"

    def _write_checkpoint(self, record: TaskRecord) -> None:
        """Persist a task record through the real atomic state-store."""
        write_checked(
            self._checkpoint_path(record.thread_id),
            task_to_json(record).encode("utf-8"),
            schema_version=CHECKPOINT_SCHEMA_VERSION,
        )

    def load_record(self, thread_id: str) -> TaskRecord:
        """Read a task record back, integrity-checked (raises IntegrityError)."""
        data = read_checked(
            self._checkpoint_path(thread_id),
            schema_version=CHECKPOINT_SCHEMA_VERSION,
        )
        return task_from_json(data.decode("utf-8"))

    # -- ledger helpers (real db.schema connection) -----------------------

    def _connect(self) -> sqlite3.Connection:
        return sqlite3.connect(
            str(self._db_path), isolation_level=None, check_same_thread=False
        )

    def ledger_row(self, question_id: str) -> sqlite3.Row | None:
        """Return the durable ``pending_questions`` row for ``question_id``."""
        conn = self._connect()
        conn.row_factory = sqlite3.Row
        try:
            return conn.execute(
                "SELECT * FROM pending_questions WHERE question_id = ?",
                (question_id,),
            ).fetchone()
        finally:
            conn.close()

    # -- pipeline operations ----------------------------------------------

    def submit(self, *, questions: list[str], deadline_in: int) -> SuspendedTask:
        """Submit a task; it advances to CLARIFY and suspends on a question.

        Writes the ledger row ``open`` *first*, then posts to the transport and
        stores the returned ``channel_ref`` (the §3.3.1 delivery order). If the
        post fails the row stays ``open`` with no ref for the reconcile loop to
        retry. The durable task record is checkpointed as ``WAITING_HUMAN``.
        """
        thread_id = new_thread_id()
        question_id = uuid.uuid4().hex
        turn = 0
        deadline_at = self._clock.stamp(self._clock.now() + int(deadline_in))

        record = TaskRecord(
            thread_id=thread_id,
            status=TaskStatus.WAITING_HUMAN,
            current_phase=Phase.CLARIFY,
            transport=getattr(self._transport, "name", "sim"),
            created_at=self._clock.stamp(),
            updated_at=self._clock.stamp(),
        )
        self._write_checkpoint(record)

        # Ledger row first (open, no channel_ref yet).
        conn = self._connect()
        try:
            conn.execute(
                "INSERT INTO pending_questions "
                "(question_id, thread_id, turn, status, transport, posted_at, deadline_at) "
                "VALUES (?, ?, ?, 'open', ?, ?, ?)",
                (
                    question_id,
                    thread_id,
                    turn,
                    record.transport,
                    self._clock.stamp(),
                    deadline_at,
                ),
            )
        finally:
            conn.close()

        # Then deliver; tolerate a lost post (row stays open, no ref).
        self._deliver(
            thread_id=thread_id,
            question_id=question_id,
            turn=turn,
            questions=questions,
            deadline_at=deadline_at,
        )
        return SuspendedTask(thread_id=thread_id, question_id=question_id, turn=turn)

    def _deliver(
        self,
        *,
        thread_id: str,
        question_id: str,
        turn: int,
        questions: list[str],
        deadline_at: str,
    ) -> str | None:
        """Post the question and store the channel_ref; tolerate post failure."""
        question_set = QuestionSet(
            thread_id=thread_id,
            question_id=question_id,
            turn=turn,
            questions=list(questions),
        )
        try:
            channel_ref = self._transport.post_question(
                thread_id=thread_id,
                question_id=question_id,
                turn=turn,
                question_set=question_set,
                deadline=deadline_at,
            )
        except Exception:
            # Lost post: row stays open with no ref; reconcile retries later.
            return None
        conn = self._connect()
        try:
            conn.execute(
                "UPDATE pending_questions SET channel_ref = ? WHERE question_id = ?",
                (channel_ref, question_id),
            )
        finally:
            conn.close()
        return channel_ref

    def reconcile(self, *, questions_by_qid: dict[str, list[str]]) -> int:
        """Retry delivery for ``open`` rows lacking a ``channel_ref`` (§3.3.1).

        Returns the number of rows for which a (re)delivery now succeeded.
        ``questions_by_qid`` supplies the question text per id (the harness
        does not persist question text on the ledger, mirroring the design's
        ledger schema which carries lifecycle, not prompt bodies).
        """
        conn = self._connect()
        conn.row_factory = sqlite3.Row
        try:
            rows = conn.execute(
                "SELECT question_id, thread_id, turn, deadline_at "
                "FROM pending_questions "
                "WHERE status = 'open' AND channel_ref IS NULL"
            ).fetchall()
        finally:
            conn.close()
        redelivered = 0
        for row in rows:
            ref = self._deliver(
                thread_id=row["thread_id"],
                question_id=row["question_id"],
                turn=row["turn"],
                questions=questions_by_qid.get(row["question_id"], []),
                deadline_at=row["deadline_at"],
            )
            if ref is not None:
                redelivered += 1
        return redelivered

    def submit_answer(self, raw: Any) -> bool:
        """Normalize ``raw`` via the transport and run the §3.3.1 compare-and-set.

        Returns ``True`` when this answer won the race (ledger rowcount 1 — the
        first valid answer, a resume is now eligible) and ``False`` when it
        lost (rowcount 0 — duplicate, late, or for a closed question, ignored).
        Drives the *committed* :func:`agent_team.db.schema.answer_question`
        ``BEGIN IMMEDIATE`` statement; the harness never re-implements it.
        """
        question_id, answer, via = self._transport.parse_answer(raw)
        conn = self._connect()
        try:
            return answer_question(
                conn,
                question_id=question_id,
                answer_json=json.dumps(answer),
                answered_via=via,
                answered_at=self._clock.stamp(),
            )
        finally:
            conn.close()

    def run_deadline_sweep(self) -> list[str]:
        """Expire every overdue ``open`` question and park its task (§3.3.1).

        A timer loop flips overdue ``open`` rows to ``expired`` via the
        committed compare-and-set (:func:`agent_team.db.schema.
        expire_question`) and applies the park policy: the task record flips to
        :attr:`TaskStatus.PARKED` / :attr:`Phase.PARKED`. Returns the list of
        ``question_id`` s expired by this sweep. An answer arriving for an
        already-expired question will lose its own compare-and-set.
        """
        now = self._clock.now()
        conn = self._connect()
        conn.row_factory = sqlite3.Row
        try:
            rows = conn.execute(
                "SELECT question_id, thread_id, deadline_at "
                "FROM pending_questions WHERE status = 'open'"
            ).fetchall()
        finally:
            conn.close()

        expired: list[str] = []
        for row in rows:
            if not self._is_overdue(row["deadline_at"], now):
                continue
            conn2 = self._connect()
            try:
                won = expire_question(conn2, question_id=row["question_id"])
            finally:
                conn2.close()
            if won:
                expired.append(row["question_id"])
                self._park(row["thread_id"])
        return expired

    @staticmethod
    def _is_overdue(deadline_at: str | None, now: int) -> bool:
        """Decode a ``SimClock``-stamped deadline and test it against ``now``."""
        if not deadline_at:
            return False
        try:
            deadline_tick = int(deadline_at.lstrip("t"))
        except ValueError:
            return False
        return now >= deadline_tick

    def _park(self, thread_id: str) -> None:
        """Flip a task record to PARKED durably (idempotent)."""
        record = self.load_record(thread_id)
        record.status = TaskStatus.PARKED
        record.current_phase = Phase.PARKED
        record.updated_at = self._clock.stamp()
        self._write_checkpoint(record)

    def resume(self, thread_id: str, question_id: str) -> ResumeOutcome:
        """Turn-guarded resume of an ``answered`` question (§3.3.1, single-flight).

        Mirrors the design's resume worker: before advancing the graph it
        checks the live checkpoint is still interrupted on this turn. The
        durable task record is the checkpoint here, so:

        * if the record is still ``WAITING_HUMAN`` on the answered question, the
          graph advances (CLARIFY -> PLAN), the record is checkpointed
          ``ACTIVE``, and ``resumed`` is True;
        * if the record already advanced (a stale/redelivered resume), the
          question is marked ``superseded`` via the committed compare-and-set
          and the resume is skipped (``superseded`` True), so a resume can
          never double-apply.
        """
        row = self.ledger_row(question_id)
        if row is None or row["status"] != "answered":
            return ResumeOutcome(
                thread_id=thread_id, resumed=False, superseded=False, new_phase=None
            )

        record = self.load_record(thread_id)
        # Turn guard: only resume if still suspended on this turn/phase.
        if (
            record.status is not TaskStatus.WAITING_HUMAN
            or record.current_phase is not Phase.CLARIFY
        ):
            conn = self._connect()
            try:
                supersede_question(conn, question_id=question_id)
            finally:
                conn.close()
            return ResumeOutcome(
                thread_id=thread_id,
                resumed=False,
                superseded=True,
                new_phase=record.current_phase,
            )

        # Apply the won answer into the durable Q&A history and advance a phase.
        answer = json.loads(row["answer_json"]) if row["answer_json"] else None
        record.qa_history.append(
            {"question_id": question_id, "turn": row["turn"], "answer": answer}
        )
        record.status = TaskStatus.ACTIVE
        record.current_phase = Phase.PLAN
        record.updated_at = self._clock.stamp()
        self._write_checkpoint(record)
        return ResumeOutcome(
            thread_id=thread_id,
            resumed=True,
            superseded=False,
            new_phase=Phase.PLAN,
        )

    # -- restart recovery -------------------------------------------------

    def reopen(self) -> SimPipeline:
        """Simulate "kill the box": return a fresh pipeline over the same disk.

        The new pipeline shares the durable SQLite ledger and the
        integrity-checked checkpoints but holds *no* in-memory state, so any
        recovery must come entirely from disk (§3.3.1 "No in-memory-only
        state."). The transport and clock are re-used (a real restart would
        re-instantiate adapters; reusing them keeps recorded posts visible to
        the assertions).
        """
        return SimPipeline(self._root, clock=self._clock, transport=self._transport)

    def startup_sweep(
        self, *, questions_by_qid: dict[str, list[str]]
    ) -> dict[str, Any]:
        """Run the §3.3.1 startup convergence sweep after a restart.

        Concretely: (1) retry delivery for ``open`` rows lacking a ref; (2)
        re-enqueue a resume for ``answered`` rows whose task is still suspended
        on that turn (idempotent via the turn guard); (3) apply the deadline
        policy for overdue ``open`` rows. Returns a summary of what converged.
        """
        redelivered = self.reconcile(questions_by_qid=questions_by_qid)

        conn = self._connect()
        conn.row_factory = sqlite3.Row
        try:
            answered = conn.execute(
                "SELECT question_id, thread_id FROM pending_questions "
                "WHERE status = 'answered'"
            ).fetchall()
        finally:
            conn.close()

        resumed: list[str] = []
        for row in answered:
            outcome = self.resume(row["thread_id"], row["question_id"])
            if outcome.resumed:
                resumed.append(row["thread_id"])

        expired = self.run_deadline_sweep()
        return {
            "redelivered": redelivered,
            "resumed": resumed,
            "expired": expired,
        }


def assert_no_integrity_error(pipeline: SimPipeline, thread_id: str) -> TaskRecord:
    """Load a record and surface :class:`IntegrityError` as an explicit failure.

    Convenience for tests that want the durable read to be part of the
    assertion (the foundation fails closed on corruption rather than returning
    junk).
    """
    try:
        return pipeline.load_record(thread_id)
    except IntegrityError as exc:  # pragma: no cover - defensive
        raise AssertionError(
            f"durable checkpoint failed integrity check: {exc}"
        ) from exc