From a01897c791f4bb6471defafd560d29d491eb2df8 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Mon, 17 Aug 2026 03:34:31 +0530 Subject: [PATCH] fix(proxy/gemini): guard CCR continuation usage against present-null counts (#3035) ## Description On the Gemini native `generateContent` path, a successful (200) response that triggers a CCR retrieval continuation is masked as a synthetic 502 when the continuation response carries a present-null usage count. `handle_gemini_generate_content` reads `usageMetadata` at three sites. The initial-response site and the non-CCR site both guard against Gemini returning a present-null count (a key present with a JSON `null`, which `.get(key, default)` returns as `None` rather than the default). The CCR-continuation site read the continuation's `usageMetadata` with a bare `.get(key, prior)`: ```python total_input_tokens = usage.get("promptTokenCount", total_input_tokens) output_tokens = usage.get("candidatesTokenCount", output_tokens) cache_read_tokens = usage.get("cachedContentTokenCount", cache_read_tokens) ``` When the continuation turn reports `"promptTokenCount": null`, `total_input_tokens` becomes `None`, and the following `uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens)` and the `total_input_tokens > 0` baseline guard raise `TypeError`. The method's outer `except Exception` then returns a 502 JSONResponse and records a provider failure, so a genuinely successful upstream turn is reported to the client as a 502. ## Fix Read the continuation usage through the same `_usage_int` guard the two sibling sites use, keeping the pre-continuation count as the fallback (`_usage_int(value, default)` returns `default` when `value is None`). Behavior is otherwise unchanged: a present, valid count is still used, and an absent count still falls back to the pre-continuation value. Fixes #3034 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature - [ ] Breaking change - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/proxy/handlers/gemini.py` (`handle_gemini_generate_content`, CCR-continuation branch): read `promptTokenCount` / `candidatesTokenCount` / `cachedContentTokenCount` through `_usage_int(..., prior)` instead of a bare `.get(key, prior)`. - `tests/test_gemini_ccr_continuation_usage.py`: drive the handler through a CCR continuation whose `usageMetadata` counts are present-null; assert the client gets 200 (not 502), no provider failure is recorded, and the pre-continuation count survives as the fallback. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`) - [x] New tests added ### Test Output ```text tests/test_gemini_ccr_continuation_usage.py 1 passed tests/test_gemini_nonjson_status.py tests/test_gemini_compression_offload.py tests/test_proxy_gemini_native_integration.py (all pass; platform-skipped cases skipped) # uvx ruff@0.15.22 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, pytest 9.1.1, ruff 0.15.22 and mypy 1.20.2 via uvx. - Exact command / steps: ran `python -m pytest tests/test_gemini_ccr_continuation_usage.py -q` (pass-after); proved fail-before by `git stash`-ing only the `gemini.py` change and re-running (the test failed with `assert 502 == 200` and the captured log `TypeError: unsupported operand type(s) for -: 'NoneType' and 'NoneType'` at `gemini.py`), then restored the fix and re-ran green; ran the surrounding Gemini suite (`test_gemini_nonjson_status.py`, `test_gemini_compression_offload.py`, `test_proxy_gemini_native_integration.py`); then `uvx ruff@0.15.22 check` and `uvx mypy@1.20.2 headroom/proxy/handlers/gemini.py`. - Observed result: with the fix a CCR continuation carrying a present-null `promptTokenCount` returns 200 to the client and records the outcome with the pre-continuation count (100) instead of raising `TypeError` and returning a synthetic 502. - Not tested: a live Gemini session that both triggers a CCR retrieval continuation and receives a present-null continuation usage payload (needs a real safety-blocked continuation). The contract is verified at the handler with the same stub pattern the existing `test_gemini_nonjson_status.py` uses. ## Runtime Rollout Safety - Rollout-managed feature(s): none. This is the always-on Gemini native `generateContent` request path, not a rollout-channel-gated feature. - Minimum rollout channel: N/A (no rollout-managed behavior). - Stable/default behavior changed: yes, as a bug fix. A CCR continuation with a present-null usage count now returns the real 200 instead of a synthetic 502; all other cases (present valid count, absent count) are unchanged. - Kill switch / disable path: N/A. There is no behavioral toggle; the change only makes the existing continuation path null-safe. - Unsafe override required: no. - Qualification impact: brings the CCR-continuation usage extraction to parity with the two sibling sites that already guard present-null counts; no routing, compression, or pricing change. - Rollback path: revert this PR; the continuation site returns to the bare `.get(key, prior)` read. ## 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 The unguarded site was introduced in #2253 (native CCR retrieval); the present-null guard on the sibling sites landed separately and did not extend to it. The fix reuses the existing `_usage_int` helper so all three Gemini usage-extraction sites now handle present-null identically. --- headroom/proxy/handlers/gemini.py | 19 ++- tests/test_gemini_ccr_continuation_usage.py | 130 ++++++++++++++++++++ 2 files changed, 146 insertions(+), 3 deletions(-) create mode 100644 tests/test_gemini_ccr_continuation_usage.py diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index dfdc3e90f..35e249b39 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -876,9 +876,22 @@ class GeminiHandlerMixin: resp_json = final_resp_json response_content = json.dumps(resp_json).encode() usage = resp_json.get("usageMetadata", {}) - total_input_tokens = usage.get("promptTokenCount", total_input_tokens) - output_tokens = usage.get("candidatesTokenCount", output_tokens) - cache_read_tokens = usage.get("cachedContentTokenCount", cache_read_tokens) + # A CCR continuation response can carry a present-null count + # (e.g. a safety-blocked continuation turn), where + # ``.get(key, prior)`` returns None rather than the prior + # value, and the ``max(0, prompt - cache_read)`` / + # ``total_input_tokens > 0`` arithmetic below would then raise + # TypeError and the outer handler would mask a successful 200 + # as a synthetic 502. Guard with ``_usage_int`` (keeping the + # pre-continuation count as the fallback), mirroring the two + # sibling extraction sites above. + total_input_tokens = _usage_int( + usage.get("promptTokenCount"), total_input_tokens + ) + output_tokens = _usage_int(usage.get("candidatesTokenCount"), output_tokens) + cache_read_tokens = _usage_int( + usage.get("cachedContentTokenCount"), cache_read_tokens + ) uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens) diff --git a/tests/test_gemini_ccr_continuation_usage.py b/tests/test_gemini_ccr_continuation_usage.py new file mode 100644 index 000000000..8cf0649a8 --- /dev/null +++ b/tests/test_gemini_ccr_continuation_usage.py @@ -0,0 +1,130 @@ +"""A Gemini CCR continuation with a present-null usage count must not 502. + +The initial-response and non-CCR extraction sites guard against Gemini +returning a *present-null* ``promptTokenCount`` (a key that is present with a +JSON ``null`` value, which ``.get(key, default)`` returns as ``None`` rather +than the default). The CCR-continuation site re-read the continuation's +``usageMetadata`` with a bare ``.get(key, prior)`` and skipped that guard, so a +present-null count on the continuation turned the ``max(0, prompt - cache_read)`` +arithmetic into ``None`` math, raised ``TypeError``, and the outer handler +masked a successful 200 as a synthetic 502. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from headroom.proxy.handlers.gemini import GeminiHandlerMixin + + +class _FakeRequest: + def __init__(self) -> None: + self.headers: dict[str, str] = {} + self.query_params: dict[str, str] = {} + self.url = SimpleNamespace(path="/v1beta/models/gemini-pro:generateContent", query="") + + +class _CcrToolCallResponse: + """Initial 200 carrying a CCR tool call and a valid promptTokenCount.""" + + status_code = 200 + content = json.dumps( + { + "candidates": [ + {"content": {"parts": [{"functionCall": {"name": "headroom_retrieve"}}]}} + ], + "usageMetadata": {"promptTokenCount": 100, "candidatesTokenCount": 5}, + } + ).encode() + headers = {"content-type": "application/json"} + + def json(self) -> object: + return json.loads(self.content) + + +class _CcrConfig: + enabled = True + + +class _CcrHandler: + """Stub CCR handler whose continuation reports a present-null usage count.""" + + config = _CcrConfig() + + def has_ccr_tool_calls(self, resp_json, provider) -> bool: # noqa: ANN001 + return True + + async def handle_response(self, resp_json, contents, native_fns, api_call_fn, provider): # noqa: ANN001, ANN201 + return { + "candidates": [{"content": {"parts": [{"text": "resolved"}]}}], + # The continuation turn omits real counts as JSON null. + "usageMetadata": { + "promptTokenCount": None, + "candidatesTokenCount": None, + "cachedContentTokenCount": None, + }, + } + + def residual_ccr_status(self, final_resp_json, provider): # noqa: ANN001, ANN201 + return None # not RESIDUAL_CCR_ERROR + + +class _FakeMetrics: + def __init__(self) -> None: + self.failed: list[str] = [] + + async def record_failed(self, *, provider: str, model: str = "") -> None: + self.failed.append(f"{provider}:{model}") + + +class _Handler(GeminiHandlerMixin): + GEMINI_API_URL = "https://gemini.example" + + def __init__(self) -> None: + self.memory_handler = None + self.rate_limiter = None + self.usage_reporter = None + self.config = SimpleNamespace( + optimize=False, + anthropic_pre_upstream_memory_context_timeout_seconds=0.1, + ) + self.metrics = _FakeMetrics() + self.ccr_response_handler = _CcrHandler() + self.outcomes: list = [] + + async def _next_request_id(self) -> str: + return "req-ccr-1" + + async def _retry_request(self, method, url, headers, body): # noqa: ANN001, ANN201 + return _CcrToolCallResponse() + + async def _record_request_outcome(self, outcome) -> None: # noqa: ANN001 + self.outcomes.append(outcome) + + async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201 + return SimpleNamespace(), 100 + + +@pytest.mark.asyncio +async def test_ccr_continuation_present_null_usage_does_not_502( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def payload(request): # noqa: ANN001, ANN201 + return {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]} + + monkeypatch.setattr("headroom.proxy.helpers._read_request_json", payload) + + handler = _Handler() + response = await handler.handle_gemini_generate_content(_FakeRequest(), "gemini-pro") + + # Before the fix this raised TypeError on the None arithmetic and the outer + # handler returned a synthetic 502 with a recorded failure. + assert response.status_code == 200 + assert handler.metrics.failed == [] + assert handler.outcomes[0].status_code == 200 + # The pre-continuation count (100) survives as the fallback. + assert handler.outcomes[0].optimized_tokens == 100 + assert handler.outcomes[0].cache_read_tokens == 0