orchestrator/agent-team/tests/test_transitions.py

"""Tests for the task_transitions recorder + graph instrumentation.

Covers (plan refs): idempotent record_entry under resume replay (BLOCK-1), the
terminal close that fills exited_at on the last node (N2), fail-soft writes that
never break the pipeline, the B4 signature-preservation guarantee of the
``_instrument`` wrapper, and an end-to-end graph drive that records rows.
"""

from __future__ import annotations

import inspect
import sqlite3
from pathlib import Path

from agent_team.db import connect, init_db, read_transitions
from agent_team.db.schema import assert_task_transitions_ready
from agent_team.db.transitions import TransitionRecorder


def _rows(db: Path, thread_id: str) -> list[dict[str, object]]:
    conn = connect(db)
    try:
        return read_transitions(conn, thread_id)
    finally:
        conn.close()


def test_record_entry_chains_and_closes_previous(tmp_path: Path) -> None:
    db = tmp_path / "led.sqlite"
    init_db(db)
    rec = TransitionRecorder(db)

    rec.record_entry(thread_id="t1", to_phase="intake", status="active")
    rec.record_entry(thread_id="t1", to_phase="clarify", status="active")

    rows = _rows(db, "t1")
    assert [r["to_phase"] for r in rows] == ["intake", "clarify"]
    # First row closed (exited_at set) when the second landed; chain carried.
    assert rows[0]["exited_at"] is not None
    assert rows[1]["exited_at"] is None
    assert rows[1]["from_phase"] == "intake"


def test_record_entry_is_idempotent_under_replay(tmp_path: Path) -> None:
    """Re-entering the SAME open node (resume replay) does not duplicate a row."""
    db = tmp_path / "led.sqlite"
    init_db(db)
    rec = TransitionRecorder(db)

    rec.record_entry(thread_id="t1", to_phase="clarify", status="active")
    rec.record_entry(thread_id="t1", to_phase="clarify", status="active")  # replay
    rec.record_entry(thread_id="t1", to_phase="clarify", status="active")  # replay

    rows = _rows(db, "t1")
    assert len(rows) == 1
    assert rows[0]["to_phase"] == "clarify"
    assert rows[0]["exited_at"] is None  # still open, single row


def test_close_terminal_stamps_exit_on_last_open_row(tmp_path: Path) -> None:
    db = tmp_path / "led.sqlite"
    init_db(db)
    rec = TransitionRecorder(db)

    rec.record_entry(thread_id="t1", to_phase="plan", status="active")
    rec.close_terminal(thread_id="t1", status="done")

    rows = _rows(db, "t1")
    assert len(rows) == 1
    assert rows[0]["exited_at"] is not None
    assert rows[0]["status"] == "done"


def test_writes_are_fail_soft(tmp_path: Path) -> None:
    """A ledger error is swallowed — recording must never raise into the graph."""
    # Point the recorder at a path that is a directory, so connect/execute fails.
    bad = tmp_path / "a_dir"
    bad.mkdir()
    rec = TransitionRecorder(bad)
    # Must not raise.
    rec.record_entry(thread_id="t1", to_phase="intake", status="active")
    rec.close_terminal(thread_id="t1", status="done")


def test_read_transitions_missing_table_returns_empty(tmp_path: Path) -> None:
    db = tmp_path / "empty.sqlite"
    # A bare DB with no agent-team schema.
    sqlite3.connect(str(db)).close()
    conn = connect(db)
    try:
        assert read_transitions(conn, "t1") == []
    finally:
        conn.close()


def test_assert_task_transitions_ready(tmp_path: Path) -> None:
    db = tmp_path / "led.sqlite"
    init_db(db)
    # Does not raise on a properly migrated DB.
    assert_task_transitions_ready(db)


# --- B4: _instrument signature preservation + graph integration --------------


def test_instrument_preserves_signature_and_forwards_config() -> None:
    """The wrapper must present fn's signature (so LangGraph injects config) and
    forward every argument verbatim (plan B4)."""
    from agent_team.graph import _instrument

    seen: dict[str, object] = {}

    def node_with_config(state: dict, config: dict) -> dict:  # type: ignore[type-arg]
        seen["state"] = state
        seen["config"] = config
        return {"status": "active"}

    class _Rec:
        def record_entry(self, **_kw: object) -> None: ...
        def close_terminal(self, **_kw: object) -> None: ...

    wrapped = _instrument("plan", node_with_config, _Rec())
    # Signature still advertises BOTH params, so LangGraph passes config.
    assert list(inspect.signature(wrapped).parameters) == ["state", "config"]
    # And the wrapper forwards them.
    wrapped({"thread_id": "t"}, {"configurable": {}})
    assert seen["config"] == {"configurable": {}}


def test_instrument_none_recorder_is_identity() -> None:
    from agent_team.graph import _instrument

    def node(state: dict) -> dict:  # type: ignore[type-arg]
        return {}

    assert _instrument("intake", node, None) is node


def test_graph_drive_records_transitions(tmp_path: Path) -> None:
    """End-to-end: a recorder-instrumented graph records rows as a task runs."""
    from langgraph.checkpoint.memory import MemorySaver

    from agent_team.graph import build_graph, resume_task, start_task

    db = tmp_path / "led.sqlite"
    init_db(db)
    rec = TransitionRecorder(db)
    graph = build_graph(MemorySaver(), transition_recorder=rec)

    tid, _ = start_task(graph, transport="cli", task="do a thing")
    # Ran INTAKE -> CLARIFY (suspended on the human gate).
    rows = _rows(db, tid)
    phases = [r["to_phase"] for r in rows]
    assert "intake" in phases
    assert "clarify" in phases

    resume_task(graph, thread_id=tid, answer="scope it")
    rows = _rows(db, tid)
    phases = [r["to_phase"] for r in rows]
    assert "plan" in phases
    # Terminal plan node (status DONE) closed its row.
    plan_row = [r for r in rows if r["to_phase"] == "plan"][-1]
    assert plan_row["exited_at"] is not None