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