From 01df2452529a86c689cf226fecd5918cc5d19676 Mon Sep 17 00:00:00 2001 From: Parideboy Date: Mon, 3 Aug 2026 08:05:44 +0200 Subject: [PATCH] fix(proxy/cost): mark estimated-basis budget records and add an enforcement policy (#2713) (#2725) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 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 --- docs/content/docs/metrics.mdx | 38 +++ headroom/cli/doctor.py | 31 +- headroom/cli/proxy.py | 16 + headroom/proxy/budget_basis_policy.py | 81 +++++ headroom/proxy/cost.py | 184 ++++++++++- headroom/proxy/handlers/anthropic.py | 2 +- headroom/proxy/models.py | 3 + headroom/proxy/server.py | 17 + headroom/settings_store.py | 16 + ...est_anthropic_pre_upstream_backpressure.py | 3 + tests/test_cli_doctor.py | 56 ++++ tests/test_cost_budget_basis.py | 306 ++++++++++++++++++ 12 files changed, 736 insertions(+), 17 deletions(-) create mode 100644 headroom/proxy/budget_basis_policy.py create mode 100644 tests/test_cost_budget_basis.py diff --git a/docs/content/docs/metrics.mdx b/docs/content/docs/metrics.mdx index 1439c6861..6c8b1412f 100644 --- a/docs/content/docs/metrics.mdx +++ b/docs/content/docs/metrics.mdx @@ -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 | diff --git a/headroom/cli/doctor.py b/headroom/cli/doctor.py index 5a0a032b4..f1f06b380 100644 --- a/headroom/cli/doctor.py +++ b/headroom/cli/doctor.py @@ -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: diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 1a4dbe7ee..30292a4e1 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -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). diff --git a/headroom/proxy/budget_basis_policy.py b/headroom/proxy/budget_basis_policy.py new file mode 100644 index 000000000..1226755ec --- /dev/null +++ b/headroom/proxy/budget_basis_policy.py @@ -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 diff --git a/headroom/proxy/cost.py b/headroom/proxy/cost.py index f2142023f..067af7dcd 100644 --- a/headroom/proxy/cost.py +++ b/headroom/proxy/cost.py @@ -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(), } diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index e255ab886..cafd72edd 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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). diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index 8a5651fe9..069867e2b 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -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 diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index a6dff74f3..04ea3cb5e 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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"), diff --git a/headroom/settings_store.py b/headroom/settings_store.py index 6a8dff200..c085c181f 100644 --- a/headroom/settings_store.py +++ b/headroom/settings_store.py @@ -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", diff --git a/tests/test_anthropic_pre_upstream_backpressure.py b/tests/test_anthropic_pre_upstream_backpressure.py index 3595fc429..777ee7c6a 100644 --- a/tests/test_anthropic_pre_upstream_backpressure.py +++ b/tests/test_anthropic_pre_upstream_backpressure.py @@ -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 diff --git a/tests/test_cli_doctor.py b/tests/test_cli_doctor.py index a8cebba3f..6bce283e8 100644 --- a/tests/test_cli_doctor.py +++ b/tests/test_cli_doctor.py @@ -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: diff --git a/tests/test_cost_budget_basis.py b/tests/test_cost_budget_basis.py new file mode 100644 index 000000000..cdb7ad1d6 --- /dev/null +++ b/tests/test_cost_budget_basis.py @@ -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"