mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
b97c7c6e99
## Description The non-streaming Gemini `generateContent` finalizer builds its `RequestOutcome` with `optimized_tokens` set to Gemini's own `promptTokenCount` (the provider's tokenizer scale, which correctly feeds billing and the dashboard), while `original_tokens` stays a local estimator count. Those two are on different rulers. Every delta the beacon derives from the pair is a same-ruler difference: `tokens_saved`, `tokens_inflated`, `attempted_input_tokens`, and the beacon's `eligible_pct` / `yield_pct`. When Gemini counts the forwarded prompt higher than our local estimator does, `attempted_input_tokens` (which is `optimized_tokens + tokens_saved`) exceeds the local `original_tokens`, and the request ships a structurally-impossible `eligible_pct > 100` plus a phantom `tokens_inflated`. This is the exact class of bug #2756 removed, on a path #2756 did not touch: it fixed the non-streaming OpenAI handler, and the streaming finalizer (`_finalize_stream_response`) already guards against it by lifting the baseline onto the provider scale. The non-streaming Gemini path had neither treatment. The fix mirrors the streaming finalizer's already-tested handling: when a provider count is present, lift the baseline to `max(original_tokens, promptTokenCount + tokens_saved)` so `attempted_input_tokens <= original_tokens` holds and `tokens_inflated` collapses to 0. It is guarded on a present count, so a null or absent `promptTokenCount` leaves the local baseline untouched and the existing zero-usage preservation test still holds. `optimized_tokens` still carries the provider count, so billing and the dashboard are unchanged. Closes # ## 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/gemini.py` (`handle_gemini_request`, non-streaming `generateContent` branch): compute `effective_original_tokens = max(original_tokens, total_input_tokens + tokens_saved)` when `total_input_tokens > 0` (else keep `original_tokens`), and pass it as the outcome's `original_tokens`. Mirrors the streaming finalizer's provider-usage handling. - `tests/test_proxy/test_gemini_savings_profile.py`: added `test_gemini_provider_count_above_local_estimate_does_not_inflate_eligible`, which drives a request where Gemini's `promptTokenCount` (150) exceeds the local post-compression count (80), and asserts `attempted_input_tokens <= original_tokens`, `tokens_inflated == 0`, the provider count is still carried in `optimized_tokens`, and the baseline is lifted to 170. ## 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 # Fail-before (source fix stashed, new test kept): tests/test_proxy/test_gemini_savings_profile.py::test_gemini_provider_count_above_local_estimate_does_not_inflate_eligible FAILED assert outcome.attempted_input_tokens <= outcome.original_tokens AssertionError: assert 170 <= 100 # Pass-after (fix applied): tests/test_proxy/test_gemini_savings_profile.py::test_gemini_provider_count_above_local_estimate_does_not_inflate_eligible PASSED # Full file + related outcome suites: tests/test_proxy/test_gemini_savings_profile.py tests/test_proxy_gemini_native_integration.py tests/test_request_outcome.py tests/test_outcome_token_scale.py 47 passed, 18 skipped # uvx ruff@0.15.17 check -> All checks passed! # uvx mypy@1.20.2 headroom/proxy/handlers/gemini.py -> Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12.11, project venv (litellm installed), pytest 9.1.1 with pytest-asyncio 1.4.0 (asyncio_mode=auto per pyproject), ruff 0.15.17 and mypy 1.20.2 via uvx. - Exact command / steps: confirmed the streaming sibling already lifts the baseline (`_finalize_stream_response` in `headroom/proxy/handlers/streaming.py` sets `effective_original_tokens = max(original_tokens, provider_input_tokens + tokens_saved)` for openai/gemini), then fail-before with `git stash push headroom/proxy/handlers/gemini.py` and `python -m pytest tests/test_proxy/test_gemini_savings_profile.py -k inflate_eligible` (the assertion fails with `170 <= 100`, i.e. eligible_pct 170%), then pass-after with `git stash pop` and rerunning (passes), then the full file plus the outcome suites (47 passed, 18 skipped). - Observed result: with Gemini reporting `promptTokenCount=150` against a local post-compression count of 80 (saved 20), the outcome now reports `original_tokens=170`, `attempted_input_tokens=170` (so `eligible_pct <= 100`) and `tokens_inflated=0`, while `optimized_tokens` stays 150 so billing and the dashboard are unchanged. Before the fix the same request reported `original_tokens=100`, `attempted_input_tokens=170` (eligible_pct 170%) and `tokens_inflated=50`. - Not tested: a live streamed call to real Gemini/Vertex (no provider credentials in this environment). The provider-count-above-local case is reproduced with a mock response mirroring Gemini's `usageMetadata` shape, and the baseline-lift it mirrors is existing, tested code on the streaming path. ## 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 - [x] I did **not** edit `CHANGELOG.md`: it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Additional Notes Docs and manual testing are N/A: this aligns the non-streaming Gemini finalizer with the already-correct streaming finalizer, no API surface change. The baseline lift is guarded on a present provider count, so the existing zero-usage preservation test (`test_gemini_zero_usage_prompt_count_is_preserved`) is unaffected: a null or zero `promptTokenCount` keeps the local baseline and leaves `optimized_tokens` at 0. |
||
|---|---|---|
| .. | ||
| 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_bedrock_sse_ping.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 | ||