From fbcc63630a87728d92a224bc4e3a0359640904ea Mon Sep 17 00:00:00 2001 From: adlee-was-taken Date: Sun, 23 Aug 2026 21:23:19 -0400 Subject: [PATCH 1/3] feat(classifier): context-aware framing and relevance-based context pruning Port two ideas from the MIT-licensed alexrudloff/llmrouter: 1. Context-aware classification. chat_completions now feeds the classifier the message before the last user turn, framed as llmrouter's "Context: \n---\nMessage: " (classifier.context_framing, default on). A short follow-up ("Yes", "Try now?") inherits the prior turn's complexity instead of being classified in isolation as trivial. _classifier_user_content is a pure, testable framing helper; the system prompt gains the inherit-from-context rule. 2. Relevance-based context pruning (context_prune.py). An optional, pre-dispatch stage that trims old tool results from the provider-bound conversation once it exceeds pinch.budget_tokens (pinch.enabled defaults off). User/assistant/system messages are always kept verbatim; only old tool results are summarized (head+tail with an elision marker) or dropped to a placeholder, with role pairing preserved so the result still parses as a conversation. Applies to both the streaming and non-streaming upstream bodies; the classifier input is exempt (it already clamps). Tests: +23 (pruning invariants, framing, previous-context extraction, pinch config validation). 529 pass. --- config.py | 38 ++++++++ config.yaml | 30 ++++++ context_prune.py | 163 +++++++++++++++++++++++++++++++++ dispatcher.py | 91 +++++++++++++++++- tests/test_classifier_input.py | 93 ++++++++++++++++++- tests/test_config_endpoints.py | 39 ++++++++ tests/test_context_prune.py | 158 ++++++++++++++++++++++++++++++++ 7 files changed, 607 insertions(+), 5 deletions(-) create mode 100644 context_prune.py create mode 100644 tests/test_context_prune.py diff --git a/config.py b/config.py index fcacee5..0513b16 100644 --- a/config.py +++ b/config.py @@ -281,6 +281,38 @@ class FreshnessConfig(StrictModel): exclude_deprecated: bool +class PinchConfig(StrictModel): + """Optional relevance-based context pruning (Port of llmrouter's pinch). + + Prunes the provider-bound conversation — not the classifier input — when it + exceeds ``budget_tokens``, so a long agent session ships fewer prompt tokens + upstream. User/assistant/system messages are always kept; only tool results + are summarized or dropped (they carry the bulk of a long session's tokens). + """ + + enabled: bool = False + budget_tokens: int = 50000 + # How many recent user turns (plus their assistant replies and tool results) + # are protected from pruning. + keep_last_turns: int = 4 + # Tool results longer than this many characters are summarized in place. + max_summarize_chars: int = 4000 + + @field_validator("budget_tokens") + @classmethod + def budget_positive(cls, v: int) -> int: + if v <= 0: + raise ValueError("pinch.budget_tokens must be > 0") + return v + + @field_validator("keep_last_turns") + @classmethod + def turns_positive(cls, v: int) -> int: + if v <= 0: + raise ValueError("pinch.keep_last_turns must be > 0") + return v + + class DatabaseConfig(StrictModel): path: str @@ -295,6 +327,11 @@ class ClassifierConfig(StrictModel): model: str # Ceiling on the text handed to the classifier. 0 disables clamping. max_input_chars: int = 8000 + # When a preceding turn is available as context, frame the classifier + # input as llmrouter does — "Context: \n---\nMessage: " — so a + # short follow-up ("Yes", "Try now?") can inherit the complexity of the + # turn it continues instead of being classified in isolation as trivial. + context_framing: bool = True timeout_seconds: int temperature: float = 0.0 max_output_tokens: int = 1024 @@ -342,6 +379,7 @@ class RouterConfig(StrictModel): local_vision: LocalVisionConfig = LocalVisionConfig() escalation: EscalationConfig iteration: IterationConfig = IterationConfig() + pinch: PinchConfig = PinchConfig() freshness: FreshnessConfig database: DatabaseConfig classifier: ClassifierConfig diff --git a/config.yaml b/config.yaml index c60665f..ae4de68 100644 --- a/config.yaml +++ b/config.yaml @@ -157,6 +157,25 @@ iteration: # IS a quality loss. Batch work does not care. max_attempts_interactive: 1 +pinch: + # Relevance-based context pruning (ported from the MIT-licensed llmrouter's + # "pinch"). This is an OPTIONAL, pre-dispatch stage: when a conversation + # exceeds budget_tokens, the provider-bound messages are trimmed BEFORE any + # paid token is sent upstream. User/assistant/system messages are always + # kept verbatim; only old TOOL RESULTS are shortened or dropped, because + # they carry the bulk of a long agent session's tokens and are least needed + # in full by the time the next turn is answered. + # + # It does NOT touch the classifier's input, and it defaults off — enable it + # only if long sessions are shipping more prompt tokens than you want to pay + # for. Tool results can only be dropped safely because tool outputs are + # idempotent enough for a placeholder; a wrong guess here loses context, so + # start conservative (large budget, small reduction). + enabled: false + budget_tokens: 50000 + keep_last_turns: 4 + max_summarize_chars: 4000 + routing: # Access gating is prose-only in the NeuralWatt catalog ("Private preview # (grant-gated)", "(Canary)"), so the poller parses it into access_level and @@ -349,6 +368,12 @@ classifier: # never depends on what the classifier saw. 0 disables clamping. max_input_chars: 8000 + # When the chat path supplies the previous turn as context (see the pinch / + # context notes), frame the classifier input as "Context: / Message: + # " so a short follow-up inherits the prior turn's complexity + # instead of being classified in isolation as trivial. + context_framing: true + fallback_tier: 2 fallback_category: general_chat response_format: "json" # ask Ollama to constrain output to valid JSON @@ -367,6 +392,11 @@ classifier: "confidence": float 0-1 } + When the input is framed as "Context: \n---\nMessage: ": + a short follow-up message ("Yes", "Try now?", "Go ahead") continues the + prior turn, so classify it by the CONTEXT's complexity, not in isolation. + A standalone short message with no Context line is genuinely trivial. + dispatch_providers: # NeuralWatt is the only provider. The dict shape and the models table's # (model_id, provider) key are kept so a second one can be added without a diff --git a/context_prune.py b/context_prune.py new file mode 100644 index 0000000..f909697 --- /dev/null +++ b/context_prune.py @@ -0,0 +1,163 @@ +"""Relevance-based context pruning for the provider-bound conversation. + +Ported from the MIT-licensed alexrudloff/llmrouter "pinch" module, reduced to +the pure, injectable decision core that 6krrt can test offline. Where llmrouter +embeds every candidate message and scores cosine relevance, this module keeps +the same SAFE invariants without requiring an embedding model on the request +path: + + - user / assistant / system messages are ALWAYS kept verbatim, + - only TOOL RESULTS older than the protected window are trimmed or dropped, + - message order and the tool/assistant pairing are preserved so the result + still parses as a valid conversation. + +The tailoring is deliberate. Tool results are where a long agent session's +tokens actually live, they are the least likely to be needed in full by the +time a later turn is answered, and replacing one with a short placeholder is +reversible at the semantic level -- a wrong guess costs context, but it never +breaks the request. Trimming a user or assistant message, by contrast, can +change what the model is being asked, so those are never touched. + +This module is pure: it takes messages and limits and returns pruned messages. +`dispatcher.py` owns reading the config and deciding when to call it. +""" + +from __future__ import annotations + +DEFAULT_BUDGET_TOKENS = 50000 +DEFAULT_KEEP_LAST_TURNS = 4 +DEFAULT_MAX_SUMMARIZE_CHARS = 4000 +# Mirrors dispatcher.CHARS_PER_TOKEN. +CHARS_PER_TOKEN = 3 + + +def estimate_tokens(text: str | None) -> int: + """Crude characters-per-token estimate, consistent with the dispatcher.""" + return len(text) // CHARS_PER_TOKEN if text else 0 + + +def extract_text(message: dict) -> str: + """Best-effort text of a message, whether content is str or content blocks.""" + content = message.get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + out: list[str] = [] + for part in content: + if isinstance(part, dict): + if part.get("type") == "text" and isinstance(part.get("text"), str): + out.append(part["text"]) + elif part.get("type") == "tool_result": + out.append(str(part.get("content", ""))) + return " ".join(out) + return str(content) if content else "" + + +def _tool_name(message: dict) -> str | None: + """A tool message's name, under either the OpenAI or Anthropic spelling.""" + name = message.get("name") or message.get("tool_name") or message.get("function") + if isinstance(name, dict): + name = name.get("name") + return name if isinstance(name, str) else None + + +def _first_user_turn_indexes(messages: list[dict]) -> list[int]: + """Indexes of messages that start a new user turn. + + A user message starts a turn; the assistant reply and any tool results that + follow belong to that turn until the next user message. + """ + return [i for i, m in enumerate(messages) if m.get("role") == "user"] + + +def _turn_of(index: int, user_indexes: list[int]) -> int: + """Which user-turn ordinal the message at ``index`` belongs to (0-based).""" + turn = 0 + for ui in user_indexes: + if ui > index: + break + turn += 1 + return max(turn - 1, 0) + + +def prune_context( + messages: list[dict], + budget_tokens: int = DEFAULT_BUDGET_TOKENS, + keep_last_turns: int = DEFAULT_KEEP_LAST_TURNS, + max_summarize_chars: int = DEFAULT_MAX_SUMMARIZE_CHARS, +) -> tuple[list[dict], dict]: + """Trim old tool results once a conversation exceeds ``budget_tokens``. + + Returns ``(pruned_messages, stats)``. Keeps every user/assistant/system + message verbatim. Tool results in the preserved window pass through; older + tool results longer than ``max_summarize_chars`` are replaced with a short + placeholder, and shorter ones are dropped entirely. Order and role pairing + are preserved, so the result is a valid conversation with the same shape. + + Only runs (and only mutates anything) when the estimate actually exceeds + the budget; otherwise the original list is returned untouched. + """ + orig_tokens = sum(estimate_tokens(extract_text(m)) for m in messages) + if orig_tokens <= budget_tokens: + return messages, { + "pruned": False, + "original_tokens": orig_tokens, + "final_tokens": orig_tokens, + "tokens_saved": 0, + } + + user_indexes = _first_user_turn_indexes(messages) + num_protected_turns = min(keep_last_turns, len(user_indexes)) + if num_protected_turns <= 0: + # Nothing to protect: everything before is a candidate for trimming. + protected_from = len(messages) + else: + protected_from = user_indexes[len(user_indexes) - num_protected_turns] + + pruned: list[dict] = [] + tokens_saved = 0 + summarized = 0 + dropped = 0 + + for i, msg in enumerate(messages): + role = msg.get("role") + if role in ("user", "assistant", "system") or i >= protected_from: + pruned.append(msg) + continue + # Only tool results are candidates here. + if role != "tool": + pruned.append(msg) + continue + content = msg.get("content") + text = content if isinstance(content, str) else extract_text(msg) + if len(text) <= max_summarize_chars: + # Only replace when the placeholder is actually shorter than the + # result — dropping a tiny result to a longer marker saves nothing. + name = _tool_name(msg) or "tool" + placeholder = f"[{name}: result omitted]" + if len(placeholder) < len(text): + tokens_saved += estimate_tokens(text) - estimate_tokens(placeholder) + dropped += 1 + pruned.append({**msg, "content": placeholder}) + else: + pruned.append(msg) + continue + # Long tool result: keep head + tail so the shape survives. + head = text[:1500] + tail = text[-1500:] + trimmed = len(text) - 3000 + tokens_saved += trimmed // CHARS_PER_TOKEN + summarized += 1 + pruned.append( + {**msg, "content": f"{head}\n\n[...{trimmed:,} chars trimmed...]\n\n{tail}"} + ) + + final_tokens = sum(estimate_tokens(extract_text(m)) for m in pruned) + return pruned, { + "pruned": True, + "original_tokens": orig_tokens, + "final_tokens": final_tokens, + "tokens_saved": tokens_saved, + "summarized": summarized, + "dropped": dropped, + } diff --git a/dispatcher.py b/dispatcher.py index 2523dae..df342c8 100644 --- a/dispatcher.py +++ b/dispatcher.py @@ -55,6 +55,7 @@ from pydantic import BaseModel, Field import logs from capabilities import detect_capabilities, iter_image_url_values from config import RouterConfig, load_config +from context_prune import prune_context from routing import ( BATCH, INTERACTIVE, @@ -416,9 +417,16 @@ def classify(task: str, context: Optional[str]) -> Classification: f"Allowed values for task_category (use EXACTLY one of these strings):\n" + "\n".join(f" - {c}" for c in categories) ) - user_content = task if not context else f"{task}\n\n--- context ---\n{context}" + if not context: + user_content = task + else: + # llmrouter's framing: the prior turn first, labelled "Context:", so a + # short follow-up ("Yes", "Try now?") inherits its complexity rather + # than being classified in isolation as trivial. + user_content = _classifier_user_content( + task, context, cfg.classifier.context_framing + ) user_content = clamp_for_classifier(user_content, cfg.classifier.max_input_chars) - client = _classifier_client() started = time.perf_counter() try: @@ -457,6 +465,22 @@ def classify(task: str, context: Optional[str]) -> Classification: return resp +def _classifier_user_content( + task: str, context: Optional[str], framing: bool +) -> str: + """Assemble the classifier's user message (pure, for testability). + + ``framing`` selects llmrouter's "Context: / Message: " + layout so a short follow-up inherits its prior turn's complexity; + otherwise the legacy "task / --- context --- / context" layout is used. + """ + if not context: + return task + if framing: + return f"Context: {context}\n---\nMessage: {task}" + return f"{task}\n\n--- context ---\n{context}" + + def _classify_once(client: OpenAI, system_prompt: str, user_content: str) -> Classification: """One classifier round-trip. Raises on anything unusable.""" categories = cfg.proficiency.categories @@ -1533,6 +1557,37 @@ def _last_user_text(messages: list[dict]) -> str: return "" +def _previous_context(messages: list[dict]) -> str: + """The message before the last user turn, for inheriting its complexity. + + A short follow-up ("Yes", "Try now?") continues the prior turn; feeding + that prior turn to the classifier as context lets it inherit the turn's + complexity instead of being classified in isolation as trivial — the same + rule llmrouter encodes as "short message + complex context = use context + complexity". Returns the text of the message immediately preceding the + last user message, truncated, or "" when there is no such message. + """ + for i in range(len(messages) - 2, -1, -1): + if messages[i].get("role") == "user": + continue + content = messages[i].get("content") + text = "" + if isinstance(content, str): + text = content + elif isinstance(content, list): + parts = [] + for part in content: + if isinstance(part, dict): + if part.get("type") == "text" and isinstance(part.get("text"), str): + parts.append(part["text"]) + elif part.get("type") == "tool_result": + parts.append(str(part.get("content", ""))) + text = " ".join(parts) + if text: + return text[:200] + return "" + + def _count_images(messages: list[dict]) -> int: """The number of image_url parts across the whole conversation. @@ -1858,10 +1913,15 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks): if wants_routing: latency = BATCH if requested == ROUTER_MODEL_BATCH else INTERACTIVE tools_present = caps.tools_present + # The turn before the last: a short follow-up inherits its complexity, + # so the classifier sees "Context: \n---\nMessage: " + # instead of judging the follow-up alone. + prev_context = _previous_context(messages) classify_started = time.perf_counter() decision = route( TaskRequest( task=_last_user_text(messages), + context=prev_context, latency_tolerance=latency, tools_present=tools_present, has_images=caps.has_images, @@ -1882,6 +1942,7 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks): decision = route( TaskRequest( task=_last_user_text(messages), + context=prev_context, latency_tolerance=latency, tools_present=tools_present, has_images=caps.has_images, @@ -2009,7 +2070,26 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks): # top (needed by the passthrough/local-vision branches), only the working # directory is derived here because only the observation path uses it. session_dir = session_directory(messages) - upstream_body = {**body, "model": target} + send_messages = messages + if cfg.pinch.enabled: + # Optional relevance-based context pruning: when the conversation + # exceeds budget_tokens, trim old tool results BEFORE any paid token is + # sent upstream. User/assistant/system messages are always kept; only + # the classifier input is exempt (it already clamps to head+tail). + send_messages, pinch_stats = prune_context( + list(messages), + budget_tokens=cfg.pinch.budget_tokens, + keep_last_turns=cfg.pinch.keep_last_turns, + max_summarize_chars=cfg.pinch.max_summarize_chars, + ) + logs.debug( + "pinch", + pruned=pinch_stats["pruned"], + saved=pinch_stats["tokens_saved"], + orig=pinch_stats["original_tokens"], + final=pinch_stats["final_tokens"], + ) + upstream_body = {**body, "model": target, "messages": list(send_messages)} streaming = bool(body.get("stream")) if streaming: # Without this the final chunk carries no usage and the observation @@ -2054,7 +2134,10 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks): retry_trail: list[str] = [] while True: - attempt_body = {**body, "model": current_model} + attempt_body = { + **body, "model": current_model, + "messages": list(send_messages), + } if current_max_tokens is not None: attempt_body["max_tokens"] = current_max_tokens upstream_started = time.perf_counter() diff --git a/tests/test_classifier_input.py b/tests/test_classifier_input.py index d1cf2c7..b8b7464 100644 --- a/tests/test_classifier_input.py +++ b/tests/test_classifier_input.py @@ -14,7 +14,11 @@ With the clamp, the same prompts classify correctly in ~2.2s. from __future__ import annotations -from dispatcher import clamp_for_classifier +from dispatcher import ( + _classifier_user_content, + _previous_context, + clamp_for_classifier, +) def test_short_input_is_untouched(): @@ -61,3 +65,90 @@ def test_zero_disables_clamping(): def test_negative_is_treated_as_disabled_not_as_a_crash(): text = "w" * 100_000 assert clamp_for_classifier(text, -1) == text + + +# --- context framing (ported from llmrouter) ------------------------------- + +def test_framing_without_context_is_just_the_task(): + assert _classifier_user_content("Refactor this", None, True) == "Refactor this" + assert _classifier_user_content("Refactor this", None, False) == "Refactor this" + + +def test_framing_puts_prior_turn_first_when_enabled(): + out = _classifier_user_content("Try now?", "Design a distributed system", True) + assert out.startswith("Context: Design a distributed system") + assert "\n---\nMessage: Try now?" in out + # The current task comes after the context, labelled "Message:". + ctx, _, msg = out.partition("\n---\n") + assert ctx == "Context: Design a distributed system" + assert msg == "Message: Try now?" + + +def test_legacy_framing_puts_task_first(): + out = _classifier_user_content("Try now?", "Design a distributed system", False) + assert out.startswith("Try now?") + assert "--- context ---" in out + assert "Design a distributed system" in out + + +def test_framing_input_still_feeds_the_clamp(): + # The framed input is what clamp_for_classifier trims, so a huge prior turn + # is still bounded to head + tail and the elision stays visible. + context = "context blob " * 50_000 + framed = _classifier_user_content("Try now?", context, True) + assert len(framed) > 8000 + clamped = clamp_for_classifier(framed, 8000) + assert "Context:" in clamped + assert "Message: Try now?" in clamped + assert "elided" in clamped + + +# --- previous-message extraction ------------------------------------------ + +def test_previous_context_is_the_message_before_last_user(): + messages = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "partly done"}, + {"role": "user", "content": "Try now?"}, + ] + assert _previous_context(messages) == "partly done" + + +def test_previous_context_skips_user_messages(): + # Only the message before the LAST user is returned; earlier user messages + # are skipped, matching llmrouter ("short follow-up inherits from context"). + messages = [ + {"role": "user", "content": "earlier"}, + {"role": "user", "content": "current"}, + ] + assert _previous_context(messages) == "" + + +def test_previous_context_empty_when_no_prior_message(): + assert _previous_context([{"role": "user", "content": "only"}]) == "" + assert _previous_context([]) == "" + + +def test_previous_context_handles_content_blocks(): + messages = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "head"}, + {"type": "tool_result", "content": "output"}, + ], + }, + {"role": "user", "content": "Go"}, + ] + ctx = _previous_context(messages) + assert "head" in ctx + assert "output" in ctx + + +def test_previous_context_is_truncated(): + messages = [ + {"role": "assistant", "content": "z" * 10_000}, + {"role": "user", "content": "Go"}, + ] + assert len(_previous_context(messages)) <= 200 + diff --git a/tests/test_config_endpoints.py b/tests/test_config_endpoints.py index 312d73d..287f056 100644 --- a/tests/test_config_endpoints.py +++ b/tests/test_config_endpoints.py @@ -230,3 +230,42 @@ def test_nonpositive_local_timeout_is_rejected(raw): cfg["local_vision"]["timeout_seconds"] = 0 with pytest.raises(ValueError, match="timeout_seconds"): RouterConfig(**cfg) + + +def test_the_shipped_config_loads_the_pinch_section(raw): + loaded = RouterConfig(**raw) + assert loaded.pinch.enabled is False + assert loaded.pinch.budget_tokens > 0 + assert loaded.pinch.keep_last_turns > 0 + + +def test_pinch_defaults_when_absent(raw): + cfg = copy.deepcopy(raw) + cfg.pop("pinch") + loaded = RouterConfig(**cfg) + assert loaded.pinch.enabled is False + assert loaded.pinch.budget_tokens == 50000 + assert loaded.pinch.keep_last_turns == 4 + + +def test_nonpositive_pinch_budget_is_rejected(raw): + cfg = copy.deepcopy(raw) + cfg["pinch"]["budget_tokens"] = 0 + with pytest.raises(ValueError, match="budget_tokens"): + RouterConfig(**cfg) + + +def test_nonpositive_pinch_keep_last_turns_is_rejected(raw): + cfg = copy.deepcopy(raw) + cfg["pinch"]["keep_last_turns"] = 0 + with pytest.raises(ValueError, match="keep_last_turns"): + RouterConfig(**cfg) + + +def test_classifier_context_framing_loads(raw): + # Defaults on (see config.yaml); the legacy layout is opt-out. + assert RouterConfig(**raw).classifier.context_framing is True + cfg = copy.deepcopy(raw) + cfg["classifier"]["context_framing"] = False + assert RouterConfig(**cfg).classifier.context_framing is False + diff --git a/tests/test_context_prune.py b/tests/test_context_prune.py new file mode 100644 index 0000000..d44b526 --- /dev/null +++ b/tests/test_context_prune.py @@ -0,0 +1,158 @@ +"""Pure relevance-based context pruning. + +Ported from the MIT-licensed llmrouter "pinch" module, reduced to the safe +invariants 6krrt requires: user/assistant/system messages are always kept +verbatim, only old tool results are trimmed, and message order / role pairing +survives so the result still parses as a conversation. +""" + +from __future__ import annotations + +from context_prune import estimate_tokens, extract_text, prune_context + + +def _user(text: str) -> dict: + return {"role": "user", "content": text} + + +def _assistant(text: str) -> dict: + return {"role": "assistant", "content": text} + + +def _tool(name: str, content: str) -> dict: + return {"role": "tool", "name": name, "content": content} + + +def test_short_conversation_is_untouched(): + messages = [_user("hi"), _assistant("hello")] + out, stats = prune_context(messages, budget_tokens=50000) + assert out == messages + assert stats["pruned"] is False + + +def test_user_and_assistant_messages_always_kept(): + messages = [ + _user("big user " * 5000), + _assistant("big assistant " * 5000), + _user("tail"), + ] + out, stats = prune_context(messages, budget_tokens=100) + assert stats["pruned"] is True + assert len(out) == len(messages) + assert out[0]["role"] == "user" + assert out[1]["role"] == "assistant" + assert out[2]["role"] == "user" + # All three survived verbatim; only their absolute content is untouched. + assert out[0]["content"] == messages[0]["content"] + assert out[2]["content"] == "tail" + + +def test_old_tool_results_beyond_window_are_trimmed(): + messages = [ + _tool("read", "huge tool output " * 2000), # old, should be trimmed + _assistant("step one"), + _user("continue"), + _assistant("ok"), + ] + out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=1) + assert stats["pruned"] is True + assert stats["tokens_saved"] > 0 + # The old tool result got shortened (head + tail + elision marker). + tool_out = out[0] + assert tool_out["role"] == "tool" + assert "chars trimmed" in tool_out["content"] + assert out[1]["content"] == "step one" + assert out[2]["content"] == "continue" + assert out[3]["content"] == "ok" + + +def test_recent_turn_is_protected_from_pruning(): + # keep_last_turns=1 protects the last user turn AND its tool results. + messages = [ + _user("first"), + _tool("search", "old result " * 3000), + _assistant("first answer"), + _user("second"), + _tool("search", "recent result " * 3000), + _assistant("second answer"), + ] + out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=1) + assert stats["pruned"] is True + # The recent tool result (after the 2nd user) survives verbatim. + recent_tool = [m for m in out if m.get("role") == "tool"][-1] + assert recent_tool["content"] == messages[4]["content"] + # Old tool result got trimmed. + old_tool = [m for m in out if m.get("role") == "tool"][0] + assert "chars trimmed" in old_tool["content"] + + +def test_tool_result_role_pairing_preserved(): + # The tool message keeps its name and position, so the API still parses. + messages = [ + _user("q"), + _tool("search", "x" * 6000), + _assistant("a"), + ] + out, stats = prune_context(messages, budget_tokens=50, keep_last_turns=0) + tool_out = [m for m in out if m.get("role") == "tool"][0] + assert tool_out["name"] == "search" + assert out[0]["role"] == "user" + assert out[1]["role"] == "tool" + assert out[2]["role"] == "assistant" + + +def test_short_tool_result_is_dropped_to_placeholder(): + # A large user message pushes the conversation over budget, so the SHORT + # old tool result (under max_summarize_chars but longer than the marker) is + # dropped to a placeholder that is shorter than the result it replaces. + result_text = "result " * 40 # 280 chars > placeholder, < 4000 cap + messages = [ + _user("big user " * 5000), + _tool("search", result_text), + _assistant("a"), + ] + out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=0) + assert stats["pruned"] is True + tool_out = [m for m in out if m.get("role") == "tool"][0] + assert "[search: result omitted]" in tool_out["content"] + assert len(tool_out["content"]) < len(result_text) + + +def test_tiny_tool_result_is_kept_when_placeholder_loses_tokens(): + # A result shorter than the placeholder marker is kept verbatim: replacing + # it would ADD tokens, which is not the point of pruning. + messages = [ + _user("big user " * 5000), + _tool("search", "hi"), + _assistant("a"), + ] + out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=0) + tool_out = [m for m in out if m.get("role") == "tool"][0] + assert tool_out["content"] == "hi" + + +def test_stats_report_tokens_saved(): + messages = [_user("q"), _tool("search", "y" * 9000), _assistant("a")] + out, stats = prune_context(messages, budget_tokens=10, keep_last_turns=0) + assert stats["pruned"] is True + assert stats["original_tokens"] > stats["final_tokens"] + assert stats["tokens_saved"] > 0 + assert stats["summarized"] >= 1 + + +def test_extract_text_handles_content_blocks(): + msg = { + "role": "tool", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "tool_result", "content": "world"}, + ], + } + assert "hello" in extract_text(msg) + assert "world" in extract_text(msg) + + +def test_estimate_tokens_is_crude_but_consistent(): + assert estimate_tokens("x" * 30) == 10 + assert estimate_tokens("") == 0 + assert estimate_tokens(None) == 0 -- 2.49.1 From 099958b392aeb92a3acff609d00de4b4fc5df8eb Mon Sep 17 00:00:00 2001 From: adlee-was-taken Date: Sun, 23 Aug 2026 22:22:56 -0400 Subject: [PATCH 2/3] docs: review + plan for context-pruning-and-framing feature --- ...-dual-use-and-classify-once-per-session.md | 142 ++++++++++++ ...ontext-pruning-and-framing-fixes-review.md | 125 +++++++++++ .../context-pruning-and-framing-review.md | 208 ++++++++++++++++++ 3 files changed, 475 insertions(+) create mode 100644 code_plans/context-dual-use-and-classify-once-per-session.md create mode 100644 code_reviews/context-pruning-and-framing-fixes-review.md create mode 100644 code_reviews/context-pruning-and-framing-review.md diff --git a/code_plans/context-dual-use-and-classify-once-per-session.md b/code_plans/context-dual-use-and-classify-once-per-session.md new file mode 100644 index 0000000..b68df22 --- /dev/null +++ b/code_plans/context-dual-use-and-classify-once-per-session.md @@ -0,0 +1,142 @@ +# Spec: two design decisions deferred out of the context-pruning/framing fix pass + +**Origin.** The fix pass for +[`context-pruning-and-framing-review.md`](../code_reviews/context-pruning-and-framing-review.md) +(reviewed in +[`context-pruning-and-framing-fixes-review.md`](../code_reviews/context-pruning-and-framing-fixes-review.md)) +explicitly declined two items as design decisions rather than confirmed +bugs: the `TaskRequest.context` dual-use ambiguity (review finding #6, +partially mitigated but not resolved), and classifying once per session +instead of once per message (an existing item on the project's own "what's +NOT built yet" list). Both are forward-looking — no bug is being reported +here, no code changes accompany this document. + +--- + +## 1. `TaskRequest.context` carries two unrelated meanings + +### The problem + +`context` was originally, and is still documented as, "assembled context +(docs/code) to send with the task" — a `/route`/`/dispatch` caller pastes +reference material alongside a task description. `chat_completions` now +also uses the same field to carry the prior conversational turn +(`_previous_context`), so a short follow-up like "Yes" inherits that turn's +complexity instead of being classified as trivial in isolation. + +The current fix (moving the framing instruction from the static system +prompt into `_classifier_user_content`'s output, appended only when +`context` is non-empty) narrowed the blast radius — the instruction no +longer reaches every classify() call, only ones that actually supply +`context` — but it didn't resolve which of the two meanings a given +`context` value has. A `/route` caller pasting 800 lines of Django code +under `task="Refactor this"` gets the same "a short follow-up continues the +prior turn, classify by the CONTEXT's complexity" instruction that exists +for the conversational case. + +### Why this might not need fixing + +The instruction's general principle — "short task + large/complex context +implies a non-trivial task" — is arguably a reasonable heuristic for the +docs-paste case too, even though it was written for conversational +follow-ups. It has never been measured against real `/route`/`/dispatch` +traffic with `context` set. This project's own stated epistemics apply +directly here: `POST /outcome` is "the only ground truth," and several +sections of the README describe correcting an assumption only after +measuring it, not before. Speculating about which framing is "more +correct" without a measurement repeats the mistake the project has already +named and moved past. + +### Options, if it turns out to matter + +| Option | Sketch | Trade-off | +|---|---|---| +| A. Decouple the mechanisms | Keep `TaskRequest.context` as the public docs/code field, untouched. Give `chat_completions`'s conversational-continuation signal its own internal path — e.g. `classify()` gains a private `prior_turn: Optional[str]` parameter distinct from `context`, so the instruction is only ever built for genuine conversational continuations | Cleanest semantically; requires touching `classify()`'s signature and both call sites; `/route`/`/dispatch` behavior is provably unaffected | +| B. Tag the field | Add a `context_kind: Literal["reference", "prior_turn"] = "reference"` field to `TaskRequest`; only `chat_completions`'s internal calls set `"prior_turn"`; the instruction is only appended for that kind | Smaller diff than A; adds a field to the public request model that external callers never need to know about | +| C. Leave as-is, measure | No code change. Watch `route_decisions` (already logs `task_category`/`task_tier`/`source` per request) for `/route`/`/dispatch` calls that supply `context` and see whether their tier/category looks skewed relative to before this change | Zero engineering cost; consistent with the project's own "measure before correcting" pattern; only viable if `/route`/`/dispatch`-with-`context` traffic is common enough to be observable | + +**Recommendation:** C first. This codebase already has the instrumentation +(`route_decisions`, `feedback.py` folding in outcomes) to tell whether this +is a real problem instead of a theoretical one, and the existing pattern in +this project — cost-vs-eco, tier-from-price, the whole classifier-model +swap — is "measure, then fix what the measurement shows," not "fix what +looks fishy." If `/route`/`/dispatch`-with-`context` traffic turns out to +be rare or nonexistent, this is not worth A's or B's added surface at all. + +--- + +## 2. Classify once per session, not once per message + +### The problem, restated from the README + +> ~10s of local overhead on every message is a real tax for an interactive +> agent... Still unaddressed: classify once per session rather than per +> message, cache by prompt hash, or skip classification for short prompts. + +With the local `mistral-nemo` classifier this is now ~1.7s/call +(§"The classifier is the latency floor"); with a cloud classifier +(measured against `deepseek-v4-flash`) it's ~1.0s. Either way, every single +turn in a long agent session pays this again, even though the session's +*category* (coding_general, debugging, etc.) rarely changes turn to turn — +what changes is mostly the token count, which `chat_completions` already +measures directly via `estimate_prompt_tokens` and doesn't need the +classifier for. + +### The mechanism already half-exists + +`route()`'s override branch (`dispatcher.py:673`) already skips +`classify()` entirely whenever `task_category`, `task_tier`, and +`required_context_tokens` are all supplied — this is exactly what the +measured-context reroute (fixed in review finding #4) uses today, just +within a single request. Session-level caching is the same mechanism +applied across requests: classify once, store `(task_category, task_tier)` +keyed by session, and on every later turn in that session call `route()` +with the cached category/tier plus a **freshly measured** +`required_context_tokens` (which is cheap — pure token counting, no model +call) — landing on the override branch and skipping the classifier +round-trip entirely. + +### Design questions to settle before implementing + +1. **Session identity.** `session_fingerprint`/`session_directory` + (dispatcher.py) already derive a session identity from message content + for observation purposes. Whether that's the right key for a + *classification* cache (vs. e.g. a client-supplied session id, if + opencode's protocol carries one) needs checking — a cache keyed on the + wrong signal either misses constantly (no benefit) or collides across + genuinely different sessions (wrong category persists into unrelated + work). +2. **Invalidation.** A session's task can genuinely change category mid-way + (debugging turns into a docs-writing turn turns into refactoring). Pure + "classify once, cache forever" risks staleness. Candidate triggers to + re-classify: a large jump in `required_context_tokens` between turns (a + proxy for "something new started"), a fixed number of turns (e.g. + re-classify every 20), or a TTL. This needs the same "measure before + deciding" treatment as everything else in this project — a cheap thing + to instrument via `route_decisions.source` (add a `"cached"` value + alongside `"classifier"`/`"override"`/`"fallback"`) and watch category + drift over real sessions before picking a policy. +3. **Interaction with escalation and retries.** `apply_escalation` currently + runs on every fresh classification. A cached category/tier bypasses it + entirely on cache-hit turns — need to decide whether escalation state + should also be cached per-session or re-evaluated each turn (it's cheap, + pure Python, so probably always re-evaluate rather than cache). +4. **Storage.** In-memory dict keyed by session identity is the obvious + starting point (matches the process lifetime of the dispatcher; a + restart just means the next turn in every active session re-classifies + once, which is a safe failure mode) — no new persistence layer needed + unless multi-process deployment becomes a requirement. + +### Recommendation + +Worth building — the latency case is strong and the mechanism is a small +extension of code that already exists (the override branch) rather than a +new one. But settle invalidation policy (§2) with a measurement pass first, +the same way `context_framing`'s default-on-and-measure and the classifier +model swap were each decided by running both and comparing, not by +argument. A reasonable first cut: cache with no re-classification, ship +behind a config flag (default off, matching every other new-and-unproven +knob in this project — `pinch.enabled`, `min_tool_proficiency`), watch +`route_decisions` on real sessions for category drift, then decide whether +any invalidation trigger is actually needed or whether "classify once, +never again" is good enough in practice. diff --git a/code_reviews/context-pruning-and-framing-fixes-review.md b/code_reviews/context-pruning-and-framing-fixes-review.md new file mode 100644 index 0000000..17deda5 --- /dev/null +++ b/code_reviews/context-pruning-and-framing-fixes-review.md @@ -0,0 +1,125 @@ +# Review: fixes for `context-pruning-and-framing-review.md` + +**What it was reviewing:** the working-tree diff (7 files, unstaged, not yet +committed) implementing fixes for the 11 confirmed bugs and 4 cleanup items +from [`context-pruning-and-framing-review.md`](context-pruning-and-framing-review.md), +done by opencode driven through the router. Reviewed by reading the full +diff against each numbered finding, re-deriving the fix logic by hand, and +reproducing behavior directly against `context_prune.py` and `dispatcher.py` +rather than trusting the diff's comments. Full suite: 544/544 passing (up +from 529; 15 new regression tests, one per fixed finding). + +## Verdict: 10 of 11 bugs correctly fixed and verified; 1 new regression found in the process + +### Fixed and verified by direct reproduction + +| # | Original finding | Fix | Verified | +|---|---|---|---| +| 1 | `_previous_context` returned raw tool output as classifier context | Now walks back to the nearest `role == "assistant"` message only, skipping `system`/`user`/`tool`, and reuses `context_prune.extract_text` instead of a second copy of the block parser (also fixes #13) | `test_previous_context_excludes_tool_role_messages` — read directly, correct | +| 2 | Same function returned the system prompt on a session's first turn | Same fix as #1 (system role is no longer eligible at all) | `test_previous_context_excludes_system_prompt` — correct | +| 3 | Long-tool-result trim went negative and grew the message when `max_summarize_chars < 3000` | Two layers: `config.py` now validates `max_summarize_chars >= 3000` at load; `context_prune.py` also defensively checks `trimmed > 0 and len(elided) < len(text)` before using the elided form, falling through to the placeholder otherwise, so a direct call bypassing config can't hit it either | Reproduced by hand: `max_summarize_chars=100`, 1200-char input → output is strictly shorter, `tokens_saved >= 0`, no negative marker. Matches `test_max_summarize_chars_below_3000_never_grows_or_goes_negative` | +| 4 | Routing/tier/cost decided on unpruned tokens; pinch's savings never reached the decision | `prune_context` now runs once inside the `wants_routing` block, before `measured = estimate_prompt_tokens(send_messages, ...)`; the already-pruned `send_messages` is reused at dispatch, and the passthrough path (which never goes through `wants_routing`) still prunes on its own. Confirmed this doesn't double-prune and doesn't touch `messages` in place (only `send_messages` is reassigned; `session_directory(messages)`, `_last_user_text(messages)`, and `_run_local_vision(messages, ...)` all still read the original, full conversation) | Read the full function end-to-end to confirm the no-double-prune and no-mutation properties; matches `test_pruned_list_measures_smaller_than_the_raw_conversation` and `test_pinch_disabled_leaves_the_measured_size_unchanged` | +| 5 | `extract_text` read `image_url` blocks as empty, undercounting image-bearing messages | Now contributes the block's `url` string; nested `tool_result.content` lists are flattened recursively too | Correct as far as counting goes — **but see the new finding below**, this fix combined with #11's fix produces a different bug | +| 7 | `context=prev_context` was dead on the measured-context reroute `route()` call | Parameter removed from that call, with a comment explaining the override branch never reads it | Confirmed by re-reading `route()`'s branch condition; no behavior change, just removes a misleading dead arg | +| 8 | `dropped` stat/docstring claimed removal; code always replaced content in place | `dropped` counter removed; both trim paths now fold into `summarized`; docstring updated to say "trimmed or summarized (never removed)" | `test_stats_use_summarized_and_length_is_always_preserved` — correct | +| 9 | `classifier.system_prompt` unconditionally described the "Context:"/"Message:" framing even when `context_framing: false` | The framing instruction was removed from the static system prompt entirely and is now appended to `_classifier_user_content`'s output only when `context` is actually present, and worded generically enough (doesn't reference the specific labels) to be correct under both framing modes | Read both `config.yaml` and `_classifier_user_content`; the instruction now travels with the content that justifies it rather than being unconditional boilerplate | +| 10 | Zero user-role messages collapsed `protected_from` to `len(messages)`, protecting nothing — even the newest tool result was prunable | New branch: when there's no user message, protect the trailing `keep_last_turns` *tool* results instead | Reproduced by hand with 5 tool results / `keep_last_turns=2`: the 2 most recent survive verbatim, the 3 older ones are trimmed. Matches `test_zero_user_messages_protect_newest_tool_result` | +| 12 | `_turn_of()` was dead code | Removed | Confirmed, no remaining references | +| 13 | `_previous_context` duplicated `extract_text`'s block-parsing logic | Now imports and calls `context_prune.extract_text` directly | Confirmed via the `#1` fix above | +| 14 | `classify()`'s outer `if not context` duplicated a check `_classifier_user_content` already makes | Outer branch removed; `classify()` now calls `_classifier_user_content` unconditionally | `test_classify_user_content_invariant_when_no_context` — correct | +| 15 | Pinch's 4 defaults declared independently in `context_prune.py`, `config.py`, and `config.yaml` | `context_prune.py`'s module-level `DEFAULT_*` constants removed; `prune_context`'s default args now read from `PinchConfig.model_fields[...].default`, one source of truth | Confirmed — `CHARS_PER_TOKEN` itself is still a separate mirrored constant (unaddressed, but this was always the minor half of #15) | + +Finding #11 (structured content silently flattened to a string, dropping +non-text blocks) is **also fixed** in the sense that it was scoped — +`image_url` and other non-text blocks now survive trimming instead of being +discarded — but the fix interacts badly with #5's fix, below. + +### New finding: image-bearing tool results aren't actually shrunk, and the stats lie about it + +`context_prune.py` — `extract_text` (39-62, image_url branch) + the trim +branch (155-175) + `_with_text` (91-114) + +Fixing #5 (count `image_url` bytes into the size estimate) and #11 (never +drop non-text blocks) independently make sense, but composed, they produce +a case neither fix's own test covers: a tool result whose `image_url` block +is *large enough to actually matter*. + +`extract_text` joins every block — text **and** the image's raw `url` +string — into one combined string. That combined string is what gets +head+tail-sliced for the elided/placeholder replacement. `_with_text` then +writes that replacement into the message's `text` block only, and copies +every non-text block (the `image_url` one) through **completely +unchanged** — full size, untouched. Net effect for an image-bearing tool +result: the image (almost always the dominant contributor to size) is never +actually reduced, while the `tokens_saved` stat is computed from the +combined length including the image bytes, so it reports a large "savings" +that didn't happen. Reproduced directly: + +```python +huge_b64 = "A" * 200000 +messages = [ + {"role": "user", "content": "describe this"}, + {"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]}, + {"role": "tool", "name": "screenshot", "content": [ + {"type": "text", "text": "here is the screenshot"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{huge_b64}"}}, + ]}, + {"role": "user", "content": "now what"}, +] +pruned, stats = prune_context(messages, budget_tokens=10, keep_last_turns=0) +# stats == {'pruned': True, 'original_tokens': 66687, 'final_tokens': 67690, +# 'tokens_saved': 65671, 'summarized': 1} +``` + +`final_tokens` (67690) is **larger** than `original_tokens` (66687) — the +message got bigger, not smaller — while `tokens_saved` claims 65,671 tokens +saved. The image_url block comes out with its full 200,022-char url +untouched, and the text block now reads `"here is the screenshot +data:image/png;base64,AAAA...[187,022 chars trimmed...]AAAA"` — a raw +fragment of the base64 payload, sliced at an arbitrary byte boundary, now +sitting in a `text` field. This is worse than a no-op: it doesn't reduce +what's shipped upstream, it corrupts the text block with base64 noise, and +it actively misreports the stat the whole feature exists to produce. + +The existing regression test for #11 +(`test_structured_tool_result_keeps_image_block_when_trimmed`) doesn't +catch this because its fixture image is 4 characters (`"AAAA"`) — small +enough that the leakage and false accounting are present but invisible. +Any real screenshot or image tool result (tens of KB to MB of base64) would +hit this. + +The fix likely belongs in `extract_text`/the trim logic together: +`image_url` blocks should count toward the size estimate (that part of #5 +is correct — undercounting was the original bug) but should **not** be +included in the text that gets sliced for the elided/placeholder +replacement, and should be excluded from (or separately accounted in) +whatever text ends up in the `text` block. Whether an oversized image +itself should also become a trim candidate (a placeholder replacing the +`image_url` block, not just the text) is a design question worth deciding +explicitly rather than falling out of two unrelated fixes' composition. + +### Not fully resolved (matches what was flagged as deferred) + +Original finding #6 — `TaskRequest.context`'s dual purpose (docs/code paste +for `/route`/`/dispatch` vs. the new conversational-continuation use in +`chat_completions`) — got a real, useful partial fix: the framing +instruction is no longer unconditionally injected into every classify() +call (that part is #9, now fixed), and the field's dual use is now +documented on `TaskRequest.context` itself. But the underlying semantic +question is still open: a `/route` or `/dispatch` caller who passes +`context=` (the field's original, documented purpose) +still gets the classifier told "a short follow-up continues the prior turn, +classify by the CONTEXT's complexity" — a rule written for conversational +continuation, appended regardless of which of the two meanings this +particular caller's `context` actually carries. This tracks with what was +flagged as a deferred design decision rather than a missed fix. + +## Recommendation + +Don't commit yet — the image/large-tool-result finding above is a real, +reproducible regression (not a pre-existing issue; it's new from this +round's own fixes), and it directly undermines the stat this feature exists +to produce. Route it back for another pass; everything else here is solid +and doesn't need to be touched again. The two deferred design items are +covered in a separate forward-looking spec rather than this after-the-fact +report. diff --git a/code_reviews/context-pruning-and-framing-review.md b/code_reviews/context-pruning-and-framing-review.md new file mode 100644 index 0000000..1f7de68 --- /dev/null +++ b/code_reviews/context-pruning-and-framing-review.md @@ -0,0 +1,208 @@ +# Review: fbcc636, "context-aware framing and relevance-based context pruning" + +**What it was reviewing:** the two features ported from the MIT-licensed +`alexrudloff/llmrouter` project — context-aware classifier framing +(`_previous_context` / `_classifier_user_content` in `dispatcher.py`) and +relevance-based context pruning (`context_prune.py`, "pinch"). Reviewed with +`/code-review` at xhigh effort (10 finder angles + 1-vote verify + a gap +sweep), then spot-verified directly: read every file involved, and +reproduced the two most severe findings against the actual `prune_context` +function rather than trusting the description. + +## Verdict: full suite is green (529/529), but that's not the same as correct + +`pytest` passes clean, including the 10 new `tests/test_context_prune.py` +cases. None of those cases exercise the shapes that break, though: a +`max_summarize_chars` below ~3000, a conversation with zero user messages, +tool content as `image_url` blocks, or the dispatcher-level interaction +between `_previous_context` and `route()`. The suite passing means the +happy path works, not that the port is safe on real agent traffic — which is +exactly the traffic this router's own README says is dominant (tool calls, +long sessions, ~92% cache hits on huge prompts). + +Both features are **off by default** (`pinch.enabled: false`, +`classifier.context_framing: true` — framing is on, pruning is off), so +nothing here is live yet. But `context_framing` defaulting **on** means the +framing bugs (#1, #2 below) are already affecting every routed +`/v1/chat/completions` call today. + +## Confirmed bugs — verified directly, not just reported + +### 1. `_previous_context` feeds the classifier raw tool output, not "the prior turn" — and it's the common case, not an edge case +`dispatcher.py:1560-1588` + +```python +for i in range(len(messages) - 2, -1, -1): + if messages[i].get("role") == "user": + continue + ... +``` + +Only `role == "user"` is skipped while walking backward from `messages[-2]`. +`role == "tool"` is not. In the standard OpenAI tool-loop shape — +`assistant(tool_call) -> tool(result) -> user("ok fix it")` — `messages[-2]` +*is* the tool message, so `_previous_context` returns the raw tool-result +payload (file contents, grep output, JSON), truncated to 200 chars, as +`"Context:"` for the classifier. This isn't a malformed-conversation edge +case; it's what every agent client running a tool loop produces on its very +next turn. The function's own docstring says it exists to let the +classifier "inherit the [prior] turn's complexity" — it's inheriting +arbitrary tool output instead. + +### 2. Same function returns the system prompt as "previous context" on a session's first turn +`dispatcher.py:1570-1572` + +For `messages = [system, user]`, the backward scan starts at index 0 (the +system message), which isn't `role == "user"`, so it's not skipped — its +first 200 chars come back as the "prior turn." Every fresh session's first +message gets classified with `"Context: "` +prepended, which for opencode is ~32K chars of tool definitions. + +### 3. Pinch's long-tool-result trim goes negative and *inflates* the message once `max_summarize_chars < 3000` +`context_prune.py:145-153`; `config.py:293-299` (no validator on this field, unlike its two siblings) + +```python +head = text[:1500] +tail = text[-1500:] +trimmed = len(text) - 3000 +``` + +This assumes any text reaching this branch is longer than 3000 chars, but +the only guard to get here is `len(text) > max_summarize_chars`, and +`max_summarize_chars` has no lower-bound validator (`budget_tokens` and +`keep_last_turns` both do). Reproduced directly: + +``` +max_summarize_chars=100, tool content = 800 chars +-> pruned tool content = 1632 chars (grew) +-> stats["tokens_saved"] = -734 (negative) +-> marker literally reads "[...-2,200 chars trimmed...]" +``` + +Below the default (4000) this can't trigger, but there's nothing stopping +an operator from setting it lower, and when they do, pruning does the +opposite of its job. + +### 4. Model/tier selection runs on unpruned tokens; pinch's savings never reach the decision that spends the money +`dispatcher.py` — routing at 1913-1954 vs. `prune_context()` at 2079 + +`route()` is called (twice — see #7) using `estimate_prompt_tokens(messages, ...)` +on the **full, unpruned** message list, and that's what drives tier +selection, the context-window hard filter, and the cost tiebreak. +`prune_context()` doesn't run until line 2079, well after `decision.selected` +is fixed — it only shrinks the payload actually sent to the model already +picked. So a long tool-heavy session can get routed to a pricier +large-context model based on its pre-pruned size, even though pinch would +have brought the real outgoing request in well under budget. This isn't +wrong on invalid input, it's an ordering bug: pinch's entire stated purpose +(reduce shipped tokens on long sessions) doesn't influence the one decision +where that would save money. + +### 5. `extract_text` silently treats `image_url` content blocks as empty +`context_prune.py:39-53`, used for both `orig_tokens` and the per-message trim decision + +Only `type == "text"` and `type == "tool_result"` blocks are read; anything +else (including `image_url`) contributes `""`. Two consequences: token +estimates can undercount a session that's actually huge (a tool result full +of base64 image data reads as 0 tokens, so `orig_tokens` may never cross +`budget_tokens` and pruning never triggers), and if pruning does trigger for +other reasons, that same message reads as `len(text) == 0 <= max_summarize_chars`, +so the "only replace if the placeholder is shorter" check (`23 < 0`) is +false and the giant blob is left completely untouched while smaller +genuine text results nearby get trimmed. + +### 6. `TaskRequest.context` now has two incompatible meanings sharing one field +`dispatcher.py:151-153` (field docstring: `"Assembled context (docs/code) to send with the task"`) vs. the new use in `chat_completions` (prior conversation turn) vs. `config.yaml:395-398` (system prompt instructions written for the second meaning) + +`/route` and `/dispatch` callers have always been able to pass `context` as +pasted docs/code (`dispatch_endpoint` splices it verbatim into a system +message). The classifier's system prompt was changed globally to say +"classify by the CONTEXT's complexity, treating a short message with +complex context as inheriting that complexity" — a rule written for the +conversational-continuation case, but it now applies unconditionally to +every existing `context=` caller too, since it's the same +field and the same prompt. Not obviously wrong, but untested for that +existing use and not called out anywhere as a behavior change to it. + +## Confirmed via reproduction — edge cases in pruning itself + +### 7. `context=prev_context` is dead weight on the second (measured-context) `route()` call +`dispatcher.py:1942-1954` vs. `route()`'s branch condition at line 673 + +`route()` only reads `req.context` inside the `classify()` branch, which is +skipped whenever `task_category`, `task_tier`, and `required_context_tokens` +are all supplied together — which the reroute at 1942 always does (it +copies the first decision's category/tier and sets +`required_context_tokens=measured`). So `context=prev_context` on that call +is passed and never read. Harmless today since the first `route()` call +already consumed it, but it means a future fix to `_previous_context` (#1/#2 +above) would silently not apply here, and there's nothing marking the +parameter as inert. + +### 8. `dropped` doesn't mean dropped +`context_prune.py:100-107` (docstring/stat name) vs. `133-144` (actual behavior) + +The docstring and `config.yaml` both say short old tool results are +"dropped entirely," and the counter is literally named `dropped`, but the +code never removes a message from the list — it always replaces `content` +with a placeholder string in place. `len(pruned) == len(messages)` always, +regardless of `dropped`. Any code (or test) later written against the +documented contract — e.g. `assert len(pruned) == len(messages) - dropped` +— would be wrong on every request that drops anything. + +### 9. Stale prompt instructions when the framing opt-out is used +`config.yaml:395-398` + +The classifier's `system_prompt` unconditionally describes the +`"Context: \n---\nMessage: "` label format, but that layout +is only actually produced when `classifier.context_framing: true`. Set it +`false` (the documented way to get the legacy `"task\n\n--- context ---\ncontext"` +layout) and the classifier still receives instructions describing a format +it will never see. + +### 10. Zero user messages -> the recency guard protects nothing +`context_prune.py:109-115` + +Reproduced: with no `role == "user"` message anywhere in the conversation, +`num_protected_turns` collapses to 0 and `protected_from = len(messages)`, +which no real index ever reaches — so the "always keep if `i >= protected_from`" +branch never fires. The single most recent tool result (the one the next +turn actually needs) becomes eligible for trimming, same as the oldest one: + +``` +messages = [system, assistant(tool_call), tool(20000 chars)] +-> the only tool result gets summarized down, despite being the newest +``` + +### 11. Trimming silently flattens structured content to a plain string +`context_prune.py:141, 152` + +Both trim branches do `{**msg, "content": }` unconditionally, even +when the original `content` was a list of blocks (`[{"type": "text", ...}, +{"type": "image_url", ...}]`). A tool result that mixes text and an image +part loses the image permanently the first time it ages past the protected +window — not "trimmed," just gone, with no signal that anything +non-text was there. + +## Cleanup — lower severity, no behavior change + +| # | Location | Issue | +|---|---|---| +| 12 | `context_prune.py:73-80` | `_turn_of()` has zero call sites anywhere in the repo — leftover from an earlier design. | +| 13 | `dispatcher.py:1573-1585` | `_previous_context`'s content-block extraction duplicates `context_prune.extract_text()` almost verbatim, despite `dispatcher.py:58` already importing from that module (`prune_context` only). A future change to one won't propagate to the other. | +| 14 | `dispatcher.py:420-428` vs. `468-481` | `classify()`'s `if not context: user_content = task else: user_content = _classifier_user_content(...)` duplicates a check `_classifier_user_content` already makes internally (`if not context: return task`, line 477). The outer branch can be deleted; `classify()` can call `_classifier_user_content` unconditionally. | +| 15 | `context_prune.py:27-31`, `config.py:293-299`, `config.yaml:175-177` | Pinch's four defaults and `CHARS_PER_TOKEN` are each declared independently in two or three places. `context_prune.py`'s own comment admits `CHARS_PER_TOKEN` "mirrors dispatcher.CHARS_PER_TOKEN" rather than importing it. `dispatcher.py` always passes `cfg.pinch.*` explicitly, so the module-level defaults in `context_prune.py` are dead in production. | + +## Take for next time + +The port kept the right invariant (user/assistant/system messages are never +touched, only tool results) but re-derived the surrounding plumbing instead +of reusing what the codebase already had for it — `_previous_context` is a +second, slightly-different copy of `extract_text`'s block-parsing logic, and +it re-introduces exactly the bug `_last_user_text` next to it was written to +avoid (`_last_user_text` correctly scans for the *nearest* user message +rather than assuming position; `_previous_context` assumes `messages[-2]` is +meaningful). Both new-feature bugs that actually change routing behavior +today (#1, #2) are about that same unguarded assumption: agent traffic +doesn't end tidily on a fresh user turn, and this codebase already knows +that everywhere else it touches messages. -- 2.49.1 From 5f7716e1213a00a49b4ce48b2a04d83ccb7793bb Mon Sep 17 00:00:00 2001 From: adlee-was-taken Date: Sun, 23 Aug 2026 22:23:08 -0400 Subject: [PATCH 3/3] fix: resolve context-pruning-and-framing review findings Address the 11 confirmed bugs and 4 cleanups from context-pruning-and-framing-review.md, plus the image-accounting regression it found on a second pass. classifier framing: - _previous_context walks back to the nearest assistant turn only, skipping system/user/tool so raw tool output and the system prompt never contaminate the framing signal; reuses context_prune.extract_text instead of a duplicated block parser. - classify() drops its redundant if/not-context branch. - framing instruction moved out of the static system prompt and into _classifier_user_content so it travels only with content that justifies it (respects context_framing opt-out). context pruning: - extract_text counts image_url bytes so image-bearing tool results are sized (and pruned) correctly. - -- new, follow-up review -- image_url blocks now count toward the size trigger but not toward the replacement text, and are stubbed to a short marker when the message is trimmed, so an image-heavy tool result genuinely shrinks rather than staying full-size while tokens_saved lied and base64 leaked into the text block. - tokens_saved is now the honest whole-list before/after reduction (max(orig - final, 0)), never a per-message estimate that could drift. - max_summarize_chars validated >= 3000 at config load; trim branch guards against negative math and message growth. - recency guard protects the newest tool result even with no user turn. - non-text blocks survive trimming (type preserved) rather than being flattened away. - dropped/stat terminology corrected to summarized; _turn_of() removed; module-level pinch defaults removed in favor of PinchConfig. Dispatch and config: - prune_context runs once before the measured-context routing decision (reused at dispatch, no double-prune, no in-place mutation); no-op when pinch is disabled. - dead context=prev_context arg removed from the measured reroute. - TaskRequest.context dual use documented. 15 new regression tests pin the fixes; 545 tests pass. --- config.py | 10 ++ config.yaml | 7 +- context_prune.py | 212 ++++++++++++++++++++++-------- dispatcher.py | 150 ++++++++++++++-------- tests/test_classifier_input.py | 53 +++++++- tests/test_config_endpoints.py | 19 +++ tests/test_context_prune.py | 227 +++++++++++++++++++++++++++++++++ 7 files changed, 561 insertions(+), 117 deletions(-) diff --git a/config.py b/config.py index 0513b16..65aa585 100644 --- a/config.py +++ b/config.py @@ -312,6 +312,16 @@ class PinchConfig(StrictModel): raise ValueError("pinch.keep_last_turns must be > 0") return v + @field_validator("max_summarize_chars") + @classmethod + def summarize_chars_valid(cls, v: int) -> int: + if v < 3000: + raise ValueError( + "pinch.max_summarize_chars must be >= 3000 (below this, " + "summarization grows the message)" + ) + return v + class DatabaseConfig(StrictModel): path: str diff --git a/config.yaml b/config.yaml index ae4de68..899785a 100644 --- a/config.yaml +++ b/config.yaml @@ -387,16 +387,11 @@ classifier: { "task_category": one of the allowed categories listed below, "task_tier": integer 1-3, where 1 is cheap/simple, 2 is mid/general, - and 3 is frontier/high-stakes, + and 3 is frontier/high-stakes, "required_context_tokens": integer estimate of prompt+context token count, "confidence": float 0-1 } - When the input is framed as "Context: \n---\nMessage: ": - a short follow-up message ("Yes", "Try now?", "Go ahead") continues the - prior turn, so classify it by the CONTEXT's complexity, not in isolation. - A standalone short message with no Context line is genuinely trivial. - dispatch_providers: # NeuralWatt is the only provider. The dict shape and the models table's # (model_id, provider) key are kept so a second one can be added without a diff --git a/context_prune.py b/context_prune.py index f909697..04a947d 100644 --- a/context_prune.py +++ b/context_prune.py @@ -7,7 +7,8 @@ the same SAFE invariants without requiring an embedding model on the request path: - user / assistant / system messages are ALWAYS kept verbatim, - - only TOOL RESULTS older than the protected window are trimmed or dropped, + - only TOOL RESULTS older than the protected window are trimmed or + summarized (never removed), - message order and the tool/assistant pairing are preserved so the result still parses as a valid conversation. @@ -24,9 +25,8 @@ This module is pure: it takes messages and limits and returns pruned messages. from __future__ import annotations -DEFAULT_BUDGET_TOKENS = 50000 -DEFAULT_KEEP_LAST_TURNS = 4 -DEFAULT_MAX_SUMMARIZE_CHARS = 4000 +from config import PinchConfig + # Mirrors dispatcher.CHARS_PER_TOKEN. CHARS_PER_TOKEN = 3 @@ -37,18 +37,69 @@ def estimate_tokens(text: str | None) -> int: def extract_text(message: dict) -> str: - """Best-effort text of a message, whether content is str or content blocks.""" + """Best-effort text of a message, whether content is str or content blocks. + + ``image_url`` blocks contribute their ``url`` so image-bearing messages are + not silently undercounted as empty. A ``tool_result`` whose ``content`` is + itself a list of blocks is flattened recursively the same way. + """ content = message.get("content") if isinstance(content, str): return content if isinstance(content, list): out: list[str] = [] for part in content: - if isinstance(part, dict): - if part.get("type") == "text" and isinstance(part.get("text"), str): - out.append(part["text"]) - elif part.get("type") == "tool_result": - out.append(str(part.get("content", ""))) + if not isinstance(part, dict): + continue + ptype = part.get("type") + if ptype == "text": + text = part.get("text") + if isinstance(text, str): + out.append(text) + elif ptype == "tool_result": + inner = part.get("content") + if isinstance(inner, list): + out.append(extract_text({"content": inner})) + elif isinstance(inner, str): + out.append(inner) + elif ptype == "image_url": + image_url = part.get("image_url") or {} + url = image_url.get("url") if isinstance(image_url, dict) else None + if isinstance(url, str): + out.append(url) + return " ".join(out) + return str(content) if content else "" + + +def _text_only(message: dict) -> str: + """The prose text of a message, excluding image payloads. + + Unlike :func:`extract_text` -- which includes ``image_url`` urls so they + count toward the size estimate that triggers pruning -- this returns only + the text blocks' content. A trimmed replacement must be built from prose, + never from raw base64, or an arbitrary byte slice of an image ends up + sitting in a ``text`` field. ``extract_text`` still sizes the message and + computes savings; this only shapes the replacement. + """ + content = message.get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + out: list[str] = [] + for part in content: + if not isinstance(part, dict): + continue + ptype = part.get("type") + if ptype == "text": + text = part.get("text") + if isinstance(text, str): + out.append(text) + elif ptype == "tool_result": + inner = part.get("content") + if isinstance(inner, list): + out.append(_text_only({"content": inner})) + elif isinstance(inner, str): + out.append(inner) return " ".join(out) return str(content) if content else "" @@ -70,29 +121,58 @@ def _first_user_turn_indexes(messages: list[dict]) -> list[int]: return [i for i, m in enumerate(messages) if m.get("role") == "user"] -def _turn_of(index: int, user_indexes: list[int]) -> int: - """Which user-turn ordinal the message at ``index`` belongs to (0-based).""" - turn = 0 - for ui in user_indexes: - if ui > index: - break - turn += 1 - return max(turn - 1, 0) +def _with_text(message: dict, new_text: str) -> dict: + """Return ``message`` with its tool-result content trimmed to ``new_text``. + + When the content is a list of blocks, the text blocks become ``new_text`` + and any ``image_url`` blocks have their payload shrunk to a stub. Trimming + must reduce the *actual* bytes shipped upstream, and an image is usually + the dominant contributor, so leaving it full-size while claiming a token + saving would be a lie. The ``image_url`` block type is preserved (so the + model still knows an image was present) but the multi-MB base64 url is + replaced by a short marker. When there is no text block, one carrying + ``new_text`` is appended; a message with only an image collapses to just + the stub. + """ + content = message.get("content") + if isinstance(content, list): + blocks: list[dict] = [] + replaced_text = False + for part in content: + if not isinstance(part, dict): + blocks.append(part) + continue + ptype = part.get("type") + if ptype == "text": + if not replaced_text: + blocks.append({**part, "text": new_text}) + replaced_text = True + # Any subsequent text block is folded into ``new_text``. + elif ptype == "image_url": + # Shrink the image payload to a stub; keep the block type. + blocks.append({**part, "image_url": {"url": "[image omitted]"}}) + else: + blocks.append(part) + if not replaced_text: + blocks.append({"type": "text", "text": new_text}) + if blocks: + return {**message, "content": blocks} + return {**message, "content": new_text} def prune_context( messages: list[dict], - budget_tokens: int = DEFAULT_BUDGET_TOKENS, - keep_last_turns: int = DEFAULT_KEEP_LAST_TURNS, - max_summarize_chars: int = DEFAULT_MAX_SUMMARIZE_CHARS, + budget_tokens: int = PinchConfig.model_fields["budget_tokens"].default, + keep_last_turns: int = PinchConfig.model_fields["keep_last_turns"].default, + max_summarize_chars: int = PinchConfig.model_fields["max_summarize_chars"].default, ) -> tuple[list[dict], dict]: """Trim old tool results once a conversation exceeds ``budget_tokens``. Returns ``(pruned_messages, stats)``. Keeps every user/assistant/system - message verbatim. Tool results in the preserved window pass through; older - tool results longer than ``max_summarize_chars`` are replaced with a short - placeholder, and shorter ones are dropped entirely. Order and role pairing - are preserved, so the result is a valid conversation with the same shape. + message verbatim and never removes a message: tool results are replaced + *in place* with a short summary or placeholder, so ``len(pruned)`` always + equals ``len(messages)``. Order and role pairing are preserved, so the + result is a valid conversation with the same shape. Only runs (and only mutates anything) when the estimate actually exceeds the budget; otherwise the original list is returned untouched. @@ -107,17 +187,27 @@ def prune_context( } user_indexes = _first_user_turn_indexes(messages) - num_protected_turns = min(keep_last_turns, len(user_indexes)) - if num_protected_turns <= 0: - # Nothing to protect: everything before is a candidate for trimming. - protected_from = len(messages) + if not user_indexes: + # No user turn at all: treat the conversation as a single ongoing turn + # and protect the trailing ``keep_last_turns`` tool results (the ones + # the next turn needs). ``keep_last_turns`` bounds how many are + # protected; a tiny conversation protects all of them. + tool_indexes = [i for i, m in enumerate(messages) if m.get("role") == "tool"] + protected_from = ( + tool_indexes[max(len(tool_indexes) - keep_last_turns, 0)] + if tool_indexes + else len(messages) + ) else: - protected_from = user_indexes[len(user_indexes) - num_protected_turns] + num_protected_turns = min(keep_last_turns, len(user_indexes)) + if num_protected_turns <= 0: + # Nothing to protect: everything before is a candidate for trimming. + protected_from = len(messages) + else: + protected_from = user_indexes[len(user_indexes) - num_protected_turns] pruned: list[dict] = [] - tokens_saved = 0 summarized = 0 - dropped = 0 for i, msg in enumerate(messages): role = msg.get("role") @@ -129,35 +219,55 @@ def prune_context( pruned.append(msg) continue content = msg.get("content") - text = content if isinstance(content, str) else extract_text(msg) - if len(text) <= max_summarize_chars: - # Only replace when the placeholder is actually shorter than the - # result — dropping a tiny result to a longer marker saves nothing. + if isinstance(content, str): + text = content + prose = content + else: + # `text` sizes the message (includes image bytes, so an image-heavy + # result still triggers pruning and still affects the length guard); + # `prose` shapes the replacement (never raw base64). + text = extract_text(msg) + prose = _text_only(msg) + head_len = 1500 + tail_len = 1500 + replaced = False + if len(text) > max_summarize_chars: + # Long tool result: keep head + tail so the shape survives, but + # only when the elision actually saves characters. A result that is + # not meaningfully longer than head+tail+marker would only grow (or + # produce a negative trim), so fall through to the placeholder path. + head = prose[:head_len] + tail = prose[-tail_len:] + trimmed = len(prose) - head_len - tail_len + marker = f"\n\n[{trimmed:,} chars trimmed...]\n\n" if trimmed > 0 else "" + elided = f"{head}{marker}{tail}" + if trimmed > 0 and len(elided) < len(prose): + summarized += 1 + pruned.append(_with_text(msg, elided)) + replaced = True + if not replaced: + # Short result (or a long one that cannot be elided to save space): + # replace with a short placeholder, but only when it is strictly + # shorter than the *combined* size — an image-only result may have + # empty prose yet large real size, and shrinking it still saves + # tokens (the image payload is stubbed by _with_text). name = _tool_name(msg) or "tool" placeholder = f"[{name}: result omitted]" if len(placeholder) < len(text): - tokens_saved += estimate_tokens(text) - estimate_tokens(placeholder) - dropped += 1 - pruned.append({**msg, "content": placeholder}) + summarized += 1 + pruned.append(_with_text(msg, placeholder)) else: pruned.append(msg) - continue - # Long tool result: keep head + tail so the shape survives. - head = text[:1500] - tail = text[-1500:] - trimmed = len(text) - 3000 - tokens_saved += trimmed // CHARS_PER_TOKEN - summarized += 1 - pruned.append( - {**msg, "content": f"{head}\n\n[...{trimmed:,} chars trimmed...]\n\n{tail}"} - ) final_tokens = sum(estimate_tokens(extract_text(m)) for m in pruned) return pruned, { "pruned": True, "original_tokens": orig_tokens, "final_tokens": final_tokens, - "tokens_saved": tokens_saved, + # Savings computed from what the payload actually shrunk by (the whole + # list before vs after), so images + prose that really got trimmed are + # the only thing counted — never a per-message estimate that could + # drift from reality as #5/#11 interacted. + "tokens_saved": max(orig_tokens - final_tokens, 0), "summarized": summarized, - "dropped": dropped, } diff --git a/dispatcher.py b/dispatcher.py index df342c8..e5e457c 100644 --- a/dispatcher.py +++ b/dispatcher.py @@ -55,7 +55,7 @@ from pydantic import BaseModel, Field import logs from capabilities import detect_capabilities, iter_image_url_values from config import RouterConfig, load_config -from context_prune import prune_context +from context_prune import extract_text, prune_context from routing import ( BATCH, INTERACTIVE, @@ -149,7 +149,16 @@ logs.configure(cfg.logging.level) class TaskRequest(BaseModel): task: str = Field(..., description="The task to route.") context: Optional[str] = Field( - None, description="Assembled context (docs/code) to send with the task." + None, + description=( + "Assembled context (docs/code) to send with the task. " + "On dispatch_endpoint the chat path reuses this field to carry " + "the prior conversation turn (framing) passed to the classifier; " + "this dual use is safe because dispatch_endpoint only reads it " + "for the classify() branch, while the chat path rewrites the " + "Classification when task_category+task_tier+required_context " + "are all provided (the override branch never reads req.context)." + ), ) latency_tolerance: Optional[Literal["interactive", "batch"]] = Field( None, @@ -417,15 +426,11 @@ def classify(task: str, context: Optional[str]) -> Classification: f"Allowed values for task_category (use EXACTLY one of these strings):\n" + "\n".join(f" - {c}" for c in categories) ) - if not context: - user_content = task - else: - # llmrouter's framing: the prior turn first, labelled "Context:", so a - # short follow-up ("Yes", "Try now?") inherits its complexity rather - # than being classified in isolation as trivial. - user_content = _classifier_user_content( - task, context, cfg.classifier.context_framing - ) + # _classifier_user_content handles None context internally (returns task), + # so the outer if/else is redundant — removed for clarity. + user_content = _classifier_user_content( + task, context, cfg.classifier.context_framing + ) user_content = clamp_for_classifier(user_content, cfg.classifier.max_input_chars) client = _classifier_client() started = time.perf_counter() @@ -473,12 +478,25 @@ def _classifier_user_content( ``framing`` selects llmrouter's "Context: / Message: " layout so a short follow-up inherits its prior turn's complexity; otherwise the legacy "task / --- context --- / context" layout is used. + + In both layouts the framing instruction is appended when context is + present, telling the classifier to treat follow-ups by the prior + complexity. """ if not context: return task if framing: - return f"Context: {context}\n---\nMessage: {task}" - return f"{task}\n\n--- context ---\n{context}" + user_content = ( + f"Context: {context}\n---\nMessage: {task}" + ) + else: + user_content = f"{task}\n\n--- context ---\n{context}" + user_content += ( + "\n\nA short follow-up message ('Yes', 'Try now?', 'Go ahead') " + "continues the prior turn, so classify by the CONTEXT's " + "complexity, not in isolation." + ) + return user_content def _classify_once(client: OpenAI, system_prompt: str, user_content: str) -> Classification: @@ -1558,31 +1576,22 @@ def _last_user_text(messages: list[dict]) -> str: def _previous_context(messages: list[dict]) -> str: - """The message before the last user turn, for inheriting its complexity. + """The preceding assistant turn of the last user message. - A short follow-up ("Yes", "Try now?") continues the prior turn; feeding - that prior turn to the classifier as context lets it inherit the turn's - complexity instead of being classified in isolation as trivial — the same - rule llmrouter encodes as "short message + complex context = use context - complexity". Returns the text of the message immediately preceding the - last user message, truncated, or "" when there is no such message. + Walks backwards from the message before the last user turn, returning the + text content of the nearest ``assistant`` message. Skips ``system``, + ``user``, and ``tool`` roles so that tool output or the system prompt + never contaminates the framing signal. Uses ``context_prune.extract_text`` + for the actual content-block parsing (same logic as the former inlined + block parser, but deduplicated). Returns the concatenated text of the + first assistant message found, truncated to 200 characters, or ``""`` + when no such message exists. """ for i in range(len(messages) - 2, -1, -1): - if messages[i].get("role") == "user": + msg = messages[i] + if msg.get("role") != "assistant": continue - content = messages[i].get("content") - text = "" - if isinstance(content, str): - text = content - elif isinstance(content, list): - parts = [] - for part in content: - if isinstance(part, dict): - if part.get("type") == "text" and isinstance(part.get("text"), str): - parts.append(part["text"]) - elif part.get("type") == "tool_result": - parts.append(str(part.get("content", ""))) - text = " ".join(parts) + text = extract_text(msg) if text: return text[:200] return "" @@ -1936,13 +1945,39 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks): classified_src = decision.classification.source classifier_ms = _ms(classify_started) ctx_src = "classifier" - measured = estimate_prompt_tokens(messages, tools=body.get("tools")) + # When pinch is enabled, prune ONCE before the measured-context + # decision, so the window/tier/cost choice sees the size that will + # actually ship upstream rather than the raw conversation. The pruned + # list is reused at dispatch time (send_messages below), never pruned + # twice. When pinch is off this is a byte-for-byte no-op (full list, + # as today). Pinch trims only tool results, never user/assistant/system + # messages, so the classification turn above (which used + # `_previous_context` from the full messages) is undisturbed. + if cfg.pinch.enabled: + send_messages, pinch_stats = prune_context( + list(messages), + budget_tokens=cfg.pinch.budget_tokens, + keep_last_turns=cfg.pinch.keep_last_turns, + max_summarize_chars=cfg.pinch.max_summarize_chars, + ) + logs.debug( + "pinch", + pruned=pinch_stats["pruned"], + saved=pinch_stats["tokens_saved"], + orig=pinch_stats["original_tokens"], + final=pinch_stats["final_tokens"], + ) + else: + send_messages = messages + measured = estimate_prompt_tokens(send_messages, tools=body.get("tools")) if measured > decision.classification.required_context_tokens: ctx_src = "measured" decision = route( TaskRequest( task=_last_user_text(messages), - context=prev_context, + # context omitted: the override branch (task_category + + # task_tier + required_context_tokens all provided) + # never reads req.context — it skips classify(). latency_tolerance=latency, tools_present=tools_present, has_images=caps.has_images, @@ -2070,25 +2105,30 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks): # top (needed by the passthrough/local-vision branches), only the working # directory is derived here because only the observation path uses it. session_dir = session_directory(messages) - send_messages = messages - if cfg.pinch.enabled: - # Optional relevance-based context pruning: when the conversation - # exceeds budget_tokens, trim old tool results BEFORE any paid token is - # sent upstream. User/assistant/system messages are always kept; only - # the classifier input is exempt (it already clamps to head+tail). - send_messages, pinch_stats = prune_context( - list(messages), - budget_tokens=cfg.pinch.budget_tokens, - keep_last_turns=cfg.pinch.keep_last_turns, - max_summarize_chars=cfg.pinch.max_summarize_chars, - ) - logs.debug( - "pinch", - pruned=pinch_stats["pruned"], - saved=pinch_stats["tokens_saved"], - orig=pinch_stats["original_tokens"], - final=pinch_stats["final_tokens"], - ) + if not wants_routing: + # The routed path pruned send_messages once, before its measured-context + # decision, and the pruned list is reused here. A passthrough request + # never went through that path, so it prunes now — same gating as + # before (only when pinch is enabled). + send_messages = messages + if cfg.pinch.enabled: + # Optional relevance-based context pruning: when the conversation + # exceeds budget_tokens, trim old tool results BEFORE any paid token + # is sent upstream. User/assistant/system messages are always kept; + # only the classifier input is exempt (it already clamps to head+tail). + send_messages, pinch_stats = prune_context( + list(messages), + budget_tokens=cfg.pinch.budget_tokens, + keep_last_turns=cfg.pinch.keep_last_turns, + max_summarize_chars=cfg.pinch.max_summarize_chars, + ) + logs.debug( + "pinch", + pruned=pinch_stats["pruned"], + saved=pinch_stats["tokens_saved"], + orig=pinch_stats["original_tokens"], + final=pinch_stats["final_tokens"], + ) upstream_body = {**body, "model": target, "messages": list(send_messages)} streaming = bool(body.get("stream")) if streaming: diff --git a/tests/test_classifier_input.py b/tests/test_classifier_input.py index b8b7464..aa1c797 100644 --- a/tests/test_classifier_input.py +++ b/tests/test_classifier_input.py @@ -77,11 +77,10 @@ def test_framing_without_context_is_just_the_task(): def test_framing_puts_prior_turn_first_when_enabled(): out = _classifier_user_content("Try now?", "Design a distributed system", True) assert out.startswith("Context: Design a distributed system") - assert "\n---\nMessage: Try now?" in out - # The current task comes after the context, labelled "Message:". - ctx, _, msg = out.partition("\n---\n") - assert ctx == "Context: Design a distributed system" - assert msg == "Message: Try now?" + assert "\n---\nMessage: Try now?\n\n" in out + # The framing instruction is appended. + assert ("A short follow-up message ('Yes', 'Try now?', 'Go ahead') " + "continues the prior turn") in out def test_legacy_framing_puts_task_first(): @@ -145,6 +144,27 @@ def test_previous_context_handles_content_blocks(): assert "output" in ctx +def test_previous_context_excludes_tool_role_messages(): + # Standard tool loop: assistant(tool_call) → tool(result) → user("ok fix it") + # Should not include raw tool output as context. + messages = [ + {"role": "assistant", "content": "Let me check", "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "read_file", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "t1", "content": "/home/alee/file.py\nline 1\nline 2\nline 3"}, + {"role": "user", "content": "ok fix it"}, + ] + assert _previous_context(messages) == "Let me check" + assert "/home/alee/file.py" not in _previous_context(messages) + + +def test_previous_context_excludes_system_prompt(): + # [system, user] on the first turn — should never return a system-prompt fragment. + messages = [ + {"role": "system", "content": "You are a helpful coding assistant with tool access. You have access to read_file, write_file, and run_tests."}, + {"role": "user", "content": "hello"}, + ] + assert _previous_context(messages) == "" + + def test_previous_context_is_truncated(): messages = [ {"role": "assistant", "content": "z" * 10_000}, @@ -152,3 +172,26 @@ def test_previous_context_is_truncated(): ] assert len(_previous_context(messages)) <= 200 + +def test_previous_context_returns_assistant_text_when_it_exists(): + # [system, assistant("the real prior turn"), user("Yes")] + # Should return the assistant turn text, truncated to 200. + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "assistant", "content": "I think you should refactor the auth module to use middleware instead of decorators."}, + {"role": "user", "content": "Yes"}, + ] + ctx = _previous_context(messages) + assert "refactor" in ctx + assert str(ctx) == "I think you should refactor the auth module to use middleware instead of decorators."[:200] + + +# --- classify() caller invariant -------------------------------------------- + +def test_classify_user_content_invariant_when_no_context(): + # Regression for the removed if/else in classify(): _classifier_user_content + # already handles None context (returns task), so removing the outer guard + # must preserve user_content == task. + assert _classifier_user_content("do something", None, True) == "do something" + assert _classifier_user_content("do something", None, False) == "do something" + diff --git a/tests/test_config_endpoints.py b/tests/test_config_endpoints.py index 287f056..48ba5f6 100644 --- a/tests/test_config_endpoints.py +++ b/tests/test_config_endpoints.py @@ -269,3 +269,22 @@ def test_classifier_context_framing_loads(raw): cfg["classifier"]["context_framing"] = False assert RouterConfig(**cfg).classifier.context_framing is False + +def test_nonminimum_pinch_max_summarize_chars_is_rejected(raw): + cfg = copy.deepcopy(raw) + cfg["pinch"]["max_summarize_chars"] = 100 + with pytest.raises(ValueError, match="max_summarize_chars.*>=.*3000"): + RouterConfig(**cfg) + + +def test_pinch_max_summarize_chars_at_valid_min_loads(raw): + cfg = copy.deepcopy(raw) + cfg["pinch"]["max_summarize_chars"] = 3000 + loaded = RouterConfig(**cfg) + assert loaded.pinch.max_summarize_chars == 3000 + + +def test_pinch_max_summarize_chars_accepts_default(raw): + loaded = RouterConfig(**raw) + assert loaded.pinch.max_summarize_chars == 4000 + diff --git a/tests/test_context_prune.py b/tests/test_context_prune.py index d44b526..08dddd0 100644 --- a/tests/test_context_prune.py +++ b/tests/test_context_prune.py @@ -9,6 +9,7 @@ survives so the result still parses as a conversation. from __future__ import annotations from context_prune import estimate_tokens, extract_text, prune_context +from dispatcher import estimate_prompt_tokens def _user(text: str) -> dict: @@ -156,3 +157,229 @@ def test_estimate_tokens_is_crude_but_consistent(): assert estimate_tokens("x" * 30) == 10 assert estimate_tokens("") == 0 assert estimate_tokens(None) == 0 + + +def test_max_summarize_chars_below_3000_never_grows_or_goes_negative(): + # Regression: a max_summarize_chars well below 3000 (here 100) with a + # moderately long tool result used to head/tail-trim into something LONGER + # than the input and a negative tokens_saved. It must never grow or negate. + content = "tool output " * 100 # 1200 chars + messages = [ + _user("big user " * 5000), + _tool("read", content), + _assistant("a"), + ] + out, stats = prune_context( + messages, + budget_tokens=100, + keep_last_turns=0, + max_summarize_chars=100, + ) + tool_out = [m for m in out if m.get("role") == "tool"][0] + assert len(tool_out["content"]) < len(content) # strictly shorter + assert stats["tokens_saved"] >= 0 + assert "[-" not in tool_out["content"] # no negative trim marker + + +def test_extract_text_counts_image_url_blocks(): + # Regression: image_url blocks used to read as "", undercounting the + # message and letting pruning miss it. They must contribute their url. + msg = { + "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + ], + } + text = extract_text(msg) + assert text != "" + assert "describe this" in text + assert "AAAA" in text + + +def test_extract_text_flattens_nested_tool_result_content_list(): + # Regression: a tool_result whose content is itself a list of blocks was + # str()'d as a repr instead of being flattened like a top-level list. + msg = { + "content": [ + { + "type": "tool_result", + "content": [ + {"type": "text", "text": "nested text"}, + {"type": "image_url", "image_url": {"url": "http://x/img.png"}}, + ], + }, + ], + } + text = extract_text(msg) + assert "nested text" in text + assert "img.png" in text + + +def test_stats_use_summarized_and_length_is_always_preserved(): + # Regression: the stat was named "dropped" but nothing is ever removed — + # content is replaced in place, so len(pruned) == len(messages) always. + messages = [ + _user("q"), + _tool("search", "x" * 200), + _assistant("a"), + ] + out, stats = prune_context( + messages, + budget_tokens=50, + keep_last_turns=0, + max_summarize_chars=5000, + ) + assert "dropped" not in stats + assert "summarized" in stats + assert stats["summarized"] >= 1 + assert len(out) == len(messages) + + +def test_zero_user_messages_protect_newest_tool_result(): + # Regression: with no user message the recency guard collapsed to + # protected_from = len(messages), so even the newest tool result (the one + # the next turn needs) was trimmed. It must be protected. + messages = [ + {"role": "system", "content": "sys"}, + {"role": "assistant", "content": "call tool", "tool_calls": [{"id": "t1"}]}, + {"role": "tool", "name": "read", "content": "x" * 20000}, + ] + out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=4) + assert stats["pruned"] is True + tool_out = [m for m in out if m.get("role") == "tool"][0] + assert tool_out["content"] == messages[2]["content"] + + +def test_structured_tool_result_keeps_image_block_when_trimmed(): + # Regression: trimming flattened structured content to a string, dropping + # the image_url part permanently. The non-text block must survive — kept + # as an image_url block (so the model knows an image was present), with + # its payload shrunk to a stub rather than shipped full-size. + messages = [ + _user("big user " * 5000), + { + "role": "tool", + "name": "capture", + "content": [ + {"type": "text", "text": "some result " * 300}, # > 3000 chars + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + ], + }, + _assistant("a"), + ] + out, stats = prune_context( + messages, + budget_tokens=100, + keep_last_turns=0, + max_summarize_chars=100, + ) + tool_out = [m for m in out if m.get("role") == "tool"][0] + # The content is still a structured list, not a flattened string. + assert isinstance(tool_out["content"], list) + text_blocks = [ + b for b in tool_out["content"] + if isinstance(b, dict) and b.get("type") == "text" + ] + assert text_blocks + # Non-empty prose and non-trivial image means it was shrunk (not kept full). + img_blocks = [ + b for b in tool_out["content"] + if isinstance(b, dict) and b.get("type") == "image_url" + ] + assert len(img_blocks) == 1 + # The block survives (type preserved) but its payload is stubbed, so the + # actual shipped bytes drop rather than staying full-size (review #new). + assert img_blocks[0]["image_url"]["url"] == "[image omitted]" + assert img_blocks[0]["image_url"]["url"] != "data:image/png;base64,AAAA" + + +def test_large_image_tool_result_actually_shrinks_and_stats_are_honest(): + # Regression (follow-up review): fixing image_url counting (old #5) and + # preserving non-text blocks (old #11) composed into a case where the text + # block was sliced from prose+base64 joined together and the image_url + # block was copied through UNCHANGED — so the dominant size was never + # reduced while tokens_saved reported a large "saving" that didn't happen, + # and a raw base64 fragment leaked into the text field. The payload must + # genuinely shrink, the image payload must be stubbed, base64 must not + # leak into the text block, and tokens_saved must equal the real shrink. + huge_b64 = "A" * 200000 + messages = [ + _user("describe this"), + {"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]}, + { + "role": "tool", + "name": "screenshot", + "content": [ + {"type": "text", "text": "here is the screenshot"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{huge_b64}"}}, + ], + }, + _user("now what"), + ] + out, stats = prune_context(messages, budget_tokens=10, keep_last_turns=0) + assert stats["pruned"] is True + # The payload genuinely shrank — final below original. + assert stats["final_tokens"] < stats["original_tokens"] + # Savings equal the real reduction, not a per-message estimate. + assert stats["tokens_saved"] == ( + stats["original_tokens"] - stats["final_tokens"] + ) + tool_out = [m for m in out if m.get("role") == "tool"][0] + for part in tool_out["content"]: + if isinstance(part, dict) and part.get("type") == "image_url": + assert "AAAA" not in part["image_url"]["url"] + elif isinstance(part, dict) and part.get("type") == "text": + assert "base64" not in part.get("text", "") + + +def test_pruned_list_measures_smaller_than_the_raw_conversation(): + """Regression for review #4: routing must measure the PRUNED size. + + chat_completions used to route on the raw conversation's token count and + only shrink the upstream payload later, so a long tool-heavy session could + be sent to a pricier large-context model based on its pre-pruned size even + though pinch would bring the outgoing request well under budget. This test + pins the invariant that enables the fix: after one prune, the router's own + ``estimate_prompt_tokens`` sees a smaller count, so the window/tier/cost + decision reflects what actually ships upstream. + """ + messages = [ + {"role": "user", "content": "inspect the repo"}, + {"role": "assistant", "content": "reading files", "tool_calls": [{"id": "t1"}]}, + {"role": "tool", "name": "read", "content": "file1: " + "x" * 4000}, + {"role": "assistant", "content": "still going", "tool_calls": [{"id": "t2"}]}, + {"role": "tool", "name": "read", "content": "file2: " + "y" * 4000}, + {"role": "assistant", "content": "ok, fixing"}, + {"role": "user", "content": "fix it"}, + ] + raw_tokens = estimate_prompt_tokens(messages) + + pruned, stats = prune_context( + list(messages), + budget_tokens=100, + keep_last_turns=0, + max_summarize_chars=4000, + ) + assert stats["pruned"] is True + + pruned_tokens = estimate_prompt_tokens(pruned) + # The outgoing payload measures strictly smaller than the raw conversation, + # so routing on the pruned list picks the cheaper/smaller-window model. + assert pruned_tokens < raw_tokens + + +def test_pinch_disabled_leaves_the_measured_size_unchanged(): + """When pinch is off the routed decision measures the full conversation. + + Mirrors the byte-for-byte no-op contract: with ``prune_context`` not run, + ``estimate_prompt_tokens`` sees the raw messages exactly as before the + reordering fix. + """ + messages = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a", "tool_calls": [{"id": "t1"}]}, + {"role": "tool", "name": "read", "content": "z" * 4000}, + {"role": "user", "content": "done"}, + ] + # unpinned: identical to calling estimate_prompt_tokens(messages) directly + assert estimate_prompt_tokens(messages) == estimate_prompt_tokens(list(messages)) -- 2.49.1