fix(proxy/gemini): None-guard token counts from usageMetadata (#2347)

## Description

The non-streaming Gemini/Vertex handler reads token counts straight from
the response's `usageMetadata`:

```python
try:
    usage = resp_json.get("usageMetadata", {})
    total_input_tokens = usage.get("promptTokenCount", optimized_tokens)
    output_tokens = usage.get("candidatesTokenCount", 0)
    cache_read_tokens = usage.get("cachedContentTokenCount", 0)
except (...):
    ...
uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens)   # OUTSIDE the try
```

`.get(key, default)` only falls back when the key is **absent**. When
`usageMetadata` carries a key with a **null** value — which Gemini can
do on a safety-blocked turn that produced no candidates — `.get` returns
`None`. That `None` then reaches:

- `max(0, total_input_tokens - cache_read_tokens)` (a `None - int` →
`TypeError`), and
- `RequestOutcome(output_tokens=...)`, whose field is `int` and which
the metrics recorder increments (`tokens_output_total += output_tokens`
→ `TypeError`).

Both run on the success (non-`except`) path, so a single such response
crashes the request and its outcome recording. The Gemini streaming path
already guards these with a `_usage_int` helper; the non-streaming path
(two sites) did not.

## Fix

Coerce the three counts with `int(... or fallback)`, matching the
streaming `_usage_int` guard and the LiteLLM usage mappings:

```python
total_input_tokens = int(usage.get("promptTokenCount", optimized_tokens) or optimized_tokens)
output_tokens = int(usage.get("candidatesTokenCount", 0) or 0)
cache_read_tokens = int(usage.get("cachedContentTokenCount", 0) or 0)
```

No change for a normal integer usage; only a `None` (or absent) value
now becomes the fallback/0. Applied to both non-streaming
usage-extraction sites in `handlers/gemini.py`.

## 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`: `int(... or fallback)`-guard
`promptTokenCount` / `candidatesTokenCount` / `cachedContentTokenCount`
at both non-streaming usage sites.
- `tests/test_proxy/test_gemini_savings_profile.py`: add a regression
driving a `generateContent` request whose
`usageMetadata.candidatesTokenCount` is `null`, asserting a 200, an
`int` `output_tokens == 0`, and `uncached_input_tokens == 20` (the
`max(0, …)` no longer raises).
- `CHANGELOG.md`: Bug Fixes entry.

## Testing

- [ ] 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
$ uvx ruff@0.15.17 check headroom/proxy/handlers/gemini.py tests/test_proxy/test_gemini_savings_profile.py
All checks passed!
$ uvx ruff@0.15.17 format --check headroom/proxy/handlers/gemini.py tests/test_proxy/test_gemini_savings_profile.py
2 files already formatted
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/handlers/gemini.py
Success: no issues found in 1 source file
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOMs this box (ML-stack import), so I
reproduced the extraction with a dependency-free script and left the
full pytest to CI.
- Exact command / steps: modeled the OLD (bare `.get`) and NEW (`int(...
or fallback)`) derivations for a blocked response
(`candidatesTokenCount: null`, valid prompt count), a null
`promptTokenCount`, a normal response, and an absent-usage response.
- Observed result: OLD raised `TypeError` at `max(0, None - …)` for a
null prompt count and left `output_tokens = None` (which crashes the
int-typed outcome/metrics recorder) for a null candidate count; NEW
produced `(20, 0)` for the blocked case, `(15, 0)` for the null-prompt
case (the `optimized_tokens` fallback), `(60, 30)` for a normal
response, and the fallbacks for absent usage. The added
`create_app`/`TestClient` test drives the handler end to end and asserts
a 200 with `int` outcome counts.
- Not tested: a live Gemini safety-blocked response; the added test uses
a mocked `_retry_request` returning a `usageMetadata` with a null count,
matching the existing Gemini test harness in this file.

## 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
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

The "unit tests pass locally" box is unchecked because a local pytest
run imports the ML stack and OOMs this box; the added test reuses the
existing `create_app`/`TestClient` + mocked-`_retry_request` harness in
`test_gemini_savings_profile.py` and runs under the normal CI pytest
job, and the behavior is corroborated by the standalone proof above.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>

headroom/proxy/handlers/gemini.py

@@ -29,6 +29,12 @@ DEFAULT_CLOUDCODE_API_URL = "https://cloudcode-pa.googleapis.com"
 ANTIGRAVITY_DAILY_API_URL = "https://daily-cloudcode-pa.sandbox.googleapis.com"
 
 
+def _usage_int(value: Any, default: int = 0) -> int:
+    if value is None:
+        return default
+    return int(value)
+
+
 class GeminiHandlerMixin:
     """Mixin providing Gemini API handler methods for HeadroomProxy."""
 
@@ -423,9 +429,15 @@ class GeminiHandlerMixin:
                 try:
                     resp_json = response.json()
                     usage = resp_json.get("usageMetadata", {})
-                    total_input_tokens = usage.get("promptTokenCount", 0)
+                    # Gemini omits or nulls these counts on some responses
-                    output_tokens = usage.get("candidatesTokenCount", 0)
+                    # (e.g. a safety-blocked turn with no candidates). A
-                    cache_read_tokens = usage.get("cachedContentTokenCount", 0)
+                    # present-null leaves .get returning None, and the max(0,
+                    # prompt - cache_read) below (and RequestOutcome's int
+                    # output_tokens) would then raise TypeError on the non-error
+                    # path. Mirrors the streaming _usage_int guard.
+                    total_input_tokens = _usage_int(usage.get("promptTokenCount"))
+                    output_tokens = _usage_int(usage.get("candidatesTokenCount"))
+                    cache_read_tokens = _usage_int(usage.get("cachedContentTokenCount"))
                 except (json.JSONDecodeError, ValueError, KeyError, TypeError, AttributeError):
                     pass
                 await self._record_request_outcome(
@@ -666,11 +678,21 @@ class GeminiHandlerMixin:
                 try:
                     resp_json = response.json()
                     usage = resp_json.get("usageMetadata", {})
-                    total_input_tokens = usage.get("promptTokenCount", optimized_tokens)
+                    # Gemini omits or nulls these counts on some responses
-                    output_tokens = usage.get("candidatesTokenCount", 0)
+                    # (e.g. a safety-blocked turn with no candidates). A
+                    # present-null leaves .get returning None, and the max(0,
+                    # prompt - cache_read) below (plus RequestOutcome's int
+                    # output_tokens) would then raise TypeError on the non-error
+                    # path. Mirrors the streaming _usage_int guard.
+                    total_input_tokens = int(
+                        optimized_tokens
+                        if usage.get("promptTokenCount") is None
+                        else usage["promptTokenCount"]
+                    )
+                    output_tokens = _usage_int(usage.get("candidatesTokenCount"))
                     # Gemini returns cachedContentTokenCount for context-cached tokens
                     # These are charged at 10-25% of the input price depending on model
-                    cache_read_tokens = usage.get("cachedContentTokenCount", 0)
+                    cache_read_tokens = _usage_int(usage.get("cachedContentTokenCount"))
                 except (json.JSONDecodeError, ValueError, KeyError, TypeError, AttributeError) as e:
                     # A non-JSON upstream body (HTML/empty error page from an
                     # overloaded Google/Vertex frontend) makes response.json()

tests/test_proxy/test_gemini_savings_profile.py

@@ -85,3 +85,117 @@ def test_gemini_generate_content_threads_savings_profile_kwargs_into_apply():
     assert captured.get("target_ratio") == 0.10
     assert captured.get("min_tokens_to_compress") == 120
     assert captured.get("compress_system_messages") is True
+
+
+def test_gemini_null_usage_counts_do_not_crash():
+    """A Gemini response whose usageMetadata carries a null token count (e.g. a
+    safety-blocked turn with no candidates) must not crash outcome recording:
+    the counts are coerced to int, not left as None."""
+    config = ProxyConfig(
+        optimize=True,
+        cache_enabled=False,
+        rate_limit_enabled=False,
+        cost_tracking_enabled=False,
+    )
+
+    def passthrough_apply(**kwargs):
+        return SimpleNamespace(
+            messages=kwargs["messages"],
+            transforms_applied=[],
+            timing={},
+            tokens_before=10,
+            tokens_after=10,
+            waste_signals=None,
+        )
+
+    resp = MagicMock()
+    resp.status_code = 200
+    resp.headers = {"content-type": "application/json"}
+    resp.content = (
+        b'{"candidates":[{"content":{"parts":[{"text":"ok"}]}}],'
+        b'"usageMetadata":{"promptTokenCount":20,"candidatesTokenCount":null}}'
+    )
+    resp.json.return_value = {
+        "candidates": [{"content": {"parts": [{"text": "ok"}]}}],
+        "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": None},
+    }
+
+    captured: dict[str, object] = {}
+
+    async def recording_outcome(outcome):  # noqa: ANN001
+        captured["outcome"] = outcome
+
+    big = "word " * 4000
+    app = create_app(config)
+    with TestClient(app) as client:
+        proxy = client.app.state.proxy
+        proxy.openai_pipeline.apply = MagicMock(side_effect=passthrough_apply)
+        proxy._retry_request = AsyncMock(return_value=resp)
+        proxy._record_request_outcome = AsyncMock(side_effect=recording_outcome)
+
+        r = client.post(
+            "/v1beta/models/gemini-2.0-flash:generateContent?key=test-key",
+            json={"contents": [{"parts": [{"text": big}]}]},
+        )
+
+    assert r.status_code == 200, r.text
+    outcome = captured["outcome"]
+    assert outcome.output_tokens == 0
+    assert isinstance(outcome.output_tokens, int)
+    # max(0, promptTokenCount - cache_read) with a null candidate count must not raise.
+    assert outcome.uncached_input_tokens == 20
+
+
+def test_gemini_zero_usage_prompt_count_is_preserved():
+    """A real zero promptTokenCount must stay zero, not fall back to estimates."""
+    config = ProxyConfig(
+        optimize=True,
+        cache_enabled=False,
+        rate_limit_enabled=False,
+        cost_tracking_enabled=False,
+    )
+
+    def passthrough_apply(**kwargs):
+        return SimpleNamespace(
+            messages=kwargs["messages"],
+            transforms_applied=[],
+            timing={},
+            tokens_before=10,
+            tokens_after=10,
+            waste_signals=None,
+        )
+
+    resp = MagicMock()
+    resp.status_code = 200
+    resp.headers = {"content-type": "application/json"}
+    resp.content = (
+        b'{"candidates":[{"content":{"parts":[{"text":"ok"}]}}],'
+        b'"usageMetadata":{"promptTokenCount":0,"candidatesTokenCount":0}}'
+    )
+    resp.json.return_value = {
+        "candidates": [{"content": {"parts": [{"text": "ok"}]}}],
+        "usageMetadata": {"promptTokenCount": 0, "candidatesTokenCount": 0},
+    }
+
+    captured: dict[str, object] = {}
+
+    async def recording_outcome(outcome):  # noqa: ANN001
+        captured["outcome"] = outcome
+
+    big = "word " * 4000
+    app = create_app(config)
+    with TestClient(app) as client:
+        proxy = client.app.state.proxy
+        proxy.openai_pipeline.apply = MagicMock(side_effect=passthrough_apply)
+        proxy._retry_request = AsyncMock(return_value=resp)
+        proxy._record_request_outcome = AsyncMock(side_effect=recording_outcome)
+
+        r = client.post(
+            "/v1beta/models/gemini-2.0-flash:generateContent?key=test-key",
+            json={"contents": [{"parts": [{"text": big}]}]},
+        )
+
+    assert r.status_code == 200, r.text
+    outcome = captured["outcome"]
+    assert outcome.optimized_tokens == 0
+    assert outcome.uncached_input_tokens == 0