orchestrator/agent-team/tests/test_clarifier_llm.py

"""Unit tests for agent_team.nodes.clarifier_llm (§3.3, §7.1 P1).

The Claude-backed clarifier callables are exercised with a FAKE invoke that
returns canned :class:`~agent_team.billing.ClaudeResult` text — no network. The
load-bearing properties under test:

* **One call per turn (memoization).** ``assess_confidence`` then
  ``generate_questions`` on the same turn must reuse a single Claude call.
* **Fail SAFE (the human gate).** A garbled / non-JSON reply must yield
  confidence ``0.0`` (never >= the 0.98 bar) and a non-empty fallback question.
* **Defensive parsing.** ```json fences and surrounding prose still parse, and
  out-of-range confidence is clamped to ``[0, 1]``.
* **Integration smoke.** The callables wire into the real
  :func:`~agent_team.nodes.clarifier.make_clarifier_node` and clear the gate
  once confidence rises across turns.
"""

from __future__ import annotations

import json
from typing import Any

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, StateGraph
from langgraph.types import Command

from agent_team.billing import BillingMode, ClaudeResult
from agent_team.nodes.clarifier import (
    DEFAULT_CONFIDENCE_THRESHOLD,
    make_clarifier_node,
)
from agent_team.nodes.clarifier_llm import (
    FALLBACK_QUESTION,
    ClaudeClarifier,
    build_claude_clarifier_callables,
)
from agent_team.task_model import Phase, PipelineState, TaskStatus

# --------------------------------------------------------------------------- #
# Fakes / helpers
# --------------------------------------------------------------------------- #


class _FakeInvoke:
    """A fake billing.claude_invoke that returns canned text and counts calls.

    ``replies`` may be a single string (returned every call) or a list of
    strings (consumed one per call, last one repeating) so a test can simulate
    rising confidence across turns.
    """

    def __init__(self, replies: str | list[str]) -> None:
        self._replies = [replies] if isinstance(replies, str) else list(replies)
        self.calls: list[dict[str, Any]] = []

    def __call__(self, prompt: str, **kw: Any) -> ClaudeResult:
        idx = min(len(self.calls), len(self._replies) - 1)
        text = self._replies[idx]
        self.calls.append({"prompt": prompt, "kw": kw})
        return ClaudeResult(text=text, mode=BillingMode.SUBSCRIPTION)


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


def _json(confidence: Any, questions: Any, rationale: str = "because") -> str:
    return json.dumps(
        {"confidence": confidence, "questions": questions, "rationale": rationale}
    )


# --------------------------------------------------------------------------- #
# High confidence: assess returns ~value AND the call is reused (memoization).
# --------------------------------------------------------------------------- #


def test_high_confidence_parsed() -> None:
    fake = _FakeInvoke(_json(0.99, []))
    clar = ClaudeClarifier(invoke=fake)
    assert clar.assess_confidence([], _state()) == 0.99


def test_single_call_per_turn_memoized() -> None:
    fake = _FakeInvoke(_json(0.99, []))
    clar = ClaudeClarifier(invoke=fake)
    qa: list[object] = []

    # Both methods called for the same turn -> exactly ONE Claude call.
    conf = clar.assess_confidence(qa, _state())
    questions = clar.generate_questions(qa, _state())

    assert conf == 0.99
    assert questions == []  # confident, no questions needed
    assert len(fake.calls) == 1


def test_memo_recomputes_when_history_grows() -> None:
    fake = _FakeInvoke([_json(0.10, ["q1"]), _json(0.99, [])])
    clar = ClaudeClarifier(invoke=fake)

    # Turn 0: one answer-less call.
    assert clar.assess_confidence([], _state()) == 0.10
    assert clar.generate_questions([], _state()) == ["q1"]
    assert len(fake.calls) == 1

    # Turn 1: history grew -> a fresh call, now confident.
    assert clar.assess_confidence(["a1"], _state()) == 0.99
    assert clar.generate_questions(["a1"], _state()) == []
    assert len(fake.calls) == 2


def test_memo_isolates_across_tasks_no_cross_gate_clear() -> None:
    """A second task at the same history length must NOT reuse task A's memo.

    Regression for the fail-OPEN collision: one ClaudeClarifier instance serves
    every task through the shared graph node, so keying the memo on history
    length alone would let task A's cached 0.99 clear task B's human gate with
    no model call. Keying on thread_id forces a fresh assessment for task B.
    """
    fake = _FakeInvoke([_json(0.99, []), _json(0.10, ["need more from B"])])
    clar = ClaudeClarifier(invoke=fake)

    # Task A (thread t-A), empty history -> confident, cached.
    assert clar.assess_confidence([], _state(thread_id="t-A")) == 0.99
    assert len(fake.calls) == 1

    # Task B (thread t-B), SAME empty history/length -> must re-assess, NOT
    # inherit A's cache, so its low confidence holds and the gate stays shut.
    assert clar.assess_confidence([], _state(thread_id="t-B")) == 0.10
    assert clar.generate_questions([], _state(thread_id="t-B")) == ["need more from B"]
    assert len(fake.calls) == 2  # a real second call happened for task B


# --------------------------------------------------------------------------- #
# Low confidence: below the bar, questions are returned.
# --------------------------------------------------------------------------- #


def test_low_confidence_returns_questions() -> None:
    fake = _FakeInvoke(_json(0.40, ["What is the scope?", "Which repo?"]))
    clar = ClaudeClarifier(invoke=fake)

    assert clar.assess_confidence([], _state()) < DEFAULT_CONFIDENCE_THRESHOLD
    assert clar.generate_questions([], _state()) == [
        "What is the scope?",
        "Which repo?",
    ]
    assert len(fake.calls) == 1


# --------------------------------------------------------------------------- #
# Malformed output: fail SAFE (0.0 confidence, non-empty fallback questions).
# --------------------------------------------------------------------------- #


def test_malformed_output_fails_safe() -> None:
    fake = _FakeInvoke("I'm sorry, I cannot help with that. <no json here>")
    clar = ClaudeClarifier(invoke=fake)

    assert clar.assess_confidence([], _state()) == 0.0
    questions = clar.generate_questions([], _state())
    assert questions == [FALLBACK_QUESTION]
    assert questions  # non-empty


def test_garbage_never_clears_the_gate() -> None:
    # The critical safety property: garbage must never read >= 0.98.
    for garbage in ["", "   ", "not json", "{broken", "[1,2,3]", "null", "42"]:
        clar = ClaudeClarifier(invoke=_FakeInvoke(garbage))
        conf = clar.assess_confidence([], _state())
        assert conf < DEFAULT_CONFIDENCE_THRESHOLD
        assert conf == 0.0
        assert clar.generate_questions([], _state()) == [FALLBACK_QUESTION]


def test_missing_confidence_key_defaults_zero() -> None:
    fake = _FakeInvoke(json.dumps({"questions": ["q?"], "rationale": "x"}))
    clar = ClaudeClarifier(invoke=fake)
    assert clar.assess_confidence([], _state()) == 0.0
    # Questions present in the reply are kept as-is.
    assert clar.generate_questions([], _state()) == ["q?"]


def test_low_confidence_empty_questions_gets_fallback() -> None:
    # Below the bar but model gave no questions -> generic fallback so the loop
    # always has something to ask.
    fake = _FakeInvoke(_json(0.20, []))
    clar = ClaudeClarifier(invoke=fake)
    assert clar.generate_questions([], _state()) == [FALLBACK_QUESTION]


# --------------------------------------------------------------------------- #
# Defensive parsing: fences and surrounding prose still parse.
# --------------------------------------------------------------------------- #


def test_json_in_code_fence_is_parsed() -> None:
    fenced = "```json\n" + _json(0.97, ["q?"]) + "\n```"
    clar = ClaudeClarifier(invoke=_FakeInvoke(fenced))
    assert clar.assess_confidence([], _state()) == 0.97
    assert clar.generate_questions([], _state()) == ["q?"]


def test_json_wrapped_in_prose_is_parsed() -> None:
    prose = (
        "Sure! Here is my assessment:\n"
        + _json(0.55, ["Clarify the deadline?"])
        + "\nLet me know if that helps."
    )
    clar = ClaudeClarifier(invoke=_FakeInvoke(prose))
    assert clar.assess_confidence([], _state()) == 0.55
    assert clar.generate_questions([], _state()) == ["Clarify the deadline?"]


def test_plain_json_fence_without_lang_is_parsed() -> None:
    fenced = "```\n" + _json(0.33, ["q?"]) + "\n```"
    clar = ClaudeClarifier(invoke=_FakeInvoke(fenced))
    assert clar.assess_confidence([], _state()) == 0.33


# --------------------------------------------------------------------------- #
# Confidence clamping into [0, 1].
# --------------------------------------------------------------------------- #


def test_confidence_above_one_is_clamped() -> None:
    clar = ClaudeClarifier(invoke=_FakeInvoke(_json(1.5, [])))
    assert clar.assess_confidence([], _state()) == 1.0


def test_confidence_below_zero_is_clamped() -> None:
    clar = ClaudeClarifier(invoke=_FakeInvoke(_json(-0.2, ["q?"])))
    assert clar.assess_confidence([], _state()) == 0.0


def test_confidence_as_string_is_coerced() -> None:
    clar = ClaudeClarifier(invoke=_FakeInvoke(_json("0.88", ["q?"])))
    assert clar.assess_confidence([], _state()) == 0.88


# --------------------------------------------------------------------------- #
# Prompt assembly pulls task/repo/context out of state.
# --------------------------------------------------------------------------- #


def test_prompt_includes_task_repo_and_qa() -> None:
    fake = _FakeInvoke(_json(0.99, []))
    clar = ClaudeClarifier(invoke=fake)
    state = _state(task="Add a webhook verifier", repo="agent-team")
    clar.assess_confidence(["prior answer"], state)

    prompt = fake.calls[0]["prompt"]
    assert "Add a webhook verifier" in prompt
    assert "agent-team" in prompt
    assert "prior answer" in prompt


# --------------------------------------------------------------------------- #
# Factory returns the exact ConfidenceAssessor / QuestionGenerator pair.
# --------------------------------------------------------------------------- #


def test_factory_returns_shared_memoized_pair() -> None:
    fake = _FakeInvoke(_json(0.45, ["q?"]))
    assess, generate = build_claude_clarifier_callables(invoke=fake)

    # Both come from one shared instance -> one call serves both this turn.
    assert assess([], _state()) == 0.45
    assert generate([], _state()) == ["q?"]
    assert len(fake.calls) == 1


# --------------------------------------------------------------------------- #
# Integration smoke: wire into the real make_clarifier_node, gate clears.
# --------------------------------------------------------------------------- #


def _build_app(node):
    graph = StateGraph(PipelineState)
    graph.add_node("clarify", node)
    graph.add_edge(START, "clarify")
    graph.add_edge("clarify", END)
    return graph.compile(checkpointer=MemorySaver())


def test_node_clears_gate_when_confidence_rises() -> None:
    # Turn 0 (no answers): low confidence, asks. Turn 1 (one answer): confident.
    fake = _FakeInvoke([_json(0.20, ["What is the goal?"]), _json(0.99, [])])
    assess, generate = build_claude_clarifier_callables(invoke=fake)
    node = make_clarifier_node(assess_confidence=assess, generate_questions=generate)
    app = _build_app(node)
    cfg = {"configurable": {"thread_id": "t-1"}}

    first = app.invoke(_state(thread_id="t-1"), cfg)
    assert "__interrupt__" in first  # suspended on the question-set

    final = app.invoke(Command(resume="ship feature X"), cfg)
    assert "__interrupt__" not in final
    assert final["qa_history"] == ["ship feature X"]
    assert final["current_phase"] == Phase.PLAN.value
    assert final["status"] == TaskStatus.ACTIVE.value


def test_node_parks_when_garbage_never_clears_gate() -> None:
    # A model that only ever emits garbage must NEVER open the gate; the loop
    # asks until the turn cap and parks (human gate stays shut).
    fake = _FakeInvoke("garbage, no json")
    assess, generate = build_claude_clarifier_callables(invoke=fake)
    from agent_team.nodes.clarifier import ClarifierConfig

    node = make_clarifier_node(
        assess_confidence=assess,
        generate_questions=generate,
        config=ClarifierConfig(max_turns=2),
    )
    app = _build_app(node)
    cfg = {"configurable": {"thread_id": "t-1"}}

    assert "__interrupt__" in app.invoke(_state(thread_id="t-1"), cfg)
    assert "__interrupt__" in app.invoke(Command(resume="a1"), cfg)
    final = app.invoke(Command(resume="a2"), cfg)

    assert "__interrupt__" not in final
    assert final["current_phase"] == Phase.PARKED.value
    assert final["status"] == TaskStatus.PARKED.value
    assert final["current_phase"] != Phase.PLAN.value