From dfd50ceccd8ff6b743bc5f23a2dff0d2ac5aa3b9 Mon Sep 17 00:00:00 2001
From: 0xbyt4 <35742124+0xbyt4@users.noreply.github.com>
Date: Sat, 28 Feb 2026 18:01:13 +0300
Subject: [PATCH] fix: preserve Gemini thought_signature in tool call messages

Gemini 3 thinking models attach extra_content with thought_signature
to function call responses. This must be echoed back on subsequent
API calls or the server rejects with a 400 error. The assistant
message builder was dropping this field, causing all Gemini 3 Flash/Pro
tool-calling flows to fail after the first function call.
---
 run_agent.py            | 18 ++++++++++++++----
 tests/test_run_agent.py | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/run_agent.py b/run_agent.py
index 59a547f0df..3a939d1617 100644
--- a/run_agent.py
+++ b/run_agent.py
@@ -1369,8 +1369,9 @@ class AIAgent:
             ]
 
         if assistant_message.tool_calls:
-            msg["tool_calls"] = [
-                {
+            tc_list = []
+            for tool_call in assistant_message.tool_calls:
+                tc_dict = {
                     "id": tool_call.id,
                     "type": tool_call.type,
                     "function": {
@@ -1378,8 +1379,17 @@ class AIAgent:
                         "arguments": tool_call.function.arguments
                     }
                 }
-                for tool_call in assistant_message.tool_calls
-            ]
+                # Preserve extra_content (e.g. Gemini thought_signature) so it
+                # is sent back on subsequent API calls.  Without this, Gemini 3
+                # thinking models reject the request with a 400 error.
+                extra = getattr(tool_call, "extra_content", None)
+                if extra is not None:
+                    # Convert Pydantic models to plain dicts for JSON safety
+                    if hasattr(extra, "model_dump"):
+                        extra = extra.model_dump()
+                    tc_dict["extra_content"] = extra
+                tc_list.append(tc_dict)
+            msg["tool_calls"] = tc_list
 
         return msg
 
diff --git a/tests/test_run_agent.py b/tests/test_run_agent.py
index 2d37039337..ad90bd270e 100644
--- a/tests/test_run_agent.py
+++ b/tests/test_run_agent.py
@@ -546,6 +546,24 @@ class TestBuildAssistantMessage:
         result = agent._build_assistant_message(msg, "stop")
         assert result["content"] == ""
 
+    def test_tool_call_extra_content_preserved(self, agent):
+        """Gemini thinking models attach extra_content with thought_signature
+        to tool calls. This must be preserved so subsequent API calls include it."""
+        tc = _mock_tool_call(name="get_weather", arguments='{"city":"NYC"}', call_id="c2")
+        tc.extra_content = {"google": {"thought_signature": "abc123"}}
+        msg = _mock_assistant_msg(content="", tool_calls=[tc])
+        result = agent._build_assistant_message(msg, "tool_calls")
+        assert result["tool_calls"][0]["extra_content"] == {
+            "google": {"thought_signature": "abc123"}
+        }
+
+    def test_tool_call_without_extra_content(self, agent):
+        """Standard tool calls (no thinking model) should not have extra_content."""
+        tc = _mock_tool_call(name="web_search", arguments='{}', call_id="c3")
+        msg = _mock_assistant_msg(content="", tool_calls=[tc])
+        result = agent._build_assistant_message(msg, "tool_calls")
+        assert "extra_content" not in result["tool_calls"][0]
+
 
 class TestFormatToolsForSystemMessage:
     def test_no_tools_returns_empty_array(self, agent):