fix(proxy/cost): count Gemini thinking tokens in output usage (#2639)

## Description

The Gemini handlers take the response's output-token count straight from
`candidatesTokenCount`:

```python
output_tokens = _usage_int(usage.get("candidatesTokenCount"))
```

For Gemini 2.5 thinking models that undercounts. Gemini reports
`candidatesTokenCount` **sometimes inclusive** of the reasoning tokens
(`thoughtsTokenCount`) and **sometimes exclusive** of them. When it is
exclusive, the thinking tokens are a separate bucket that is still
billed at the output rate, so dropping them makes `output_tokens` (and
therefore the output cost that flows through `record_tokens` ->
`estimate_cost`) too low. The gap grows with reasoning effort.

litellm handles exactly this: it adds `thoughtsTokenCount` to completion
tokens unless `promptTokenCount + candidatesTokenCount ==
totalTokenCount` (its `is_candidate_token_count_inclusive` check). The
Headroom handlers had no equivalent.

## Fix

Add `gemini_output_tokens(usage_meta)` in
`headroom/proxy/token_counting.py`:

- No `thoughtsTokenCount` (the common non-2.5 case): return
`candidatesTokenCount` unchanged.
- `promptTokenCount + candidatesTokenCount == totalTokenCount`:
candidates already include thoughts, return `candidatesTokenCount`.
- Otherwise: return `candidatesTokenCount + thoughtsTokenCount`.

This mirrors litellm's rule and is robust to missing or null fields.
Wire it into the native Gemini handler (both the generate and count
paths), the streaming usage extractors, and the OpenAI-compatible
passthrough usage normalizer, so every Gemini usage path counts output
the same way.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature
- [ ] Breaking change
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring

## Changes Made

- `headroom/proxy/token_counting.py`: add `gemini_output_tokens()`.
- `headroom/proxy/handlers/gemini.py`: use it for `output_tokens` on
both response paths.
- `headroom/proxy/handlers/streaming.py`: use it in the two Gemini
streaming usage extractors.
- `headroom/proxy/handlers/openai.py`: use it in
`_passthrough_usage_from_json` (Gemini-shaped usage).
- `tests/test_proxy_handler_helpers.py`: unit test for
`gemini_output_tokens` (inclusive / exclusive / no-thinking / empty) and
a `_passthrough_usage_from_json` test that thinking tokens land in
`output_tokens`.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check` / `ruff format --check`)
- [x] Type checking passes (`mypy`)
- [x] New tests added for the fix
- [ ] Manual testing performed

### Test Output

```text
$ python -m pytest tests/test_proxy_handler_helpers.py -k "gemini_output_tokens or thinking or vertex_usage_metadata" -q
3 passed

$ python -m pytest tests/test_proxy_gemini_native_integration.py tests/test_proxy/test_gemini_savings_profile.py tests/test_proxy_handler_helpers.py -q
38 passed, 18 skipped

# with the wiring reverted, the passthrough test fails (output_tokens is 200, not 700):
$ git stash push headroom/proxy/handlers/openai.py && \
    python -m pytest tests/test_proxy_handler_helpers.py -k passthrough_usage_counts_gemini_thinking -q
1 failed

$ uvx ruff@0.15.17 check headroom/proxy/token_counting.py headroom/proxy/handlers/gemini.py headroom/proxy/handlers/streaming.py headroom/proxy/handlers/openai.py tests/test_proxy_handler_helpers.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/token_counting.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: called `gemini_output_tokens` on an exclusive
usage (`prompt=1000, candidates=200, thoughts=500, total=1700`), an
inclusive usage (`candidates=700, total=1700`), a no-thinking usage, and
`{}`; drove `_passthrough_usage_from_json` with a thinking usage; then
reverted the handler wiring and re-ran the passthrough test.
- Observed result: exclusive returns 700 (200 visible plus 500
thinking), inclusive returns 700, no-thinking returns the candidates
count, empty returns 0; `_passthrough_usage_from_json` reports
`output_tokens=700`. With the wiring reverted it reports 200 (the
undercount). Verified against litellm's documented rule.
- Not tested: a live Gemini 2.5 request end to end (the accounting is
verified at the usage-extraction boundary against litellm's reference
logic).

## 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
- [ ] I have updated the CHANGELOG.md if applicable

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

headroom/proxy/handlers/gemini.py

@@ -22,6 +22,7 @@ from headroom.proxy.auth_mode import classify_client
 from headroom.proxy.compression_decision import CompressionDecision
 from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS, extract_tags
 from headroom.proxy.outcome import RequestOutcome
+from headroom.proxy.token_counting import gemini_output_tokens
 
 logger = logging.getLogger("headroom.proxy")
 
@@ -462,7 +463,9 @@ class GeminiHandlerMixin:
                     # 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"))
+                    output_tokens = gemini_output_tokens(
+                        usage
+                    )  # includes thinking tokens (2.5-family)
                     cache_read_tokens = _usage_int(usage.get("cachedContentTokenCount"))
                 except (json.JSONDecodeError, ValueError, KeyError, TypeError, AttributeError):
                     pass
@@ -714,7 +717,9 @@ class GeminiHandlerMixin:
                         if usage.get("promptTokenCount") is None
                         else usage["promptTokenCount"]
                     )
-                    output_tokens = _usage_int(usage.get("candidatesTokenCount"))
+                    output_tokens = gemini_output_tokens(
+                        usage
+                    )  # includes thinking tokens (2.5-family)
                     # Gemini returns cachedContentTokenCount for context-cached tokens
                     # These are charged at 10-25% of the input price depending on model
                     cache_read_tokens = _usage_int(usage.get("cachedContentTokenCount"))

headroom/proxy/handlers/openai.py

@@ -75,6 +75,7 @@ from headroom.proxy.passthrough import (
     custom_base_passthrough_telemetry as _custom_base_passthrough_telemetry,
 )
 from headroom.proxy.project_context import classify_project, set_current_project
+from headroom.proxy.token_counting import gemini_output_tokens
 
 logger = logging.getLogger("headroom.proxy")
 
@@ -331,7 +332,7 @@ def _passthrough_usage_from_json(payload: Any) -> dict[str, int]:
     if isinstance(usage_meta, dict):
         return {
             "input_tokens": _usage_int(usage_meta.get("promptTokenCount")),
-            "output_tokens": _usage_int(usage_meta.get("candidatesTokenCount")),
+            "output_tokens": gemini_output_tokens(usage_meta),
             "cache_read_input_tokens": _usage_int(usage_meta.get("cachedContentTokenCount")),
         }
 

headroom/proxy/handlers/streaming.py

@@ -18,6 +18,7 @@ from headroom.proxy.helpers import (
     jitter_delay_ms,
     retry_after_ms,
 )
+from headroom.proxy.token_counting import gemini_output_tokens
 
 if TYPE_CHECKING:
     from fastapi.responses import Response, StreamingResponse
@@ -224,7 +225,7 @@ class StreamingMixin:
                     usage_meta = data.get("usageMetadata")
                     if usage_meta:
                         usage["input_tokens"] = usage_meta.get("promptTokenCount", 0)
-                        usage["output_tokens"] = usage_meta.get("candidatesTokenCount", 0)
+                        usage["output_tokens"] = gemini_output_tokens(usage_meta)
                         # Gemini also has cachedContentTokenCount for context caching
                         usage["cache_read_input_tokens"] = usage_meta.get(
                             "cachedContentTokenCount", 0
@@ -342,7 +343,7 @@ class StreamingMixin:
                 usage_meta = data.get("usageMetadata")
                 if usage_meta:
                     usage_found["input_tokens"] = usage_meta.get("promptTokenCount", 0)
-                    usage_found["output_tokens"] = usage_meta.get("candidatesTokenCount", 0)
+                    usage_found["output_tokens"] = gemini_output_tokens(usage_meta)
                     usage_found["cache_read_input_tokens"] = usage_meta.get(
                         "cachedContentTokenCount", 0
                     )

headroom/proxy/token_counting.py

@@ -77,3 +77,33 @@ async def count_texts_offloaded(owner: Any, model: Any, texts: Any) -> tuple[Any
     return await _count_offloaded(
         owner, model, lambda counter: sum(counter.count_text(text) for text in text_list)
     )
+
+
+def gemini_output_tokens(usage_meta: dict[str, Any]) -> int:
+    """Output-token count for a Gemini ``usageMetadata``, including thinking tokens.
+
+    Gemini reports ``candidatesTokenCount`` sometimes inclusive of the
+    ``thoughtsTokenCount`` (2.5-family reasoning) and sometimes exclusive of it.
+    When ``promptTokenCount + candidatesTokenCount != totalTokenCount`` the
+    thinking tokens are a separate bucket and must be added, or the output cost
+    (billed at the output rate) is undercounted. Mirrors litellm's
+    ``is_candidate_token_count_inclusive`` rule. Robust to missing/null fields.
+    """
+
+    def _int(value: Any) -> int:
+        try:
+            return max(int(value), 0)
+        except (TypeError, ValueError):
+            return 0
+
+    candidates = _int(usage_meta.get("candidatesTokenCount"))
+    thoughts = _int(usage_meta.get("thoughtsTokenCount"))
+    if thoughts <= 0:
+        return candidates
+    prompt = _int(usage_meta.get("promptTokenCount"))
+    total = _int(usage_meta.get("totalTokenCount"))
+    # Inclusive iff prompt + candidates already equals total; otherwise the
+    # thinking tokens are a separate bucket that belongs in the output count.
+    if prompt + candidates == total:
+        return candidates
+    return candidates + thoughts

tests/test_proxy_handler_helpers.py

@@ -381,6 +381,51 @@ def test_passthrough_usage_normalizes_vertex_usage_metadata() -> None:
     }
 
 
+def test_gemini_output_tokens_includes_thinking_when_exclusive() -> None:
+    """Gemini 2.5 thinking: when prompt + candidates != total, thoughtsTokenCount
+    is a separate output bucket and must be added, or output cost undercounts."""
+    from headroom.proxy.token_counting import gemini_output_tokens
+
+    exclusive = {
+        "promptTokenCount": 1000,
+        "candidatesTokenCount": 200,
+        "thoughtsTokenCount": 500,
+        "totalTokenCount": 1700,
+    }
+    assert gemini_output_tokens(exclusive) == 700  # 200 visible + 500 thinking
+
+    # Inclusive: candidatesTokenCount already covers thoughts (prompt+cand==total).
+    inclusive = {
+        "promptTokenCount": 1000,
+        "candidatesTokenCount": 700,
+        "thoughtsTokenCount": 500,
+        "totalTokenCount": 1700,
+    }
+    assert gemini_output_tokens(inclusive) == 700
+
+    # No thinking tokens: just the candidates count (common non-2.5 case).
+    assert gemini_output_tokens({"candidatesTokenCount": 42, "totalTokenCount": 100}) == 42
+    # Robust to empty / missing fields.
+    assert gemini_output_tokens({}) == 0
+
+
+def test_passthrough_usage_counts_gemini_thinking_tokens() -> None:
+    """_passthrough_usage_from_json must include thinking tokens in output_tokens."""
+    usage = _passthrough_usage_from_json(
+        {
+            "usageMetadata": {
+                "promptTokenCount": 1000,
+                "candidatesTokenCount": 200,
+                "thoughtsTokenCount": 500,
+                "totalTokenCount": 1700,
+                "cachedContentTokenCount": 100,
+            }
+        }
+    )
+    assert usage["output_tokens"] == 700
+    assert usage["input_tokens"] == 1000
+
+
 def test_vertex_passthrough_records_usage_metadata_for_dashboard() -> None:
     handler = object.__new__(HeadroomProxy)
     handler.http_client = _VertexUsageClient()