nisse/nisse-dev/tests/test_context.py

"""The context budget (spec A10)."""

import io
from pathlib import Path
from typing import Any

import pytest

from nisse_dev.canonical import (
    AssistantText,
    CanonicalEvent,
    GateOpened,
    PlanSubmitted,
    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.executor import LLMExecutor
from nisse_dev.llm import RoleMap
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):
        self.calls.append({"events": events, "max_tokens": max_tokens})
        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({"executor": "llama3.2", "planner": "claude-haiku-4-5"})
    with pytest.raises(UnknownLimits, match="executor='undeclared'"):
        CATALOG.require_limits({"executor": "undeclared", "planner": "llama3.2"})


# --- the executor sees the approved plan and what came after ---------------


def test_the_executor_sees_only_its_milestone(tmp_path: Path) -> None:
    store = StateStore(tmp_path / "state.db", oplog=OpLog("t", io.StringIO()))
    for event in (
        UserMessage(text="use greet: a long planning conversation", source="t"),
        PlanSubmitted(plan="arch plan", model="m"),
        UserMessage(text="approve", source="t"),
        PlanSubmitted(plan="milestone 1: write greet.py", model="m"),
        GateOpened(gate="milestone-1", reason="plan submitted"),
        UserMessage(text="start milestone 1", source="t"),
    ):
        store.append_event(event)

    class Workspace:
        def describe(self) -> str:
            return "workspace"

    seen: dict[str, Any] = {}

    class Client:
        def complete(self, events, tools, model, system=None):
            seen["events"] = [e.event for e in events]
            seen["system"] = system
            return [AssistantText(text="done", model=model)]

    roles = RoleMap(executor="m", planner="m", delegate="m", consultant="m")
    LLMExecutor(Client(), store, roles, Workspace()).run(1)

    texts = [getattr(e, "text", None) for e in seen["events"]]
    assert "start milestone 1" in texts
    assert "use greet: a long planning conversation" not in texts
    assert not any(isinstance(e, PlanSubmitted) for e in seen["events"])
    # The plan rides in the system prompt, which the budget never drops.
    assert "The approved plan:\nmilestone 1: write greet.py" in seen["system"]
    assert "arch plan" not in seen["system"]


def test_a_request_that_cannot_fit_stops_the_milestone(tmp_path: Path) -> None:
    # Never sent: the executor stops and asks the operator, like any failed
    # call, rather than handing a server a prompt it would silently cut.
    store = StateStore(tmp_path / "state.db", oplog=OpLog("t", io.StringIO()))
    store.append_event(UserMessage(text="start milestone 1", source="t"))

    class Workspace:
        def describe(self) -> str:
            return "workspace " + "w" * 20_000  # a system prompt too big to fit

    inner = RecordingClient()
    client = BudgetedClient(inner, lambda m: (4096, 1024))
    roles = RoleMap(executor="m", planner="m", delegate="m", consultant="m")
    outcome = LLMExecutor(client, store, roles, Workspace()).run(1)

    assert outcome.status == "escalate"
    assert "fit" in outcome.note
    assert inner.calls == []