From 8ade61be46be6097ec8f922cfa28fd0688a7dc60 Mon Sep 17 00:00:00 2001 From: Giles Odigwe Date: Thu, 6 Aug 2026 12:59:07 -0700 Subject: [PATCH] Python: Restore Gemini thought_signature on approval replays Gemini 3.x rejects a request whose functionCall parts lack a thought_signature. The signature was carried as base64 protected_data on a text_reasoning content and re-attached by adjacency, which requires the carrier to immediately precede its call. An approval round trip replays the call with no carrier at all, so the next turn failed with a 400. Track signatures in a bounded per-client call_id map populated at parse time from the resolved call_id, and backfill only when the emitted part has no signature. Also stop clearing the held signature on contents that emit no Part, so an approval response or an unsigned thought summary between the carrier and its call no longer drops it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: dd0909cd-c7c3-42cb-aef1-1e9a3e64d917 --- .../agent_framework_gemini/_chat_client.py | 74 ++++++++++- .../gemini/tests/test_gemini_client.py | 118 +++++++++++++++++- 2 files changed, 187 insertions(+), 5 deletions(-) diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index cd053711ca..f24064be00 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -6,6 +6,7 @@ import json import logging import sys +from collections import OrderedDict from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence from typing import Any, ClassVar, Generic, cast from uuid import uuid4 @@ -184,6 +185,7 @@ class GoogleGeminiSettings(TypedDict, total=False): _GEMINI_API_BASE_URL = "https://generativelanguage.googleapis.com" _VERTEX_AI_BASE_URL = "https://aiplatform.googleapis.com" +_DEFAULT_MAX_THOUGHT_SIGNATURES = 256 def _resolve_vertexai_mode(client: genai.Client, *, fallback: bool | None = None) -> bool: @@ -327,6 +329,7 @@ def __init__( env_file_encoding: str | None = None, client: genai.Client | None = None, additional_properties: dict[str, Any] | None = None, + max_tracked_thought_signatures: int = _DEFAULT_MAX_THOUGHT_SIGNATURES, ) -> None: """Create a raw Gemini chat client. @@ -347,7 +350,14 @@ def __init__( env_file_encoding: Encoding for the ``.env`` file. client: Pre-built ``genai.Client`` instance. When provided, connector auth settings are not required. additional_properties: Extra properties stored on the client instance. + max_tracked_thought_signatures: Maximum number of Gemini 3 thought signatures retained + for replay, keyed by call ID. Least-recently-used entries are evicted beyond this. + + Raises: + ValueError: If ``max_tracked_thought_signatures`` is less than 1. """ + if max_tracked_thought_signatures < 1: + raise ValueError("max_tracked_thought_signatures must be greater than 0.") settings = load_settings( GeminiSettings, env_prefix="GEMINI_", @@ -408,6 +418,8 @@ def __init__( self._vertexai = _resolve_vertexai_mode(self._genai_client, fallback=configured_vertexai) self._service_url = _resolve_service_url(self._genai_client, vertexai=self._vertexai) self.model = google_settings.get("model") or settings.get("model") + self.max_tracked_thought_signatures = max_tracked_thought_signatures + self._thought_signature_cache: OrderedDict[str, bytes] = OrderedDict() super().__init__(additional_properties=additional_properties) @@ -677,23 +689,35 @@ def _convert_message_contents( if content.type == "text_reasoning": # Gemini 3's thought_signature travels as base64 protected_data on reasoning content; # hold it for the function call it precedes (reasoning is not sent back as a Part). - pending_signature = None + # Reasoning without protected_data (a thought summary) must not clear a held + # signature, otherwise the pairing breaks on ordering alone. encoded_signature = content.protected_data if isinstance(encoded_signature, str) and encoded_signature: + pending_signature = None try: pending_signature = base64.b64decode(encoded_signature, validate=True) except ValueError: logger.warning("Ignoring malformed thought_signature on reasoning content") continue - # A signature applies only to a function call immediately following its reasoning content. - thought_signature = pending_signature - pending_signature = None + thought_signature: bytes | None = None + if content.type == "function_call": + # A held signature belongs to the next function call in the message. + thought_signature = pending_signature + pending_signature = None + elif content.type in ("text", "function_result", "data", "uri"): + # Content that emits its own Part breaks the reasoning-to-call pairing. Content that + # emits nothing (approval requests and responses) is left transparent. + pending_signature = None match content.type: case "text": parts.append(types.Part(text=content.text or "")) case "function_call": call_id = content.call_id or self._generate_tool_call_id() raw_part = content.raw_representation + if thought_signature is None: + # No adjacent carrier: fall back to the signature recorded for this call at + # parse time. Covers replays where the carrier was dropped in transit. + thought_signature = self._recall_thought_signature(call_id) if ( content.informational_only and isinstance(raw_part, types.Part) @@ -1178,6 +1202,10 @@ def _parse_parts(self, parts: Sequence[types.Part]) -> list[Content]: protected_data=base64.b64encode(part.thought_signature).decode("utf-8") ) ) + # Also key it by the resolved call_id so the signature survives surfaces that + # drop the reasoning carrier. Captured here, not from the raw Part, because the + # raw part's id may be absent and replaced by a generated fallback above. + self._remember_thought_signature(call_id, part.thought_signature) contents.append( Content.from_function_call( call_id=call_id, @@ -1253,6 +1281,37 @@ def _generate_tool_call_id() -> str: """ return f"tool-call-{uuid4().hex}" + # region Thought signature tracking + + def _remember_thought_signature(self, call_id: str, signature: bytes) -> None: + """Record a thought signature so a later replay of the same call can be re-signed. + + Correlating by ``call_id`` keeps the signature recoverable when the reasoning content that + carries it is dropped by a surface (for example an approval round trip through a client). + + Args: + call_id: The resolved framework call ID of the function call the signature belongs to. + signature: The opaque Gemini thought signature bytes. + """ + cache = self._thought_signature_cache + cache[call_id] = signature + cache.move_to_end(call_id) + while len(cache) > self.max_tracked_thought_signatures: + cache.popitem(last=False) + + def _recall_thought_signature(self, call_id: str) -> bytes | None: + """Look up a previously recorded thought signature for a function call. + + Args: + call_id: The framework call ID of the function call being serialized. + + Returns: + The recorded signature bytes, or ``None`` when nothing was recorded for this call. + """ + return self._thought_signature_cache.get(call_id) + + # endregion + class GeminiChatClient( FunctionInvocationLayer[GeminiChatOptionsT], @@ -1290,6 +1349,7 @@ def __init__( additional_properties: dict[str, Any] | None = None, middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None, function_invocation_configuration: FunctionInvocationConfiguration | None = None, + max_tracked_thought_signatures: int = _DEFAULT_MAX_THOUGHT_SIGNATURES, ) -> None: """Create a Gemini chat client. @@ -1309,6 +1369,11 @@ def __init__( additional_properties: Extra properties stored on the client instance. middleware: Optional middleware chain applied to every call. function_invocation_configuration: Optional configuration for the function invocation loop. + max_tracked_thought_signatures: Maximum number of Gemini 3 thought signatures retained + for replay, keyed by call ID. Least-recently-used entries are evicted beyond this. + + Raises: + ValueError: If ``max_tracked_thought_signatures`` is less than 1. """ super().__init__( api_key=api_key, @@ -1323,4 +1388,5 @@ def __init__( additional_properties=additional_properties, middleware=middleware, function_invocation_configuration=function_invocation_configuration, + max_tracked_thought_signatures=max_tracked_thought_signatures, ) diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index 7dbda00ca9..e750c193da 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -162,12 +162,13 @@ async def _async_iter(items: list[Any]): def _make_gemini_client( model: str | None = "gemini-2.5-flash", mock_client: MagicMock | None = None, + **kwargs: Any, ) -> tuple[GeminiChatClient, MagicMock]: """Return a (GeminiChatClient, mock_genai_client) pair.""" mock = mock_client or MagicMock() mock._api_client.vertexai = False mock._api_client._http_options.base_url = "https://generativelanguage.googleapis.com/" - client = GeminiChatClient(client=mock, model=model) + client = GeminiChatClient(client=mock, model=model, **kwargs) return client, mock @@ -860,6 +861,121 @@ def test_reconstructed_function_call_signature_survives_round_trip() -> None: assert parts[-1].thought_signature == b"sig-123" +def test_replayed_call_is_signed_by_call_id_when_the_carrier_is_dropped() -> None: + """The approval replay history has no carrier at all; the call is still re-signed by call_id.""" + client, _ = _make_gemini_client() + parsed = client._parse_parts([ + _make_part(function_call=("call-1", "get_weather", {"location": "Paris"}), thought_signature=b"sig-123") + ]) + call = parsed[1] + assert call.call_id is not None + # The history an approval round trip produces: the call rebuilt from the function_call nested in + # the approval content, with neither a reasoning carrier nor the original raw Part. + replay = [Content.from_function_call(call_id=call.call_id, name="get_weather", arguments={"location": "Paris"})] + assert [content.type for content in replay] == ["function_call"] + assert not any(content.type == "text_reasoning" and content.protected_data for content in replay) + assert not isinstance(replay[0].raw_representation, types.Part) + + parts = client._convert_message_contents(replay, {}) + + assert len(parts) == 1 + assert parts[0].thought_signature == b"sig-123" + + +def test_call_id_backfill_uses_the_generated_id_when_gemini_omits_one() -> None: + """Signatures are keyed by the resolved call_id, so the generated-id fallback path is covered.""" + client, _ = _make_gemini_client() + part = _make_part(function_call=(None, "get_weather", {"location": "Paris"}), thought_signature=b"sig-123") + + parsed = client._parse_parts([part]) + call = parsed[1] + assert call.call_id is not None + assert call.call_id.startswith("tool-call-") + replay = [Content.from_function_call(call_id=call.call_id, name="get_weather", arguments={"location": "Paris"})] + + parts = client._convert_message_contents(replay, {}) + + assert parts[0].thought_signature == b"sig-123" + + +@pytest.mark.parametrize( + "intervening", + [ + pytest.param( + Content.from_function_approval_response( + True, + id="call-1", + function_call=Content.from_function_call(call_id="call-1", name="get_weather", arguments={}), + ), + id="approval_response", + ), + pytest.param(Content.from_text_reasoning(text="a thought summary"), id="unsigned_reasoning"), + ], +) +def test_signature_survives_content_between_the_carrier_and_the_call(intervening: Content) -> None: + """Content that emits no Part must not break the reasoning-to-call pairing.""" + client, _ = _make_gemini_client() + reasoning = Content.from_text_reasoning(protected_data=base64.b64encode(b"sig-123").decode("utf-8")) + call = Content.from_function_call(call_id="call-1", name="get_weather", arguments={"location": "Paris"}) + + parts = client._convert_message_contents([reasoning, intervening, call], {}) + + assert parts[-1].function_call is not None + assert parts[-1].thought_signature == b"sig-123" + + +def test_call_id_backfill_never_overrides_a_signature_already_on_the_part() -> None: + """Backfill only fills gaps, so a raw Part's own signature always wins.""" + client, _ = _make_gemini_client() + client._parse_parts([ + _make_part(function_call=("call-1", "get_weather", {"location": "Paris"}), thought_signature=b"stale-sig") + ]) + raw_part = types.Part( + function_call=types.FunctionCall(id="call-1", name="get_weather", args={"location": "Paris"}), + thought_signature=b"fresh-sig", + ) + call = Content.from_function_call( + call_id="call-1", + name="get_weather", + arguments={"location": "Paris"}, + raw_representation=raw_part, + ) + + parts = client._convert_message_contents([call], {}) + + assert parts[0].thought_signature == b"fresh-sig" + + +def test_thought_signature_cache_is_bounded() -> None: + """The per-client signature cache must not grow without limit on long conversations.""" + client, _ = _make_gemini_client() + overflow = client.max_tracked_thought_signatures + 5 + + for index in range(overflow): + client._parse_parts([_make_part(function_call=(f"call-{index}", "get_weather", {}), thought_signature=b"sig")]) + + cache = client._thought_signature_cache + assert len(cache) == client.max_tracked_thought_signatures + assert "call-0" not in cache + assert f"call-{overflow - 1}" in cache + + +def test_max_tracked_thought_signatures_is_configurable() -> None: + """The retention bound is a constructor option, so hosts can tune it per client.""" + client, _ = _make_gemini_client(max_tracked_thought_signatures=2) + + for index in range(3): + client._parse_parts([_make_part(function_call=(f"call-{index}", "get_weather", {}), thought_signature=b"sig")]) + + assert list(client._thought_signature_cache) == ["call-1", "call-2"] + + +def test_max_tracked_thought_signatures_rejects_non_positive_values() -> None: + """A bound below 1 would evict every signature immediately, so it is rejected up front.""" + with pytest.raises(ValueError, match="max_tracked_thought_signatures"): + _make_gemini_client(max_tracked_thought_signatures=0) + + def test_server_side_tool_call_part_is_informational_only() -> None: """Server-side Gemini tool calls are transcript content, not local function invocation requests.""" client, _ = _make_gemini_client()