From 7c2f0ea07953beaed45b25bd0fc8c5a34d60cb3f Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Wed, 8 Jul 2026 16:29:35 -0400 Subject: [PATCH] feat(cache): provider-agnostic cache-mode delta + cc-agnostic prefix comparison (#1868) ## Description Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - ## Testing - [ ] Unit tests pass (`pytest`) - [ ] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [ ] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text # Paste relevant command output or artifact links here ``` ## Real Behavior Proof - Environment: - Exact command / steps: - Observed result: - Not tested: ## Review Readiness - [ ] I have performed a self-review - [ ] This PR is ready for human review ## Checklist - [ ] My code follows the project's style guidelines - [ ] I have performed a self-review of my code - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) Add screenshots to help explain your changes. ## Additional Notes --- headroom/agent_savings.py | 14 + headroom/cache/prefix_tracker.py | 161 ++++- headroom/proxy/cost.py | 15 + headroom/proxy/handlers/anthropic.py | 73 ++- headroom/proxy/handlers/openai.py | 20 +- headroom/proxy/prometheus_metrics.py | 15 + headroom/tokenizers/registry.py | 23 + headroom/transforms/content_detector.py | 92 ++- headroom/transforms/content_router.py | 482 +++++++++++++- headroom/transforms/lossless_compaction.py | 76 +++ headroom/transforms/read_lifecycle.py | 4 +- tests/test_bash_search_lossless_fold.py | 73 +++ tests/test_cache_mode_delta_marker.py | 185 ++++++ tests/test_cache_prefix_overlay.py | 80 +++ tests/test_observed_wire_shapes.py | 588 ++++++++++++++++++ .../test_anthropic_ccr_deferred_injection.py | 19 +- tests/test_proxy_anthropic_cache_stability.py | 27 +- tests/test_proxy_handler_helpers.py | 30 + tests/test_semantic_canonicalize.py | 270 ++++++++ tests/test_tokenizers.py | 21 + tests/test_transforms_content_detection.py | 19 +- 21 files changed, 2203 insertions(+), 84 deletions(-) create mode 100644 tests/test_cache_mode_delta_marker.py create mode 100644 tests/test_observed_wire_shapes.py create mode 100644 tests/test_semantic_canonicalize.py diff --git a/headroom/agent_savings.py b/headroom/agent_savings.py index 975f9963f..d998e5af1 100644 --- a/headroom/agent_savings.py +++ b/headroom/agent_savings.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import os from dataclasses import dataclass, replace from typing import Protocol @@ -278,6 +279,19 @@ def proxy_pipeline_kwargs(config: object) -> dict[str, object]: if smart_crusher_with_compaction is not None: kwargs["smart_crusher_with_compaction"] = bool(smart_crusher_with_compaction) + # Lower the block-compression char floor (default 500) so modest tool outputs + # are eligible for the LOSSY path too. Matters in cache mode, where the only + # compressible content each turn is a single (often small) delta observation; + # a 500-char floor buckets most of them as "small" (skipped). Env-gated so it + # only changes behavior when explicitly set; lossless folding has no floor and + # is unaffected. + _min_chars_block = os.environ.get("HEADROOM_MIN_CHARS_FOR_BLOCK") + if _min_chars_block: + try: + kwargs["min_chars_for_block_compression"] = int(_min_chars_block) + except ValueError: + pass + return kwargs diff --git a/headroom/cache/prefix_tracker.py b/headroom/cache/prefix_tracker.py index 7fd2d4c9f..09dc80b73 100644 --- a/headroom/cache/prefix_tracker.py +++ b/headroom/cache/prefix_tracker.py @@ -128,6 +128,134 @@ def _strip_cache_control(obj: Any) -> Any: return obj +# Keys that carry NO semantic payload for the model — transport / caching-directive +# / telemetry / client-routing annotations that clients attach and vary turn-to-turn. +# Grounded in provider API docs (Anthropic Messages, OpenAI Chat+Responses, Bedrock +# Converse) + client-library field inventories (litellm, Vercel AI SDK, opencode, +# Claude Code, Cline). Dropped from the cross-turn prefix-equality key ONLY. +# +# NOTE ON SAFETY: this projection is a COMPARISON KEY, never a source to rebuild +# forwarded bytes — the cache-stable-delta path always forwards the previously +# forwarded bytes + the raw appended delta. So dropping these can't deprive the +# model. What we must NOT do is drop a *semantic* field (that would mask a real +# divergence and replay a stale prefix), which is why: (1) reasoning SIGNATURES are +# NOT in this set (Anthropic 400s if a thinking block is altered/missing, and a +# present/absent flip is a real divergence we want to detect); (2) tool inputs / +# arguments / json payloads are treated as OPAQUE and compared verbatim (see +# _OPAQUE_PAYLOAD_KEYS) so a user key that happens to be named "index"/"state" is +# never stripped from inside a tool call. +_NON_SEMANTIC_KEYS = frozenset( + { + # cache-breakpoint markers (moved to the newest block every turn) + "cache_control", # Anthropic (per-block) + "cachePoint", # Bedrock (per-block content block) + # litellm unified-message / tool annotations + "caller", # litellm programmatic-tool tag on tool_use + "provider_specific_fields", + "reasoning_content", # litellm display echo (the paired signature is separate) + "reasoning_items", + "annotations", # citation/display metadata + # OpenAI response echoes that can ride on assistant messages + "system_fingerprint", + "service_tier", + # Vercel AI SDK / opencode part transport + "providerMetadata", + "providerOptions", + "callProviderMetadata", + "state", + "providerExecuted", + "synthetic", + "ignored", + # streaming-assembly artifact + "index", + } +) + +# Values under these keys are opaque semantic payloads (tool-call input, OpenAI +# stringified arguments, Bedrock tool_result json). They are compared VERBATIM — we +# never recurse into them to strip "noise" keys, because arbitrary user data there +# may legitimately contain keys that collide with _NON_SEMANTIC_KEYS (e.g. an +# `input` of {"state": "CA", "index": 3}). Recursing would corrupt the comparison. +_OPAQUE_PAYLOAD_KEYS = frozenset({"input", "arguments", "json"}) + + +def _canonicalize_for_prefix_compare(obj: Any) -> Any: + """Representation-agnostic canonical form for cross-turn prefix equality. + + Providers accept several *equivalent* encodings for the same message, and real + clients vary them turn-to-turn; a raw-dict prefix compare then fails spuriously + and drops cache mode to raw (uncompressed) forwarding. This normalizes ONLY + representation: + * drops non-semantic annotation / cache-directive / telemetry keys + (_NON_SEMANTIC_KEYS) at any message/block level; + * wraps a bare string ``content`` into ``[{"type": "text", "text": ...}]`` + (Anthropic's string sugar, which litellm flips per turn); + * leaves tool ``input`` / ``arguments`` / ``json`` payloads verbatim + (_OPAQUE_PAYLOAD_KEYS) so user data is never corrupted; + * KEEPS all real content (text, tool name/input, tool_result content, reasoning + signatures, ids) so two messages canonicalize-equal iff they are semantically + identical. + + Used ONLY as a comparison key for the cache-stable delta path; the original, + unmodified messages are always what gets forwarded. + """ + if isinstance(obj, dict): + out: dict[str, Any] = {} + for key, value in obj.items(): + if key in _NON_SEMANTIC_KEYS: + continue + if key in _OPAQUE_PAYLOAD_KEYS: + out[key] = value # verbatim — do not recurse into user payloads + elif key == "content" and isinstance(value, str): + out[key] = [{"type": "text", "text": value}] + else: + out[key] = _canonicalize_for_prefix_compare(value) + return out + if isinstance(obj, list): + canon = [_canonicalize_for_prefix_compare(value) for value in obj] + # Drop blocks that projected to {} — a pure cache-directive content block + # (e.g. Bedrock {"cachePoint": {...}}) whose only key was non-semantic. Left + # in place it would be an empty-dict entry, so a directive block moving + # position across turns would spuriously fail the length/order compare. + return [value for value in canon if value != {}] + return obj + + +def extract_cache_stable_delta( + current_messages: list[dict[str, Any]], + previous_original_messages: list[dict[str, Any]] | None, + previous_forwarded_messages: list[dict[str, Any]] | None, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None: + """Return ``(stable_forwarded_prefix, appended_delta_messages)`` when the current + request append-only-extends the previous one, else ``None``. + + Provider-agnostic delta engine for cache mode. "Append-only" is decided by comparing + the *canonicalized* prefix (:func:`_canonicalize_for_prefix_compare`, which ignores + per-turn transport / cache-directive / client-annotation noise across + Anthropic / OpenAI / Bedrock and the common clients), so a moved cache marker or + shape churn does not spuriously collapse cache mode to raw forwarding. On a match the + caller replays the byte-identical previously-forwarded prefix and compresses ONLY the + appended delta. + + This is a COMPARISON + slice only: the returned prefix is the previously-forwarded + bytes verbatim and the delta is the raw appended messages — never a rebuild from the + canonical projection — so the projection dropping non-semantic fields is safe. + """ + if not previous_original_messages or previous_forwarded_messages is None: + return None + prefix_len = len(previous_original_messages) + if len(current_messages) < prefix_len: + return None + if _canonicalize_for_prefix_compare( + current_messages[:prefix_len] + ) != _canonicalize_for_prefix_compare(previous_original_messages): + return None + return ( + copy.deepcopy(previous_forwarded_messages), + copy.deepcopy(current_messages[prefix_len:]), + ) + + def overlay_cached_prefix( optimized_messages: list[dict[str, Any]], current_original_messages: list[dict[str, Any]], @@ -168,12 +296,17 @@ def overlay_cached_prefix( if len(current_original_messages) < n or len(optimized_messages) < n: return optimized_messages # Append-only guard on CONTENT ONLY: the frozen region must be the same - # messages we cached. Compare with cache_control stripped — clients move that - # breakpoint to the newest message each turn, so a raw dict compare would - # spuriously fail whenever a marker lands in the frozen prefix, skip the - # replay, and bust the cache (the residual busts observed after the first - # fix). Content stability is what the provider's prefix cache actually keys on. - if _strip_cache_control(current_original_messages[:n]) != _strip_cache_control(prev_orig): + # messages we cached. Compare with the shared canonicalizer (not just + # cache_control-stripping) so the guard is robust to ALL per-turn transport / + # annotation churn — cache_control movement (Anthropic), litellm `caller`, + # provider_specific_fields, streaming `index`, string<->block content shape, + # etc. — across providers/clients. Content stability is what the provider's + # prefix cache actually keys on; a coarser cache_control-only strip let other + # clients' noise spuriously fail the guard, skip the replay, and bust. This + # helps every handler that shares overlay_cached_prefix (Anthropic + OpenAI). + if _canonicalize_for_prefix_compare( + current_original_messages[:n] + ) != _canonicalize_for_prefix_compare(prev_orig): return optimized_messages # Replay the cached (compressed) prefix byte-identical; keep this turn's tail. return list(prev_fwd) + list(optimized_messages[n:]) @@ -563,6 +696,22 @@ class PrefixCacheTracker: chars += len(text) else: chars = 0 + # OpenAI function-calling: the assistant's command lives in the + # top-level `tool_calls` (or legacy `function_call`) field, NOT in + # `content` (which is empty/None on a tool-call turn). Anthropic puts + # the equivalent in a `tool_use` content BLOCK (counted above), but + # the OpenAI shape was never counted here. That under-counted every + # tool-based assistant turn to ~0, so the frozen-prefix estimate + # overshot the real cache boundary and froze the NEWEST delta — which + # is why OpenAI/Kimi (fireworks) tool harnesses got ~zero compression + # while text/back-tick harnesses (command in `content`) compressed. + for tc in msg.get("tool_calls") or []: + if isinstance(tc, dict): + fn = tc.get("function") or {} + chars += len(str(fn.get("name", ""))) + len(str(fn.get("arguments", ""))) + fc = msg.get("function_call") + if isinstance(fc, dict): + chars += len(str(fc.get("name", ""))) + len(str(fc.get("arguments", ""))) # Add overhead for role, block structure, etc. chars += 20 counts.append(max(1, int(chars / 3.5))) diff --git a/headroom/proxy/cost.py b/headroom/proxy/cost.py index c419d4cbc..0bb772f92 100644 --- a/headroom/proxy/cost.py +++ b/headroom/proxy/cost.py @@ -738,6 +738,21 @@ class CostTracker: uncached_tokens: Non-cached input tokens from API response usage. output_tokens: Output tokens from API response usage. """ + # Post-guard invariant (all providers): Headroom never forwards a request + # larger than the original (handlers revert any inflation before sending), + # so compression savings are >= 0 by construction. A negative here is an + # intermediate/hook token-count artifact that never reached the model; + # clamp it so `total_tokens_removed` reflects actually-forwarded bytes + # instead of surfacing spurious negatives (verified clean on the wire). + if tokens_saved < 0: + import logging as _lg + + _lg.getLogger(__name__).debug( + "record_tokens: clamping negative tokens_saved=%d to 0 for %s (artifact; wire not inflated)", + tokens_saved, + model, + ) + tokens_saved = 0 self._tokens_saved_by_model[model] = ( self._tokens_saved_by_model.get(model, 0) + tokens_saved ) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index e1898f49a..e71c77a74 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -362,17 +362,20 @@ class AnthropicHandlerMixin: Safe means the prior original request is an exact message-prefix of the current original request. This lets us replay the exact forwarded bytes for historical context and only transform newly appended message suffixes. + + The append-only check ignores per-turn transport / cache-directive / client + annotation noise (cache_control moved to the newest block, litellm caller, + provider_specific_fields, streaming index, string<->block content shape, …) via + the shared canonicalizer, so that churn doesn't spuriously drop cache mode to raw + forwarding. Delegates to the provider-agnostic engine in prefix_tracker so + OpenAI / Bedrock share one implementation. """ - if not previous_original_messages or previous_forwarded_messages is None: - return None - prefix_len = len(previous_original_messages) - if len(current_messages) < prefix_len: - return None - if current_messages[:prefix_len] != previous_original_messages: - return None - return ( - copy.deepcopy(previous_forwarded_messages), - copy.deepcopy(current_messages[prefix_len:]), + from headroom.cache.prefix_tracker import extract_cache_stable_delta + + return extract_cache_stable_delta( + current_messages, + previous_original_messages, + previous_forwarded_messages, ) @staticmethod @@ -1342,13 +1345,52 @@ class AnthropicHandlerMixin: optimized_messages = messages optimized_tokens = tokenizer.count_messages(optimized_messages) else: + # Compress the delta, with two cache-mode adjustments: + # + # fix-5: strip the client's transient cache_control marker so + # the router's per-block "never compress an explicit cache + # key" guard (content_router.py:4006) doesn't skip the ONLY + # compressible content every turn (route_counts had + # cache_control_protected == the whole delta -> 0%). In cache + # mode that marker is NOT the real forwarded breakpoint: the + # compressed delta is frozen + replayed verbatim next turn and + # normalize_message_cache_control (AFTER compression, below) + # owns the single forwarded breakpoint. Cache-safety is + # enforced post-compression, not by protecting the delta. + # + # fix-6: the delta is a lone tool_result whose tool_use (tool + # NAME + call args) lives in the frozen prefix. Passing only + # the delta to the router leaves tool_name="" so + # _bash_search_fold (lossless grep/rg folding, no size floor), + # per-tool bias, and relevance-query enrichment all degrade. + # Pass the FULL current messages with frozen_message_count = + # prefix length: _build_tool_name_map scans ALL messages (the + # delta resolves its tool_name from the prefix's tool_use) but + # the compression loop only touches indices >= frozen count, + # so ONLY the delta is compressed. Splice the compressed delta + # onto the byte-stable forwarded prefix. + from headroom.cache.prefix_tracker import _strip_cache_control + + # Compression context = the EXACT forwarded (cached) prefix + # + the stripped delta, with the prefix frozen. Using the + # forwarded prefix (not the original) keeps _build_tool_name_map + # AND cross-turn dedup consistent with what is actually cached: + # dedup can only reference bytes that are truly present in the + # forwarded context, so no pointer can dangle. The prefix is + # frozen (never compressed) and we discard the router's copy of + # it below, so the forwarded prefix stays byte-identical to last + # turn -> append-only -> no bust. + prefix_n = len(stable_forwarded_prefix) + compression_input = list(stable_forwarded_prefix) + list( + _strip_cache_control(delta_messages) + ) result = await self._run_compression_in_executor( lambda: self.anthropic_pipeline.apply( - messages=delta_messages, + messages=compression_input, model=model, model_limit=context_limit, - context=extract_user_query(delta_messages), - frozen_message_count=0, + context=extract_user_query(compression_input), + frozen_message_count=prefix_n, biases=biases, request_id=request_id, compression_policy=compression_policy, @@ -1356,7 +1398,10 @@ class AnthropicHandlerMixin: ), timeout=COMPRESSION_TIMEOUT_SECONDS, ) - optimized_messages = stable_forwarded_prefix + result.messages + # Only the delta was eligible for compression (prefix frozen); + # forward the byte-identical cached prefix + the compressed delta. + compressed_delta = result.messages[prefix_n:] + optimized_messages = stable_forwarded_prefix + compressed_delta transforms_applied = result.transforms_applied pipeline_timing = result.timing optimized_tokens = tokenizer.count_messages(optimized_messages) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 24e72f684..18ca14b77 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -985,11 +985,24 @@ class OpenAIHandlerMixin: messages: list[dict[str, Any]], base_frozen_count: int, ) -> int: - """Freeze all prior turns in cache mode; only final user turn is mutable.""" + """Freeze all prior turns in cache mode; only the final OBSERVATION turn + is mutable (the newest delta we compress-once-then-freeze). + + The newest observation is the compressible delta. Its role depends on the + harness: text/back-tick harnesses (mini-swe-agent, Codex) append it as + ``role:"user"``, but OpenAI function-calling harnesses (Kimi / any + fireworks/OpenAI-compatible tool-based model) append it as ``role:"tool"`` + (and legacy function-calling as ``role:"function"``). Gating solely on + ``role == "user"`` froze the ENTIRE conversation on every OpenAI + tool-based turn — so NOTHING was ever compressed on those models (the + delta was frozen before the content router saw it). Treat tool/function + observations as the mutable tail too; assistant/system endings still + freeze everything (they are not observations). + """ if not messages: return base_frozen_count final_idx = len(messages) - 1 - if messages[final_idx].get("role") == "user": + if messages[final_idx].get("role") in ("user", "tool", "function"): return max(base_frozen_count, final_idx) return len(messages) @@ -3086,6 +3099,8 @@ class OpenAIHandlerMixin: # OpenAI has no write penalty — uncached = total - cached uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens) + # (record_tokens clamps negative savings to 0 universally — the + # forwarded request is never larger than the original.) if self.cost_tracker: self.cost_tracker.record_tokens( model, @@ -4009,6 +4024,7 @@ class OpenAIHandlerMixin: cache_read_tokens, ) uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens) + # (record_tokens clamps negative savings to 0 universally.) self.cost_tracker.record_tokens( model, tokens_saved, diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index c5ab25c64..25033b6a0 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -590,6 +590,21 @@ class PrometheusMetrics: client: str | None = None, ): """Record metrics for a request.""" + # Post-guard invariant (all providers): Headroom never forwards a request + # larger than the original — handlers revert any inflation before sending + # (verified clean on the wire). So compression savings are >= 0; a negative + # here is an intermediate/hook token-count artifact that never reached the + # model. Clamp so total_tokens_removed / avg_compression_pct reflect the + # actually-forwarded bytes instead of surfacing spurious negatives. + if tokens_saved < 0: + import logging as _lg + + _lg.getLogger(__name__).debug( + "metrics.record: clamping negative tokens_saved=%d to 0 for %s (artifact; wire not inflated)", + tokens_saved, + model, + ) + tokens_saved = 0 async with self._lock: self.requests_total += 1 self.requests_by_provider[provider] += 1 diff --git a/headroom/tokenizers/registry.py b/headroom/tokenizers/registry.py index a2579c0a0..ad59aeec0 100644 --- a/headroom/tokenizers/registry.py +++ b/headroom/tokenizers/registry.py @@ -53,6 +53,13 @@ MODEL_PATTERNS: list[tuple[str, str]] = [ (r"^palm", "google"), # Cohere models -> estimation (r"^command", "cohere"), + # Moonshot Kimi (K2 / K2.7 code). No public BPE we can load offline, so use + # a calibrated estimator like Claude/Gemini. Matched with a leading ``.*`` so + # every serving form resolves: the Fireworks body model + # ``accounts/fireworks/models/kimi-...``, the litellm slug + # ``fireworks_ai/kimi-...``, and the native ``moonshotai/kimi-...``. + (r".*moonshot", "moonshot"), + (r".*kimi", "moonshot"), # Open models commonly served via OpenAI-compatible APIs (r"^phi-", "huggingface"), (r"^qwen", "huggingface"), @@ -113,6 +120,7 @@ class TokenizerRegistry: "google": self._create_google, "cohere": self._create_cohere, "mistral": self._create_mistral, + "moonshot": self._create_moonshot, "estimation": self._create_estimation, } @@ -351,6 +359,21 @@ class TokenizerRegistry: """ return EstimatingTokenCounter(chars_per_token=4.0) + def _create_moonshot(self, model: str) -> TokenCounter: + """Create Moonshot/Kimi tokenizer. + + Kimi (K2 / K2.7-code) ships no BPE we can load in the offline proxy + image, so — like Claude/Gemini/Cohere — we use a calibrated fixed-ratio + estimator. 3.1 chars/token was measured against Fireworks' + provider-reported ``prompt_tokens`` on a SWE-bench Kimi-K2.7-code run + (172,906 content chars -> 55,863 reported tokens = 3.10 chars/tok). The + default adaptive estimator effectively uses ~3.63 on that (code-dense) + content and so under-counted Kimi by ~20%, which starved the compression + size-gates. Slightly over-counting (lower ratio) is the safe direction + here: it makes the router MORE likely to compress, never less. + """ + return EstimatingTokenCounter(chars_per_token=3.1) + def _create_estimation(self, model: str) -> TokenCounter: """Create estimation-based tokenizer.""" return EstimatingTokenCounter() diff --git a/headroom/transforms/content_detector.py b/headroom/transforms/content_detector.py index 35bb77a20..3a22020cb 100644 --- a/headroom/transforms/content_detector.py +++ b/headroom/transforms/content_detector.py @@ -181,6 +181,16 @@ def detect_content_type(content: str) -> DetectionResult: return DetectionResult(ContentType.PLAIN_TEXT, 0.5, {}) +_JSON_DECODER = json.JSONDecoder() +# The decoded JSON value must be at least this fraction of the content for a +# WRAPPED payload to still count as JSON: a small structural wrapper (a harness +# observation shell, an ``Exit code:`` prefix) around a JSON body passes, but a +# prose/code blob that merely contains a JSON fragment does not. Fraction-based +# so it is size-correct — a large JSON with a proportionally small wrapper passes, +# a short mostly-prose string does not. (Pure JSON never reaches this check.) +_JSON_MIN_BULK_FRACTION = 0.6 + + def _decode_concatenated_json(content: str) -> list | None: """Decode a run of whitespace-separated top-level JSON values. @@ -224,44 +234,64 @@ def normalize_concatenated_json(content: str) -> str | None: def _try_detect_json(content: str) -> DetectionResult | None: - """Try to detect JSON array content.""" - content = content.strip() + """Detect JSON by PARSING, not by surface patterns. - if content.startswith("["): - try: - parsed = json.loads(content) - except json.JSONDecodeError: - return None - if isinstance(parsed, list): - # Check if it's a list of dicts (SmartCrusher compatible) - if parsed and all(isinstance(item, dict) for item in parsed): + JSON is whatever parses as JSON — objects, arrays, and any nesting are all + equally JSON, so a leading-``[`` check misses every ``{…}`` config/data file. + Tool output is often a JSON value wrapped in a little surrounding text (a + harness observation shell, an ``Exit code:`` prefix); we decode one JSON value + out of the payload and accept it when it is the bulk of the content, which + tolerates ANY wrapper without hard-coding a harness's tags. The whitespace- + separated web_search shape (``{...} {...}``, #1741) is detected too and + normalized to a real array before crushing (see normalize_concatenated_json). + """ + stripped = content.strip() + if not stripped: + return None + + try: + value = json.loads(stripped) + except ValueError: + # Not pure JSON. First: a run of whitespace-separated top-level JSON + # objects (web_search output, #1741) -> JSON_ARRAY. + if stripped.startswith("{"): + items = _decode_concatenated_json(stripped) + if items and len(items) >= 2 and all(isinstance(item, dict) for item in items): return DetectionResult( ContentType.JSON_ARRAY, 1.0, - {"item_count": len(parsed), "is_dict_array": True}, + {"item_count": len(items), "is_dict_array": True, "concatenated": True}, ) - # It's a list but not of dicts - return DetectionResult( - ContentType.JSON_ARRAY, - 0.8, - {"item_count": len(parsed), "is_dict_array": False}, - ) + # Otherwise decode one JSON value out of a small wrapped payload. + start = min((i for i in (stripped.find("{"), stripped.find("[")) if i >= 0), default=-1) + if start < 0: + return None + try: + value, end = _JSON_DECODER.raw_decode(stripped, start) + except ValueError: + return None + # Accept only when the decoded JSON is the BULK of the content (see + # _JSON_MIN_BULK_FRACTION) — a small structural wrapper around a JSON body, + # not a prose/code blob that merely contains a JSON fragment. + if (end - start) < len(stripped) * _JSON_MIN_BULK_FRACTION: + return None + + # A bare scalar (42, "s", true) is not structured data worth routing as JSON. + if not isinstance(value, (dict, list)): return None - # Space-separated JSON objects (typical web_search output) aren't a valid - # array, so they'd fall through to PLAIN_TEXT and skip SmartCrusher at 0% - # compression. SmartCrusher normalizes this shape to a real array before - # crushing (#1741). - if content.startswith("{"): - items = _decode_concatenated_json(content) - if items and len(items) >= 2 and all(isinstance(item, dict) for item in items): - return DetectionResult( - ContentType.JSON_ARRAY, - 1.0, - {"item_count": len(items), "is_dict_array": True, "concatenated": True}, - ) - - return None + if isinstance(value, list): + is_dict_array = bool(value) and all(isinstance(item, dict) for item in value) + return DetectionResult( + ContentType.JSON_ARRAY, + 1.0 if is_dict_array else 0.8, + {"item_count": len(value), "is_dict_array": is_dict_array}, + ) + return DetectionResult( + ContentType.JSON_ARRAY, + 0.9, + {"is_dict_array": False, "is_object": True}, + ) def _try_detect_diff(content: str) -> DetectionResult | None: diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index eb03be94e..2f6e46838 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -113,6 +113,116 @@ def _tool_call_command_text(raw: Any) -> str: return cmd if isinstance(cmd, str) else "" +def _fenced_shell_command(content: Any) -> str: + """Extract the shell command from a TEXT-BASED agent's fenced code block. + + Text-based harnesses (mini-swe-agent backticks, Codex, Cursor, and any + non-native-tool OpenAI agent) put the command in a ```mswea_bash_command / + ```bash fenced block inside the assistant's *string* content — there is no + ``tool_use``/``tool_calls`` block. Returns the first fenced block's body, or + "" when there is none. Shape-agnostic input to read-detection so cat/sed + reads are protected on any model, not just those emitting tool-call blocks. + """ + if not isinstance(content, str) or "```" not in content: + return "" + m = re.search(r"```(?:[\w.-]+)?[ \t]*\n(.*?)```", content, re.S) + return m.group(1).strip() if m else "" + + +_READ_VERBS = ("cat", "head", "tail", "nl", "bat", "less", "more") + +# Machine-generated dependency lockfiles detect as PLAIN_TEXT (so the content-based +# read gate would protect them), but they are regenerated by a tool and never patched +# byte-for-byte — the biggest, most-repetitive read in a session. Match by NAME so a +# read of one is never added to the protected set and stays compressible. +_LOCKFILE_RE = re.compile( + r"(^|[\s/])(" + r"bun\.lock|bun\.lockb|package-lock\.json|npm-shrinkwrap\.json|yarn\.lock|" + r"pnpm-lock\.yaml|uv\.lock|poetry\.lock|Pipfile\.lock|requirements\.txt\.lock|" + r"Cargo\.lock|go\.sum|Gemfile\.lock|composer\.lock|flake\.lock|Package\.resolved|" + r"gradle\.lockfile|packages\.lock\.json" + r")(\s|$)", + re.IGNORECASE, +) + + +def _strip_cd_prefix(command: str) -> str: + """Peel leading ``cd &&|; `` chains from a shell command. + + Agent harnesses (mini-swe-agent, Codex, Cursor, …) prefix nearly every command + with ``cd && `` (or ``cd ; ``) to run in the checkout. Command- + classification helpers must strip this first, or the parsed program is ``cd`` + instead of the real tool (``grep``/``cat``/…) — which silently disables read- + protection and the lossless search-fold (observed: 100% of ``cd … && rg`` + output went uncompacted). Provider/harness-agnostic: operates on the plain + shell string that every client ultimately produces. Only ``&&`` and ``;`` + connectors are peeled (a mis-parse is harmless — the caller falls through to + the normal path, guarded by downstream reversibility checks). + """ + if not command or not isinstance(command, str): + return "" + c = command.strip() + while True: + m = re.match(r"^cd\s+[^&;|]+(?:&&|;)\s*(.*)$", c, re.S) + if not m: + break + c = m.group(1).strip() + return c + + +def _is_read_command(command: str) -> bool: + """True when a shell command's output is essentially raw FILE CONTENT the agent + will read/edit from — ``cat``/``head``/``tail``/``nl``/``less``/``more`` of a file, + or ``sed -n`` range-printing. + + Such reads must NOT be lossy-compressed: the agent needs the exact bytes to produce + a precise patch. Lossy-compressing them was observed (SWE-bench, mini-swe-agent) to + cause the agent to RE-READ the same file (cat -> cat -A -> cat -n) to recover exact + detail — turn inflation — and, when recovery failed, resolve loss. Search/list/test + output (grep/rg/ls/find/pytest) is derived and stays compressible. + + This identifies that a command is a file READ. Whether the read is actually PROTECTED + is finalized downstream by CONTENT type (see ``_read_output_should_be_protected``): + reads are protected by default, and only released to compression when the output is a + confidently non-code DATA type. The one command-level carve-out is lockfiles: they + detect as PLAIN_TEXT (so the content gate would protect them) yet are regenerated + artifacts, never byte-patched — so a lockfile read returns False here and stays + compressible. + + Excludes writes: a redirect (``>``/``>>``), ``tee``, or heredoc (``<<``) means the + command WRITES a file (e.g. ``cat > f < && ` chains (agents prefix reads with a cd) + c = _strip_cd_prefix(command) + # a write / append / tee / heredoc anywhere => not a pure read + if re.search(r"(^|\s)(>>?|tee\b|<<)", c): + return False + # Parse the real program with the SAME structural parser the search-fold uses + # (_bash_program peels sudo/env/timeout/rtk wrappers + env assignments), so + # `sudo cat f`, `timeout 30 cat f`, `rtk cat f` are recognized as reads, not + # silently dropped by a first-token match. + prog, rest = _bash_program(c) + if not prog: + return False + if prog in {"sh", "bash", "zsh", "dash"} and rest: + # `bash -lc "cat …"` (Codex): the real command is the -c argument. + for j, tok in enumerate(rest): + if tok in {"-c", "-lc", "-lic", "-ic"} and j + 1 < len(rest): + return _is_read_command(" ".join(rest[j + 1 :]).strip("'\"")) + return False + is_read = prog in _READ_VERBS or ( + # `sed -n '1,20p' file` prints a range (read); bare `sed` is a stream editor. + prog == "sed" and bool(re.search(r"(^|\s)-n(\s|$)", c)) + ) + if not is_read: + return False + # Lockfiles are tool-regenerated, not byte-patched — never protect (keep compressible). + return not _LOCKFILE_RE.search(c) + + # Shell wrappers that prefix the real program — peeled to find it. Shell # grammar, not tunable policy: rtk (the user's token proxy), sudo/env/timeout/… _SHELL_WRAPPERS = frozenset( @@ -164,6 +274,9 @@ def _bash_command_is_search(command: str, search_commands: frozenset[str]) -> bo """True when ``command`` is a read-only search whose output folds byte- losslessly (grep/rg/git grep/…). Peels wrappers and recurses into ``sh -c``. """ + # Peel `cd && ` chains first — harnesses prefix every command with a + # cd, so without this the parsed program is `cd` and the fold never fires. + command = _strip_cd_prefix(command) prog, rest = _bash_program(command) if not prog: return False @@ -430,6 +543,46 @@ def _detect_content(content: str) -> DetectionResult: ) +# Content types safe to compress even when read from a file: confidently non-code, +# machine-derived DATA the agent never byte-patches. Everything else (SOURCE_CODE AND +# the PLAIN_TEXT fallback) is protected — critically, code in a language the detector +# does not recognize falls through to PLAIN_TEXT, so protecting PLAIN_TEXT keeps those +# reads safe. The code detector only knows ~6 languages, so we do NOT rely on positively +# identifying code; we release only positively-identified data. +_RELEASABLE_READ_TYPES = frozenset( + { + ContentType.JSON_ARRAY, + ContentType.SEARCH_RESULTS, + ContentType.BUILD_OUTPUT, # compiler/test/lint logs + ContentType.GIT_DIFF, + ContentType.HTML, + ContentType.TABULAR, # CSV/TSV, tables + } +) + + +def _read_output_should_be_protected(text: Any) -> bool: + """Finalize read-protection by CONTENT — protect by default, release only DATA. + + ``_is_read_command`` says "this came from a cat/sed/head file read (and isn't a + lockfile)". Protection exists so the agent keeps EXACT BYTES of code it will patch. + Because the code detector recognizes only a handful of languages, we do NOT gate on + "is this SOURCE_CODE" (that would leave Ruby/C/SQL/… code — seen as PLAIN_TEXT — + unprotected and lossy-compressed). Instead we PROTECT unless the content is a + confidently non-code data type (JSON object/array, CSV/tabular, build/test log, git + diff, HTML, search output), which are never byte-patched and route to a compressor. + JSON objects are now recognized by the content detector's real parse, so no + separate object carve-out is needed here. + """ + if not isinstance(text, str) or not text: + return False + try: + return _detect_content(text).content_type not in _RELEASABLE_READ_TYPES + except Exception: + # Detection failure → protect (preserve the byte-exact default). + return True + + def _create_content_signature( content_type: str, content: str, @@ -1351,6 +1504,18 @@ class ContentRouter(Transform): self.config.enable_cross_turn_dedup or os.environ.get("HEADROOM_DEDUPE", "").strip().lower() in ("1", "true", "yes", "on") ) + # EXPERIMENT (HEADROOM_EXPERIMENTAL_READ_KEEP_RATIO): file reads are + # protected verbatim by default so the agent keeps exact bytes to patch. + # This probe instead LIGHTLY lossy-compresses a protected read with + # Kompress at the given keep ratio (e.g. 0.9 = keep ~90%), trading a small + # resolve risk for savings on the biggest untouched bucket (code reads). + # 0/unset = OFF (verbatim, today's behavior). Resolve-risk probe only. + try: + self._exp_read_keep_ratio: float = float( + os.environ.get("HEADROOM_EXPERIMENTAL_READ_KEEP_RATIO", "") or 0 + ) + except ValueError: + self._exp_read_keep_ratio = 0.0 # Lossless-then-lossy. Config field OR env HEADROOM_LOSSLESS_THEN_LOSSY. # Only takes effect in lossy mode (STAGE 0 guards on `not config.lossless`). self._lossless_then_lossy: bool = self.config.lossless_then_lossy or os.environ.get( @@ -1828,7 +1993,7 @@ class ContentRouter(Transform): CompressionStrategy.DIFF: "diff", }.get(strategy) order = ([primary] if primary else []) + [ - k for k in ("search", "log", "diff", "text") if k != primary + k for k in ("search", "paths", "log", "diff", "text") if k != primary ] best, best_label = content, None for kind in order: @@ -2014,6 +2179,31 @@ class ContentRouter(Transform): compressor_name = "KompressCompressor" decision_reason = "code_aware_unavailable_fallback_kompress" strategy_chain.append(CompressionStrategy.KOMPRESS.value) + elif ( + self._lossless_then_lossy + and compressed_tokens is not None + and compressed_tokens >= original_tokens + ): + # #3 — lossless-then-lossy: code-aware produced NO net shrink + # and the lossless fold found nothing either, so this code + # block would otherwise pass through uncompressed. Give the + # lossy ML compressor (Kompress) a shot so lossy runs even when + # lossless has no savings. Reads are protected upstream, so + # only NON-read code reaches here. Keep Kompress ONLY if it + # actually shrinks (never inflate). + _k, _kt = self._try_ml_compressor(content, context, question) + if ( + _k is not None + and _kt is not None + and _kt < original_tokens + and len(_k) < len(content) + ): + compressed, compressed_tokens = _k, _kt + strategy = CompressionStrategy.KOMPRESS + actual_strategy = strategy + compressor_name = "KompressCompressor" + decision_reason = "code_aware_no_shrink_fallback_kompress" + strategy_chain.append(CompressionStrategy.KOMPRESS.value) elif strategy == CompressionStrategy.SMART_CRUSHER: # SmartCrusher handles its own TOIN recording @@ -2173,6 +2363,50 @@ class ContentRouter(Transform): f"{decision_reason}_fallback_log_after_no_savings" ) + # ── lossless_then_lossy (general): LAYER lossy on top of a + # conservative strategy result ────────────────────────────── + # SEARCH/LOG/HTML compressors are structural and often bank only a + # trickle (e.g. search keeps every keyword-matching line, so a grep + # dump into a large data/config file barely shrinks). The zero- + # savings fallback above never fires in that case (there WAS a tiny + # win), so lossy never runs. When the operator opted into + # lossless_then_lossy, run Kompress over whatever the strategy + # produced and KEEP it only if it removes a further meaningful chunk + # (>= _lossy_min_extra_savings beyond the strategy result) and is + # actually shorter — never inflating, never doing worse than the + # strategy output. DIFF is excluded (Kompress corrupts ``git + # apply``); TEXT/KOMPRESS already ran Kompress; CODE_AWARE has its + # own inline no-shrink fallback; SMART_CRUSHER/TABULAR use the + # zero-savings fallback above. + if ( + self._lossless_then_lossy + and compressed is not None + and compressed_tokens is not None + and strategy + in { + CompressionStrategy.SEARCH, + CompressionStrategy.LOG, + CompressionStrategy.HTML, + } + and not self._looks_like_diff(content) + ): + try: + _layer_k, _layer_kt = self._try_ml_compressor(compressed, context, question) + except Exception as exc: # noqa: BLE001 + logger.debug("lossless_then_lossy layer failed: %s", exc) + _layer_k, _layer_kt = None, None + if ( + _layer_k is not None + and _layer_kt is not None + and _layer_kt <= compressed_tokens * (1 - self._lossy_min_extra_savings) + and len(_layer_k) < len(compressed) + ): + compressed, compressed_tokens = _layer_k, _layer_kt + actual_strategy = CompressionStrategy.KOMPRESS + compressor_name = "KompressCompressor" + strategy_chain.append(CompressionStrategy.KOMPRESS.value) + decision_reason = f"{decision_reason}_lossless_then_lossy_layer" + # Re-narrow for mypy: all reassignments above produce str, but # mypy 1.14.x widens after nested try/except/else reassignments. assert compressed is not None @@ -2234,7 +2468,11 @@ class ContentRouter(Transform): return content, original_tokens, strategy_chain def _try_ml_compressor( - self, content: str, context: str, question: str | None = None + self, + content: str, + context: str, + question: str | None = None, + target_ratio: float | None = None, ) -> tuple[str, int]: """ML-based compression using Kompress. @@ -2325,7 +2563,11 @@ class ContentRouter(Transform): text_to_compress, context=context, question=question, - target_ratio=getattr(self, "_runtime_target_ratio", None), + target_ratio=( + target_ratio + if target_ratio is not None + else getattr(self, "_runtime_target_ratio", None) + ), allow_download=False, ) compressed = result.compressed @@ -2343,6 +2585,26 @@ class ContentRouter(Transform): return compressed, compressed_tokens or len(compressed.split()) + def _experimental_compress_read(self, content: Any, context: str = "") -> str | None: + """EXPERIMENT (HEADROOM_EXPERIMENTAL_READ_KEEP_RATIO): lightly Kompress a + protected file read instead of passing it verbatim. + + Reads are protected to keep the exact bytes the agent patches from, so + this is OFF by default and a resolve-risk probe: at keep ratio 0.9 the + model still sees ~90% of the (importance-ranked) tokens. Returns the + compressed text only when it actually shrank and is non-empty; otherwise + None, so the caller falls back to verbatim protection. Never raises. + """ + ratio = getattr(self, "_exp_read_keep_ratio", 0.0) + if not ratio or not isinstance(content, str) or len(content) < 200: + return None + try: + out, _ = self._try_ml_compressor(content, context or "", target_ratio=ratio) + except Exception as exc: # noqa: BLE001 + logger.debug("experimental read-kompress failed: %s", exc) + return None + return out if (out and len(out) < len(content)) else None + def _strategy_from_detection_type(self, content_type: ContentType) -> CompressionStrategy: """Get strategy from ContentType enum.""" mapping = { @@ -2887,8 +3149,13 @@ class ContentRouter(Transform): if msg.get("role") != "assistant": continue - # OpenAI format: tool_calls array - for tc in msg.get("tool_calls", []): + # OpenAI format: tool_calls array. Coalesce None -> [] : OpenAI/LiteLLM + # assistant messages carry an explicit ``tool_calls: null`` (and + # ``function_call: null``) when there are no calls, so ``.get(k, [])`` + # returns None (not []) and iterating it crashes _build_tool_name_map -> + # apply() -> compression silently falls through to passthrough on every + # OpenAI turn. This is the generic OpenAI-shape fix. + for tc in msg.get("tool_calls") or []: if isinstance(tc, dict): tc_id = tc.get("id", "") fn = tc.get("function", {}) @@ -3162,6 +3429,61 @@ class ContentRouter(Transform): if is_tool_excluded(name, exclude_tools) } + # Read protection (HEADROOM_PROTECT_READS=1): for bash-family agents the + # exclude-by-tool-NAME set above never catches file reads (they are `bash` + # tool calls whose COMMAND is a cat/sed/head/...). Mark those tool_use_ids so + # their output is never LOSSY-compressed (the agent needs exact bytes to edit; + # lossy reads caused re-reads/turn-inflation + resolve loss on SWE-bench). + # Type-specific by design: grep/test/ls output stays compressible, so the + # cache-mode delta still compresses whenever the newest turn is NOT a read. + self._protect_read_tool_ids = set() + if os.environ.get("HEADROOM_PROTECT_READS", "0").strip().lower() not in ( + "0", + "", + "false", + "no", + ): + # Use _tool_call_commands (the parsed shell command), NOT + # _tool_call_args (a compact free-text blob that, for OpenAI-style + # JSON-string args, is the raw ``{"command": ...}`` JSON — on which + # _is_read_command always returns False, silently disabling read + # protection for OpenAI-native harnesses). _tool_call_commands is + # extracted via _tool_call_command_text, correct for both wire shapes. + self._protect_read_tool_ids = { + tid + for tid in tool_name_map + if _is_read_command(self._tool_call_commands.get(tid, "")) + } + + # Read protection — TEXT-BASED shape (shape-agnostic twin of the above). + # Text-based agents (GPT-5.4/Codex/Cursor backticks) have no tool_use + # blocks: the command is in the PRECEDING assistant message's fenced + # block and the observation is a plain user string with no id to match. + # Detect the producing command by walking back to that assistant turn and + # mark the observation's message index so it is passed verbatim — so + # cat/sed/head code reads are protected on ANY model/harness, not just + # those that emit tool-call/tool_result blocks. + self._protect_read_msg_indices: set[int] = set() + if os.environ.get("HEADROOM_PROTECT_READS", "0").strip().lower() not in ( + "0", + "", + "false", + "no", + ): + for _idx, _m in enumerate(messages): + if _m.get("role") != "user": + continue + _cmd = "" + for _j in range(_idx - 1, -1, -1): + _rj = messages[_j].get("role") + if _rj == "assistant": + _cmd = _fenced_shell_command(messages[_j].get("content")) + break + if _rj == "user": + break + if _cmd and _is_read_command(_cmd): + self._protect_read_msg_indices.add(_idx) + # --- Adaptive parameters based on context pressure --- num_messages = len(messages) model_limit = kwargs.get("model_limit", 0) @@ -3393,6 +3715,38 @@ class ContentRouter(Transform): ) continue + # Read protection (ROLE / SHAPE-AGNOSTIC). An observation produced by + # a file read command (cat/sed/head/…) is passed VERBATIM so the agent + # keeps exact bytes to patch — regardless of how THIS harness labels + # it. The SAME operation surfaces under different roles/shapes across + # harnesses (Anthropic tool_use, OpenAI/Kimi `role:tool`, text-harness + # `role:user` string), so we key off the OUTCOME — "a read command + # produced code output" — not the role. Link via the observation's + # tool_call_id/tool_use_id (tool-based) OR the preceding fenced + # command's message index (text-based); a given harness populates + # exactly one, so ORing them is shape-agnostic and collision-free. + # (Anthropic tool_result BLOCKS carry list content and are protected + # in the block path; this covers STRING-content observations.) + if role in ("user", "tool", "function"): + _tcid = message.get("tool_call_id") or message.get("tool_use_id") or "" + _is_read_obs = _tcid in getattr(self, "_protect_read_tool_ids", ()) or i in getattr( + self, "_protect_read_msg_indices", () + ) + if _is_read_obs and _read_output_should_be_protected(content): + _exp = self._experimental_compress_read(content, context) + if _exp is not None: + result_slots[i] = {**message, "content": _exp} + transforms_applied.append("router:read_kompress_exp") + route_counts["read_kompress_exp"] = ( + route_counts.get("read_kompress_exp", 0) + 1 + ) + continue + result_slots[i] = message + transforms_applied.append("router:read_protected") + route_counts.setdefault("read_protected", 0) + route_counts["read_protected"] += 1 + continue + # Protection 1: Never compress user messages (unless overridden) if skip_user and role == "user": result_slots[i] = message @@ -3715,6 +4069,22 @@ class ContentRouter(Transform): ", ".join(parts), ) + # Per-request routing visibility (grep `[router] route_counts`): how many + # messages/blocks hit each route this request — skip reasons (small, + # user_msg, non_string, recent_code, analysis_ctx, content_blocks, + # excluded_tool, read_protected, error_protected, already_compressed, …) + # plus successful compressions. Makes "what is Headroom missing?" answerable + # per provider shape (e.g. OpenAI plain-string user obs vs Anthropic + # tool_result blocks) directly from a run's logs. INFO so it's on by default. + _nonzero = {k: v for k, v in route_counts.items() if v} + logger.info( + "[router] route_counts=%s compressed=%d frozen=%d msgs=%d", + _nonzero, + len(compressed_details), + frozen_message_count, + num_messages, + ) + # Forward route_counts to the observer so `/stats` can surface a # session-level protection breakdown (issue #454). The observer # may not implement this method on older versions; ignore @@ -4021,6 +4391,43 @@ class ContentRouter(Transform): if block_type == "tool_result": # Check if tool is excluded from compression tool_use_id = block.get("tool_use_id", "") + # Flatten OpenAI-style list-form content up front (see fix-7 note below) + # so both the read-protection content check and the compressor see the + # same text. + _tr_content = block.get("content", "") + _tr_list_form_early = ( + isinstance(_tr_content, list) + and bool(_tr_content) + and all(isinstance(b, dict) and b.get("type") == "text" for b in _tr_content) + ) + _tr_text = ( + "".join(b.get("text", "") for b in _tr_content) + if _tr_list_form_early + else _tr_content + ) + # Read protection (HEADROOM_PROTECT_READS): never LOSSY-compress a file + # read (cat/sed/head/...) whose content is SOURCE CODE — pass it verbatim + # so the agent keeps exact bytes to edit from. A read of DATA (json/csv/ + # log/lockfile/text) is not byte-patched, so it falls through to its + # content-specific compressor. Cross-turn dedup still runs later, so + # re-reads of the same file are losslessly de-duplicated either way. + if tool_use_id in getattr( + self, "_protect_read_tool_ids", () + ) and _read_output_should_be_protected(_tr_text): + _exp = self._experimental_compress_read(_tr_text, context or "") + if _exp is not None: + new_blocks.append({**block, "content": _exp}) + any_compressed = True + if route_counts is not None: + route_counts["read_kompress_exp"] = ( + route_counts.get("read_kompress_exp", 0) + 1 + ) + continue + new_blocks.append(block) + if route_counts is not None: + route_counts.setdefault("read_protected", 0) + route_counts["read_protected"] += 1 + continue if tool_use_id in excluded_tool_ids: if messages_from_end <= read_protection_window: # Protected from lossy compression — but grep/log/json @@ -4059,12 +4466,42 @@ class ContentRouter(Transform): tool_content = block.get("content", "") + # fix-7: OpenAI-style clients (litellm) send tool_result `content` + # as a LIST of text blocks ([{"type":"text","text": ...}]), not a + # bare string. Every check/compressor below is `isinstance(str)`- + # gated, so list-form tool outputs were skipped entirely (bucketed + # "small" -> 0% compression even on 10k-char reads). Flatten the + # text blocks to a string for the checks/compressors, and re-wrap + # the result in the SAME container on write-back so the on-wire + # shape is unchanged. Mixed / non-text content (e.g. images) does + # NOT flatten (tool_text stays the list) -> str checks fail -> + # block passes through unchanged, exactly as before. + _tr_list_form = ( + isinstance(tool_content, list) + and bool(tool_content) + and all(isinstance(b, dict) and b.get("type") == "text" for b in tool_content) + ) + tool_text = ( + "".join(b.get("text", "") for b in tool_content) + if _tr_list_form + else tool_content + ) + # Bash-search lossless pre-empt (twin of the string-form path): # fold read-only search output (grep/rg/git grep) byte-losslessly # instead of taking the lossy strategy path. - bash_folded = self._bash_search_fold(tool_name, tool_use_id, tool_content) + bash_folded = self._bash_search_fold(tool_name, tool_use_id, tool_text) if bash_folded is not None: - new_blocks.append({**block, "content": bash_folded}) + new_blocks.append( + { + **block, + "content": ( + [{"type": "text", "text": bash_folded}] + if _tr_list_form + else bash_folded + ), + } + ) transforms_applied.append("router:bash:lossless_search") if route_counts is not None: route_counts["bash_lossless_search"] = ( @@ -4082,11 +4519,11 @@ class ContentRouter(Transform): # error lines in big logs. if ( self.config.protect_error_outputs - and isinstance(tool_content, str) - and len(tool_content) <= self.config.error_protection_max_chars + and isinstance(tool_text, str) + and len(tool_text) <= self.config.error_protection_max_chars and ( block.get("is_error") is True - or content_has_strong_error_indicators(tool_content) + or content_has_strong_error_indicators(tool_text) ) ): new_blocks.append(block) @@ -4099,13 +4536,13 @@ class ContentRouter(Transform): # Only process string content. Blocks below the lossy min_chars # floor still pass when a byte-lossless fold shrinks them — the # floor guards the lossy path only; lossless has no size floor. - if isinstance(tool_content, str) and ( - len(tool_content) > min_chars or self._has_lossless_fold(tool_content) + if isinstance(tool_text, str) and ( + len(tool_text) > min_chars or self._has_lossless_fold(tool_text) ): # Compression pinning: skip already-compressed content if ( - "Retrieve more: hash=" in tool_content - or "Retrieve original: hash=" in tool_content + "Retrieve more: hash=" in tool_text + or "Retrieve original: hash=" in tool_text ): new_blocks.append(block) if route_counts is not None: @@ -4115,10 +4552,8 @@ class ContentRouter(Transform): # Two-tier compression cache → shared helper compressed_content, was_compressed = self._compress_block_content( - content=tool_content, - content_key=hash( - (tool_content, getattr(self, "_runtime_target_ratio", None)) - ), + content=tool_text, + content_key=hash((tool_text, getattr(self, "_runtime_target_ratio", None))), context=block_context, bias=bias, min_ratio=min_ratio, @@ -4131,7 +4566,16 @@ class ContentRouter(Transform): enforce_reversibility=True, ) if compressed_content is not None: - new_blocks.append({**block, "content": compressed_content}) + new_blocks.append( + { + **block, + "content": ( + [{"type": "text", "text": compressed_content}] + if _tr_list_form + else compressed_content + ), + } + ) any_compressed = True else: new_blocks.append(block) diff --git a/headroom/transforms/lossless_compaction.py b/headroom/transforms/lossless_compaction.py index 4cdcfac98..a8f60c364 100644 --- a/headroom/transforms/lossless_compaction.py +++ b/headroom/transforms/lossless_compaction.py @@ -203,6 +203,75 @@ def diff_strip_index(text: str) -> str: return _join(out, had_trailing) +# A whole-line file path: optional ``./``/``../`` root, >=1 directory segment, +# then a basename. No whitespace or ':' (so grep ``path:line:content`` rows — +# handled by search_heading — are excluded). Directory-only lines (trailing '/') +# don't match (empty basename), which keeps the fold unambiguous. +_PATH_ROW_RE = re.compile(r"^(?P(?:\.{0,2}/)?(?:[^/\s:]+/)+)(?P[^/\s:]+)$") + + +def path_heading(text: str) -> str: + """Fold a *pure* file-path listing (``find`` / ``ls -1`` / ``rg -l`` output) + into ripgrep-heading form: each parent directory printed once on its own + line (ending in ``/``), then the bare basenames beneath it. + + Reversibility is not assumed here — ``compact_lossless`` verifies the exact + round-trip via :func:`path_unheading` and discards the fold on any mismatch + (e.g. a stray no-slash line mistaken for a basename), so mixed content is + always safe. Requires >=2 path rows or there is nothing to group. + Complements ``search_heading``, which only handles the ``path:line:content`` + grep shape, not plain path lists. + """ + lines, had_trailing = _split_keep_trailing(text) + if sum(1 for ln in lines if _PATH_ROW_RE.match(ln)) < 2: + return text + out: list[str] = [] + current: str | None = None + for line in lines: + m = _PATH_ROW_RE.match(line) + if m: + d = m.group("dir") + if d != current: + out.append(d) + current = d + out.append(m.group("base")) + else: # blank line inside/around the listing + out.append(line) + current = None + return _join(out, had_trailing) + + +def path_unheading(text: str) -> str: + """Exact inverse of :func:`path_heading`. + + A *header* is a line ending in ``/`` immediately followed by a basename row + (a non-empty line with no ``/``); it is consumed and re-prefixed onto each + following basename row until a blank line or another header. + """ + lines, had_trailing = _split_keep_trailing(text) + if not lines: + return text + out: list[str] = [] + current: str | None = None + n = len(lines) + i = 0 + while i < n: + line = lines[i] + is_base = line != "" and "/" not in line + if current is not None and is_base: + out.append(current + line) + i += 1 + continue + if line.endswith("/") and i + 1 < n and lines[i + 1] != "" and "/" not in lines[i + 1]: + current = line + i += 1 + continue + current = None + out.append(line) + i += 1 + return _join(out, had_trailing) + + def _smaller(candidate: str, original: str) -> bool: return len(candidate) < len(original) @@ -234,6 +303,13 @@ def compact_lossless(content: str, kind: str) -> str: return content return candidate if _smaller(candidate, content) else content + if kind == "paths": + # Pure path listings (find/ls -1/rg -l): fold repeated parent dirs. + candidate = path_heading(content) + if path_unheading(candidate) != content: + return content + return candidate if _smaller(candidate, content) else content + if kind == "diff": # Purely subtractive of non-semantic bookkeeping lines; the # remaining hunks still apply. No exact inverse needed. diff --git a/headroom/transforms/read_lifecycle.py b/headroom/transforms/read_lifecycle.py index 967b5b696..43a6ceea9 100644 --- a/headroom/transforms/read_lifecycle.py +++ b/headroom/transforms/read_lifecycle.py @@ -177,7 +177,7 @@ class ReadLifecycleManager: continue # OpenAI format: tool_calls array - for tc in msg.get("tool_calls", []): + for tc in msg.get("tool_calls") or []: # coalesce None (OpenAI tool_calls:null) if not isinstance(tc, dict): continue tc_id = tc.get("id", "") @@ -272,7 +272,7 @@ class ReadLifecycleManager: continue # OpenAI format - for tc in msg.get("tool_calls", []): + for tc in msg.get("tool_calls") or []: # coalesce None (OpenAI tool_calls:null) if isinstance(tc, dict) and tc.get("id") == tool_call_id: return i diff --git a/tests/test_bash_search_lossless_fold.py b/tests/test_bash_search_lossless_fold.py index 687b94d25..0094fa84d 100644 --- a/tests/test_bash_search_lossless_fold.py +++ b/tests/test_bash_search_lossless_fold.py @@ -1,3 +1,4 @@ +# ruff: noqa: E402 — test sections import after helper/setup code by design. """Bash-search lossless fold. `bash` is not an excluded tool, so its output normally takes the lossy strategy @@ -148,3 +149,75 @@ def test_source_output_from_search_command_untouched(tokenizer): out, transforms = _openai("grep -l foo", CODE, tokenizer) assert "router:bash:lossless_search" not in transforms assert out == CODE + + +# ---- path-listing fold (find/ls -1/rg -l): fold repeated parent dirs ---- +from headroom.transforms.lossless_compaction import ( + compact_lossless as _cl, +) +from headroom.transforms.lossless_compaction import ( + path_heading as _ph, +) +from headroom.transforms.lossless_compaction import ( + path_unheading as _puh, +) + + +def test_path_fold_roundtrip_and_shrinks_pure_list(): + c = "./suma/apps/ext/core.py\n./suma/apps/ext/dao.py\n./suma/apps/other/x.py" + folded = _cl(c, "paths") + assert _puh(_ph(c)) == c # exact inverse + assert len(folded) < len(c) # shrinks + assert folded != c + + +def test_path_fold_safe_passthrough_on_non_path_shapes(): + # grep path:line:content is the search fold's job, not paths -> unchanged + assert _cl("a/b.py:12:def f\na/b.py:15:x", "paths") == "a/b.py:12:def f\na/b.py:15:x" + # trailing-slash dir entries and single paths -> unchanged + assert _cl("./a/b/\n./a/c/", "paths") == "./a/b/\n./a/c/" + assert _cl("./only/one.py", "paths") == "./only/one.py" + + +def test_path_fold_mixed_content_roundtrips_or_passes_through(): + # a non-path no-slash line among paths must never corrupt: compact_lossless + # verifies and returns original if the fold isn't exactly reversible. + c = "./a/b/f.py\n./a/b/g.py\nsome log line\n./a/b/h.py" + out = _cl(c, "paths") + assert _puh(_ph(out)) == out or out == c # never corrupts + # simplest invariant: decoding whatever we emit reconstructs the input + assert _puh(_ph(c)) == c or _cl(c, "paths") == c + + +# ---- EXPERIMENT: HEADROOM_EXPERIMENTAL_READ_KEEP_RATIO (light Kompress on reads) ---- +def test_experimental_read_keep_ratio_flag_and_gating(monkeypatch): + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + + # OFF by default -> verbatim (helper returns None, no compression attempted) + monkeypatch.delenv("HEADROOM_EXPERIMENTAL_READ_KEEP_RATIO", raising=False) + r_off = ContentRouter(ContentRouterConfig()) + assert r_off._exp_read_keep_ratio == 0.0 + assert r_off._experimental_compress_read("x" * 500) is None + + # ON -> calls Kompress at the ratio; keeps result only if it actually shrank + monkeypatch.setenv("HEADROOM_EXPERIMENTAL_READ_KEEP_RATIO", "0.9") + r_on = ContentRouter(ContentRouterConfig()) + assert r_on._exp_read_keep_ratio == 0.9 + seen = {} + + def fake_ml(content, context, question=None, target_ratio=None): + seen["ratio"] = target_ratio + return content[: len(content) // 2], 10 # pretend it shrank + + monkeypatch.setattr(r_on, "_try_ml_compressor", fake_ml) + out = r_on._experimental_compress_read("y" * 500, "ctx") + assert out is not None and len(out) < 500 # adopted (shrank) + assert seen["ratio"] == 0.9 # ratio threaded through + + # no-shrink -> None (fall back to verbatim protection) + monkeypatch.setattr( + r_on, "_try_ml_compressor", lambda c, ctx, question=None, target_ratio=None: (c, 1) + ) + assert r_on._experimental_compress_read("z" * 500) is None + # sub-floor content never attempted + assert r_on._experimental_compress_read("short") is None diff --git a/tests/test_cache_mode_delta_marker.py b/tests/test_cache_mode_delta_marker.py new file mode 100644 index 000000000..096cda818 --- /dev/null +++ b/tests/test_cache_mode_delta_marker.py @@ -0,0 +1,185 @@ +"""fix-3: the cache-mode delta path must ignore moved cache_control markers. + +Cache mode replays the exact previously-forwarded bytes for history and +compresses ONLY the newly appended delta (compress-once-then-freeze). The gate +is ``AnthropicHandler._extract_cache_stable_delta``: it only engages when the +prior original request is a message-prefix of the current one. + +Real clients (litellm, Claude Code) move the ephemeral cache_control breakpoint +to the newest message every turn, so a historical message carries the marker on +one turn and not the next. The original raw-dict prefix compare therefore failed +every turn, dropping cache mode to RAW (uncompressed) forwarding -- byte-stable +(0 busts) but 0% compression. Observed directly on the mini-swe-agent cache-mode +run: avg_compression_pct=0.0 on every instance, orig==opt on every turn. + +These tests pin the scenario against the real handler method and prove the +cache_control-agnostic compare lets the delta engage while the replayed prefix +stays byte-identical (so the provider prefix still hits). +""" + +import copy + +from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin + +delta = AnthropicHandlerMixin._extract_cache_stable_delta + + +def B(role, text, cc=False): + """Anthropic block-style message; cache_control lives on a content block.""" + blk = {"type": "text", "text": text} + if cc: + blk["cache_control"] = {"type": "ephemeral"} + return {"role": role, "content": [blk]} + + +# Turn t: client marked the (then-newest) msg2. We forwarded it verbatim. +PREV_ORIG = [B("user", "sys+task"), B("assistant", "ok"), B("user", "obs-1", cc=True)] +PREV_FWD = copy.deepcopy(PREV_ORIG) +# Turn t+1: appended act-2 + obs-2 and MOVED the marker off msg2 onto the newest. +CUR = [ + B("user", "sys+task"), + B("assistant", "ok"), + B("user", "obs-1"), # marker gone + B("assistant", "act-2"), + B("user", "obs-2", cc=True), # marker moved here +] + + +def test_moved_marker_engages_delta_not_raw_fallback(): + out = delta(CUR, PREV_ORIG, PREV_FWD) + assert out is not None, "moved marker must NOT force raw fallback" + stable_prefix, appended = out + # The replayed prefix is byte-identical to what we forwarded (and the + # provider cached) last turn -> the prefix hits instead of busting. + assert stable_prefix == PREV_FWD + # Only the two newly appended messages are handed to compression. + assert len(appended) == 2 + assert appended[0]["content"][0]["text"] == "act-2" + assert appended[1]["content"][0]["text"] == "obs-2" + + +def test_control_marker_not_moved_also_engages(): + # Same append, marker left on the historical msg2: engages either way. + cur = [ + B("user", "sys+task"), + B("assistant", "ok"), + B("user", "obs-1", cc=True), + B("assistant", "act-2"), + B("user", "obs-2"), + ] + assert delta(cur, PREV_ORIG, PREV_FWD) is not None + + +def test_real_content_divergence_still_falls_back(): + # Safety preserved: a genuinely different historical message (not just a + # moved marker) must still bail to raw -- we never replay stale content. + cur = [ + B("user", "sys+task"), + B("assistant", "DIFFERENT"), # content actually changed + B("user", "obs-1"), + B("assistant", "act-2"), + ] + assert delta(cur, PREV_ORIG, PREV_FWD) is None + + +def test_cold_start_returns_none(): + assert delta(CUR, None, None) is None + assert delta(CUR, [], []) is None + + +def test_shorter_current_returns_none(): + assert delta([B("user", "sys+task")], PREV_ORIG, PREV_FWD) is None + + +# ── fix-4: tool_result content shape (string <-> [{type:text}]) ──────────────── +def _tr(tool_use_id, text, as_string): + content = text if as_string else [{"type": "text", "text": text}] + return { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": tool_use_id, "content": content}], + } + + +def test_tool_result_string_vs_block_engages(): + # Stored previous_original holds the block-list form; client resends the SAME + # tool_result as a bare string. These are Anthropic-equivalent -> must engage. + prev_orig = [ + B("user", "task"), + _tr("t1", "0\n\n", as_string=False), + ] + prev_fwd = copy.deepcopy(prev_orig) + cur = [ + B("user", "task"), + _tr("t1", "0\n\n", as_string=True), # string form + B("assistant", "next-action"), + _tr("t2", "done", as_string=True), + ] + out = delta(cur, prev_orig, prev_fwd) + assert out is not None, "tool_result string-vs-block must NOT force raw fallback" + stable_prefix, appended = out + assert stable_prefix == prev_fwd # replay the byte-identical cached prefix + assert len(appended) == 2 # only the new assistant + tool_result are the delta + + +def test_tool_result_different_text_still_falls_back(): + # Safety: same shape-normalization must NOT hide a genuine content change. + prev_orig = [B("user", "task"), _tr("t1", "OUTPUT-A", as_string=False)] + prev_fwd = copy.deepcopy(prev_orig) + cur = [B("user", "task"), _tr("t1", "OUTPUT-B", as_string=True), B("assistant", "x")] + assert delta(cur, prev_orig, prev_fwd) is None + + +def test_tool_use_caller_annotation_ignored(): + # mini-swe-agent/litellm adds a non-semantic `caller` tag to tool_use blocks + # on the stored copy but not on the re-sent wire message -> must still engage. + def _asst(with_caller): + tu = {"type": "tool_use", "id": "tu1", "name": "bash", "input": {"command": "ls"}} + if with_caller: + tu["caller"] = {"type": "direct"} + return {"role": "assistant", "content": [tu]} + + prev_orig = [B("user", "task"), _asst(with_caller=True)] # stored: has caller + prev_fwd = copy.deepcopy(prev_orig) + cur = [ + B("user", "task"), + _asst(with_caller=False), + _tr("t1", "out", as_string=True), + ] # wire: no caller + out = delta(cur, prev_orig, prev_fwd) + assert out is not None, "a client-only `caller` annotation must not force raw fallback" + assert out[0] == prev_fwd + assert len(out[1]) == 1 + + +def test_tool_use_different_input_still_falls_back(): + # Safety: a real change to the tool command must still bail. + def _asst(cmd): + return { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "tu1", "name": "bash", "input": {"command": cmd}} + ], + } + + prev_orig = [B("user", "task"), _asst("ls")] + prev_fwd = copy.deepcopy(prev_orig) + cur = [B("user", "task"), _asst("rm -rf /"), _tr("t1", "out", as_string=True)] + assert delta(cur, prev_orig, prev_fwd) is None + + +def test_combined_marker_move_and_tool_result_shape_engages(): + # The real mini-swe-agent situation: marker moved AND tool_result shape differs. + prev_orig = [ + B("user", "task", cc=True), + _tr("t1", "obs-1", as_string=False), + ] + prev_fwd = copy.deepcopy(prev_orig) + cur = [ + B("user", "task"), # marker gone + _tr("t1", "obs-1", as_string=True), # shape flipped to string + B("assistant", "act", cc=True), # marker moved to newest + ] + out = delta(cur, prev_orig, prev_fwd) + assert out is not None + assert out[0] == prev_fwd + assert len(out[1]) == 1 diff --git a/tests/test_cache_prefix_overlay.py b/tests/test_cache_prefix_overlay.py index e66a5ec62..8cf97df37 100644 --- a/tests/test_cache_prefix_overlay.py +++ b/tests/test_cache_prefix_overlay.py @@ -1,3 +1,4 @@ +# ruff: noqa: E402 — test sections import after helper/setup code by design. """overlay_cached_prefix: freeze must forward the CACHED (compressed) bytes. The freeze path can emit the agent's ORIGINAL bytes for a frozen message, but @@ -77,3 +78,82 @@ def test_cache_hit_property_prefix_matches_last_forward(): out = overlay_cached_prefix(OPTIMIZED_BUGGY, CUR_ORIG, PREV_ORIG, PREV_FWD) n = len(PREV_FWD) assert out[:n] == PREV_FWD # exact byte-identical prefix → provider cache hit + + +# ============================================================================ +# OpenAI function-calling frozen-count: tool_calls must be counted (Kimi bug) +# ============================================================================ +# _estimate_message_tokens only counted `content` + Anthropic content-blocks, +# never OpenAI top-level `tool_calls`. So a function-calling assistant turn +# (content None, command in tool_calls) estimated to ~0, the frozen-prefix +# estimate overshot the real cache boundary, and the NEWEST delta got frozen — +# giving OpenAI/Kimi tool harnesses ~zero compression. These lock in the fix. +import json as _json + +from headroom.cache.prefix_tracker import PrefixCacheTracker, PrefixFreezeConfig + + +def _openai_asst(cmd): + return { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "bash", "arguments": _json.dumps({"command": cmd})}, + } + ], + } + + +def test_estimate_counts_openai_tool_calls(): + est = PrefixCacheTracker._estimate_message_tokens + cmd = "cd /tmp/core && cat suma/apps/underwriting/followup/service.py" + with_calls = est([_openai_asst(cmd)])[0] + # empty content + no tool_calls counted => only the +20 overhead (~5 tok) + bare = est([{"role": "assistant", "content": None}])[0] + assert with_calls > bare + 5, (with_calls, bare) # the command is now counted + # legacy function_call shape too + fc = est( + [ + { + "role": "assistant", + "content": None, + "function_call": {"name": "bash", "arguments": _json.dumps({"command": cmd})}, + } + ] + )[0] + assert fc > bare + 5, (fc, bare) + + +def test_frozen_count_leaves_openai_tool_delta_mutable(): + # A tool-based turn: cached prefix (system+task+prior tool obs) then a NEW + # assistant tool_call + its observation. After update_from_response reports + # the prefix cached, the frozen count must NOT swallow the newest delta. + trk = PrefixCacheTracker("openai", PrefixFreezeConfig(min_cached_tokens=10)) + msgs = [ + {"role": "system", "content": "s" * 400}, + {"role": "user", "content": "task " * 200}, + _openai_asst("cd /tmp/core && rg -n foo ."), + {"role": "tool", "tool_call_id": "c1", "content": "hit\n" * 300}, # cached prefix ends here + _openai_asst("cd /tmp/core && cat foo.py"), # NEW delta (assistant) + { + "role": "tool", + "tool_call_id": "c1", + "content": "code\n" * 400, + }, # NEW delta (observation) + ] + counts = PrefixCacheTracker._estimate_message_tokens(msgs) + # cache_read ~= the first 4 messages' real tokens (prefix cached) + cached_prefix_tokens = sum(counts[:4]) + trk.update_from_response( + cache_read_tokens=cached_prefix_tokens, + cache_write_tokens=0, + messages=msgs, + message_token_counts=counts, + ) + frozen = trk.get_frozen_message_count() + # must freeze ~the cached prefix (<=4), NOT the whole 6 (which would freeze + # the newest observation delta and block all compression). + assert frozen <= 4, f"frozen={frozen} swallowed the delta (len={len(msgs)})" diff --git a/tests/test_observed_wire_shapes.py b/tests/test_observed_wire_shapes.py new file mode 100644 index 000000000..d7670cdc9 --- /dev/null +++ b/tests/test_observed_wire_shapes.py @@ -0,0 +1,588 @@ +# ruff: noqa: E402, E731 — test sections import after setup; lambdas are test stubs. +"""Cache-mode delta engagement against REAL observed wire shapes + extension scaffolding. + +Section 1 pins the exact message shapes we observed on the wire during the SWE-bench +mini-swe-agent + litellm -> Anthropic run (captured via the proxy's per-turn DELTA-DIAG): +these are the shapes that broke the naive prefix compare and that fix-4..7 + the +generalized canonicalizer now handle. They are our validated path and MUST stay green. + +Section 2 is EXTENSION-READY coverage for provider/client shapes we researched (OpenAI +Chat + Responses, Bedrock Converse, Vercel-AI-SDK/opencode) but have NOT yet exercised +end-to-end. The shared canonicalizer is provider-agnostic, so these already pass at the +*comparison* layer; the comments mark what additional *handler* wiring each provider +still needs for full delta-only compression (see the per-provider TODOs). + +Everything here is comparison-layer only (no Modal / no provider calls). +""" + +import copy + +from headroom.cache.prefix_tracker import ( + _canonicalize_for_prefix_compare as CANON, +) +from headroom.cache.prefix_tracker import ( + extract_cache_stable_delta as delta, +) + + +def _eq(a, b): + return CANON(a) == CANON(b) + + +# ============================================================================ +# Section 1 — OBSERVED: mini-swe-agent + litellm -> Anthropic wire (validated) +# ============================================================================ +# Real shapes from the run's DELTA-DIAG. mini emits a bash action; litellm converts +# the OpenAI-ish history to Anthropic blocks on the wire, and (turn-to-turn) it: +# (a) moves the ephemeral cache_control marker to the newest block, +# (b) attaches `caller: {type: direct}` to tool_use on the stored copy, +# (c) flips tool_result.content between a bare string and [{type:text,text}], +# while the observation payload itself (`N\n…`) +# is unchanged. All three must be ignored by the prefix compare. + +_RC = "0\n\n./suma/apps/foo.py\n" # real observation form + + +def _asst_tooluse(with_caller: bool, cc: bool): + tu = { + "type": "tool_use", + "id": "toolu_01ABC", + "name": "bash", + "input": {"command": 'cd /tmp/core && rg -l "safe_math" --type py | head'}, + } + if with_caller: + tu["caller"] = {"type": "direct"} # litellm programmatic-tool tag + if cc: + tu["cache_control"] = {"type": "ephemeral"} + return {"role": "assistant", "content": [tu]} + + +def _tool_result(as_string: bool, cc: bool): + content = _RC if as_string else [{"type": "text", "text": _RC}] + block = {"type": "tool_result", "tool_use_id": "toolu_01ABC", "content": content} + if cc: + block["cache_control"] = {"type": "ephemeral"} + return {"role": "user", "content": [block]} + + +def test_observed_caller_present_vs_absent_ignored(): + assert _eq( + _asst_tooluse(with_caller=True, cc=False), _asst_tooluse(with_caller=False, cc=False) + ) + + +def test_observed_tool_result_string_vs_block_ignored(): + assert _eq(_tool_result(as_string=True, cc=False), _tool_result(as_string=False, cc=False)) + + +def test_observed_moved_cache_control_ignored(): + # marker on tool_use one turn, on tool_result the next + assert _eq(_asst_tooluse(with_caller=True, cc=True), _asst_tooluse(with_caller=True, cc=False)) + + +def test_observed_thinking_block_stable_signature(): + # mini turns carry an Anthropic thinking block with a stable signature; unchanged + # across resend -> stays equal (and a signature change would be a real divergence). + th = lambda: { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "", "signature": "Eo4CCmMIDxgCKkD..."}, + {"type": "text", "text": "Let me find the tool."}, + {"type": "tool_use", "id": "toolu_01ABC", "name": "bash", "input": {"command": "ls"}}, + ], + } + assert _eq(th(), th()) + + +def test_observed_full_turn_delta_engages(): + # The exact failure the run hit: prev stored the assistant with `caller` + + # cache_control on the newest block; this turn re-sends the same assistant WITHOUT + # caller, tool_result as a STRING, and the marker MOVED to the new observation. + # After fix-4..7 + generalized canon, the delta must engage (replay prefix + 1 delta). + prev_orig = [ + {"role": "user", "content": [{"type": "text", "text": "task"}]}, + _asst_tooluse(with_caller=True, cc=True), + ] + prev_fwd = copy.deepcopy(prev_orig) + cur = [ + {"role": "user", "content": [{"type": "text", "text": "task"}]}, + _asst_tooluse(with_caller=False, cc=False), + _tool_result(as_string=True, cc=True), + ] + out = delta(cur, prev_orig, prev_fwd) + assert out is not None, "observed litellm churn must NOT force raw fallback" + stable_prefix, appended = out + assert stable_prefix == prev_fwd # replay the byte-identical cached prefix + assert len(appended) == 1 # only the new tool_result is the delta + + +def test_observed_genuine_command_change_still_diverges(): + # Safety: a real change to the bash command must still fail the compare. + prev = [ + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "t", "name": "bash", "input": {"command": "ls"}} + ], + } + ] + cur = [ + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "t", "name": "bash", "input": {"command": "rm -rf /"}} + ], + } + ] + assert not _eq(prev[0], cur[0]) + + +# ============================================================================ +# Section 2 — EXTENSION-READY: researched shapes not yet exercised end-to-end +# ============================================================================ +# The generalized canonicalizer is provider-agnostic, so these pass at the COMPARISON +# layer today. Each block notes the additional HANDLER wiring still required for full +# delta-only compression on that provider (tracked as follow-ups). + + +# ---- OpenAI Chat Completions ------------------------------------------------ +# Tool result is a `role:"tool"` message with STRING content; assistant tool call is +# `tool_calls[].function{name, arguments(JSON string)}`; automatic prefix caching (NO +# cache_control marker). Noise seen on echoes: system_fingerprint/service_tier, and +# streaming `index` on tool_calls. +# EXTENSION TODO (handler): openai.py cache mode currently does overlay + frozen-count, +# NOT delta-only compression. Wire it to extract_cache_stable_delta (marker policy = none). +def test_ext_openai_tool_calls_index_and_fingerprint_ignored(): + a = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "index": 0, + "function": {"name": "bash", "arguments": '{"command":"ls"}'}, + } + ], + "system_fingerprint": "fp_a", + "service_tier": "default", + } + b = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": '{"command":"ls"}'}, + } + ], + } + assert _eq(a, b) + # but different arguments (opaque JSON string) must diverge + c = copy.deepcopy(b) + c["tool_calls"][0]["function"]["arguments"] = '{"command":"pwd"}' + assert not _eq(b, c) + + +# ---- OpenAI Responses API --------------------------------------------------- +# function_call / function_call_output linked by call_id (not id); reasoning items carry +# a VERBATIM `encrypted_content` that must round-trip. `summary` is display-only. +# EXTENSION TODO (handler): same delta-path wiring as Chat; ensure reasoning items are +# treated as content (encrypted_content kept in the identity — already is, generically). +def test_ext_openai_responses_encrypted_content_is_the_semantic_carrier(): + # The verbatim reasoning token is what matters: a change must diverge (never masked), + # identical must equate. + diff = { + "role": "assistant", + "content": [{"type": "reasoning", "id": "rs_1", "encrypted_content": "ENC_DIFFERENT"}], + } + base = { + "role": "assistant", + "content": [{"type": "reasoning", "id": "rs_1", "encrypted_content": "ENC1"}], + } + same = { + "role": "assistant", + "content": [{"type": "reasoning", "id": "rs_1", "encrypted_content": "ENC1"}], + } + assert not _eq(diff, base) + assert _eq(base, same) + # EXTENSION TODO: `summary` (display-only per OpenAI docs) and the reasoning item + # `id` are NOT yet in _NON_SEMANTIC_KEYS. If a client varies them per turn, the + # compare falls back to raw (safe: 0 compression, no stale replay). Add them to the + # deny-list when the OpenAI Responses delta path is wired and we've confirmed on a + # captured wire trace that they are non-load-bearing. + + +# ---- Bedrock Converse ------------------------------------------------------- +# camelCase; toolUse/toolResult keyed (no `type`); toolResult.content allows {json} +# (structured!) + a `status`; cachePoint is a standalone content block; reasoningContent +# carries a verbatim signature. +# EXTENSION TODO (handler): bedrock.py bypasses compression in cache mode. Wire a +# cachePoint delta path (marker policy = strip/relocate cachePoint) to the shared engine. +def test_ext_bedrock_cachepoint_ignored_json_and_status_kept(): + a = { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "tu1", + "content": [{"json": {"ok": True, "n": 1}}], + "status": "success", + } + }, + {"cachePoint": {"type": "default"}}, + ], + } + b = { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "tu1", + "content": [{"json": {"ok": True, "n": 1}}], + "status": "success", + } + } + ], + } + assert _eq(a, b) # cachePoint block dropped + c = copy.deepcopy(b) + c["content"][0]["toolResult"]["content"][0]["json"]["n"] = 2 + assert not _eq(b, c) # opaque json payload compared verbatim + d = copy.deepcopy(b) + d["content"][0]["toolResult"]["status"] = "error" + assert not _eq(b, d) # status is semantic + + +# ---- Vercel AI SDK / opencode ---------------------------------------------- +# Parts-based; reasoning signature lives in providerMetadata.anthropic.signature; parts +# carry `state`/`providerExecuted`/`step-start` transport. NOTE: the proxy sees the +# PROVIDER wire (post-AI-SDK-serialization), so providerMetadata typically does not reach +# us — but we drop it defensively. EXTENSION TODO: if we ever ingest pre-wire AI-SDK +# messages, ensure the signature is lifted from providerMetadata into the identity. +def test_ext_aisdk_provider_metadata_and_state_ignored(): + a = { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "ok", + "state": "done", + "providerMetadata": {"anthropic": {"x": 1}}, + "providerExecuted": True, + } + ], + } + b = {"role": "assistant", "content": [{"type": "text", "text": "ok"}]} + assert _eq(a, b) + + +# ============================================================================ +# Section 3 — EXTENSION HOOKS (documented, not yet implemented) +# ============================================================================ +# When wiring a new provider to the shared delta engine, add here: +# * a per-provider marker policy test (Anthropic cache_control / Bedrock cachePoint +# stripped from the delta before compression; OpenAI: none); +# * a round-trip test that the forwarded prefix stays byte-identical across a real +# multi-turn fixture for that provider (byte-level; ideally sourced from a captured +# HEADROOM_LOG_MESSAGES trace of the `inspect` non-litellm harness); +# * a tool-shape compression test (OpenAI role:tool with tool safeguards; Bedrock +# toolResult json). These live in the content_router tests once fix-7 is generalized +# beyond the Anthropic tool_result block path. + + +# ============================================================================ +# Section 4 — read protection (HEADROOM_PROTECT_READS): never lossy-compress reads +# ============================================================================ +from headroom.transforms.content_router import _is_read_command as _isread + + +def test_read_command_classifier(): + reads = [ + "cat foo.py", + "cat -n foo.py", + "cd /x && cat a.py", + "cd /x && cat -A a.py | head -60", + "sed -n '1,50p' f.py", + "head -100 f.py", + "tail -20 log", + "nl f.py", + ] + non = [ + "cat > f.py <<'EOF'\nx\nEOF", + "cat a >> b", + "echo x | tee f", + "sed -i 's/a/b/' f", + "sed 's/a/b/' f", + "rg -l x --type py", + "grep -rn x .", + "ls -la", + "python -c 'x'", + "git diff -- f", + "swebench-pytest-lite t/", + "", + None, + ] + assert all(_isread(c) for c in reads), [c for c in reads if not _isread(c)] + assert not any(_isread(c) for c in non), [c for c in non if _isread(c)] + + +# ============================================================================ +# Section 5 — command classification is harness-agnostic (Bug A + Bug B) +# ============================================================================ +# Two bugs that silently disabled compression/protection on real harnesses. +# These lock in the fixes and assert they hold across the command-prefix and +# tool-call wire shapes different harnesses/providers emit. +from headroom.transforms.content_router import ( + _bash_command_is_search as _issearch, +) +from headroom.transforms.content_router import ( + _is_read_command as _isread2, +) +from headroom.transforms.content_router import ( + _strip_cd_prefix as _stripcd, +) +from headroom.transforms.content_router import ( + _tool_call_command_text as _cmdtext, +) + +_SEARCH = frozenset({"grep", "rg", "ag", "fgrep", "egrep", "ripgrep"}) + + +def test_bugA_cd_prefixed_search_detected_all_harnesses(): + # Harnesses run every command inside the checkout: `cd && ` + # (mini-swe-agent, most) or `cd ; ` (some Codex configs). Before + # the fix, _bash_program read the program as `cd` -> search fold never fired. + for cmd in [ + "cd /tmp/core && rg -l safe_math --type py", + "cd /tmp/core && grep -rn foo suma/", + "cd /repo; grep -n bar .", # semicolon connector + "cd /a && cd b && rg pat", # chained cds + "grep -rn x .", # no prefix (regression) + "rg pattern src/", + ]: + assert _issearch(cmd, _SEARCH), f"search not detected: {cmd!r}" + # non-search must stay non-search even with a cd prefix + for cmd in ["cd /x && cat a.py", "cd /x && python -c 'x'", "cd /x && ls -la"]: + assert not _issearch(cmd, _SEARCH), f"false search: {cmd!r}" + + +def test_bugA_strip_cd_prefix_shapes(): + assert _stripcd("cd /tmp/core && rg x") == "rg x" + assert _stripcd("cd /repo; grep x") == "grep x" + assert _stripcd("cd a && cd b && grep x") == "grep x" + assert _stripcd("grep x .") == "grep x ." # nothing to strip + assert _stripcd("") == "" and _stripcd(None) == "" # defensive + + +def test_openai_tool_calls_none_does_not_crash_and_still_compresses(): + # OpenAI/LiteLLM assistant messages carry an explicit `tool_calls: None` (and + # `function_call: None`) when there are no calls. `msg.get("tool_calls", [])` + # returns None (not []), so iterating it crashed _build_tool_name_map -> + # apply() -> compression silently fell through to PASSTHROUGH on every OpenAI + # turn (observed on GPT-5.4 text-based: only ~2/24 requests compressed, net + # token inflation). This asserts the coalesce fix: no crash, map builds. + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + + cr = ContentRouter(ContentRouterConfig()) + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + { + "role": "assistant", + "content": "THOUGHT: look\n```bash\ncd /r && cat x.py\n```", + "tool_calls": None, + "function_call": None, + }, # <- the OpenAI shape + { + "role": "user", + "content": "0\n\n" + "x\n" * 200 + "", + }, + ] + name_map = cr._build_tool_name_map(msgs) # must not raise + assert isinstance(name_map, dict) + + +def test_text_based_read_protection_shape_agnostic(monkeypatch=None): + # Text-based agents (GPT-5.4/Codex/Cursor) have NO tool_use/tool_result blocks: + # the command is a fenced block in the assistant STRING, the observation is a + # plain user string. Read-protection must still fire off the *preceding + # command* so cat/sed code reads are passed verbatim on ANY model/harness. + import os + + from headroom.tokenizers.registry import get_tokenizer + from headroom.transforms.content_router import ( + ContentRouter, + ContentRouterConfig, + _fenced_shell_command, + ) + from headroom.transforms.read_lifecycle import ReadLifecycleConfig + + assert ( + _fenced_shell_command("T\n```mswea_bash_command\ncd /r && cat x.py\n```") + == "cd /r && cat x.py" + ) + assert _fenced_shell_command("no fence here") == "" + os.environ["HEADROOM_PROTECT_READS"] = "1" + tok = get_tokenizer("gpt-4o") + big_code = "def f():\n" + " x = 1\n" * 300 + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + { + "role": "assistant", + "content": "T\n```mswea_bash_command\ncd /r && cat a.py\n```", + "tool_calls": None, + }, # READ command + { + "role": "user", + "content": "0\n\n" + big_code + "", + }, + { + "role": "assistant", + "content": "T\n```mswea_bash_command\ncd /r && grep -rn foo .\n```", + "tool_calls": None, + }, # SEARCH command + { + "role": "user", + "content": "0\n\n" + + ("a.py:1:foo\n" * 300) + + "", + }, + ] + r = ContentRouter( + ContentRouterConfig( + skip_user_messages=False, read_lifecycle=ReadLifecycleConfig(enabled=False) + ) + ) + r.apply( + [dict(m) for m in msgs], + tok, + frozen_message_count=0, + context="", + compress_user_messages=True, + protect_recent=0, + min_tokens_to_compress=25, + ) + # the observation AFTER the cat (index 3) must be read-protected; the grep one (5) must not + assert 3 in r._protect_read_msg_indices, r._protect_read_msg_indices + assert 5 not in r._protect_read_msg_indices, r._protect_read_msg_indices + + +def test_bugB_read_detection_across_tool_call_wire_shapes(): + # The SAME read action, as each provider/harness serializes its tool call. + # _tool_call_command_text must recover the shell command from all of them so + # read-protection fires regardless of client. (Bug B: the old path fed the + # raw OpenAI JSON blob to _is_read_command, which always returned False.) + import json + + anthropic_input = {"command": "cd /tmp/core && cat suma/x.py"} # Anthropic: dict + openai_args = json.dumps({"command": "cd /tmp/core && cat suma/x.py"}) # OpenAI: JSON string + codex_list = {"command": ["cat", "suma/x.py"]} # Codex: argv list + for raw in (anthropic_input, openai_args, codex_list): + assert _isread2(_cmdtext(raw)), f"read not detected from {raw!r}" + # a search command from any shape must NOT be read-protected (stays compressible) + assert not _isread2(_cmdtext({"command": "cd /x && rg pat"})) + assert not _isread2(_cmdtext(json.dumps({"command": "grep -rn x ."}))) + + +# ============================================================================ +# Section 6 — JSON-OBJECT reads are releasable (detector now parses, not [-only) +# ============================================================================ +# `_try_detect_json` used to recognize only JSON *arrays* ([...]); a JSON +# *object* ({...}) — celery.json / package.json / most config+data — fell +# through to PLAIN_TEXT and got read-PROTECTED (never compressed). The detector +# now decides JSON by PARSING (objects, arrays, and a JSON value inside a small +# bounded wrapper), so these lock in that an object read is released for +# compression across wrapped/unwrapped shapes while source code stays protected. +from headroom.transforms.content_router import ( + _read_output_should_be_protected as _protect, +) + + +def test_json_object_read_is_releasable_all_shapes(): + big_obj = ( + "{\n" + + ",\n".join(f' "suma.apps.task_{i}": {{"queue": "q", "rate": {i}}}' for i in range(40)) + + "\n}" + ) + wrapped = "0\n\n" + big_obj + "\n" + # object — raw and harness-wrapped — is RELEASED (not protected) for compression + assert _protect(big_obj) is False + assert _protect(wrapped) is False + # a JSON array is likewise releasable + assert _protect('[{"a": 1}, {"a": 2}]') is False + # genuine source code (even with a dict literal) stays PROTECTED + assert _protect("def f():\n return {1: 2}\n" * 20) is True + + +def test_read_protection_releases_json_object_but_protects_code(): + big_obj = "{\n" + ",\n".join(f' "k{i}": {{"v": {i}}}' for i in range(60)) + "\n}" + wrapped_obj = "0\n\n" + big_obj + "\n" + py = "0\n\n" + ("def f():\n x = 1\n" * 60) + "" + assert _protect(wrapped_obj) is False # config/data object → RELEASE (compressible) + assert _protect(py) is True # source code → PROTECT (byte-exact) + + +def test_read_protection_role_agnostic_openai_role_tool(): + # Kimi / fireworks (OpenAI function-calling): the read command is in the + # assistant tool_calls and the observation is a `role:tool` STRING message. + # Read-protection was gated on role=='user' (+ tool_result blocks), so these + # role:tool reads slipped through UNPROTECTED. This locks in the role-agnostic + # fix: a read observation is protected by OUTCOME (read command -> code), + # whatever role the harness stamps on it. + import json as _json + import os + + from headroom.tokenizers.registry import get_tokenizer + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + from headroom.transforms.read_lifecycle import ReadLifecycleConfig + + os.environ["HEADROOM_PROTECT_READS"] = "1" + os.environ.pop("HEADROOM_EXPERIMENTAL_READ_KEEP_RATIO", None) # protection, not the experiment + tok = get_tokenizer("gpt-4o") + code = "def f():\n" + " x = 1\n" * 300 + + def tc(cid, cmd): + return { + "id": cid, + "type": "function", + "function": {"name": "bash", "arguments": _json.dumps({"command": cmd})}, + } + + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + {"role": "assistant", "content": "", "tool_calls": [tc("c_read", "cd /r && cat a.py")]}, + {"role": "tool", "tool_call_id": "c_read", "content": code}, # READ (role:tool) -> protect + { + "role": "assistant", + "content": "", + "tool_calls": [tc("c_grep", "cd /r && grep -rn foo .")], + }, + { + "role": "tool", + "tool_call_id": "c_grep", + "content": "a.py:1:foo\n" * 300, + }, # SEARCH -> not protected + ] + r = ContentRouter( + ContentRouterConfig( + skip_user_messages=False, read_lifecycle=ReadLifecycleConfig(enabled=False) + ) + ) + out = r.apply( + [dict(m) for m in msgs], + tok, + frozen_message_count=0, + context="", + compress_user_messages=True, + protect_recent=0, + min_tokens_to_compress=25, + ) + assert "c_read" in r._protect_read_tool_ids, ( + "read cmd must be identified from OpenAI tool_calls" + ) + # the role:tool READ observation is protected verbatim (was the bug: unprotected) + assert out.messages[3]["content"] == code, "role:tool code read must be protected verbatim" diff --git a/tests/test_proxy/test_anthropic_ccr_deferred_injection.py b/tests/test_proxy/test_anthropic_ccr_deferred_injection.py index fce7dade2..da6311043 100644 --- a/tests/test_proxy/test_anthropic_ccr_deferred_injection.py +++ b/tests/test_proxy/test_anthropic_ccr_deferred_injection.py @@ -1042,8 +1042,16 @@ def test_cache_mode_existing_retrieve_tool_compresses_only_the_unfrozen_delta( def _fake_apply(**kwargs): captured.setdefault("compression_calls", []).append(kwargs["messages"]) + captured["frozen_message_count"] = kwargs.get("frozen_message_count") + # fix-6 contract: the compressor is handed the frozen forwarded + # prefix + the delta and only compresses indices >= + # frozen_message_count (so the delta's tool_name resolves from the + # prefix). Mirror it: pass the frozen prefix through, compress the tail. + fz = kwargs.get("frozen_message_count") or 0 + msgs = kwargs["messages"] return SimpleNamespace( - messages=[ + messages=list(msgs[:fz]) + + [ { "role": "user", "content": ( @@ -1094,7 +1102,14 @@ def test_cache_mode_existing_retrieve_tool_compresses_only_the_unfrozen_delta( assert response.status_code == 200 assert len(captured.get("compression_calls", [])) == 1 - assert captured["compression_calls"][0] == [original_messages[1]] + # fix-6 contract: the compressor receives the frozen forwarded prefix + # (the previously-forwarded compressed message) + the raw delta, with + # frozen_message_count = prefix length so ONLY the delta is compressed. + assert captured["compression_calls"][0] == [ + previous_forwarded_messages[0], + original_messages[1], + ] + assert captured["frozen_message_count"] == 1 forwarded = captured["body"] assert forwarded["messages"] == [ previous_forwarded_messages[0], diff --git a/tests/test_proxy_anthropic_cache_stability.py b/tests/test_proxy_anthropic_cache_stability.py index 086ba83ad..2e00d3f96 100644 --- a/tests/test_proxy_anthropic_cache_stability.py +++ b/tests/test_proxy_anthropic_cache_stability.py @@ -1046,8 +1046,16 @@ def test_cache_mode_reuses_prior_forwarded_prefix_and_compresses_only_new_suffix def _fake_apply(**kwargs): captured["calls"].append(kwargs["messages"]) + captured["frozen_message_count"] = kwargs.get("frozen_message_count") + # fix-6 contract: the compressor is handed the frozen forwarded prefix + # + the new delta and only compresses indices >= frozen_message_count + # (so a lone tool_result can resolve its tool_name from the prefix). + # Mirror the real router: pass the frozen prefix through verbatim and + # compress only the tail — the handler splices result.messages[prefix_n:]. + fz = kwargs.get("frozen_message_count") or 0 + msgs = kwargs["messages"] return SimpleNamespace( - messages=[{"role": "user", "content": "COMPRESSED_TURN3"}], + messages=list(msgs[:fz]) + [{"role": "user", "content": "COMPRESSED_TURN3"}], transforms_applied=["fake:delta"], timing={}, tokens_before=40, @@ -1094,7 +1102,22 @@ def test_cache_mode_reuses_prior_forwarded_prefix_and_compresses_only_new_suffix ) assert response.status_code == 200 - assert captured["calls"] == [[{"role": "user", "content": "turn3"}]] + # fix-6 contract: the compressor receives the frozen FORWARDED prefix + # (with COMPRESSED_TURN2, the byte-stable cached form) + the raw new + # delta (turn3), so tool_name resolution / dedup stay consistent with + # what is actually cached. frozen_message_count = prefix length pins + # compression to the delta ONLY — the prefix is never re-compressed. + assert captured["calls"] == [ + [ + {"role": "user", "content": "turn1"}, + {"role": "assistant", "content": "turn1-assistant"}, + {"role": "user", "content": "COMPRESSED_TURN2"}, + {"role": "assistant", "content": "turn2-assistant"}, + {"role": "user", "content": "turn3"}, + ] + ] + assert captured["frozen_message_count"] == 4 # only the delta (turn3) is compressed + # Forwarded body = byte-identical cached prefix + the compressed delta. assert captured["body"]["messages"] == [ {"role": "user", "content": "turn1"}, {"role": "assistant", "content": "turn1-assistant"}, diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index 1efe1a8fc..e55662c75 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -955,3 +955,33 @@ class TestHasNewCcrMarkers: ) is False ) + + +def test_strict_frozen_count_tool_and_function_tail_are_mutable(): + # OpenAI function-calling harnesses (Kimi / fireworks) end each turn with a + # role:"tool" (or legacy role:"function") observation — NOT role:"user". + # Gating the mutable tail on role=="user" froze the whole conversation on + # every such turn => zero compression. Tool/function observations must be + # treated as the mutable delta (freeze all-but-last), like a user obs. + from headroom.proxy.handlers.openai import OpenAIHandlerMixin as M + + # role:tool tail -> only the last message is mutable (frozen = final_idx) + assert ( + M._strict_previous_turn_frozen_count( + [{"role": "user"}, {"role": "assistant"}, {"role": "tool"}], 0 + ) + == 2 + ) + assert ( + M._strict_previous_turn_frozen_count( + [{"role": "user"}, {"role": "assistant"}, {"role": "function"}], 0 + ) + == 2 + ) + # assistant/system tail is NOT an observation -> freeze everything + assert ( + M._strict_previous_turn_frozen_count( + [{"role": "user"}, {"role": "tool"}, {"role": "assistant"}], 0 + ) + == 3 + ) diff --git a/tests/test_semantic_canonicalize.py b/tests/test_semantic_canonicalize.py new file mode 100644 index 000000000..92a7006b6 --- /dev/null +++ b/tests/test_semantic_canonicalize.py @@ -0,0 +1,270 @@ +"""Generalized cross-turn prefix canonicalizer (`_canonicalize_for_prefix_compare`). + +The delta path decides "is this turn an append-only extension of the last?" by +comparing the canonicalized prefix. Clients attach non-semantic annotations that +vary turn-to-turn (cache_control moved to the newest block, litellm `caller`, +provider_specific_fields, AI-SDK providerMetadata, streaming `index`, string vs +block content). The canonicalizer must ignore all of those, while NEVER dropping a +semantic field (which would mask a real divergence -> stale replay). + +Two messages canonicalize-equal IFF they are semantically identical. These tests +pin: (1) each noise field is ignored, across Anthropic/OpenAI/Bedrock shapes; +(2) semantic differences are still detected; (3) reasoning signatures are kept; +(4) opaque tool payloads (input/arguments/json) are compared verbatim so user data +containing keys like `state`/`index` is never corrupted. +""" + +from headroom.cache.prefix_tracker import _canonicalize_for_prefix_compare as C + + +def eq(a, b): + return C(a) == C(b) + + +# ── noise is ignored (equal despite it) ─────────────────────────────────────── +def test_cache_control_ignored_anthropic(): + a = { + "role": "user", + "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}], + } + b = {"role": "user", "content": [{"type": "text", "text": "hi"}]} + assert eq(a, b) + + +def test_cachepoint_and_caller_ignored(): + a = { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "t1", + "name": "bash", + "input": {"cmd": "ls"}, + "caller": {"type": "direct"}, + }, + {"cachePoint": {"type": "default"}}, + ], + } + b = { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "t1", "name": "bash", "input": {"cmd": "ls"}}, + {"cachePoint": {"type": "default"}}, + ], + } + assert eq(a, b) + + +def test_litellm_and_aisdk_noise_ignored(): + a = { + "role": "assistant", + "content": "ok", + "provider_specific_fields": {"x": 1}, + "reasoning_content": "...", + "annotations": [{"u": "url"}], + "system_fingerprint": "fp_1", + "service_tier": "default", + } + b = { + "role": "assistant", + "content": "ok", + "provider_specific_fields": {"x": 999}, + "system_fingerprint": "fp_2", + } + assert eq(a, b) + + +def test_streaming_index_and_state_ignored_at_block_level(): + a = { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "t1", + "name": "b", + "input": {"c": 1}, + "index": 2, + "state": "output-available", + "providerMetadata": {"a": 1}, + } + ], + } + b = { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "b", "input": {"c": 1}}], + } + assert eq(a, b) + + +def test_string_content_normalized_to_block(): + assert eq( + {"role": "user", "content": "hello"}, + {"role": "user", "content": [{"type": "text", "text": "hello"}]}, + ) + + +def test_tool_result_string_vs_block_equal(): + a = { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "out"}], + } + b = { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t1", + "content": [{"type": "text", "text": "out"}], + } + ], + } + assert eq(a, b) + + +# ── semantic differences ARE detected (not masked) ──────────────────────────── +def test_different_text_detected(): + assert not eq({"role": "user", "content": "A"}, {"role": "user", "content": "B"}) + + +def test_different_tool_input_detected(): + a = { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "b", "input": {"cmd": "ls"}}], + } + b = { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "b", "input": {"cmd": "rm -rf /"}}], + } + assert not eq(a, b) + + +def test_different_role_detected(): + assert not eq({"role": "user", "content": "x"}, {"role": "assistant", "content": "x"}) + + +def test_reasoning_signature_preserved_and_compared(): + # Same thinking text, DIFFERENT signature -> genuinely different (must not equate). + a = { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "", "signature": "SIG_A"}], + } + b = { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "", "signature": "SIG_B"}], + } + assert not eq(a, b) + # Same signature but cache_control noise differs -> equal. + c = { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "", + "signature": "SIG_A", + "cache_control": {"type": "ephemeral"}, + } + ], + } + assert eq(a, c) + + +def test_thinking_present_absent_flip_detected(): + # The litellm/opencode persistence bug: a thinking block dropped on a later turn + # is a REAL divergence and must fail the compare (raw fallback, never stale replay). + a = { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "", "signature": "S"}, + {"type": "text", "text": "ok"}, + ], + } + b = {"role": "assistant", "content": [{"type": "text", "text": "ok"}]} + assert not eq(a, b) + + +# ── the opaque-payload safety trap: noise-named keys inside user data ────────── +def test_opaque_input_with_colliding_keys_not_corrupted(): + # `state`/`index` are noise keys at BLOCK level, but here they are legitimate + # tool-input DATA. They must be compared verbatim, so different inputs differ. + a = { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "t1", "name": "set", "input": {"state": "CA", "index": 3}} + ], + } + b = { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "t1", "name": "set", "input": {"state": "NY", "index": 3}} + ], + } + assert not eq(a, b), "tool input with keys named like noise must NOT be stripped/equated" + + +def test_opaque_arguments_string_verbatim(): + a = { + "role": "assistant", + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "f", "arguments": '{"index": 1}'}} + ], + "content": None, + } + b = { + "role": "assistant", + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "f", "arguments": '{"index": 2}'}} + ], + "content": None, + } + assert not eq(a, b) + + +def test_bedrock_toolresult_json_payload_verbatim(): + a = { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "content": [{"json": {"state": "ok", "n": 1}}], + "status": "success", + } + } + ], + } + b = { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "content": [{"json": {"state": "ok", "n": 2}}], + "status": "success", + } + } + ], + } + assert not eq(a, b) + + +def test_bedrock_cachepoint_and_reasoning_signature(): + # cachePoint (noise) ignored; reasoningText.signature (semantic) compared. + a = { + "role": "assistant", + "content": [ + {"reasoningContent": {"reasoningText": {"text": "r", "signature": "BSIG"}}}, + {"cachePoint": {"type": "default"}}, + ], + } + b = { + "role": "assistant", + "content": [{"reasoningContent": {"reasoningText": {"text": "r", "signature": "BSIG"}}}], + } + assert eq(a, b) + c = { + "role": "assistant", + "content": [ + {"reasoningContent": {"reasoningText": {"text": "r", "signature": "DIFFERENT"}}} + ], + } + assert not eq(a, c) diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index 61b61ea61..371a084f6 100644 --- a/tests/test_tokenizers.py +++ b/tests/test_tokenizers.py @@ -299,6 +299,27 @@ class TestTokenizerRegistry: tokenizer = get_tokenizer("unknown-model-xyz") assert isinstance(tokenizer, EstimatingTokenCounter) + def test_get_kimi_moonshot_calibrated_estimator(self): + """Kimi/Moonshot resolves to the calibrated (3.1 chars/tok) estimator + across every serving form — Fireworks body, litellm slug, native — so + the size-gates aren't starved by the ~20% under-count of the default + adaptive estimator (measured on a SWE-bench Kimi-K2.7-code run).""" + for m in ( + "accounts/fireworks/models/kimi-k2p7-code", # Fireworks body model + "fireworks_ai/kimi-k2p7-code-high", # litellm slug + "moonshotai/Kimi-K2-Instruct", # native + "KIMI-K2P7-CODE", # case-insensitive + ): + tk = get_tokenizer(m) + assert isinstance(tk, EstimatingTokenCounter), m + assert tk._fixed_ratio == 3.1, f"{m}: expected 3.1, got {tk._fixed_ratio}" + # calibrated estimate must beat the default adaptive on Kimi-like code + # (which the default under-counts): denser ratio -> more tokens. + code = 'def f(x):\n return {"a": 1, "b": [2, 3]}\n' * 200 + kimi = get_tokenizer("fireworks_ai/kimi-k2p7-code-high").count_text(code) + default = get_tokenizer("unknown-model-xyz").count_text(code) + assert kimi > default, (kimi, default) + def test_get_with_specific_backend(self): """Test forcing specific backend.""" tokenizer = get_tokenizer("any-model", backend="estimation") diff --git a/tests/test_transforms_content_detection.py b/tests/test_transforms_content_detection.py index d209ecabc..1201c2a60 100644 --- a/tests/test_transforms_content_detection.py +++ b/tests/test_transforms_content_detection.py @@ -56,7 +56,12 @@ def test_json_detection_distinguishes_dict_arrays_and_other_lists() -> None: assert empty_result is not None assert empty_result.metadata == {"item_count": 0, "is_dict_array": False} - assert _try_detect_json('{"id": 1}') is None + # JSON OBJECTS are recognized too (config/data files are ``{…}``, not arrays). + object_result = _try_detect_json('{"id": 1}') + assert object_result is not None + assert object_result.content_type is ContentType.JSON_ARRAY + assert object_result.metadata == {"is_dict_array": False, "is_object": True} + assert _try_detect_json("[not valid json") is None assert is_json_array_of_dicts('[{"id": 1}]') is True assert is_json_array_of_dicts('["value"]') is False @@ -82,12 +87,14 @@ def test_space_separated_json_objects_detected_as_array() -> None: assert _try_detect_json(newline_sep).content_type is ContentType.JSON_ARRAY -def test_space_separated_json_detection_is_conservative() -> None: - # A single object is not an array — must not be claimed. - assert _try_detect_json('{"id": 1}') is None - # Objects interleaved with prose are not clean concatenated JSON. +def test_json_detection_is_liberal_but_bulk_gated() -> None: + # Liberal (parse-based): a lone JSON object IS structured data worth routing — + # config/data files are ``{...}`` (chosen over the earlier conservative stance + # when this PR merged with the concatenated-JSON detector, #1742). + assert _try_detect_json('{"id": 1}').content_type is ContentType.JSON_ARRAY + # But a JSON fragment that is only a minority of the content (prose or a loose + # scalar around it) is NOT claimed — the decoded value must be the bulk. assert _try_detect_json('{"id": 1} then some prose {"id": 2}') is None - # Scalars/strings between objects disqualify the run of dicts. assert _try_detect_json('{"id": 1} "loose string"') is None