Python: preserve Gemini 3 thought_signature across function-call replays#7095
Python: preserve Gemini 3 thought_signature across function-call replays#7095giles17 wants to merge 3 commits into
Conversation
Gemini 3 requires the opaque thought_signature attached to each functionCall part to be echoed back on every replay of that call, or the request is rejected with 400 INVALID_ARGUMENT. The signature previously survived only via raw_representation, so any layer that reconstructs a FunctionCallContent (e.g. harness tool approval) dropped it and broke the next step of the tool loop. Capture the signature into additional_properties on parse and replay it when building the Gemini Part, independent of raw_representation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 33233834-bc6e-4ad2-a6f3-6f1d6e57b1d2
There was a problem hiding this comment.
Pull request overview
This PR updates the Python Gemini chat client to preserve Gemini 3 thought_signature metadata across function-call replays, so reconstructed tool calls (e.g., harness approval flows) don’t fail with 400 INVALID_ARGUMENT when the prior functionCall is replayed.
Changes:
- Capture
thought_signaturefrom incoming GeminiPartobjects during_parse_partsintoContent.additional_properties. - Replay
thought_signaturefromContent.additional_propertieswhen rebuilding outgoing GeminifunctionCallparts in_convert_message_contents. - Add unit tests covering capture/replay behavior and a parse→rebuild round-trip.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| python/packages/gemini/agent_framework_gemini/_chat_client.py | Captures and replays thought_signature for function-call parts across reconstructed-call paths. |
| python/packages/gemini/tests/test_gemini_client.py | Adds unit tests validating thought signature capture and replay semantics. |
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||
Content.additional_properties is serialized via json.dumps(message.to_dict()) by history providers (e.g. RedisHistoryProvider), which fails on raw bytes. Store the thought_signature as a base64 string on parse and decode it back to bytes when building the Gemini Part. Also narrow call_id/name in the round-trip test to satisfy the type checkers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 33233834-bc6e-4ad2-a6f3-6f1d6e57b1d2
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Automated Code Review
Reviewers: 5 | Confidence: 85%
✓ Correctness
The implementation is correct. It properly captures Gemini 3's thought_signature as a base64-encoded string during parsing (ensuring JSON serializability for history providers), and decodes it back to bytes when constructing the Gemini Part for replay. The backfill logic for raw Parts that lack a signature, the new-Part construction path, and the no-op paths when no signature is present are all handled correctly. The previous review feedback about bytes vs base64 serialization has been fully addressed.
✓ Security Reliability
The PR correctly preserves Gemini 3 thought_signature across function-call replays using base64 encoding in additional_properties. The implementation addresses the previously-resolved review comments about JSON serializability. The main reliability concern is that base64.b64decode at line 698 has no error handling — a corrupted value in persisted history would crash the entire agent tool loop rather than gracefully degrading, unlike similar decode operations elsewhere in this file that use try/except.
✓ Test Coverage
The test coverage for the new thought_signature preservation logic is comprehensive and well-structured. Tests cover: capture on parse, JSON serialization round-trip, no-op when absent, replay from additional_properties (reconstructed calls), no-op replay, and full parse→reconstruct→replay round-trip. Assertions are meaningful (checking specific decoded values rather than just 'no exception'). The only untested branch is the defensive backfill at line 701-702 in _convert_message_contents, where a raw Part exists with function_call but has lost its thought_signature while additional_properties still carries one — this cannot occur in normal _parse_parts flow and is purely protective code.
✓ Failure Modes
The PR correctly implements thought_signature preservation across function-call replays using base64 encoding for JSON serializability. The parse path stores signatures as base64 strings, and the replay path decodes them back to bytes. Both the raw-Part-present and raw-Part-absent code paths are handled correctly with appropriate None guards. The backfill logic (line 701-702) correctly defers to the raw Part's own signature when available. No silent failures, swallowed exceptions, or operational failure modes are introduced by this diff.
✓ Design Approach
The fix is targeted at the right layer, but the replay path currently trusts
Content.additional_properties["thought_signature"]too loosely. BecauseContent.additional_propertiesis explicitly untyped (dict[str, Any]), decoding any truthy value withbase64.b64decode(...)can silently turn malformed persisted data into the wrong bytes (for example, non-base64 text can decode tob''or arbitrary bytes) and reintroduce the same invalid replay this PR is trying to prevent.
Automated review by giles17's agents
Guard the untyped additional_properties value with an isinstance(str) check and decode with validate=True, degrading gracefully (warn + drop the signature) on malformed data instead of raising binascii.Error mid tool loop. Matches the defensive base64 handling already used for data URIs in this file. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 33233834-bc6e-4ad2-a6f3-6f1d6e57b1d2
| # Capture Gemini 3's thought_signature (on the Part, not the FunctionCall) as a base64 | ||
| # string so it stays JSON-serializable when messages are persisted and can be replayed. | ||
| additional_properties = ( | ||
| {"thought_signature": base64.b64encode(part.thought_signature).decode("utf-8")} |
There was a problem hiding this comment.
we should use protected_data as part of reasoning_content for this.
Motivation & Context
Gemini 3 models return an opaque, cryptographically-signed
thought_signatureon eachfunctionCallpart and require it echoed back on every later turn that replays that call (Google docs). A request that replays a function call without its signature is rejected with400 INVALID_ARGUMENT.agent_framework_geminipreserved the signature only when the replayedFunctionCallContentstill carried its originaltypes.Partasraw_representation. Any layer that reconstructs a call fromcall_id/name/arguments — the harnessToolApprovalMiddlewarereleasing an auto-approved tool, analways_requireMCP tool, a history rebuilt from a transcript — produced a freshFunctionCallContentwithout that raw part, so the client emitted a signature-lessfunctionCallpart on the next step and the run died. In practice a Gemini 3 harness agent broke on the second step of essentially any tool loop. Disabling thinking does not help — the signature is required regardless ofThinkingConfig.Description & Review Guide
_parse_parts: when afunctionCallpart carries athought_signature(a field ontypes.Part, notFunctionCall), capture it intoContent.additional_properties["thought_signature"]so it survives content transformations that dropraw_representation._convert_message_contents: replay the signature fromadditional_propertieswhen building the GeminiPart— in the reconstructed-call path, and as a backfill when a rawPartis present but lacks one. Applied unconditionally, independent ofThinkingConfig.additional_properties(the harness reconstruction path), no-op replay, and a full parse -> rebuild -> replay round-trip.always_require, rebuilt history) no longer fail with400 INVALID_ARGUMENTon the continuation step. No public API change; the signature is stored under a bareadditional_propertieskey, matching the existing OpenAI (server_label,fc_id) convention.MESSAGES_SNAPSHOT) is out of scope here —thought_signatureisbytes, so serialized hosts need the ag-ui adapters to carry it separately; tracked as a follow-up. This client-level fix covers the common in-process multi-step tool loop.Related Issue
Fixes #6963
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.