From 12f9f58cb3dcfc67af1238424d404d8dd9bad1dd Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Wed, 12 Aug 2026 10:18:43 +0530 Subject: [PATCH] fix(backends/litellm): None-guard core token counts in OpenAI usage block (#2324) ## Description `LiteLLMBackend.send_openai_message` builds the OpenAI-shape response body. The core token counts are copied straight off LiteLLM's `Usage` object with no guard, even though the cache fields immediately below already use the defensive `int(getattr(..., 0) or 0)` form: ```python usage_block: dict[str, Any] = { "prompt_tokens": response.usage.prompt_tokens, "completion_tokens": response.usage.completion_tokens, "total_tokens": response.usage.total_tokens, } # Defensive getattr right below: cache_read = int(getattr(response.usage, "cache_read_input_tokens", 0) or 0) cache_write = int(getattr(response.usage, "cache_creation_input_tokens", 0) or 0) ``` A provider can leave any of `prompt_tokens` / `completion_tokens` / `total_tokens` as `None` on the `Usage` object. That `None` then lands in `response.body["usage"]`, and the backend-routed OpenAI handler reads it straight into arithmetic and the outcome ledger: ```python usage = backend_response.body.get("usage", {}) output_tokens = usage.get("completion_tokens", 0) # present key -> None, not the default total_input_tokens = usage.get("prompt_tokens", optimized_tokens) # present key -> None ... uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens - cache_write_tokens) # None - int -> TypeError ... RequestOutcome(..., output_tokens=output_tokens, ...) # declared int; None crashes recording (e.g. prometheus += ) ``` So an OpenAI-format request routed through a `--backend` (Bedrock / Vertex / LiteLLM) whose provider returns a `None` count crashes on the `max(0, None - ...)` subtraction, or later in outcome recording. `.get(key, default)` does not help here because the key is present with a `None` value, so the default never applies. This is the same class of bug as the Anthropic-shape mapping and is fixed the same way. ## Fix Coerce the three counts to `int` with the same defensive form already used for the cache fields two lines down: ```python usage_block: dict[str, Any] = { "prompt_tokens": int(getattr(response.usage, "prompt_tokens", 0) or 0), "completion_tokens": int(getattr(response.usage, "completion_tokens", 0) or 0), "total_tokens": int(getattr(response.usage, "total_tokens", 0) or 0), } ``` No change for a normal integer usage; only a `None` (or absent) value now becomes `0`. ## 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/backends/litellm.py`: `int`-coerce `prompt_tokens` / `completion_tokens` / `total_tokens` in the `send_openai_message` usage block. - `tests/test_backends/test_litellm_cache_stats.py`: add `test_none_core_counts_coerced_to_zero`, driving `send_openai_message` with a `None`-count usage and asserting the block emits `int` `0`s. - `CHANGELOG.md`: Bug Fixes entry. ## Testing - [ ] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uvx ruff@0.15.17 check headroom/backends/litellm.py tests/test_backends/test_litellm_cache_stats.py All checks passed! $ uvx ruff@0.15.17 format --check headroom/backends/litellm.py tests/test_backends/test_litellm_cache_stats.py 2 files already formatted $ uvx mypy@1.20.2 --ignore-missing-imports headroom/backends/litellm.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx mypy@1.20.2`. A full `pytest` OOMs this box (ML-stack import), so I reproduced the field logic with a dependency-free script and left the full pytest to CI. - Exact command / steps: modeled the OLD (bare copy) and NEW (`int(getattr(..., 0) or 0)`) field derivations for a `None` count, an integer, and zero, then simulated the two downstream operations the handler performs: `output_tokens += ...` and `max(0, prompt_tokens - read - write)`. - Observed result: OLD produced `None` and both downstream operations raised `TypeError`; NEW produced `0` and both succeeded; an integer count passed through unchanged. The added unit test drives `send_openai_message` end to end (mocked `acompletion`) and asserts the block emits `int` `0`s. - Not tested: a live LiteLLM/Bedrock request that returns `None` counts; the added test reuses the existing `_FakeUsage` / `_make_response` / mocked-`acompletion` harness in `test_litellm_cache_stats.py`. ## 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 or that my feature works - [ ] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes The "unit tests pass locally" box is unchecked because a local pytest run imports the ML stack and OOMs this box; the added test reuses the existing mocked-`acompletion` harness in `test_litellm_cache_stats.py` and runs under the normal CI pytest job, and the behavior is corroborated by the standalone proof above. --------- Co-authored-by: JerrettDavis --- headroom/backends/litellm.py | 12 ++++++--- .../test_backends/test_litellm_cache_stats.py | 27 +++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 6e3afd7e4..01c73b032 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -1389,10 +1389,16 @@ class LiteLLMBackend(Backend): # cache_creation_tokens for the OpenAI nested dialect. Surface both # so PrefixCacheTracker.update_from_response on the backend-routed # path observes a stable shape instead of branching on key presence. + # None-guard the core counts (same defensive style as the cache + # fields just below). A provider can leave any of these None on the + # Usage object; emitting None here flows into the OpenAI-shape body, + # and the backend-routed OpenAI handler reads them straight into + # arithmetic and RequestOutcome (output_tokens=..., and + # max(0, prompt_tokens - ...)), which raises TypeError on None. usage_block: dict[str, Any] = { - "prompt_tokens": response.usage.prompt_tokens, - "completion_tokens": response.usage.completion_tokens, - "total_tokens": response.usage.total_tokens, + "prompt_tokens": int(getattr(response.usage, "prompt_tokens", 0) or 0), + "completion_tokens": int(getattr(response.usage, "completion_tokens", 0) or 0), + "total_tokens": int(getattr(response.usage, "total_tokens", 0) or 0), } # Defensive getattr: LiteLLM only attaches these top-level attrs diff --git a/tests/test_backends/test_litellm_cache_stats.py b/tests/test_backends/test_litellm_cache_stats.py index 083a7a5d1..4ec6e057f 100644 --- a/tests/test_backends/test_litellm_cache_stats.py +++ b/tests/test_backends/test_litellm_cache_stats.py @@ -215,3 +215,30 @@ async def test_no_cache_fields_means_no_cache_keys_in_usage_block() -> None: assert "cache_read_input_tokens" not in body_usage assert "cache_creation_input_tokens" not in body_usage assert "prompt_tokens_details" not in body_usage + + +async def test_none_core_counts_coerced_to_zero() -> None: + """A provider can leave prompt/completion/total token counts None on the + Usage object. The OpenAI-shape usage block must emit ints, not None, so the + backend-routed OpenAI handler (which reads these straight into arithmetic + and RequestOutcome) does not crash with a TypeError.""" + usage = _FakeUsage( + prompt_tokens=None, # type: ignore[arg-type] + completion_tokens=None, # type: ignore[arg-type] + total_tokens=None, # type: ignore[arg-type] + ) + response = _make_response(usage) + + backend = _make_backend() + with patch("headroom.backends.litellm.acompletion", new_callable=AsyncMock) as mock_acomp: + mock_acomp.return_value = response + result = await backend.send_openai_message(_request_body(), {}) + + body_usage = result.body["usage"] + assert body_usage["prompt_tokens"] == 0 + assert body_usage["completion_tokens"] == 0 + assert body_usage["total_tokens"] == 0 + assert all( + isinstance(body_usage[k], int) + for k in ("prompt_tokens", "completion_tokens", "total_tokens") + )