orchestrator/agent-team/tests/test_graph.py

"""Unit tests for agent_team.graph (Plane-2 P1 LangGraph wiring; §3.3, §7.1).

These exercise the P1 skeleton + human gate wiring:

* the pure node functions (intake/clarify-author/plan) in isolation,
* graph assembly + edge topology,
* the suspend-on-interrupt / resume-with-Command mechanic end to end,
* the thread_id-keyed driver seam (start/resume/get_state/pending_question),
* that the foundation contracts (PipelineState / Phase / TaskStatus /
  QuestionSet) are imported verbatim and round-trip through the wiring.

An in-memory checkpointer is injected (the SQLite checkpointer is the
production store, D9, not constructed in pre-deploy scaffolding).
"""

from __future__ import annotations

import pytest

# InMemorySaver is the modern name; fall back to MemorySaver on older langgraph.
try:  # pragma: no cover - import shim
    from langgraph.checkpoint.memory import InMemorySaver as _Saver
except ImportError:  # pragma: no cover - import shim
    from langgraph.checkpoint.memory import MemorySaver as _Saver

from agent_team import graph as graph_mod
from agent_team.graph import (
    CLARIFY,
    INTAKE,
    P1_PHASE_SEQUENCE,
    PLAN,
    build_graph,
    build_sqlite_checkpointer,
    clarify_node,
    get_pipeline_state,
    intake_node,
    pending_question,
    plan_node,
    plan_phase,
    resume_task,
    start_task,
    thread_config,
)
from agent_team.task_model import Phase, PipelineState, TaskStatus
from agent_team.transport import QuestionSet


@pytest.fixture()
def compiled():
    """A graph compiled with a fresh in-memory checkpointer per test."""
    return build_graph(checkpointer=_Saver())


# --- Module surface / constants. -------------------------------------------


def test_node_name_constants_are_distinct() -> None:
    assert len({INTAKE, CLARIFY, PLAN}) == 3


def test_p1_phase_sequence_stops_at_plan() -> None:
    # P1 ends at an approved plan — no BUILD/VERIFY in the wired sequence (§7.1).
    assert P1_PHASE_SEQUENCE == (Phase.INTAKE, Phase.CLARIFY, Phase.PLAN)
    assert Phase.BUILD not in P1_PHASE_SEQUENCE
    assert Phase.VERIFY not in P1_PHASE_SEQUENCE


# --- Pure node behaviour. ---------------------------------------------------


def test_intake_node_activates_and_advances_to_clarify() -> None:
    out = intake_node(PipelineState(thread_id="t", current_phase=Phase.INTAKE.value))
    assert out["status"] == TaskStatus.ACTIVE.value
    assert out["current_phase"] == Phase.CLARIFY.value
    assert out["updated_at"]


def test_plan_node_lands_approved_plan_and_finishes() -> None:
    out = plan_node(PipelineState(thread_id="t", qa_history=[{"answer": "x"}]))
    assert out["status"] == TaskStatus.DONE.value
    assert out["current_phase"] == Phase.DONE.value
    assert out["plan"]["approved"] is True


def test_plan_phase_counts_qa_turns() -> None:
    state = PipelineState(qa_history=[{"answer": "a"}, {"answer": "b"}])
    plan = plan_phase(state)
    assert plan["qa_turns"] == 2
    assert plan["approved"] is True


def test_plan_phase_handles_empty_history() -> None:
    assert plan_phase(PipelineState())["qa_turns"] == 0


def test_clarify_node_suspends_rather_than_falling_through() -> None:
    # Called bare (no running graph), interrupt() refuses to return a value:
    # it raises because there is no runnable context to suspend into. This
    # confirms clarify_node genuinely suspends rather than falling through to
    # its post-interrupt return.
    with pytest.raises(RuntimeError):
        clarify_node(PipelineState(thread_id="t", transport="slack"))


# --- Graph assembly. --------------------------------------------------------


def test_build_graph_without_checkpointer_compiles() -> None:
    # An uncheckpointed graph still compiles (used only for straight-through
    # smoke paths); the driver requires a checkpointer for suspend/resume.
    assert build_graph() is not None


def test_build_graph_with_checkpointer_compiles(compiled) -> None:
    assert compiled is not None


def test_graph_nodes_present(compiled) -> None:
    nodes = set(compiled.get_graph().nodes)
    assert {INTAKE, CLARIFY, PLAN} <= nodes


# --- Suspend / resume end to end. ------------------------------------------


def test_start_task_suspends_on_human_gate(compiled) -> None:
    thread_id, state = start_task(compiled, transport="slack")
    # The task ran INTAKE then suspended at CLARIFY's interrupt().
    assert "__interrupt__" in state
    payload = pending_question(compiled, thread_id=thread_id)
    assert payload is not None
    assert payload["thread_id"] == thread_id
    assert payload["transport"] == "slack"
    assert payload["turn"] == 0
    assert payload["deadline"]


def test_pending_question_carries_foundation_questionset(compiled) -> None:
    thread_id, _ = start_task(compiled, transport="slack")
    payload = pending_question(compiled, thread_id=thread_id)
    qset = payload["question_set"]
    # Verbatim foundation contract — not a redefinition.
    assert isinstance(qset, QuestionSet)
    assert qset.thread_id == thread_id
    assert qset.question_id == payload["question_id"]
    assert qset.turn == 0
    assert qset.questions  # non-empty question-set


def test_resume_drives_task_to_done(compiled) -> None:
    thread_id, _ = start_task(compiled, transport="slack")
    final = resume_task(compiled, thread_id=thread_id, answer={"text": "do the thing"})
    assert final["status"] == TaskStatus.DONE.value
    assert final["current_phase"] == Phase.DONE.value
    assert final["plan"]["approved"] is True


def test_answer_is_recorded_in_qa_history(compiled) -> None:
    thread_id, _ = start_task(compiled, transport="slack")
    answer = {"text": "scope is X"}
    final = resume_task(compiled, thread_id=thread_id, answer=answer)
    assert len(final["qa_history"]) == 1
    assert final["qa_history"][0]["answer"] == answer
    assert final["qa_history"][0]["turn"] == 0


def test_question_id_is_stable_across_resume(compiled) -> None:
    # The clarifier node re-executes on resume; the question_id must NOT change
    # between the id delivered at suspend (the ledger key) and the one recorded
    # in qa_history, or the §3.3.1 identity contract breaks.
    thread_id, _ = start_task(compiled, transport="slack")
    delivered = pending_question(compiled, thread_id=thread_id)["question_id"]
    final = resume_task(compiled, thread_id=thread_id, answer="ok")
    assert final["qa_history"][0]["question_id"] == delivered


def test_no_pending_question_after_completion(compiled) -> None:
    thread_id, _ = start_task(compiled, transport="slack")
    resume_task(compiled, thread_id=thread_id, answer="ok")
    assert pending_question(compiled, thread_id=thread_id) is None


def test_get_pipeline_state_reflects_suspend_then_done(compiled) -> None:
    thread_id, _ = start_task(compiled, transport="slack")
    mid = get_pipeline_state(compiled, thread_id=thread_id)
    # Suspended ON the clarifier gate: INTAKE already advanced the phase to
    # CLARIFY, and the clarifier's post-interrupt write (-> PLAN) has NOT yet
    # committed because the node is paused at interrupt(). Task is mid-flight.
    assert mid["current_phase"] == Phase.CLARIFY.value
    assert mid["status"] == TaskStatus.ACTIVE.value
    resume_task(compiled, thread_id=thread_id, answer="ok")
    done = get_pipeline_state(compiled, thread_id=thread_id)
    assert done["status"] == TaskStatus.DONE.value
    assert done["current_phase"] == Phase.DONE.value


# --- Thread isolation (§3.3.1 P1 exit criterion (d)). ----------------------


def test_two_tasks_suspend_and_resume_independently(compiled) -> None:
    t1, _ = start_task(compiled, transport="slack")
    t2, _ = start_task(compiled, transport="github")

    assert t1 != t2
    p1 = pending_question(compiled, thread_id=t1)
    p2 = pending_question(compiled, thread_id=t2)
    assert p1["transport"] == "slack"
    assert p2["transport"] == "github"
    assert p1["question_id"] != p2["question_id"]

    # Resume only t1; t2 must remain suspended on its own gate.
    f1 = resume_task(compiled, thread_id=t1, answer="answer-1")
    assert f1["status"] == TaskStatus.DONE.value
    assert pending_question(compiled, thread_id=t2) is not None

    f2 = resume_task(compiled, thread_id=t2, answer="answer-2")
    assert f2["status"] == TaskStatus.DONE.value
    assert f2["qa_history"][0]["answer"] == "answer-2"


def test_explicit_thread_id_is_honoured(compiled) -> None:
    tid, _ = start_task(compiled, thread_id="fixed-thread", transport="slack")
    assert tid == "fixed-thread"
    assert pending_question(compiled, thread_id="fixed-thread") is not None


# --- Durable resume across a fresh graph object (P1 exit criterion (a)). ----


def test_resume_works_on_a_new_graph_over_shared_checkpointer() -> None:
    # Simulates a process restart: a NEW compiled graph object built over the
    # SAME checkpointer must resume a task suspended by the first graph object.
    saver = _Saver()
    g1 = build_graph(checkpointer=saver)
    thread_id, _ = start_task(g1, transport="slack")

    g2 = build_graph(checkpointer=saver)  # "after restart"
    assert pending_question(g2, thread_id=thread_id) is not None
    final = resume_task(g2, thread_id=thread_id, answer="post-restart")
    assert final["status"] == TaskStatus.DONE.value
    assert final["qa_history"][0]["answer"] == "post-restart"


# --- Driver-seam helpers. ---------------------------------------------------


def test_thread_config_shape() -> None:
    assert thread_config("abc") == {"configurable": {"thread_id": "abc"}}


def test_start_task_mints_unique_thread_ids(compiled) -> None:
    t1, _ = start_task(compiled, transport="slack")
    t2, _ = start_task(compiled, transport="slack")
    assert t1 != t2


# --- Production checkpointer factory. --------------------------------------


def test_build_sqlite_checkpointer_missing_dep_raises_runtimeerror(
    monkeypatch, tmp_path
) -> None:
    # When the optional langgraph-checkpoint-sqlite package is absent, the
    # factory must fail loudly with a clear RuntimeError, never silently run
    # uncheckpointed. Force the ImportError path deterministically.
    import builtins

    real_import = builtins.__import__

    def _blocking_import(name, *args, **kwargs):
        if name == "langgraph.checkpoint.sqlite":
            raise ImportError("blocked for test")
        return real_import(name, *args, **kwargs)

    monkeypatch.setattr(builtins, "__import__", _blocking_import)
    with pytest.raises(RuntimeError, match="SQLite checkpointer"):
        build_sqlite_checkpointer(tmp_path / "state.db")


def test_build_sqlite_checkpointer_builds_when_dep_present(tmp_path) -> None:
    # If the optional package IS installed, the factory returns a checkpointer
    # over the DB path. Skip cleanly where it's absent (pre-deploy scaffolding).
    pytest.importorskip("langgraph.checkpoint.sqlite")
    saver = build_sqlite_checkpointer(tmp_path / "nested" / "state.db")
    assert saver is not None
    assert (tmp_path / "nested").is_dir()


# --- Module import hygiene. -------------------------------------------------


def test_module_imports_without_optional_sqlite_dep() -> None:
    # The module-level import of graph must not pull in the optional SQLite
    # checkpointer (that import is deferred into build_sqlite_checkpointer).
    assert hasattr(graph_mod, "build_graph")
    assert hasattr(graph_mod, "build_sqlite_checkpointer")