Support MCP Elicitation, Sampling, and Roots via new McpSessionHost (#6833)

Co-authored-by: Tyler Payne <[email protected]>
Co-authored-by: Claude <[email protected]>
Co-authored-by: Eric Zhu <[email protected]>

Author: 4 people

python/packages/autogen-ext/examples/mcp_example_server.py

@@ -0,0 +1,75 @@
+"""From: https://github.com/modelcontextprotocol/python-sdk?tab=readme-ov-file#elicitation"""
+
+from pathlib import Path
+
+from mcp import SamplingMessage
+from mcp.server.fastmcp import Context, FastMCP
+from mcp.types import TextContent
+from pydantic import BaseModel, Field
+
+mcp = FastMCP(name="Elicitation Example")
+
+
+class BookingPreferences(BaseModel):
+    """Schema for collecting user preferences."""
+
+    checkAlternative: bool = Field(description="Would you like to check another time?")
+    alternativeTime: str = Field(
+        description="Alternative time.",
+    )
+
+
+@mcp.tool()
+async def book_table(
+    time: str,
+    party_size: int,
+    ctx: Context,
+) -> str:
+    """Book a table with time availability check."""
+    # time unavailable - ask user for alternative
+    result = await ctx.elicit(
+        message=(f"No tables available for {party_size} at {time}. Would you like to try another time?"),
+        schema=BookingPreferences,
+    )
+
+    if result.action == "accept" and result.data:
+        if result.data.checkAlternative:
+            return f"[SUCCESS] Booked for {result.data.alternativeTime}"
+        return "[CANCELLED] No booking made"
+    return "[CANCELLED] Booking cancelled"
+
+
+@mcp.tool()
+async def list_dir(path: Path, ctx: Context) -> list[str]:
+    """List the files and directories in path"""
+    roots = await ctx.session.list_roots()
+    for root in roots.roots:
+        root_path = root.uri.path
+        if root_path:
+            root_path = Path(root_path)
+            try:
+                _ = path.relative_to(root_path)
+                return ["Downloads", "Documents", "image.png", "presentation.pptx"]
+            except ValueError:
+                # Skip relative_to failure
+                pass
+    raise ValueError(f"Cannot list_dir in {path} because it is not a child of the available roots.")
+
+
+@mcp.tool()
+async def generate_poem(topic: str, ctx: Context) -> str:
+    poem = await ctx.session.create_message(
+        [SamplingMessage(role="user", content=TextContent(type="text", text=f"Write a poem about {topic}."))],
+        max_tokens=100,
+        system_prompt="You are a very creative poet.",
+        temperature=0.8,
+        stop_sequences=["\n\n"],
+    )
+    if isinstance(poem.content, TextContent):
+        return poem.content.text
+    else:
+        raise TypeError(f"Unrecognized message response type {type(poem.content).__name__}")
+
+
+if __name__ == "__main__":
+    mcp.run("stdio")

python/packages/autogen-ext/examples/mcp_session_host_example.py

@@ -0,0 +1,251 @@
+"""
+Interactive MCP Host Capabilities Demo
+
+This example demonstrates advanced MCP host capabilities including:
+- Sampling: Language model text generation requests from MCP server back to host
+- Elicitation: Interactive user input collection through command-line
+- Roots: File system root listing
+
+The demo is fully interactive and allows you to communicate directly with
+the MCP server through the command line interface.
+"""
+
+import argparse
+import asyncio
+import json
+import logging
+import sys
+from pathlib import Path
+
+import yaml
+from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
+from autogen_agentchat.conditions import MaxMessageTermination
+from autogen_agentchat.teams import RoundRobinGroupChat
+from autogen_agentchat.ui import Console
+from autogen_core.models import ChatCompletionClient
+from autogen_ext.models.openai import OpenAIChatCompletionClient
+from autogen_ext.tools.mcp import (
+    ChatCompletionClientSampler,
+    McpSessionHost,
+    McpWorkbench,
+    StaticRootsProvider,
+    StdioElicitor,
+    StdioServerParams,
+)
+from mcp.types import Root
+from pydantic import FileUrl
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(message)s",  # Clean format for demo output
+)
+logger = logging.getLogger(__name__)
+logging.getLogger("autogen_core").setLevel(logging.WARNING)
+
+
+def load_model_client_from_config(config_path: str) -> ChatCompletionClient:
+    """Load a ChatCompletionClient from a JSON or YAML config file.
+
+    Args:
+        config_path: Path to the JSON or YAML config file
+
+    Returns:
+        ChatCompletionClient: Loaded model client
+
+    Raises:
+        FileNotFoundError: If config file doesn't exist
+        ValueError: If config format is invalid or unsupported file type
+    """
+    config_file = Path(config_path)
+    if not config_file.exists():
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    # Load config based on file extension
+    if config_file.suffix.lower() == ".json":
+        with open(config_file, "r") as f:
+            config_data = json.load(f)
+    elif config_file.suffix.lower() in [".yml", ".yaml"]:
+        with open(config_file, "r") as f:
+            config_data = yaml.safe_load(f)
+    else:
+        raise ValueError(f"Unsupported config file type: {config_file.suffix}. Use .json, .yml, or .yaml")
+
+    if not isinstance(config_data, dict):
+        raise ValueError("Config file must contain a JSON/YAML object")
+
+    logger.info(f"📄 Loading ChatCompletionClient from config: {config_path}")
+    return ChatCompletionClient.load_component(config_data)
+
+
+async def interactive_mcp_demo(config_path: str | None = None):
+    """Interactive MCP host capabilities demo with command-line interface."""
+    logger.info("🌟 Interactive MCP Host Capabilities Demo")
+    logger.info("=" * 60)
+    logger.info("This demo showcases MCP server-to-host communication:")
+    logger.info("• Sampling: MCP server requests language model generation")
+    logger.info("• Elicitation: MCP server requests user input via AgentElicitor")
+    logger.info("• Roots: MCP server lists available file system roots")
+    logger.info("=" * 60)
+
+    # Setup model client for sampling
+    if config_path:
+        logger.info("⚙️ Loading model client from config file...")
+        model_client = load_model_client_from_config(config_path)
+    else:
+        logger.info("⚙️ Setting up default OpenAI model client (gpt-4)...")
+        model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")
+
+    other_assistant = AssistantAgent(
+        "booking_assistant",
+        model_client=model_client,
+        description="An AI assistant who helps a user book 5pm reservations.",
+    )
+
+    sampler = ChatCompletionClientSampler(model_client)
+
+    # Start runtime and create AgentElicitor that targets the UserProxy
+    logger.info("🎯 Creating StdioElicitor...")
+    elicitor = StdioElicitor()
+
+    roots = StaticRootsProvider(
+        [Root(uri=FileUrl("file:///home"), name="Home"), Root(uri=FileUrl("file:///tmp"), name="Tmp")]
+    )
+
+    # Create host with all capabilities including elicitation
+    logger.info("🏠 Creating MCP session host with sampling, elicitation, and roots support...")
+    host = McpSessionHost(
+        sampler=sampler,  # Support sampling via model clicent
+        elicitor=elicitor,  # Support elicitation via booking_assistant
+        roots=roots, # support roots in /home and /tmp
+    )
+
+    # Setup workbench with host
+    logger.info("🔧 Creating MCP Workbench for mcp_example_server...")
+    mcp_workbench = McpWorkbench(
+        server_params=StdioServerParams(
+            command=sys.executable,
+            args=[str(Path(__file__).parent / "mcp_example_server.py")],
+            read_timeout_seconds=60,
+        ),
+        host=host,
+    )
+
+    # Create assistant with MCP capabilities
+    assistant = AssistantAgent(
+        "mcp_assistant",
+        model_client=model_client,
+        workbench=mcp_workbench,
+        description="An AI assistant with access to MCP tools that can request sampling and elicitation from the host",
+    )
+
+    # Create RoundRobinGroupChat with the agents
+    logger.info("🔄 Setting up RoundRobinGroupChat...")
+    team = RoundRobinGroupChat(
+        [assistant, other_assistant], termination_condition=MaxMessageTermination(max_messages=2)
+    )
+
+    # Run the team with the initial task
+    tasks = ["Book a table for 2 at 7pm", "Generate a poem about computer protocols.", "ls /home", "ls /bin"]
+    for task in tasks:
+        await team.reset()
+        result = await Console(team.run_stream(task=task))
+
+        logger.info("💬 Team conversation:")
+        for message in result.messages:
+            header = f"--- {type(message).__name__.upper()} ---"
+            logger.info(header)
+            logger.info(message.model_dump_json(indent=2))
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Interactive MCP Host Capabilities Demo with AgentElicitor",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Run with default OpenAI GPT-4 client
+  python mcp_elicitation_example.py
+
+  # Run with custom model client from config
+  python mcp_elicitation_example.py --config model_config.json
+  python mcp_elicitation_example.py --config model_config.yaml
+
+Config file format (JSON/YAML):
+{
+  "component_type": "OpenAIChatCompletionClient",
+  "model": "gpt-4",
+  "api_key": "your-api-key"
+}
+        """,
+    )
+
+    parser.add_argument(
+        "--config",
+        "-c",
+        type=str,
+        help="Path to JSON or YAML config file containing ChatCompletionClient configuration",
+    )
+
+    return parser.parse_args()
+
+
+async def main():
+    """
+    Run the interactive MCP host capabilities demonstration.
+
+    This demo allows direct command-line interaction with an MCP-enabled assistant
+    that can use tools requiring host-side capabilities like sampling and elicitation.
+    """
+    args = parse_arguments()
+
+    try:
+        await interactive_mcp_demo(config_path=args.config)
+    except KeyboardInterrupt:
+        logger.info("\n👋 Demo interrupted by user. Goodbye!")
+    except Exception as e:
+        logger.error(f"❌ Error running demo: {e}")
+        logger.info("Troubleshooting tips:")
+        logger.info("1. Install the everything server:")
+        logger.info("   npm install -g @modelcontextprotocol/server-everything")
+        logger.info("2. Ensure your OpenAI API key is configured")
+        logger.info("3. Check that Node.js and npx are available in your PATH")
+        logger.info("4. Make sure you have internet connectivity for npm package download")
+        if args.config:
+            logger.info("5. Check that your config file exists and contains valid ChatCompletionClient configuration")
+
+
+if __name__ == "__main__":
+    """
+    Interactive MCP Host Capabilities Demo with AgentElicitor
+
+    This demo provides a command-line interface to interact with an MCP-enabled
+    assistant that demonstrates advanced host capabilities:
+
+    🔄 Sampling: MCP server can request language model text generation from the host
+    ❓ Elicitation: MCP server can request interactive user input via AgentElicitor → UserProxy
+    📁 Roots: MCP server can request file system root listings from the host
+
+    Key Features:
+    - AgentElicitor routes elicitation requests from MCP server to UserProxyAgent
+    - Full bidirectional communication between MCP server and AutoGen agents
+    - Interactive command-line interface for real-time demonstration
+
+    Prerequisites:
+    1. Install the everything reference server:
+       npm install -g @modelcontextprotocol/server-everything
+    2. Set up your OpenAI API key (required for sampling capability)
+    3. Ensure Node.js and npx are available in your PATH
+
+    Usage:
+    - Run with default model: python mcp_elicitation_example.py
+    - Run with custom model: python mcp_elicitation_example.py --config model.json
+    - Interact through the command-line interface
+    - Ask the assistant to use MCP tools that demonstrate host capabilities
+    - Watch elicitation requests get routed through the AgentElicitor
+    - Type 'quit' to exit the interactive session
+    """
+
+    # Run the interactive demo
+    asyncio.run(main())

python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py

@@ -1,6 +1,19 @@
 from ._actor import McpSessionActor
 from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams
 from ._factory import mcp_server_tools
+from ._host import (
+    ChatCompletionClientSampler,
+    ChatCompletionClientSamplerConfig,
+    Elicitor,
+    McpSessionHost,
+    RootsProvider,
+    Sampler,
+    StaticRootsProvider,
+    StaticRootsProviderConfig,
+    StdioElicitor,
+    StdioElicitorConfig,
+    StreamElicitor,
+)
 from ._session import create_mcp_server_session
 from ._sse import SseMcpToolAdapter
 from ._stdio import StdioMcpToolAdapter
@@ -19,4 +32,15 @@
     "McpServerParams",
     "mcp_server_tools",
     "McpWorkbench",
+    "Elicitor",
+    "StdioElicitor",
+    "StdioElicitorConfig",
+    "StreamElicitor",
+    "RootsProvider",
+    "StaticRootsProvider",
+    "StaticRootsProviderConfig",
+    "McpSessionHost",
+    "ChatCompletionClientSampler",
+    "ChatCompletionClientSamplerConfig",
+    "Sampler",
 ]