diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index bfb4f67005c..232f8e8ea2f 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/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, diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index d3ad53c5abb..f2af3cfd7ae 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/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()