github-featured/chopratejas--headroom
2
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix(proxy/bedrock): report uncached input tokens from backend usage, not the live-zone count (#2318)

Browse Source 
## 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`.
This commit is contained in:
Abhay Singh
2026-08-12 10:17:33 +05:30
committed by GitHub
parent 1f2c681c0b
commit c19e412b33
2 changed files with 152 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
headroom/proxy/handlers/anthropic.py
+22 -3
View File
@@ -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
tests/test_backend_nonstreaming_cache_metrics.py
+130
View File
@@ -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