Expand MCP Workbench to support more MCP Client features (#6785)

Co-authored-by: Tyler Payne <tylerpayne@microsoft.com>
Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
Co-authored-by: Victor Dibia <victordibia@microsoft.com>

Author: 4 people

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

@@ -1,27 +1,87 @@
 import asyncio
 import atexit
+import base64
+import io
+import logging
 from typing import Any, Coroutine, Dict, Mapping, TypedDict
 
-from autogen_core import Component, ComponentBase
+from autogen_core import Component, ComponentBase, ComponentModel, Image
+from autogen_core.models import (
+    AssistantMessage,
+    ChatCompletionClient,
+    LLMMessage,
+    ModelInfo,
+    SystemMessage,
+    UserMessage,
+)
+from PIL import Image as PILImage
 from pydantic import BaseModel
 from typing_extensions import Self
 
-from mcp.types import CallToolResult, ListToolsResult
+from mcp import types as mcp_types
+from mcp.client.session import ClientSession
+from mcp.shared.context import RequestContext
 
 from ._config import McpServerParams
 from ._session import create_mcp_server_session
 
-McpResult = Coroutine[Any, Any, ListToolsResult] | Coroutine[Any, Any, CallToolResult]
+logger = logging.getLogger(__name__)
+
+McpResult = (
+    Coroutine[Any, Any, mcp_types.ListToolsResult]
+    | Coroutine[Any, Any, mcp_types.CallToolResult]
+    | Coroutine[Any, Any, mcp_types.ListPromptsResult]
+    | Coroutine[Any, Any, mcp_types.ListResourcesResult]
+    | Coroutine[Any, Any, mcp_types.ListResourceTemplatesResult]
+    | Coroutine[Any, Any, mcp_types.ReadResourceResult]
+    | Coroutine[Any, Any, mcp_types.GetPromptResult]
+)
 McpFuture = asyncio.Future[McpResult]
 
 
+def _parse_sampling_content(
+    content: mcp_types.TextContent | mcp_types.ImageContent | mcp_types.AudioContent, model_info: ModelInfo
+) -> str | Image:
+    """Convert MCP content types to Autogen content types."""
+    if content.type == "text":
+        return content.text
+    elif content.type == "image":
+        if not model_info["vision"]:
+            raise ValueError("Sampling model does not support image content.")
+        # Decode base64 image data and create PIL Image
+        image_data = base64.b64decode(content.data)
+        pil_image = PILImage.open(io.BytesIO(image_data))
+        return Image.from_pil(pil_image)
+    else:
+        raise ValueError(f"Unsupported content type: {content.type}")
+
+
+def _parse_sampling_message(message: mcp_types.SamplingMessage, model_info: ModelInfo) -> LLMMessage:
+    """Convert MCP sampling messages to Autogen messages."""
+    content = _parse_sampling_content(message.content, model_info=model_info)
+    if message.role == "user":
+        return UserMessage(
+            source="user",
+            content=[content],
+        )
+    elif message.role == "assistant":
+        assert isinstance(content, str), "Assistant messages only support string content."
+        return AssistantMessage(
+            source="assistant",
+            content=content,
+        )
+    else:
+        raise ValueError(f"Unrecognized message role: {message.role}")
+
+
 class McpActorArgs(TypedDict):
     name: str | None
     kargs: Mapping[str, Any]
 
 
 class McpSessionActorConfig(BaseModel):
     server_params: McpServerParams
+    model_client: ComponentModel | Dict[str, Any] | None = None
 
 
 class McpSessionActor(ComponentBase[BaseModel], Component[McpSessionActorConfig]):
@@ -33,16 +93,22 @@ class McpSessionActor(ComponentBase[BaseModel], Component[McpSessionActorConfig]
 
     # model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def __init__(self, server_params: McpServerParams) -> None:
+    def __init__(self, server_params: McpServerParams, model_client: ChatCompletionClient | None = None) -> None:
         self.server_params: McpServerParams = server_params
+        self._model_client = model_client
         self.name = "mcp_session_actor"
         self.description = "MCP session actor"
         self._command_queue: asyncio.Queue[Dict[str, Any]] = asyncio.Queue()
         self._actor_task: asyncio.Task[Any] | None = None
         self._shutdown_future: asyncio.Future[Any] | None = None
         self._active = False
+        self._initialize_result: mcp_types.InitializeResult | None = None
         atexit.register(self._sync_shutdown)
 
+    @property
+    def initialize_result(self) -> mcp_types.InitializeResult | None:
+        return self._initialize_result
+
     async def initialize(self) -> None:
         if not self._active:
             self._active = True
@@ -54,17 +120,28 @@ async def call(self, type: str, args: McpActorArgs | None = None) -> McpFuture:
         if self._actor_task and self._actor_task.done():
             raise RuntimeError("MCP actor task crashed", self._actor_task.exception())
         fut: asyncio.Future[McpFuture] = asyncio.Future()
-        if type in {"list_tools", "shutdown"}:
+        if type in {"list_tools", "list_prompts", "list_resources", "list_resource_templates", "shutdown"}:
             await self._command_queue.put({"type": type, "future": fut})
             res = await fut
-        elif type == "call_tool":
+        elif type in {"call_tool", "read_resource", "get_prompt"}:
             if args is None:
-                raise ValueError("args is required for call_tool")
+                raise ValueError(f"args is required for {type}")
             name = args.get("name", None)
             kwargs = args.get("kargs", {})
-            if name is None:
+            if type == "call_tool" and name is None:
                 raise ValueError("name is required for call_tool")
-            await self._command_queue.put({"type": type, "name": name, "args": kwargs, "future": fut})
+            elif type == "read_resource":
+                uri = kwargs.get("uri", None)
+                if uri is None:
+                    raise ValueError("uri is required for read_resource")
+                await self._command_queue.put({"type": type, "uri": uri, "future": fut})
+            elif type == "get_prompt":
+                if name is None:
+                    raise ValueError("name is required for get_prompt")
+                prompt_args = kwargs.get("arguments", None)
+                await self._command_queue.put({"type": type, "name": name, "args": prompt_args, "future": fut})
+            else:  # call_tool
+                await self._command_queue.put({"type": type, "name": name, "args": kwargs, "future": fut})
             res = await fut
         else:
             raise ValueError(f"Unknown command type: {type}")
@@ -79,11 +156,64 @@ async def close(self) -> None:
         await self._actor_task
         self._active = False
 
+    async def _sampling_callback(
+        self,
+        context: RequestContext[ClientSession, Any],
+        params: mcp_types.CreateMessageRequestParams,
+    ) -> mcp_types.CreateMessageResult | mcp_types.ErrorData:
+        """Handle sampling requests using the provided model client."""
+        if self._model_client is None:
+            # Return an error when no model client is available
+            return mcp_types.ErrorData(
+                code=mcp_types.INVALID_REQUEST,
+                message="No model client available for sampling.",
+                data=None,
+            )
+
+        llm_messages: list[LLMMessage] = []
+
+        try:
+            if params.systemPrompt:
+                llm_messages.append(SystemMessage(content=params.systemPrompt))
+
+            for mcp_message in params.messages:
+                llm_messages.append(_parse_sampling_message(mcp_message, model_info=self._model_client.model_info))
+
+        except Exception as e:
+            return mcp_types.ErrorData(
+                code=mcp_types.INVALID_PARAMS,
+                message="Error processing sampling messages.",
+                data=f"{type(e).__name__}: {e}",
+            )
+
+        try:
+            result = await self._model_client.create(messages=llm_messages)
+
+            content = result.content
+            if not isinstance(content, str):
+                content = str(content)
+
+            return mcp_types.CreateMessageResult(
+                role="assistant",
+                content=mcp_types.TextContent(type="text", text=content),
+                model=self._model_client.model_info["family"],
+                stopReason=result.finish_reason,
+            )
+        except Exception as e:
+            return mcp_types.ErrorData(
+                code=mcp_types.INTERNAL_ERROR,
+                message="Error sampling from model client.",
+                data=f"{type(e).__name__}: {e}",
+            )
+
     async def _run_actor(self) -> None:
         result: McpResult
         try:
-            async with create_mcp_server_session(self.server_params) as session:
-                await session.initialize()
+            async with create_mcp_server_session(
+                self.server_params, sampling_callback=self._sampling_callback
+            ) as session:
+                # Save the initialize result
+                self._initialize_result = await session.initialize()
                 while True:
                     cmd = await self._command_queue.get()
                     if cmd["type"] == "shutdown":
@@ -95,15 +225,47 @@ async def _run_actor(self) -> None:
                             cmd["future"].set_result(result)
                         except Exception as e:
                             cmd["future"].set_exception(e)
+                    elif cmd["type"] == "read_resource":
+                        try:
+                            result = session.read_resource(uri=cmd["uri"])
+                            cmd["future"].set_result(result)
+                        except Exception as e:
+                            cmd["future"].set_exception(e)
+                    elif cmd["type"] == "get_prompt":
+                        try:
+                            result = session.get_prompt(name=cmd["name"], arguments=cmd["args"])
+                            cmd["future"].set_result(result)
+                        except Exception as e:
+                            cmd["future"].set_exception(e)
                     elif cmd["type"] == "list_tools":
                         try:
                             result = session.list_tools()
                             cmd["future"].set_result(result)
                         except Exception as e:
                             cmd["future"].set_exception(e)
+                    elif cmd["type"] == "list_prompts":
+                        try:
+                            result = session.list_prompts()
+                            cmd["future"].set_result(result)
+                        except Exception as e:
+                            cmd["future"].set_exception(e)
+                    elif cmd["type"] == "list_resources":
+                        try:
+                            result = session.list_resources()
+                            cmd["future"].set_result(result)
+                        except Exception as e:
+                            cmd["future"].set_exception(e)
+                    elif cmd["type"] == "list_resource_templates":
+                        try:
+                            result = session.list_resource_templates()
+                            cmd["future"].set_result(result)
+                        except Exception as e:
+                            cmd["future"].set_exception(e)
         except Exception as e:
             if self._shutdown_future and not self._shutdown_future.done():
                 self._shutdown_future.set_exception(e)
+            else:
+                logger.exception("Exception in MCP actor task")
         finally:
             self._active = False
             self._actor_task = None

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

@@ -3,6 +3,7 @@
 from typing import AsyncGenerator
 
 from mcp import ClientSession
+from mcp.client.session import SamplingFnT
 from mcp.client.sse import sse_client
 from mcp.client.stdio import stdio_client
 from mcp.client.streamable_http import streamablehttp_client
@@ -12,7 +13,7 @@
 
 @asynccontextmanager
 async def create_mcp_server_session(
-    server_params: McpServerParams,
+    server_params: McpServerParams, sampling_callback: SamplingFnT | None = None
 ) -> AsyncGenerator[ClientSession, None]:
     """Create an MCP client session for the given server parameters."""
     if isinstance(server_params, StdioServerParams):
@@ -21,6 +22,7 @@ async def create_mcp_server_session(
                 read_stream=read,
                 write_stream=write,
                 read_timeout_seconds=timedelta(seconds=server_params.read_timeout_seconds),
+                sampling_callback=sampling_callback,
             ) as session:
                 yield session
     elif isinstance(server_params, SseServerParams):
@@ -29,6 +31,7 @@ async def create_mcp_server_session(
                 read_stream=read,
                 write_stream=write,
                 read_timeout_seconds=timedelta(seconds=server_params.sse_read_timeout),
+                sampling_callback=sampling_callback,
             ) as session:
                 yield session
     elif isinstance(server_params, StreamableHttpServerParams):
@@ -47,5 +50,6 @@ async def create_mcp_server_session(
                 read_stream=read,
                 write_stream=write,
                 read_timeout_seconds=timedelta(seconds=server_params.sse_read_timeout),
+                sampling_callback=sampling_callback,
             ) as session:
                 yield session