From 54419ad8b851299f7cc387cdc49e4457cd199b8e Mon Sep 17 00:00:00 2001 From: JerrettDavis Date: Sat, 4 Apr 2026 14:32:07 -0500 Subject: [PATCH] Rebrand proxy modes to token/cache and harden cache-mode stability --- benchmarks/proxy_mode_benchmark.py | 282 ++++++++++++++++++ docs/benchmarks.md | 3 + docs/proxy.md | 17 ++ headroom/ccr/mcp_server.py | 2 +- headroom/cli/proxy.py | 24 +- headroom/proxy/cost.py | 8 +- headroom/proxy/handlers/anthropic.py | 84 +++++- headroom/proxy/handlers/openai.py | 5 +- headroom/proxy/models.py | 5 +- headroom/proxy/modes.py | 51 ++++ headroom/proxy/server.py | 34 ++- tests/test_proxy_anthropic_cache_stability.py | 149 ++++++++- tests/test_proxy_mode_benchmark.py | 19 ++ tests/test_proxy_modes.py | 30 ++ tests/test_token_headroom_mode.py | 2 +- 15 files changed, 671 insertions(+), 44 deletions(-) create mode 100644 benchmarks/proxy_mode_benchmark.py create mode 100644 headroom/proxy/modes.py create mode 100644 tests/test_proxy_mode_benchmark.py create mode 100644 tests/test_proxy_modes.py diff --git a/benchmarks/proxy_mode_benchmark.py b/benchmarks/proxy_mode_benchmark.py new file mode 100644 index 000000000..ce1961bb7 --- /dev/null +++ b/benchmarks/proxy_mode_benchmark.py @@ -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() diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 1288544ca..4db1d0303 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -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 ``` diff --git a/docs/proxy.md b/docs/proxy.md index 90b7d5920..125caf8b5 100644 --- a/docs/proxy.md +++ b/docs/proxy.md @@ -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 | diff --git a/headroom/ccr/mcp_server.py b/headroom/ccr/mcp_server.py index 690e3b1de..b9bbb684a 100644 --- a/headroom/ccr/mcp_server.py +++ b/headroom/ccr/mcp_server.py @@ -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}") diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 44fc7c370..306fb06ec 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -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: diff --git a/headroom/proxy/cost.py b/headroom/proxy/cost.py index a1d341a09..d2f1859ea 100644 --- a/headroom/proxy/cost.py +++ b/headroom/proxy/cost.py @@ -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%." ) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 515516901..cd733b437 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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 diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 949e8be92..3f9b70d38 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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 diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index a140fc0f9..ff594f4ae 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -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 diff --git a/headroom/proxy/modes.py b/headroom/proxy/modes.py new file mode 100644 index 000000000..0dab93a33 --- /dev/null +++ b/headroom/proxy/modes.py @@ -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 diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index deceef9a7..08dd7d1da 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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 diff --git a/tests/test_proxy_anthropic_cache_stability.py b/tests/test_proxy_anthropic_cache_stability.py index ed55a45ed..6481dc4f5 100644 --- a/tests/test_proxy_anthropic_cache_stability.py +++ b/tests/test_proxy_anthropic_cache_stability.py @@ -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] diff --git a/tests/test_proxy_mode_benchmark.py b/tests/test_proxy_mode_benchmark.py new file mode 100644 index 000000000..b2fd2dd58 --- /dev/null +++ b/tests/test_proxy_mode_benchmark.py @@ -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 diff --git a/tests/test_proxy_modes.py b/tests/test_proxy_modes.py new file mode 100644 index 000000000..a3fa06544 --- /dev/null +++ b/tests/test_proxy_modes.py @@ -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 diff --git a/tests/test_token_headroom_mode.py b/tests/test_token_headroom_mode.py index 7c8b818f1..ef63b50fb 100644 --- a/tests/test_token_headroom_mode.py +++ b/tests/test_token_headroom_mode.py @@ -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,