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 all 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
16 changes: 10 additions & 6 deletions agentscope-core/src/main/java/io/agentscope/core/model/ChatUsage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@
* Represents token usage information for chat completion responses.
*
* <p>This immutable data class tracks the number of tokens used during a chat completion,
* including input tokens (prompt), output tokens (generated response), cached input tokens, and
* execution time.
* including input tokens (prompt), output tokens (generated response, including reasoning/thinking
* tokens when the provider reports them separately), cached input tokens, and execution time.
*/
public class ChatUsage {

Expand All @@ -39,7 +39,8 @@ public class ChatUsage {
* #ChatUsage(int, int, int, double)} with {@code cachedTokens} defaulting to {@code 0}.
*
* @param inputTokens the number of tokens used for the input/prompt
* @param outputTokens the number of tokens used for the output/generated response
* @param outputTokens the number of tokens used for model-generated output, including
* reasoning/thinking tokens when reported separately by the provider
* @param time the execution time in seconds
*/
public ChatUsage(int inputTokens, int outputTokens, double time) {
Expand All @@ -50,7 +51,8 @@ public ChatUsage(int inputTokens, int outputTokens, double time) {
* Creates a new ChatUsage instance.
*
* @param inputTokens the number of tokens used for the input/prompt
* @param outputTokens the number of tokens used for the output/generated response
* @param outputTokens the number of tokens used for model-generated output, including
* reasoning/thinking tokens when reported separately by the provider
* @param cachedTokens the number of input tokens served from the prompt cache (a subset of
* {@code inputTokens}); {@code 0} when the provider does not report cache information
* @param time the execution time in seconds
Expand Down Expand Up @@ -79,7 +81,8 @@ public int getInputTokens() {
/**
* Gets the number of output tokens used.
*
* @return the number of tokens used for the output/generated response
* @return the number of model-generated output tokens, including reasoning/thinking tokens when
* reported separately by the provider
*/
public int getOutputTokens() {
return outputTokens;
Expand Down Expand Up @@ -148,7 +151,8 @@ public Builder inputTokens(int inputTokens) {
/**
* Sets the number of output tokens.
*
* @param outputTokens the number of tokens used for the output/generated response
* @param outputTokens the number of tokens used for model-generated output, including
* reasoning/thinking tokens when reported separately by the provider
* @return this builder instance
*/
public Builder outputTokens(int outputTokens) {
Expand Down
7 changes: 6 additions & 1 deletion ...-dashscope/src/main/java/io/agentscope/extensions/model/dashscope/dto/DashScopeUsage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,12 @@ public class DashScopeUsage {
@JsonProperty("input_tokens")
private Integer inputTokens;

/** Number of tokens in the output. */
/**
* Number of tokens in the output.
*
* <p>DashScope includes reasoning tokens in this total; any separately reported reasoning-token
* count is a subset rather than an additional amount.
*/
@JsonProperty("output_tokens")
private Integer outputTokens;

Expand Down
28 changes: 23 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,32 @@ 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 if (metadata.totalTokenCount().isPresent()) {
int totalTokens = metadata.totalTokenCount().get();
int reportedOutputTokens = totalTokens - inputTokens;
if (reportedOutputTokens < 0) {
log.debug(
"Gemini usage totalTokenCount ({}) is smaller than input token"
+ " count ({}); clamping outputTokens to zero",
totalTokens,
inputTokens);
}
outputTokens = Math.max(0, reportedOutputTokens);
} else {
outputTokens = 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.

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

}

usage =
ChatUsage.builder()
Expand Down
99 changes: 94 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 @@ -260,9 +264,10 @@ void testParseUsageMetadataReadsCachedContentTokenCount() {
// cachedContentTokenCount 是 promptTokenCount 的子集(Gemini SDK 文档:
// promptTokenCount 包含 cachedContentTokenCount),故 prompt 必须 > cached
.promptTokenCount(500)
.toolUsePromptTokenCount(300)
.candidatesTokenCount(60)
.thoughtsTokenCount(10)
.totalTokenCount(560)
.totalTokenCount(870)
.cachedContentTokenCount(300)
.build();

Expand All @@ -277,6 +282,89 @@ void testParseUsageMetadataReadsCachedContentTokenCount() {

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

@Test
void testUsageTokenAccountingFromIssuePayloadJson() {
// Preserve the issue #3033 reproduction in provider wire format so SDK field mapping cannot
// drift together with the parser's arithmetic fixtures.
GenerateContentResponse response =
GenerateContentResponse.fromJson(
"""
{
"usageMetadata": {
"promptTokenCount": 500,
"candidatesTokenCount": 120,
"toolUsePromptTokenCount": 300,
"thoughtsTokenCount": 10,
"totalTokenCount": 930
}
}
""");

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

assertNotNull(usage);
assertEquals(800, usage.getInputTokens());
assertEquals(130, usage.getOutputTokens());
assertEquals(930, usage.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 +379,7 @@ void testParseEmptyResponse() {
// Verify
assertNotNull(chatResponse);
assertEquals("response-empty", chatResponse.getId());
assertNull(chatResponse.getUsage());
assertEquals(0, chatResponse.getContent().size());
}

Expand Down
Loading