Python: Restore Gemini thought_signature on approval replays

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,
         )

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()