Fix tool calls on gateway

src/ai/models/ai_gateway/stream.py

@@ -158,9 +158,16 @@ async def _build_request_body(
 # ---------------------------------------------------------------------------
 
 
-def _expand_tool_call(data: dict[str, Any]) -> list[events_.Event]:
-    """Expand a complete ``tool-call`` part into Start + Delta + End."""
+def _expand_tool_call(
+    data: dict[str, Any], streamed_tool_ids: set[str]
+) -> list[events_.Event]:
+    """Expand a complete ``tool-call`` part into Start + Delta + End.
+
+    Returns empty when the tool was already streamed via ``tool-input-*``.
+    """
     tc_id = data.get("toolCallId", "")
+    if tc_id in streamed_tool_ids:
+        return []
     tool_name = data.get("toolName", "")
     tool_input = data.get("input", "")
     args_str = tool_input if isinstance(tool_input, str) else json.dumps(tool_input)
@@ -198,7 +205,9 @@ def _parse_usage(data: Any) -> usage_.Usage:
     )
 
 
-def _parse_stream_part(data: dict[str, Any]) -> list[events_.Event]:
+def _parse_stream_part(
+    data: dict[str, Any], streamed_tool_ids: set[str]
+) -> list[events_.Event]:
     """Convert a ``LanguageModelV3StreamPart`` to public events."""
     match data.get("type", ""):
         case "text-start":
@@ -230,9 +239,11 @@ def _parse_stream_part(data: dict[str, Any]) -> list[events_.Event]:
             return [events_.ReasoningEnd(block_id=data.get("id", "reasoning"))]
 
         case "tool-input-start":
+            tcid = data.get("id", "")
+            streamed_tool_ids.add(tcid)
             return [
                 events_.ToolStart(
-                    tool_call_id=data.get("id", ""),
+                    tool_call_id=tcid,
                     tool_name=data.get("toolName", ""),
                 )
             ]
@@ -249,7 +260,7 @@ def _parse_stream_part(data: dict[str, Any]) -> list[events_.Event]:
             return [events_.ToolEnd(tool_call_id=data.get("id", ""))]
 
         case "tool-call":
-            return _expand_tool_call(data)
+            return _expand_tool_call(data, streamed_tool_ids)
 
         case "file":
             return [
@@ -313,8 +324,9 @@ async def stream(
                 )
 
             yield events_.StreamStart()
+            streamed_tool_ids: set[str] = set()
             async for data in _common.parse_sse_lines(response):
-                for event in _parse_stream_part(data):
+                for event in _parse_stream_part(data, streamed_tool_ids):
                     yield event
     except errors.GatewayError:
         raise

tests/models/ai_gateway/test_protocol.py

@@ -235,7 +235,7 @@ class TestParseStreamPartComplex:
     def test_text_delta_uses_textDelta_key(self) -> None:
         """The gateway sends ``textDelta`` (camelCase), not ``delta``."""
         events = stream_mod._parse_stream_part(
-            {"type": "text-delta", "id": "t1", "textDelta": "Hello"}
+            {"type": "text-delta", "id": "t1", "textDelta": "Hello"}, set()
         )
         assert isinstance(events[0], events_.TextDelta)
         assert events[0].chunk == "Hello"
@@ -249,7 +249,8 @@ def test_tool_call_expands_to_three_events(self) -> None:
                 "toolCallId": "tc-1",
                 "toolName": "get_weather",
                 "input": {"city": "SF"},
-            }
+            },
+            set(),
         )
         assert len(events) == 3
         assert isinstance(events[0], events_.ToolStart)
@@ -258,6 +259,24 @@ def test_tool_call_expands_to_three_events(self) -> None:
         assert json.loads(events[1].chunk) == {"city": "SF"}
         assert isinstance(events[2], events_.ToolEnd)
 
+    def test_tool_call_skipped_when_already_streamed(self) -> None:
+        """A ``tool-call`` that duplicates a streamed tool is dropped."""
+        seen: set[str] = set()
+        stream_mod._parse_stream_part(
+            {"type": "tool-input-start", "id": "tc-1", "toolName": "get_weather"},
+            seen,
+        )
+        events = stream_mod._parse_stream_part(
+            {
+                "type": "tool-call",
+                "toolCallId": "tc-1",
+                "toolName": "get_weather",
+                "input": {"city": "SF"},
+            },
+            seen,
+        )
+        assert events == []
+
     def test_finish_flat_usage(self) -> None:
         events = stream_mod._parse_stream_part(
             {
@@ -267,7 +286,8 @@ def test_finish_flat_usage(self) -> None:
                     "prompt_tokens": 10,
                     "completion_tokens": 20,
                 },
-            }
+            },
+            set(),
         )
         done = events[0]
         assert isinstance(done, events_.StreamEnd)
@@ -293,7 +313,8 @@ def test_finish_v3_nested_usage(self) -> None:
                         "reasoning": 30,
                     },
                 },
-            }
+            },
+            set(),
         )
         done = events[0]
         assert isinstance(done, events_.StreamEnd)
@@ -311,7 +332,8 @@ def test_file_part(self) -> None:
                 "id": "f1",
                 "mediaType": "image/png",
                 "data": "iVBORw0KGgo=",
-            }
+            },
+            set(),
         )
         assert len(events) == 1
         assert isinstance(events[0], events_.FileEvent)
@@ -321,14 +343,16 @@ def test_file_part(self) -> None:
 
     def test_file_part_defaults(self) -> None:
         """A minimal ``file`` part uses sensible defaults."""
-        events = stream_mod._parse_stream_part({"type": "file", "data": "somedata"})
+        events = stream_mod._parse_stream_part(
+            {"type": "file", "data": "somedata"}, set()
+        )
         assert len(events) == 1
         assert isinstance(events[0], events_.FileEvent)
         assert events[0].media_type == "application/octet-stream"
 
     def test_unknown_types_produce_no_events(self) -> None:
         for t in ("stream-start", "raw", "response-metadata", "banana"):
-            assert stream_mod._parse_stream_part({"type": t}) == []
+            assert stream_mod._parse_stream_part({"type": t}, set()) == []
 
 
 # ---------------------------------------------------------------------------