akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Rebrand proxy modes to token/cache and harden cache-mode stability

Browse source
This commit is contained in:
JerrettDavis 2026-04-04 14:32:07 -05:00
parent d78fdfe02d
commit 54419ad8b8
15 changed files with 671 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

282
benchmarks/proxy_mode_benchmark.py Normal file
View file

@ -0,0 +1,282 @@
#!/usr/bin/env python3
"""Local benchmark for proxy run modes (no API calls).
Compares:
- baseline: no compression
- token mode: prioritize compression
- cache mode: preserve prior-turn prefix stability
Includes an optional real-test harness printout for Claude Code, but does not
invoke external APIs unless the user does so manually.
"""
from __future__ import annotations
import argparse
import copy
import json
import logging
from dataclasses import dataclass
from typing import Any
from headroom.cache.compression_cache import CompressionCache
from headroom.cache.prefix_tracker import PrefixCacheTracker
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
from headroom.proxy.models import ProxyConfig
from headroom.proxy.modes import PROXY_MODE_CACHE, PROXY_MODE_TOKEN
from headroom.proxy.server import HeadroomProxy
from headroom.tokenizers import get_tokenizer
from headroom.utils import extract_user_query
MODEL = "claude-sonnet-4-6"
@dataclass
class ModeBenchmarkResult:
mode: str
total_original_tokens: int = 0
total_sent_tokens: int = 0
total_tokens_saved: int = 0
total_cache_read_tokens: int = 0
total_cache_write_tokens: int = 0
total_uncached_tokens: int = 0
@property
def compression_pct(self) -> float:
if self.total_original_tokens <= 0:
return 0.0
return self.total_tokens_saved / self.total_original_tokens * 100.0
@property
def cache_hit_pct(self) -> float:
total = self.total_cache_read_tokens + self.total_cache_write_tokens + self.total_uncached_tokens
if total <= 0:
return 0.0
return self.total_cache_read_tokens / total * 100.0
def _build_tool_result(turn: int, rows: int = 240) -> str:
payload = []
for i in range(rows):
payload.append(
{
"id": f"{turn:02d}-{i:04d}",
"status": "ok" if i % 37 else "warning",
"service": "auth-api" if i % 2 else "gateway",
"latency_ms": 100 + (i % 13),
"hint": "retry with exponential backoff" if i % 89 == 0 else "none",
}
)
return json.dumps(payload)
def _build_conversation(turn: int) -> list[dict[str, Any]]:
messages: list[dict[str, Any]] = []
for t in range(1, turn):
messages.extend(
[
{"role": "user", "content": f"Analyze tool output turn {t} and summarize anomalies."},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": f"tool-{t}",
"content": _build_tool_result(t),
}
],
},
{"role": "assistant", "content": f"Turn {t} acknowledged."},
]
)
# Current turn: user request + fresh tool output, no assistant response yet.
messages.extend(
[
{"role": "user", "content": f"Analyze tool output turn {turn} and summarize anomalies."},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": f"tool-{turn}",
"content": _build_tool_result(turn),
}
],
},
]
)
return messages
def _common_prefix_tokens(
prev: list[dict[str, Any]], curr: list[dict[str, Any]], tokenizer: Any
) -> tuple[int, list[int]]:
common = 0
counts: list[int] = []
for msg in curr:
counts.append(tokenizer.count_message(msg))
for i, (a, b) in enumerate(zip(prev, curr)):
if a != b:
break
common += counts[i]
return common, counts
def _make_proxy(mode: str) -> HeadroomProxy:
cfg = ProxyConfig(
mode=mode,
optimize=True,
image_optimize=False,
smart_routing=False,
code_aware_enabled=False,
read_lifecycle=False,
intelligent_context=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
)
return HeadroomProxy(cfg)
def _simulate_mode(turns: int, mode: str) -> ModeBenchmarkResult:
tokenizer = get_tokenizer(MODEL)
result = ModeBenchmarkResult(mode=mode)
if mode == "baseline":
prev_forwarded: list[dict[str, Any]] = []
for turn in range(1, turns + 1):
messages = _build_conversation(turn)
before = tokenizer.count_messages(messages)
common, counts = _common_prefix_tokens(prev_forwarded, messages, tokenizer)
uncached = max(0, before - common)
result.total_original_tokens += before
result.total_sent_tokens += before
result.total_cache_read_tokens += common
result.total_cache_write_tokens += 0
result.total_uncached_tokens += uncached
prev_forwarded = copy.deepcopy(messages)
return result
proxy = _make_proxy(mode)
prefix_tracker = PrefixCacheTracker("anthropic")
comp_cache = CompressionCache()
prev_forwarded = []
for turn in range(1, turns + 1):
messages = _build_conversation(turn)
before = tokenizer.count_messages(messages)
frozen = prefix_tracker.get_frozen_message_count()
if mode == PROXY_MODE_CACHE:
frozen = AnthropicHandlerMixin._strict_previous_turn_frozen_count(messages, frozen)
working = messages
if mode == PROXY_MODE_TOKEN:
working = comp_cache.apply_cached(messages)
frozen = min(frozen, comp_cache.compute_frozen_count(messages))
context_limit = proxy.anthropic_provider.get_context_limit(MODEL)
pipeline_result = proxy.anthropic_pipeline.apply(
messages=working,
model=MODEL,
model_limit=context_limit,
context=extract_user_query(working),
frozen_message_count=frozen,
)
forwarded = pipeline_result.messages
if mode == PROXY_MODE_TOKEN:
comp_cache.update_from_result(messages, forwarded)
if mode == PROXY_MODE_CACHE:
forwarded, _ = AnthropicHandlerMixin._restore_frozen_prefix(
messages, forwarded, frozen_message_count=frozen
)
after = tokenizer.count_messages(forwarded)
common, msg_counts = _common_prefix_tokens(prev_forwarded, forwarded, tokenizer)
uncached = max(0, after - common)
result.total_original_tokens += before
result.total_sent_tokens += after
result.total_tokens_saved += max(0, before - after)
result.total_cache_read_tokens += common
result.total_uncached_tokens += uncached
prefix_tracker.update_from_response(
cache_read_tokens=common,
cache_write_tokens=uncached,
messages=forwarded,
message_token_counts=msg_counts,
)
result.total_cache_write_tokens += uncached
prev_forwarded = copy.deepcopy(forwarded)
return result
def run_local_benchmark(turns: int = 12) -> dict[str, ModeBenchmarkResult]:
return {
"baseline": _simulate_mode(turns, "baseline"),
PROXY_MODE_TOKEN: _simulate_mode(turns, PROXY_MODE_TOKEN),
PROXY_MODE_CACHE: _simulate_mode(turns, PROXY_MODE_CACHE),
}
def _print_results(results: dict[str, ModeBenchmarkResult]) -> None:
print(
"\nMode benchmark (higher compression + higher cache_hit is better for total cost):\n"
"mode orig_tok sent_tok saved_tok compression cache_hit uncached_tok"
)
for key in ("baseline", PROXY_MODE_TOKEN, PROXY_MODE_CACHE):
r = results[key]
print(
f"{r.mode:<9} {r.total_original_tokens:>9,} {r.total_sent_tokens:>10,} "
f"{r.total_tokens_saved:>10,} {r.compression_pct:>10.1f}% "
f"{r.cache_hit_pct:>9.1f}% {r.total_uncached_tokens:>12,}"
)
token = results[PROXY_MODE_TOKEN]
cache = results[PROXY_MODE_CACHE]
print("\nDelta (cache - token):")
print(f" cache_hit_pct: {cache.cache_hit_pct - token.cache_hit_pct:+.1f}%")
print(f" compression_pct: {cache.compression_pct - token.compression_pct:+.1f}%")
print(f" uncached_tokens: {cache.total_uncached_tokens - token.total_uncached_tokens:+,}")
def _print_real_harness() -> None:
print("\nReal test harness (manual; optional, not executed by this benchmark):")
print(" 1) Start proxy in cache mode: HEADROOM_MODE=cache headroom proxy --port 8787")
print(" 2) Start proxy in token mode: HEADROOM_MODE=token headroom proxy --port 8787")
print(" 3) Run Claude Code against each:")
print(" ANTHROPIC_BASE_URL=http://localhost:8787 claude")
print(" 4) Compare /stats prefix_cache and compression sections per run.")
def main() -> None:
logging.getLogger("headroom").setLevel(logging.WARNING)
logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
parser = argparse.ArgumentParser(description="Local benchmark for proxy token/cache modes")
parser.add_argument("--turns", type=int, default=12, help="Conversation turns to simulate")
parser.add_argument(
"--show-real-harness",
action="store_true",
help="Print manual steps for optional Claude Code real testing",
)
args = parser.parse_args()
results = run_local_benchmark(turns=args.turns)
_print_results(results)
if args.show_real_harness:
_print_real_harness()
if __name__ == "__main__":
main()

3
docs/benchmarks.md
View file

@ -203,4 +203,7 @@ pytest tests/test_evals/ -v -s
# Run compression benchmark
python -c "from headroom import compress; print(compress([{'role':'user','content':'test'}]))"
# Run local proxy mode benchmark (no API calls)
python benchmarks/proxy_mode_benchmark.py --turns 12 --show-real-harness
```

17
docs/proxy.md
View file

@ -31,6 +31,7 @@ Telemetry is enabled by default. Opt out with `HEADROOM_TELEMETRY=off` or `headr
|--------|---------|-------------|
| `--host` | `127.0.0.1` | Host to bind to |
| `--port` | `8787` | Port to bind to |
| `--mode` | `token` | Run mode: `token` (maximize compression) or `cache` (freeze prior turns) |
| `--no-optimize` | `false` | Disable optimization (passthrough mode) |
| `--no-cache` | `false` | Disable semantic caching |
| `--no-rate-limit` | `false` | Disable rate limiting |
@ -38,6 +39,22 @@ Telemetry is enabled by default. Opt out with `HEADROOM_TELEMETRY=off` or `headr
| `--budget` | None | Daily budget limit in USD |
| `--openai-api-url` | `https://api.openai.com` | Custom OpenAI API URL endpoint |
### Run Modes
Headroom proxy has two explicit run modes:
- `token` mode: prioritize token reduction. Prior history may be rewritten when that improves compression.
- `cache` mode: prioritize provider prefix cache stability. Prior turns are frozen; only the newest turn is mutable.
Set via CLI or env:
```bash
headroom proxy --mode token
HEADROOM_MODE=cache headroom proxy
```
Legacy values (`token_headroom`, `cost_savings`) are still accepted as aliases.
### Context Management Options
| Option | Default | Description |

2
headroom/ccr/mcp_server.py
View file

@ -87,7 +87,7 @@ def _format_session_summary(summary: dict[str, Any], local_stats: dict[str, Any]
lines.append("Headroom Session Summary")
lines.append("=" * 40)
mode = summary.get("mode", "token_headroom")
mode = summary.get("mode", "token")
api_reqs = summary.get("api_requests", 0)
model = summary.get("primary_model", "unknown")
lines.append(f"Mode: {mode} | {api_reqs} API requests | {model}")

24
headroom/cli/proxy.py
View file

@ -4,6 +4,8 @@ import os
import click
from headroom.proxy.modes import PROXY_MODE_TOKEN, normalize_proxy_mode
from .main import main
@ -13,8 +15,22 @@ from .main import main
@click.option(
"--mode",
default=None,
type=click.Choice(["cost_savings", "token_headroom"]),
help="Optimization mode: token_headroom (compress for session extension) or cost_savings (preserve prefix cache). Default: token_headroom. Env: HEADROOM_MODE",
type=click.Choice(
[
"token",
"cache",
"token_mode",
"cache_mode",
"token_savings",
"cost_savings",
"token_headroom",
]
),
help=(
"Optimization mode: token (prioritize compression) or cache "
"(freeze prior turns for prefix-cache stability). "
"Legacy aliases are accepted. Default: token. Env: HEADROOM_MODE"
),
)
@click.option("--no-optimize", is_flag=True, help="Disable optimization (passthrough mode)")
@click.option("--no-cache", is_flag=True, help="Disable semantic caching")
@ -195,7 +211,9 @@ def proxy(
effective_anyllm_provider = os.environ.get("HEADROOM_ANYLLM_PROVIDER") or anyllm_provider
# Resolve mode: CLI flag > env var > default
effective_mode: str = mode or os.environ.get("HEADROOM_MODE") or "token_headroom"
effective_mode: str = normalize_proxy_mode(
mode or os.environ.get("HEADROOM_MODE") or PROXY_MODE_TOKEN
)
# Telemetry opt-out: --no-telemetry flag sets the env var
if no_telemetry:

8
headroom/proxy/cost.py
View file

@ -13,6 +13,8 @@ from collections import deque
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any
from headroom.proxy.modes import PROXY_MODE_CACHE
if TYPE_CHECKING:
from headroom.proxy.prometheus_metrics import PrometheusMetrics
@ -314,10 +316,10 @@ def build_session_summary(
},
}
# Add tip if token_headroom mode would help
if proxy.config.mode == "cost_savings" and uncompressed_reasons["prefix_frozen"] > 10:
# Add tip if token mode would help
if proxy.config.mode == PROXY_MODE_CACHE and uncompressed_reasons["prefix_frozen"] > 10:
summary["tip"] = (
"Most requests are prefix-frozen. Set HEADROOM_MODE=token_headroom "
"Most requests are prefix-frozen. Set HEADROOM_MODE=token "
"to compress frozen messages and extend your session by ~25-35%."
)

84
headroom/proxy/handlers/anthropic.py
View file

@ -6,6 +6,7 @@ Contains all Anthropic Messages API handlers including batch operations.
from __future__ import annotations
import asyncio
import copy
import json
import logging
import time
@ -129,6 +130,31 @@ class AnthropicHandlerMixin:
return max(base_frozen_count, final_idx)
return len(messages)
@staticmethod
def _restore_frozen_prefix(
original_messages: list[dict[str, Any]],
candidate_messages: list[dict[str, Any]],
*,
frozen_message_count: int,
) -> tuple[list[dict[str, Any]], int]:
"""Force frozen prefix bytes to match the original request exactly."""
if frozen_message_count <= 0 or not original_messages:
return candidate_messages, 0
frozen = min(frozen_message_count, len(original_messages))
restored = list(candidate_messages)
# Defensive: if a transform dropped prefix messages, restore them.
if len(restored) < frozen:
return list(original_messages[:frozen]) + restored, frozen
changed = 0
for idx in range(frozen):
if restored[idx] != original_messages[idx]:
restored[idx] = original_messages[idx]
changed += 1
return restored, changed
async def handle_anthropic_messages(
self,
request: Request,
@ -143,8 +169,10 @@ class AnthropicHandlerMixin:
from headroom.proxy.helpers import (
MAX_MESSAGE_ARRAY_LENGTH,
MAX_REQUEST_BODY_SIZE,
_get_image_compressor,
_read_request_json,
)
from headroom.proxy.modes import is_cache_mode, is_token_mode
from headroom.proxy.models import RequestLog
from headroom.tokenizers import get_tokenizer
from headroom.utils import extract_user_query
@ -182,6 +210,7 @@ class AnthropicHandlerMixin:
)
model = body.get("model", "unknown")
messages = body.get("messages", [])
original_client_messages = copy.deepcopy(messages)
# Validate message array size
if len(messages) > MAX_MESSAGE_ARRAY_LENGTH:
@ -306,21 +335,25 @@ class AnthropicHandlerMixin:
session_id = self.session_tracker_store.compute_session_id(request, model, messages)
prefix_tracker = self.session_tracker_store.get_or_create(session_id, "anthropic")
frozen_message_count = prefix_tracker.get_frozen_message_count()
frozen_message_count = self._strict_previous_turn_frozen_count(
messages,
frozen_message_count,
)
if is_cache_mode(self.config.mode):
frozen_message_count = self._strict_previous_turn_frozen_count(
original_client_messages,
frozen_message_count,
)
# Image compression (cache-safe): only compress latest non-frozen user turn.
# Rewriting historical image bytes can invalidate Anthropic prompt caches.
if self.config.image_optimize and messages and not _bypass:
compressor = _get_image_compressor()
if compressor and compressor.has_images(messages):
messages = self._compress_latest_user_turn_images_cache_safe(
messages,
frozen_message_count=frozen_message_count,
compressor=compressor,
)
if is_cache_mode(self.config.mode):
messages = self._compress_latest_user_turn_images_cache_safe(
messages,
frozen_message_count=frozen_message_count,
compressor=compressor,
)
else:
messages = compressor.compress(messages, provider="anthropic")
if compressor.last_result:
logger.info(
f"Image compression: {compressor.last_result.technique.value} "
@ -343,7 +376,7 @@ class AnthropicHandlerMixin:
else None
)
if self.config.mode == "token_headroom":
if is_token_mode(self.config.mode):
comp_cache = self._get_compression_cache(session_id)
# Zone 1: Swap cached compressed versions into working copy
@ -353,10 +386,6 @@ class AnthropicHandlerMixin:
# Safety: never freeze beyond provider-confirmed cached prefix.
cache_frozen_count = comp_cache.compute_frozen_count(messages)
frozen_message_count = min(frozen_message_count, cache_frozen_count)
frozen_message_count = self._strict_previous_turn_frozen_count(
messages,
frozen_message_count,
)
result = await asyncio.wait_for(
asyncio.to_thread(
@ -614,6 +643,18 @@ class AnthropicHandlerMixin:
# The echo changes every turn, invalidating the cached prefix.
# To re-enable, uncomment and set query_echo_enabled on ProxyConfig.
if is_cache_mode(self.config.mode):
optimized_messages, restored_count = self._restore_frozen_prefix(
original_client_messages,
optimized_messages,
frozen_message_count=frozen_message_count,
)
if restored_count > 0:
logger.warning(
f"[{request_id}] Restored {restored_count} frozen prefix message(s) "
"to preserve cache stability"
)
# Update body
body["messages"] = optimized_messages
if tools is not None:
@ -1151,6 +1192,7 @@ class AnthropicHandlerMixin:
from fastapi.responses import JSONResponse, Response
from headroom.ccr import CCRToolInjector
from headroom.proxy.modes import is_cache_mode
from headroom.proxy.helpers import MAX_REQUEST_BODY_SIZE, _read_request_json
from headroom.utils import extract_user_query
@ -1220,6 +1262,7 @@ class AnthropicHandlerMixin:
if canonical_tools is not None:
canonical_params["tools"] = self._sort_tools_deterministically(canonical_tools)
messages = params.get("messages", [])
original_messages = copy.deepcopy(messages)
model = params.get("model", "unknown")
if not messages or not self.config.optimize:
@ -1235,15 +1278,26 @@ class AnthropicHandlerMixin:
# Apply optimization
try:
context_limit = self.anthropic_provider.get_context_limit(model)
frozen_message_count = (
self._strict_previous_turn_frozen_count(original_messages, 0)
if is_cache_mode(self.config.mode)
else 0
)
result = self.anthropic_pipeline.apply(
messages=messages,
model=model,
model_limit=context_limit,
context=extract_user_query(messages),
frozen_message_count=self._strict_previous_turn_frozen_count(messages, 0),
frozen_message_count=frozen_message_count,
)
optimized_messages = result.messages
if is_cache_mode(self.config.mode):
optimized_messages, _ = self._restore_frozen_prefix(
original_messages,
optimized_messages,
frozen_message_count=frozen_message_count,
)
for k, v in result.timing.items():
pipeline_timing[k] = pipeline_timing.get(k, 0.0) + v
# Use pipeline's token counts for consistency with pipeline logs

5
headroom/proxy/handlers/openai.py
View file

@ -41,6 +41,7 @@ class OpenAIHandlerMixin:
MAX_REQUEST_BODY_SIZE,
_read_request_json,
)
from headroom.proxy.modes import is_token_mode
from headroom.tokenizers import get_tokenizer
from headroom.utils import extract_user_query
@ -191,7 +192,7 @@ class OpenAIHandlerMixin:
try:
context_limit = self.openai_provider.get_context_limit(model)
if self.config.mode == "token_headroom":
if is_token_mode(self.config.mode):
comp_cache = self._get_compression_cache(openai_session_id)
# Zone 1: Swap cached compressed versions
@ -217,7 +218,7 @@ class OpenAIHandlerMixin:
if result.messages != working_messages:
comp_cache.update_from_result(messages, result.messages)
# Always use pipeline result in token_headroom mode
# Always use pipeline result in token mode
optimized_messages = result.messages
transforms_applied = result.transforms_applied
pipeline_timing = result.timing

5
headroom/proxy/models.py
View file

@ -86,8 +86,9 @@ class ProxyConfig:
bedrock_profile: str | None = None
anyllm_provider: str = "openai"
# Optimization mode
mode: str = "token_headroom"
# Optimization mode: "token" (rewrite for max compression) or
# "cache" (freeze prior turns for prefix-cache stability).
mode: str = "token"
# Optimization
optimize: bool = True

51
headroom/proxy/modes.py Normal file
View file

@ -0,0 +1,51 @@
"""Proxy run mode helpers.
Canonical modes:
- token: prioritize compression (history may be rewritten for max savings)
- cache: prioritize provider prefix cache stability (freeze prior turns)
"""
from __future__ import annotations
import logging
logger = logging.getLogger("headroom.proxy")
PROXY_MODE_TOKEN = "token"
PROXY_MODE_CACHE = "cache"
_MODE_ALIASES = {
"token": PROXY_MODE_TOKEN,
"token_mode": PROXY_MODE_TOKEN,
"token_savings": PROXY_MODE_TOKEN,
"token_headroom": PROXY_MODE_TOKEN,
"cache": PROXY_MODE_CACHE,
"cache_mode": PROXY_MODE_CACHE,
"cost_savings": PROXY_MODE_CACHE,
}
def normalize_proxy_mode(mode: str | None, *, default: str = PROXY_MODE_TOKEN) -> str:
"""Normalize a user-provided proxy mode to canonical token/cache values."""
key = (mode or "").strip().lower()
if not key:
return default
normalized = _MODE_ALIASES.get(key)
if normalized is None:
logger.warning("Unknown HEADROOM_MODE '%s', falling back to '%s'", mode, default)
return default
if key != normalized:
logger.info("HEADROOM_MODE alias '%s' normalized to '%s'", mode, normalized)
return normalized
def is_token_mode(mode: str | None) -> bool:
"""Return True when mode resolves to token mode."""
return normalize_proxy_mode(mode) == PROXY_MODE_TOKEN
def is_cache_mode(mode: str | None) -> bool:
"""Return True when mode resolves to cache mode."""
return normalize_proxy_mode(mode) == PROXY_MODE_CACHE

34
headroom/proxy/server.py
View file

@ -102,6 +102,12 @@ from headroom.proxy.helpers import (
_read_request_json, # noqa: F401
_setup_file_logging, # noqa: F401
)
from headroom.proxy.modes import (
PROXY_MODE_CACHE,
PROXY_MODE_TOKEN,
is_token_mode,
normalize_proxy_mode,
)
from headroom.proxy.memory_handler import MemoryConfig, MemoryHandler
# Data models (extracted to headroom/proxy/models.py for maintainability)
@ -162,6 +168,7 @@ class HeadroomProxy(
def __init__(self, config: ProxyConfig):
self.config = config
self.config.mode = normalize_proxy_mode(self.config.mode)
# Override ANTHROPIC_API_URL with config if set
# Strip trailing /v1 or /v1/ to avoid double-path (e.g., .../v1/v1/models)
@ -224,8 +231,8 @@ 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":
# Token mode: allow compression of older excluded-tool results
if is_token_mode(config.mode):
router_config.protect_recent_reads_fraction = 0.3
transforms = [
CacheAligner(CacheAlignerConfig(enabled=False)),
@ -303,7 +310,7 @@ class HeadroomProxy(
)
)
# Compression cache store for token_headroom mode (session-scoped)
# Compression cache store for token mode (session-scoped)
self._compression_caches: dict[str, CompressionCache] = {}
self.logger = (
@ -514,17 +521,16 @@ 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 = "token_headroom"
self.config.mode = normalize_proxy_mode(self.config.mode)
logger.info(f"Mode: {self.config.mode}")
if self.config.mode == "token_headroom":
if self.config.mode == PROXY_MODE_TOKEN:
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")
if self.config.mode == PROXY_MODE_CACHE:
logger.info(" Prefix freeze: strict (all prior turns immutable)")
logger.info(" Mutations: latest turn only")
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(
@ -1019,9 +1025,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
except Exception:
logger.warning("Failed to log /stats summary payload")
# Compression cache stats (token_headroom mode)
# Compression cache stats (token mode)
compression_cache_stats: dict = {}
if proxy.config.mode == "token_headroom" and proxy._compression_caches:
if proxy.config.mode == PROXY_MODE_TOKEN and proxy._compression_caches:
total_entries = 0
total_hits = 0
total_misses = 0
@ -1033,7 +1039,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
total_misses += s.get("misses", 0)
total_tokens_saved += s.get("total_tokens_saved", 0)
compression_cache_stats = {
"mode": "token_headroom",
"mode": PROXY_MODE_TOKEN,
"active_sessions": len(proxy._compression_caches),
"total_entries": total_entries,
"total_hits": total_hits,
@ -1042,7 +1048,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
"total_tokens_saved": total_tokens_saved,
}
else:
compression_cache_stats = {"mode": "token_headroom"}
compression_cache_stats = {"mode": PROXY_MODE_TOKEN}
# Build unified savings summary (all layers)
compression_tokens = m.tokens_saved_total
@ -2327,7 +2333,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", "token_headroom"),
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_TOKEN)),
)
# Get worker and concurrency settings

149
tests/test_proxy_anthropic_cache_stability.py
View file

@ -253,12 +253,12 @@ def test_append_context_does_not_touch_previous_turns_if_last_message_not_user()
assert result[0]["content"] == "previous user turn"
def test_token_headroom_freeze_is_capped_by_prefix_tracker() -> None:
def test_token_mode_freeze_is_capped_by_prefix_tracker() -> None:
captured = {}
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "token_headroom"
proxy.config.mode = "token"
proxy.config.image_optimize = False
fake_tracker = _FakePrefixTracker(frozen_count=1)
@ -453,7 +453,7 @@ def test_previous_turns_always_frozen_only_final_turn_mutable() -> None:
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "cost_savings"
proxy.config.mode = "cache"
proxy.config.image_optimize = False
fake_tracker = _FakePrefixTracker(frozen_count=0)
@ -515,6 +515,7 @@ def test_batch_optimization_freezes_previous_turns_only() -> None:
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "cache"
proxy.config.image_optimize = False
proxy.config.ccr_inject_tool = False
@ -566,3 +567,145 @@ def test_batch_optimization_freezes_previous_turns_only() -> None:
assert response.status_code == 200
assert captured["frozen_message_count"] == 2
def test_token_mode_does_not_force_freeze_all_previous_turns() -> None:
captured = {}
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "token"
proxy.config.image_optimize = False
fake_tracker = _FakePrefixTracker(frozen_count=0)
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: "stable-session"
proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker
class _FakeCompressionCache:
def apply_cached(self, messages): # noqa: ANN001
return messages
def compute_frozen_count(self, messages): # noqa: ANN001
return 0
def update_from_result(self, originals, compressed): # noqa: ANN001
return None
proxy._get_compression_cache = lambda session_id: _FakeCompressionCache()
def _fake_apply(**kwargs):
captured["frozen_message_count"] = kwargs.get("frozen_message_count")
return SimpleNamespace(
messages=kwargs["messages"],
transforms_applied=[],
timing={},
tokens_before=70,
tokens_after=70,
waste_signals=None,
)
proxy.anthropic_pipeline.apply = _fake_apply
async def _fake_retry(method, url, headers, body, stream=False): # noqa: ANN001
return httpx.Response(
200,
json={
"id": "msg_tok_1",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "ok"}],
"usage": {
"input_tokens": 70,
"output_tokens": 3,
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
},
},
)
proxy._retry_request = _fake_retry
response = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"messages": [
{"role": "user", "content": "turn1"},
{"role": "assistant", "content": "turn1-assistant"},
{"role": "user", "content": "current turn"},
],
},
)
assert response.status_code == 200
assert captured["frozen_message_count"] == 0
def test_cache_mode_restores_frozen_prefix_if_transform_mutates_history() -> None:
captured = {}
with _make_proxy_client() as client:
proxy = client.app.state.proxy
proxy.config.optimize = True
proxy.config.mode = "cache"
proxy.config.image_optimize = False
fake_tracker = _FakePrefixTracker(frozen_count=0)
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: "stable-session"
proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker
original_messages = [
{"role": "user", "content": "turn1"},
{"role": "assistant", "content": "turn1-assistant"},
{"role": "user", "content": "current turn"},
]
def _fake_apply(**kwargs):
mutated = list(kwargs["messages"])
mutated[0] = {**mutated[0], "content": "MUTATED_PREFIX"}
return SimpleNamespace(
messages=mutated,
transforms_applied=["fake:mutated"],
timing={},
tokens_before=80,
tokens_after=70,
waste_signals=None,
)
proxy.anthropic_pipeline.apply = _fake_apply
async def _fake_retry(method, url, headers, body, stream=False): # noqa: ANN001
captured["body"] = body
return httpx.Response(
200,
json={
"id": "msg_cache_1",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "ok"}],
"usage": {
"input_tokens": 70,
"output_tokens": 3,
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
},
},
)
proxy._retry_request = _fake_retry
response = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"messages": original_messages,
},
)
assert response.status_code == 200
sent_messages = captured["body"]["messages"]
assert sent_messages[0] == original_messages[0]
assert sent_messages[1] == original_messages[1]

19
tests/test_proxy_mode_benchmark.py Normal file
View file

@ -0,0 +1,19 @@
"""Tests for local token/cache mode benchmark harness."""
from benchmarks.proxy_mode_benchmark import run_local_benchmark
def test_local_mode_benchmark_shows_compression_and_cache_tradeoff() -> None:
results = run_local_benchmark(turns=6)
baseline = results["baseline"]
token = results["token"]
cache = results["cache"]
assert token.total_tokens_saved > 0
assert cache.total_tokens_saved > 0
assert token.total_sent_tokens < baseline.total_sent_tokens
assert cache.total_sent_tokens < baseline.total_sent_tokens
# Cache mode should preserve prefix better than token mode.
assert cache.total_cache_read_tokens >= token.total_cache_read_tokens

30
tests/test_proxy_modes.py Normal file
View file

@ -0,0 +1,30 @@
"""Tests for proxy token/cache mode normalization."""
from headroom.proxy.modes import (
PROXY_MODE_CACHE,
PROXY_MODE_TOKEN,
is_cache_mode,
is_token_mode,
normalize_proxy_mode,
)
def test_proxy_mode_normalizes_canonical_values() -> None:
assert normalize_proxy_mode("token") == PROXY_MODE_TOKEN
assert normalize_proxy_mode("cache") == PROXY_MODE_CACHE
def test_proxy_mode_normalizes_legacy_aliases() -> None:
assert normalize_proxy_mode("token_headroom") == PROXY_MODE_TOKEN
assert normalize_proxy_mode("token_savings") == PROXY_MODE_TOKEN
assert normalize_proxy_mode("cost_savings") == PROXY_MODE_CACHE
assert normalize_proxy_mode("cache_mode") == PROXY_MODE_CACHE
def test_proxy_mode_invalid_falls_back_to_default() -> None:
assert normalize_proxy_mode("wat", default=PROXY_MODE_CACHE) == PROXY_MODE_CACHE
def test_proxy_mode_predicates() -> None:
assert is_token_mode("token_headroom") is True
assert is_cache_mode("cost_savings") is True

2
tests/test_token_headroom_mode.py
View file

@ -1,4 +1,4 @@
"""Integration tests for token_headroom mode.
"""Integration tests for token mode (legacy token_headroom behavior).
Tests the CompressionCache working across simulated multi-turn conversations,
verifying the critical invariants: no message injection, correct frozen counts,