From a033ac4176b09c716905aa0f45ae317e954f0eb9 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Mon, 3 Aug 2026 22:41:10 -0700 Subject: [PATCH] fix(cost): send litellm the total prompt so --budget stops seeing $0 (#2757) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description **`--budget` has been silently inert on any cache-warm request** — which is the normal case in an agent session. This is a disabled control, not a metrics bug. `record_tokens` passed only the **uncached slice** as litellm's `prompt_tokens`. Measured, `litellm.cost_per_token` charges: ``` (prompt_tokens - cache_read - cache_creation) * input_rate + cache_read * read_rate + cache_creation * write_rate ``` So `prompt_tokens` is the **whole** prompt and litellm removes the cached parts itself. Handing it the uncached slice drives the input term **negative** as soon as anything is cached. `estimate_cost` ends with `float(total) if total > 0 else None`, so it returned `None`, no `CostEntry` was appended, and `check_budget()` saw **$0**. | model | 100k prompt, 80k cached — old call | booked | | --- | --- | --- | | `gpt-5` | **-$0.065000** | None | | `gpt-4o-mini` | **-$0.003000** | None | | `claude-sonnet-4-5` | **-$0.156000** | None | Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made Fixing the total alone would **over-charge OpenAI**, because two bugs are entangled here. OpenAI exposes no cache-write counter, so `_infer_openai_cache_write_tokens` uses the uncached portion as a write proxy. At two of the four inference sites (`openai.py:4304`, `openai.py:5413`) `uncached_input_tokens` is derived by subtracting **only** `cache_read`, so `cache_write_tokens` and `uncached_tokens` are **the same tokens**. Summing all three would double-count the prompt and charge a write premium OpenAI does not have. (The other two sites — `openai.py:3969` and `streaming.py:2024` — subtract both, so their buckets are genuinely disjoint; those are left alone.) - `cost.py` — `record_tokens` now passes `uncached + cache_read + cache_write` as the prompt total. - `cost.py` — new `cache_inferred: bool = False` parameter. When set, the inferred write is excluded from **both** the prompt total and the write premium. The default preserves behaviour for every provider that reports disjoint buckets. - `outcome.py` — plumbs `outcome.cache_inferred` through. The field already existed on `RequestOutcome` for the dashboard; it just never reached cost. - `handlers/openai.py` — sets `cache_inferred=True` at the two outcome sites whose buckets genuinely duplicate. ## Testing - [x] Unit tests pass (new file) - [x] Linting passes (`ruff check` + `format --check`, pinned 0.15.17) - [x] New tests added - [x] Manual testing performed ### Test Output ```text $ uvx ruff@0.15.17 check headroom/proxy/cost.py headroom/proxy/outcome.py headroom/proxy/handlers/openai.py All checks passed! $ uvx ruff@0.15.17 format --check 3 files already formatted $ pytest tests/test_cost_budget_total_prompt.py -q 5 passed in 0.26s ``` The 5 new tests assert on the **arguments handed to `estimate_cost`** rather than on dollar values, so they pin the contract that broke without depending on litellm's pricing tables — or on litellm being installed at all. ## Real Behavior Proof - **Environment:** macOS 26.4 arm64, isolated git worktree off `upstream/main`. **litellm's actual formula**, established by probe rather than assumed: ```text rates (claude-sonnet-4-5): input=3e-06 read=3e-07 write=3.75e-06 total=100k, 80k read, 5k write -> $0.087750 hypothesis (p-r)*ir + r*rr + w*wr = $0.102750 x hypothesis (p-r-w)*ir + r*rr + w*wr = $0.087750 <- matches ``` **Before → after:** ```text Anthropic, disjoint: uncached=900 read=48000 write=1500 OLD prompt=900 raw=-0.125775 booked=None <-- BUDGET BLIND NEW prompt=50400 raw= 0.022725 booked=0.022725 OpenAI, inferred write == uncached: uncached=20000 read=80000 write=20000 OLD prompt=20000 raw=-0.090000 booked=None <-- BUDGET BLIND NEW prompt=100000, write excluded raw= 0.035000 booked=0.035 truth=0.035000 ``` The OpenAI "after" equals the hand-computed truth `20,000*input + 80,000*read_rate` exactly. - **Not fully tested locally.** The cost/outcome suites show an **identical 10-failure set** on this branch and on clean `upstream/main` in the same throwaway env — all `ModuleNotFoundError: headroom._core`, the compiled Rust extension this machine cannot currently build. So they are environment-only, not regressions. One of them, `test_funnel_passes_canonical_record_tokens_shape`, covers the `record_tokens` call shape this PR changes, so **CI is the authoritative check for that one**. ## 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] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] I did **not** edit `CHANGELOG.md` ## Context Found during a tokenizer-consistency audit that also turned up: HuggingFace-routed models counting a 6,000-char message as **2 tokens**, `gpt-5`/`o4-mini` falling to a char estimator, and the Kompress size gate missing its own cap by 24%. Those are separate PRs — different subsystems, different risk. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- headroom/proxy/cost.py | 30 +++++- headroom/proxy/handlers/openai.py | 2 + headroom/proxy/outcome.py | 1 + tests/test_cost_budget_total_prompt.py | 124 +++++++++++++++++++++++++ tests/test_request_outcome.py | 5 + 5 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 tests/test_cost_budget_total_prompt.py diff --git a/headroom/proxy/cost.py b/headroom/proxy/cost.py index 067af7dcd..435c09fd4 100644 --- a/headroom/proxy/cost.py +++ b/headroom/proxy/cost.py @@ -809,6 +809,7 @@ class CostTracker: cache_write_1h_tokens: int = 0, uncached_tokens: int = 0, output_tokens: int = 0, + cache_inferred: bool = False, ): """Record token counts per model and accumulate request cost for budget enforcement. @@ -820,6 +821,12 @@ class CostTracker: cache_write_tokens: Cache write tokens from API response usage. uncached_tokens: Non-cached input tokens from API response usage. output_tokens: Output tokens from API response usage. + cache_inferred: True when ``cache_write_tokens`` was DERIVED from the + uncached portion rather than reported by the provider (OpenAI + exposes no write counter). Such a value is the same tokens as + ``uncached_tokens``, so it is excluded from the billed prompt + total and from the write premium. Defaults False, which preserves + behaviour for providers that report disjoint buckets. """ # Post-guard invariant (all providers): Headroom never forwards a request # larger than the original (handlers revert any inflation before sending), @@ -864,8 +871,27 @@ class CostTracker: # 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. + # ``litellm.cost_per_token`` wants the TOTAL prompt in ``prompt_tokens``: + # measured, it charges + # (prompt - cache_read - cache_creation) * input_rate + # + cache_read * read_rate + # + cache_creation * write_rate + # Passing only the uncached slice therefore drives the input term + # NEGATIVE once anything was cached, and ``estimate_cost`` returns None on + # a non-positive total — so no CostEntry was appended and ``check_budget()`` + # saw $0. Every cache-warm request, i.e. the normal case in an agent + # session, was booking zero spend and the budget could never trip. + # Measured before this fix, 100k prompt with 80k cached: + # gpt-5 $-0.065, gpt-4o-mini $-0.003, claude-sonnet-4-5 $-0.156. + # + # An INFERRED cache-write (OpenAI exposes no write counter, so the + # uncached portion is used as a write proxy) is the SAME tokens as + # ``uncached_tokens``. Adding it to the total would double-count the + # prompt, and charging it at the write premium would invent a cost OpenAI + # does not have — so it is excluded from both. + effective_cache_write = 0 if cache_inferred else cache_write_tokens basis = COST_BASIS_MEASURED - input_tokens = uncached_tokens + input_tokens = uncached_tokens + cache_read_tokens + effective_cache_write if not (uncached_tokens or cache_read_tokens or cache_write_tokens): input_tokens = tokens_sent basis = COST_BASIS_ESTIMATED @@ -875,7 +901,7 @@ class CostTracker: input_tokens=input_tokens, output_tokens=output_tokens, cache_read_tokens=cache_read_tokens, - cache_write_tokens=cache_write_tokens, + cache_write_tokens=effective_cache_write, ) if cost is not None: self._costs.append(CostEntry(datetime.now(), cost, basis)) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 736455386..3f0adc904 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -4432,6 +4432,7 @@ class OpenAIHandlerMixin: cache_read_tokens=cache_read_tokens, cache_write_tokens=cache_write_tokens, uncached_input_tokens=uncached_input_tokens, + cache_inferred=True, total_latency_ms=total_latency, overhead_ms=optimization_latency, pipeline_timing=pipeline_timing, @@ -5451,6 +5452,7 @@ class OpenAIHandlerMixin: cache_read_tokens=cache_read_tokens, cache_write_tokens=cache_write_tokens, uncached_input_tokens=uncached_input_tokens, + cache_inferred=True, total_latency_ms=total_latency, overhead_ms=optimization_latency, transforms_applied=tuple(transforms_applied), diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index 66ec1e6e0..f25b99a7b 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -476,6 +476,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None: cache_write_5m_tokens=outcome.cache_write_5m_tokens, cache_write_1h_tokens=outcome.cache_write_1h_tokens, uncached_tokens=outcome.uncached_input_tokens, + cache_inferred=outcome.cache_inferred, output_tokens=outcome.output_tokens, ) diff --git a/tests/test_cost_budget_total_prompt.py b/tests/test_cost_budget_total_prompt.py new file mode 100644 index 000000000..451d09e32 --- /dev/null +++ b/tests/test_cost_budget_total_prompt.py @@ -0,0 +1,124 @@ +"""``record_tokens`` must hand litellm the TOTAL prompt, not the uncached slice. + +`litellm.cost_per_token` charges:: + + (prompt_tokens - cache_read - cache_creation) * input_rate + + cache_read * read_rate + + cache_creation * write_rate + +so ``prompt_tokens`` is the whole prompt and litellm removes the cached parts +itself. Passing only the uncached slice drove the input term NEGATIVE as soon as +anything was cached; ``estimate_cost`` returns None on a non-positive total, no +``CostEntry`` was appended, and ``check_budget()`` therefore saw $0. Every +cache-warm request — the normal case in an agent session — booked zero spend, so +``--budget`` could never trip. Measured against real litellm pricing on a 100k +prompt with 80k cached: gpt-5 -$0.065, gpt-4o-mini -$0.003, +claude-sonnet-4-5 -$0.156. + +These assert on the arguments handed to ``estimate_cost`` rather than on dollar +values, so they pin the contract that broke without depending on litellm's +pricing tables (or on litellm being installed). +""" + +from __future__ import annotations + +from headroom.proxy.cost import COST_BASIS_MEASURED, CostTracker + + +def _tracker_capturing_cost(**kwargs) -> tuple[CostTracker, dict]: + """A tracker whose ``estimate_cost`` records its kwargs and returns a real cost.""" + tracker = CostTracker(**kwargs) + seen: dict = {} + + def fake_estimate_cost(**kw): + seen.update(kw) + return 0.5 # non-None, so a CostEntry is appended + + tracker.estimate_cost = fake_estimate_cost # type: ignore[method-assign] + return tracker, seen + + +def test_disjoint_buckets_send_the_summed_total_as_prompt_tokens() -> None: + """Anthropic reports uncached / read / creation as three disjoint buckets.""" + tracker, seen = _tracker_capturing_cost() + + tracker.record_tokens( + "claude-sonnet-4-5", + tokens_saved=0, + tokens_sent=1_000, + cache_read_tokens=48_000, + cache_write_tokens=1_500, + uncached_tokens=900, + ) + + assert seen["input_tokens"] == 900 + 48_000 + 1_500 + # The write premium still applies — those tokens really were written. + assert seen["cache_write_tokens"] == 1_500 + assert seen["cache_read_tokens"] == 48_000 + + +def test_inferred_write_is_excluded_from_the_total_and_the_premium() -> None: + """OpenAI exposes no write counter, so the write value IS the uncached tokens. + + Counting it again would double the prompt, and charging it at a write premium + would invent a cost OpenAI does not have. + """ + tracker, seen = _tracker_capturing_cost() + + tracker.record_tokens( + "gpt-5", + tokens_saved=0, + tokens_sent=1_000, + cache_read_tokens=80_000, + cache_write_tokens=20_000, # inferred: identical to uncached_tokens + uncached_tokens=20_000, + cache_inferred=True, + ) + + assert seen["input_tokens"] == 20_000 + 80_000, "inferred write must not be added" + assert seen["cache_write_tokens"] == 0, "inferred write must not be charged a premium" + + +def test_a_cache_warm_request_actually_books_spend() -> None: + """The regression itself: the budget must see this request.""" + tracker, _ = _tracker_capturing_cost(budget_limit_usd=100.0) + + tracker.record_tokens( + "claude-sonnet-4-5", + tokens_saved=0, + tokens_sent=1_000, + cache_read_tokens=48_000, + cache_write_tokens=1_500, + uncached_tokens=900, + ) + + assert len(tracker._costs) == 1, "cache-warm request booked no spend — budget is blind" + assert tracker._costs[0].basis == COST_BASIS_MEASURED + + +def test_no_usage_breakdown_still_falls_back_to_tokens_sent() -> None: + """Pre-existing estimated-basis fallback must be untouched by this change.""" + tracker, seen = _tracker_capturing_cost() + + tracker.record_tokens("gpt-4o", tokens_saved=0, tokens_sent=4_242) + + assert seen["input_tokens"] == 4_242 + assert len(tracker._costs) == 1 + assert tracker._costs[0].basis != COST_BASIS_MEASURED + + +def test_cache_inferred_defaults_false_so_reporting_providers_are_unchanged() -> None: + """Callers that never pass the flag keep the disjoint-bucket arithmetic.""" + tracker, seen = _tracker_capturing_cost() + + tracker.record_tokens( + "claude-sonnet-4-5", + tokens_saved=0, + tokens_sent=1_000, + cache_read_tokens=10, + cache_write_tokens=20, + uncached_tokens=30, + ) + + assert seen["input_tokens"] == 60 + assert seen["cache_write_tokens"] == 20 diff --git a/tests/test_request_outcome.py b/tests/test_request_outcome.py index e9368bf3e..09936a094 100644 --- a/tests/test_request_outcome.py +++ b/tests/test_request_outcome.py @@ -296,6 +296,11 @@ async def test_funnel_passes_canonical_record_tokens_shape() -> None: "cache_write_1h_tokens": 20, "uncached_tokens": 0, "output_tokens": 50, + # Tells cost whether cache_write_tokens was REPORTED by the provider or + # derived from the uncached portion. An inferred value is the same tokens + # as uncached_tokens, so counting it in the billed prompt total would + # double it. Defaults False for providers with disjoint buckets. + "cache_inferred": False, }