fix: provide unique toolUseId for gemini models (strands-agents#1201)

Co-authored-by: spicadust <spicastre@gmail.com>
Co-authored-by: Patrick Gray <pgrayy@amazon.com>

Author: 3 people

src/strands/models/gemini.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import mimetypes
+import secrets
 from collections.abc import AsyncGenerator
 from typing import Any, TypedDict, TypeVar, cast
 
@@ -86,6 +87,7 @@ def __init__(
 
         self._custom_client = client
         self.client_args = client_args or {}
+        self._tool_use_id_to_name: dict[str, str] = {}
 
         # Validate gemini_tools if provided
         if "gemini_tools" in self.config:
@@ -173,10 +175,13 @@ def _format_request_content_part(self, content: ContentBlock) -> genai.types.Par
             return genai.types.Part(text=content["text"])
 
         if "toolResult" in content:
+            tool_use_id = content["toolResult"]["toolUseId"]
+            function_name = self._tool_use_id_to_name.get(tool_use_id, tool_use_id)
+
             return genai.types.Part(
                 function_response=genai.types.FunctionResponse(
-                    id=content["toolResult"]["toolUseId"],
-                    name=content["toolResult"]["toolUseId"],
+                    id=tool_use_id,
+                    name=function_name,
                     response={
                         "output": [
                             tool_result_content
@@ -191,6 +196,12 @@ def _format_request_content_part(self, content: ContentBlock) -> genai.types.Par
             )
 
         if "toolUse" in content:
+            # Store the mapping from toolUseId to name for later use in toolResult formatting.
+            # This mapping is built as we format the request, ensuring that when we encounter
+            # toolResult blocks (which come after toolUse blocks in the message history),
+            # we can look up the function name.
+            self._tool_use_id_to_name[content["toolUse"]["toolUseId"]] = content["toolUse"]["name"]
+
             return genai.types.Part(
                 function_call=genai.types.FunctionCall(
                     args=content["toolUse"]["input"],
@@ -317,16 +328,16 @@ def _format_chunk(self, event: dict[str, Any]) -> StreamEvent:
             case "content_start":
                 match event["data_type"]:
                     case "tool":
-                        # Note: toolUseId is the only identifier available in a tool result. However, Gemini requires
-                        #       that name be set in the equivalent FunctionResponse type. Consequently, we assign
-                        #       function name to toolUseId in our tool use block. And another reason, function_call is
-                        #       not guaranteed to have id populated.
+                        function_call = event["data"].function_call
+                        # Use Gemini's provided ID or generate one if missing
+                        tool_use_id = function_call.id or f"tooluse_{secrets.token_urlsafe(16)}"
+
                         return {
                             "contentBlockStart": {
                                 "start": {
                                     "toolUse": {
-                                        "name": event["data"].function_call.name,
-                                        "toolUseId": event["data"].function_call.name,
+                                        "name": function_call.name,
+                                        "toolUseId": tool_use_id,
                                     },
                                 },
                             },
@@ -417,6 +428,7 @@ async def stream(
             ModelThrottledException: If the request is throttled by Gemini.
         """
         request = self._format_request(messages, tool_specs, system_prompt, self.config.get("params"))
+        self._tool_use_id_to_name.clear()
 
         client = self._get_client().aio
 

tests/strands/models/test_gemini.py

@@ -360,6 +360,71 @@ async def test_stream_request_with_tool_results(gemini_client, 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_results_preserving_name(gemini_client, model, model_id):
+    messages = [
+        {
+            "role": "assistant",
+            "content": [
+                {
+                    "toolUse": {
+                        "toolUseId": "t1",
+                        "name": "tool_1",
+                        "input": {},
+                    },
+                },
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "toolResult": {
+                        "toolUseId": "t1",
+                        "status": "success",
+                        "content": [{"text": "done"}],
+                    },
+                },
+            ],
+        },
+    ]
+    await anext(model.stream(messages))
+
+    exp_request = {
+        "config": {
+            "tools": [{"function_declarations": []}],
+        },
+        "contents": [
+            {
+                "parts": [
+                    {
+                        "function_call": {
+                            "args": {},
+                            "id": "t1",
+                            "name": "tool_1",
+                        },
+                    },
+                ],
+                "role": "model",
+            },
+            {
+                "parts": [
+                    {
+                        "function_response": {
+                            "id": "t1",
+                            "name": "tool_1",
+                            "response": {"output": [{"text": "done"}]},
+                        },
+                    },
+                ],
+                "role": "user",
+            },
+        ],
+        "model": model_id,
+    }
+    gemini_client.aio.models.generate_content_stream.assert_called_with(**exp_request)
+
+
 @pytest.mark.asyncio
 async def test_stream_request_with_empty_content(gemini_client, model, model_id):
     messages = [
@@ -459,7 +524,7 @@ async def test_stream_response_tool_use(gemini_client, model, messages, agenerat
     exp_chunks = [
         {"messageStart": {"role": "assistant"}},
         {"contentBlockStart": {"start": {}}},
-        {"contentBlockStart": {"start": {"toolUse": {"name": "calculator", "toolUseId": "calculator"}}}},
+        {"contentBlockStart": {"start": {"toolUse": {"name": "calculator", "toolUseId": "c1"}}}},
         {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"expression": "2+2"}'}}}},
         {"contentBlockStop": {}},
         {"contentBlockStop": {}},