mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
commit
4bda5a31be
7 changed files with 1171 additions and 43 deletions
3
headroom/cache/__init__.py
vendored
3
headroom/cache/__init__.py
vendored
|
|
@ -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
206
headroom/cache/compression_cache.py
vendored
Normal 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))
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
283
tests/test_compression_cache.py
Normal file
283
tests/test_compression_cache.py
Normal 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
|
||||
307
tests/test_token_headroom_mode.py
Normal file
307
tests/test_token_headroom_mode.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue