Gemini: preserve thought signature across streaming tool calls for Gemini 3

[cpunion]

Problem

Gemini 3 validates thought_signature for tool calling. In streaming mode the signature may arrive only on the first functionCall chunk or in an empty final text part, and signed parts must not be merged. Missing signatures or reordered parts can lead to 400s or degraded tool-calling quality.

Summary

• Preserve thought_signature while aggregating streaming text, including empty final text parts.
• Avoid merging signed text parts and keep unsigned text aggregation stable.
• Aggregate streaming function-call arguments by call ID; handle missing IDs with deterministic ordering.
• Treat empty functionCall chunks as valid end markers when a streamed call is pending.
• Flush all pending streaming function calls on Close if a stream ends unexpectedly.
• Add/annotate tests with doc links, quotes, and scenario summaries.
• Stabilize session/database tests by pinning time.Local to UTC (test-only).
• Update agenttool test expectation to reflect multi-part text aggregation.

Doc Coverage / Scenarios (per test)

• TestStreamAggregator

  • Doc: https://ai.google.dev/gemini-api/docs/thought-signatures
  • Quote: "It is advisable to parse the entire request until the finish_reason is returned by the model."
  • Summary: Streaming consumers parse all chunks before final aggregation.

• TestStreamAggregatorStreamingFunctionCallArguments

  • Doc: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling
  • Quote: "Intermediate responses will contain a functionCall object with partialArgs and willContinue fields."
  • Summary: Streamed function-call args arrive across chunks and must be aggregated.

• TestStreamAggregatorFinalizesFunctionCallWhenEmptyChunkArrives

  • Doc: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling
  • Quote: ""functionCall": {}"
  • Summary: Empty functionCall chunk marks the end of streaming args.

• TestStreamAggregatorParallelFunctionCallsPreserveOrderAndSignature

  • Doc: https://ai.google.dev/gemini-api/docs/thought-signatures
  • Quote: "the thought_signature is attached only to the first functionCall part."
  • Summary: Parallel calls keep order and only the first carries the signature.

• TestStreamAggregatorSequentialFunctionCallsPreserveSignatures

  • Doc: https://ai.google.dev/gemini-api/docs/thought-signatures
  • Quote: "each function call will have a signature and you must pass all signatures back."
  • Summary: Each sequential call must keep its own signature.

• TestStreamAggregatorPreservesSignatureOnEmptyFinalTextPart

  • Doc: https://ai.google.dev/gemini-api/docs/thought-signatures
  • Quote: "the model may return the thought signature in a part with an empty text content part."
  • Summary: Preserve empty final text parts that carry signatures.

• TestStreamAggregatorDoesNotMergeSignedTextPart

  • Doc: https://ai.google.dev/gemini-api/docs/function_calling
  • Quote: "Don't merge a Part containing a signature with one that does not."
  • Summary: Signed parts remain distinct from adjacent text.

• TestStreamAggregatorInterleavedStreamingFunctionCalls

  • Doc: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling
  • Quote: "If true, more partialArgs are expected for the overall function call in subsequent streamed responses."
  • Summary: Streaming args may span multiple chunks; isolate each call's args.

• TestStreamAggregatorInterleavedFunctionCallsWithoutIDs

  • Doc: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling
  • Quote: "Intermediate responses will contain a functionCall object with partialArgs and willContinue fields."
  • Summary: Streaming examples omit IDs, so clients should not depend on them.

• TestStreamAggregatorCloseFlushesMultiplePendingFunctionCalls

  • Doc: https://ai.google.dev/gemini-api/docs/thought-signatures
  • Quote: "You must return this signature in the exact part where it was received when sending the conversation history back to the model."
  • Summary: Pending function-call parts must be preserved to avoid dropping signatures.

• TestStreamAggregatorMixedTextAndFunctionCallOrder

  • Doc: https://ai.google.dev/gemini-api/docs/function_calling
  • Quote: "include the responses in the same order as they were requested."
  • Summary: Preserve ordering when mixing text and function calls.

Tests

• go test ./internal/llminternal

Notes / Limitations

• When call IDs are missing and multiple streamed calls are active, we fall back to deterministic ordering; ambiguous same-name concurrent calls may still be indistinguishable.

[gemini-code-assist]

Summary of Changes

Hello @cpunion, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request focuses on ensuring the consistent preservation of the ThoughtSignature field within the Gemini framework, particularly during the aggregation of streaming function calls. This is crucial for maintaining the integrity of internal model states across partial responses. The changes include robust handling of function call arguments and comprehensive test coverage to validate this behavior, alongside an improvement in database test reliability through standardized time handling.

Highlights

• Thought Signature Preservation: The ThoughtSignature field is now correctly preserved when aggregating streaming function calls, including scenarios with empty continuation chunks.
• Enhanced Test Coverage: New tests have been added to thoroughly verify that the ThoughtSignature is maintained throughout the request processing pipeline and when stored and retrieved from the session database.
• Database Test Stability: Database tests have been stabilized by explicitly setting the local time zone to UTC, which prevents inconsistencies related to time zone differences during test execution.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively adds support for preserving the thought_signature across streaming tool calls, which is a valuable enhancement. The changes are well-structured, and I appreciate the comprehensive test coverage added to verify this new behavior at different layers, including stream aggregation, request processing, and the database round-trip. The stabilization of database tests by pinning the timezone to UTC is also a great improvement for test reliability. I have a couple of suggestions to improve the robustness of the new JSON path parsing logic.

[cpunion]

/gemini review

[gemini-code-assist]

Code Review

This is a substantial and well-executed pull request that significantly enhances the handling of Gemini's streaming tool calls, with a key focus on preserving thought_signature. The refactoring of stream_aggregator.go is extensive and thoughtfully implemented to support complex scenarios like parallel and interleaved function calls. The accompanying test suite is exceptionally thorough, covering a wide range of edge cases and linking back to the official documentation, which adds a great deal of confidence in the changes. The improvement to stabilize database tests by setting a fixed timezone is also a welcome addition for test reliability. I have one suggestion regarding the new JSONPath parsing logic to improve its robustness.