Merge feat/token-headroom-mode: dual-mode optimization + clean stats

Token headroom mode (HEADROOM_MODE=token_headroom) compresses older
messages to extend session length, trading prefix cache cost savings
for token reduction. Content-addressed CompressionCache avoids
re-compression across turns. Works for both Anthropic and OpenAI.

Also: claude-opus-4-6 model entry, clean session summary in /stats
and MCP headroom_stats tool.
This commit is contained in:
chopratejas 2026-03-19 12:56:39 -07:00
commit 4bda5a31be
7 changed files with 1171 additions and 43 deletions

View file

@ -36,6 +36,7 @@ from .base import (
CacheStrategy,
OptimizationContext,
)
from .compression_cache import CompressionCache
from .dynamic_detector import (
DetectorConfig,
DynamicCategory,
@ -74,6 +75,8 @@ __all__ = [
# Semantic caching
"SemanticCacheLayer",
"SemanticCache",
# Compression cache (token headroom mode)
"CompressionCache",
# Prefix cache tracking
"PrefixCacheTracker",
"PrefixFreezeConfig",

206
headroom/cache/compression_cache.py vendored Normal file
View file

@ -0,0 +1,206 @@
"""Content-addressed compression cache with LRU eviction.
Used in "token headroom mode" to avoid re-compressing messages across turns.
Maps original content hashes to their compressed versions.
"""
from __future__ import annotations
import copy
import hashlib
import json
import logging
from collections import OrderedDict
from dataclasses import dataclass
logger = logging.getLogger(__name__)
@dataclass
class _CacheEntry:
"""Internal cache entry storing compressed text and metadata."""
compressed: str
tokens_saved: int
def _is_tool_result_message(msg: dict) -> bool:
"""Check if a message is a tool result in Anthropic or OpenAI format."""
# OpenAI format: role="tool"
if msg.get("role") == "tool":
return True
# Anthropic format: role="user" with content list containing tool_result blocks
content = msg.get("content")
if isinstance(content, list):
return any(
isinstance(block, dict) and block.get("type") == "tool_result" for block in content
)
return False
def _extract_tool_result_content(msg: dict) -> str | None:
"""Extract text content from a tool result message (both formats)."""
# OpenAI format
if msg.get("role") == "tool":
content = msg.get("content")
return content if isinstance(content, str) else None
# Anthropic format
content = msg.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
inner = block.get("content")
if isinstance(inner, str):
return inner
return None
def _swap_tool_result_content(msg: dict, new_content: str) -> dict:
"""Deep copy msg and replace tool result content with new_content."""
new_msg = copy.deepcopy(msg)
# OpenAI format
if new_msg.get("role") == "tool":
new_msg["content"] = new_content
return new_msg
# Anthropic format
content = new_msg.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
block["content"] = new_content
break
return new_msg
class CompressionCache:
"""Content-addressed cache mapping content hashes to compressed versions.
Uses an OrderedDict for O(1) LRU eviction. Entries are evicted oldest-first
when the cache exceeds max_entries.
"""
def __init__(self, max_entries: int = 10000) -> None:
self.max_entries = max_entries
self._cache: OrderedDict[str, _CacheEntry] = OrderedDict()
self._hits: int = 0
self._misses: int = 0
self._total_tokens_saved: int = 0
def get_compressed(self, hash: str) -> str | None:
"""Retrieve compressed content by hash, refreshing LRU position on hit."""
entry = self._cache.get(hash)
if entry is None:
self._misses += 1
return None
self._hits += 1
self._cache.move_to_end(hash)
return entry.compressed
def store_compressed(self, hash: str, compressed: str, tokens_saved: int) -> None:
"""Store a compressed version keyed by content hash.
If the hash already exists, the entry is overwritten and moved to the
end (most recently used). When the cache exceeds max_entries, the oldest
entry is evicted.
"""
if hash in self._cache:
old_entry = self._cache[hash]
self._total_tokens_saved -= old_entry.tokens_saved
del self._cache[hash]
self._cache[hash] = _CacheEntry(compressed=compressed, tokens_saved=tokens_saved)
self._total_tokens_saved += tokens_saved
while len(self._cache) > self.max_entries:
_, evicted = self._cache.popitem(last=False)
self._total_tokens_saved -= evicted.tokens_saved
def get_stats(self) -> dict:
"""Return cache statistics."""
return {
"entries": len(self._cache),
"hits": self._hits,
"misses": self._misses,
"tokens_saved": self._total_tokens_saved,
}
@staticmethod
def content_hash(content: str | list) -> str:
"""Compute a truncated SHA-256 hash for string or list content.
For list content (Anthropic-format messages with type/text/content fields),
the list is JSON-serialized with sorted keys for deterministic hashing.
"""
if isinstance(content, list):
raw = json.dumps(content, sort_keys=True, ensure_ascii=False)
else:
raw = content
return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
def compute_frozen_count(self, messages: list[dict]) -> int:
"""Count consecutive stable messages from the start.
A message is stable if it is a plain user/assistant/system message,
an assistant message with tool_use blocks, or a tool_result whose
content hash is already in the cache. The first unstable tool_result
(cache miss) stops the count.
"""
count = 0
for msg in messages:
if _is_tool_result_message(msg):
content = _extract_tool_result_content(msg)
if content is not None:
h = self.content_hash(content)
if h not in self._cache:
break
else:
# tool_result with non-string content; treat as unstable
break
# Regular user/assistant/system messages and assistant+tool_use
# are always stable — fall through.
count += 1
return count
def apply_cached(self, messages: list[dict]) -> list[dict]:
"""Return a new list with cached compressions swapped into tool results.
Never mutates the input list or any message dict within it.
Output always has the same length as input.
"""
result: list[dict] = []
for msg in messages:
if _is_tool_result_message(msg):
content = _extract_tool_result_content(msg)
if content is not None:
h = self.content_hash(content)
compressed = self.get_compressed(h)
if compressed is not None:
result.append(_swap_tool_result_content(msg, compressed))
continue
result.append(msg)
return result
def update_from_result(self, originals: list[dict], compressed: list[dict]) -> None:
"""Cache new compressions by comparing original and compressed messages.
Index-aligned: for each position, if both are tool results and the
content differs, store the mapping original_hash -> compressed_content.
"""
if len(originals) != len(compressed):
logger.warning(
"update_from_result: length mismatch (originals=%d, compressed=%d), skipping",
len(originals),
len(compressed),
)
return
for orig, comp in zip(originals, compressed):
orig_content = _extract_tool_result_content(orig)
comp_content = _extract_tool_result_content(comp)
if orig_content is None or comp_content is None:
continue
if orig_content == comp_content:
continue
h = self.content_hash(orig_content)
tokens_saved = len(orig_content) // 4 - len(comp_content) // 4
self.store_compressed(h, comp_content, tokens_saved=max(tokens_saved, 0))

View file

@ -69,6 +69,85 @@ logger = logging.getLogger("headroom.ccr.mcp")
DEFAULT_PROXY_URL = os.environ.get("HEADROOM_PROXY_URL", "http://127.0.0.1:8787")
def _format_session_summary(summary: dict[str, Any], local_stats: dict[str, Any]) -> str:
"""Format the proxy summary + local MCP stats into clean readable text."""
lines: list[str] = []
lines.append("Headroom Session Summary")
lines.append("=" * 40)
mode = summary.get("mode", "cost_savings")
api_reqs = summary.get("api_requests", 0)
model = summary.get("primary_model", "unknown")
lines.append(f"Mode: {mode} | {api_reqs} API requests | {model}")
lines.append("")
# Compression section
comp = summary.get("compression", {})
n_compressed = comp.get("requests_compressed", 0)
if n_compressed > 0:
lines.append(f"Compression ({n_compressed} requests compressed):")
lines.append(f" Avg compression: {comp.get('avg_compression_pct', 0)}%")
best = comp.get("best_compression_pct", 0)
detail = comp.get("best_detail", "")
if best > 0:
lines.append(f" Best compression: {best}% ({detail})")
removed = comp.get("total_tokens_removed", 0)
lines.append(f" Tokens removed: {removed:,}")
else:
lines.append("Compression: no requests compressed yet")
lines.append("")
# Uncompressed reasons
uncomp = summary.get("uncompressed_requests", {})
if uncomp:
total_uncomp = sum(uncomp.values())
lines.append(f"Uncompressed requests ({total_uncomp}):")
reason_labels = {
"prefix_frozen": "Prefix-frozen (cached by provider)",
"too_small": "Too small (< 500 tokens)",
"passthrough": "Passthrough (token counting)",
"no_compressible_content": "No compressible content (user/assistant only)",
}
for key, count in uncomp.items():
label = reason_labels.get(key, key)
lines.append(f" {label}: {count}")
lines.append("")
# Cost section
cost = summary.get("cost", {})
without = cost.get("without_headroom_usd", 0)
with_hr = cost.get("with_headroom_usd", 0)
saved = cost.get("total_saved_usd", 0)
pct = cost.get("savings_pct", 0)
if without > 0:
lines.append("Cost Impact:")
lines.append(f" Without Headroom: ${without:.2f}")
lines.append(f" With Headroom: ${with_hr:.2f}")
lines.append(f" You saved: ${saved:.2f} ({pct}%)")
breakdown = cost.get("breakdown", {})
cache_s = breakdown.get("cache_savings_usd", 0)
comp_s = breakdown.get("compression_savings_usd", 0)
if cache_s > 0 or comp_s > 0:
lines.append(f" Cache savings: ${cache_s:.2f}")
lines.append(f" Compression savings: ${comp_s:.2f}")
lines.append("")
# MCP-local stats (compressions done by MCP tool directly)
local_compressions = local_stats.get("compressions", 0)
local_saved = local_stats.get("total_tokens_saved", 0)
if local_compressions > 0:
lines.append(f"MCP Tool: {local_compressions} compressions, {local_saved:,} tokens saved")
lines.append("")
# Tip
tip = summary.get("tip")
if tip:
lines.append(f"Tip: {tip}")
return "\n".join(lines)
# Session-scoped TTL: content persists for the session (1 hour), not 5 minutes.
# The MCP server process lives as long as the coding session.
MCP_SESSION_TTL = 3600
@ -542,45 +621,55 @@ class HeadroomMCPServer:
"estimated_cost_saved_usd": round(all_saved * 3.0 / 1_000_000, 4),
}
# Fetch proxy stats (prefix cache hits, etc.) if proxy is reachable
# Fetch proxy stats and format summary if proxy is reachable
if self.check_proxy and HTTPX_AVAILABLE:
proxy_stats = await self._fetch_proxy_stats()
if proxy_stats:
stats["proxy"] = proxy_stats
proxy_data = await self._fetch_full_proxy_stats()
if proxy_data:
summary = proxy_data.get("summary")
if summary:
# Return clean formatted summary instead of raw JSON
formatted = _format_session_summary(summary, stats)
return [TextContent(type="text", text=formatted)]
# Fallback: add proxy stats to local stats
proxy_stats = self._extract_proxy_stats(proxy_data)
if proxy_stats:
stats["proxy"] = proxy_stats
return [TextContent(type="text", text=json.dumps(stats, indent=2))]
async def _fetch_proxy_stats(self) -> dict[str, Any] | None:
"""Fetch stats from the proxy, including prefix cache hit info."""
async def _fetch_full_proxy_stats(self) -> dict[str, Any] | None:
"""Fetch full stats from the proxy (includes summary)."""
try:
if self._http_client is None:
self._http_client = httpx.AsyncClient(timeout=15.0)
response = await self._http_client.get(f"{self.proxy_url}/stats")
if response.status_code != 200:
return None
data = response.json()
# Extract the most useful fields
result: dict[str, Any] = {}
if "requests_total" in data:
result["requests_total"] = data["requests_total"]
if "tokens_saved_total" in data:
result["tokens_saved_total"] = data["tokens_saved_total"]
# Prefix cache stats
cache = data.get("cache", data.get("caching", {}))
if cache:
result["cache"] = {
"hits": cache.get("hits", cache.get("cache_hits", 0)),
"misses": cache.get("misses", cache.get("cache_misses", 0)),
"hit_rate": cache.get("hit_rate", cache.get("cache_hit_rate", 0)),
}
# Cost tracking
cost = data.get("cost", {})
if cost:
result["cost_saved_usd"] = cost.get("total_saved", cost.get("saved", 0))
return result if result else None
result: dict[str, Any] = response.json()
return result
except Exception:
return None
@staticmethod
def _extract_proxy_stats(data: dict[str, Any]) -> dict[str, Any] | None:
"""Extract key fields from full proxy stats (fallback when no summary)."""
result: dict[str, Any] = {}
if "requests_total" in data:
result["requests_total"] = data["requests_total"]
if "tokens_saved_total" in data:
result["tokens_saved_total"] = data["tokens_saved_total"]
cache = data.get("cache", data.get("caching", {}))
if cache:
result["cache"] = {
"hits": cache.get("hits", cache.get("cache_hits", 0)),
"misses": cache.get("misses", cache.get("cache_misses", 0)),
"hit_rate": cache.get("hit_rate", cache.get("cache_hit_rate", 0)),
}
cost = data.get("cost", {})
if cost:
result["cost_saved_usd"] = cost.get("total_saved", cost.get("saved", 0))
return result if result else None
async def run_stdio(self) -> None:
"""Run the server with stdio transport."""
async with stdio_server() as (read_stream, write_stream):

View file

@ -45,6 +45,8 @@ _UNKNOWN_MODEL_WARNINGS: set[str] = set()
# Anthropic model context limits
# All Claude 3+ models have 200K context
ANTHROPIC_CONTEXT_LIMITS: dict[str, int] = {
# Claude 4.6 (Opus 4.6) - 1M context
"claude-opus-4-6": 1000000,
# Claude 4.5 (Opus 4.5)
"claude-opus-4-5-20251101": 200000,
# Claude 4 (Sonnet 4, Haiku 4)
@ -70,6 +72,8 @@ ANTHROPIC_CONTEXT_LIMITS: dict[str, int] = {
# NOTE: These are ESTIMATES. Always verify against actual Anthropic billing.
# Last updated: 2025-01-14
ANTHROPIC_PRICING: dict[str, dict[str, float]] = {
# Claude 4.6 (Opus tier pricing)
"claude-opus-4-6": {"input": 15.00, "output": 75.00, "cached_input": 1.50},
# Claude 4.5 (Opus tier pricing)
"claude-opus-4-5-20251101": {"input": 15.00, "output": 75.00, "cached_input": 1.50},
# Claude 4 (Sonnet/Haiku tier pricing)

View file

@ -39,6 +39,7 @@ from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
if TYPE_CHECKING:
from ..cache.compression_cache import CompressionCache
from ..memory.tracker import ComponentStats, MemoryTracker
import httpx
@ -411,6 +412,116 @@ def _merge_cost_stats(
}
def _build_session_summary(
proxy: HeadroomProxy,
metrics: Any,
prefix_cache_stats: dict,
cli_tokens_avoided: int,
total_tokens_before: int,
) -> dict[str, Any]:
"""Build a human-readable session summary from metrics and request logs.
This is the headline view users see first in /stats designed to answer
"is Headroom working?" at a glance.
"""
# Analyze per-request compression from the logger
compressed_requests: list[dict] = []
uncompressed_reasons: dict[str, int] = {
"prefix_frozen": 0,
"too_small": 0,
"passthrough": 0,
"no_compressible_content": 0,
}
if proxy.logger:
for entry in proxy.logger._logs:
if entry.model and "count_tokens" in entry.model:
uncompressed_reasons["passthrough"] += 1
continue
if entry.tokens_saved > 0 and entry.savings_percent > 0:
compressed_requests.append(
{
"savings_pct": round(entry.savings_percent, 1),
"tokens_saved": entry.tokens_saved,
"original": entry.input_tokens_original,
"optimized": entry.input_tokens_optimized,
}
)
elif entry.input_tokens_original > 0:
# Categorize why it wasn't compressed
transforms = entry.transforms_applied or []
if not transforms:
# Pipeline returned unchanged — likely all frozen
uncompressed_reasons["prefix_frozen"] += 1
elif all("excluded" in t or "protected" in t for t in transforms):
uncompressed_reasons["no_compressible_content"] += 1
elif entry.input_tokens_original < 500:
uncompressed_reasons["too_small"] += 1
else:
uncompressed_reasons["prefix_frozen"] += 1
# Compute compression stats for requests that DID compress
avg_compression = 0.0
best_compression = 0.0
best_detail = ""
if compressed_requests:
avg_compression = round(
sum(r["savings_pct"] for r in compressed_requests) / len(compressed_requests),
1,
)
best = max(compressed_requests, key=lambda r: r["savings_pct"])
best_compression = best["savings_pct"]
best_detail = f"{best['original']:,}{best['optimized']:,} tokens"
# Cost summary
cost_stats = proxy.cost_tracker.stats() if proxy.cost_tracker else {}
cost_without = cost_stats.get("cost_without_headroom_usd", 0.0)
cache_net = prefix_cache_stats.get("totals", {}).get("net_savings_usd", 0.0)
compression_savings = cost_stats.get("savings_usd", 0.0) if cost_stats else 0.0
total_saved_usd = round(compression_savings + cache_net, 2)
cost_with = round(cost_without - total_saved_usd, 2) if cost_without else 0.0
savings_pct_cost = round(total_saved_usd / cost_without * 100, 1) if cost_without > 0 else 0.0
# Primary models used
models = dict(metrics.requests_by_model)
primary_model = max(models, key=lambda k: models[k]) if models else "unknown"
api_requests = sum(v for k, v in models.items() if "count_tokens" not in k)
# Build the summary
summary: dict[str, Any] = {
"mode": proxy.config.mode,
"api_requests": api_requests,
"primary_model": primary_model,
"compression": {
"requests_compressed": len(compressed_requests),
"avg_compression_pct": avg_compression,
"best_compression_pct": best_compression,
"best_detail": best_detail,
"total_tokens_removed": metrics.tokens_saved_total,
},
"uncompressed_requests": {k: v for k, v in uncompressed_reasons.items() if v > 0},
"cost": {
"without_headroom_usd": round(cost_without, 2),
"with_headroom_usd": cost_with,
"total_saved_usd": total_saved_usd,
"savings_pct": savings_pct_cost,
"breakdown": {
"cache_savings_usd": round(cache_net, 2),
"compression_savings_usd": round(compression_savings, 2),
},
},
}
# Add tip if token_headroom mode would help
if proxy.config.mode == "cost_savings" and uncompressed_reasons["prefix_frozen"] > 10:
summary["tip"] = (
"Most requests are prefix-frozen. Set HEADROOM_MODE=token_headroom "
"to compress frozen messages and extend your session by ~25-35%."
)
return summary
# Maximum request body size (100MB - increased to support image-heavy requests)
MAX_REQUEST_BODY_SIZE = 100 * 1024 * 1024
@ -492,6 +603,11 @@ class ProxyConfig:
bedrock_profile: str | None = None # AWS profile (optional)
anyllm_provider: str = "openai" # any-llm provider (openai, mistral, groq, etc.)
# Optimization mode: "cost_savings" (default) or "token_headroom"
# cost_savings: preserve prefix cache for cost reduction
# token_headroom: compress older messages for session extension
mode: str = "cost_savings"
# Optimization
optimize: bool = True
image_optimize: bool = True # Compress images using trained ML router
@ -1422,6 +1538,9 @@ class HeadroomProxy:
tool_profiles=config.tool_profiles,
read_lifecycle=ReadLifecycleConfig(enabled=config.read_lifecycle),
)
# Token headroom mode: allow compression of older excluded-tool results
if config.mode == "token_headroom":
router_config.protect_recent_reads_fraction = 0.3
transforms = [
CacheAligner(CacheAlignerConfig(enabled=True)),
ContentRouter(router_config),
@ -1501,6 +1620,9 @@ class HeadroomProxy:
)
)
# Compression cache store for token_headroom mode (session-scoped)
self._compression_caches: dict[str, CompressionCache] = {}
self.logger = (
RequestLogger(
log_file=config.log_file,
@ -1618,6 +1740,14 @@ class HeadroomProxy:
)
self.memory_handler = MemoryHandler(memory_config)
def _get_compression_cache(self, session_id: str) -> CompressionCache:
"""Get or create a CompressionCache for a session."""
if session_id not in self._compression_caches:
from headroom.cache.compression_cache import CompressionCache
self._compression_caches[session_id] = CompressionCache()
return self._compression_caches[session_id]
def _setup_llmlingua(self, config: ProxyConfig, transforms: list) -> str:
"""Set up LLMLingua compression if enabled.
@ -1698,6 +1828,17 @@ class HeadroomProxy:
)
logger.info("Headroom Proxy started")
logger.info(f"Optimization: {'ENABLED' if self.config.optimize else 'DISABLED'}")
if self.config.mode not in ("cost_savings", "token_headroom"):
logger.warning(
f"Unknown HEADROOM_MODE '{self.config.mode}', falling back to 'cost_savings'"
)
self.config.mode = "cost_savings"
logger.info(f"Mode: {self.config.mode}")
if self.config.mode == "token_headroom":
logger.info(" Prefix freeze: re-freeze after compression")
logger.info(" Read protection window: 30%% of excluded-tool messages")
logger.info(" CCR TTL: extended for session lifetime")
logger.info(" Compression cache: active")
logger.info(f"Caching: {'ENABLED' if self.config.cache_enabled else 'DISABLED'}")
logger.info(f"Rate Limiting: {'ENABLED' if self.config.rate_limit_enabled else 'DISABLED'}")
logger.info(
@ -2071,23 +2212,55 @@ class HeadroomProxy:
if self.config.optimize and messages:
try:
context_limit = self.anthropic_provider.get_context_limit(model)
result = self.anthropic_pipeline.apply(
messages=messages,
model=model,
model_limit=context_limit,
frozen_message_count=frozen_message_count,
biases=self.config.hooks.compute_biases(messages, _hook_ctx)
biases = (
self.config.hooks.compute_biases(messages, _hook_ctx)
if self.config.hooks
else None,
else None
)
if result.messages != messages:
if self.config.mode == "token_headroom":
comp_cache = self._get_compression_cache(session_id)
# Zone 1: Swap cached compressed versions into working copy
working_messages = comp_cache.apply_cached(messages)
# Re-freeze boundary: consecutive stable messages from start
frozen_message_count = comp_cache.compute_frozen_count(messages)
result = self.anthropic_pipeline.apply(
messages=working_messages,
model=model,
model_limit=context_limit,
frozen_message_count=frozen_message_count,
biases=biases,
)
# Cache newly compressed messages (index-aligned diff)
if result.messages != working_messages:
comp_cache.update_from_result(messages, result.messages)
# Always use pipeline result — Zone 1 swaps are already applied
optimized_messages = result.messages
transforms_applied = result.transforms_applied
pipeline_timing = result.timing
# Use pipeline's token counts for consistency with pipeline logs
original_tokens = result.tokens_before
optimized_tokens = result.tokens_after
else:
result = self.anthropic_pipeline.apply(
messages=messages,
model=model,
model_limit=context_limit,
frozen_message_count=frozen_message_count,
biases=biases,
)
if result.messages != messages:
optimized_messages = result.messages
transforms_applied = result.transforms_applied
pipeline_timing = result.timing
original_tokens = result.tokens_before
optimized_tokens = result.tokens_after
if result.waste_signals:
waste_signals_dict = result.waste_signals.to_dict()
except Exception as e:
@ -4604,19 +4777,49 @@ class HeadroomProxy:
if self.config.optimize and messages:
try:
context_limit = self.openai_provider.get_context_limit(model)
result = self.openai_pipeline.apply(
messages=messages,
model=model,
model_limit=context_limit,
frozen_message_count=openai_frozen_count,
biases=_hook_biases,
)
if result.messages != messages:
if self.config.mode == "token_headroom":
comp_cache = self._get_compression_cache(openai_session_id)
# Zone 1: Swap cached compressed versions
working_messages = comp_cache.apply_cached(messages)
# Re-freeze boundary
openai_frozen_count = comp_cache.compute_frozen_count(messages)
result = self.openai_pipeline.apply(
messages=working_messages,
model=model,
model_limit=context_limit,
frozen_message_count=openai_frozen_count,
biases=_hook_biases,
)
if result.messages != working_messages:
comp_cache.update_from_result(messages, result.messages)
# Always use pipeline result in token_headroom mode
optimized_messages = result.messages
transforms_applied = result.transforms_applied
pipeline_timing = result.timing
original_tokens = result.tokens_before
optimized_tokens = result.tokens_after
else:
result = self.openai_pipeline.apply(
messages=messages,
model=model,
model_limit=context_limit,
frozen_message_count=openai_frozen_count,
biases=_hook_biases,
)
if result.messages != messages:
optimized_messages = result.messages
transforms_applied = result.transforms_applied
pipeline_timing = result.timing
original_tokens = result.tokens_before
optimized_tokens = result.tokens_after
if result.waste_signals:
waste_signals_dict = result.waste_signals.to_dict()
except Exception as e:
@ -6305,12 +6508,43 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
# Calculate total tokens before compression
total_tokens_before = m.tokens_input_total + m.tokens_saved_total
# Build human-readable summary
summary = _build_session_summary(
proxy, m, prefix_cache_stats, cli_tokens_avoided, total_tokens_before
)
# Compression cache stats (token_headroom mode)
compression_cache_stats: dict = {}
if proxy.config.mode == "token_headroom" and proxy._compression_caches:
total_entries = 0
total_hits = 0
total_misses = 0
total_tokens_saved = 0
for cache in proxy._compression_caches.values():
s = cache.get_stats()
total_entries += s.get("entries", 0)
total_hits += s.get("hits", 0)
total_misses += s.get("misses", 0)
total_tokens_saved += s.get("total_tokens_saved", 0)
compression_cache_stats = {
"mode": "token_headroom",
"active_sessions": len(proxy._compression_caches),
"total_entries": total_entries,
"total_hits": total_hits,
"total_misses": total_misses,
"hit_rate": round(total_hits / max(1, total_hits + total_misses) * 100, 1),
"total_tokens_saved": total_tokens_saved,
}
else:
compression_cache_stats = {"mode": "cost_savings"}
# Build unified savings summary (all layers)
compression_tokens = m.tokens_saved_total
cache_net_usd = prefix_cache_stats.get("totals", {}).get("net_savings_usd", 0.0)
total_tokens_all_layers = compression_tokens + cli_tokens_avoided
return {
"summary": summary,
"savings": {
"total_tokens": total_tokens_all_layers,
"by_layer": {
@ -6396,6 +6630,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
"compressed_tokens_cached": compression_stats.get("total_compressed_tokens", 0),
"ccr_retrievals": compression_stats.get("total_retrievals", 0),
},
"compression_cache": compression_cache_stats,
"telemetry": {
"enabled": telemetry_stats.get("enabled", False),
"total_compressions": telemetry_stats.get("total_compressions", 0),
@ -7606,6 +7841,7 @@ if __name__ == "__main__":
max_keepalive_connections=_get_env_int("HEADROOM_MAX_KEEPALIVE", args.max_keepalive),
http2=not args.no_http2 and _get_env_bool("HEADROOM_HTTP2", True),
tool_profiles=tool_profiles if tool_profiles else None,
mode=_get_env_str("HEADROOM_MODE", "cost_savings"),
)
# Get worker and concurrency settings

View file

@ -0,0 +1,283 @@
"""Tests for CompressionCache with LRU eviction."""
from __future__ import annotations
import pytest
from headroom.cache.compression_cache import CompressionCache
@pytest.fixture
def cache() -> CompressionCache:
return CompressionCache()
@pytest.fixture
def small_cache() -> CompressionCache:
return CompressionCache(max_entries=3)
class TestCompressionCache:
def test_cache_miss_returns_none(self, cache: CompressionCache) -> None:
h = CompressionCache.content_hash("some content")
assert cache.get_compressed(h) is None
def test_store_and_retrieve(self, cache: CompressionCache) -> None:
content = "hello world this is a long message"
h = CompressionCache.content_hash(content)
cache.store_compressed(h, "hello world...compressed", tokens_saved=15)
assert cache.get_compressed(h) == "hello world...compressed"
def test_different_content_different_hash(self) -> None:
h1 = CompressionCache.content_hash("content A")
h2 = CompressionCache.content_hash("content B")
assert h1 != h2
def test_overwrite_same_hash(self, cache: CompressionCache) -> None:
h = CompressionCache.content_hash("some content")
cache.store_compressed(h, "v1", tokens_saved=10)
cache.store_compressed(h, "v2", tokens_saved=20)
assert cache.get_compressed(h) == "v2"
def test_stats_tracking(self, cache: CompressionCache) -> None:
h = CompressionCache.content_hash("content")
cache.store_compressed(h, "compressed", tokens_saved=5)
# One hit
cache.get_compressed(h)
# One miss
cache.get_compressed("nonexistent")
stats = cache.get_stats()
assert stats["hits"] == 1
assert stats["misses"] == 1
assert stats["entries"] == 1
assert stats["tokens_saved"] == 5
def test_eviction_at_max_entries(self, small_cache: CompressionCache) -> None:
h1 = CompressionCache.content_hash("a")
h2 = CompressionCache.content_hash("b")
h3 = CompressionCache.content_hash("c")
h4 = CompressionCache.content_hash("d")
small_cache.store_compressed(h1, "ca", tokens_saved=1)
small_cache.store_compressed(h2, "cb", tokens_saved=1)
small_cache.store_compressed(h3, "cc", tokens_saved=1)
# Adding a 4th should evict the oldest (h1)
small_cache.store_compressed(h4, "cd", tokens_saved=1)
assert small_cache.get_compressed(h1) is None
assert small_cache.get_compressed(h2) == "cb"
assert small_cache.get_compressed(h4) == "cd"
def test_access_refreshes_lru(self, small_cache: CompressionCache) -> None:
h1 = CompressionCache.content_hash("a")
h2 = CompressionCache.content_hash("b")
h3 = CompressionCache.content_hash("c")
h4 = CompressionCache.content_hash("d")
small_cache.store_compressed(h1, "ca", tokens_saved=1)
small_cache.store_compressed(h2, "cb", tokens_saved=1)
small_cache.store_compressed(h3, "cc", tokens_saved=1)
# Access h1 to refresh it
small_cache.get_compressed(h1)
# Adding h4 should evict h2 (oldest untouched), not h1
small_cache.store_compressed(h4, "cd", tokens_saved=1)
assert small_cache.get_compressed(h1) == "ca"
assert small_cache.get_compressed(h2) is None
assert small_cache.get_compressed(h4) == "cd"
def test_content_hash_list_content(self) -> None:
"""content_hash handles Anthropic-format list content."""
list_content = [
{"type": "text", "text": "hello"},
{"type": "text", "text": "world"},
]
h = CompressionCache.content_hash(list_content)
assert isinstance(h, str)
assert len(h) == 16
# Same content produces same hash
assert CompressionCache.content_hash(list_content) == h
def test_content_hash_string_length(self) -> None:
h = CompressionCache.content_hash("test")
assert len(h) == 16
class TestCompressionCacheFrozenCount:
def test_empty_cache_returns_zero(self, cache: CompressionCache) -> None:
assert cache.compute_frozen_count([]) == 0
def test_user_assistant_always_stable(self, cache: CompressionCache) -> None:
messages = [
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi there"},
{"role": "user", "content": "how are you"},
]
assert cache.compute_frozen_count(messages) == 3
def test_tool_result_with_cache_hit_is_stable(self, cache: CompressionCache) -> None:
tool_content = "tool output data"
h = CompressionCache.content_hash(tool_content)
cache.store_compressed(h, "compressed tool output", tokens_saved=5)
messages = [
{"role": "user", "content": "do something"},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "t1", "name": "my_tool", "input": {}}],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}],
},
]
assert cache.compute_frozen_count(messages) == 3
def test_tool_result_cache_miss_stops_frozen(self, cache: CompressionCache) -> None:
messages = [
{"role": "user", "content": "hello"},
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "t1", "content": "uncached stuff"}
],
},
{"role": "user", "content": "follow up"},
]
assert cache.compute_frozen_count(messages) == 1
def test_frozen_count_with_dropped_messages(self, cache: CompressionCache) -> None:
cached_content = "cached tool output"
h = CompressionCache.content_hash(cached_content)
cache.store_compressed(h, "compressed", tokens_saved=3)
messages = [
{"role": "user", "content": "start"},
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "t1", "content": cached_content}
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t2", "content": "not cached"}],
},
]
assert cache.compute_frozen_count(messages) == 2
class TestCompressionCacheApplyAndUpdate:
def test_apply_cached_swaps_tool_results(self, cache: CompressionCache) -> None:
original_content = "big tool output"
h = CompressionCache.content_hash(original_content)
cache.store_compressed(h, "small output", tokens_saved=5)
messages = [
{"role": "user", "content": "hi"},
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "t1", "content": original_content}
],
},
]
result = cache.apply_cached(messages)
assert result[1]["content"][0]["content"] == "small output"
def test_apply_cached_preserves_uncached_messages(self, cache: CompressionCache) -> None:
messages = [
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "world"},
]
result = cache.apply_cached(messages)
assert result[0] is messages[0]
assert result[1] is messages[1]
def test_apply_cached_never_adds_messages(self, cache: CompressionCache) -> None:
# Store something in cache that doesn't correspond to any message
cache.store_compressed("orphan_hash", "orphan_value", tokens_saved=1)
messages = [
{"role": "user", "content": "hello"},
{"role": "assistant", "content": "hi"},
]
result = cache.apply_cached(messages)
assert len(result) == len(messages)
def test_update_from_result_caches_changes(self, cache: CompressionCache) -> None:
originals = [
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "t1", "content": "original output"}
],
},
]
compressed = [
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "t1", "content": "compressed output"}
],
},
]
cache.update_from_result(originals, compressed)
h = CompressionCache.content_hash("original output")
assert cache.get_compressed(h) == "compressed output"
def test_update_from_result_ignores_unchanged(self, cache: CompressionCache) -> None:
originals = [
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "t1", "content": "same content"}
],
},
]
compressed = [
{
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "t1", "content": "same content"}
],
},
]
cache.update_from_result(originals, compressed)
h = CompressionCache.content_hash("same content")
assert cache.get_compressed(h) is None
def test_apply_does_not_modify_original_messages(self, cache: CompressionCache) -> None:
original_content = "big tool output"
h = CompressionCache.content_hash(original_content)
cache.store_compressed(h, "small output", tokens_saved=5)
msg = {
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": original_content}],
}
messages = [msg]
cache.apply_cached(messages)
# Original must be untouched
assert msg["content"][0]["content"] == original_content
def test_openai_format_tool_result(self, cache: CompressionCache) -> None:
original_content = "openai tool output"
h = CompressionCache.content_hash(original_content)
cache.store_compressed(h, "compressed openai", tokens_saved=4)
messages = [
{"role": "tool", "tool_call_id": "tc1", "content": original_content},
]
result = cache.apply_cached(messages)
assert result[0]["content"] == "compressed openai"
# Original untouched
assert messages[0]["content"] == original_content

View file

@ -0,0 +1,307 @@
"""Integration tests for token_headroom mode.
Tests the CompressionCache working across simulated multi-turn conversations,
verifying the critical invariants: no message injection, correct frozen counts,
proper handling of both Anthropic and OpenAI formats, and correct behavior
when Claude Code drops messages.
"""
import copy
from headroom.cache.compression_cache import CompressionCache
def _make_user_msg(text: str) -> dict:
return {"role": "user", "content": text}
def _make_assistant_msg(text: str) -> dict:
return {"role": "assistant", "content": text}
def _make_tool_use_msg(tool_id: str, name: str) -> dict:
return {
"role": "assistant",
"content": [{"type": "tool_use", "id": tool_id, "name": name, "input": {}}],
}
def _make_tool_result_msg(tool_id: str, content: str) -> dict:
"""Anthropic-format tool result."""
return {
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": tool_id, "content": content}],
}
def _make_openai_tool_msg(tool_call_id: str, content: str) -> dict:
"""OpenAI-format tool result."""
return {"role": "tool", "tool_call_id": tool_call_id, "content": content}
def _large_code_content(n: int = 200) -> str:
"""Generate realistic Python code content."""
parts = ["import os\nimport sys\nfrom typing import List, Dict\n\n"]
for i in range(n // 10):
parts.append(
f"def function_{i}(arg: str) -> str:\n"
f' """Docstring for function {i}."""\n'
f" result = arg.strip()\n"
f" for j in range({i}):\n"
f" result += str(j)\n"
f" return result\n\n"
)
return "".join(parts)
class TestMultiTurnCompression:
"""Simulate multi-turn conversations to verify compression cascade."""
def test_first_turn_nothing_cached(self):
"""On first turn, no cache hits, frozen count is minimal."""
cache = CompressionCache()
messages = [
_make_user_msg("Read file.py"),
_make_tool_use_msg("t1", "Read"),
_make_tool_result_msg("t1", _large_code_content(100)),
]
frozen = cache.compute_frozen_count(messages)
# user (stable) + tool_use (stable) + tool_result (miss) → 2
assert frozen == 2
def test_second_turn_cache_hits(self):
"""After caching, same content gets cache hits."""
cache = CompressionCache()
code = _large_code_content(100)
compressed = "# compressed version"
# Simulate first turn: pipeline compressed the code
h = CompressionCache.content_hash(code)
cache.store_compressed(h, compressed, tokens_saved=500)
# Second turn: same messages
messages = [
_make_user_msg("Read file.py"),
_make_tool_use_msg("t1", "Read"),
_make_tool_result_msg("t1", code),
_make_user_msg("now edit it"),
]
frozen = cache.compute_frozen_count(messages)
# All 4 messages stable (user, tool_use, tool_result cached, user)
assert frozen == 4
# apply_cached should swap the content
result = cache.apply_cached(messages)
tool_result = result[2]["content"][0]
assert tool_result["content"] == compressed
def test_multi_turn_waterfall(self):
"""Messages age out progressively across turns."""
cache = CompressionCache()
# Build conversation with 3 read results
code_a = "code A " * 200
code_b = "code B " * 200
code_c = "code C " * 200
messages = [
_make_user_msg("read A"),
_make_tool_result_msg("t1", code_a),
_make_user_msg("read B"),
_make_tool_result_msg("t2", code_b),
_make_user_msg("read C"),
_make_tool_result_msg("t3", code_c),
]
# Turn 1: nothing cached
frozen = cache.compute_frozen_count(messages)
assert frozen == 1 # only first user msg
# Simulate pipeline compressing A and B (not C — in protection window)
cache.store_compressed(CompressionCache.content_hash(code_a), "ca", tokens_saved=100)
cache.store_compressed(CompressionCache.content_hash(code_b), "cb", tokens_saved=100)
# Turn 2: A and B cached, C still uncached
frozen = cache.compute_frozen_count(messages)
# user(stable) + tool_result_A(cached) + user(stable) + tool_result_B(cached) + user(stable) + tool_result_C(miss)
assert frozen == 5
# Now cache C too
cache.store_compressed(CompressionCache.content_hash(code_c), "cc", tokens_saved=100)
# Turn 3: all cached
frozen = cache.compute_frozen_count(messages)
assert frozen == 6 # all stable
class TestNoMessageInjection:
"""Critical invariant: proxy never adds messages."""
def test_output_length_equals_input(self):
cache = CompressionCache()
messages = [
_make_user_msg("hello"),
_make_tool_result_msg("t1", _large_code_content(50)),
_make_user_msg("bye"),
]
result = cache.apply_cached(messages)
assert len(result) == len(messages)
def test_orphan_cache_entries_not_injected(self):
"""Cache entries with no matching message are NOT injected."""
cache = CompressionCache()
cache.store_compressed("orphan_hash_1", "orphan content 1", tokens_saved=100)
cache.store_compressed("orphan_hash_2", "orphan content 2", tokens_saved=200)
messages = [_make_user_msg("hello")]
result = cache.apply_cached(messages)
assert len(result) == 1
assert result[0]["content"] == "hello"
def test_input_not_mutated(self):
"""apply_cached must NOT mutate the input list or messages."""
cache = CompressionCache()
code = "original code content"
h = CompressionCache.content_hash(code)
cache.store_compressed(h, "compressed", tokens_saved=50)
messages = [
_make_tool_result_msg("t1", code),
]
original = copy.deepcopy(messages)
cache.apply_cached(messages)
assert messages == original
class TestClaudeCodeDropsMessages:
"""When Claude Code drops messages via its own context management."""
def test_dropped_messages_not_readded(self):
cache = CompressionCache()
content_a = "content A " * 100
content_b = "content B " * 100
cache.store_compressed(CompressionCache.content_hash(content_a), "ca", tokens_saved=100)
cache.store_compressed(CompressionCache.content_hash(content_b), "cb", tokens_saved=100)
# CC dropped the message with content_b
messages = [
_make_user_msg("hello"),
_make_tool_result_msg("t1", content_a),
_make_user_msg("continue"),
]
result = cache.apply_cached(messages)
assert len(result) == 3 # NOT 4
def test_frozen_count_breaks_at_gap(self):
"""Dropped cached message creates a gap that stops frozen count."""
cache = CompressionCache()
content_a = "content A " * 100
content_c = "content C " * 100
cache.store_compressed(CompressionCache.content_hash(content_a), "ca", tokens_saved=100)
cache.store_compressed(CompressionCache.content_hash(content_c), "cc", tokens_saved=100)
# CC dropped content_b, content_c is still here but preceded by uncached gap
messages = [
_make_tool_result_msg("t1", content_a),
_make_tool_result_msg("t2", "UNCACHED content_b replacement"),
_make_tool_result_msg("t3", content_c),
]
frozen = cache.compute_frozen_count(messages)
# t1 (cached, stable), t2 (NOT cached, stop)
assert frozen == 1
class TestOpenAIFormat:
"""Verify OpenAI-format tool messages work correctly."""
def test_openai_tool_result_cached(self):
cache = CompressionCache()
content = "large openai output " * 100
h = CompressionCache.content_hash(content)
cache.store_compressed(h, "compressed openai output", tokens_saved=300)
messages = [
_make_user_msg("run command"),
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "tc1",
"type": "function",
"function": {"name": "bash", "arguments": "{}"},
}
],
},
_make_openai_tool_msg("tc1", content),
]
result = cache.apply_cached(messages)
assert len(result) == 3
assert result[2]["content"] == "compressed openai output"
def test_openai_frozen_count(self):
cache = CompressionCache()
content = "openai tool output " * 100
h = CompressionCache.content_hash(content)
cache.store_compressed(h, "compressed", tokens_saved=200)
messages = [
_make_user_msg("hello"),
_make_openai_tool_msg("tc1", content),
_make_openai_tool_msg("tc2", "uncached content"),
]
frozen = cache.compute_frozen_count(messages)
# user (stable), tool tc1 (cached), tool tc2 (miss)
assert frozen == 2
class TestUpdateFromResult:
"""Verify update_from_result correctly caches compression results."""
def test_caches_compressed_anthropic(self):
cache = CompressionCache()
original_content = "long original " * 100
compressed_content = "short compressed"
originals = [
_make_user_msg("hello"),
_make_tool_result_msg("t1", original_content),
]
compressed = [
_make_user_msg("hello"),
_make_tool_result_msg("t1", compressed_content),
]
cache.update_from_result(originals, compressed)
h = CompressionCache.content_hash(original_content)
assert cache.get_compressed(h) == compressed_content
def test_caches_compressed_openai(self):
cache = CompressionCache()
original_content = "long openai output " * 100
compressed_content = "short compressed"
originals = [_make_openai_tool_msg("tc1", original_content)]
compressed = [_make_openai_tool_msg("tc1", compressed_content)]
cache.update_from_result(originals, compressed)
h = CompressionCache.content_hash(original_content)
assert cache.get_compressed(h) == compressed_content
def test_length_mismatch_no_crash(self):
"""If pipeline somehow changes message count, don't crash."""
cache = CompressionCache()
originals = [_make_user_msg("a"), _make_user_msg("b")]
compressed = [_make_user_msg("a")] # shorter
# Should not raise, just log warning
cache.update_from_result(originals, compressed)
assert cache.get_stats()["entries"] == 0
def test_unchanged_content_not_cached(self):
cache = CompressionCache()
msg = _make_user_msg("same content")
cache.update_from_result([msg], [msg])
assert cache.get_stats()["entries"] == 0