feat: Add web ui for multi_mcp ADK agent

Author: andrewginns

Makefile

@@ -9,4 +9,7 @@ leaderboard:
 	uv run -- streamlit run agents_mcp_usage/multi_mcp/eval_multi_mcp/merbench_ui.py
 
 adk_basic_ui:
-	cd agents_mcp_usage/basic_mcp && uv run adk web
+	uv run adk web agents_mcp_usage/basic_mcp
+
+adk_multi_ui:
+	uv run adk web agents_mcp_usage/multi_mcp

agents_mcp_usage/multi_mcp/README.md

@@ -21,11 +21,14 @@ Agents utilising multiple MCP servers can be dramatically more complex than an A
    # Run the Google ADK multi-MCP example
    uv run agents_mcp_usage/multi_mcp/multi_mcp_use/adk_mcp.py
 
+   ## Launch ADK web UI for visual interaction
+   make adk_multi_ui
+   
    # Run the multi-MCP evaluation
    uv run agents_mcp_usage/multi_mcp/eval_multi_mcp/evals_pydantic_mcp.py
 
    # Run multi-model benchmarking
-   uv run agents_mcp_usage/multi_mcp/eval_multi_mcp/run_multi_evals.py --models "gemini-2.5-pro,gemini-2.0-flash" --runs 5 --parallel
+   uv run agents_mcp_usage/multi_mcp/eval_multi_mcp/run_multi_evals.py --models "gemini-2.5-pro-preview-06-05,gemini-2.0-flash" --runs 5 --parallel
 
    # Launch the evaluation dashboard
    uv run streamlit run agents_mcp_usage/multi_mcp/eval_multi_mcp/merbench_ui.py
@@ -250,18 +253,18 @@ The `run_multi_evals.py` script enables systematic comparison across multiple LL
 ```bash
 # Parallel benchmarking across multiple models
 uv run agents_mcp_usage/multi_mcp/eval_multi_mcp/run_multi_evals.py \
-  --models "gemini-2.5-pro,gemini-2.0-flash,gemini-2.5-flash-preview-04-17" \
+  --models "gemini-2.5-pro-preview-06-05,gemini-2.0-flash,gemini-2.5-flash" \
   --runs 5 \
   --parallel \
   --timeout 600 \
   --output-dir ./benchmark_results
 
 # Sequential execution with custom judge model
 uv run agents_mcp_usage/multi_mcp/eval_multi_mcp/run_multi_evals.py \
-  --models "gemini-2.5-pro,claude-3-opus" \
+  --models "gemini-2.5-pro-preview-06-05,claude-3-opus" \
   --runs 3 \
   --sequential \
-  --judge-model "gemini-2.5-pro" \
+  --judge-model "gemini-2.5-pro-preview-06-05" \
   --output-dir ./comparative_analysis
 ```
 
@@ -295,7 +298,7 @@ Key features:
 
 **File:** `multi_mcp_use/adk_mcp.py`
 
-This example demonstrates how to use multiple MCP servers with Google's Agent Development Kit (ADK).
+This example demonstrates how to use multiple MCP servers with Google's Agent Development Kit (ADK 1.3.0).
 
 ```bash
 uv run agents_mcp_usage/multi_mcp/multi_mcp_use/adk_mcp.py
@@ -304,10 +307,17 @@ uv run agents_mcp_usage/multi_mcp/multi_mcp_use/adk_mcp.py
 Key features:
 - Uses Google's ADK framework with Gemini model
 - Connects to both Python MCP server and Python Mermaid validator
-- Demonstrates proper connection management with contextlib.AsyncExitStack
+- Uses new ADK 1.3.0 patterns for toolset management
+- Implements proper resource tracking for MCP connections
 - Shows how to handle asynchronous MCP tool integration
+- Supports ADK web UI through module exports and callback-based tool attachment
 - Uses a simple test case that utilizes both MCP servers in a single query
 
+To run with the ADK web UI:
+```bash
+make adk_multi_ui
+```
+
 ### Multi-MCP Evaluation
 
 **File:** `eval_multi_mcp/evals_pydantic_mcp.py`

agents_mcp_usage/multi_mcp/multi_mcp_use/__init__.py

@@ -0,0 +1,4 @@
+# Multi MCP usage examples package
+
+# Allow discovery of the ADK agent via the ADK web UI
+from .adk_mcp import root_agent

agents_mcp_usage/multi_mcp/multi_mcp_use/adk_mcp.py

@@ -8,14 +8,19 @@
 - RunConfig for request limiting and control
 - Enhanced event tracking and metrics collection
 - Graceful resource cleanup and Logfire instrumentation
+- Web UI integration via ADK's callback system
 
 Compatible with ADK v1.3.0+ and showcases production-ready patterns
 for complex agent architectures involving multiple tool sources.
+
+This module provides an ADK agent that can be used both in the ADK web UI
+and directly from the command line. The agent uses multiple MCP servers
+to access different external functionalities.
 """
 
 import asyncio
 import os
-import time
+from typing import List, Tuple, Any
 
 import logfire
 from dotenv import load_dotenv
@@ -42,22 +47,26 @@
 logfire.configure(send_to_logfire="if-token-present", service_name="adk-multi-mcp")
 logfire.instrument_mcp()
 
+# Global variable to store toolset instances for cleanup
+_ACTIVE_TOOLSETS = []
+
+# Flag to track if tools have been attached
+TOOLS_ATTACHED = False
 
-async def get_tools_async() -> tuple[list, list]:
+
+async def get_tools_async() -> Tuple[List[Any], List[MCPToolset]]:
     """Initializes connections to MCP servers and returns their tools.
 
     This function connects to the local example server and the mermaid
-    validator server, and returns the combined tools and a combined exit
-    stack for cleanup.
+    validator server, and returns the combined tools and a list of toolsets
+    for cleanup.
 
     Returns:
         A tuple containing the list of all tools and the list of toolsets for cleanup.
     """
+    global _ACTIVE_TOOLSETS
     print("Connecting to MCP servers...")
 
-    # Keep track of toolsets for cleanup (ADK v1.3.0+ API)
-    toolsets = []
-
     # Set up MCP server connections
     local_server = StdioServerParameters(
         command="uv",
@@ -76,25 +85,111 @@ async def get_tools_async() -> tuple[list, list]:
         ],
     )
 
-    # Connect to local python MCP server (ADK v1.3.0+ API)
-    local_connection = StdioConnectionParams(server_params=local_server)
-    local_toolset = MCPToolset(connection_params=local_connection)
-    local_tools = await local_toolset.get_tools()
-    toolsets.append(local_toolset)
-    print(f"Connected to local python MCP server. Found {len(local_tools)} tools.")
+    local_toolset = None
+    mermaid_toolset = None
+    toolsets = []
+
+    try:
+        # Connect to local python MCP server
+        local_connection = StdioConnectionParams(server_params=local_server)
+        local_toolset = MCPToolset(connection_params=local_connection)
+        local_tools = await local_toolset.get_tools()
+        toolsets.append(local_toolset)
+        _ACTIVE_TOOLSETS.append(local_toolset)
+        print(f"Connected to local python MCP server. Found {len(local_tools)} tools.")
+
+        # Connect to mermaid MCP server
+        mermaid_connection = StdioConnectionParams(server_params=mermaid_server)
+        mermaid_toolset = MCPToolset(connection_params=mermaid_connection)
+        mermaid_tools = await mermaid_toolset.get_tools()
+        toolsets.append(mermaid_toolset)
+        _ACTIVE_TOOLSETS.append(mermaid_toolset)
+        print(f"Connected to mermaid MCP server. Found {len(mermaid_tools)} tools.")
+
+        # Combine tools from both servers
+        all_tools = local_tools + mermaid_tools
+        print(f"Total tools available: {len(all_tools)}")
+
+        return all_tools, toolsets
+
+    except Exception as e:
+        # Clean up in case of initialisation error
+        if local_toolset in toolsets:
+            try:
+                await local_toolset.close()
+                toolsets.remove(local_toolset)
+                if local_toolset in _ACTIVE_TOOLSETS:
+                    _ACTIVE_TOOLSETS.remove(local_toolset)
+            except Exception:
+                pass
+
+        if mermaid_toolset in toolsets:
+            try:
+                await mermaid_toolset.close()
+                toolsets.remove(mermaid_toolset)
+                if mermaid_toolset in _ACTIVE_TOOLSETS:
+                    _ACTIVE_TOOLSETS.remove(mermaid_toolset)
+            except Exception:
+                pass
+
+        raise e
+
+
+async def cleanup_toolsets():
+    """Clean up any active MCP toolset connections."""
+    global _ACTIVE_TOOLSETS
 
-    # Connect to npx mermaid MCP server (ADK v1.3.0+ API)
-    mermaid_connection = StdioConnectionParams(server_params=mermaid_server)
-    mermaid_toolset = MCPToolset(connection_params=mermaid_connection)
-    mermaid_tools = await mermaid_toolset.get_tools()
-    toolsets.append(mermaid_toolset)
-    print(f"Connected to npx mermaid MCP server. Found {len(mermaid_tools)} tools.")
+    for toolset in _ACTIVE_TOOLSETS:
+        try:
+            await toolset.close()
+            print("MCP toolset connection closed.")
+        except asyncio.CancelledError:
+            print("MCP cleanup cancelled - this is normal")
+        except Exception as e:
+            print(f"Warning: Error during toolset cleanup: {e}")
 
-    # Combine tools from both servers
-    all_tools = local_tools + mermaid_tools
-    print(f"Total tools available: {len(all_tools)}")
+    _ACTIVE_TOOLSETS = []
 
-    return all_tools, toolsets
+
+# Define a before_agent_callback to attach tools
+async def attach_tools_callback(callback_context):
+    """Callback to attach tools to the agent before it runs.
+
+    Args:
+        callback_context: The callback context from ADK.
+
+    Returns:
+        None: The callback doesn't modify the content.
+    """
+    await ensure_tools_attached()
+    return None
+
+
+# This is the agent that will be imported by the ADK web UI
+root_agent = LlmAgent(
+    model="gemini-2.0-flash",
+    name="multi_mcp_adk_assistant",
+    instruction="You are an assistant that uses multiple MCP servers to help users. You have access to both a Python MCP server with time tools and a Mermaid diagram validator.",
+    before_agent_callback=attach_tools_callback,  # This ensures tools are attached
+)
+
+
+# Function to dynamically attach tools to the agent
+async def ensure_tools_attached():
+    """Ensures that tools are attached to the agent before it's used."""
+    global TOOLS_ATTACHED
+
+    if not TOOLS_ATTACHED:
+        try:
+            tools, _ = await get_tools_async()
+            print(f"✓ Connected to MCP servers. Found {len(tools)} tools.")
+            # Update the agent's tools
+            root_agent.tools = tools
+            TOOLS_ATTACHED = True
+        except Exception as e:
+            print(f"Error attaching MCP tools: {e}")
+            # Set empty tools to avoid errors
+            root_agent.tools = []
 
 
 async def main(query: str = "Hi!", request_limit: int = 5) -> None:
@@ -108,17 +203,9 @@ async def main(query: str = "Hi!", request_limit: int = 5) -> None:
         query: The query to run the agent with.
         request_limit: The maximum number of LLM calls allowed.
     """
-    toolsets = []
     try:
-        # Get tools from MCP servers
-        tools, toolsets = await get_tools_async()
-
-        # Create agent
-        agent = LlmAgent(
-            model="gemini-2.0-flash",
-            name="multi_mcp_adk",
-            tools=tools,
-        )
+        # Ensure tools are attached to the agent
+        await ensure_tools_attached()
 
         # Create async session service
         session_service = InMemorySessionService()
@@ -130,10 +217,10 @@ async def main(query: str = "Hi!", request_limit: int = 5) -> None:
         # Create a RunConfig with a limit for LLM calls (500 is the default)
         run_config = RunConfig(max_llm_calls=request_limit)
 
-        # Create runner
+        # Create runner using the globally defined agent
         runner = Runner(
             app_name="multi_mcp_adk",
-            agent=agent,
+            agent=root_agent,
             session_service=session_service,
         )
 
@@ -158,14 +245,7 @@ async def main(query: str = "Hi!", request_limit: int = 5) -> None:
     finally:
         # Clean up MCP toolsets to prevent asyncio shutdown errors
         print("Cleaning up MCP connections...")
-        for i, toolset in enumerate(toolsets):
-            try:
-                await toolset.close()
-                print(f"Toolset {i + 1} closed successfully.")
-            except asyncio.CancelledError:
-                print(f"Toolset {i + 1} cleanup cancelled - this is normal")
-            except Exception as e:
-                print(f"Warning during cleanup of toolset {i + 1}: {e}")
+        await cleanup_toolsets()
         print("MCP cleanup completed.")