From 313c290df96ca58a19ea0f79c67f5b71bb5f4d60 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Mon, 20 Jul 2026 10:45:22 +0530 Subject: [PATCH] fix(proxy/openai): None-guard usage token counts on the chat path (#2431) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `handle_openai_chat` reads token counts from the response usage to record metrics and update the prefix tracker: ```python usage = backend_response.body.get("usage", {}) output_tokens = usage.get("completion_tokens", 0) total_input_tokens = usage.get("prompt_tokens", optimized_tokens) ``` `.get(key, default)` only falls back when the key is **absent**. When an OpenAI-compatible backend emits a key with a **null** value (providers do this on a stopped or empty turn, the same shape that caused the Gemini crash in #2347), `.get` returns `None`. That `None` then flows into: - `_infer_openai_cache_write_tokens(total_input_tokens, cache_read_tokens)` → `max(input_tokens - cache_read_tokens, 0)` (a `None - int` → `TypeError`), - `uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens - cache_write_tokens)`, and - `RequestOutcome(output_tokens=..., optimized_tokens=...)`, whose fields are `int` and which the metrics recorder increments. Both chat usage-extraction sites are affected. On the direct-provider branch the arithmetic runs **outside** the surrounding `try`, so a single such response raises an uncaught `TypeError` and 500s the request; on the backend branch it corrupts outcome recording. ## Fix Coerce the three counts with the existing module-level `_usage_int` guard (`max(int(value), 0)`, 0 on failure) at both sites, matching the streaming path, the already-guarded cache keys in the same block (`usage.get("cache_read_input_tokens", 0) or 0`), and the Gemini fix in #2347. A normal integer usage is unchanged; only a null (or absent) value now becomes the fallback/0. `prompt_tokens` keeps its `optimized_tokens` fallback so our own input estimate is used when the count is missing. ## 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/openai.py`: `_usage_int`-guard `completion_tokens` / `prompt_tokens` / `cached_tokens` at both non-streaming usage-extraction sites in `handle_openai_chat`. - `tests/test_proxy/test_openai_chat_savings_profile.py`: regression driving a `/v1/chat/completions` request whose backend usage reports null `prompt_tokens` / `completion_tokens`, asserting a 200 instead of a crash. ## Testing - [x] 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 $ python -m pytest tests/test_proxy/test_openai_chat_savings_profile.py -q 2 passed # with the fix reverted, the new test fails (the null-usage response 500s): $ git stash push -- headroom/proxy/handlers/openai.py $ python -m pytest tests/test_proxy/test_openai_chat_savings_profile.py::test_chat_completions_survives_null_usage_token_counts -q 1 failed $ uvx ruff@0.15.17 check headroom/proxy/handlers/openai.py tests/test_proxy/test_openai_chat_savings_profile.py All checks passed! $ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/handlers/openai.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, project venv (`uv sync --extra proxy`), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`, pytest in the venv. - Exact command / steps: ran the new FastAPI `TestClient` regression, which drives the real `handle_openai_chat` through a mock backend returning `usage: {prompt_tokens: null, completion_tokens: null, total_tokens: null}`; then reverted only `openai.py` and re-ran the same test. - Observed result: with the fix the request returns 200; with the fix reverted the same request fails (the null count reaches the `max(...)` arithmetic and outcome recording). Ran against the actual handler via the app. - Not tested: a live third-party OpenAI-compatible gateway emitting null usage. ## 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 - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable --- headroom/proxy/handlers/openai.py | 22 +++++-- .../test_openai_chat_savings_profile.py | 61 +++++++++++++++++++ 2 files changed, 78 insertions(+), 5 deletions(-) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 455132142..39926a430 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -3665,8 +3665,15 @@ class OpenAIHandlerMixin: # cache stats from the LAST upstream call. total_latency = (time.time() - start_time) * 1000 usage = backend_response.body.get("usage", {}) - output_tokens = usage.get("completion_tokens", 0) - total_input_tokens = usage.get("prompt_tokens", optimized_tokens) + # `.get(key, default)` only falls back when the key is + # absent; a present-but-null count (some OpenAI-compatible + # backends emit these on a stopped/empty turn) would return + # None and crash the downstream `max(...)` arithmetic and the + # int-typed outcome/metrics. `_usage_int` coerces both cases, + # matching the streaming path and the guarded cache keys below + # (same class as the gemini fix in #2347). + output_tokens = _usage_int(usage.get("completion_tokens")) + total_input_tokens = _usage_int(usage.get("prompt_tokens")) or optimized_tokens # Cache stats: prefer the Anthropic/Bedrock top-level # keys when present (authoritative). Fall back to @@ -3976,12 +3983,17 @@ class OpenAIHandlerMixin: try: resp_json = response.json() usage = resp_json.get("usage", {}) - total_input_tokens = usage.get("prompt_tokens", optimized_tokens) - output_tokens = usage.get("completion_tokens", 0) + # Coerce present-but-null counts: the arithmetic below + # (`_infer_openai_cache_write_tokens`, `max(...)`) runs + # outside this try, so a null `prompt_tokens`/`cached_tokens` + # would otherwise raise an uncaught TypeError and 500 the + # request (same class as the gemini fix in #2347). + total_input_tokens = _usage_int(usage.get("prompt_tokens")) or optimized_tokens + output_tokens = _usage_int(usage.get("completion_tokens")) # OpenAI returns cached_tokens in prompt_tokens_details # These are charged at 50% of the input price prompt_details = usage.get("prompt_tokens_details") or {} - cache_read_tokens = prompt_details.get("cached_tokens", 0) + cache_read_tokens = _usage_int(prompt_details.get("cached_tokens")) except (KeyError, TypeError, AttributeError) as e: logger.debug( f"[{request_id}] Failed to extract cached tokens from OpenAI response: {e}" diff --git a/tests/test_proxy/test_openai_chat_savings_profile.py b/tests/test_proxy/test_openai_chat_savings_profile.py index 3cf8ddb2a..827d2d569 100644 --- a/tests/test_proxy/test_openai_chat_savings_profile.py +++ b/tests/test_proxy/test_openai_chat_savings_profile.py @@ -50,6 +50,67 @@ def _make_mock_backend() -> MagicMock: return backend +def _make_mock_backend_with_usage(usage: dict) -> MagicMock: + backend = MagicMock() + backend.name = "anyllm-openai" + backend.send_openai_message = AsyncMock( + return_value=BackendResponse( + body={ + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": usage, + }, + status_code=200, + headers={"content-type": "application/json"}, + ) + ) + return backend + + +def test_chat_completions_survives_null_usage_token_counts(): + """A backend that reports present-but-null token counts must not 500. + + `.get(key, default)` returns None for a null value, and the chat path + feeds those counts into `max(...)`/int-typed metrics. Without coercion a + single such response crashes the request and its outcome recording + (same class as the gemini fix in #2347). + """ + config = ProxyConfig( + optimize=True, + cache_enabled=False, + rate_limit_enabled=False, + backend="anyllm", + anyllm_provider="openai", + ) + + # prompt_tokens / completion_tokens present but explicitly null. + mock_backend = _make_mock_backend_with_usage( + {"prompt_tokens": None, "completion_tokens": None, "total_tokens": None} + ) + with patch("headroom.proxy.server.AnyLLMBackend", return_value=mock_backend): + app = create_app(config) + with TestClient(app) as client: + resp = client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + "stream": False, + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert resp.status_code == 200, resp.text + + def test_chat_completions_threads_savings_profile_kwargs_into_apply(): """With HEADROOM_SAVINGS_PROFILE=agent-90, the chat path must pass the profile knobs (compress_user_messages, target_ratio, ...) to apply()."""