Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
guslegend0510 wants to merge 12 commits into from
Closed

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

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
ac9790c
fix(gemini): correct tool-use and thinking token accounting
guslegend Sep 8, 2026
fa76732
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 8, 2026
617184b
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 9, 2026
de1f986
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 9, 2026
7d3e86a
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 9, 2026
8601299
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 10, 2026
54a0be1
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 10, 2026
1786d1d
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 10, 2026
47da7d5
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 10, 2026
6ae69bb
Merge remote-tracking branch 'origin/main' into HEAD
guslegend Sep 11, 2026
f13f8c1
fix(model): clarify reasoning token accounting
guslegend Sep 11, 2026
7f0cff1
refactor(gemini): make usage fallbacks explicit
guslegend Sep 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 15 additions & 5 deletions ...i/src/main/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParser.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,14 +94,24 @@ public ChatResponse parseResponse(GenerateContentResponse response, Instant star
if (response.usageMetadata().isPresent()) {
GenerateContentResponseUsageMetadata metadata = response.usageMetadata().get();

int inputTokens = metadata.promptTokenCount().orElse(0);
// Server-side tool results are fed back to the model as additional input.
int inputTokens =
metadata.promptTokenCount().orElse(0)
+ metadata.toolUsePromptTokenCount().orElse(0);

@oss-maintainer oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

int cachedTokens = metadata.cachedContentTokenCount().orElse(0);
int totalOutputTokens = metadata.candidatesTokenCount().orElse(0);
int thinkingTokens = metadata.thoughtsTokenCount().orElse(0);

// Output tokens exclude thinking tokens (following DashScope behavior)
// In Gemini, candidatesTokenCount includes thinking, so we subtract it
int outputTokens = totalOutputTokens - thinkingTokens;
// Gemini reports candidate and thinking tokens separately; both are output.
// The total already includes thinking, so do not add it again in the fallback.
int outputTokens;
if (metadata.candidatesTokenCount().isPresent()) {
outputTokens = metadata.candidatesTokenCount().get() + thinkingTokens;

@oss-maintainer oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

} else {
outputTokens =
metadata.totalTokenCount()
.map(total -> Math.max(0, total - inputTokens))

@oss-maintainer oss-maintainer Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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".

.orElse(thinkingTokens);
}

usage =
ChatUsage.builder()
Expand Down
72 changes: 67 additions & 5 deletions ...c/test/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParserTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

import com.google.genai.types.Candidate;
Expand All @@ -38,6 +39,8 @@
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.CsvSource;

/**
* Unit tests for GeminiResponseParser.
Expand Down Expand Up @@ -216,9 +219,9 @@ void testParseUsageMetadata() {
GenerateContentResponseUsageMetadata usageMetadata =
GenerateContentResponseUsageMetadata.builder()
.promptTokenCount(100)
.candidatesTokenCount(60) // Includes thinking
.candidatesTokenCount(60) // Excludes thinking
.thoughtsTokenCount(10) // Thinking tokens
.totalTokenCount(160)
.totalTokenCount(170)
.build();

GenerateContentResponse response =
Expand All @@ -238,8 +241,9 @@ void testParseUsageMetadata() {
// Input tokens = promptTokenCount
assertEquals(100, usage.getInputTokens());

// Output tokens = candidatesTokenCount - thoughtsTokenCount
assertEquals(50, usage.getOutputTokens());
// Output tokens include both candidate and thinking tokens.
assertEquals(70, usage.getOutputTokens());
assertEquals(170, usage.getTotalTokens());

// Time should be > 0
assertTrue(usage.getTime() >= 0);
Expand All @@ -262,7 +266,7 @@ void testParseUsageMetadataReadsCachedContentTokenCount() {
.promptTokenCount(500)
.candidatesTokenCount(60)
.thoughtsTokenCount(10)
.totalTokenCount(560)
.totalTokenCount(570)
.cachedContentTokenCount(300)
.build();

Expand All @@ -277,6 +281,63 @@ void testParseUsageMetadataReadsCachedContentTokenCount() {

assertNotNull(chatResponse.getUsage());
assertEquals(300, chatResponse.getUsage().getCachedTokens());
assertEquals(500, chatResponse.getUsage().getInputTokens());
assertEquals(70, chatResponse.getUsage().getOutputTokens());
assertEquals(570, chatResponse.getUsage().getTotalTokens());
}

@ParameterizedTest(name = "{0}")
@CsvSource({
"server-side tools and thinking, 500, 300, 120, 10, 930, 800, 130",
"server-side tools without thinking, 500, 300, 120, , 920, 800, 120",
"thinking exceeds candidates, 100, , 10, 60, 170, 100, 70",
"missing candidates with tools, 500, 300, , 10, 930, 800, 130",
"missing candidates without tools, 100, , , 10, 170, 100, 70",
"explicit zero candidates, 100, , 0, 10, 170, 100, 10",
"missing total with candidates, 100, , 60, 10, , 100, 70",
"missing total and candidates, 100, , , 10, , 100, 10",
"total smaller than input, 100, 50, , , 120, 150, 0",
"prompt only, 100, , , , , 100, 0",
"no prompt with tool input, , 300, 120, 10, 430, 300, 130",
"empty metadata, , , , , , 0, 0"
})
void testUsageTokenAccounting(
String scenario,
Integer prompt,
Integer toolPrompt,
Integer candidates,
Integer thoughts,
Integer total,
int expectedInput,
int expectedOutput) {
GenerateContentResponseUsageMetadata.Builder metadata =
GenerateContentResponseUsageMetadata.builder();
if (prompt != null) {
metadata.promptTokenCount(prompt);
}
if (toolPrompt != null) {
metadata.toolUsePromptTokenCount(toolPrompt);
}
if (candidates != null) {
metadata.candidatesTokenCount(candidates);
}
if (thoughts != null) {
metadata.thoughtsTokenCount(thoughts);
}
if (total != null) {
metadata.totalTokenCount(total);
}
// Streaming responses may carry usage without candidate content.
GenerateContentResponse response =
GenerateContentResponse.builder().usageMetadata(metadata.build()).build();

ChatUsage usage = parser.parseResponse(response, startTime).getUsage();

assertNotNull(usage);
assertEquals(expectedInput, usage.getInputTokens(), scenario);
assertEquals(expectedOutput, usage.getOutputTokens(), scenario);
assertEquals(expectedInput + expectedOutput, usage.getTotalTokens(), scenario);
assertEquals(0, usage.getCachedTokens());
}

@Test
Expand All @@ -291,6 +352,7 @@ void testParseEmptyResponse() {
// Verify
assertNotNull(chatResponse);
assertEquals("response-empty", chatResponse.getId());
assertNull(chatResponse.getUsage());
assertEquals(0, chatResponse.getContent().size());
}

Expand Down
Loading