fix: propagate reasoningSignature on Gemini tool use (strands-agents#1703)

afarntrog · web-flow · commit 2281a204db41 · 2026-02-16T15:03:40.000-05:00
diff --git a/src/strands/event_loop/streaming.py b/src/strands/event_loop/streaming.py
@@ -186,6 +186,8 @@ def handle_content_block_start(event: ContentBlockStartEvent) -> dict[str, Any]:
         current_tool_use["toolUseId"] = tool_use_data["toolUseId"]
         current_tool_use["name"] = tool_use_data["name"]
         current_tool_use["input"] = ""
+        if "reasoningSignature" in tool_use_data:
+            current_tool_use["reasoningSignature"] = tool_use_data["reasoningSignature"]
 
     return current_tool_use
 
@@ -286,6 +288,8 @@ def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]:
             name=tool_use_name,
             input=current_tool_use["input"],
         )
+        if "reasoningSignature" in current_tool_use:
+            tool_use["reasoningSignature"] = current_tool_use["reasoningSignature"]
         content.append({"toolUse": tool_use})
         state["current_tool_use"] = {}
 
diff --git a/src/strands/models/gemini.py b/src/strands/models/gemini.py
@@ -3,6 +3,7 @@
 - Docs: https://ai.google.dev/api
 """
 
+import base64
 import json
 import logging
 import mimetypes
@@ -14,7 +15,7 @@
 from google import genai
 from typing_extensions import Required, Unpack, override
 
-from ..types.content import ContentBlock, Messages
+from ..types.content import ContentBlock, ContentBlockStartToolUse, Messages
 from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
 from ..types.streaming import StreamEvent
 from ..types.tools import ToolChoice, ToolSpec
@@ -173,7 +174,7 @@ def _format_request_content_part(
             return genai.types.Part(
                 text=content["reasoningContent"]["reasoningText"]["text"],
                 thought=True,
-                thought_signature=thought_signature.encode("utf-8") if thought_signature else None,
+                thought_signature=base64.b64decode(thought_signature) if thought_signature else None,
             )
 
         if "text" in content:
@@ -202,14 +203,18 @@ def _format_request_content_part(
             )
 
         if "toolUse" in content:
-            tool_use_id_to_name[content["toolUse"]["toolUseId"]] = content["toolUse"]["name"]
+            tool_use_id = content["toolUse"]["toolUseId"]
+            tool_use_id_to_name[tool_use_id] = content["toolUse"]["name"]
+
+            reasoning_signature = content["toolUse"].get("reasoningSignature")
 
             return genai.types.Part(
                 function_call=genai.types.FunctionCall(
                     args=content["toolUse"]["input"],
-                    id=content["toolUse"]["toolUseId"],
+                    id=tool_use_id,
                     name=content["toolUse"]["name"],
                 ),
+                thought_signature=base64.b64decode(reasoning_signature) if reasoning_signature else None,
             )
 
         raise TypeError(f"content_type=<{next(iter(content))}> | unsupported type")
@@ -349,13 +354,18 @@ def _format_chunk(self, event: dict[str, Any]) -> StreamEvent:
                         # Use Gemini's provided ID or generate one if missing
                         tool_use_id = function_call.id or f"tooluse_{secrets.token_urlsafe(16)}"
 
+                        tool_use_start: ContentBlockStartToolUse = {
+                            "name": function_call.name,
+                            "toolUseId": tool_use_id,
+                        }
+                        if event["data"].thought_signature:
+                            tool_use_start["reasoningSignature"] = base64.b64encode(
+                                event["data"].thought_signature
+                            ).decode("ascii")
                         return {
                             "contentBlockStart": {
                                 "start": {
-                                    "toolUse": {
-                                        "name": function_call.name,
-                                        "toolUseId": tool_use_id,
-                                    },
+                                    "toolUse": tool_use_start,
                                 },
                             },
                         }
@@ -379,7 +389,11 @@ def _format_chunk(self, event: dict[str, Any]) -> StreamEvent:
                                     "reasoningContent": {
                                         "text": event["data"].text,
                                         **(
-                                            {"signature": event["data"].thought_signature.decode("utf-8")}
+                                            {
+                                                "signature": base64.b64encode(event["data"].thought_signature).decode(
+                                                    "ascii"
+                                                )
+                                            }
                                             if event["data"].thought_signature
                                             else {}
                                         ),
diff --git a/src/strands/types/content.py b/src/strands/types/content.py
@@ -8,7 +8,7 @@
 
 from typing import Literal
 
-from typing_extensions import TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 from .citations import CitationsContentBlock
 from .media import DocumentContent, ImageContent, VideoContent
@@ -129,10 +129,12 @@ class ContentBlockStartToolUse(TypedDict):
     Attributes:
         name: The name of the tool that the model is requesting to use.
         toolUseId: The ID for the tool request.
+        reasoningSignature: Token that ties the model's reasoning to this tool call.
     """
 
     name: str
     toolUseId: str
+    reasoningSignature: NotRequired[str]
 
 
 class ContentBlockStart(TypedDict, total=False):
diff --git a/src/strands/types/tools.py b/src/strands/types/tools.py
@@ -58,11 +58,13 @@ class ToolUse(TypedDict):
             Can be any JSON-serializable type.
         name: The name of the tool to invoke.
         toolUseId: A unique identifier for this specific tool use request.
+        reasoningSignature: Token that ties the model's reasoning to this tool call.
     """
 
     input: Any
     name: str
     toolUseId: str
+    reasoningSignature: NotRequired[str]
 
 
 class ToolResultContent(TypedDict, total=False):
diff --git a/tests/strands/event_loop/test_streaming.py b/tests/strands/event_loop/test_streaming.py
@@ -125,6 +125,10 @@ def test_handle_message_start():
             {"start": {"toolUse": {"toolUseId": "test", "name": "test"}}},
             {"toolUseId": "test", "name": "test", "input": ""},
         ),
+        (
+            {"start": {"toolUse": {"toolUseId": "test", "name": "test", "reasoningSignature": "YWJj"}}},
+            {"toolUseId": "test", "name": "test", "input": "", "reasoningSignature": "YWJj"},
+        ),
     ],
 )
 def test_handle_content_block_start(chunk: ContentBlockStartEvent, exp_tool_use):
@@ -310,6 +314,39 @@ def test_handle_content_block_delta(event: ContentBlockDeltaEvent, event_type, s
                 "redactedContent": b"",
             },
         ),
+        # Tool Use - With reasoningSignature
+        (
+            {
+                "content": [],
+                "current_tool_use": {
+                    "toolUseId": "123",
+                    "name": "test",
+                    "input": '{"key": "value"}',
+                    "reasoningSignature": "YWJj",
+                },
+                "text": "",
+                "reasoningText": "",
+                "citationsContent": [],
+                "redactedContent": b"",
+            },
+            {
+                "content": [
+                    {
+                        "toolUse": {
+                            "toolUseId": "123",
+                            "name": "test",
+                            "input": {"key": "value"},
+                            "reasoningSignature": "YWJj",
+                        }
+                    }
+                ],
+                "current_tool_use": {},
+                "text": "",
+                "reasoningText": "",
+                "citationsContent": [],
+                "redactedContent": b"",
+            },
+        ),
         # Tool Use - Missing input
         (
             {
diff --git a/tests/strands/models/test_gemini.py b/tests/strands/models/test_gemini.py
@@ -203,7 +203,7 @@ async def test_stream_request_with_reasoning(gemini_client, model, model_id):
                 {
                     "reasoningContent": {
                         "reasoningText": {
-                            "signature": "abc",
+                            "signature": "YWJj",  # base64 of "abc"
                             "text": "reasoning_text",
                         },
                     },
@@ -260,6 +260,51 @@ async def test_stream_request_with_tool_spec(gemini_client, model, model_id, too
 
 @pytest.mark.asyncio
 async def test_stream_request_with_tool_use(gemini_client, model, model_id):
+    """Test toolUse with reasoningSignature is sent as function_call with thought_signature."""
+    messages = [
+        {
+            "role": "assistant",
+            "content": [
+                {
+                    "toolUse": {
+                        "toolUseId": "c1",
+                        "name": "calculator",
+                        "input": {"expression": "2+2"},
+                        "reasoningSignature": "YWJj",  # base64 of "abc"
+                    },
+                },
+            ],
+        },
+    ]
+    await anext(model.stream(messages))
+
+    exp_request = {
+        "config": {
+            "tools": [{"function_declarations": []}],
+        },
+        "contents": [
+            {
+                "parts": [
+                    {
+                        "function_call": {
+                            "args": {"expression": "2+2"},
+                            "id": "c1",
+                            "name": "calculator",
+                        },
+                        "thought_signature": "YWJj",
+                    },
+                ],
+                "role": "model",
+            },
+        ],
+        "model": model_id,
+    }
+    gemini_client.aio.models.generate_content_stream.assert_called_with(**exp_request)
+
+
+@pytest.mark.asyncio
+async def test_stream_request_with_tool_use_no_reasoning_signature(gemini_client, model, model_id):
+    """Test toolUse without reasoningSignature is sent as function_call without thought_signature."""
     messages = [
         {
             "role": "assistant",
@@ -532,6 +577,55 @@ async def test_stream_response_tool_use(gemini_client, model, messages, agenerat
     assert tru_chunks == exp_chunks
 
 
+@pytest.mark.asyncio
+async def test_stream_response_tool_use_with_thought_signature(gemini_client, model, messages, agenerator, alist):
+    """Test that tool use responses with thought_signature include reasoningSignature."""
+    gemini_client.aio.models.generate_content_stream.return_value = agenerator(
+        [
+            genai.types.GenerateContentResponse(
+                candidates=[
+                    genai.types.Candidate(
+                        content=genai.types.Content(
+                            parts=[
+                                genai.types.Part(
+                                    function_call=genai.types.FunctionCall(
+                                        args={"expression": "2+2"},
+                                        id="c1",
+                                        name="calculator",
+                                    ),
+                                    thought_signature=b"abc",
+                                ),
+                            ],
+                        ),
+                        finish_reason="STOP",
+                    ),
+                ],
+                usage_metadata=genai.types.GenerateContentResponseUsageMetadata(
+                    prompt_token_count=1,
+                    total_token_count=3,
+                ),
+            ),
+        ]
+    )
+
+    tru_chunks = await alist(model.stream(messages))
+    exp_chunks = [
+        {"messageStart": {"role": "assistant"}},
+        {
+            "contentBlockStart": {
+                "start": {
+                    "toolUse": {"name": "calculator", "toolUseId": "c1", "reasoningSignature": "YWJj"},
+                },
+            },
+        },
+        {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"expression": "2+2"}'}}}},
+        {"contentBlockStop": {}},
+        {"messageStop": {"stopReason": "tool_use"}},
+        {"metadata": {"usage": {"inputTokens": 1, "outputTokens": 2, "totalTokens": 3}, "metrics": {"latencyMs": 0}}},
+    ]
+    assert tru_chunks == exp_chunks
+
+
 @pytest.mark.asyncio
 async def test_stream_response_reasoning(gemini_client, model, messages, agenerator, alist):
     gemini_client.aio.models.generate_content_stream.return_value = agenerator(
@@ -563,7 +657,7 @@ async def test_stream_response_reasoning(gemini_client, model, messages, agenera
     exp_chunks = [
         {"messageStart": {"role": "assistant"}},
         {"contentBlockStart": {"start": {}}},
-        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "abc", "text": "test reason"}}}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "YWJj", "text": "test reason"}}}},
         {"contentBlockStop": {}},
         {"messageStop": {"stopReason": "end_turn"}},
         {"metadata": {"usage": {"inputTokens": 1, "outputTokens": 2, "totalTokens": 3}, "metrics": {"latencyMs": 0}}},
@@ -622,7 +716,11 @@ async def test_stream_response_reasoning_and_text(gemini_client, model, messages
     exp_chunks = [
         {"messageStart": {"role": "assistant"}},
         {"contentBlockStart": {"start": {}}},
-        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig1", "text": "thinking about math"}}}},
+        {
+            "contentBlockDelta": {
+                "delta": {"reasoningContent": {"signature": "c2lnMQ==", "text": "thinking about math"}}
+            }
+        },
         {"contentBlockStop": {}},
         {"contentBlockStart": {"start": {}}},
         {"contentBlockDelta": {"delta": {"text": "2 + 2 = 4"}}},
