feat(cache): provider-agnostic cache-mode delta + cc-agnostic prefix comparison (#1868)

## Description

<!-- Briefly explain the change and why it is needed. -->

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

<!-- Check what you actually ran, then paste the real command output
below. -->

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

<!-- Mention any N/A checklist items, tradeoffs, follow-ups, or
maintainer context. -->
This commit is contained in:
Tejas Chopra 2026-07-08 16:29:35 -04:00 committed by GitHub
parent 38074888ac
commit 7c2f0ea079
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 2203 additions and 84 deletions

View file

@ -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

View file

@ -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)))

View file

@ -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
)

View file

@ -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)

View file

@ -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,

View file

@ -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

View file

@ -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()

View file

@ -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:

View file

@ -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 <dir> &&|; `` chains from a shell command.
Agent harnesses (mini-swe-agent, Codex, Cursor, ) prefix nearly every command
with ``cd <repo> && `` (or ``cd <repo>; ``) 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 <<EOF``), and a bare ``sed`` (without ``-n``)
is a stream edit neither is a read.
"""
if not command or not isinstance(command, str):
return False
# strip leading `cd <dir> && ` 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 <dir> && ` 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)

View file

@ -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<dir>(?:\.{0,2}/)?(?:[^/\s:]+/)+)(?P<base>[^/\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.

View file

@ -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

View file

@ -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

View file

@ -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", "<returncode>0</returncode>\n<output>\n</output>", as_string=False),
]
prev_fwd = copy.deepcopy(prev_orig)
cur = [
B("user", "task"),
_tr("t1", "<returncode>0</returncode>\n<output>\n</output>", as_string=True), # string form
B("assistant", "next-action"),
_tr("t2", "<output>done</output>", 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

View file

@ -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)})"

View file

@ -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 (`<returncode>N</returncode>\n<output>…</output>`)
# is unchanged. All three must be ignored by the prefix compare.
_RC = "<returncode>0</returncode>\n<output>\n./suma/apps/foo.py\n</output>" # 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 <repo> && <tool>`
# (mini-swe-agent, most) or `cd <repo>; <tool>` (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": "<returncode>0</returncode>\n<output>\n" + "x\n" * 200 + "</output>",
},
]
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": "<returncode>0</returncode>\n<output>\n" + big_code + "</output>",
},
{
"role": "assistant",
"content": "T\n```mswea_bash_command\ncd /r && grep -rn foo .\n```",
"tool_calls": None,
}, # SEARCH command
{
"role": "user",
"content": "<returncode>0</returncode>\n<output>\n"
+ ("a.py:1:foo\n" * 300)
+ "</output>",
},
]
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 = "<returncode>0</returncode>\n<output>\n" + big_obj + "\n</output>"
# 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 = "<returncode>0</returncode>\n<output>\n" + big_obj + "\n</output>"
py = "<returncode>0</returncode>\n<output>\n" + ("def f():\n x = 1\n" * 60) + "</output>"
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"

View file

@ -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],

View file

@ -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"},

View file

@ -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
)

View file

@ -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)

View file

@ -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")

View file

@ -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