fix(gemini): correct tool-use and thinking token accounting

[guslegend0510]

Fixes #3033

Gemini usage parsing currently omits server-side tool-use prompt tokens from input and subtracts thinking tokens from candidate output. This undercounts usage and can produce negative output counts.

This change:

• Includes tool-use prompt tokens in input usage.
• Adds thinking tokens to candidate output.
• Falls back to total minus input when candidate counts are absent, clamping negative results to zero.
• Preserves reported thinking tokens when both candidate and total counts are absent.
• Keeps cached tokens as a subset of input and leaves the public ChatUsage API unchanged.

For the issue's example, input/output counts change from 500/110 to 800/130, matching the reported total of 930.

Validation:

• Regression tests reproduce 12 failures before the fix and pass afterward.
• Added 12 parameterized cases covering tool use, thinking, missing fields, explicit zero counts, and fallback behavior.
• Spotless and git diff --check passed.
• Full Core, Harness, and Gemini test suites: 3,457 tests, 14 skipped, zero failures or errors.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[oss-maintainer]

Summary

Direction is right: server-side tool_use_prompt_token_count belongs in input usage, and subtracting thinking from candidatesTokenCount was wrong once the SDK reports thinking separately. The parameterized matrix over absent/present fields is a good way to pin the fallback down. Filed as COMMENT because the change alters a cross-provider convention and should land with that convention written down.

Findings

• [Warning] GeminiResponseParser.java:108 — outputTokens now includes thinking for Gemini only; confirm/align the DashScope convention and document it, or per-provider cost and budget aggregation silently diverges.
• [Info] GeminiResponseParser.java:100 — keep cachedTokens a subset of the widened inputTokens.
• [Info] GeminiResponseParser.java:112 — the clamp hides provider inconsistency; a debug log would make it diagnosable.

Suggestions

Add at least one test driven by a captured real usageMetadata payload (the current fixtures moved with the formula, which makes the expected totals circular), and state in the description whether the same accounting correction applies to the other model extensions.

---

Automated review by github-manager-bot

[oss-maintainer]

This flips what ChatUsage.outputTokens means for Gemini: thinking used to be excluded (candidates - thinking, per the removed DashScope-parity comment) and is now included (candidates + thinking). Consumers that sum usage across providers — cost accounting, budget/context-compaction decisions in the harness, tracing exporters — will now read Gemini's output differently from DashScope's. Please confirm the cross-provider convention explicitly: does the DashScope path include reasoning/thinking tokens in outputTokens or not? If it excludes them, either this PR or the other provider needs a matching follow-up, and the convention should be documented on ChatUsage.getOutputTokens() — otherwise Gemini becomes internally consistent while the framework becomes less so. Related: the pre-existing fixture was edited alongside the formula (total 160 -> 170), so it now encodes the new assumption rather than an observed payload.

[oss-maintainer]

toolUsePromptTokenCount is folded into inputTokens while cachedTokens is passed through unchanged. ChatUsage's contract states cachedTokens is a subset of inputTokens. That still holds when both come from the prompt, but the new no prompt with tool input case in the test matrix (prompt absent, toolUsePrompt 300) is exactly the shape where a cached-token overlap would go unnoticed. Worth one sentence confirming the cached/prompt relationship when only tool-use prompt tokens are reported.

[oss-maintainer]

Math.max(0, total - inputTokens) clamps a provider-side inconsistency to 0, which is a safe default but erases the signal: outputTokens == 0 from a clamp is indistinguishable from a genuinely empty completion. Since this is the fallback branch, a log.debug when the clamp actually engages would let operators tell "provider reported no output" apart from "provider reported inconsistent usage".

[oss-maintainer]

Summary

Fixes Gemini token usage accounting (#3033): promptTokenCount + toolUsePromptTokenCount now count as input, candidatesTokenCount + thoughtsTokenCount as output, with a totalTokenCount - input fallback clamped at zero and logged at debug. The reasoning is correct and matches how Gemini actually reports usage — thinking tokens are part of output, and server-side tool prompts are part of input. The wire-format regression test added here is the right idea: parsing an actual usageMetadata payload means the assertion cannot drift along with the fixture arithmetic, which is exactly how the original bug survived the existing tests.

Left COMMENT rather than APPROVE mainly because of the duplicate-PR question below, plus one code-shape nit.

Findings

• [Warning] GeminiResponseParser.java:110 — outputTokens = thinkingTokens; is unconditionally overwritten at line 111 whenever totalTokenCount is present, so it only survives in the no-total-count case. Correct, but it reads like a leftover and would break subtly if someone added an else. #3034 expresses the same three cases as an Optional chain with no dead assignment, and computes identical numbers.
• [Info] GeminiResponseParser.java:114 — the clamp plus debug log is the right treatment. Flagging separately: adding toolUsePromptTokenCount to inputTokens changes what ChatUsage.getInputTokens() reports for every tool-using Gemini response, and getTotalTokens() is derived as input + output rather than read from the provider. Anything comparing total tokens against a budget/context-window limit will see different values than before, and usage persisted in existing sessions will not be comparable to new records. Worth stating in the PR description; the cachedTokens-subset-of-inputTokens invariant still holds (updated test: cached=300 against input=800).

Cross-PR Note

#3034 (fix(gemini): correct token usage accounting, @liugy789) fixes the same issue #3033 in the same two files, with the same semantics — input = prompt + tool-use prompt, output = candidates + thoughts, fallback = total - input clamped at zero. The two differ only in implementation shape and in this PR's extra doc/ChatUsage clarifications and toolUsePromptTokenCount fixture update; #3034 adds somewhat broader parser-level fallback coverage (thinking-only, missing-candidate).

Both PRs will conflict on GeminiResponseParser.java and its test, and merging both is not possible. Could one of you coordinate, or a maintainer pick? Happy to re-review whichever survives.

Tests

Good coverage on the parser itself. DashScopeUsage and ChatUsage changes here are javadoc-only, so no additional tests needed for those.

---

Automated review by github-manager-bot

[oss-maintainer]

[Warning] This initializer is dead on arrival: outputTokens is set to thinkingTokens here and then unconditionally overwritten on line 111 whenever totalTokenCount is present, so the assignment only survives in the no-total-count case.

That is presumably deliberate, but written as an assignment-plus-override it reads like a leftover, and a later edit that adds an else branch would silently change the fallback. Expressing it as a single conditional (or an Optional chain) makes the three cases — candidate count present, total count present, neither present — mutually exclusive by construction rather than by ordering.

Note #3034 solves the same problem with .map(...).orElseGet(...) and has no dead assignment, while computing identical numbers to this version. Since the two PRs are otherwise interchangeable on behaviour, aligning on one of the two shapes would keep them from drifting.

[oss-maintainer]

[Info] Clamping to zero and logging at debug is the right call — a provider reporting totalTokenCount < inputTokens is a data-quality problem, not something the caller can act on.

Worth confirming that adding toolUsePromptTokenCount into inputTokens here is what you want downstream consumers to see: ChatUsage.getInputTokens() now includes server-side tool prompt tokens for every Gemini response that uses them, and getTotalTokens() is derived as input + output rather than taken from the provider's totalTokenCount. Budget/cost accounting and any context-window or compaction check that compares getTotalTokens() against a configured limit will shift for tool-using calls, and previously stored sessions will not be comparable to new ones. A line in the PR description naming that would save reviewers the derivation; the cachedTokens-is-a-subset-of-inputTokens invariant still holds (the updated test has cached=300 against input=800).

[oss-maintainer]

Summary

Re-reviewed after 7f0cff1e ("refactor(gemini): make usage fallbacks explicit"). The new commit directly resolves the item from my previous review: outputTokens is no longer assigned and then unconditionally overwritten — the three provider shapes are now an explicit if / else if / else chain, which is the shape I asked for. Everything else in the PR still holds up: server-side toolUsePromptTokenCount folded into input, candidatesTokenCount + thoughtsTokenCount as output, the total - input fallback clamped and logged, and the convention written down in ChatUsage / DashScopeUsage javadoc so the cross-provider semantics are no longer implicit. The added GenerateContentResponse.fromJson wire-format test plus the 12-row parameterized matrix is the right way to stop the fixtures from drifting with the arithmetic.

Remaining two comments are non-blocking info nits (state the "components win over provider total" rule at line 108; cover the thinkingTokens-only branch with a matrix row).

Still filing as COMMENT rather than APPROVE, for one process reason only: #3034 targets the same accounting bug in the same file and is still open, with an equivalent fix. Both look correct now, so this is a maintainer pick, not a code objection — whoever merges should close the other to avoid the second-order conflict. Once that is settled this is good to go from my side.

---

Automated review by github-manager-bot

[oss-maintainer]

Re-review of 7f0cff1e: the previous dead-assignment concern is resolved — the three cases are now an explicit if / else if / else chain, so outputTokens is assigned exactly once per path. Thanks for the quick turnaround.

One residual: in this primary branch totalTokenCount is ignored entirely, and ChatUsage.getTotalTokens() is derived as input + output. So when the provider reports a totalTokenCount that disagrees with the components, the disagreement is invisible to callers (the "total smaller than input" row in the new matrix is exactly such a case, where the derived total is 150 while the provider said 120). That is a defensible choice, but it is now an implicit rule rather than a stated one — a one-line comment here (// prefer component counts; provider total is only used as a fallback) would keep the next reader from "fixing" it by adding another clamp.

[oss-maintainer]

This branch is only reachable when candidatesTokenCount, totalTokenCount are both absent and thoughtsTokenCount is present, so outputTokens == thinkingTokens and any non-reasoning output is recorded as 0. Correct given the fields available, but it is the one path with no test row (missing total and candidates leaves thoughts empty, so it lands on the else with 0). Adding a row with candidates/total absent and thoughts=10 would pin the intent down.

[jujn]

fixed in: #3034