mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
38074888ac
commit
7c2f0ea079
21 changed files with 2203 additions and 84 deletions
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
161
headroom/cache/prefix_tracker.py
vendored
161
headroom/cache/prefix_tracker.py
vendored
|
|
@ -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)))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
185
tests/test_cache_mode_delta_marker.py
Normal file
185
tests/test_cache_mode_delta_marker.py
Normal 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
|
||||
|
|
@ -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)})"
|
||||
|
|
|
|||
588
tests/test_observed_wire_shapes.py
Normal file
588
tests/test_observed_wire_shapes.py
Normal 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"
|
||||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
270
tests/test_semantic_canonicalize.py
Normal file
270
tests/test_semantic_canonicalize.py
Normal 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)
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue