fix(gemini): replay thought signatures for tool calls

docs/reference/provider-capabilities.md

@@ -91,6 +91,11 @@ jobs, see [Building With Deferred Delivery](../building-with-deferred-delivery.m
   when the selected model and reasoning mode support it. Pollux forwards both
   `reasoning_effort` and `reasoning_budget_tokens`; model-specific acceptance
   is enforced by the Gemini API.
+- Gemini function-call thought signatures and call IDs are preserved in opaque
+  `continuation` state and replayed with tool results. This applies to both
+  `interact()` and `stream()` and is required for Gemini 3 tool loops. Keep the
+  continuation untouched until the pending tool exchange is complete; portable
+  application-authored `history` does not contain provider signatures.
 
 ### OpenAI
 

src/pollux/providers/gemini.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import asyncio
+import base64
 import contextlib
+from dataclasses import replace
 from datetime import datetime
 import json
 import logging
@@ -48,6 +50,7 @@
 logger = logging.getLogger(__name__)
 
 _GEMINI_BATCH_INLINE_LIMIT_BYTES = 20_000_000
+_GEMINI_FUNCTION_CALLS_KEY = "gemini_function_calls"
 
 
 def _provider_hint_payload(
@@ -306,15 +309,23 @@ def _build_contents(
         self,
         parts: list[Any],
         history: list[Message] | None,
+        provider_state: dict[str, object] | None = None,
     ) -> Any:
         """Build Gemini contents from history + current parts."""
         from google.genai import types
 
         input_contents: list[Any] = []
         call_id_to_name: dict[str, str] = {}
+        history_states = (
+            provider_state.get("history", [])
+            if isinstance(provider_state, dict)
+            else []
+        )
+        if not isinstance(history_states, list):
+            history_states = []
 
         if history:
-            for item in history:
+            for item_index, item in enumerate(history):
                 if item.role == "tool":
                     call_id = item.tool_call_id
                     name = "unknown_tool"
@@ -334,15 +345,25 @@ def _build_contents(
                         types.Content(
                             role="user",
                             parts=[
-                                types.Part.from_function_response(
-                                    name=name,
-                                    response=parsed_content,
+                                types.Part(
+                                    function_response=types.FunctionResponse(
+                                        id=call_id,
+                                        name=name,
+                                        response=parsed_content,
+                                    )
                                 )
                             ],
                         )
                     )
                 elif item.role == "assistant":
                     ast_parts: list[Any] = []
+                    item_state = (
+                        history_states[item_index]
+                        if item_index < len(history_states)
+                        and isinstance(history_states[item_index], dict)
+                        else None
+                    )
+                    signatures = _function_call_signatures(item_state)
                     if item.content:
                         ast_parts.append(types.Part.from_text(text=item.content))
                     if item.tool_calls:
@@ -353,8 +374,13 @@ def _build_contents(
                             except Exception:
                                 args_dict = {}
                             ast_parts.append(
-                                types.Part.from_function_call(
-                                    name=tc.name, args=args_dict
+                                types.Part(
+                                    function_call=types.FunctionCall(
+                                        id=tc.id,
+                                        name=tc.name,
+                                        args=args_dict,
+                                    ),
+                                    thought_signature=signatures.get(tc.id),
                                 )
                             )
                     if ast_parts:
@@ -430,9 +456,11 @@ async def submit_deferred(
                         config,
                         upload_cache=upload_cache,
                     )
-                    history, _prev, _ps = _compile.prior_turns(inp)
+                    history, _prev, provider_state = _compile.prior_turns(inp)
                     inlined_request = types.InlinedRequest(
-                        contents=self._build_contents(parts, history or None),
+                        contents=self._build_contents(
+                            parts, history or None, provider_state
+                        ),
                         config=types.GenerateContentConfig(
                             **self._build_config_kwargs(parts, snapshot, requirements)
                         ),
@@ -621,9 +649,9 @@ async def generate(
         from google.genai import types
 
         parts = _compile.request_parts(snapshot, input)
-        history, _previous_response_id, _provider_state = _compile.prior_turns(input)
+        history, _previous_response_id, provider_state = _compile.prior_turns(input)
         config_kwargs = self._build_config_kwargs(parts, snapshot, requirements)
-        contents = self._build_contents(parts, history or None)
+        contents = self._build_contents(parts, history or None, provider_state)
 
         try:
             response = await client.aio.models.generate_content(
@@ -656,20 +684,22 @@ async def stream_generate(
     ) -> AsyncIterator[ProviderStreamChunk]:
         """Stream normalized deltas from Gemini's streaming content endpoint.
 
-        Gemini replays conversation via history (no provider_state), so thought
-        text is surfaced for display only. Function calls arrive whole rather
-        than fragmented, so each one is emitted as a single-fragment tool-call
-        delta with its own slot index.
+        Gemini replays conversation via history. Thought text is surfaced for
+        display, while signed function-call state is retained opaquely for the
+        next continuation turn. Function calls arrive whole rather than
+        fragmented, so each one is emitted as a single-fragment tool-call delta
+        with its own slot index.
         """
         client = self._get_client()
         from google.genai import types
 
         parts = _compile.request_parts(snapshot, input)
-        history, _previous_response_id, _provider_state = _compile.prior_turns(input)
+        history, _previous_response_id, provider_state = _compile.prior_turns(input)
         config_kwargs = self._build_config_kwargs(parts, snapshot, requirements)
-        contents = self._build_contents(parts, history or None)
+        contents = self._build_contents(parts, history or None, provider_state)
 
         tool_call_index = 0
+        function_call_state: list[dict[str, str]] = []
         stream: Any = None
         try:
             stream = await client.aio.models.generate_content_stream(
@@ -681,7 +711,23 @@ async def stream_generate(
                 for chunk in self._stream_response_to_chunks(response, tool_call_index):
                     if chunk.tool_calls:
                         tool_call_index += len(chunk.tool_calls)
+                    if chunk.provider_state:
+                        entries = chunk.provider_state.get(_GEMINI_FUNCTION_CALLS_KEY)
+                        if isinstance(entries, list):
+                            function_call_state.extend(
+                                entry for entry in entries if isinstance(entry, dict)
+                            )
+                            remaining_state = dict(chunk.provider_state)
+                            remaining_state.pop(_GEMINI_FUNCTION_CALLS_KEY, None)
+                            chunk = replace(
+                                chunk,
+                                provider_state=remaining_state or None,
+                            )
                     yield chunk
+            if function_call_state:
+                yield ProviderStreamChunk(
+                    provider_state={_GEMINI_FUNCTION_CALLS_KEY: function_call_state}
+                )
         except asyncio.CancelledError:
             raise
         except Exception as e:
@@ -717,6 +763,7 @@ def _stream_response_to_chunks(
                         _field(function_call, "id") or f"call_{uuid.uuid4().hex[:8]}"
                     )
                     arguments = json.dumps(_field(function_call, "args") or {})
+                    state_entry = _function_call_state_entry(str(call_id), part)
                     chunks.append(
                         ProviderStreamChunk(
                             tool_calls=(
@@ -726,7 +773,12 @@ def _stream_response_to_chunks(
                                     name=str(_field(function_call, "name", "")),
                                     arguments=arguments,
                                 ),
-                            )
+                            ),
+                            provider_state=(
+                                {_GEMINI_FUNCTION_CALLS_KEY: [state_entry]}
+                                if state_entry is not None
+                                else None
+                            ),
                         )
                     )
                     local_tool_count += 1
@@ -1238,6 +1290,26 @@ def _parse_response(self, response: Any) -> ProviderResponse:
                     )
                 )
 
+        function_call_parts: list[Any] = []
+        candidates = _field(response, "candidates", []) or []
+        if isinstance(candidates, list) and candidates:
+            content = _field(candidates[0], "content")
+            candidate_parts = (
+                _field(content, "parts", []) if content is not None else []
+            )
+            if isinstance(candidate_parts, list):
+                function_call_parts = [
+                    part
+                    for part in candidate_parts
+                    if _field(part, "function_call") is not None
+                ]
+
+        function_call_state = [
+            state
+            for call, part in zip(tool_calls, function_call_parts, strict=False)
+            if (state := _function_call_state_entry(call.id, part)) is not None
+        ]
+
         finish_reason: str | None = None
         reasoning_parts: list[str] = []
         artifacts: dict[str, Any] = {}
@@ -1273,10 +1345,52 @@ def _parse_response(self, response: Any) -> ProviderResponse:
             tool_calls=tool_calls if tool_calls else None,
             response_id=None,
             finish_reason=finish_reason,
+            provider_state=(
+                {_GEMINI_FUNCTION_CALLS_KEY: function_call_state}
+                if function_call_state
+                else None
+            ),
             artifacts=artifacts or None,
         )
 
 
+def _function_call_state_entry(call_id: str, part: Any) -> dict[str, str] | None:
+    """Encode one Gemini function-call thought signature for continuation replay."""
+    signature = _field(part, "thought_signature")
+    if isinstance(signature, bytes) and signature:
+        encoded = base64.b64encode(signature).decode("ascii")
+    elif isinstance(signature, str) and signature:
+        encoded = signature
+    else:
+        return None
+    return {"id": call_id, "thought_signature": encoded}
+
+
+def _function_call_signatures(
+    provider_state: object,
+) -> dict[str, bytes]:
+    """Decode preserved Gemini signatures, keyed by normalized tool-call ID."""
+    if not isinstance(provider_state, dict):
+        return {}
+    entries = provider_state.get(_GEMINI_FUNCTION_CALLS_KEY)
+    if not isinstance(entries, list):
+        return {}
+
+    signatures: dict[str, bytes] = {}
+    for entry in entries:
+        if not isinstance(entry, dict):
+            continue
+        call_id = entry.get("id")
+        encoded = entry.get("thought_signature")
+        if not isinstance(call_id, str) or not isinstance(encoded, str):
+            continue
+        try:
+            signatures[call_id] = base64.b64decode(encoded, validate=True)
+        except ValueError:
+            continue
+    return signatures
+
+
 def _field(obj: Any, key: str, default: Any = None) -> Any:
     """Read a field from either an SDK object or a raw JSON dict."""
     if isinstance(obj, dict):

tests/providers/test_gemini_contract.py

@@ -838,6 +838,39 @@ def test_gemini_parse_response_extracts_tool_calls() -> None:
     assert result.tool_calls[0].arguments == '{"city": "NYC"}'
 
 
+def test_gemini_parse_response_preserves_function_call_thought_signature() -> None:
+    """Gemini function-call signatures survive in opaque replay state."""
+    provider = GeminiProvider("test-key")
+
+    fc = _obj(id="call_abc", name="get_weather", args={"city": "NYC"})
+    function_call_part = _obj(
+        function_call=fc,
+        thought_signature=b"signed-thought",
+        text=None,
+        thought=False,
+    )
+    fake_response = _obj(
+        text="",
+        parsed=None,
+        usage_metadata=None,
+        function_calls=[fc],
+        candidates=[
+            _obj(content=_obj(parts=[function_call_part]), finish_reason="STOP")
+        ],
+    )
+
+    result = provider._parse_response(fake_response)
+
+    assert result.provider_state == {
+        "gemini_function_calls": [
+            {
+                "id": "call_abc",
+                "thought_signature": "c2lnbmVkLXRob3VnaHQ=",
+            }
+        ]
+    }
+
+
 def test_gemini_parse_response_generates_id_when_missing() -> None:
     """Gemini may omit fc.id; a stable synthetic ID should be generated."""
     provider = GeminiProvider("test-key")

tests/providers/test_gemini_stream.py

@@ -29,6 +29,7 @@ def _part(**kwargs: Any) -> SimpleNamespace:
     kwargs.setdefault("text", None)
     kwargs.setdefault("thought", False)
     kwargs.setdefault("function_call", None)
+    kwargs.setdefault("thought_signature", None)
     return SimpleNamespace(**kwargs)
 
 
@@ -54,7 +55,8 @@ def _thinking_tool_chunks() -> list[SimpleNamespace]:
                 _part(
                     function_call=SimpleNamespace(
                         id=None, name="get_weather", args={"city": "NYC"}
-                    )
+                    ),
+                    thought_signature=b"stream-signature",
                 )
             ],
             finish="STOP",
@@ -154,6 +156,11 @@ async def test_gemini_stream_through_interaction_assembles_output() -> None:
     assert done.usage.total_tokens == 15
     assert done.usage.reasoning_tokens == 2
     assert done.metrics.finish_reason == "stop"
+    assert done.continuation is not None
+    state = done.continuation.to_jsonable()["provider_state"]
+    assert state["gemini_function_calls"][0]["thought_signature"] == (
+        "c3RyZWFtLXNpZ25hdHVyZQ=="
+    )
 
 
 @pytest.mark.asyncio