mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
313c290df9
## 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
|
||
|---|---|---|
| .. | ||
| test_anthropic_buffered_timeout.py | ||
| test_anthropic_ccr_deferred_injection.py | ||
| test_anthropic_ccr_raise.py | ||
| test_anthropic_streaming_ccr_retrieve.py | ||
| test_anthropic_upstream_header.py | ||
| test_background_compression.py | ||
| test_bedrock_passthrough.py | ||
| test_cc_switch_reconciler.py | ||
| test_ccr_frozen_prefix_coupling.py | ||
| test_compression_failure_action.py | ||
| test_compression_timeout_config.py | ||
| test_compute_turn_id.py | ||
| test_gemini_savings_profile.py | ||
| test_header_safe_transforms.py | ||
| test_mcp_stats_aggregation.py | ||
| test_model_router.py | ||
| test_model_router_wiring.py | ||
| test_openai_backend_path.py | ||
| test_openai_chat_savings_profile.py | ||
| test_openai_responses_ccr.py | ||
| test_openai_stream_usage_option.py | ||
| test_openai_transport_path_prefix.py | ||
| test_openai_upstream_header.py | ||
| test_phase3_byte_identity.py | ||
| test_request_logger.py | ||
| test_settings_fresh_process_precedence.py | ||
| test_settings_store.py | ||
| test_transformations_feed.py | ||