From 96a2522541f1dfdae2f73f8613eb754a7e3e4df3 Mon Sep 17 00:00:00 2001 From: White-Mouse <15983334+White-Mouse@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:44:15 +0800 Subject: [PATCH] Python: bind AG-UI tool arguments to call IDs --- .../_event_converters.py | 36 ++++++++++++-- .../ag-ui/tests/ag_ui/test_ag_ui_client.py | 48 +++++++++++++++++++ .../tests/ag_ui/test_event_converters.py | 45 +++++++++++++++++ 3 files changed, 126 insertions(+), 3 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py b/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py index 3510661a355..122e7d6ffbe 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py +++ b/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: diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index 0d385779497..7fd71f6ff2d 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/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"}] diff --git a/python/packages/ag-ui/tests/ag_ui/test_event_converters.py b/python/packages/ag-ui/tests/ag_ui/test_event_converters.py index f4bc3ddc4f1..fae715d97ef 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_event_converters.py +++ b/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 == ""