From 1a84784a82471bfaa17b51043fdeb28e93d01391 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Fri, 31 Jul 2026 21:47:11 -0400 Subject: [PATCH 1/5] feat(proxy): add an opt-in pre-forward context budget guard (#2649) --- docs/content/docs/configuration.mdx | 15 + headroom/providers/anthropic.py | 32 +- headroom/proxy/context_budget_policy.py | 127 ++++ headroom/proxy/handlers/anthropic.py | 89 +++ tests/test_proxy_context_budget.py | 809 ++++++++++++++++++++++++ wiki/configuration.md | 15 + 6 files changed, 1086 insertions(+), 1 deletion(-) create mode 100644 headroom/proxy/context_budget_policy.py create mode 100644 tests/test_proxy_context_budget.py diff --git a/docs/content/docs/configuration.mdx b/docs/content/docs/configuration.mdx index c4bf1bc8f..13b8c28ac 100644 --- a/docs/content/docs/configuration.mdx +++ b/docs/content/docs/configuration.mdx @@ -276,6 +276,8 @@ headroom proxy --learn --min-evidence 3 | `HEADROOM_TELEMETRY` | Set to `on` for **local-only** usage stats (powers your own `/stats` and dashboard; nothing is sent externally) | `off` | | `HEADROOM_STATELESS` | Set to `true` to disable filesystem writes | `false` | | `HEADROOM_MODEL_LIMITS` | Custom model config (JSON string or file path) | -- | +| `HEADROOM_CONTEXT_LIMIT_MODE` | Context budget guard mode. `observe` (default) logs over-limit finalized requests and forwards them; `reject` returns a local 400 before any upstream attempt. Effective only when a model limit is declared in `HEADROOM_MODEL_LIMITS` or `models.json`. | `observe` | +| `HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN` | Token reserve subtracted from the declared model context limit before comparing against the finalized request token count (in addition to `max_tokens`). A non-negative integer. | `0` | | `HEADROOM_BASE_URL` | Base URL of the Headroom proxy (TypeScript SDK) | `http://localhost:8787` | | `HEADROOM_API_KEY` | Optional API key for authenticated Headroom endpoints (TypeScript SDK) | -- | | `HEADROOM_CONFIG_DIR` | Canonical config (read-mostly) root. Derives `models.json` and per-plugin config paths when set. | `~/.headroom/config` | @@ -446,6 +448,19 @@ Settings are resolved in this order (later overrides earlier): 3. `HEADROOM_MODEL_LIMITS` environment variable 4. SDK constructor arguments +### Context budget guard + +When a model's context limit is declared in `HEADROOM_MODEL_LIMITS` or `models.json`, the proxy evaluates every finalized Anthropic `/v1/messages` request against that limit before forwarding. The threshold is `declared_limit - max(HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN, max_tokens)`. + +In the default `observe` mode the proxy logs an over-limit request at WARNING and forwards it unchanged. Set `HEADROOM_CONTEXT_LIMIT_MODE=reject` to have the proxy return a local 400 before any upstream attempt, which stops the provider's own 400 from triggering a client retry storm. Unconfigured installs are unaffected; the guard is a no-op when no limit is declared for the model. + +```bash +export HEADROOM_MODEL_LIMITS='{"context_limits":{"step-router-v1":262144}}' +export HEADROOM_CONTEXT_LIMIT_MODE=reject +export HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN=4096 +headroom proxy +``` + ### Pattern-Based Inference Unknown models are automatically inferred from naming patterns: diff --git a/headroom/providers/anthropic.py b/headroom/providers/anthropic.py index 107fc53c5..3c9a49b4b 100644 --- a/headroom/providers/anthropic.py +++ b/headroom/providers/anthropic.py @@ -551,12 +551,16 @@ class AnthropicProvider(Provider): # Load from config file and env var custom_config = _load_custom_model_config() + # Capture operator-declared limits BEFORE they merge into the inference + # table so their provenance remains distinguishable (#2649). + self._operator_context_limits: dict[str, int] = dict(custom_config["context_limits"]) self._context_limits.update(custom_config["context_limits"]) self._pricing.update(custom_config["pricing"]) - # Explicit overrides take precedence + # Explicit overrides take precedence, and are operator-supplied too. if context_limits: self._context_limits.update(context_limits) + self._operator_context_limits.update(context_limits) @property def name(self) -> str: @@ -640,6 +644,32 @@ class AnthropicProvider(Provider): self._context_limits[model] = limit return limit + def get_operator_context_limit(self, model: str) -> int | None: + """Return the operator-declared context limit for *model*, or None. + + Exact-match only: tries the raw model id first, then the sanitized id. + Returns None when the operator declared nothing for this model — never + falls back to inferred or default values. No caching side effects. + """ + if model in self._operator_context_limits: + return self._operator_context_limits[model] + sanitized = sanitize_anthropic_model_id(model) + if sanitized in self._operator_context_limits: + return self._operator_context_limits[sanitized] + return None + + def has_raw_operator_context_limit(self, model: str) -> bool: + """Whether the operator declared a limit against this exact, unsanitized id. + + `wrap --1m` appends a ``[1m]`` suffix so Claude Code emits the + ``context-1m`` beta, and :func:`sanitize_anthropic_model_id` strips it. + A caller deciding whether the declared window covers the 1M variant + needs to know the declaration matched the suffixed id itself, which + :meth:`get_operator_context_limit` cannot express because it falls back + to the sanitized id. + """ + return model in self._operator_context_limits + def _warn_unknown_model(self, model: str, limit: int, reason: str) -> None: """Warn about unknown model (once per model).""" global _UNKNOWN_MODEL_WARNINGS diff --git a/headroom/proxy/context_budget_policy.py b/headroom/proxy/context_budget_policy.py new file mode 100644 index 000000000..815742a5e --- /dev/null +++ b/headroom/proxy/context_budget_policy.py @@ -0,0 +1,127 @@ +"""Pre-forward context budget policy for the Anthropic messages handler. + +Pure-logic module: no token counting, no message inspection, no I/O beyond +os.environ in the two resolver helpers. Does not import FastAPI, the handler, +or the provider. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class BudgetDecision: + mode: str + counted_tokens: int + declared_limit: int | None + reserve: int + threshold: int + overage: int + should_reject: bool + reason: str + # Retained for diagnostic access; not used in evaluation logic. + _extra: dict = field(default_factory=dict, compare=False) + + +def resolve_mode() -> str: + """Read HEADROOM_CONTEXT_LIMIT_MODE; default 'observe'. + + Accepts 'observe' and 'reject' (case-insensitive). + Raises ValueError naming the bad value and the accepted set. + """ + raw = os.environ.get("HEADROOM_CONTEXT_LIMIT_MODE", "observe").strip().lower() + if raw not in ("observe", "reject"): + raise ValueError( + f"HEADROOM_CONTEXT_LIMIT_MODE={raw!r} is not accepted; " + "accepted values: 'observe', 'reject'" + ) + return raw + + +def resolve_safety_margin() -> int: + """Read HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN; default 0. + + Raises ValueError on non-integer or negative value. + """ + raw = os.environ.get("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0").strip() + try: + val = int(raw) + except ValueError as exc: + raise ValueError(f"HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN={raw!r} is not an integer") from exc + if val < 0: + raise ValueError(f"HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN={val} must be >= 0") + return val + + +def evaluate( + *, + counted_tokens: int, + declared_limit: int | None, + max_output_tokens: int, + mode: str, + safety_margin: int, +) -> BudgetDecision: + """Evaluate the context budget for a finalized request. + + Branches (in order): + 1. declared_limit is None -> reason='no_declared_limit', should_reject=False + 2. threshold <= 0 -> reason='degenerate_threshold', should_reject=False + 3. counted_tokens <= threshold -> reason='under_threshold', should_reject=False + 4. else -> reason='over_threshold', should_reject=(mode=='reject') + + reserve = max(safety_margin, max_output_tokens) + threshold = declared_limit - reserve + overage = max(0, counted_tokens - threshold) + """ + if declared_limit is None: + return BudgetDecision( + mode=mode, + counted_tokens=counted_tokens, + declared_limit=None, + reserve=0, + threshold=0, + overage=0, + should_reject=False, + reason="no_declared_limit", + ) + + reserve = max(safety_margin, max_output_tokens) + threshold = declared_limit - reserve + + if threshold <= 0: + return BudgetDecision( + mode=mode, + counted_tokens=counted_tokens, + declared_limit=declared_limit, + reserve=reserve, + threshold=threshold, + overage=0, + should_reject=False, + reason="degenerate_threshold", + ) + + if counted_tokens <= threshold: + return BudgetDecision( + mode=mode, + counted_tokens=counted_tokens, + declared_limit=declared_limit, + reserve=reserve, + threshold=threshold, + overage=0, + should_reject=False, + reason="under_threshold", + ) + + overage = counted_tokens - threshold + return BudgetDecision( + mode=mode, + counted_tokens=counted_tokens, + declared_limit=declared_limit, + reserve=reserve, + threshold=threshold, + overage=overage, + should_reject=(mode == "reject"), + reason="over_threshold", + ) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index e255ab886..f4a9fc9d5 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2586,6 +2586,95 @@ class AnthropicHandlerMixin: ): headers["anthropic-beta"] = _client_beta_value + # Context budget guard (#2649). Applied after all input shaping + # so the count covers the finalized body. Only operator-declared + # limits are accepted; inferred defaults never drive a refusal. + # Fail-open: any exception inside the block forwards the request + # unchanged. + try: + from headroom.proxy.context_budget_policy import ( + evaluate as _cbp_evaluate, + ) + from headroom.proxy.context_budget_policy import ( + resolve_mode as _cbp_resolve_mode, + ) + from headroom.proxy.context_budget_policy import ( + resolve_safety_margin as _cbp_resolve_safety_margin, + ) + + _cbp_should_guard = True + # Short-circuit: bypass header skips all Headroom behaviour. + if _bypass: + _cbp_should_guard = False + + if _cbp_should_guard: + _cbp_mode = _cbp_resolve_mode() + _cbp_margin = _cbp_resolve_safety_margin() + _cbp_declared = self.anthropic_provider.get_operator_context_limit(model) + _cbp_max_out = int(body.get("max_tokens") or 0) + + # Degrade to observe when the outbound anthropic-beta carries + # context-1m and no raw-id declaration exists: the effective + # window is then unknowable because sanitize_anthropic_model_id + # strips the [1m] suffix and the declared limit would apply to + # the base model, not the 1M variant. + _cbp_outbound_beta = headers.get("anthropic-beta", "") + _cbp_has_context1m = "context-1m" in _cbp_outbound_beta + if _cbp_has_context1m and not ( + self.anthropic_provider.has_raw_operator_context_limit(model) + ): + _cbp_mode = "observe" + + _cbp_decision = _cbp_evaluate( + counted_tokens=optimized_tokens, + declared_limit=_cbp_declared, + max_output_tokens=_cbp_max_out, + mode=_cbp_mode, + safety_margin=_cbp_margin, + ) + + if _cbp_decision.reason == "over_threshold": + logger.warning( + "[%s] context_budget_guard: model=%s declared_limit=%d " + "reserve=%d threshold=%d counted=%d overage=%d mode=%s", + request_id, + model, + _cbp_decision.declared_limit, + _cbp_decision.reserve, + _cbp_decision.threshold, + _cbp_decision.counted_tokens, + _cbp_decision.overage, + _cbp_mode, + ) + + if _cbp_decision.should_reject: + await _finalize_pre_upstream() + return JSONResponse( + status_code=400, + content={ + "type": "error", + "error": { + "type": "invalid_request_error", + "message": ( + f"Request exceeds the declared context limit for {model}: " + f"{_cbp_decision.counted_tokens} tokens counted, " + f"{_cbp_decision.threshold} available " + f"(declared {_cbp_decision.declared_limit}, " + f"reserve {_cbp_decision.reserve}). " + "Set HEADROOM_CONTEXT_LIMIT_MODE=observe to log only." + ), + }, + }, + ) + except Exception: + # Any failure in limit lookup, mode resolution, or evaluation + # forwards the request unchanged. + logger.debug( + "[%s] context_budget_guard: exception during evaluation, forwarding", + request_id, + exc_info=True, + ) + # Forward request - use Bedrock backend if configured, otherwise direct API if self.anthropic_backend is not None: # Route through Bedrock backend diff --git a/tests/test_proxy_context_budget.py b/tests/test_proxy_context_budget.py new file mode 100644 index 000000000..33d3c7cd7 --- /dev/null +++ b/tests/test_proxy_context_budget.py @@ -0,0 +1,809 @@ +"""Context budget guard tests for headroom target #2649. + +Test matrix follows the Required Proof Matrix in the execution prompt: + reproduction base-fails/head-passes contract + mode_* / variant_* every Fault Scope variant is classified + preservation unconfigured requests forward unchanged + negative_space at-threshold / observe-over / bypassed-over all forward + production_route real POST /v1/messages through create_app + contract_isolation policy module has no FastAPI/handler/provider imports +""" + +from __future__ import annotations + +import asyncio +import importlib +import json +import os +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import anyio +from fastapi import Request +from fastapi.testclient import TestClient + +# --------------------------------------------------------------------------- # +# Shared helpers # +# --------------------------------------------------------------------------- # + + +class _DummyTokenizer: + """Token counter that returns a configurable value.""" + + def __init__(self, count: int = 1) -> None: + self._count = count + + def count_messages(self, messages) -> int: + return self._count + + def count_text(self, text: str) -> int: + return self._count + + def count(self, messages) -> int: + return self._count + + +class _DummyMetrics: + def __init__(self) -> None: + self.stage_timings: list = [] + + async def record_request(self, **kwargs) -> None: + return None + + async def record_stage_timings(self, path, timings) -> None: + self.stage_timings.append((path, timings)) + + async def record_failed(self, **kwargs) -> None: + return None + + async def record_rate_limited(self, **kwargs) -> None: + return None + + def record_compression_failed(self, reason: str) -> None: + return None + + +def _stub_response(status: int = 200) -> Any: + body = json.dumps( + { + "id": "msg_test", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "ok"}], + "model": "step-router-v1", + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + ) + ns = SimpleNamespace() + ns.status_code = status + ns.headers = {"content-type": "application/json"} + ns.content = body.encode() + ns.text = body + ns.json = lambda: json.loads(body) + return ns + + +def _make_anthropic_provider(operator_limit: int | None = None) -> Any: + limits: dict[str, int] = {} + if operator_limit is not None: + limits["step-router-v1"] = operator_limit + + def _get_operator_limit(model: str) -> int | None: + if model in limits: + return limits[model] + from headroom.providers.anthropic import sanitize_anthropic_model_id + + sanitized = sanitize_anthropic_model_id(model) + return limits.get(sanitized) + + return SimpleNamespace( + get_context_limit=lambda model: 200_000, + get_operator_context_limit=_get_operator_limit, + has_raw_operator_context_limit=lambda model: model in limits, + _operator_context_limits=limits, + ) + + +class _DummyHandler: + """Minimal AnthropicHandlerMixin with configurable upstream stub.""" + + from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin + + ANTHROPIC_API_URL = "https://api.anthropic.com" + + def __init__( + self, + *, + operator_limit: int | None = None, + token_count: int = 1, + bypass: bool = False, + mode_header: str | None = None, + ) -> None: + from headroom.proxy.models import ProxyConfig + + # Inherit the mixin by embedding via __class__ trick is messy; + # subclass it inline. + self._operator_limit = operator_limit + self._token_count = token_count + self._bypass_header = bypass + self.upstream_calls: list[dict] = [] + + self.rate_limiter = None + self.metrics = _DummyMetrics() + self.config = ProxyConfig( + optimize=False, + image_optimize=False, + retry_max_attempts=1, + retry_base_delay_ms=1, + retry_max_delay_ms=1, + connect_timeout_seconds=10, + mode="token", + cache_enabled=False, + rate_limit_enabled=False, + fallback_enabled=False, + fallback_provider=None, + prefix_freeze_enabled=False, + memory_enabled=False, + ) + self.usage_reporter = None + self.anthropic_provider = _make_anthropic_provider(operator_limit) + self.anthropic_pipeline = SimpleNamespace(apply=MagicMock()) + self.anthropic_backend = None + self.cost_tracker = None + self.memory_handler = None + self.cache = None + self.security = None + self.ccr_context_tracker = None + self.ccr_injector = None + self.ccr_response_handler = None + self.ccr_feedback = None + self.ccr_batch_processor = None + self.ccr_mcp_server = None + self.traffic_learner = None + self.tool_injector = None + self.read_lifecycle_manager = None + self.logger = SimpleNamespace(log=lambda *a, **k: None) + self.request_logger = self.logger + self.usage_observer = None + self.image_compressor = None + self.session_tracker_store = SimpleNamespace( + compute_session_id=lambda *a, **k: "sess-budget-test", + get_or_create=lambda *a, **k: SimpleNamespace( + _cached_token_count=0, + get_frozen_message_count=lambda: 0, + get_last_original_messages=lambda: [], + get_last_forwarded_messages=lambda: [], + update_from_response=lambda *a, **k: None, + record_request=lambda *a, **k: None, + ), + resolve_tracker=lambda *a, **k: SimpleNamespace( + _cached_token_count=0, + get_frozen_message_count=lambda: 0, + get_last_original_messages=lambda: [], + get_last_forwarded_messages=lambda: [], + update_from_response=lambda *a, **k: None, + record_request=lambda *a, **k: None, + ), + ) + self.anthropic_pre_upstream_sem = None + self.anthropic_pre_upstream_concurrency = 0 + import concurrent.futures as _cf + import threading as _threading + + self._compression_executor = _cf.ThreadPoolExecutor(max_workers=2) + self.compression_max_workers = 2 + self._compression_in_flight = 0 + self._compression_in_flight_max = 0 + self._compression_leaked_threads = 0 + self._compression_metrics_lock = _threading.Lock() + self._background_compression_enabled = False + + async def _run_compression_in_executor(self, fn, *, timeout): + loop = asyncio.get_running_loop() + future = loop.run_in_executor(self._compression_executor, fn) + return await asyncio.wait_for(future, timeout=timeout) + + async def _record_request_outcome(self, outcome) -> None: + from headroom.proxy.outcome import emit_request_outcome + + await emit_request_outcome(self, outcome) + + async def _next_request_id(self) -> str: + return "req-budget-test" + + def _extract_tags(self, headers): + return {} + + async def _retry_request(self, method, url, headers, body, **_kwargs): + self.upstream_calls.append({"method": method, "url": url}) + return _stub_response() + + def _get_compression_cache(self, session_id): + return SimpleNamespace( + apply_cached=lambda m: m, + compute_frozen_count=lambda m: 0, + mark_stable_from_messages=lambda *a, **k: None, + should_defer_compression=lambda h: False, + mark_stable=lambda h: None, + content_hash=lambda c: "h", + update_from_result=lambda *a, **k: None, + _cache={}, + _stable_hashes=set(), + ) + + def _extract_anthropic_cache_ttl_metrics(self, usage): + return (0, 0) + + +def _make_handler_subclass() -> type: + """Return a concrete class inheriting both _DummyHandler and AnthropicHandlerMixin.""" + from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin + + class _BudgetHandler(_DummyHandler, AnthropicHandlerMixin): + pass + + return _BudgetHandler + + +def _build_request( + body: dict, + headers: dict[str, str] | None = None, +) -> Request: + h = {"authorization": "Bearer sk-ant-api-test"} + if headers: + h.update(headers) + payload = json.dumps(body).encode("utf-8") + + async def receive(): + return {"type": "http.request", "body": payload, "more_body": False} + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "https", + "path": "/v1/messages", + "raw_path": b"/v1/messages", + "query_string": b"", + "headers": [(k.lower().encode(), v.encode()) for k, v in h.items()], + "client": ("127.0.0.1", 12345), + "server": ("testserver", 443), + } + return Request(scope, receive) + + +# --------------------------------------------------------------------------- # +# Policy unit tests (contract_isolation) # +# --------------------------------------------------------------------------- # + + +def test_contract_isolation_no_fastapi_import(): + """The policy module must not import FastAPI, the handler, or the provider.""" + + src = importlib.import_module("headroom.proxy.context_budget_policy") + # Walk module's __dict__ for imported symbols that would indicate a bad dep. + forbidden = {"fastapi", "headroom.proxy.handlers", "headroom.providers"} + for mod_name in list(src.__dict__.keys()): + obj = src.__dict__[mod_name] + if hasattr(obj, "__module__") and obj.__module__ is not None: + for f in forbidden: + assert not obj.__module__.startswith(f), ( + f"Policy module imported {obj.__module__!r} (matches forbidden prefix {f!r})" + ) + + +def test_contract_isolation_standalone_evaluate_matches_handler(monkeypatch): + """evaluate() produces the same result whether called directly or via the handler.""" + from headroom.proxy.context_budget_policy import evaluate, resolve_mode, resolve_safety_margin + + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + + mode = resolve_mode() + margin = resolve_safety_margin() + # Standalone call + direct = evaluate( + counted_tokens=270_000, + declared_limit=262_144, + max_output_tokens=8_192, + mode=mode, + safety_margin=margin, + ) + assert direct.should_reject is True + assert direct.reason == "over_threshold" + + +# --------------------------------------------------------------------------- # +# Policy evaluate branches (mode_ / variant_) # +# --------------------------------------------------------------------------- # + + +def test_mode_no_declared_limit(): + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=300_000, + declared_limit=None, + max_output_tokens=8_192, + mode="reject", + safety_margin=0, + ) + assert d.reason == "no_declared_limit" + assert d.should_reject is False + + +def test_mode_observe_under_threshold(): + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=100_000, + declared_limit=262_144, + max_output_tokens=8_192, + mode="observe", + safety_margin=0, + ) + assert d.reason == "under_threshold" + assert d.should_reject is False + + +def test_mode_observe_over_threshold(): + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=260_000, + declared_limit=262_144, + max_output_tokens=8_192, + mode="observe", + safety_margin=0, + ) + assert d.reason == "over_threshold" + assert d.should_reject is False # observe => no rejection + + +def test_mode_reject_over_threshold(): + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=260_000, + declared_limit=262_144, + max_output_tokens=8_192, + mode="reject", + safety_margin=0, + ) + assert d.reason == "over_threshold" + assert d.should_reject is True + + +def test_variant_degenerate_threshold(): + """When max_output_tokens >= declared_limit the threshold is <= 0: no rejection.""" + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=100_000, + declared_limit=8_192, + max_output_tokens=8_192, # equal to declared_limit -> threshold=0 + mode="reject", + safety_margin=0, + ) + assert d.reason == "degenerate_threshold" + assert d.should_reject is False + + +def test_variant_degenerate_threshold_max_output_exceeds(): + """max_output_tokens > declared_limit is also degenerate.""" + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=100_000, + declared_limit=8_192, + max_output_tokens=16_000, + mode="reject", + safety_margin=0, + ) + assert d.reason == "degenerate_threshold" + assert d.should_reject is False + + +def test_variant_safety_margin_adds_reserve(): + from headroom.proxy.context_budget_policy import evaluate + + # declared_limit=262144, max_output_tokens=8192, safety_margin=10000 + # reserve = max(10000, 8192) = 10000 + # threshold = 262144 - 10000 = 252144 + # count = 253000 > 252144 => over + d = evaluate( + counted_tokens=253_000, + declared_limit=262_144, + max_output_tokens=8_192, + mode="reject", + safety_margin=10_000, + ) + assert d.reason == "over_threshold" + assert d.reserve == 10_000 + assert d.threshold == 252_144 + assert d.overage == 253_000 - 252_144 + + +def test_variant_bypass_not_evaluated(monkeypatch): + """The bypass variant skips evaluation entirely — tested via handler.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=300_000) + + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + }, + {"x-headroom-bypass": "true"}, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + anyio.run(handler.handle_anthropic_messages, req) + # Guard must not fire: upstream call must have happened + assert len(handler.upstream_calls) == 1 + + +def test_variant_context1m_without_raw_declaration_degrades_to_observe(monkeypatch): + """context-1m beta without raw-id declaration degrades to observe even in reject mode.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + # Operator declared 'claude-opus-4' (sanitized form) but NOT 'claude-opus-4[1m]' + BudgetHandler = _make_handler_subclass() + + # No [1m]-keyed declaration in _operator_context_limits + handler = BudgetHandler(operator_limit=None, token_count=300_000) + # Override _operator_context_limits to have the sanitized key but not raw + handler.anthropic_provider._operator_context_limits = {} + handler.anthropic_provider.get_operator_context_limit = lambda m: None + handler.anthropic_provider.has_raw_operator_context_limit = lambda m: False + + req = _build_request( + { + "model": "claude-opus-4[1m]", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + }, + # client carries context-1m beta + {"anthropic-beta": "context-1m"}, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + anyio.run(handler.handle_anthropic_messages, req) + # Degrades to observe: request forwards (not rejected) + assert len(handler.upstream_calls) == 1 + + +def test_variant_guard_internal_error_forwards(monkeypatch): + """Any exception inside the guard block forwards the request unchanged.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + import headroom.proxy.context_budget_policy as _pol + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=300_000) + + def _bad_evaluate(**kwargs): + raise RuntimeError("synthetic policy error") + + monkeypatch.setattr(_pol, "evaluate", _bad_evaluate) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + }, + ) + anyio.run(handler.handle_anthropic_messages, req) + # Forwarded despite policy error + assert len(handler.upstream_calls) == 1 + + +# --------------------------------------------------------------------------- # +# Preservation (preservation) # +# --------------------------------------------------------------------------- # + + +def test_preservation_unconfigured_install_forwards(monkeypatch): + """Without any operator limit, requests forward byte-identically.""" + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=None, token_count=999_999) + + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + }, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(999_999)) + anyio.run(handler.handle_anthropic_messages, req) + assert len(handler.upstream_calls) == 1 + + +def test_preservation_get_context_limit_unchanged(): + """get_context_limit behavior is identical before and after the change.""" + from headroom.providers.anthropic import AnthropicProvider + + provider = AnthropicProvider( + context_limits={"test-model": 200_000}, + warn=False, + ) + assert provider.get_context_limit("test-model") == 200_000 + + +def test_preservation_get_operator_context_limit_no_declaration(): + """get_operator_context_limit returns None for undeclared models.""" + from headroom.providers.anthropic import AnthropicProvider + + provider = AnthropicProvider(warn=False) + # step-router-v1 is not in the built-in table so operator has not declared it + result = provider.get_operator_context_limit("step-router-v1") + assert result is None + + +def test_preservation_get_operator_context_limit_declared(): + """get_operator_context_limit returns the declared value for a declared model.""" + import json + + from headroom.providers.anthropic import AnthropicProvider + + limits = json.dumps({"context_limits": {"step-router-v1": 262_144}}) + with patch.dict(os.environ, {"HEADROOM_MODEL_LIMITS": limits}): + provider = AnthropicProvider(warn=False) + + assert provider.get_operator_context_limit("step-router-v1") == 262_144 + # Must not affect get_context_limit behavior + assert provider.get_context_limit("step-router-v1") == 262_144 + + +def test_preservation_no_message_mutation(monkeypatch): + """The guard never mutates body['messages'], body['system'], or body['tools'].""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "observe") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=300_000) + + original_messages = [{"role": "user", "content": "hi there"}] + req = _build_request( + {"model": "step-router-v1", "messages": original_messages, "max_tokens": 8_192}, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + anyio.run(handler.handle_anthropic_messages, req) + # Body forwarded: upstream was called (observe mode, not rejected) + assert len(handler.upstream_calls) == 1 + # The guard must not have mutated messages; the request reached upstream unchanged. + + +# --------------------------------------------------------------------------- # +# Negative space (negative_space) # +# --------------------------------------------------------------------------- # + + +def test_negative_space_at_threshold_forwards(monkeypatch): + """A request exactly at threshold forwards unchanged in both modes.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + from headroom.proxy.context_budget_policy import evaluate + + # declared=262144, max_out=8192 => threshold=253952; count=253952 (at threshold) + declared = 262_144 + max_out = 8_192 + threshold = declared - max_out + d = evaluate( + counted_tokens=threshold, + declared_limit=declared, + max_output_tokens=max_out, + mode="reject", + safety_margin=0, + ) + assert d.reason == "under_threshold" + assert d.should_reject is False + + +def test_negative_space_observe_over_threshold_still_forwards(monkeypatch): + """In observe mode an over-threshold request still reaches upstream.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "observe") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=300_000) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "big payload"}], + "max_tokens": 8_192, + }, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + anyio.run(handler.handle_anthropic_messages, req) + assert len(handler.upstream_calls) == 1 + + +def test_negative_space_bypassed_over_threshold_in_reject_mode_forwards(monkeypatch): + """A bypassed over-threshold request in reject mode still reaches upstream.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=300_000) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + }, + {"x-headroom-bypass": "true"}, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + anyio.run(handler.handle_anthropic_messages, req) + assert len(handler.upstream_calls) == 1 + + +def test_variant_context1m_with_raw_declaration_still_enforces(monkeypatch): + """A sticky context-1m beta must not disarm a model the operator declared by raw id. + + `anthropic-beta` is session-sticky (`get_session_beta_tracker`), so a later + request can carry `context-1m` it never sent. The degrade-to-observe branch + keys on whether the operator declared the raw id, so a declared model keeps + enforcing even when the beta arrives from the session baseline. + """ + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=300_000) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + }, + {"anthropic-beta": "context-1m"}, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + resp = anyio.run(handler.handle_anthropic_messages, req) + + assert resp.status_code == 400 + assert len(handler.upstream_calls) == 0 + + +# --------------------------------------------------------------------------- # +# Reproduction (reproduction) # +# --------------------------------------------------------------------------- # + + +def test_reproduction_over_limit_rejected_locally_with_zero_upstream_calls(monkeypatch): + """head: over-threshold request in reject mode returns 400 with zero upstream calls.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=270_000) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "a" * 1000}], + "max_tokens": 8_192, + }, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(270_000)) + resp = anyio.run(handler.handle_anthropic_messages, req) + assert resp.status_code == 400, f"Expected 400, got {resp.status_code}" + body = json.loads(resp.body) + assert body["type"] == "error" + assert body["error"]["type"] == "invalid_request_error" + assert "step-router-v1" in body["error"]["message"] + # Zero upstream calls + assert len(handler.upstream_calls) == 0 + + +# --------------------------------------------------------------------------- # +# Production route (production_route) # +# --------------------------------------------------------------------------- # + + +def test_production_route_guard_fires_on_create_app(monkeypatch): + """Guard fires on a real POST /v1/messages through create_app. + + The upstream transport is stubbed so no live API call is made. + Falsifiability: removing the guard by patching get_operator_context_limit + to return None makes the request reach the (stubbed) upstream instead. + """ + import json as _json + + import headroom.tokenizers as _tk + from headroom.proxy.models import ProxyConfig + from headroom.proxy.server import create_app + + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + memory_enabled=False, + ) + app = create_app(config) + proxy = app.state.proxy + + # Patch the tokenizer to return a token count above the declared limit. + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(270_000)) + + # Declare a limit on the provider (step-router-v1 at 262144). + # The guard uses get_operator_context_limit; patch it to return 262144. + original_get_op = proxy.anthropic_provider.get_operator_context_limit + proxy.anthropic_provider._operator_context_limits = {"step-router-v1": 262_144} + proxy.anthropic_provider.get_operator_context_limit = lambda m: ( + proxy.anthropic_provider._operator_context_limits.get(m) + ) + + # Stub upstream so it records calls instead of hitting the network. + upstream_calls: list[dict] = [] + + async def _stub_retry(self_inner, method, url, headers, body, **kwargs): + upstream_calls.append({"method": method, "url": url}) + return _stub_response() + + import headroom.proxy.server as _srv + + monkeypatch.setattr(_srv.HeadroomProxy, "_retry_request", _stub_retry) + + req_body = _json.dumps( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "over limit payload"}], + "max_tokens": 8_192, + } + ) + + with TestClient(app) as client: + # --- HEAD: guard fires, returns 400, zero upstream calls --- + resp = client.post( + "/v1/messages", + content=req_body.encode(), + headers={ + "authorization": "Bearer sk-ant-api-test", + "content-type": "application/json", + }, + ) + assert resp.status_code == 400, f"Expected 400, got {resp.status_code}: {resp.text}" + pre_calls = len(upstream_calls) + assert pre_calls == 0, f"Guard should have prevented upstream call; got {pre_calls}" + + # --- MUTATION CHECK: removing guard (no declared limit) forwards request --- + proxy.anthropic_provider.get_operator_context_limit = lambda m: None + proxy.anthropic_provider._operator_context_limits = {} + upstream_calls.clear() + + resp2 = client.post( + "/v1/messages", + content=req_body.encode(), + headers={ + "authorization": "Bearer sk-ant-api-test", + "content-type": "application/json", + }, + ) + # Without a declared limit the guard is inert; upstream should be called. + assert len(upstream_calls) >= 1, ( + "Mutation check failed: removing declared limit should allow upstream call" + ) + assert resp2.status_code != 400 or len(upstream_calls) >= 1 + + # Restore for cleanup + proxy.anthropic_provider.get_operator_context_limit = original_get_op diff --git a/wiki/configuration.md b/wiki/configuration.md index 3c594e525..5bc94f7ea 100644 --- a/wiki/configuration.md +++ b/wiki/configuration.md @@ -220,6 +220,8 @@ Some settings can be configured via environment variables: | Variable | Description | Default | |----------|-------------|---------| | `HEADROOM_MODEL_LIMITS` | Custom model config (JSON string or file path) | - | +| `HEADROOM_CONTEXT_LIMIT_MODE` | Context budget guard mode. `observe` (default) logs over-limit finalized requests and forwards them; `reject` returns a local 400 before any upstream attempt. Effective only when a model limit is declared in `HEADROOM_MODEL_LIMITS` or `models.json`. | `observe` | +| `HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN` | Token reserve subtracted from the declared model context limit before comparing against the finalized request token count (in addition to `max_tokens`). A non-negative integer. | `0` | | `HEADROOM_CONFIG_DIR` | Canonical config (read-mostly) root. Derives `models.json` and per-plugin config paths when set. | `~/.headroom/config` | | `HEADROOM_WORKSPACE_DIR` | Canonical workspace (read-write state) root. Derives savings ledger, memory DB, logs, TOIN, subscription state, and more when set. | `~/.headroom` | | `HEADROOM_SAVINGS_PATH` | Full path to the proxy savings JSON ledger. Always wins when set. | derived from `${HEADROOM_WORKSPACE_DIR}` | @@ -340,6 +342,19 @@ export HEADROOM_MODEL_LIMITS='{"anthropic":{"context_limits":{"claude-new":20000 export HEADROOM_MODEL_LIMITS=/path/to/models.json ``` +### Context budget guard + +When a model's context limit is declared in `HEADROOM_MODEL_LIMITS` or `models.json`, the proxy evaluates every finalized Anthropic `/v1/messages` request against that limit before forwarding. The threshold is `declared_limit - max(HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN, max_tokens)`. + +In the default `observe` mode the proxy logs an over-limit request at WARNING and forwards it unchanged. Set `HEADROOM_CONTEXT_LIMIT_MODE=reject` to have the proxy return a local 400 before any upstream attempt, which stops the provider's own 400 from triggering a client retry storm. Unconfigured installs are unaffected; the guard is a no-op when no limit is declared for the model. + +```bash +export HEADROOM_MODEL_LIMITS='{"context_limits":{"step-router-v1":262144}}' +export HEADROOM_CONTEXT_LIMIT_MODE=reject +export HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN=4096 +headroom proxy +``` + ### Pattern-Based Inference Unknown models are automatically inferred from naming patterns: From 5d172c77c8bcfc05530f5ca2e9467bf538ed94d8 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Sat, 1 Aug 2026 06:30:48 -0400 Subject: [PATCH 2/5] fix(proxy): count the finalized Anthropic input before rejection --- headroom/proxy/handlers/anthropic.py | 18 ++- tests/test_proxy_context_budget.py | 186 ++++++++++++++++++++++++++- 2 files changed, 194 insertions(+), 10 deletions(-) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index f4a9fc9d5..7859bec52 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2610,7 +2610,7 @@ class AnthropicHandlerMixin: if _cbp_should_guard: _cbp_mode = _cbp_resolve_mode() _cbp_margin = _cbp_resolve_safety_margin() - _cbp_declared = self.anthropic_provider.get_operator_context_limit(model) + _cbp_declared = self.anthropic_provider.get_operator_context_limit(raw_model) _cbp_max_out = int(body.get("max_tokens") or 0) # Degrade to observe when the outbound anthropic-beta carries @@ -2621,12 +2621,24 @@ class AnthropicHandlerMixin: _cbp_outbound_beta = headers.get("anthropic-beta", "") _cbp_has_context1m = "context-1m" in _cbp_outbound_beta if _cbp_has_context1m and not ( - self.anthropic_provider.has_raw_operator_context_limit(model) + self.anthropic_provider.has_raw_operator_context_limit(raw_model) ): _cbp_mode = "observe" + _cbp_final_messages = body.get("messages", optimized_messages) + if body.get("system") is not None: + _cbp_final_messages = [ + {"role": "system", "content": body["system"]}, + *_cbp_final_messages, + ] + _cbp_counted_tokens = tokenizer.count_messages(_cbp_final_messages) + if body.get("tools"): + _cbp_counted_tokens += tokenizer.count_text( + json.dumps(body["tools"], default=str) + ) + _cbp_decision = _cbp_evaluate( - counted_tokens=optimized_tokens, + counted_tokens=_cbp_counted_tokens, declared_limit=_cbp_declared, max_output_tokens=_cbp_max_out, mode=_cbp_mode, diff --git a/tests/test_proxy_context_budget.py b/tests/test_proxy_context_budget.py index 33d3c7cd7..6eb7c00af 100644 --- a/tests/test_proxy_context_budget.py +++ b/tests/test_proxy_context_budget.py @@ -20,6 +20,7 @@ from typing import Any from unittest.mock import MagicMock, patch import anyio +import pytest from fastapi import Request from fastapi.testclient import TestClient @@ -44,6 +45,19 @@ class _DummyTokenizer: return self._count +class _FinalizedBodyTokenizer(_DummyTokenizer): + """Expose system and tool content only when the final body is counted.""" + + def count_messages(self, messages) -> int: + count = self._count + if any(message.get("role") == "system" for message in messages): + count += 50_000 + return count + + def count_text(self, text: str) -> int: + return 50_000 if "large-tool-schema" in text else 0 + + class _DummyMetrics: def __init__(self) -> None: self.stage_timings: list = [] @@ -129,6 +143,7 @@ class _DummyHandler: self._token_count = token_count self._bypass_header = bypass self.upstream_calls: list[dict] = [] + self.upstream_bodies: list[dict] = [] self.rate_limiter = None self.metrics = _DummyMetrics() @@ -218,6 +233,7 @@ class _DummyHandler: async def _retry_request(self, method, url, headers, body, **_kwargs): self.upstream_calls.append({"method": method, "url": url}) + self.upstream_bodies.append(json.loads(body) if isinstance(body, bytes) else body) return _stub_response() def _get_compression_cache(self, session_id): @@ -315,6 +331,22 @@ def test_contract_isolation_standalone_evaluate_matches_handler(monkeypatch): assert direct.should_reject is True assert direct.reason == "over_threshold" + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=270_000) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + } + ) + import headroom.tokenizers as _tk + + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(270_000)) + response = anyio.run(handler.handle_anthropic_messages, req) + assert response.status_code == 400 + assert len(handler.upstream_calls) == 0 + # --------------------------------------------------------------------------- # # Policy evaluate branches (mode_ / variant_) # @@ -363,6 +395,38 @@ def test_mode_observe_over_threshold(): assert d.should_reject is False # observe => no rejection +def test_variant_observe_over_threshold_logs_decision_fields(monkeypatch, caplog): + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "observe") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=260_000) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + } + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(260_000)) + + anyio.run(handler.handle_anthropic_messages, req) + + assert len(handler.upstream_calls) == 1 + warning = next( + record.getMessage() + for record in caplog.records + if "context_budget_guard" in record.getMessage() + ) + assert "declared_limit=262144" in warning + assert "reserve=8192" in warning + assert "threshold=253952" in warning + assert "counted=260000" in warning + assert "overage=6048" in warning + assert "mode=observe" in warning + + def test_mode_reject_over_threshold(): from headroom.proxy.context_budget_policy import evaluate @@ -377,6 +441,27 @@ def test_mode_reject_over_threshold(): assert d.should_reject is True +def test_mode_resolvers_default_and_invalid(monkeypatch): + from headroom.proxy.context_budget_policy import resolve_mode, resolve_safety_margin + + monkeypatch.delenv("HEADROOM_CONTEXT_LIMIT_MODE", raising=False) + monkeypatch.delenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", raising=False) + assert resolve_mode() == "observe" + assert resolve_safety_margin() == 0 + + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "invalid") + with pytest.raises(ValueError, match="accepted values"): + resolve_mode() + + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "not-an-int") + with pytest.raises(ValueError, match="not an integer"): + resolve_safety_margin() + + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "-1") + with pytest.raises(ValueError, match="must be >= 0"): + resolve_safety_margin() + + def test_variant_degenerate_threshold(): """When max_output_tokens >= declared_limit the threshold is <= 0: no rejection.""" from headroom.proxy.context_budget_policy import evaluate @@ -519,16 +604,16 @@ def test_preservation_unconfigured_install_forwards(monkeypatch): BudgetHandler = _make_handler_subclass() handler = BudgetHandler(operator_limit=None, token_count=999_999) - req = _build_request( - { - "model": "step-router-v1", - "messages": [{"role": "user", "content": "hi"}], - "max_tokens": 8_192, - }, - ) + original_body = { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + } + req = _build_request(original_body) monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(999_999)) anyio.run(handler.handle_anthropic_messages, req) assert len(handler.upstream_calls) == 1 + assert handler.upstream_bodies[-1] == original_body def test_preservation_get_context_limit_unchanged(): @@ -584,6 +669,9 @@ def test_preservation_no_message_mutation(monkeypatch): # Body forwarded: upstream was called (observe mode, not rejected) assert len(handler.upstream_calls) == 1 # The guard must not have mutated messages; the request reached upstream unchanged. + assert handler.upstream_bodies[-1]["messages"] == original_messages + assert "system" not in handler.upstream_bodies[-1] + assert "tools" not in handler.upstream_bodies[-1] # --------------------------------------------------------------------------- # @@ -680,6 +768,90 @@ def test_variant_context1m_with_raw_declaration_still_enforces(monkeypatch): assert len(handler.upstream_calls) == 0 +def test_variant_context1m_suffixed_declaration_still_enforces(monkeypatch): + """A declaration keyed by the [1m] model id survives handler sanitization.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=None, token_count=300_000) + handler.anthropic_provider._operator_context_limits["claude-opus-4[1m]"] = 262_144 + req = _build_request( + { + "model": "claude-opus-4[1m]", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + }, + {"anthropic-beta": "context-1m"}, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + + response = anyio.run(handler.handle_anthropic_messages, req) + + assert response.status_code == 400 + assert len(handler.upstream_calls) == 0 + + +def test_variant_finalized_body_counts_system_and_tools(monkeypatch): + """Reject mode counts top-level system and tool input after shaping.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=220_000) + req = _build_request( + { + "model": "step-router-v1", + "system": "system instructions", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"name": "large-tool-schema", "input_schema": {"type": "object"}}], + "max_tokens": 8_192, + } + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _FinalizedBodyTokenizer(220_000)) + + response = anyio.run(handler.handle_anthropic_messages, req) + + assert response.status_code == 400 + assert len(handler.upstream_calls) == 0 + + +def test_variant_output_shaper_mutation_is_counted_after_recount(monkeypatch): + """The guard sees a system mutation made after the metrics recount.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1") + import headroom.proxy.output_savings as _savings + import headroom.proxy.output_shaper as _shaper + import headroom.tokenizers as _tk + + monkeypatch.setattr(_savings, "assign_arm", lambda *args: "treatment") + monkeypatch.setattr(_shaper, "resolve_verbosity_level", lambda settings: (2, "test")) + + def _mutate_body(body, settings, *, level_override): + body["system"] = "new system content" + return SimpleNamespace(changed=True, labels=["test-shaper"]) + + monkeypatch.setattr(_shaper, "shape_request", _mutate_body) + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=220_000) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + } + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _FinalizedBodyTokenizer(220_000)) + + response = anyio.run(handler.handle_anthropic_messages, req) + + assert response.status_code == 400 + assert len(handler.upstream_calls) == 0 + + # --------------------------------------------------------------------------- # # Reproduction (reproduction) # # --------------------------------------------------------------------------- # From 1d2a1ebef906e22820f17bf24760752f9331e127 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Sat, 1 Aug 2026 07:03:03 -0400 Subject: [PATCH 3/5] fix(proxy): make budget guard configuration failures visible --- docs/content/docs/configuration.mdx | 2 +- headroom/proxy/handlers/anthropic.py | 29 ++++++++---- tests/test_proxy_context_budget.py | 67 ++++++++++++++++++++++++---- wiki/configuration.md | 2 +- 4 files changed, 82 insertions(+), 18 deletions(-) diff --git a/docs/content/docs/configuration.mdx b/docs/content/docs/configuration.mdx index 13b8c28ac..44d19f4c6 100644 --- a/docs/content/docs/configuration.mdx +++ b/docs/content/docs/configuration.mdx @@ -277,7 +277,7 @@ headroom proxy --learn --min-evidence 3 | `HEADROOM_STATELESS` | Set to `true` to disable filesystem writes | `false` | | `HEADROOM_MODEL_LIMITS` | Custom model config (JSON string or file path) | -- | | `HEADROOM_CONTEXT_LIMIT_MODE` | Context budget guard mode. `observe` (default) logs over-limit finalized requests and forwards them; `reject` returns a local 400 before any upstream attempt. Effective only when a model limit is declared in `HEADROOM_MODEL_LIMITS` or `models.json`. | `observe` | -| `HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN` | Token reserve subtracted from the declared model context limit before comparing against the finalized request token count (in addition to `max_tokens`). A non-negative integer. | `0` | +| `HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN` | Token reserve subtracted from the declared model context limit before comparing against the finalized request token count. The reserve is the larger of this value and `max_tokens`. A non-negative integer. | `0` | | `HEADROOM_BASE_URL` | Base URL of the Headroom proxy (TypeScript SDK) | `http://localhost:8787` | | `HEADROOM_API_KEY` | Optional API key for authenticated Headroom endpoints (TypeScript SDK) | -- | | `HEADROOM_CONFIG_DIR` | Canonical config (read-mostly) root. Derives `models.json` and per-plugin config paths when set. | `~/.headroom/config` | diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 7859bec52..f7a1aa7f9 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2602,15 +2602,12 @@ class AnthropicHandlerMixin: resolve_safety_margin as _cbp_resolve_safety_margin, ) - _cbp_should_guard = True # Short-circuit: bypass header skips all Headroom behaviour. - if _bypass: - _cbp_should_guard = False - - if _cbp_should_guard: + if not _bypass: + _cbp_declared = self.anthropic_provider.get_operator_context_limit(raw_model) + if not _bypass and _cbp_declared is not None: _cbp_mode = _cbp_resolve_mode() _cbp_margin = _cbp_resolve_safety_margin() - _cbp_declared = self.anthropic_provider.get_operator_context_limit(raw_model) _cbp_max_out = int(body.get("max_tokens") or 0) # Degrade to observe when the outbound anthropic-beta carries @@ -2619,11 +2616,17 @@ class AnthropicHandlerMixin: # strips the [1m] suffix and the declared limit would apply to # the base model, not the 1M variant. _cbp_outbound_beta = headers.get("anthropic-beta", "") - _cbp_has_context1m = "context-1m" in _cbp_outbound_beta + _cbp_has_context1m = "context-1m" in _cbp_outbound_beta.lower() if _cbp_has_context1m and not ( self.anthropic_provider.has_raw_operator_context_limit(raw_model) ): _cbp_mode = "observe" + logger.warning( + "[%s] context_budget_guard: context-1m beta has no raw model " + "declaration for %s; forwarding in observe mode", + request_id, + raw_model, + ) _cbp_final_messages = body.get("messages", optimized_messages) if body.get("system") is not None: @@ -2648,7 +2651,7 @@ class AnthropicHandlerMixin: if _cbp_decision.reason == "over_threshold": logger.warning( "[%s] context_budget_guard: model=%s declared_limit=%d " - "reserve=%d threshold=%d counted=%d overage=%d mode=%s", + "reserve=%d threshold=%d counted=%d overage=%d mode=%s %s", request_id, model, _cbp_decision.declared_limit, @@ -2657,6 +2660,8 @@ class AnthropicHandlerMixin: _cbp_decision.counted_tokens, _cbp_decision.overage, _cbp_mode, + "; set HEADROOM_CONTEXT_LIMIT_MODE=reject to refuse locally or " + "adjust HEADROOM_MODEL_LIMITS", ) if _cbp_decision.should_reject: @@ -2678,6 +2683,14 @@ class AnthropicHandlerMixin: }, }, ) + except ValueError as exc: + logger.warning( + "[%s] context_budget_guard: invalid configuration (%s); forwarding " + "unchanged. Set HEADROOM_CONTEXT_LIMIT_MODE to observe or reject and " + "use a non-negative HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN.", + request_id, + exc, + ) except Exception: # Any failure in limit lookup, mode resolution, or evaluation # forwards the request unchanged. diff --git a/tests/test_proxy_context_budget.py b/tests/test_proxy_context_budget.py index 6eb7c00af..fc2130f3f 100644 --- a/tests/test_proxy_context_budget.py +++ b/tests/test_proxy_context_budget.py @@ -144,6 +144,7 @@ class _DummyHandler: self._bypass_header = bypass self.upstream_calls: list[dict] = [] self.upstream_bodies: list[dict] = [] + self.upstream_original_body_bytes: list[bytes | None] = [] self.rate_limiter = None self.metrics = _DummyMetrics() @@ -234,6 +235,7 @@ class _DummyHandler: async def _retry_request(self, method, url, headers, body, **_kwargs): self.upstream_calls.append({"method": method, "url": url}) self.upstream_bodies.append(json.loads(body) if isinstance(body, bytes) else body) + self.upstream_original_body_bytes.append(_kwargs.get("original_body_bytes")) return _stub_response() def _get_compression_cache(self, session_id): @@ -535,7 +537,7 @@ def test_variant_bypass_not_evaluated(monkeypatch): assert len(handler.upstream_calls) == 1 -def test_variant_context1m_without_raw_declaration_degrades_to_observe(monkeypatch): +def test_variant_context1m_without_raw_declaration_degrades_to_observe(monkeypatch, caplog): """context-1m beta without raw-id declaration degrades to observe even in reject mode.""" monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") @@ -544,12 +546,9 @@ def test_variant_context1m_without_raw_declaration_degrades_to_observe(monkeypat # Operator declared 'claude-opus-4' (sanitized form) but NOT 'claude-opus-4[1m]' BudgetHandler = _make_handler_subclass() - # No [1m]-keyed declaration in _operator_context_limits handler = BudgetHandler(operator_limit=None, token_count=300_000) - # Override _operator_context_limits to have the sanitized key but not raw - handler.anthropic_provider._operator_context_limits = {} - handler.anthropic_provider.get_operator_context_limit = lambda m: None - handler.anthropic_provider.has_raw_operator_context_limit = lambda m: False + # The sanitized declaration must be present so this is a true context-1m false positive. + handler.anthropic_provider._operator_context_limits.update({"claude-opus-4": 262_144}) req = _build_request( { @@ -557,13 +556,14 @@ def test_variant_context1m_without_raw_declaration_degrades_to_observe(monkeypat "messages": [{"role": "user", "content": "hi"}], "max_tokens": 8_192, }, - # client carries context-1m beta - {"anthropic-beta": "context-1m"}, + # Configured header casing must not change context-1m detection. + {"anthropic-beta": "Context-1M"}, ) monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) anyio.run(handler.handle_anthropic_messages, req) # Degrades to observe: request forwards (not rejected) assert len(handler.upstream_calls) == 1 + assert any("no raw model declaration" in record.getMessage() for record in caplog.records) def test_variant_guard_internal_error_forwards(monkeypatch): @@ -592,6 +592,32 @@ def test_variant_guard_internal_error_forwards(monkeypatch): assert len(handler.upstream_calls) == 1 +def test_variant_invalid_guard_configuration_forwards_with_warning(monkeypatch, caplog): + """Invalid operator values fail open and identify the configuration fix.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "invalid") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=300_000) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + } + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(300_000)) + + anyio.run(handler.handle_anthropic_messages, req) + + assert len(handler.upstream_calls) == 1 + assert any( + "invalid configuration" in record.getMessage() + and "forwarding unchanged" in record.getMessage() + for record in caplog.records + ) + + # --------------------------------------------------------------------------- # # Preservation (preservation) # # --------------------------------------------------------------------------- # @@ -614,6 +640,7 @@ def test_preservation_unconfigured_install_forwards(monkeypatch): anyio.run(handler.handle_anthropic_messages, req) assert len(handler.upstream_calls) == 1 assert handler.upstream_bodies[-1] == original_body + assert handler.upstream_original_body_bytes[-1] == json.dumps(original_body).encode() def test_preservation_get_context_limit_unchanged(): @@ -674,6 +701,30 @@ def test_preservation_no_message_mutation(monkeypatch): assert "tools" not in handler.upstream_bodies[-1] +def test_preservation_system_and_tools_forward_unchanged(monkeypatch): + """Observe mode forwards top-level system and tools without mutation.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "observe") + import headroom.tokenizers as _tk + + original_body = { + "model": "step-router-v1", + "system": "system instructions", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"name": "tool", "input_schema": {"type": "object"}}], + "max_tokens": 8_192, + } + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=262_144, token_count=10) + req = _build_request(original_body) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(10)) + + anyio.run(handler.handle_anthropic_messages, req) + + assert len(handler.upstream_calls) == 1 + assert handler.upstream_bodies[-1] == original_body + assert handler.upstream_original_body_bytes[-1] == json.dumps(original_body).encode() + + # --------------------------------------------------------------------------- # # Negative space (negative_space) # # --------------------------------------------------------------------------- # diff --git a/wiki/configuration.md b/wiki/configuration.md index 5bc94f7ea..a4c4b9a1b 100644 --- a/wiki/configuration.md +++ b/wiki/configuration.md @@ -221,7 +221,7 @@ Some settings can be configured via environment variables: |----------|-------------|---------| | `HEADROOM_MODEL_LIMITS` | Custom model config (JSON string or file path) | - | | `HEADROOM_CONTEXT_LIMIT_MODE` | Context budget guard mode. `observe` (default) logs over-limit finalized requests and forwards them; `reject` returns a local 400 before any upstream attempt. Effective only when a model limit is declared in `HEADROOM_MODEL_LIMITS` or `models.json`. | `observe` | -| `HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN` | Token reserve subtracted from the declared model context limit before comparing against the finalized request token count (in addition to `max_tokens`). A non-negative integer. | `0` | +| `HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN` | Token reserve subtracted from the declared model context limit before comparing against the finalized request token count. The reserve is the larger of this value and `max_tokens`. A non-negative integer. | `0` | | `HEADROOM_CONFIG_DIR` | Canonical config (read-mostly) root. Derives `models.json` and per-plugin config paths when set. | `~/.headroom/config` | | `HEADROOM_WORKSPACE_DIR` | Canonical workspace (read-write state) root. Derives savings ledger, memory DB, logs, TOIN, subscription state, and more when set. | `~/.headroom` | | `HEADROOM_SAVINGS_PATH` | Full path to the proxy savings JSON ledger. Always wins when set. | derived from `${HEADROOM_WORKSPACE_DIR}` | From da3f10f18cb3436bc77efd5dfc26b87fb7a096b6 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Sat, 1 Aug 2026 08:36:17 -0400 Subject: [PATCH 4/5] test(proxy): cover operator context limit lookup branches --- tests/test_proxy_context_budget.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/test_proxy_context_budget.py b/tests/test_proxy_context_budget.py index fc2130f3f..955919c60 100644 --- a/tests/test_proxy_context_budget.py +++ b/tests/test_proxy_context_budget.py @@ -679,6 +679,32 @@ def test_preservation_get_operator_context_limit_declared(): assert provider.get_context_limit("step-router-v1") == 262_144 +def test_preservation_get_operator_context_limit_sanitized_variant(): + """A sanitized lookup finds the declaration for a styled model id.""" + from headroom.providers.anthropic import AnthropicProvider + + provider = AnthropicProvider( + context_limits={"claude-opus-4": 262_144}, + warn=False, + ) + + assert provider.get_operator_context_limit("claude-opus-4[1m]") == 262_144 + assert provider.has_raw_operator_context_limit("claude-opus-4[1m]") is False + + +def test_preservation_has_raw_operator_context_limit(): + """A declaration keyed by the styled id is recognized as raw.""" + from headroom.providers.anthropic import AnthropicProvider + + provider = AnthropicProvider( + context_limits={"claude-opus-4[1m]": 1_000_000}, + warn=False, + ) + + assert provider.get_operator_context_limit("claude-opus-4[1m]") == 1_000_000 + assert provider.has_raw_operator_context_limit("claude-opus-4[1m]") is True + + def test_preservation_no_message_mutation(monkeypatch): """The guard never mutates body['messages'], body['system'], or body['tools'].""" monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "observe") From cb9101fedfabb9e40d050fb92270af24b3d95bcf Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Tue, 4 Aug 2026 05:06:09 -0400 Subject: [PATCH 5/5] fix(proxy): treat a non-positive budget threshold as over budget in reject mode --- headroom/proxy/context_budget_policy.py | 13 ++- headroom/proxy/handlers/anthropic.py | 2 +- tests/test_proxy_context_budget.py | 146 +++++++++++++++++++++++- 3 files changed, 154 insertions(+), 7 deletions(-) diff --git a/headroom/proxy/context_budget_policy.py b/headroom/proxy/context_budget_policy.py index 815742a5e..0a25c4964 100644 --- a/headroom/proxy/context_budget_policy.py +++ b/headroom/proxy/context_budget_policy.py @@ -67,7 +67,8 @@ def evaluate( Branches (in order): 1. declared_limit is None -> reason='no_declared_limit', should_reject=False - 2. threshold <= 0 -> reason='degenerate_threshold', should_reject=False + 2. threshold <= 0 -> reason='degenerate_threshold', overage reported, + should_reject=(mode=='reject') 3. counted_tokens <= threshold -> reason='under_threshold', should_reject=False 4. else -> reason='over_threshold', should_reject=(mode=='reject') @@ -91,14 +92,20 @@ def evaluate( threshold = declared_limit - reserve if threshold <= 0: + # The reserved output alone consumes the whole declared window, so no + # request can fit by construction. In reject mode that is over budget: + # refusing keeps the proxy's promise that an impossible request never + # reaches upstream. Observe mode still logs and forwards. The overage + # uses the same counted - threshold formula as the over-threshold branch, + # which is positive whenever the threshold is non-positive. return BudgetDecision( mode=mode, counted_tokens=counted_tokens, declared_limit=declared_limit, reserve=reserve, threshold=threshold, - overage=0, - should_reject=False, + overage=counted_tokens - threshold, + should_reject=(mode == "reject"), reason="degenerate_threshold", ) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index f7a1aa7f9..32b7d5a50 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2648,7 +2648,7 @@ class AnthropicHandlerMixin: safety_margin=_cbp_margin, ) - if _cbp_decision.reason == "over_threshold": + if _cbp_decision.reason in ("over_threshold", "degenerate_threshold"): logger.warning( "[%s] context_budget_guard: model=%s declared_limit=%d " "reserve=%d threshold=%d counted=%d overage=%d mode=%s %s", diff --git a/tests/test_proxy_context_budget.py b/tests/test_proxy_context_budget.py index 955919c60..138c1bad5 100644 --- a/tests/test_proxy_context_budget.py +++ b/tests/test_proxy_context_budget.py @@ -465,7 +465,12 @@ def test_mode_resolvers_default_and_invalid(monkeypatch): def test_variant_degenerate_threshold(): - """When max_output_tokens >= declared_limit the threshold is <= 0: no rejection.""" + """Equal max_output_tokens (threshold=0) is over budget in reject mode. + + A request whose reserved output consumes the entire window cannot fit by + construction, so reject mode refuses it instead of forwarding (review + 4850625429 regression). + """ from headroom.proxy.context_budget_policy import evaluate d = evaluate( @@ -476,11 +481,13 @@ def test_variant_degenerate_threshold(): safety_margin=0, ) assert d.reason == "degenerate_threshold" - assert d.should_reject is False + assert d.threshold == 0 + assert d.should_reject is True + assert d.overage == 100_000 # counted_tokens - threshold def test_variant_degenerate_threshold_max_output_exceeds(): - """max_output_tokens > declared_limit is also degenerate.""" + """max_output_tokens > declared_limit is over budget in reject mode.""" from headroom.proxy.context_budget_policy import evaluate d = evaluate( @@ -491,7 +498,140 @@ def test_variant_degenerate_threshold_max_output_exceeds(): safety_margin=0, ) assert d.reason == "degenerate_threshold" + assert d.threshold == 8_192 - 16_000 + assert d.should_reject is True + assert d.overage == 100_000 - d.threshold + + +def test_variant_degenerate_threshold_observe_forwards(): + """Observe mode keeps a non-positive threshold non-rejecting and logs.""" + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=100_000, + declared_limit=8_192, + max_output_tokens=8_192, + mode="observe", + safety_margin=0, + ) + assert d.reason == "degenerate_threshold" assert d.should_reject is False + assert d.overage == 100_000 + + +def test_variant_degenerate_threshold_safety_margin_reject(): + """safety_margin >= declared_limit is over budget in reject mode.""" + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=100_000, + declared_limit=8_192, + max_output_tokens=4_096, + mode="reject", + safety_margin=8_192, # equal to declared_limit -> threshold <= 0 + ) + assert d.reason == "degenerate_threshold" + assert d.threshold <= 0 + assert d.should_reject is True + assert d.overage == 100_000 - d.threshold + + +def test_variant_degenerate_threshold_safety_margin_observe(): + """Degenerate safety_margin in observe mode stays non-rejecting.""" + from headroom.proxy.context_budget_policy import evaluate + + d = evaluate( + counted_tokens=100_000, + declared_limit=8_192, + max_output_tokens=4_096, + mode="observe", + safety_margin=8_192, + ) + assert d.reason == "degenerate_threshold" + assert d.should_reject is False + + +def test_handler_degenerate_max_tokens_rejects_with_zero_upstream(monkeypatch): + """Handler: max_tokens >= declared_limit in reject mode returns 400 locally. + + Review 4850625429 scenario: declared_limit=200_000, max_output_tokens=200_000, + counted_tokens=1, mode=reject. The threshold is 0, which must count as over + budget so the request is refused before any upstream attempt. + """ + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=200_000, token_count=1) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 200_000, # equals declared_limit -> threshold=0 + }, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(1)) + + resp = anyio.run(handler.handle_anthropic_messages, req) + + assert resp.status_code == 400, f"Expected local 400, got {resp.status_code}" + body = json.loads(resp.body) + assert "step-router-v1" in body["error"]["message"] + assert len(handler.upstream_calls) == 0 + + +def test_handler_degenerate_safety_margin_rejects_with_zero_upstream(monkeypatch): + """Handler: safety_margin >= declared_limit in reject mode returns 400 locally.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "reject") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "200000") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=200_000, token_count=1) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8_192, + }, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(1)) + + resp = anyio.run(handler.handle_anthropic_messages, req) + + assert resp.status_code == 400, f"Expected local 400, got {resp.status_code}" + assert len(handler.upstream_calls) == 0 + + +def test_handler_degenerate_observe_logs_and_forwards(monkeypatch, caplog): + """Handler: observe mode logs a degenerate threshold and still forwards.""" + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_MODE", "observe") + monkeypatch.setenv("HEADROOM_CONTEXT_LIMIT_SAFETY_MARGIN", "0") + import headroom.tokenizers as _tk + + BudgetHandler = _make_handler_subclass() + handler = BudgetHandler(operator_limit=200_000, token_count=1) + req = _build_request( + { + "model": "step-router-v1", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 200_000, + }, + ) + monkeypatch.setattr(_tk, "get_tokenizer", lambda m: _DummyTokenizer(1)) + + anyio.run(handler.handle_anthropic_messages, req) + + assert len(handler.upstream_calls) == 1 + warning = next( + record.getMessage() + for record in caplog.records + if "context_budget_guard" in record.getMessage() + ) + assert "declared_limit=200000" in warning + assert "threshold=0" in warning + assert "overage=1" in warning def test_variant_safety_margin_adds_reserve():