orchestrator/agent-team/tests/test_p3_async_resume.py

"""End-to-end async resume-on-CI-complete (design §4 Decision 2, P3 BLOCK-1/BLOCK-2).

These tests prove the "suspended forever" defect is gone: a task that dispatches
and suspends at VERIFY awaiting CI is, on a later ``tick()``, RESUMED once its run
reaches a terminal conclusion and TIMEOUT-PARKED once ``dispatched_at + timeout``
elapses with no terminal result. They drive the REAL machinery — the durable
``ci_pending_provider`` (:meth:`Coordinator._enumerate_ci_pending`, which walks the
LangGraph checkpointer), the real CI-watcher sweep, and the real turn-guarded
resume worker (:meth:`ResumeWorker.resume_ci`) — over a real compiled LangGraph app
whose VERIFY node suspends with the production awaiting-CI interrupt payload.

No network: the CI poll seam is injected. The graph is a faithful minimal stand-in
for the production BUILD → DISPATCH → VERIFY shape — DISPATCH writes the trusted
``run_id`` / ``dispatched_at`` watermarks; VERIFY ``interrupt()``s with
``{"awaiting_ci": True, "run_id": ...}`` exactly as
:func:`agent_team.nodes.build_verify_subgraph._await_ci` does — so the durable
enumeration, the suspend, and the resume are the real ones.
"""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, TypedDict

import pytest

pytest.importorskip("langgraph")

from langgraph.checkpoint.memory import InMemorySaver  # noqa: E402
from langgraph.graph import END, START, StateGraph  # noqa: E402
from langgraph.types import interrupt  # noqa: E402

from agent_team.ci_watcher import CiPollResult, CiWatchAction  # noqa: E402
from agent_team.coordinator import Coordinator  # noqa: E402
from agent_team.db.schema import init_db  # noqa: E402
from agent_team.resume_worker import ResumeWorker  # noqa: E402
from agent_team.transport.base import Transport  # noqa: E402


# --------------------------------------------------------------------------- #
# Test doubles
# --------------------------------------------------------------------------- #


class _Transport(Transport):
    """Record-only transport (no network)."""

    def post_question(self, **_kwargs: Any) -> str:  # type: ignore[override]
        return "fake:q"

    def parse_answer(self, raw: Any) -> tuple[str, Any, str]:
        return raw["question_id"], raw["answer"], "fake"


class _S(TypedDict, total=False):
    run_id: str
    dispatched_at: str
    ci_results: Any
    status: str
    current_phase: str


def _build_dispatch_verify_app(saver: InMemorySaver, *, dispatched_at: str) -> Any:
    """Compile a real DISPATCH → VERIFY graph that suspends awaiting CI.

    DISPATCH writes the trusted ``run_id`` / ``dispatched_at`` watermarks (as the
    production dispatch node does). VERIFY suspends via ``interrupt()`` with the
    production awaiting-CI payload while there is a run but no terminal
    ``ci_results``; on resume it falls through (the CI-watcher re-drives it).
    """

    def dispatch(state: _S) -> dict[str, Any]:
        return {"run_id": "R-1", "dispatched_at": dispatched_at}

    def verify(state: _S) -> dict[str, Any]:
        run_id = state.get("run_id")
        ci = state.get("ci_results")
        if run_id and ci is None:
            # Production payload shape (build_verify_subgraph._await_ci). On resume
            # the production node RE-FETCHES the authenticated conclusion; this
            # stand-in uses the resume value the CI-watcher passes (the terminal
            # poll result) as that re-fetched conclusion so VERIFY can advance.
            ci = interrupt({"awaiting_ci": True, "run_id": run_id})
        return {"status": "done", "current_phase": "done", "ci_results": ci}

    g: StateGraph = StateGraph(_S)
    g.add_node("dispatch", dispatch)
    g.add_node("verify", verify)
    g.add_edge(START, "dispatch")
    g.add_edge("dispatch", "verify")
    g.add_edge("verify", END)
    return g.compile(checkpointer=saver)


def _coordinator_over(
    db_path: Path,
    app: Any,
    saver: InMemorySaver,
    *,
    poller: Any,
    timeout: timedelta | None = None,
    alarm_hook: Any = None,
) -> Coordinator:
    """A Coordinator whose graph/resume-worker are the supplied real app.

    ``_enumerate_ci_pending`` (the durable provider) is bound as the
    ``ci_pending_provider`` so the watcher is fed REAL suspended threads off the
    real checkpointer — exactly the live serve wiring.
    """
    coord = Coordinator(
        db_path=db_path,
        transport=_Transport(),
        build_checkpointer=lambda _p: saver,
        ci_poller=poller,
        ci_timeout=timeout,
        alarm_hook=alarm_hook,
    )
    coord._graph = app
    coord._resume_worker = ResumeWorker(app, _connect(db_path))
    coord._ci_pending_provider = coord._enumerate_ci_pending
    return coord


def _connect(db_path: Path) -> Any:
    from agent_team.db.schema import connect

    return connect(db_path)


def _thread_cfg(thread_id: str) -> dict[str, Any]:
    return {"configurable": {"thread_id": thread_id}}


@pytest.fixture()
def db_path(tmp_path: Path) -> Path:
    path = tmp_path / "state" / "agent_team.sqlite"
    init_db(path)
    return path


# --------------------------------------------------------------------------- #
# (a) provider enumeration: includes suspended-at-VERIFY, excludes advanced
# --------------------------------------------------------------------------- #


def test_provider_enumerates_suspended_and_excludes_advanced(db_path: Path) -> None:
    """``_enumerate_ci_pending`` returns the thread suspended at VERIFY awaiting CI
    and EXCLUDES a resumed/done thread and a never-dispatched (no run) thread."""
    saver = InMemorySaver()
    app = _build_dispatch_verify_app(saver, dispatched_at=_now_iso())
    coord = _coordinator_over(
        db_path, app, saver, poller=lambda t: CiPollResult.pending()
    )

    # t-wait: dispatch + suspend at VERIFY awaiting CI (still pending).
    app.invoke({}, _thread_cfg("t-wait"))
    # t-done: dispatch, suspend, then RESUME with a terminal CI result -> advances
    # past VERIFY to DONE (no awaiting-CI interrupt left).
    from langgraph.types import Command

    app.invoke({}, _thread_cfg("t-done"))
    app.invoke(Command(resume={"conclusion": "success"}), _thread_cfg("t-done"))

    pending = coord._enumerate_ci_pending()
    thread_ids = {p.thread_id for p in pending}

    assert "t-wait" in thread_ids  # suspended at VERIFY awaiting CI -> included
    assert "t-done" not in thread_ids  # advanced past the gate -> excluded
    # The included task carries the trusted dispatch watermarks the watcher keys
    # off (so the poll + timeout have a run to act on).
    waiting = next(p for p in pending if p.thread_id == "t-wait")
    assert waiting.run_id == "R-1"
    assert waiting.dispatched_at is not None


# --------------------------------------------------------------------------- #
# (b) end-to-end: tick() RESUMES on terminal, TIMEOUT-PARKS on no-terminal
# --------------------------------------------------------------------------- #


def test_tick_resumes_suspended_task_on_terminal_conclusion(db_path: Path) -> None:
    """A task suspended at VERIFY is RESUMED by a later tick() once its run is
    terminal — proving 'suspended forever' is gone."""
    saver = InMemorySaver()
    app = _build_dispatch_verify_app(saver, dispatched_at=_now_iso())
    coord = _coordinator_over(
        db_path,
        app,
        saver,
        poller=lambda t: CiPollResult.terminal({"run_id": "R-1", "conclusion": "ok"}),
    )

    # Dispatch + suspend at VERIFY.
    app.invoke({}, _thread_cfg("t1"))
    snap = app.get_state(_thread_cfg("t1"))
    assert snap.next == ("verify",)  # genuinely suspended awaiting CI

    # A subsequent tick() runs the CI-watch sweep -> resume the suspended task.
    report = coord._ci_watch()
    assert report is not None
    assert report.resumed == 1
    assert report.outcomes[0].action is CiWatchAction.RESUMED

    # The thread advanced past VERIFY (no longer suspended).
    snap2 = app.get_state(_thread_cfg("t1"))
    assert snap2.next == ()
    assert snap2.values.get("status") == "done"


def test_tick_timeout_parks_suspended_task_when_run_never_terminates(
    db_path: Path,
) -> None:
    """A task whose run never reaches a terminal conclusion within the timeout is
    TIMEOUT-PARKED by a later tick() (never waits forever)."""
    saver = InMemorySaver()
    # Dispatched long ago: dispatched_at + timeout has already elapsed.
    old = (datetime.now(timezone.utc) - timedelta(hours=2)).isoformat()
    app = _build_dispatch_verify_app(saver, dispatched_at=old)

    alarms: list[str] = []
    coord = _coordinator_over(
        db_path,
        app,
        saver,
        poller=lambda t: CiPollResult.pending(),  # never terminal
        timeout=timedelta(minutes=30),
        alarm_hook=alarms.append,
    )

    app.invoke({}, _thread_cfg("t-slow"))
    assert app.get_state(_thread_cfg("t-slow")).next == ("verify",)

    report = coord._ci_watch()
    assert report is not None
    assert report.parked_timeout == 1
    assert report.outcomes[0].action is CiWatchAction.PARKED_TIMEOUT
    # Durably parked + ALARMed (surfaced, not silently spun on).
    parked = app.get_state(_thread_cfg("t-slow"))
    assert parked.values.get("status") == "parked"
    assert alarms == ["t-slow"]


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()