"""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 == []