nisse/nisse-dev/tests/test_context.py

"""The context budget (architecture L6)."""

import io
from pathlib import Path
from typing import Any

import pytest

from nisse_dev.canonical import (
    AssistantText,
    CanonicalEvent,
    StoredEvent,
    ToolCall,
    ToolResult,
    UserMessage,
)
from nisse_dev.catalog import ModelCatalog, UnknownLimits
from nisse_dev.context import (
    Budget,
    BudgetedClient,
    ContextOverflow,
    estimate_tokens,
    fit,
    head_and_tail,
)
from nisse_dev.loop.driver import Driver
from nisse_dev.oplog import OpLog
from nisse_dev.state import StateStore


def stored(events: list[CanonicalEvent]) -> list[StoredEvent]:
    return [StoredEvent(id=i + 1, ts="t", event=e) for i, e in enumerate(events)]


def turn(n: int, output: str) -> list[CanonicalEvent]:
    """One executor turn: a command and its (possibly long) output."""
    return [
        ToolCall(call_id=f"c{n}", tool="run_command", arguments={"command": "make"}),
        ToolResult(call_id=f"c{n}", result=output, success=True),
    ]


def transcript(turns: int, output_chars: int) -> list[StoredEvent]:
    events: list[CanonicalEvent] = [UserMessage(text="start milestone 1", source="t")]
    for n in range(turns):
        events += turn(n, f"output {n} " + "x" * output_chars)
    return stored(events)


def budget(window: int, keep: int = 2) -> Budget:
    return Budget(
        window=window, max_output=500, keep_recent_turns=keep, safety_margin=0.1
    )


def results(events: list[StoredEvent]) -> dict[str, str]:
    return {
        e.event.call_id: e.event.result
        for e in events
        if isinstance(e.event, ToolResult)
    }


# --- keeping the end of long output ---------------------------------------


def test_head_and_tail_keeps_the_failure_summary_at_the_end() -> None:
    output = (
        "exit code: 1\nstdout:\n"
        + "collected 400 items\n" * 500
        + "FAILED tests/test_app.py::test_greet - AssertionError\n"
        + "==== 1 failed, 399 passed ====\n"
    )
    kept = head_and_tail(output, 1000)
    assert len(kept) < 1100
    assert kept.startswith("exit code: 1")  # the head keeps the exit code
    assert "FAILED tests/test_app.py::test_greet" in kept
    assert "1 failed, 399 passed" in kept
    assert "characters omitted" in kept


def test_short_text_is_untouched() -> None:
    assert head_and_tail("fine", 1000) == "fine"


# --- fitting, step by step -------------------------------------------------


def test_a_transcript_that_fits_is_sent_unchanged() -> None:
    events = transcript(turns=3, output_chars=100)
    assert fit(events, "system", None, budget(100_000)) == events


def test_step_one_cuts_older_tool_results_and_keeps_recent_ones_whole() -> None:
    events = transcript(turns=6, output_chars=3000)
    fitted = fit(events, "system", None, budget(4000, keep=2))

    kept = results(fitted)
    assert len(kept) == 6  # nothing dropped yet
    for older in ("c0", "c1", "c2", "c3"):
        assert "characters omitted" in kept[older]
    for recent in ("c4", "c5"):
        assert "characters omitted" not in kept[recent]


def test_step_two_drops_the_oldest_turns_and_says_so() -> None:
    events = transcript(turns=20, output_chars=3000)
    fitted = fit(events, "system", None, budget(4000, keep=2))

    first = fitted[0].event
    assert isinstance(first, UserMessage)
    assert "omitted to fit the context window" in first.text
    kept = results(fitted)
    assert "c0" not in kept  # the oldest went
    assert {"c18", "c19"} <= kept.keys()  # the recent turns stayed
    assert "characters omitted" not in kept["c19"]
    # Every surviving result still has its call: a provider rejects an
    # answer to a call it never saw.
    calls = {e.event.call_id for e in fitted if isinstance(e.event, ToolCall)}
    assert kept.keys() <= calls
    assert estimate_tokens(fitted, "system", None) <= budget(4000).for_prompt


def test_step_three_refuses_what_cannot_fit() -> None:
    # Even the two most recent turns are bigger than the whole window.
    events = transcript(turns=3, output_chars=20_000)
    with pytest.raises(ContextOverflow, match="last 2 turns"):
        fit(events, "system", None, budget(2000, keep=2))


def test_a_huge_system_prompt_alone_overflows() -> None:
    events = transcript(turns=1, output_chars=10)
    with pytest.raises(ContextOverflow):
        fit(events, "s" * 50_000, None, budget(4096))


def test_the_output_ceiling_and_margin_come_out_of_the_window() -> None:
    b = Budget(window=4096, max_output=1024, keep_recent_turns=6, safety_margin=0.1)
    assert b.for_prompt == int(4096 * 0.9) - 1024


# --- the client wrapper ----------------------------------------------------


class RecordingClient:
    def __init__(self, reported_input: int = 0) -> None:
        self.calls: list[dict[str, Any]] = []
        self.last_usage = {"input": reported_input, "output": 5}

    def complete(self, events, tools, model, system=None, *, max_tokens=None, **kwargs):
        self.calls.append({"events": events, "max_tokens": max_tokens, **kwargs})
        return [AssistantText(text="ok", model=model)]


def test_the_ceiling_is_the_smaller_of_config_and_model() -> None:
    inner = RecordingClient()
    client = BudgetedClient(inner, lambda m: (4096, 1024), max_output_tokens=16384)
    client.complete(transcript(1, 10), None, "llama3.2")
    assert inner.calls[0]["max_tokens"] == 1024

    client = BudgetedClient(inner, lambda m: (200_000, 64_000), max_output_tokens=16384)
    client.complete(transcript(1, 10), None, "claude-haiku-4-5")
    assert inner.calls[1]["max_tokens"] == 16384


def test_the_wrapper_sends_the_fitted_events() -> None:
    inner = RecordingClient()
    client = BudgetedClient(
        inner, lambda m: (4000, 500), keep_recent_turns=2, safety_margin=0.1
    )
    client.complete(transcript(20, 3000), None, "m", system="system")
    sent = inner.calls[0]["events"]
    assert "omitted to fit" in sent[0].event.text


def test_reported_usage_only_ever_raises_the_estimate() -> None:
    events = transcript(1, 10)
    estimate = estimate_tokens(events, None, None)

    under = RecordingClient(reported_input=estimate * 2)  # we undercounted
    client = BudgetedClient(under, lambda m: (100_000, None))
    client.complete(events, None, "m")
    assert client._factors["m"] == pytest.approx(2.0, rel=0.05)

    over = RecordingClient(reported_input=estimate // 2)  # we overcounted
    under_then_over = BudgetedClient(over, lambda m: (100_000, None))
    under_then_over._factors["m"] = 2.0
    under_then_over.complete(events, None, "m")
    assert under_then_over._factors["m"] == 2.0  # never loosened


def test_last_usage_passes_through() -> None:
    client = BudgetedClient(RecordingClient(reported_input=7), lambda m: (100, None))
    assert client.last_usage["input"] == 7


# --- the limits come from LiteLLM's catalog --------------------------------


CATALOG = ModelCatalog.from_payload(
    {
        "data": [
            {
                "model_name": "claude-haiku-4-5",
                "litellm_params": {"model": "anthropic/claude-haiku-4-5"},
                "model_info": {"max_input_tokens": 200000, "max_output_tokens": 64000},
            },
            {
                "model_name": "llama3.2",
                "litellm_params": {"model": "ollama/llama3.2"},
                "model_info": {"max_input_tokens": 4096, "max_output_tokens": 1024},
            },
            {
                "model_name": "undeclared",
                "litellm_params": {"model": "ollama/qwen3"},
                "model_info": {"max_input_tokens": None},
            },
        ]
    }
)


def test_the_catalog_reads_each_models_limits() -> None:
    facts = CATALOG.facts("llama3.2")
    assert (facts.max_input_tokens, facts.max_output_tokens) == (4096, 1024)
    assert CATALOG.facts("undeclared").max_input_tokens is None


def test_a_role_on_a_model_without_a_window_is_refused() -> None:
    CATALOG.require_limits({"driver": "llama3.2", "helper": "claude-haiku-4-5"})
    with pytest.raises(UnknownLimits, match="driver='undeclared'"):
        CATALOG.require_limits({"driver": "undeclared", "helper": "llama3.2"})


def test_a_request_that_cannot_fit_ends_the_run(tmp_path: Path) -> None:
    # Never sent: the driver's run fails and the item waits on the operator,
    # like any failed call, rather than a server getting a prompt it would
    # silently cut.
    store = StateStore(tmp_path / "state.db", oplog=OpLog("t", io.StringIO()))
    store.transcript("t", "operator").append(UserMessage(text="go", source="t"))
    inner = RecordingClient()
    client = BudgetedClient(inner, lambda m: (4096, 1024))
    outcome = Driver(client, "m").run(
        store.transcript("t", "driver"),
        system="s" * 20_000,  # a system prompt too big to fit
        tools=[],
        workspace=None,
    )

    assert outcome.kind == "failed"
    assert "fit" in outcome.note
    assert inner.calls == []


# --- tool_choice is only asked of servers that enforce it -------------------

FINISH = [{"type": "function", "function": {"name": "finish", "parameters": {}}}]


def test_required_is_asked_only_of_models_declared_to_enforce_it() -> None:
    # Measured: qwen3-14b under "required" cannot stop and rambles to the
    # output ceiling, so a model that is not declared is never asked.
    inner = RecordingClient()
    client = BudgetedClient(
        inner, lambda m: (32768, 8192), enforces_tool_choice=lambda m: m == "local-fast"
    )
    client.complete(transcript(1, 10), FINISH, "local-fast")
    client.complete(transcript(1, 10), FINISH, "local-coder")
    client.complete(transcript(1, 10), None, "local-fast")
    assert [c["require_tool_call"] for c in inner.calls] == [True, False, False]


def test_the_catalog_reads_supports_tool_choice() -> None:
    catalog = ModelCatalog.from_payload(
        {
            "data": [
                {"model_name": "a", "model_info": {"supports_tool_choice": True}},
                {"model_name": "b", "model_info": {"supports_tool_choice": False}},
                {"model_name": "c", "model_info": {}},
            ]
        }
    )
    assert [catalog.facts(m).supports_tool_choice for m in "abc"] == [True, False, None]


def test_a_fallback_answer_is_recognised() -> None:
    facts = ModelCatalog.from_payload(
        {
            "data": [
                {
                    "model_name": "claude-haiku-4-5",
                    "litellm_params": {"model": "anthropic/claude-haiku-4-5"},
                }
            ]
        }
    ).facts("claude-haiku-4-5")
    assert facts.answered("claude-haiku-4-5")  # LiteLLM's alias
    assert facts.answered("claude-haiku-4-5-20251001")  # the provider's own name
    assert facts.answered(None)  # nothing reported
    assert not facts.answered("claude-sonnet-5")  # somebody else answered