From c19e412b3356d80dece001887d4ff48b6fd5150b Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Wed, 12 Aug 2026 10:17:33 +0530 Subject: [PATCH] fix(proxy/bedrock): report uncached input tokens from backend usage, not the live-zone count (#2318) ## Description On the buffered Anthropic-backend path (Bedrock / Vertex / LiteLLM-anthropic, non-streaming) the proxy reports `uncached_input_tokens` as `0` for essentially every cached multi-turn request. The handler reads the backend's Anthropic-shaped `usage` and then re-derives the uncached count from a re-tokenized live-zone count: ```python usage = backend_response.body.get("usage", {}) ... attempted_input_tokens = tokenizer.count_messages( original_client_messages[frozen_message_count:] # the LIVE ZONE only ) ... uncached_input_tokens = max(0, attempted_input_tokens - cr_tokens - cw_tokens) ``` `attempted_input_tokens` is deliberately the **live-zone** token count (the new-turn messages after the frozen prefix), kept as the denominator for the active-compression ratio. It is not the full request size. Subtracting the whole-request cache metrics (`cache_read` + `cache_creation`) from it is nonsensical: on any turn whose cached prefix is larger than the new turn -- the normal multi-turn case -- `attempted_input_tokens - cr - cw` goes negative and `max(0, ...)` clamps it to `0`. So the uncached input, which feeds the cost/uncached dashboards, is reported as `0`. Meanwhile the backend already computes the correct value. `_anthropic_usage_from_litellm` (added in #1345) sets: ```python "input_tokens": max(prompt_tokens - cache_read - cache_write, 0), ``` i.e. `usage.input_tokens` is exactly the uncached input, in Anthropic semantics. The direct-API path already uses it (`uncached_input_tokens = usage.get("input_tokens", 0)`); the backend path was the one re-deriving it. ## Fix Prefer the backend's `usage.input_tokens`, matching the direct-API path -- but guard on the backend actually reporting it, so a backend that omits `input_tokens` (or sends `null`) does not silently record `uncached=0`: ```python _reported_input_tokens = usage.get("input_tokens") if _reported_input_tokens is not None: uncached_input_tokens = int(_reported_input_tokens) else: # Backend did not report it: fall back to the live-zone derivation, # which is never worse than the previous behaviour. uncached_input_tokens = max(0, attempted_input_tokens - cr_tokens - cw_tokens) ``` A plain `usage.get("input_tokens", 0)` would have re-introduced the `0` on any backend that doesn't translate the prompt-token field; the guard keeps the authoritative value when present and the old estimate otherwise. `attempted_input_tokens` is unchanged and still used as the compression-ratio denominator. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/proxy/handlers/anthropic.py`: set `uncached_input_tokens` from `usage.input_tokens` on the buffered anthropic-backend path when the backend reports it; otherwise fall back to the prior live-zone derivation. - `tests/test_backend_nonstreaming_cache_metrics.py`: added two tests driving the buffered path -- one asserting the recorded `RequestOutcome.uncached_input_tokens == usage.input_tokens` (with a live zone far smaller than the cache), and one asserting that when the backend omits `input_tokens` the value falls back to the non-zero live-zone derivation instead of collapsing to `0`. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`) - [x] New tests added for the fix and the fallback guard ### Test Output ```text # Fail-before, primary fix (old max(0, attempted - cr - cw)): tests/..::test_anthropic_backend_nonstreaming_uncached_from_usage_input_tokens -> uncached=0, expected 1000 (FAIL) # Fail-before, safety guard (naive usage.get("input_tokens", 0)): tests/..::test_anthropic_backend_nonstreaming_uncached_falls_back_when_input_tokens_absent -> assert 0 > 0 (FAIL) # Pass-after (guarded fix): tests/test_backend_nonstreaming_cache_metrics.py 6 passed # uvx ruff@0.15.17 check -> All checks passed! # uvx mypy@1.20.2 headroom/proxy/handlers/anthropic.py -> Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12.11, project venv, pytest 9.1.1, ruff 0.15.17 and mypy 1.20.2 via uvx. - Steps: drove the buffered anthropic-backend path end to end via `create_app` + FastAPI `TestClient` with a mock `AnyLLMBackend` returning an Anthropic-shaped body, and spied on `HeadroomProxy._record_request_outcome` to capture the recorded `RequestOutcome`. With `usage.input_tokens=1000`, `cache_read=500`, `cache_write=200` and a two-token live zone, the old derivation recorded `uncached=0`; the fix records `1000`. With `input_tokens` omitted from `usage`, the naive default records `0` while the guarded fallback records the non-zero live-zone count. - Observed result: `RequestOutcome.uncached_input_tokens` now reflects the real uncached input on cached backend turns, and never regresses below the previous estimate when a backend omits the field. - Not tested: a live Bedrock/Vertex call (no cloud credentials here). The value flows through the same `RequestOutcome` funnel the proxy uses for cost/telemetry, exercised directly. ## 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 - [ ] 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 - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md`: it is generated by release-please from my Conventional Commit PR title ## Additional Notes Rebased onto current `main` and squashed to a single commit. The fallback guard is the only behavioural difference from a plain "use `usage.input_tokens`" change: it ensures the fix cannot regress a backend that doesn't report the field back to `uncached=0`. --- headroom/proxy/handlers/anthropic.py | 25 +++- ...test_backend_nonstreaming_cache_metrics.py | 130 ++++++++++++++++++ 2 files changed, 152 insertions(+), 3 deletions(-) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 9e79e3a27..362209efc 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2863,9 +2863,28 @@ class AnthropicHandlerMixin: cw_5m_tokens, cw_1h_tokens = self._extract_anthropic_cache_ttl_metrics( usage ) - uncached_input_tokens = max( - 0, attempted_input_tokens - cr_tokens - cw_tokens - ) + # Prefer the backend's Anthropic-shaped ``input_tokens``, + # which is already the uncached count (prompt tokens minus + # cache read/creation; see ``_anthropic_usage_from_litellm`` + # / #1345), matching the direct-API path below. Re-deriving + # it as ``attempted_input_tokens - cache`` is wrong whenever + # the backend reports it: ``attempted_input_tokens`` is the + # live-zone tokenizer count kept for the compression-ratio + # denominator, not the full request size, so on any turn + # whose cached prefix is larger than the new turn it + # underflows and ``max(0, ...)`` clamps uncached to 0. + # + # Guard on the backend actually reporting it: a backend that + # omits ``input_tokens`` (or sends null) must not silently + # report uncached=0 — fall back to the live-zone derivation, + # which is never worse than the previous behaviour. + _reported_input_tokens = usage.get("input_tokens") + if _reported_input_tokens is not None: + uncached_input_tokens = int(_reported_input_tokens) + else: + uncached_input_tokens = max( + 0, attempted_input_tokens - cr_tokens - cw_tokens + ) # Update prefix cache tracker for next turn. Mirrors the # direct-Anthropic-API branch below (~line 3011) — without diff --git a/tests/test_backend_nonstreaming_cache_metrics.py b/tests/test_backend_nonstreaming_cache_metrics.py index 9dbbb6c77..6eb87e44d 100644 --- a/tests/test_backend_nonstreaming_cache_metrics.py +++ b/tests/test_backend_nonstreaming_cache_metrics.py @@ -362,3 +362,133 @@ def test_anthropic_backend_nonstreaming_perf_zeros_when_upstream_omits_cache_usa handler = log_handle[0] cr, cw, chp = _find_perf_record(handler.records) assert (cr, cw, chp) == (0, 0, 0) + + +def test_anthropic_backend_nonstreaming_uncached_from_usage_input_tokens() -> None: + """The buffered anthropic-backend path must report uncached input tokens + from the backend's ``usage.input_tokens`` (which is already prompt minus + cache), not re-derive it from the live-zone tokenizer count. + + The old code computed ``uncached = attempted_input_tokens - cache_read - + cache_write``, where ``attempted_input_tokens`` is the small live-zone token + count kept for the compression-ratio denominator. On any turn whose cached + prefix is larger than the new turn that underflows to 0, so uncached input + was reported as 0. Here ``usage.input_tokens=1000`` while the live zone is a + couple of tokens, so the two behaviours are distinguishable. + """ + from headroom.proxy.server import HeadroomProxy + + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + backend="anyllm", + anyllm_provider="anthropic", + ) + body = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet-20241022", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 1000, + "output_tokens": 50, + "cache_read_input_tokens": 500, + "cache_creation_input_tokens": 200, + }, + } + backend = _make_anthropic_backend(body) + + captured: list[Any] = [] + orig_record = HeadroomProxy._record_request_outcome + + async def _spy(self, outcome): # noqa: ANN001, ANN202 + captured.append(outcome) + return await orig_record(self, outcome) + + with ( + patch("headroom.proxy.server.AnyLLMBackend", return_value=backend), + patch.object(HeadroomProxy, "_record_request_outcome", _spy), + ): + app = create_app(config) + with TestClient(app) as client: + resp = client.post( + "/v1/messages", + json={ + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 64, + }, + headers={"x-api-key": "sk-ant-test", "anthropic-version": "2023-06-01"}, + ) + assert resp.status_code == 200, resp.text[:200] + + assert captured, "expected a recorded RequestOutcome" + outcome = captured[-1] + assert outcome.uncached_input_tokens == 1000 + assert outcome.cache_read_tokens == 500 + assert outcome.cache_write_tokens == 200 + + +def test_anthropic_backend_nonstreaming_uncached_falls_back_when_input_tokens_absent() -> None: + """When the backend omits ``input_tokens``, uncached must NOT collapse to 0. + + ``usage.input_tokens`` is authoritative when present, but a backend that + does not report it must fall back to the live-zone derivation rather than + silently record uncached=0 (which is what ``usage.get("input_tokens", 0)`` + would do). With no cache counters the derivation is just the live-zone + tokenizer count, so the recorded value is a positive estimate, not 0. + """ + from headroom.proxy.server import HeadroomProxy + + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + backend="anyllm", + anyllm_provider="anthropic", + ) + # No ``input_tokens`` in usage — only output. A real backend that fails to + # translate the prompt-token field lands here. + body = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet-20241022", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": {"output_tokens": 50}, + } + backend = _make_anthropic_backend(body) + + captured: list[Any] = [] + orig_record = HeadroomProxy._record_request_outcome + + async def _spy(self, outcome): # noqa: ANN001, ANN202 + captured.append(outcome) + return await orig_record(self, outcome) + + with ( + patch("headroom.proxy.server.AnyLLMBackend", return_value=backend), + patch.object(HeadroomProxy, "_record_request_outcome", _spy), + ): + app = create_app(config) + with TestClient(app) as client: + resp = client.post( + "/v1/messages", + json={ + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "user", "content": "hello there"}], + "max_tokens": 64, + }, + headers={"x-api-key": "sk-ant-test", "anthropic-version": "2023-06-01"}, + ) + assert resp.status_code == 200, resp.text[:200] + + assert captured, "expected a recorded RequestOutcome" + outcome = captured[-1] + # No input_tokens reported and no cache: the live-zone derivation yields the + # non-zero token count of the new turn, never the 0 the naive default gave. + assert outcome.uncached_input_tokens > 0