Python: fix(ag-ui): add MCP tool support for AG-UI approval flows (#3212)

* add MCP tool support for AG-UI approval flows

* use attribute in place of property

python/packages/ag-ui/tests/test_orchestrators.py

@@ -3,10 +3,17 @@
 """Tests for AG-UI orchestrators."""
 
 from collections.abc import AsyncGenerator
-from types import SimpleNamespace
 from typing import Any
+from unittest.mock import MagicMock
 
-from agent_framework import AgentResponseUpdate, FunctionInvocationConfiguration, TextContent, ai_function
+from agent_framework import (
+    AgentResponseUpdate,
+    BaseChatClient,
+    ChatAgent,
+    FunctionInvocationConfiguration,
+    TextContent,
+    ai_function,
+)
 
 from agent_framework_ag_ui._agent import AgentConfig
 from agent_framework_ag_ui._orchestrators import DefaultOrchestrator, ExecutionContext
@@ -18,56 +25,53 @@ def server_tool() -> str:
     return "server"
 
 
-class DummyAgent:
-    """Minimal agent stub to capture run_stream parameters."""
+def _create_mock_chat_agent(
+    tools: list[Any] | None = None,
+    response_format: Any = None,
+    capture_tools: list[Any] | None = None,
+    capture_messages: list[Any] | None = None,
+) -> ChatAgent:
+    """Create a ChatAgent with mocked chat client for testing.
 
-    def __init__(self) -> None:
-        self.default_options: dict[str, Any] = {"tools": [server_tool], "response_format": None}
-        self.tools = [server_tool]
-        self.chat_client = SimpleNamespace(
-            function_invocation_configuration=FunctionInvocationConfiguration(),
-        )
-        self.seen_tools: list[Any] | None = None
+    Args:
+        tools: Tools to configure on the agent.
+        response_format: Response format to configure.
+        capture_tools: If provided, tools passed to run_stream will be appended here.
+        capture_messages: If provided, messages passed to run_stream will be appended here.
+    """
+    mock_chat_client = MagicMock(spec=BaseChatClient)
+    mock_chat_client.function_invocation_configuration = FunctionInvocationConfiguration()
 
-    async def run_stream(
-        self,
+    agent = ChatAgent(
+        chat_client=mock_chat_client,
+        tools=tools or [server_tool],
+        response_format=response_format,
+    )
+
+    # Create a mock run_stream that captures parameters and yields a simple response
+    async def mock_run_stream(
         messages: list[Any],
         *,
-        thread: Any,
+        thread: Any = None,
         tools: list[Any] | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[AgentResponseUpdate, None]:
-        self.seen_tools = tools
+        if capture_tools is not None and tools is not None:
+            capture_tools.extend(tools)
+        if capture_messages is not None:
+            capture_messages.extend(messages)
         yield AgentResponseUpdate(contents=[TextContent(text="ok")], role="assistant")
 
+    # Patch the run_stream method
+    agent.run_stream = mock_run_stream  # type: ignore[method-assign]
 
-class RecordingAgent:
-    """Agent stub that captures messages passed to run_stream."""
-
-    def __init__(self) -> None:
-        self.chat_options = SimpleNamespace(tools=[], response_format=None)
-        self.tools: list[Any] = []
-        self.chat_client = SimpleNamespace(
-            function_invocation_configuration=FunctionInvocationConfiguration(),
-        )
-        self.seen_messages: list[Any] | None = None
-
-    async def run_stream(
-        self,
-        messages: list[Any],
-        *,
-        thread: Any,
-        tools: list[Any] | None = None,
-        **kwargs: Any,
-    ) -> AsyncGenerator[AgentResponseUpdate, None]:
-        self.seen_messages = messages
-        yield AgentResponseUpdate(contents=[TextContent(text="ok")], role="assistant")
+    return agent
 
 
 async def test_default_orchestrator_merges_client_tools() -> None:
     """Client tool declarations are merged with server tools before running agent."""
-
-    agent = DummyAgent()
+    captured_tools: list[Any] = []
+    agent = _create_mock_chat_agent(tools=[server_tool], capture_tools=captured_tools)
     orchestrator = DefaultOrchestrator()
 
     input_data = {
@@ -100,8 +104,8 @@ async def test_default_orchestrator_merges_client_tools() -> None:
     async for event in orchestrator.run(context):
         events.append(event)
 
-    assert agent.seen_tools is not None
-    tool_names = [getattr(tool, "name", "?") for tool in agent.seen_tools]
+    assert len(captured_tools) > 0
+    tool_names = [getattr(tool, "name", "?") for tool in captured_tools]
     assert "server_tool" in tool_names
     assert "get_weather" in tool_names
     assert agent.chat_client.function_invocation_configuration.additional_tools
@@ -109,8 +113,7 @@ async def test_default_orchestrator_merges_client_tools() -> None:
 
 async def test_default_orchestrator_with_camel_case_ids() -> None:
     """Client tool is able to extract camelCase IDs."""
-
-    agent = DummyAgent()
+    agent = _create_mock_chat_agent()
     orchestrator = DefaultOrchestrator()
 
     input_data = {
@@ -143,8 +146,7 @@ async def test_default_orchestrator_with_camel_case_ids() -> None:
 
 async def test_default_orchestrator_with_snake_case_ids() -> None:
     """Client tool is able to extract snake_case IDs."""
-
-    agent = DummyAgent()
+    agent = _create_mock_chat_agent()
     orchestrator = DefaultOrchestrator()
 
     input_data = {
@@ -177,8 +179,8 @@ async def test_default_orchestrator_with_snake_case_ids() -> None:
 
 async def test_state_context_injected_when_tool_call_state_mismatch() -> None:
     """State context should be injected when current state differs from tool call args."""
-
-    agent = RecordingAgent()
+    captured_messages: list[Any] = []
+    agent = _create_mock_chat_agent(tools=[], capture_messages=captured_messages)
     orchestrator = DefaultOrchestrator()
 
     tool_recipe = {"title": "Salad", "special_preferences": []}
@@ -215,9 +217,9 @@ async def test_state_context_injected_when_tool_call_state_mismatch() -> None:
     async for _event in orchestrator.run(context):
         pass
 
-    assert agent.seen_messages is not None
+    assert len(captured_messages) > 0
     state_messages = []
-    for msg in agent.seen_messages:
+    for msg in captured_messages:
         role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
         if role_value != "system":
             continue
@@ -230,8 +232,8 @@ async def test_state_context_injected_when_tool_call_state_mismatch() -> None:
 
 async def test_state_context_not_injected_when_tool_call_matches_state() -> None:
     """State context should be skipped when tool call args match current state."""
-
-    agent = RecordingAgent()
+    captured_messages: list[Any] = []
+    agent = _create_mock_chat_agent(tools=[], capture_messages=captured_messages)
     orchestrator = DefaultOrchestrator()
 
     input_data = {
@@ -264,9 +266,9 @@ async def test_state_context_not_injected_when_tool_call_matches_state() -> None
     async for _event in orchestrator.run(context):
         pass
 
-    assert agent.seen_messages is not None
+    assert len(captured_messages) > 0
     state_messages = []
-    for msg in agent.seen_messages:
+    for msg in captured_messages:
         role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
         if role_value != "system":
             continue

python/packages/ag-ui/tests/test_tooling.py

@@ -1,8 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from types import SimpleNamespace
+from unittest.mock import MagicMock
 
-from agent_framework_ag_ui._orchestration._tooling import merge_tools, register_additional_client_tools
+from agent_framework import ChatAgent, ai_function
+
+from agent_framework_ag_ui._orchestration._tooling import (
+    collect_server_tools,
+    merge_tools,
+    register_additional_client_tools,
+)
 
 
 class DummyTool:
@@ -11,6 +17,30 @@ class DummyTool:
         self.declaration_only = True
 
 
+class MockMCPTool:
+    """Mock MCP tool that simulates connected MCP tool with functions."""
+
+    def __init__(self, functions: list[DummyTool], is_connected: bool = True) -> None:
+        self.functions = functions
+        self.is_connected = is_connected
+
+
+@ai_function
+def regular_tool() -> str:
+    """Regular tool for testing."""
+    return "result"
+
+
+def _create_chat_agent_with_tool(tool_name: str = "regular_tool") -> ChatAgent:
+    """Create a ChatAgent with a mocked chat client and a simple tool.
+
+    Note: tool_name parameter is kept for API compatibility but the tool
+    will always be named 'regular_tool' since ai_function uses the function name.
+    """
+    mock_chat_client = MagicMock()
+    return ChatAgent(chat_client=mock_chat_client, tools=[regular_tool])
+
+
 def test_merge_tools_filters_duplicates() -> None:
     server = [DummyTool("a"), DummyTool("b")]
     client = [DummyTool("b"), DummyTool("c")]
@@ -23,14 +53,79 @@ def test_merge_tools_filters_duplicates() -> None:
 
 
 def test_register_additional_client_tools_assigns_when_configured() -> None:
-    class Fic:
-        def __init__(self) -> None:
-            self.additional_tools = None
+    """register_additional_client_tools should set additional_tools on the chat client."""
+    from agent_framework import BaseChatClient, FunctionInvocationConfiguration
 
-    holder = SimpleNamespace(function_invocation_configuration=Fic())
-    agent = SimpleNamespace(chat_client=holder)
+    mock_chat_client = MagicMock(spec=BaseChatClient)
+    mock_chat_client.function_invocation_configuration = FunctionInvocationConfiguration()
+
+    agent = ChatAgent(chat_client=mock_chat_client)
 
     tools = [DummyTool("x")]
     register_additional_client_tools(agent, tools)
 
-    assert holder.function_invocation_configuration.additional_tools == tools
+    assert mock_chat_client.function_invocation_configuration.additional_tools == tools
+
+
+def test_collect_server_tools_includes_mcp_tools_when_connected() -> None:
+    """MCP tool functions should be included when the MCP tool is connected."""
+    mcp_function1 = DummyTool("mcp_function_1")
+    mcp_function2 = DummyTool("mcp_function_2")
+    mock_mcp = MockMCPTool([mcp_function1, mcp_function2], is_connected=True)
+
+    agent = _create_chat_agent_with_tool("regular_tool")
+    agent.mcp_tools = [mock_mcp]
+
+    tools = collect_server_tools(agent)
+
+    names = [getattr(t, "name", None) for t in tools]
+    assert "regular_tool" in names
+    assert "mcp_function_1" in names
+    assert "mcp_function_2" in names
+    assert len(tools) == 3
+
+
+def test_collect_server_tools_excludes_mcp_tools_when_not_connected() -> None:
+    """MCP tool functions should be excluded when the MCP tool is not connected."""
+    mcp_function = DummyTool("mcp_function")
+    mock_mcp = MockMCPTool([mcp_function], is_connected=False)
+
+    agent = _create_chat_agent_with_tool("regular_tool")
+    agent.mcp_tools = [mock_mcp]
+
+    tools = collect_server_tools(agent)
+
+    names = [getattr(t, "name", None) for t in tools]
+    assert "regular_tool" in names
+    assert "mcp_function" not in names
+    assert len(tools) == 1
+
+
+def test_collect_server_tools_works_with_no_mcp_tools() -> None:
+    """collect_server_tools should work when there are no MCP tools."""
+    agent = _create_chat_agent_with_tool("regular_tool")
+
+    tools = collect_server_tools(agent)
+
+    names = [getattr(t, "name", None) for t in tools]
+    assert "regular_tool" in names
+    assert len(tools) == 1
+
+
+def test_collect_server_tools_with_mcp_tools_via_public_property() -> None:
+    """collect_server_tools should access MCP tools via the public mcp_tools property."""
+    mcp_function = DummyTool("mcp_function")
+    mock_mcp = MockMCPTool([mcp_function], is_connected=True)
+
+    agent = _create_chat_agent_with_tool("regular_tool")
+    agent.mcp_tools = [mock_mcp]
+
+    # Verify the public property works
+    assert agent.mcp_tools == [mock_mcp]
+
+    tools = collect_server_tools(agent)
+
+    names = [getattr(t, "name", None) for t in tools]
+    assert "regular_tool" in names
+    assert "mcp_function" in names
+    assert len(tools) == 2