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:
    state = _state(
        plan={"task": "t", "phases": [{"name": "old", "steps": ["x"]}]},
        review_verdicts=[
            {"decision": "REQUEST_CHANGES", "notes": "Phase 1 missing rollback."}
        ],
    )
    prompt = build_plan_prompt(state)
    assert "Reviewer feedback" in prompt
    assert "Phase 1 missing rollback." in prompt
    assert "Previous plan" in prompt
    assert '"old"' in prompt


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"}))