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/agent_framework_ag_ui/_orchestration/_tooling.py

@@ -3,59 +3,85 @@
 """Tool handling helpers."""
 
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from agent_framework import BaseChatClient, ChatAgent
+from agent_framework import BaseChatClient
+
+if TYPE_CHECKING:
+    from agent_framework import AgentProtocol
 
 logger = logging.getLogger(__name__)
 
 
-def collect_server_tools(agent: Any) -> list[Any]:
-    """Collect server tools from ChatAgent or duck-typed agent."""
-    if isinstance(agent, ChatAgent):
-        tools_from_agent = agent.default_options.get("tools")
-        server_tools = list(tools_from_agent) if tools_from_agent else []
-        logger.info(f"[TOOLS] Agent has {len(server_tools)} configured tools")
-        for tool in server_tools:
-            tool_name = getattr(tool, "name", "unknown")
-            approval_mode = getattr(tool, "approval_mode", None)
-            logger.info(f"[TOOLS]   - {tool_name}: approval_mode={approval_mode}")
-        return server_tools
+def _collect_mcp_tool_functions(mcp_tools: list[Any]) -> list[Any]:
+    """Extract functions from connected MCP tools.
 
-    try:
-        default_options_attr = getattr(agent, "default_options", None)
-        if default_options_attr is not None:
-            if isinstance(default_options_attr, dict):
-                return default_options_attr.get("tools") or []
-            return getattr(default_options_attr, "tools", None) or []
-    except AttributeError:
+    Args:
+        mcp_tools: List of MCP tool instances.
+
+    Returns:
+        List of functions from connected MCP tools.
+    """
+    functions: list[Any] = []
+    for mcp_tool in mcp_tools:
+        if getattr(mcp_tool, "is_connected", False) and hasattr(mcp_tool, "functions"):
+            functions.extend(mcp_tool.functions)
+    return functions
+
+
+def collect_server_tools(agent: "AgentProtocol") -> list[Any]:
+    """Collect server tools from an agent.
+
+    This includes both regular tools from default_options and MCP tools.
+    MCP tools are stored separately for lifecycle management but their
+    functions need to be included for tool execution during approval flows.
+
+    Args:
+        agent: Agent instance to collect tools from. Works with ChatAgent
+            or any agent with default_options and optional mcp_tools attributes.
+
+    Returns:
+        List of tools including both regular tools and connected MCP tool functions.
+    """
+    # Get tools from default_options
+    default_options = getattr(agent, "default_options", None)
+    if default_options is None:
         return []
-    return []
+
+    tools_from_agent = default_options.get("tools") if isinstance(default_options, dict) else None
+    server_tools = list(tools_from_agent) if tools_from_agent else []
+
+    # Include functions from connected MCP tools (only available on ChatAgent)
+    mcp_tools = getattr(agent, "mcp_tools", None)
+    if mcp_tools:
+        server_tools.extend(_collect_mcp_tool_functions(mcp_tools))
+
+    logger.info(f"[TOOLS] Agent has {len(server_tools)} configured tools")
+    for tool in server_tools:
+        tool_name = getattr(tool, "name", "unknown")
+        approval_mode = getattr(tool, "approval_mode", None)
+        logger.info(f"[TOOLS]   - {tool_name}: approval_mode={approval_mode}")
+    return server_tools
 
 
-def register_additional_client_tools(agent: Any, client_tools: list[Any] | None) -> None:
-    """Register client tools as additional declaration-only tools to avoid server execution."""
+def register_additional_client_tools(agent: "AgentProtocol", client_tools: list[Any] | None) -> None:
+    """Register client tools as additional declaration-only tools to avoid server execution.
+
+    Args:
+        agent: Agent instance to register tools on. Works with ChatAgent
+            or any agent with a chat_client attribute.
+        client_tools: List of client tools to register.
+    """
     if not client_tools:
         return
 
-    if isinstance(agent, ChatAgent):
-        chat_client = agent.chat_client
-        if isinstance(chat_client, BaseChatClient) and chat_client.function_invocation_configuration is not None:
-            chat_client.function_invocation_configuration.additional_tools = client_tools
-            logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)")
+    chat_client = getattr(agent, "chat_client", None)
+    if chat_client is None:
         return
 
-    try:
-        chat_client_attr = getattr(agent, "chat_client", None)
-        if chat_client_attr is not None:
-            fic = getattr(chat_client_attr, "function_invocation_configuration", None)
-            if fic is not None:
-                fic.additional_tools = client_tools  # type: ignore[attr-defined]
-                logger.debug(
-                    f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)"
-                )
-    except AttributeError:
-        return
+    if isinstance(chat_client, BaseChatClient) and chat_client.function_invocation_configuration is not None:
+        chat_client.function_invocation_configuration.additional_tools = client_tools
+        logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)")
 
 
 def merge_tools(server_tools: list[Any], client_tools: list[Any] | None) -> list[Any] | None: