mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
fix(cost): send litellm the total prompt so --budget stops seeing $0 (#2757)
## 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 <same three> 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)
This commit is contained in:
Tejas Chopra
GitHub
parent
06add9e9d8
commit
a033ac4176
5 changed files with 160 additions and 2 deletions
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
124
tests/test_cost_budget_total_prompt.py
Normal file
124
tests/test_cost_budget_total_prompt.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue