Skip to content

fix: add reasoning tokens into token metrics in stream mode for gemini models#2259

Merged
aabchoo merged 6 commits into
envoyproxy:mainfrom
hustxiayang:fix-gemini-streaming-reasoning-tokens
Jul 2, 2026
Merged

fix: add reasoning tokens into token metrics in stream mode for gemini models#2259
aabchoo merged 6 commits into
envoyproxy:mainfrom
hustxiayang:fix-gemini-streaming-reasoning-tokens

Conversation

@hustxiayang

@hustxiayang hustxiayang commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

Description
The previous implementation on tokens tracking is like

if chunk.UsageMetadata.CandidatesTokenCount >= 0 {
				tokenUsage.SetOutputTokens(uint32(chunk.UsageMetadata.CandidatesTokenCount)) //nolint:gosec
			}

The output tokens should include both the reasoning tokens and the candidates tokens. Thus fix this by reusing the the variable usage, just like non-streaming mode. Also updated the tests to include reasoning tokens

Signed-off-by: yxia216 <yxia216@bloomberg.net>
@hustxiayang
hustxiayang requested a review from a team as a code owner June 18, 2026 14:17
@dosubot dosubot Bot added the size:S This PR changes 10-29 lines, ignoring generated files. label Jun 18, 2026
@hustxiayang

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request refactors how token usage is extracted and recorded during streaming responses in the OpenAI to GCP Vertex AI translator, switching from chunk.UsageMetadata to direct fields on the usage object. However, accessing nested fields like usage.PromptTokensDetails.CachedTokens and usage.CompletionTokensDetails.ReasoningTokens without nil checks introduces a risk of nil pointer dereference panics if those details are not populated in the stream. It is recommended to add defensive nil checks for these fields.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment on lines +255 to +256
tokenUsage.SetCachedInputTokens(uint32(usage.PromptTokensDetails.CachedTokens)) //nolint:gosec
tokenUsage.SetReasoningTokens(uint32(usage.CompletionTokensDetails.ReasoningTokens)) //nolint:gosec

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

high

Directly accessing usage.PromptTokensDetails and usage.CompletionTokensDetails without nil checks can lead to a nil pointer dereference panic if these fields are not populated in the streaming chunk. Adding defensive nil checks (similar to the non-streaming implementation) prevents potential runtime panics.

			if usage.PromptTokensDetails != nil {
				tokenUsage.SetCachedInputTokens(uint32(usage.PromptTokensDetails.CachedTokens)) //nolint:gosec
			}
			if usage.CompletionTokensDetails != nil {
				tokenUsage.SetReasoningTokens(uint32(usage.CompletionTokensDetails.ReasoningTokens)) //nolint:gosec
			}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

in function geminiUsageToOpenAIUsage, these values are always set

@hustxiayang hustxiayang changed the title fix: reasoning tokens into token metrics in stream mode for gemini models fix: add reasoning tokens into token metrics in stream mode for gemini models Jun 18, 2026
@codecov-commenter

codecov-commenter commented Jun 18, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.88%. Comparing base (08ba26e) to head (36e01f1).

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2259      +/-   ##
==========================================
- Coverage   84.90%   84.88%   -0.02%     
==========================================
  Files         144      144              
  Lines       21259    21254       -5     
==========================================
- Hits        18049    18042       -7     
- Misses       2131     2132       +1     
- Partials     1079     1080       +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@hustxiayang

Copy link
Copy Markdown
Contributor Author

@nacx @aabchoo @yuzisun could you please help to review this fix? Thanks!

@aabchoo aabchoo left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks for the fix!

@hustxiayang

Copy link
Copy Markdown
Contributor Author

/retest

@aabchoo
aabchoo merged commit 64ba6c5 into envoyproxy:main Jul 2, 2026
55 of 57 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:S This PR changes 10-29 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants