orchestrator/agent-team/tests/test_planner.py

"""Unit tests for agent_team.nodes.planner (§3.3, §7.1 P2)."""

from __future__ import annotations

import json
from typing import Any

import pytest

from agent_team import billing
from agent_team.billing import BillingMode, ClaudeResult
from agent_team.nodes.planner import (
    MAX_PLAN_REVISIONS,
    PlannerError,
    build_plan_prompt,
    parse_plan,
    plan_node,
)
from agent_team.task_model import Phase, PipelineState, TaskStatus

# --------------------------------------------------------------------------- #
# Fixtures / helpers
# --------------------------------------------------------------------------- #

_VALID_PLAN = {
    "summary": "Bump the dependency and update the lockfile.",
    "phases": [
        {"name": "Phase 1 — bump", "steps": ["edit requirements", "run tests"]},
        {"name": "Phase 2 — verify", "steps": ["open draft PR"]},
    ],
}


@pytest.fixture(autouse=True)
def _restore_invoker():
    """Restore the module invoker after each test (mirrors test_billing)."""
    original = billing._invoker
    yield
    billing._invoker = original


def _bind_invoker(reply: str) -> list[dict[str, Any]]:
    """Bind a fake Claude invoker returning ``reply``; capture its calls."""
    calls: list[dict[str, Any]] = []

    def _fake(prompt: str, *, mode: BillingMode, **kw: Any) -> ClaudeResult:
        calls.append({"prompt": prompt, "mode": mode, "kw": kw})
        return ClaudeResult(text=reply, mode=mode, usage={"input_tokens": 1})

    billing.set_invoker(_fake)
    return calls


def _state(**overrides: Any) -> PipelineState:
    base: PipelineState = PipelineState(
        thread_id="t-1",
        status=TaskStatus.ACTIVE.value,
        current_phase=Phase.PLAN.value,
        qa_history=[],
        review_verdicts=[],
    )
    base.update(overrides)  # type: ignore[typeddict-item]
    return base


# --------------------------------------------------------------------------- #
# parse_plan
# --------------------------------------------------------------------------- #


def test_parse_plan_valid() -> None:
    plan = parse_plan(json.dumps(_VALID_PLAN))
    assert plan["summary"].startswith("Bump")
    assert len(plan["phases"]) == 2
    assert plan["phases"][0]["name"] == "Phase 1 — bump"
    assert plan["phases"][0]["steps"] == ["edit requirements", "run tests"]


def test_parse_plan_strips_code_fence() -> None:
    fenced = "```json\n" + json.dumps(_VALID_PLAN) + "\n```"
    plan = parse_plan(fenced)
    assert len(plan["phases"]) == 2


def test_parse_plan_strips_bare_code_fence() -> None:
    fenced = "```\n" + json.dumps(_VALID_PLAN) + "\n```"
    assert len(parse_plan(fenced)["phases"]) == 2


def test_parse_plan_trims_step_whitespace_and_drops_blanks() -> None:
    raw = {"phases": [{"name": "p", "steps": ["  a  ", "", "  "]}]}
    plan = parse_plan(json.dumps(raw))
    assert plan["phases"][0]["steps"] == ["a"]
    assert plan["summary"] == ""


def test_parse_plan_empty_raises() -> None:
    with pytest.raises(PlannerError, match="empty"):
        parse_plan("   ")


def test_parse_plan_invalid_json_raises() -> None:
    with pytest.raises(PlannerError, match="valid JSON"):
        parse_plan("not json {")


def test_parse_plan_non_object_raises() -> None:
    with pytest.raises(PlannerError, match="must be an object"):
        parse_plan(json.dumps([1, 2, 3]))


def test_parse_plan_missing_phases_raises() -> None:
    with pytest.raises(PlannerError, match="non-empty 'phases'"):
        parse_plan(json.dumps({"summary": "x"}))


def test_parse_plan_empty_phases_raises() -> None:
    with pytest.raises(PlannerError, match="non-empty 'phases'"):
        parse_plan(json.dumps({"phases": []}))


def test_parse_plan_phase_missing_name_raises() -> None:
    with pytest.raises(PlannerError, match="missing a non-empty 'name'"):
        parse_plan(json.dumps({"phases": [{"steps": ["x"]}]}))


def test_parse_plan_phase_without_steps_raises() -> None:
    with pytest.raises(PlannerError, match="no steps"):
        parse_plan(json.dumps({"phases": [{"name": "p", "steps": []}]}))


def test_parse_plan_phase_with_only_blank_steps_raises() -> None:
    with pytest.raises(PlannerError, match="no non-empty steps"):
        parse_plan(json.dumps({"phases": [{"name": "p", "steps": ["", "  "]}]}))


def test_parse_plan_phase_not_object_raises() -> None:
    with pytest.raises(PlannerError, match="phase 1 must be an object"):
        parse_plan(json.dumps({"phases": ["just a string"]}))


# --------------------------------------------------------------------------- #
# build_plan_prompt
# --------------------------------------------------------------------------- #


def test_build_prompt_includes_task_and_qa() -> None:
    state = _state(
        plan={"task": "Upgrade requests to 2.32"},
        qa_history=[{"question": "Pin exact?", "answer": "Yes, exact."}],
    )
    prompt = build_plan_prompt(state)
    assert "Upgrade requests to 2.32" in prompt
    assert "Pin exact?" in prompt
    assert "Yes, exact." in prompt
    assert "PLANNER" in prompt


def test_build_prompt_task_fallback_from_top_level() -> None:
    state = _state(task="Top-level task ask")  # type: ignore[typeddict-unknown-key]
    assert "Top-level task ask" in build_plan_prompt(state)


def test_build_prompt_handles_string_qa_entries() -> None:
    state = _state(plan={"task": "t"}, qa_history=["freeform note"])
    assert "freeform note" in build_plan_prompt(state)


def test_build_prompt_no_task_uses_placeholder() -> None:
    prompt = build_plan_prompt(_state())
    assert "(no task description provided)" in prompt


def test_build_prompt_loopback_includes_feedback_and_prior_plan() -> None:
    # Use the REAL verdict shape that review_loop.ReviewResult.to_dict() writes
    # ("verdict" + "findings") — NOT the old "decision"/"notes" keys, which never
    # existed on a real verdict and silently produced empty re-plan feedback.
    state = _state(
        plan={"task": "t", "phases": [{"name": "old", "steps": ["x"]}]},
        review_verdicts=[
            {
                "verdict": "request_changes",
                "outcome": "loop_back",
                "round_index": 1,
                "findings": "Phase 1 missing rollback.",
            }
        ],
    )
    prompt = build_plan_prompt(state)
    assert "Reviewer feedback" in prompt
    assert "Phase 1 missing rollback." in prompt  # the findings reached the re-plan
    assert "Previous plan" in prompt
    assert '"old"' in prompt


def test_review_feedback_reads_real_verdict_keys() -> None:
    # Regression: the planner must read the producer's keys (verdict/findings).
    # The bug read decision/notes/comment -> empty feedback -> the planner kept
    # re-introducing the rejected flaw -> review cap -> park.
    from agent_team.nodes.planner import _format_review_feedback

    rendered = _format_review_feedback(
        [{"verdict": "request_changes", "findings": "Add a teardown fixture."}]
    )
    assert "Add a teardown fixture." in rendered
    assert "request_changes" in rendered


def test_build_prompt_no_feedback_omits_review_sections() -> None:
    prompt = build_plan_prompt(_state(plan={"task": "t"}))
    assert "Reviewer feedback" not in prompt
    assert "Previous plan" not in prompt


# --------------------------------------------------------------------------- #
# plan_node — happy path
# --------------------------------------------------------------------------- #


def test_plan_node_produces_plan_and_advances_to_review() -> None:
    _bind_invoker(json.dumps(_VALID_PLAN))
    out = plan_node(_state(plan={"task": "do a thing"}))
    assert out["current_phase"] == Phase.REVIEW.value
    assert out["status"] == TaskStatus.ACTIVE.value
    assert out["plan"]["phases"][0]["name"] == "Phase 1 — bump"
    assert out["plan"]["revision"] == 0


def test_plan_node_returns_partial_state_only() -> None:
    _bind_invoker(json.dumps(_VALID_PLAN))
    out = plan_node(_state(plan={"task": "x"}))
    # A node returns only the keys it owns; it must not echo thread_id.
    assert set(out.keys()) == {"plan", "current_phase", "status"}


def test_plan_node_forwards_config_to_billing_seam() -> None:
    calls = _bind_invoker(json.dumps(_VALID_PLAN))
    plan_node(_state(plan={"task": "x"}), config={"billing_mode": "api"})
    assert calls[0]["mode"] is BillingMode.API


def test_plan_node_sends_task_into_prompt() -> None:
    calls = _bind_invoker(json.dumps(_VALID_PLAN))
    plan_node(_state(plan={"task": "UNIQUE-TASK-MARKER"}))
    assert "UNIQUE-TASK-MARKER" in calls[0]["prompt"]


def test_plan_node_garbled_reply_raises() -> None:
    _bind_invoker("not json at all")
    with pytest.raises(PlannerError):
        plan_node(_state(plan={"task": "x"}))


# --------------------------------------------------------------------------- #
# plan_node — loop-back / convergence bound (§3.3, §6.6)
# --------------------------------------------------------------------------- #


def test_plan_node_replans_on_loopback_and_counts_revision() -> None:
    _bind_invoker(json.dumps(_VALID_PLAN))
    state = _state(
        plan={"task": "x"},
        review_verdicts=[{"decision": "REQUEST_CHANGES", "notes": "fix it"}],
    )
    out = plan_node(state)
    assert out["current_phase"] == Phase.REVIEW.value
    assert out["plan"]["revision"] == 1


def test_plan_node_parks_after_max_revisions() -> None:
    # Invoker bound but must NOT be called once we are over the bound.
    calls = _bind_invoker(json.dumps(_VALID_PLAN))
    verdicts = [
        {"decision": "REQUEST_CHANGES", "notes": f"round {i}"}
        for i in range(MAX_PLAN_REVISIONS)
    ]
    out = plan_node(_state(plan={"task": "x"}, review_verdicts=verdicts))
    assert out["status"] == TaskStatus.PARKED.value
    assert out["current_phase"] == Phase.PARKED.value
    assert "plan" not in out
    assert calls == []  # no Claude budget burned past the bound


def test_plan_node_does_not_park_just_below_bound() -> None:
    _bind_invoker(json.dumps(_VALID_PLAN))
    verdicts = [
        {"decision": "REQUEST_CHANGES", "notes": f"round {i}"}
        for i in range(MAX_PLAN_REVISIONS - 1)
    ]
    out = plan_node(_state(plan={"task": "x"}, review_verdicts=verdicts))
    assert out["current_phase"] == Phase.REVIEW.value
    assert out["plan"]["revision"] == MAX_PLAN_REVISIONS - 1


def test_plan_node_ignores_non_request_changes_verdicts_for_bound() -> None:
    # APPROVE/other verdicts must not count toward the park bound.
    _bind_invoker(json.dumps(_VALID_PLAN))
    verdicts = [{"decision": "APPROVE"}] * (MAX_PLAN_REVISIONS + 2)
    out = plan_node(_state(plan={"task": "x"}, review_verdicts=verdicts))
    assert out["current_phase"] == Phase.REVIEW.value
    assert out["plan"]["revision"] == 0


def test_plan_node_counts_string_request_changes_verdicts() -> None:
    _bind_invoker(json.dumps(_VALID_PLAN))
    verdicts: list[Any] = ["REQUEST_CHANGES"] * MAX_PLAN_REVISIONS
    out = plan_node(_state(plan={"task": "x"}, review_verdicts=verdicts))
    assert out["status"] == TaskStatus.PARKED.value


def test_plan_node_unconfigured_invoker_raises() -> None:
    # The foundation seam fails loudly when no invoker is wired.
    billing._invoker = billing._unconfigured_invoker
    with pytest.raises(RuntimeError, match="no invoker bound"):
        plan_node(_state(plan={"task": "x"}))


# --------------------------------------------------------------------------- #
# plan_node — turn headroom + auto-retry-once (Phase A)
# --------------------------------------------------------------------------- #


def test_plan_node_passes_max_turns_to_invoke_seam() -> None:
    # The single-shot Claude default (1 turn) is flaky; the planner asks for
    # headroom so the model can FINISH its JSON.
    calls = _bind_invoker(json.dumps(_VALID_PLAN))
    plan_node(_state(plan={"task": "x"}))
    assert calls[0]["kw"].get("max_turns") == 4


def _bind_sequenced_invoker(
    replies: list[Any],
) -> list[dict[str, Any]]:
    """Bind an invoker that, per call, raises if the next item is an Exception
    or returns it as the reply text otherwise. Captures its calls."""
    calls: list[dict[str, Any]] = []
    queue = list(replies)

    def _fake(prompt: str, *, mode: BillingMode, **kw: Any) -> ClaudeResult:
        calls.append({"prompt": prompt, "mode": mode, "kw": kw})
        item = queue.pop(0)
        if isinstance(item, Exception):
            raise item
        return ClaudeResult(text=item, mode=mode, usage={"input_tokens": 1})

    billing.set_invoker(_fake)
    return calls


def test_plan_node_retries_once_on_turn_cap_then_succeeds() -> None:
    # First call exhausts the single turn; the retry produces a valid plan.
    calls = _bind_sequenced_invoker(
        [
            Exception(
                "Claude Code returned an error result: "
                "Reached maximum number of turns (1)"
            ),
            json.dumps(_VALID_PLAN),
        ]
    )
    out = plan_node(_state(plan={"task": "x"}))
    assert out["current_phase"] == Phase.REVIEW.value
    assert out["plan"]["phases"][0]["name"] == "Phase 1 — bump"
    assert len(calls) == 2  # exactly one retry


def test_plan_node_retries_once_on_empty_reply_then_succeeds() -> None:
    # An EMPTY model reply is transient (PlannerError "empty"): retry once.
    calls = _bind_sequenced_invoker(["   ", json.dumps(_VALID_PLAN)])
    out = plan_node(_state(plan={"task": "x"}))
    assert out["current_phase"] == Phase.REVIEW.value
    assert len(calls) == 2  # one retry


def test_plan_node_does_not_retry_on_malformed_json() -> None:
    # A non-empty but invalid reply is DETERMINISTIC: a retry would reproduce
    # it, so the planner fails fast on the first attempt.
    calls = _bind_sequenced_invoker(["not json {", json.dumps(_VALID_PLAN)])
    with pytest.raises(PlannerError):
        plan_node(_state(plan={"task": "x"}))
    assert len(calls) == 1  # no second call


def test_build_prompt_includes_no_tools_json_only_instruction() -> None:
    prompt = build_plan_prompt(_state(plan={"task": "x"}))
    assert "Do NOT use any tools" in prompt
    assert "ONLY the JSON object" in prompt