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

fix(proxy/cost): mark estimated-basis budget records and add an enforcement policy (#2713) (#2725)

Browse source 
## Description

`CostTracker.check_budget()` is a hard spend control — the Anthropic
handler refuses the request with a 429 once the period budget is gone.
The ledger that control reads could not tell a measured dollar from a
guessed one.

When a provider response carries no input-token breakdown,
`record_tokens()` substitutes Headroom's own `tokens_sent` estimate for
the input count so input cost isn't silently dropped from the budget.
That fallback is the right call, but the resulting record was
byte-identical to a provider-measured one: no field, no log line, no
separation in `/stats`. `RequestOutcome.uncached_input_tokens` defaults
to `0`, so any route whose response omits usage lands on this branch in
production. An estimate can drift in either direction, so a budget check
could pass after real spend had already gone over — with nothing saying
the decision rested on an estimate.

This keeps the fallback and makes it visible, then lets operators decide
what an estimate is allowed to do to a hard limit.

Closes #2713

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- New `headroom/proxy/budget_basis_policy.py` (pure policy module,
matching the existing `*_policy.py` convention): the
`measured`/`estimated` basis constants, the `count`/`ignore`/`block`
policy values, and `resolve_estimated_basis_policy()` (explicit value →
`HEADROOM_BUDGET_ESTIMATED_BASIS` → `count`; an unknown value warns once
and falls back rather than failing proxy startup).
- `headroom/proxy/cost.py`: ledger entries are now `CostEntry(timestamp,
cost_usd, basis)` instead of a bare tuple; `record_tokens()` marks the
fallback branch `estimated` and logs one WARNING per model (deduped the
same way pricing warnings are, per #2504 — an unguarded warning on this
path fires once per request for a provider that never reports usage);
new `period_cost_breakdown()` and an optional `basis` filter on
`get_period_cost()`; new `budget_denial_detail()` builds the 429 body
where the ledger lives; `check_budget()` honors the policy while keeping
its `(allowed, remaining)` signature.
- `stats()` gains `budget_estimated_basis` (the active policy) and
`budget_basis` (the period split: `total_usd`, `measured_usd`,
`estimated_usd`, `estimated_pct`, `records`, `estimated_records`).
`merge_cost_stats()` already spreads `**cost_stats`, so both reach
`/stats["cost"]` with no extra plumbing.
- Operator knob wired through every config layer:
`ProxyConfig.budget_estimated_basis` (`models.py`), the Click
`--budget-estimated-basis` option with `envvar=` (`cli/proxy.py`), the
argparse `--budget-estimated-basis` flag (`server.py`, `default=None` so
the env var stays reachable), and a `SettingField` in the `Budget` group
(`settings_store.py`).
- `headroom/proxy/handlers/anthropic.py`: the 429 body now comes from
`budget_denial_detail()`, which names how much of the period's spend was
booked from an estimate and distinguishes "you overspent" from "I refuse
to enforce a hard limit on a guess".
- `headroom/cli/doctor.py`: the budget check stays **PASS** and appends
the estimated share (and the policy, when it isn't the default). No new
WARN state — a provider that never reports usage would otherwise sit at
a permanent WARN. Every new read is `.get()` + type-guarded so `doctor`
still works against an older running proxy.
- `docs/content/docs/metrics.mdx`: a "Measured vs Estimated Spend"
subsection with the `/stats` shape and the three policy values.
- Tests: new `tests/test_cost_budget_basis.py` (20 tests) plus 4 new
`doctor` tests; `tests/test_anthropic_pre_upstream_backpressure.py`'s
cost-tracker double gained `budget_denial_detail()` to match the
handler's duck-typed contract.

### Policy values

| `HEADROOM_BUDGET_ESTIMATED_BASIS` | Effect on the hard limit |
|---|---|
| `count` (default) | Unchanged behavior — estimated spend consumes the
budget. |
| `ignore` | Booked and reported, but only measured spend enforces. |
| `block` | Fail closed — refuse rather than enforce a hard limit on a
guess. |

Default enforcement is unchanged. `CHANGELOG.md` is untouched.

## Testing

- [x] Unit tests pass (`pytest`) — every test covering the changed
modules; see `Not tested` for this machine's pre-existing environment
failures
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`) — clean on every file this
PR touches
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ python -m pytest tests/test_cost_budget_basis.py tests/test_cost_tracker_counterfactual.py tests/test_cost_pricing_warning_dedup.py -q
31 passed

$ python -m pytest tests/test_cli_doctor.py -q
72 passed

$ python -m pytest tests/test_anthropic_pre_upstream_backpressure.py -q
25 passed

$ python -m pytest tests/test_proxy_settings_endpoints.py tests/test_proxy/test_settings_store.py -q
50 passed

# full suite (see "Not tested" below for the excluded modules and the pre-existing failures)
$ python -m pytest -q
...
tests\test_cost_budget_basis.py ....................                     [ 25%]
tests\test_cost_pricing_warning_dedup.py ...                             [ 25%]
tests\test_cost_tracker_counterfactual.py ........                       [ 25%]
...
217 failed, 8807 passed, 657 skipped, 5318 warnings, 59 errors in 716.30s (0:11:56)

# same failing files re-run on clean upstream/main with the change stashed -> identical count
$ git stash push -u -- headroom tests docs
$ python -m pytest tests/test_log_compressor.py tests/test_cache/test_client_integration.py \
    tests/test_fsutil.py tests/test_savings_ledger.py tests/test_ccr_mcp_http.py \
    tests/test_router_registry_dispatch.py tests/test_proxy_savings_history.py \
    tests/test_text_compressors.py tests/test_builtin_compressor_adapters.py \
    tests/test_cli_proxy_env.py -q
73 failed, 182 passed in 34.82s     # 40+16+2+2+1+1+1+4+3+3 = 73, matching the run above

$ python -m mypy headroom --ignore-missing-imports --python-version 3.13
# 12 errors, all in release_version.py / ccr/mcp_server.py / memory/mcp_server.py
# (stale local `mcp` stubs) — none in any file this PR touches

$ python -m ruff check headroom/proxy/budget_basis_policy.py headroom/proxy/cost.py headroom/proxy/server.py headroom/proxy/models.py headroom/proxy/handlers/anthropic.py headroom/cli/proxy.py headroom/cli/doctor.py headroom/settings_store.py tests/test_cost_budget_basis.py tests/test_cli_doctor.py tests/test_anthropic_pre_upstream_backpressure.py
All checks passed!

$ python -m ruff format --check <same 11 files>
11 files already formatted
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.13.11, branch
`fix/budget-estimated-basis-2713` off `upstream/main` @ `232fb49c`,
`PYTHONPATH` pointed at the working tree so the repo copy of `headroom`
is imported rather than the installed one.
- Exact command / steps: ran the repro script from the issue body
verbatim, then extended it to print `stats()["budget_basis"]` for both
trackers, to construct the same tracker with
`estimated_basis_policy="block"` and with `"ignore"`, and to record
twice against the same model to check the warning dedup. Separately
drove `headroom doctor`'s `check_budget` against stub `/stats` payloads
(mixed basis, all-measured, non-default policy, and an older proxy that
omits the new keys).
- Observed result: the issue's two figures are unchanged, so the
fallback still works — no breakdown `$0.008100`, with breakdown
`$0.005100`, ratio `1.59x`. The two are now separable: the no-breakdown
tracker reports `{'total_usd': 0.0081, 'measured_usd': 0.0,
'estimated_usd': 0.0081, 'estimated_pct': 100.0, 'records': 1,
'estimated_records': 1}` and the with-breakdown tracker reports
`estimated_usd: 0.0, estimated_pct: 0.0, estimated_records: 0`. One
`WARNING headroom.proxy: budget basis estimated: no usage breakdown from
provider for gpt-4o-mini — input cost booked from Headroom's own token
count` fires across repeated records, not one per request. With
`policy=block`, `check_budget()` returns `(False, 0.0)` and the 429
detail reads `Budget enforcement blocked for daily period: $0.0081 of
$0.0081 was booked from Headroom's own token estimate because the
provider returned no usage breakdown, and
HEADROOM_BUDGET_ESTIMATED_BASIS=block refuses to enforce a budget on an
estimate. Set it to 'count' or 'ignore' to serve these requests.` With
`policy=ignore`, `check_budget()` returns `(True, 0.0001)` while the
spend is still booked and reported (`0.7506`). `doctor` prints `pass
$10.0/daily budget enforced — 62% of period spend ($1.2400) booked from
Headroom token estimates`, appends `— estimated-basis policy: block` for
a non-default policy, and degrades to the plain `$10.0/daily budget
enforced` against a proxy that doesn't report the new fields.
`--budget-estimated-basis [count|ignore|block]` shows in `headroom proxy
--help`; the argparse path resolves the env var when the flag is absent
and an explicit flag wins over the env.
- Not tested: no live end-to-end run against a real provider that omits
usage in its response — the estimated basis was exercised through
`record_tokens()` directly, which is the single funnel
`emit_request_outcome()` uses. The `settings_store` field was not
exercised through the settings UI. The full-suite run above excludes
three things this machine cannot run, none of which touch the changed
files: `tests/test_hermes_passthrough_compression.py` (`respx` not
installed), `tests/test_memory/test_embedder_mps_serialization.py`
(`sentence_transformers` pins `tokenizers<=0.23.0`, local has `0.23.1`),
and `tests/test_cli/` (its subprocess-spawning tests wedge against a
leftover local proxy on :8787; each file passes in isolation, e.g.
`test_wrap_bridge.py` 7/7). Its 217 failures are all pre-existing
environment breakage — a stale local Rust `_core` build
(`test_log_compressor.py`, `test_text_compressors.py`,
`test_builtin_compressor_adapters.py`, `test_cli_proxy_env.py`, the
`test_transforms*` files) and the broken `sentence_transformers` install
(`tests/test_memory/*`, `test_memory_system.py`,
`test_sqlite_graph_store.py`) — with zero overlap with the modules this
PR changes; the stashed baseline above reproduces them 1:1. CI is the
authority for a green full suite.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes — the
only local failures are pre-existing and reproduce with the change
stashed
- [x] I did **not** edit `CHANGELOG.md` — it is generated by
release-please from my Conventional Commit PR title (a CI guard enforces
this)

## Additional Notes

The estimated-basis WARNING is deduped per model rather than emitted per
request, following the precedent set by #2504 for pricing warnings — the
whole point of this code path is that it fires on every request for a
provider that never reports usage, so an unguarded `logger.warning`
would flood `proxy.log`.

`headroom doctor` deliberately stays PASS. A WARN would be permanent,
not actionable, for anyone whose provider simply doesn't report usage;
the note tells them the number, and the `block` policy is there for
operators who want the hard failure.

`check_budget()` keeps its `(allowed, remaining)` signature and its
default `count` semantics, so
`tests/test_cost_tracker_counterfactual.py` — including
`test_budget_input_cost_counted_without_usage_breakdown`, the contract
that the fallback keeps working — passes unmodified.

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Parideboy 2026-08-03 08:05:44 +02:00 • committed by GitHub
parent 3eb0122068
commit 01df245252
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 736 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

38
docs/content/docs/metrics.mdx
View file

@ -252,6 +252,44 @@ headroom proxy --budget 10.00
When the budget is exceeded, requests return a budget exceeded error, the `/stats` endpoint shows budget status, and logs indicate the budget state.
### Measured vs Estimated Spend
Every cost record carries a *basis* — where its input-token count came from. When a provider response includes a usage breakdown, the basis is `measured`. When it doesn't, Headroom substitutes its own `tokens_sent` count so input cost isn't dropped from the budget, and the record's basis is `estimated`. Headroom logs one warning per model the first time this happens.
`/stats` keeps the two separable under `cost.budget_basis`:
```json
{
"cost": {
"budget_limit_usd": 10.0,
"budget_period": "daily",
"budget_estimated_basis": "count",
"budget_basis": {
"total_usd": 3.1400,
"measured_usd": 2.9000,
"estimated_usd": 0.2400,
"estimated_pct": 7.6,
"records": 412,
"estimated_records": 31
}
}
}
```
An estimate can drift in either direction, so you choose what it does to the hard limit:
```bash
headroom proxy --budget 10.00 --budget-estimated-basis count # default
```
| Value | Effect |
|-------|--------|
| `count` | Estimated spend consumes the budget like measured spend. The default; matches historical behavior. |
| `ignore` | Estimated spend is still booked and reported, but only provider-reported spend consumes the budget. |
| `block` | Refuse requests once the period holds any estimated spend, rather than enforcing a hard limit against a guess. |
Env: `HEADROOM_BUDGET_ESTIMATED_BASIS`. `headroom doctor` reports the estimated share alongside the budget check.
## Key Metrics to Monitor
| Metric | What It Tells You | Target |

31
headroom/cli/doctor.py
View file

@ -426,7 +426,36 @@ def check_budget(stats: dict[str, Any] | None) -> CheckResult:
hint="set one: headroom proxy --budget 10 (env: HEADROOM_BUDGET)",
)
period = cost.get("budget_period", "daily")
return CheckResult(name=name, status=PASS, summary=f"${limit}/{period} budget enforced")
summary = f"${limit}/{period} budget enforced"
return CheckResult(name=name, status=PASS, summary=summary + _estimated_basis_note(cost))
def _estimated_basis_note(cost: dict[str, Any]) -> str:
"""Describe how much of the period's spend was booked from a token estimate.
Informational, never a WARN: a provider that simply never reports a usage
breakdown would otherwise sit at a permanent warning. Every read is
defensive so `doctor` still works against a proxy predating these fields.
"""
note = ""
basis = cost.get("budget_basis")
if isinstance(basis, dict):
estimated_usd = basis.get("estimated_usd")
estimated_pct = basis.get("estimated_pct")
if isinstance(estimated_usd, (int, float)) and estimated_usd > 0:
pct = f"{estimated_pct:.0f}% " if isinstance(estimated_pct, (int, float)) else ""
note += (
f" — {pct}of period spend (${estimated_usd:.4f}) "
"booked from Headroom token estimates"
)
# Reported independently of the breakdown: a non-default policy changes how
# the budget is enforced and should surface even if the split is missing.
policy = cost.get("budget_estimated_basis")
if isinstance(policy, str) and policy and policy != "count":
note += f" — estimated-basis policy: {policy}"
return note
def check_deployments(manifests: list[Any], probe: Any = probe_json) -> CheckResult | None:

16
headroom/cli/proxy.py
View file

@ -534,6 +534,20 @@ def dashboard(port: int, no_open: bool) -> None:
"Env: HEADROOM_BUDGET_PERIOD."
),
)
@click.option(
"--budget-estimated-basis",
type=click.Choice(["count", "ignore", "block"]),
default="count",
envvar="HEADROOM_BUDGET_ESTIMATED_BASIS",
help=(
"What to do with spend booked from Headroom's own token estimate, which "
"happens when a provider response carries no input-token breakdown. "
"count: it consumes the budget like measured spend (default). "
"ignore: only provider-reported spend consumes the budget. "
"block: refuse requests rather than enforce a hard limit on an estimate. "
"Env: HEADROOM_BUDGET_ESTIMATED_BASIS."
),
)
# Code-aware compression (AST-based, requires `pip install headroom-ai[code]`).
# Pair of flags so users can override the env-var default in either direction.
# We resolve HEADROOM_CODE_AWARE_ENABLED in the body (not via Click's envvar=),
@ -941,6 +955,7 @@ def proxy(
codex_wire_debug_dir: str | None,
budget: float | None,
budget_period: str,
budget_estimated_basis: str,
code_aware_flag: bool | None,
disable_kompress: bool,
disable_kompress_fallback: bool,
@ -1235,6 +1250,7 @@ def proxy(
or os.environ.get("HEADROOM_LOG_MESSAGES", "").lower() in ("true", "1", "yes", "on"),
budget_limit_usd=budget,
budget_period=cast(Literal["hourly", "daily", "monthly"], budget_period),
budget_estimated_basis=cast(Literal["count", "ignore", "block"], budget_estimated_basis),
# Code-aware compression resolution:
# 1. Explicit --code-aware / --no-code-aware always wins.
# 2. Otherwise read HEADROOM_CODE_AWARE_ENABLED (truthy = on).

81
headroom/proxy/budget_basis_policy.py Normal file
View file

@ -0,0 +1,81 @@
"""Policy for how estimate-based cost records feed budget enforcement.
Every cost record the proxy books carries a *basis*: where the input-token
count came from. ``measured`` means the provider reported a usage breakdown;
``estimated`` means the response carried none and Headroom substituted its own
``tokens_sent`` count so input cost wasn't silently dropped from the budget
(#2713).
The substitution is the right default — dropping input cost entirely would
under-enforce far worse — but a budget is a hard control, and an estimate can
drift in either direction. This module names the two bases and the three
policies an operator can pick for what estimated spend does to the budget:
* ``count`` — estimated spend counts toward the limit (the historical
behavior, and still the default; nothing changes except that the record is
now marked, logged, and separable in ``/stats``).
* ``ignore`` — the record is still booked to the ledger and reported, but only
provider-measured spend consumes the budget. Fails open on the estimate.
* ``block`` — fail closed: once the period holds any estimated spend, refuse
rather than enforce a hard limit against a guess.
"""
from __future__ import annotations
import logging
from collections.abc import Mapping
logger = logging.getLogger("headroom.proxy")
# Provenance of the input-token count behind a booked cost record.
COST_BASIS_MEASURED = "measured"
COST_BASIS_ESTIMATED = "estimated"
# What estimated spend does to budget enforcement.
BUDGET_BASIS_COUNT = "count"
BUDGET_BASIS_IGNORE = "ignore"
BUDGET_BASIS_BLOCK = "block"
VALID_POLICIES = (BUDGET_BASIS_COUNT, BUDGET_BASIS_IGNORE, BUDGET_BASIS_BLOCK)
DEFAULT_POLICY = BUDGET_BASIS_COUNT
ENV_VAR = "HEADROOM_BUDGET_ESTIMATED_BASIS"
# A typo'd knob must not fail proxy startup, but it also must not warn on every
# resolve. Bounded by the number of distinct bad values seen (realistically 1).
_warned_invalid_policies: set[str] = set()
def resolve_estimated_basis_policy(
configured: str | None = None,
env: Mapping[str, str] | None = None,
) -> str:
"""Resolve the estimated-basis budget policy.
Precedence: explicit ``configured`` value, then ``env[ENV_VAR]``, then
:data:`DEFAULT_POLICY`. An unrecognized value falls back to the default and
warns once — a misconfigured knob should never stop the proxy from booting,
but it must not silently change enforcement either.
"""
raw = configured
if raw is None and env is not None:
raw = env.get(ENV_VAR)
if raw is None:
return DEFAULT_POLICY
candidate = raw.strip().lower()
if candidate in VALID_POLICIES:
return candidate
if not candidate:
return DEFAULT_POLICY
if candidate not in _warned_invalid_policies:
_warned_invalid_policies.add(candidate)
logger.warning(
"Unknown %s=%r — falling back to %r (valid: %s)",
ENV_VAR,
raw,
DEFAULT_POLICY,
", ".join(VALID_POLICIES),
)
return DEFAULT_POLICY

184
headroom/proxy/cost.py
View file

@ -13,8 +13,17 @@ import logging
import math
from collections import deque
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, NamedTuple
from headroom.proxy.budget_basis_policy import (
BUDGET_BASIS_BLOCK,
BUDGET_BASIS_IGNORE,
COST_BASIS_ESTIMATED,
COST_BASIS_MEASURED,
DEFAULT_POLICY,
ENV_VAR,
resolve_estimated_basis_policy,
)
from headroom.proxy.modes import PROXY_MODE_CACHE
if TYPE_CHECKING:
@ -60,6 +69,40 @@ def _warn_pricing_once(model: str, message: str) -> None:
logger.warning(message)
# A route whose responses never carry a usage breakdown hits the estimated-basis
# fallback on *every* request, so the warning is deduped per model for the same
# reason pricing warnings are (#2504): one line per model per process, not one
# per request. The distinction stays permanently visible in /stats regardless.
_warned_estimated_basis_models: set[str] = set()
def _warn_estimated_basis_once(model: str) -> None:
"""Warn the first time ``model`` books a cost against Headroom's estimate."""
if model in _warned_estimated_basis_models:
return
_warned_estimated_basis_models.add(model)
logger.warning(
"budget basis estimated: no usage breakdown from provider for %s — "
"input cost booked from Headroom's own token count",
model,
)
class CostEntry(NamedTuple):
"""One booked cost, with the provenance of the input count behind it.
``basis`` is :data:`~headroom.proxy.budget_basis_policy.COST_BASIS_MEASURED`
when the provider reported a usage breakdown, and ``COST_BASIS_ESTIMATED``
when it didn't and Headroom's own ``tokens_sent`` stood in for the input
count. Budget enforcement is a hard control, so the two must stay separable
in the ledger rather than collapsing into an undifferentiated dollar figure.
"""
timestamp: datetime
cost_usd: float
basis: str
# Provider-specific cache discount multipliers (what fraction of input price)
# Used to calculate dollar savings from prefix caching
_CACHE_ECONOMICS = {
@ -644,12 +687,21 @@ class CostTracker:
# get_period_cost() undercounts and check_budget() silently under-enforces.
COST_RETENTION_HOURS = 744 # 31 days
def __init__(self, budget_limit_usd: float | None = None, budget_period: str = "daily"):
def __init__(
self,
budget_limit_usd: float | None = None,
budget_period: str = "daily",
estimated_basis_policy: str = DEFAULT_POLICY,
):
self.budget_limit_usd = budget_limit_usd
self.budget_period = budget_period
# What estimated-basis spend does to enforcement. Normalized here so a
# bad value degrades to the default instead of quietly disabling the
# budget. See headroom.proxy.budget_basis_policy.
self.estimated_basis_policy = resolve_estimated_basis_policy(estimated_basis_policy)
# Cost tracking - using deque for efficient left-side removal
self._costs: deque[tuple[datetime, float]] = deque(maxlen=self.MAX_COST_ENTRIES)
self._costs: deque[CostEntry] = deque(maxlen=self.MAX_COST_ENTRIES)
self._last_prune_time: datetime = datetime.now()
# Token savings per model (exact, no dollar estimation)
@ -743,7 +795,7 @@ class CostTracker:
cutoff = now - timedelta(hours=self.COST_RETENTION_HOURS)
# Remove entries from the left (oldest) while they're older than cutoff
while self._costs and self._costs[0][0] < cutoff:
while self._costs and self._costs[0].timestamp < cutoff:
self._costs.popleft()
def record_tokens(
@ -807,9 +859,17 @@ class CostTracker:
# When the call site had no API usage breakdown (all cache/uncached
# fields are 0), fall back to tokens_sent so input cost isn't
# silently dropped from the budget.
#
# That fallback is a guess, and check_budget() is a hard control, so the
# record is stamped ``estimated`` and warned about once per model (#2713).
# The fallback behaviour itself is unchanged — the estimate is now
# labelled rather than indistinguishable from provider-reported usage.
basis = COST_BASIS_MEASURED
input_tokens = uncached_tokens
if not (uncached_tokens or cache_read_tokens or cache_write_tokens):
input_tokens = tokens_sent
basis = COST_BASIS_ESTIMATED
_warn_estimated_basis_once(model)
cost = self.estimate_cost(
model=model,
input_tokens=input_tokens,
@ -818,31 +878,121 @@ class CostTracker:
cache_write_tokens=cache_write_tokens,
)
if cost is not None:
self._costs.append((datetime.now(), cost))
self._costs.append(CostEntry(datetime.now(), cost, basis))
self._prune_old_costs()
def get_period_cost(self) -> float:
"""Get cost for current budget period."""
def _period_cutoff(self) -> datetime:
"""Start of the current budget period."""
now = datetime.now()
if self.budget_period == "hourly":
cutoff = now - timedelta(hours=1)
elif self.budget_period == "daily":
cutoff = now.replace(hour=0, minute=0, second=0, microsecond=0)
else: # monthly
cutoff = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
return now - timedelta(hours=1)
if self.budget_period == "daily":
return now.replace(hour=0, minute=0, second=0, microsecond=0)
# monthly
return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
return sum(cost for ts, cost in self._costs if ts >= cutoff)
def get_period_cost(self, basis: str | None = None) -> float:
"""Get cost for current budget period.
With no argument this is the total spend booked in the period,
regardless of how each record's input count was derived. Pass a basis
(``"measured"`` / ``"estimated"``) to get just that slice.
"""
cutoff = self._period_cutoff()
return sum(
entry.cost_usd
for entry in self._costs
if entry.timestamp >= cutoff and (basis is None or entry.basis == basis)
)
def period_cost_breakdown(self) -> dict[str, Any]:
"""Split the period's booked spend by how its input count was derived.
``estimated_usd`` is spend whose input token count came from Headroom's
own estimate because the provider returned no usage breakdown. Keeping
it separable is the point: a budget refusal driven by a guess should be
distinguishable from one driven by provider-reported usage (#2713).
"""
cutoff = self._period_cutoff()
measured_usd = 0.0
estimated_usd = 0.0
records = 0
estimated_records = 0
for entry in self._costs:
if entry.timestamp < cutoff:
continue
records += 1
if entry.basis == COST_BASIS_ESTIMATED:
estimated_usd += entry.cost_usd
estimated_records += 1
else:
measured_usd += entry.cost_usd
total_usd = measured_usd + estimated_usd
return {
"period": self.budget_period,
"policy": self.estimated_basis_policy,
"total_usd": total_usd,
"measured_usd": measured_usd,
"estimated_usd": estimated_usd,
"estimated_pct": round(estimated_usd / total_usd * 100, 1) if total_usd > 0 else 0.0,
"records": records,
"estimated_records": estimated_records,
}
def check_budget(self) -> tuple[bool, float]:
"""Check if within budget. Returns (allowed, remaining)."""
"""Check if within budget. Returns (allowed, remaining).
How estimated-basis spend participates is governed by
``estimated_basis_policy``: ``count`` (default) enforces against total
spend exactly as before, ``ignore`` enforces against provider-measured
spend only, and ``block`` refuses outright once the period holds any
estimated spend rather than enforcing a hard limit against a guess.
"""
if self.budget_limit_usd is None:
return True, float("inf")
period_cost = self.get_period_cost()
breakdown = self.period_cost_breakdown()
if self.estimated_basis_policy == BUDGET_BASIS_BLOCK and breakdown["estimated_usd"] > 0:
return False, 0.0
if self.estimated_basis_policy == BUDGET_BASIS_IGNORE:
period_cost = breakdown["measured_usd"]
else:
period_cost = breakdown["total_usd"]
remaining = self.budget_limit_usd - period_cost
return remaining > 0, max(0, remaining)
def budget_denial_detail(self) -> str:
"""Human-readable reason a request was refused on budget grounds.
Built here rather than in the handler so the message can name what the
ledger actually knows — specifically how much of the period's spend was
booked from Headroom's own token estimate.
"""
breakdown = self.period_cost_breakdown()
estimated_usd = breakdown["estimated_usd"]
if self.estimated_basis_policy == BUDGET_BASIS_BLOCK and estimated_usd > 0:
return (
f"Budget enforcement blocked for {self.budget_period} period: "
f"${estimated_usd:.4f} of ${breakdown['total_usd']:.4f} was booked from "
"Headroom's own token estimate because the provider returned no usage "
f"breakdown, and {ENV_VAR}=block refuses to enforce a budget on an "
"estimate. Set it to 'count' or 'ignore' to serve these requests."
)
detail = f"Budget exceeded for {self.budget_period} period"
if estimated_usd > 0:
detail += (
f" (${estimated_usd:.4f} of ${breakdown['total_usd']:.4f} booked from "
f"Headroom token estimates, not provider-reported usage)"
)
return detail
def _get_list_price(self, model: str) -> float | None:
"""Get list input price per 1M tokens for a model."""
litellm = _get_litellm_module()
@ -956,4 +1106,8 @@ class CostTracker:
# `headroom doctor` can report whether a budget is set.
"budget_limit_usd": self.budget_limit_usd,
"budget_period": self.budget_period,
"budget_estimated_basis": self.estimated_basis_policy,
# Period spend split by input-count provenance, so estimate-derived
# spend stays separable from provider-reported spend (#2713).
"budget_basis": self.period_cost_breakdown(),
}

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

@ -868,7 +868,7 @@ class AnthropicHandlerMixin:
await _finalize_pre_upstream()
raise HTTPException(
status_code=429,
detail=f"Budget exceeded for {self.config.budget_period} period",
detail=self.cost_tracker.budget_denial_detail(),
)
# Memory: Get user ID when memory is enabled (fallback to "default" for simple DevEx).

3
headroom/proxy/models.py
View file

@ -319,6 +319,9 @@ class ProxyConfig:
cost_tracking_enabled: bool = True
budget_limit_usd: float | None = None
budget_period: Literal["hourly", "daily", "monthly"] = "daily"
# What spend booked from Headroom's own token estimate (provider returned no
# usage breakdown) does to budget enforcement. See budget_basis_policy.
budget_estimated_basis: Literal["count", "ignore", "block"] = "count"
# Logging
log_requests: bool = True

17
headroom/proxy/server.py
View file

@ -113,6 +113,7 @@ from headroom.proxy import runtime_env
from headroom.proxy.audit import is_auditable_path, record_admin_action
from headroom.proxy.auth_mode import should_stamp_codex_client
from headroom.proxy.background_compression import BackgroundCompressor
from headroom.proxy.budget_basis_policy import resolve_estimated_basis_policy
# =============================================================================
# Extracted modules (re-exported for backward compatibility)
@ -797,6 +798,7 @@ class HeadroomProxy(
CostTracker(
budget_limit_usd=config.budget_limit_usd,
budget_period=config.budget_period,
estimated_basis_policy=config.budget_estimated_basis,
)
if config.cost_tracking_enabled
else None
@ -5472,6 +5474,17 @@ if __name__ == "__main__":
# Cost
parser.add_argument("--budget", type=float, help="Budget limit in USD")
parser.add_argument("--budget-period", choices=["hourly", "daily", "monthly"], default="daily")
parser.add_argument(
"--budget-estimated-basis",
choices=["count", "ignore", "block"],
default=None,
help=(
"What spend booked from Headroom's own token estimate (provider returned no "
"usage breakdown) does to the budget: count it (default), ignore it, or block "
"requests rather than enforce the limit on an estimate. "
"Env: HEADROOM_BUDGET_ESTIMATED_BASIS."
),
)
# Logging
parser.add_argument("--log-file", help="Log file path")
@ -5574,6 +5587,10 @@ if __name__ == "__main__":
rate_limit_tokens_per_minute=_get_env_int("HEADROOM_TPM", args.tpm),
budget_limit_usd=args.budget,
budget_period=args.budget_period,
budget_estimated_basis=cast(
Literal["count", "ignore", "block"],
resolve_estimated_basis_policy(args.budget_estimated_basis, os.environ),
),
log_file=_get_env_str("HEADROOM_LOG_FILE", args.log_file)
if args.log_file
else os.environ.get("HEADROOM_LOG_FILE"),

16
headroom/settings_store.py
View file

@ -204,6 +204,22 @@ SETTINGS: tuple[SettingField, ...] = (
help="Period the budget applies to.",
tier="basic",
),
SettingField(
"HEADROOM_BUDGET_ESTIMATED_BASIS",
"budget_estimated_basis",
"Estimated-basis spend",
"Budget",
"enum",
default="count",
choices=("count", "ignore", "block"),
help=(
"What spend booked from Headroom's own token estimate does to the budget "
"when a provider response carries no input-token breakdown. count: it "
"consumes the budget. ignore: only provider-reported spend does. block: "
"refuse requests rather than enforce a hard limit on an estimate."
),
tier="advanced",
),
# --- Networking (baked into the install manifest on supervised deploys) ---
SettingField(
"HEADROOM_HOST",

3
tests/test_anthropic_pre_upstream_backpressure.py
View file

@ -852,6 +852,9 @@ class _CostTrackerBlock:
def check_budget(self):
return False, 0
def budget_denial_detail(self):
return "Budget exceeded for daily period"
def record_tokens(self, *a, **k):
return None

56
tests/test_cli_doctor.py
View file

@ -444,6 +444,62 @@ class TestBudget:
assert result.status == PASS
assert "$10.0/daily" in result.summary
def test_estimated_basis_share_is_reported_without_warning(self):
"""#2713: spend booked from a token estimate is surfaced, not warned on.
A provider that never reports a usage breakdown would otherwise sit at a
permanent WARN, so this stays informational.
"""
result = check_budget(
{
"cost": {
"budget_limit_usd": 10.0,
"budget_period": "daily",
"budget_estimated_basis": "count",
"budget_basis": {"estimated_usd": 1.24, "estimated_pct": 62.3},
}
}
)
assert result.status == PASS
assert "62% of period spend ($1.2400)" in result.summary
assert "Headroom token estimates" in result.summary
def test_all_measured_spend_adds_no_note(self):
result = check_budget(
{
"cost": {
"budget_limit_usd": 10.0,
"budget_period": "daily",
"budget_estimated_basis": "count",
"budget_basis": {"estimated_usd": 0.0, "estimated_pct": 0.0},
}
}
)
assert result.summary == "$10.0/daily budget enforced"
def test_non_default_basis_policy_is_named(self):
result = check_budget(
{
"cost": {
"budget_limit_usd": 10.0,
"budget_period": "daily",
"budget_estimated_basis": "block",
}
}
)
assert "estimated-basis policy: block" in result.summary
def test_missing_basis_fields_degrade_quietly(self):
"""`doctor` must still work against a proxy predating these fields."""
result = check_budget({"cost": {"budget_limit_usd": 10.0, "budget_period": "daily"}})
assert result.status == PASS
assert result.summary == "$10.0/daily budget enforced"
malformed = check_budget(
{"cost": {"budget_limit_usd": 10.0, "budget_period": "daily", "budget_basis": "nope"}}
)
assert malformed.status == PASS
@dataclass
class _FakeManifest:

306
tests/test_cost_budget_basis.py Normal file
View file

@ -0,0 +1,306 @@
"""Budget records must say whether their input count was measured or estimated.
#2713: when a provider response carries no input-token breakdown,
``record_tokens`` substitutes Headroom's own ``tokens_sent`` for the input
count. The fallback is right — dropping input cost would under-enforce far
worse — but the resulting record used to be indistinguishable from a
provider-measured one, so ``check_budget`` (a hard control) could refuse or
allow on the strength of a guess with nothing saying so.
These tests pin the marking, the deduped warning, the separable ledger, and
the three operator policies.
"""
from __future__ import annotations
import logging
import pytest
from tests._dotenv import (
autouse_apply_env,
importorskip_no_env_leak,
load_env_overrides,
)
_env_overrides = load_env_overrides()
apply_dotenv = autouse_apply_env(_env_overrides)
importorskip_no_env_leak("litellm")
MODEL = "claude-sonnet-4-20250514"
@pytest.fixture(autouse=True)
def _reset_warning_dedup():
"""The per-model warn-once set is module-global; keep tests order-independent."""
import headroom.proxy.cost as cost_mod
cost_mod._warned_estimated_basis_models.clear()
yield
cost_mod._warned_estimated_basis_models.clear()
def _tracker(**kwargs):
from headroom.proxy.server import CostTracker
return CostTracker(**kwargs)
# ── Basis marking ────────────────────────────────────────────────────
def test_missing_usage_breakdown_books_estimated_basis():
"""No breakdown → the record is marked estimated and reported as such."""
ct = _tracker(budget_limit_usd=100.0)
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=50_000, output_tokens=1_000)
basis = ct.stats()["budget_basis"]
assert basis["estimated_usd"] > 0
assert basis["measured_usd"] == 0
assert basis["estimated_records"] == 1
assert basis["estimated_pct"] == 100.0
def test_provider_breakdown_books_measured_basis():
"""A reported breakdown → measured; nothing lands in the estimated bucket."""
ct = _tracker(budget_limit_usd=100.0)
ct.record_tokens(
MODEL,
tokens_saved=0,
tokens_sent=50_000,
uncached_tokens=30_000,
output_tokens=1_000,
)
basis = ct.stats()["budget_basis"]
assert basis["measured_usd"] > 0
assert basis["estimated_usd"] == 0
assert basis["estimated_records"] == 0
assert basis["estimated_pct"] == 0.0
def test_cache_read_only_response_counts_as_measured():
"""A fully cache-read turn reports usage, so it is not an estimate."""
ct = _tracker(budget_limit_usd=100.0)
ct.record_tokens(
MODEL,
tokens_saved=0,
tokens_sent=50_000,
cache_read_tokens=40_000,
output_tokens=1_000,
)
assert ct.stats()["budget_basis"]["estimated_usd"] == 0
def test_mixed_records_stay_separable_and_sum_to_total():
"""Default policy is unchanged: the budget still sees every booked dollar."""
ct = _tracker(budget_limit_usd=100.0)
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=50_000, output_tokens=1_000)
ct.record_tokens(
MODEL,
tokens_saved=0,
tokens_sent=50_000,
uncached_tokens=30_000,
output_tokens=1_000,
)
basis = ct.stats()["budget_basis"]
assert basis["records"] == 2
assert basis["estimated_records"] == 1
assert basis["measured_usd"] > 0
assert basis["estimated_usd"] > 0
assert basis["total_usd"] == pytest.approx(basis["measured_usd"] + basis["estimated_usd"])
# Regression guard: `count` (the default) enforces against total spend
# exactly as it did before this change.
assert ct.get_period_cost() == pytest.approx(basis["total_usd"])
def test_get_period_cost_can_filter_by_basis():
ct = _tracker(budget_limit_usd=100.0)
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=50_000, output_tokens=1_000)
ct.record_tokens(
MODEL, tokens_saved=0, tokens_sent=50_000, uncached_tokens=30_000, output_tokens=1_000
)
measured = ct.get_period_cost("measured")
estimated = ct.get_period_cost("estimated")
assert measured > 0
assert estimated > 0
assert ct.get_period_cost() == pytest.approx(measured + estimated)
# ── Warning ──────────────────────────────────────────────────────────
def test_estimated_basis_warns_once_per_model(caplog):
"""A route that never reports usage must not flood proxy.log (cf. #2504)."""
ct = _tracker(budget_limit_usd=100.0)
with caplog.at_level(logging.WARNING, logger="headroom.proxy"):
for _ in range(5):
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=10_000, output_tokens=100)
hits = [r for r in caplog.records if "budget basis estimated" in r.getMessage()]
assert len(hits) == 1
assert MODEL in hits[0].getMessage()
def test_distinct_models_each_warn_once(caplog):
ct = _tracker(budget_limit_usd=100.0)
other = "claude-haiku-4-5-20251001"
with caplog.at_level(logging.WARNING, logger="headroom.proxy"):
for model in (MODEL, MODEL, other, other):
ct.record_tokens(model, tokens_saved=0, tokens_sent=10_000, output_tokens=100)
msgs = [r.getMessage() for r in caplog.records if "budget basis estimated" in r.getMessage()]
assert sum(MODEL in m for m in msgs) == 1
assert sum(other in m for m in msgs) == 1
def test_measured_records_do_not_warn(caplog):
ct = _tracker(budget_limit_usd=100.0)
with caplog.at_level(logging.WARNING, logger="headroom.proxy"):
ct.record_tokens(
MODEL, tokens_saved=0, tokens_sent=10_000, uncached_tokens=9_000, output_tokens=100
)
assert not [r for r in caplog.records if "budget basis estimated" in r.getMessage()]
# ── Enforcement policies ─────────────────────────────────────────────
def test_count_policy_lets_estimated_spend_exhaust_the_budget():
"""Default: an estimate consumes the budget, as it always has."""
ct = _tracker(budget_limit_usd=0.0001, estimated_basis_policy="count")
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=500_000, output_tokens=10_000)
allowed, remaining = ct.check_budget()
assert not allowed
assert remaining == 0
def test_ignore_policy_keeps_estimated_spend_out_of_enforcement():
"""`ignore`: the record is still booked and reported, but doesn't enforce."""
ct = _tracker(budget_limit_usd=0.0001, estimated_basis_policy="ignore")
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=500_000, output_tokens=10_000)
allowed, _remaining = ct.check_budget()
assert allowed
# Still visible in the ledger — ignored for enforcement, not dropped.
assert ct.stats()["budget_basis"]["estimated_usd"] > 0
def test_ignore_policy_still_enforces_measured_spend():
ct = _tracker(budget_limit_usd=0.0001, estimated_basis_policy="ignore")
ct.record_tokens(
MODEL,
tokens_saved=0,
tokens_sent=500_000,
uncached_tokens=500_000,
output_tokens=10_000,
)
allowed, _remaining = ct.check_budget()
assert not allowed
def test_block_policy_refuses_once_any_estimated_spend_exists():
"""`block`: fail closed rather than enforce a hard limit on a guess."""
ct = _tracker(budget_limit_usd=1_000_000.0, estimated_basis_policy="block")
allowed, _remaining = ct.check_budget()
assert allowed # nothing booked yet
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=1_000, output_tokens=10)
allowed, remaining = ct.check_budget()
assert not allowed
assert remaining == 0.0
def test_block_policy_is_inert_without_a_budget_limit():
"""No limit configured means no hard control to protect."""
ct = _tracker(budget_limit_usd=None, estimated_basis_policy="block")
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=1_000, output_tokens=10)
allowed, remaining = ct.check_budget()
assert allowed
assert remaining == float("inf")
def test_block_policy_allows_purely_measured_traffic():
ct = _tracker(budget_limit_usd=1_000_000.0, estimated_basis_policy="block")
ct.record_tokens(
MODEL, tokens_saved=0, tokens_sent=1_000, uncached_tokens=900, output_tokens=10
)
allowed, _remaining = ct.check_budget()
assert allowed
def test_invalid_policy_falls_back_to_count():
ct = _tracker(budget_limit_usd=0.0001, estimated_basis_policy="nonsense")
assert ct.estimated_basis_policy == "count"
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=500_000, output_tokens=10_000)
allowed, _remaining = ct.check_budget()
assert not allowed
# ── Denial message ───────────────────────────────────────────────────
def test_denial_detail_names_the_estimated_share():
ct = _tracker(budget_limit_usd=0.0001)
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=500_000, output_tokens=10_000)
detail = ct.budget_denial_detail()
assert "Budget exceeded for daily period" in detail
assert "Headroom token estimates" in detail
def test_denial_detail_unchanged_for_purely_measured_spend():
ct = _tracker(budget_limit_usd=0.0001)
ct.record_tokens(
MODEL,
tokens_saved=0,
tokens_sent=500_000,
uncached_tokens=500_000,
output_tokens=10_000,
)
assert ct.budget_denial_detail() == "Budget exceeded for daily period"
def test_block_denial_is_distinguishable_from_overspend():
ct = _tracker(budget_limit_usd=1_000_000.0, estimated_basis_policy="block")
ct.record_tokens(MODEL, tokens_saved=0, tokens_sent=1_000, output_tokens=10)
detail = ct.budget_denial_detail()
assert "Budget enforcement blocked" in detail
assert "HEADROOM_BUDGET_ESTIMATED_BASIS=block" in detail
# ── Policy resolver ──────────────────────────────────────────────────
def test_resolver_precedence_and_fallback():
from headroom.proxy.budget_basis_policy import resolve_estimated_basis_policy
env = {"HEADROOM_BUDGET_ESTIMATED_BASIS": "ignore"}
assert resolve_estimated_basis_policy("block", env) == "block" # explicit wins
assert resolve_estimated_basis_policy(None, env) == "ignore" # env next
assert resolve_estimated_basis_policy(None, {}) == "count" # default
assert resolve_estimated_basis_policy(None, None) == "count"
assert resolve_estimated_basis_policy("BLOCK", None) == "block" # normalized
assert resolve_estimated_basis_policy("nope", None) == "count" # fallback
def test_stats_reports_the_active_policy():
ct = _tracker(budget_limit_usd=10.0, estimated_basis_policy="block")
assert ct.stats()["budget_estimated_basis"] == "block"
assert ct.stats()["budget_basis"]["policy"] == "block"