Python: Improve AG-UI tool call argument protocol conformance

python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py

@@ -31,6 +31,16 @@ def __init__(self) -> None:
         self.thread_id: str | None = None
         self.run_id: str | None = None
 
+    @staticmethod
+    def _get_tool_call_id(event: dict[str, Any]) -> str | None:
+        """Return the tool call ID from either AG-UI field spelling."""
+        tool_call_id = event.get("toolCallId")
+        if tool_call_id is None:
+            tool_call_id = event.get("tool_call_id")
+        if tool_call_id is None:
+            return None
+        return str(tool_call_id)
+
     def convert_event(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
         """Convert a single AG-UI event to ChatResponseUpdate.
 
@@ -129,7 +139,7 @@ def _handle_text_message_end(self, event: dict[str, Any]) -> ChatResponseUpdate
 
     def _handle_tool_call_start(self, event: dict[str, Any]) -> ChatResponseUpdate:
         """Handle TOOL_CALL_START event."""
-        self.current_tool_call_id = event.get("toolCallId")
+        self.current_tool_call_id = self._get_tool_call_id(event)
         self.current_tool_name = event.get("toolName") or event.get("toolCallName") or event.get("tool_call_name")
         self.accumulated_tool_args = ""
 
@@ -144,8 +154,20 @@ def _handle_tool_call_start(self, event: dict[str, Any]) -> ChatResponseUpdate:
             ],
         )
 
-    def _handle_tool_call_args(self, event: dict[str, Any]) -> ChatResponseUpdate:
+    def _handle_tool_call_args(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
         """Handle TOOL_CALL_ARGS event."""
+        event_tool_call_id = self._get_tool_call_id(event)
+        if event_tool_call_id is not None:
+            if self.current_tool_call_id and event_tool_call_id != self.current_tool_call_id:
+                logger.warning(
+                    "Ignoring TOOL_CALL_ARGS for toolCallId=%s while current toolCallId=%s",
+                    event_tool_call_id,
+                    self.current_tool_call_id,
+                )
+                return None
+            if not self.current_tool_call_id:
+                self.current_tool_call_id = event_tool_call_id
+
         delta = event.get("delta", "")
         self.accumulated_tool_args += delta
 
@@ -162,7 +184,15 @@ def _handle_tool_call_args(self, event: dict[str, Any]) -> ChatResponseUpdate:
 
     def _handle_tool_call_end(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
         """Handle TOOL_CALL_END event."""
-        self.accumulated_tool_args = ""
+        event_tool_call_id = self._get_tool_call_id(event)
+        if (
+            self.current_tool_call_id is None
+            or event_tool_call_id is None
+            or event_tool_call_id == self.current_tool_call_id
+        ):
+            self.current_tool_call_id = None
+            self.current_tool_name = None
+            self.accumulated_tool_args = ""
         return None
 
     def _handle_tool_call_result(self, event: dict[str, Any]) -> ChatResponseUpdate:

python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py

@@ -441,6 +441,54 @@ async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str
                     break
         assert found, "Expected to find function_call content for my_tool"
 
+    async def test_tool_call_args_id_mismatch_does_not_execute_current_client_tool(
+        self, monkeypatch: MonkeyPatch
+    ) -> None:
+        """Mismatched TOOL_CALL_ARGS must not be rebound to the latest client tool."""
+        executed: list[int] = []
+
+        @tool
+        def danger_tool(amount: int) -> str:
+            """Record an invocation for the regression assertion."""
+            executed.append(amount)
+            return f"danger={amount}"
+
+        call_count = 0
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                mock_events = [
+                    {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+                    {"type": "TOOL_CALL_START", "toolCallId": "safe", "toolName": "safe_tool"},
+                    {"type": "TOOL_CALL_START", "toolCallId": "danger", "toolName": "danger_tool"},
+                    {"type": "TOOL_CALL_ARGS", "toolCallId": "safe", "delta": '{"amount": 100}'},
+                    {"type": "TOOL_CALL_END", "toolCallId": "safe"},
+                    {"type": "TOOL_CALL_END", "toolCallId": "danger"},
+                    {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+                ]
+            else:
+                mock_events = [
+                    {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_2"},
+                    {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "done"},
+                    {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_2"},
+                ]
+
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        response = await client.get_response(
+            [Message(role="user", contents=["Test"])],
+            options={"tools": [danger_tool]},
+        )
+
+        assert response.text == "done"
+        assert executed == []
+
     async def test_interrupt_options_transmission(self, monkeypatch: MonkeyPatch) -> None:
         """Interrupt option fields are forwarded to the HTTP service."""
         available_interrupts = [{"id": "req_1", "type": "request_info"}]

python/packages/ag-ui/tests/ag_ui/test_event_converters.py

@@ -6,6 +6,7 @@
 from typing import Any, cast
 
 import pytest
+from agent_framework import ChatResponse
 
 from agent_framework_ag_ui._event_converters import AGUIEventConverter
 
@@ -422,3 +423,47 @@ def test_multiple_tool_calls(self) -> None:
         assert len(non_none_updates) == 4
         assert non_none_updates[0].contents[0].name == "search"
         assert non_none_updates[2].contents[0].name == "fetch"
+
+    def test_tool_call_args_must_match_current_tool_call_id(self) -> None:
+        """TOOL_CALL_ARGS for another call must not be rebound to the current tool."""
+        converter = AGUIEventConverter()
+
+        events = [
+            {"type": "TOOL_CALL_START", "toolCallId": "safe", "toolName": "safe_tool"},
+            {"type": "TOOL_CALL_START", "toolCallId": "danger", "toolName": "danger_tool"},
+            {"type": "TOOL_CALL_ARGS", "toolCallId": "safe", "delta": '{"amount": 100}'},
+            {"type": "TOOL_CALL_END", "toolCallId": "safe"},
+            {"type": "TOOL_CALL_END", "toolCallId": "danger"},
+        ]
+
+        updates = [update for event in events if (update := converter.convert_event(event)) is not None]
+        response = ChatResponse.from_updates(updates)
+        function_calls = [
+            (content.call_id, content.name, content.arguments)
+            for message in response.messages
+            for content in message.contents
+            if content.type == "function_call"
+        ]
+
+        assert ("danger", "danger_tool", "") in function_calls
+        assert ("danger", "danger_tool", '{"amount": 100}') not in function_calls
+        assert all(call[2] != '{"amount": 100}' for call in function_calls)
+
+    def test_tool_call_end_must_match_current_tool_call_id(self) -> None:
+        """TOOL_CALL_END for another call must not clear the current call state."""
+        converter = AGUIEventConverter()
+
+        converter.convert_event({"type": "TOOL_CALL_START", "toolCallId": "danger", "toolName": "danger_tool"})
+        converter.convert_event({"type": "TOOL_CALL_ARGS", "toolCallId": "danger", "delta": '{"amount":'})
+        update = converter.convert_event({"type": "TOOL_CALL_END", "toolCallId": "safe"})
+
+        assert update is None
+        assert converter.current_tool_call_id == "danger"
+        assert converter.current_tool_name == "danger_tool"
+        assert converter.accumulated_tool_args == '{"amount":'
+
+        converter.convert_event({"type": "TOOL_CALL_END", "toolCallId": "danger"})
+
+        assert converter.current_tool_call_id is None
+        assert converter.current_tool_name is None
+        assert converter.accumulated_tool_args == ""