Python: preserve Gemini 3 thought_signature across function-call replays

python/packages/gemini/agent_framework_gemini/_chat_client.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import base64
 import json
 import logging
 import sys
@@ -663,7 +664,22 @@ def _convert_message_contents(
             A list of Gemini Part objects representing the message contents.
         """
         parts: list[types.Part] = []
+        pending_signature: bytes | None = None
         for content in 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
+                encoded_signature = content.protected_data
+                if isinstance(encoded_signature, str) and encoded_signature:
+                    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
             match content.type:
                 case "text":
                     parts.append(types.Part(text=content.text or ""))
@@ -691,10 +707,15 @@ def _convert_message_contents(
                         name=content.name or "",
                         args=content.parse_arguments() or {},
                     )
+                    # Echo the signature from the preceding reasoning content, backfilling only when
+                    # the raw Part lacks one.
                     if isinstance(raw_part, types.Part) and raw_part.function_call is not None:
-                        parts.append(raw_part.model_copy(update={"function_call": function_call}, deep=True))
+                        replayed_part = raw_part.model_copy(update={"function_call": function_call}, deep=True)
+                        if replayed_part.thought_signature is None and thought_signature is not None:
+                            replayed_part.thought_signature = thought_signature
+                        parts.append(replayed_part)
                     else:
-                        parts.append(types.Part(function_call=function_call))
+                        parts.append(types.Part(function_call=function_call, thought_signature=thought_signature))
                 case "function_result":
                     raw_part = content.raw_representation
                     if isinstance(raw_part, types.Part) and raw_part.tool_response is not None:
@@ -1122,6 +1143,14 @@ def _parse_parts(self, parts: Sequence[types.Part]) -> list[Content]:
                 else:
                     call_id = self._generate_tool_call_id()
                     logger.debug("function_call missing id; generated fallback call_id=%r", call_id)
+                # Surface Gemini 3's thought_signature as reasoning content preceding the call so it
+                # survives when the call is later reconstructed without its raw Part.
+                if part.thought_signature is not None:
+                    contents.append(
+                        Content.from_text_reasoning(
+                            protected_data=base64.b64encode(part.thought_signature).decode("utf-8")
+                        )
+                    )
                 contents.append(
                     Content.from_function_call(
                         call_id=call_id,

python/packages/gemini/tests/test_gemini_client.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
+import base64
 import datetime
+import json
 import logging
 import os
 from typing import Any, cast
@@ -46,6 +48,7 @@ def _make_part(
     text: str | None = None,
     thought: bool = False,
     function_call: tuple[str | None, str, dict[str, Any]] | None = None,
+    thought_signature: bytes | None = None,
     tool_call: tuple[str | None, types.ToolType, dict[str, Any]] | None = None,
     tool_response: tuple[str | None, types.ToolType, dict[str, Any]] | None = None,
     executable_code: str | None = None,
@@ -57,6 +60,7 @@ def _make_part(
         text: Text content of the part.
         thought: Whether this is a thinking/reasoning part.
         function_call: Tuple of (id, name, args) if this is a function call part.
+        thought_signature: Opaque Gemini 3 thought signature attached to the part, if any.
         tool_call: Tuple of (id, tool_type, args) if this is a server-side tool call part.
         tool_response: Tuple of (id, tool_type, response) if this is a server-side tool response part.
         executable_code: Source code string for a code execution part.
@@ -65,6 +69,7 @@ def _make_part(
     part = MagicMock()
     part.text = text
     part.thought = thought
+    part.thought_signature = thought_signature
     part.function_response = None
     part.tool_call = None
     part.tool_response = None
@@ -741,6 +746,110 @@ def test_function_call_part_preserves_thought_signature_from_raw_part() -> None:
     assert parts[0].function_call.args == {"location": "Paris"}
 
 
+def test_function_call_part_captures_thought_signature_as_reasoning_content() -> None:
+    """Parsing a functionCall part surfaces its thought_signature as preceding reasoning content."""
+    client, _ = _make_gemini_client()
+
+    contents = client._parse_parts([
+        _make_part(function_call=("call-1", "get_weather", {"location": "Paris"}), thought_signature=b"sig-123")
+    ])
+
+    assert len(contents) == 2
+    assert contents[0].type == "text_reasoning"
+    assert contents[0].protected_data == base64.b64encode(b"sig-123").decode("utf-8")
+    assert contents[1].type == "function_call"
+    assert contents[1].call_id == "call-1"
+
+
+def test_captured_thought_signature_is_json_serializable() -> None:
+    """The captured signature must survive json.dumps(message.to_dict()) used by history providers."""
+    client, _ = _make_gemini_client()
+    contents = client._parse_parts([
+        _make_part(function_call=("call-1", "get_weather", {"location": "Paris"}), thought_signature=b"sig-123")
+    ])
+    message = Message(role="assistant", contents=contents)
+
+    serialized = json.dumps(message.to_dict())
+
+    restored = Message.from_dict(json.loads(serialized))
+    parts = client._convert_message_contents(restored.contents, {})
+    # Reasoning content is not sent back as a part; only the function call is, carrying the signature.
+    assert len(parts) == 1
+    assert parts[0].thought_signature == b"sig-123"
+
+
+def test_function_call_part_without_thought_signature_emits_no_reasoning_content() -> None:
+    """A functionCall part without a signature produces only the function_call content."""
+    client, _ = _make_gemini_client()
+
+    contents = client._parse_parts([_make_part(function_call=("call-1", "get_weather", {"location": "Paris"}))])
+
+    assert len(contents) == 1
+    assert contents[0].type == "function_call"
+
+
+def test_reconstructed_function_call_replays_thought_signature_from_reasoning_content() -> None:
+    """A function call rebuilt without its raw Part still replays the signature from reasoning content."""
+    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, call], {})
+
+    assert len(parts) == 1
+    assert parts[0].thought_signature == b"sig-123"
+    assert parts[0].function_call is not None
+    assert parts[0].function_call.id == "call-1"
+
+
+def test_function_call_without_thought_signature_replays_without_one() -> None:
+    """A function call lacking any signature produces a part with no thought_signature."""
+    client, _ = _make_gemini_client()
+    content = Content.from_function_call(call_id="call-1", name="get_weather", arguments={"location": "Paris"})
+
+    parts = client._convert_message_contents([content], {})
+
+    assert len(parts) == 1
+    assert parts[0].thought_signature is None
+
+
+def test_malformed_thought_signature_is_ignored_gracefully(caplog: pytest.LogCaptureFixture) -> None:
+    """A corrupted or non-string reasoning signature degrades to no signature instead of crashing."""
+    client, _ = _make_gemini_client()
+    corrupted = Content.from_text_reasoning(protected_data="not valid base64!!!")
+    call1 = Content.from_function_call(call_id="call-1", name="get_weather", arguments={"location": "Paris"})
+    non_string = Content("text_reasoning", protected_data=cast(Any, 123))
+    call2 = Content.from_function_call(call_id="call-2", name="get_weather", arguments={"location": "Paris"})
+
+    with caplog.at_level(logging.WARNING):
+        parts = client._convert_message_contents([corrupted, call1, non_string, call2], {})
+
+    assert len(parts) == 2
+    assert parts[0].thought_signature is None
+    assert parts[1].thought_signature is None
+    assert "malformed thought_signature" in caplog.text
+
+
+def test_reconstructed_function_call_signature_survives_round_trip() -> None:
+    """Parse yields reasoning+call; a rebuilt call (raw Part dropped) still replays the signature."""
+    client, _ = _make_gemini_client()
+
+    parsed = client._parse_parts([
+        _make_part(function_call=("call-1", "get_weather", {"location": "Paris"}), thought_signature=b"sig-123")
+    ])
+    reasoning = parsed[0]
+    fc = parsed[1]
+    assert fc.call_id is not None
+    assert fc.name is not None
+    # Simulate a layer that reconstructs the call from call_id/name/arguments, dropping raw_representation,
+    # while the sibling reasoning content is preserved in history.
+    rebuilt_call = Content.from_function_call(call_id=fc.call_id, name=fc.name, arguments=fc.arguments)
+
+    parts = client._convert_message_contents([reasoning, rebuilt_call], {})
+
+    assert parts[-1].thought_signature == b"sig-123"
+
+
 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()