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.
2026-08-17 03:34:31 +05:30
parent 9d370592b0
commit a01897c791
2 changed files with 146 additions and 3 deletions
@@ -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)

@@ -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