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

Merge cb9101fedf into 8884d87378

Browse source
This commit is contained in:
Rod Boev 2026-08-27 05:00:34 -04:00 • committed by GitHub
commit 02520fe799
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1507 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

15
docs/content/docs/configuration.mdx
View file

@ -332,6 +332,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. 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` |
@ -504,6 +506,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:

32
headroom/providers/anthropic.py
View file

@ -628,12 +628,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:
@ -725,6 +729,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

134
headroom/proxy/context_budget_policy.py Normal file
View file

@ -0,0 +1,134 @@
"""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', 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')
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:
# 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=counted_tokens - threshold,
should_reject=(mode == "reject"),
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",
)

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

@ -3248,6 +3248,120 @@ 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,
)
# Short-circuit: bypass header skips all Headroom behaviour.
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_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.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:
_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=_cbp_counted_tokens,
declared_limit=_cbp_declared,
max_output_tokens=_cbp_max_out,
mode=_cbp_mode,
safety_margin=_cbp_margin,
)
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",
request_id,
model,
_cbp_decision.declared_limit,
_cbp_decision.reserve,
_cbp_decision.threshold,
_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:
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 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.
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
#
# An extension may have published a per-request routing decision on

1198
tests/test_proxy_context_budget.py Normal file
View file

File diff suppressed because it is too large Load diff

15
wiki/configuration.md
View file

@ -225,6 +225,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. 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}` |
@ -345,6 +347,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: