first commit

This commit is contained in:
DigiJ
2026-03-13 12:56:43 -07:00
commit 159cf9fcfe
309 changed files with 64584 additions and 0 deletions

View File

@@ -0,0 +1,22 @@
from ._actor import McpSessionActor
from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams
from ._factory import mcp_server_tools
from ._session import create_mcp_server_session
from ._sse import SseMcpToolAdapter
from ._stdio import StdioMcpToolAdapter
from ._streamable_http import StreamableHttpMcpToolAdapter
from ._workbench import McpWorkbench
__all__ = [
"create_mcp_server_session",
"McpSessionActor",
"StdioMcpToolAdapter",
"StdioServerParams",
"SseMcpToolAdapter",
"SseServerParams",
"StreamableHttpMcpToolAdapter",
"StreamableHttpServerParams",
"McpServerParams",
"mcp_server_tools",
"McpWorkbench",
]

View File

@@ -0,0 +1,310 @@
import asyncio
import atexit
import base64
import io
import logging
from typing import Any, Coroutine, Dict, Mapping, TypedDict
from agentdhal_core import Component, ComponentBase, ComponentModel, Image
from agentdhal_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 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
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]):
component_type = "mcp_session_actor"
component_config_schema = McpSessionActorConfig
component_provider_override = "agentdhal_extensions.tools.mcp.McpSessionActor"
server_params: McpServerParams
# model_config = ConfigDict(arbitrary_types_allowed=True)
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
self._actor_task = asyncio.create_task(self._run_actor())
async def call(self, type: str, args: McpActorArgs | None = None) -> McpFuture:
if not self._active:
raise RuntimeError("MCP Actor not running, call initialize() first")
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", "list_prompts", "list_resources", "list_resource_templates", "shutdown"}:
await self._command_queue.put({"type": type, "future": fut})
res = await fut
elif type in {"call_tool", "read_resource", "get_prompt"}:
if args is None:
raise ValueError(f"args is required for {type}")
name = args.get("name", None)
kwargs = args.get("kargs", {})
if type == "call_tool" and name is None:
raise ValueError("name is required for call_tool")
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}")
return res
async def close(self) -> None:
if not self._active or self._actor_task is None:
return
self._shutdown_future = asyncio.Future()
await self._command_queue.put({"type": "shutdown", "future": self._shutdown_future})
await self._shutdown_future
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, 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":
cmd["future"].set_result("ok")
break
elif cmd["type"] == "call_tool":
try:
result = session.call_tool(name=cmd["name"], arguments=cmd["args"])
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
def _sync_shutdown(self) -> None:
if not self._active or self._actor_task is None:
return
try:
loop = asyncio.get_event_loop()
except RuntimeError:
# No loop available — interpreter is likely shutting down
return
if loop.is_closed():
return
if loop.is_running():
loop.create_task(self.close())
else:
loop.run_until_complete(self.close())
def _to_config(self) -> McpSessionActorConfig:
"""
Convert the adapter to its configuration representation.
Returns:
McpSessionConfig: The configuration of the adapter.
"""
return McpSessionActorConfig(server_params=self.server_params)
@classmethod
def _from_config(cls, config: McpSessionActorConfig) -> Self:
"""
Create an instance of McpSessionActor from its configuration.
Args:
config (McpSessionConfig): The configuration of the adapter.
Returns:
McpSessionActor: An instance of SseMcpToolAdapter.
"""
return cls(server_params=config.server_params)

View File

@@ -0,0 +1,190 @@
import asyncio
import builtins
import json
from abc import ABC
from typing import Any, Dict, Generic, Sequence, Type, TypeVar
from agentdhal_core import CancellationToken
from agentdhal_core.tools import BaseTool
from agentdhal_core.utils import schema_to_pydantic_model
from pydantic import BaseModel
from pydantic.networks import AnyUrl
from mcp import ClientSession, Tool
from mcp.types import AudioContent, ContentBlock, EmbeddedResource, ImageContent, ResourceLink, TextContent
from ._config import McpServerParams
from ._session import create_mcp_server_session
TServerParams = TypeVar("TServerParams", bound=McpServerParams)
class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
"""
Base adapter class for MCP tools to make them compatible with AutoGen.
Args:
server_params (TServerParams): Parameters for the MCP server connection.
tool (Tool): The MCP tool to wrap.
"""
component_type = "tool"
def __init__(self, server_params: TServerParams, tool: Tool, session: ClientSession | None = None) -> None:
self._tool = tool
self._server_params = server_params
self._session = session
# Extract name and description
name = tool.name
description = tool.description or ""
# Create the input model from the tool's schema
input_model = schema_to_pydantic_model(tool.inputSchema)
# Use Any as return type since MCP tool returns can vary
return_type: Type[Any] = object
super().__init__(input_model, return_type, name, description)
async def run(self, args: BaseModel, cancellation_token: CancellationToken) -> Any:
"""
Run the MCP tool with the provided arguments.
Args:
args (BaseModel): The arguments to pass to the tool.
cancellation_token (CancellationToken): Token to signal cancellation.
Returns:
Any: The result of the tool execution.
Raises:
Exception: If the operation is cancelled or the tool execution fails.
"""
# Convert the input model to a dictionary
# Exclude unset values to avoid sending them to the MCP servers which may cause errors
# for many servers.
kwargs = args.model_dump(exclude_unset=True)
if self._session is not None:
# If a session is provided, use it directly.
session = self._session
return await self._run(args=kwargs, cancellation_token=cancellation_token, session=session)
async with create_mcp_server_session(self._server_params) as session:
await session.initialize()
return await self._run(args=kwargs, cancellation_token=cancellation_token, session=session)
def _normalize_payload_to_content_list(self, payload: Sequence[ContentBlock]) -> list[ContentBlock]:
"""
Normalizes a raw tool output payload into a list of content items.
- If payload is already a sequence of ContentBlock items, it's converted to a list and returned.
- If payload is a single ContentBlock item, it's wrapped in a list.
- If payload is a string, it's wrapped in [TextContent(text=payload)].
- Otherwise, the payload is stringified and wrapped in [TextContent(text=str(payload))].
"""
if isinstance(payload, Sequence) and all(
isinstance(item, (TextContent, ImageContent, EmbeddedResource, AudioContent, ResourceLink))
for item in payload
):
return list(payload)
elif isinstance(payload, (TextContent, ImageContent, EmbeddedResource, AudioContent, ResourceLink)):
return [payload]
elif isinstance(payload, str):
return [TextContent(text=payload, type="text")]
else:
return [TextContent(text=str(payload), type="text")]
async def _run(self, args: Dict[str, Any], cancellation_token: CancellationToken, session: ClientSession) -> Any:
exceptions_to_catch: tuple[Type[BaseException], ...]
if hasattr(builtins, "ExceptionGroup"):
exceptions_to_catch = (asyncio.CancelledError, builtins.ExceptionGroup)
else:
exceptions_to_catch = (asyncio.CancelledError,)
try:
if cancellation_token.is_cancelled():
raise asyncio.CancelledError("Operation cancelled")
result_future = asyncio.ensure_future(session.call_tool(name=self._tool.name, arguments=args))
cancellation_token.link_future(result_future)
result = await result_future
normalized_content_list = self._normalize_payload_to_content_list(result.content)
if result.isError:
serialized_error_message = self.return_value_as_string(normalized_content_list)
raise Exception(serialized_error_message)
return normalized_content_list
except exceptions_to_catch:
# Re-raise these specific exception types directly.
raise
@classmethod
async def from_server_params(cls, server_params: TServerParams, tool_name: str) -> "McpToolAdapter[TServerParams]":
"""
Create an instance of McpToolAdapter from server parameters and tool name.
Args:
server_params (TServerParams): Parameters for the MCP server connection.
tool_name (str): The name of the tool to wrap.
Returns:
McpToolAdapter[TServerParams]: An instance of McpToolAdapter.
Raises:
ValueError: If the tool with the specified name is not found.
"""
async with create_mcp_server_session(server_params) as session:
await session.initialize()
tools_response = await session.list_tools()
matching_tool = next((t for t in tools_response.tools if t.name == tool_name), None)
if matching_tool is None:
raise ValueError(
f"Tool '{tool_name}' not found, available tools: {', '.join([t.name for t in tools_response.tools])}"
)
return cls(server_params=server_params, tool=matching_tool)
def return_value_as_string(self, value: list[Any]) -> str:
"""Return a string representation of the result."""
def serialize_item(item: Any) -> dict[str, Any]:
if isinstance(item, (TextContent, ImageContent, AudioContent)):
dumped = item.model_dump()
# Remove the 'meta' field if it exists and is None (for backward compatibility)
if dumped.get("meta") is None:
dumped.pop("meta", None)
return dumped
elif isinstance(item, EmbeddedResource):
type = item.type
resource = {}
for key, val in item.resource.model_dump().items():
# Skip 'meta' field if it's None (for backward compatibility)
if key == "meta" and val is None:
continue
if isinstance(val, AnyUrl):
resource[key] = str(val)
else:
resource[key] = val
dumped_annotations = item.annotations.model_dump() if item.annotations else None
# Remove 'meta' from annotations if it exists and is None
if dumped_annotations and dumped_annotations.get("meta") is None:
dumped_annotations.pop("meta", None)
return {"type": type, "resource": resource, "annotations": dumped_annotations}
elif isinstance(item, ResourceLink):
dumped = item.model_dump()
# Remove the 'meta' field if it exists and is None (for backward compatibility)
if dumped.get("meta") is None:
dumped.pop("meta", None)
# Convert AnyUrl to string for JSON serialization
if "uri" in dumped and isinstance(dumped["uri"], AnyUrl):
dumped["uri"] = str(dumped["uri"])
return dumped
else:
return {}
return json.dumps([serialize_item(item) for item in value])

View File

@@ -0,0 +1,42 @@
from typing import Any, Literal
from pydantic import BaseModel, Field
from typing_extensions import Annotated
from mcp import StdioServerParameters
class StdioServerParams(StdioServerParameters):
"""Parameters for connecting to an MCP server over STDIO."""
type: Literal["StdioServerParams"] = "StdioServerParams"
read_timeout_seconds: float = 5
class SseServerParams(BaseModel):
"""Parameters for connecting to an MCP server over SSE."""
type: Literal["SseServerParams"] = "SseServerParams"
url: str # The SSE endpoint URL.
headers: dict[str, Any] | None = None # Optional headers to include in requests.
timeout: float = 5 # HTTP timeout for regular operations.
sse_read_timeout: float = 60 * 5 # Timeout for SSE read operations.
class StreamableHttpServerParams(BaseModel):
"""Parameters for connecting to an MCP server over Streamable HTTP."""
type: Literal["StreamableHttpServerParams"] = "StreamableHttpServerParams"
url: str # The endpoint URL.
headers: dict[str, Any] | None = None # Optional headers to include in requests.
timeout: float = 30.0 # HTTP timeout for regular operations in seconds.
sse_read_timeout: float = 300.0 # Timeout for SSE read operations in seconds.
terminate_on_close: bool = True
McpServerParams = Annotated[
StdioServerParams | SseServerParams | StreamableHttpServerParams, Field(discriminator="type")
]

View File

@@ -0,0 +1,214 @@
from mcp import ClientSession
from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams
from ._session import create_mcp_server_session
from ._sse import SseMcpToolAdapter
from ._stdio import StdioMcpToolAdapter
from ._streamable_http import StreamableHttpMcpToolAdapter
async def mcp_server_tools(
server_params: McpServerParams,
session: ClientSession | None = None,
) -> list[StdioMcpToolAdapter | SseMcpToolAdapter | StreamableHttpMcpToolAdapter]:
"""Creates a list of MCP tool adapters that can be used with AutoGen agents.
.. warning::
Only connect to trusted MCP servers, especially when using
`StdioServerParams` as it executes commands in the local environment.
This factory function connects to an MCP server and returns adapters for all available tools.
The adapters can be directly assigned to an AutoGen agent's tools list.
.. note::
To use this function, you need to install `mcp` extra for the `autogen-ext` package.
.. code-block:: bash
pip install -U "agentdhal-ext[mcp]"
Args:
server_params (McpServerParams): Connection parameters for the MCP server.
Can be either StdioServerParams for command-line tools or
SseServerParams and StreamableHttpServerParams for HTTP/SSE services.
session (ClientSession | None): Optional existing session to use. This is used
when you want to reuse an existing connection to the MCP server. The session
will be reused when creating the MCP tool adapters.
Returns:
list[StdioMcpToolAdapter | SseMcpToolAdapter | StreamableHttpMcpToolAdapter]:
A list of tool adapters ready to use with AutoGen agents.
Examples:
**Local file system MCP service over standard I/O example:**
Install the filesystem server package from npm (requires Node.js 16+ and npm).
.. code-block:: bash
npm install -g @modelcontextprotocol/server-filesystem
Create an agent that can use all tools from the local filesystem MCP server.
.. code-block:: python
import asyncio
from pathlib import Path
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_extensions.tools.mcp import StdioServerParams, mcp_server_tools
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_core import CancellationToken
async def main() -> None:
# Setup server params for local filesystem access
desktop = str(Path.home() / "Desktop")
server_params = StdioServerParams(
command="npx.cmd", args=["-y", "@modelcontextprotocol/server-filesystem", desktop]
)
# Get all available tools from the server
tools = await mcp_server_tools(server_params)
# Create an agent that can use all the tools
agent = AssistantAgent(
name="file_manager",
model_client=OpenAIChatCompletionClient(model="gpt-4"),
tools=tools, # type: ignore
)
# The agent can now use any of the filesystem tools
await agent.run(task="Create a file called test.txt with some content", cancellation_token=CancellationToken())
if __name__ == "__main__":
asyncio.run(main())
**Local fetch MCP service over standard I/O example:**
Install the `mcp-server-fetch` package.
.. code-block:: bash
pip install mcp-server-fetch
Create an agent that can use the `fetch` tool from the local MCP server.
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_extensions.tools.mcp import StdioServerParams, mcp_server_tools
async def main() -> None:
# Get the fetch tool from mcp-server-fetch.
fetch_mcp_server = StdioServerParams(command="uvx", args=["mcp-server-fetch"])
tools = await mcp_server_tools(fetch_mcp_server)
# Create an agent that can use the fetch tool.
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent = AssistantAgent(name="fetcher", model_client=model_client, tools=tools, reflect_on_tool_use=True) # type: ignore
# Let the agent fetch the content of a URL and summarize it.
result = await agent.run(task="Summarize the content of https://en.wikipedia.org/wiki/Seattle")
print(result.messages[-1])
asyncio.run(main())
**Sharing an MCP client session across multiple tools:**
You can create a single MCP client session and share it across multiple tools.
This is sometimes required when the server maintains a session state
(e.g., a browser state) that should be reused for multiple requests.
The following example show how to create a single MCP client session
to a local `Playwright <https://github.com/microsoft/playwright-mcp>`_
server and share it across multiple tools.
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import TextMentionTermination
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_agentchat.ui import Console
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_extensions.tools.mcp import StdioServerParams, create_mcp_server_session, mcp_server_tools
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o", parallel_tool_calls=False) # type: ignore
params = StdioServerParams(
command="npx",
args=["@playwright/mcp@latest"],
read_timeout_seconds=60,
)
async with create_mcp_server_session(params) as session:
await session.initialize()
tools = await mcp_server_tools(server_params=params, session=session)
print(f"Tools: {[tool.name for tool in tools]}")
agent = AssistantAgent(
name="Assistant",
model_client=model_client,
tools=tools, # type: ignore
)
termination = TextMentionTermination("TERMINATE")
team = RoundRobinGroupChat([agent], termination_condition=termination)
await Console(
team.run_stream(
task="Go to https://ekzhu.com/, visit the first link in the page, then tell me about the linked page."
)
)
asyncio.run(main())
**Remote MCP service over SSE example:**
.. code-block:: python
from agentdhal_extensions.tools.mcp import SseServerParams, mcp_server_tools
async def main() -> None:
# Setup server params for remote service
server_params = SseServerParams(url="https://api.example.com/mcp", headers={"Authorization": "Bearer token"})
# Get all available tools
tools = await mcp_server_tools(server_params)
# Create an agent with all tools
agent = AssistantAgent(name="tool_user", model_client=OpenAIChatCompletionClient(model="gpt-4"), tools=tools) # type: ignore
For more examples and detailed usage, see the samples directory in the package repository.
"""
if session is None:
async with create_mcp_server_session(server_params) as temp_session:
await temp_session.initialize()
tools = await temp_session.list_tools()
else:
tools = await session.list_tools()
if isinstance(server_params, StdioServerParams):
return [StdioMcpToolAdapter(server_params=server_params, tool=tool, session=session) for tool in tools.tools]
elif isinstance(server_params, SseServerParams):
return [SseMcpToolAdapter(server_params=server_params, tool=tool, session=session) for tool in tools.tools]
elif isinstance(server_params, StreamableHttpServerParams):
return [
StreamableHttpMcpToolAdapter(server_params=server_params, tool=tool, session=session)
for tool in tools.tools
]
raise ValueError(f"Unsupported server params type: {type(server_params)}")

View File

@@ -0,0 +1,55 @@
from contextlib import asynccontextmanager
from datetime import timedelta
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
from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams
@asynccontextmanager
async def create_mcp_server_session(
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):
async with stdio_client(server_params) as (read, write):
async with ClientSession(
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):
async with sse_client(**server_params.model_dump(exclude={"type"})) as (read, write):
async with ClientSession(
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):
# Convert float seconds to timedelta for the streamablehttp_client
params_dict = server_params.model_dump(exclude={"type"})
params_dict["timeout"] = timedelta(seconds=server_params.timeout)
params_dict["sse_read_timeout"] = timedelta(seconds=server_params.sse_read_timeout)
async with streamablehttp_client(**params_dict) as (
read,
write,
session_id_callback, # type: ignore[assignment, unused-variable]
):
# TODO: Handle session_id_callback if needed
async with ClientSession(
read_stream=read,
write_stream=write,
read_timeout_seconds=timedelta(seconds=server_params.sse_read_timeout),
sampling_callback=sampling_callback,
) as session:
yield session

View File

@@ -0,0 +1,116 @@
from agentdhal_core import Component
from pydantic import BaseModel
from typing_extensions import Self
from mcp import ClientSession, Tool
from ._base import McpToolAdapter
from ._config import SseServerParams
class SseMcpToolAdapterConfig(BaseModel):
"""Configuration for the MCP tool adapter."""
server_params: SseServerParams
tool: Tool
class SseMcpToolAdapter(
McpToolAdapter[SseServerParams],
Component[SseMcpToolAdapterConfig],
):
"""
Allows you to wrap an MCP tool running over Server-Sent Events (SSE) and make it available to AutoGen.
This adapter enables using MCP-compatible tools that communicate over HTTP with SSE
with AutoGen agents. Common use cases include integrating with remote MCP services,
cloud-based tools, and web APIs that implement the Model Context Protocol (MCP).
.. note::
To use this class, you need to install `mcp` extra for the `autogen-ext` package.
.. code-block:: bash
pip install -U "agentdhal-ext[mcp]"
Args:
server_params (SseServerParameters): Parameters for the MCP server connection,
including URL, headers, and timeouts.
tool (Tool): The MCP tool to wrap.
session (ClientSession, optional): The MCP client session to use. If not provided,
it will create a new session. This is useful for testing or when you want to
manage the session lifecycle yourself.
Examples:
Use a remote translation service that implements MCP over SSE to create tools
that allow AutoGen agents to perform translations:
.. code-block:: python
import asyncio
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_extensions.tools.mcp import SseMcpToolAdapter, SseServerParams
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.ui import Console
from agentdhal_core import CancellationToken
async def main() -> None:
# Create server params for the remote MCP service
server_params = SseServerParams(
url="https://api.example.com/mcp",
headers={"Authorization": "Bearer your-api-key", "Content-Type": "application/json"},
timeout=30, # Connection timeout in seconds
)
# Get the translation tool from the server
adapter = await SseMcpToolAdapter.from_server_params(server_params, "translate")
# Create an agent that can use the translation tool
model_client = OpenAIChatCompletionClient(model="gpt-4")
agent = AssistantAgent(
name="translator",
model_client=model_client,
tools=[adapter],
system_message="You are a helpful translation assistant.",
)
# Let the agent translate some text
await Console(
agent.run_stream(task="Translate 'Hello, how are you?' to Spanish", cancellation_token=CancellationToken())
)
if __name__ == "__main__":
asyncio.run(main())
"""
component_config_schema = SseMcpToolAdapterConfig
component_provider_override = "agentdhal_extensions.tools.mcp.SseMcpToolAdapter"
def __init__(self, server_params: SseServerParams, tool: Tool, session: ClientSession | None = None) -> None:
super().__init__(server_params=server_params, tool=tool, session=session)
def _to_config(self) -> SseMcpToolAdapterConfig:
"""
Convert the adapter to its configuration representation.
Returns:
SseMcpToolAdapterConfig: The configuration of the adapter.
"""
return SseMcpToolAdapterConfig(server_params=self._server_params, tool=self._tool)
@classmethod
def _from_config(cls, config: SseMcpToolAdapterConfig) -> Self:
"""
Create an instance of SseMcpToolAdapter from its configuration.
Args:
config (SseMcpToolAdapterConfig): The configuration of the adapter.
Returns:
SseMcpToolAdapter: An instance of SseMcpToolAdapter.
"""
return cls(server_params=config.server_params, tool=config.tool)

View File

@@ -0,0 +1,74 @@
from agentdhal_core import Component
from pydantic import BaseModel
from typing_extensions import Self
from mcp import ClientSession, Tool
from ._base import McpToolAdapter
from ._config import StdioServerParams
class StdioMcpToolAdapterConfig(BaseModel):
"""Configuration for the MCP tool adapter."""
server_params: StdioServerParams
tool: Tool
class StdioMcpToolAdapter(
McpToolAdapter[StdioServerParams],
Component[StdioMcpToolAdapterConfig],
):
"""Allows you to wrap an MCP tool running over STDIO and make it available to AutoGen.
This adapter enables using MCP-compatible tools that communicate over standard input/output
with AutoGen agents. Common use cases include wrapping command-line tools and local services
that implement the Model Context Protocol (MCP).
.. note::
To use this class, you need to install `mcp` extra for the `autogen-ext` package.
.. code-block:: bash
pip install -U "agentdhal-ext[mcp]"
Args:
server_params (StdioServerParams): Parameters for the MCP server connection,
including command to run and its arguments
tool (Tool): The MCP tool to wrap
session (ClientSession, optional): The MCP client session to use. If not provided,
a new session will be created. This is useful for testing or when you want to
manage the session lifecycle yourself.
See :func:`~agentdhal_extensions.tools.mcp.mcp_server_tools` for examples.
"""
component_config_schema = StdioMcpToolAdapterConfig
component_provider_override = "agentdhal_extensions.tools.mcp.StdioMcpToolAdapter"
def __init__(self, server_params: StdioServerParams, tool: Tool, session: ClientSession | None = None) -> None:
super().__init__(server_params=server_params, tool=tool, session=session)
def _to_config(self) -> StdioMcpToolAdapterConfig:
"""
Convert the adapter to its configuration representation.
Returns:
StdioMcpToolAdapterConfig: The configuration of the adapter.
"""
return StdioMcpToolAdapterConfig(server_params=self._server_params, tool=self._tool)
@classmethod
def _from_config(cls, config: StdioMcpToolAdapterConfig) -> Self:
"""
Create an instance of StdioMcpToolAdapter from its configuration.
Args:
config (StdioMcpToolAdapterConfig): The configuration of the adapter.
Returns:
StdioMcpToolAdapter: An instance of StdioMcpToolAdapter.
"""
return cls(server_params=config.server_params, tool=config.tool)

View File

@@ -0,0 +1,121 @@
from agentdhal_core import Component
from pydantic import BaseModel
from typing_extensions import Self
from mcp import ClientSession, Tool
from ._base import McpToolAdapter
from ._config import StreamableHttpServerParams
class StreamableHttpMcpToolAdapterConfig(BaseModel):
"""Configuration for the MCP tool adapter."""
server_params: StreamableHttpServerParams
tool: Tool
class StreamableHttpMcpToolAdapter(
McpToolAdapter[StreamableHttpServerParams],
Component[StreamableHttpMcpToolAdapterConfig],
):
"""
Allows you to wrap an MCP tool running over Streamable HTTP and make it available to AutoGen.
This adapter enables using MCP-compatible tools that communicate over Streamable HTTP
with AutoGen agents. Common use cases include integrating with remote MCP services,
cloud-based tools, and web APIs that implement the Model Context Protocol (MCP).
.. note::
To use this class, you need to install `mcp` extra for the `autogen-ext` package.
.. code-block:: bash
pip install -U "agentdhal-ext[mcp]"
Args:
server_params (StreamableHttpServerParams): Parameters for the MCP server connection,
including URL, headers, and timeouts.
tool (Tool): The MCP tool to wrap.
session (ClientSession, optional): The MCP client session to use. If not provided,
it will create a new session. This is useful for testing or when you want to
manage the session lifecycle yourself.
Examples:
Use a remote translation service that implements MCP over Streamable HTTP to
create tools that allow AutoGen agents to perform translations:
.. code-block:: python
import asyncio
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_extensions.tools.mcp import StreamableHttpMcpToolAdapter, StreamableHttpServerParams
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.ui import Console
from agentdhal_core import CancellationToken
async def main() -> None:
# Create server params for the remote MCP service
server_params = StreamableHttpServerParams(
url="https://api.example.com/mcp",
headers={"Authorization": "Bearer your-api-key", "Content-Type": "application/json"},
timeout=30.0, # HTTP timeout in seconds
sse_read_timeout=300.0, # SSE read timeout in seconds (5 minutes)
terminate_on_close=True,
)
# Get the translation tool from the server
adapter = await StreamableHttpMcpToolAdapter.from_server_params(server_params, "translate")
# Create an agent that can use the translation tool
model_client = OpenAIChatCompletionClient(model="gpt-4")
agent = AssistantAgent(
name="translator",
model_client=model_client,
tools=[adapter],
system_message="You are a helpful translation assistant.",
)
# Let the agent translate some text
await Console(
agent.run_stream(task="Translate 'Hello, how are you?' to Spanish", cancellation_token=CancellationToken())
)
if __name__ == "__main__":
asyncio.run(main())
"""
component_config_schema = StreamableHttpMcpToolAdapterConfig
component_provider_override = "agentdhal_extensions.tools.mcp.StreamableHttpMcpToolAdapter"
def __init__(
self, server_params: StreamableHttpServerParams, tool: Tool, session: ClientSession | None = None
) -> None:
super().__init__(server_params=server_params, tool=tool, session=session)
def _to_config(self) -> StreamableHttpMcpToolAdapterConfig:
"""
Convert the adapter to its configuration representation.
Returns:
StreamableHttpMcpToolAdapterConfig: The configuration of the adapter.
"""
return StreamableHttpMcpToolAdapterConfig(server_params=self._server_params, tool=self._tool)
@classmethod
def _from_config(cls, config: StreamableHttpMcpToolAdapterConfig) -> Self:
"""
Create an instance of StreamableHttpMcpToolAdapter from its configuration.
Args:
config (StreamableHttpMcpToolAdapterConfig): The configuration of the adapter.
Returns:
StreamableHttpMcpToolAdapter: An instance of StreamableHttpMcpToolAdapter.
"""
return cls(server_params=config.server_params, tool=config.tool)

View File

@@ -0,0 +1,518 @@
import asyncio
import builtins
import warnings
from typing import Any, Dict, List, Literal, Mapping, Optional
from agentdhal_core import CancellationToken, Component, ComponentModel, Image, trace_tool_span
from agentdhal_core.models import ChatCompletionClient
from agentdhal_core.tools import (
ImageResultContent,
ParametersSchema,
TextResultContent,
ToolOverride,
ToolResult,
ToolSchema,
Workbench,
)
from pydantic import BaseModel, Field
from typing_extensions import Self
from mcp.types import (
CallToolResult,
EmbeddedResource,
GetPromptResult,
ImageContent,
ListPromptsResult,
ListResourcesResult,
ListResourceTemplatesResult,
ListToolsResult,
ReadResourceResult,
TextContent,
)
from ._actor import McpSessionActor
from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams
class McpWorkbenchConfig(BaseModel):
server_params: McpServerParams
tool_overrides: Dict[str, ToolOverride] = Field(default_factory=dict)
model_client: ComponentModel | Dict[str, Any] | None = None
class McpWorkbenchState(BaseModel):
type: Literal["McpWorkBenchState"] = "McpWorkBenchState"
class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
"""A workbench that wraps an MCP server and provides an interface
to list and call tools provided by the server.
.. warning::
Only connect to trusted MCP servers, especially when using
`StdioServerParams` as it executes commands in the local environment.
This workbench should be used as a context manager to ensure proper
initialization and cleanup of the underlying MCP session.
.. list-table:: MCP Support
:header-rows: 1
:widths: 30 70
* - MCP Capability
- Supported Features
* - Tools
- list_tools, call_tool
* - Resources
- list_resources, read_resource
* - ResourceTemplates
- list_resource_templates, read_resource_template
* - Prompts
- list_prompts, get_prompt
* - Sampling
- Optional support via model_client
* - Roots
- not supported
* - Ellicitation
- not supported
Args:
server_params (McpServerParams): The parameters to connect to the MCP server.
This can be either a :class:`StdioServerParams` or :class:`SseServerParams`.
tool_overrides (Optional[Dict[str, ToolOverride]]): Optional mapping of original tool
names to override configurations for name and/or description. This allows
customizing how server tools appear to consumers while maintaining the underlying
tool functionality.
model_client: Optional chat completion client to handle sampling requests
from MCP servers that support the sampling capability. This allows MCP
servers to request text generation from a language model during tool
execution. If not provided, sampling requests will return an error.
Raises:
ValueError: If there are conflicts in tool override names.
Examples:
Here is a simple example of how to use the workbench with a `mcp-server-fetch` server:
.. code-block:: python
import asyncio
from agentdhal_extensions.tools.mcp import McpWorkbench, StdioServerParams
async def main() -> None:
params = StdioServerParams(
command="uvx",
args=["mcp-server-fetch"],
read_timeout_seconds=60,
)
# You can also use `start()` and `stop()` to manage the session.
async with McpWorkbench(server_params=params) as workbench:
tools = await workbench.list_tools()
print(tools)
result = await workbench.call_tool(tools[0]["name"], {"url": "https://github.com/"})
print(result)
asyncio.run(main())
Example of using tool overrides:
.. code-block:: python
import asyncio
from agentdhal_extensions.tools.mcp import McpWorkbench, StdioServerParams
from agentdhal_core.tools import ToolOverride
async def main() -> None:
params = StdioServerParams(
command="uvx",
args=["mcp-server-fetch"],
read_timeout_seconds=60,
)
# Override the fetch tool's name and description
overrides = {
"fetch": ToolOverride(name="web_fetch", description="Enhanced web fetching tool with better error handling")
}
async with McpWorkbench(server_params=params, tool_overrides=overrides) as workbench:
tools = await workbench.list_tools()
# The tool will now appear as "web_fetch" with the new description
print(tools)
# Call the overridden tool
result = await workbench.call_tool("web_fetch", {"url": "https://github.com/"})
print(result)
asyncio.run(main())
Example of using the workbench with the `GitHub MCP Server <https://github.com/github/github-mcp-server>`_:
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.ui import Console
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_extensions.tools.mcp import McpWorkbench, StdioServerParams
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano")
server_params = StdioServerParams(
command="docker",
args=[
"run",
"-i",
"--rm",
"-e",
"GITHUB_PERSONAL_ACCESS_TOKEN",
"ghcr.io/github/github-mcp-server",
],
env={
"GITHUB_PERSONAL_ACCESS_TOKEN": "ghp_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX",
},
)
async with McpWorkbench(server_params) as mcp:
agent = AssistantAgent(
"github_assistant",
model_client=model_client,
workbench=mcp,
reflect_on_tool_use=True,
model_client_stream=True,
)
await Console(agent.run_stream(task="Is there a repository named Autogen"))
asyncio.run(main())
Example of using the workbench with the `Playwright MCP Server <https://github.com/microsoft/playwright-mcp>`_:
.. code-block:: python
# First run `npm install -g @playwright/mcp@latest` to install the MCP server.
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_agentchat.conditions import TextMessageTermination
from agentdhal_agentchat.ui import Console
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_extensions.tools.mcp import McpWorkbench, StdioServerParams
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano")
server_params = StdioServerParams(
command="npx",
args=[
"@playwright/mcp@latest",
"--headless",
],
)
async with McpWorkbench(server_params) as mcp:
agent = AssistantAgent(
"web_browsing_assistant",
model_client=model_client,
workbench=mcp,
model_client_stream=True,
)
team = RoundRobinGroupChat(
[agent],
termination_condition=TextMessageTermination(source="web_browsing_assistant"),
)
await Console(team.run_stream(task="Find out how many contributors for the microsoft/autogen repository"))
asyncio.run(main())
"""
component_provider_override = "agentdhal_extensions.tools.mcp.McpWorkbench"
component_config_schema = McpWorkbenchConfig
def __init__(
self,
server_params: McpServerParams,
tool_overrides: Optional[Dict[str, ToolOverride]] = None,
model_client: ChatCompletionClient | None = None,
) -> None:
self._server_params = server_params
self._tool_overrides = tool_overrides or {}
self._model_client = model_client
# Build reverse mapping from override names to original names for call_tool
self._override_name_to_original: Dict[str, str] = {}
for original_name, override in self._tool_overrides.items():
override_name = override.name
if override_name and override_name != original_name:
# Check for conflicts with other override names
if override_name in self._override_name_to_original:
existing_original = self._override_name_to_original[override_name]
raise ValueError(
f"Tool override name '{override_name}' is used by multiple tools: "
f"'{existing_original}' and '{original_name}'. Override names must be unique."
)
self._override_name_to_original[override_name] = original_name
# self._session: ClientSession | None = None
self._actor: McpSessionActor | None = None
self._actor_loop: asyncio.AbstractEventLoop | None = None
self._read = None
self._write = None
@property
def server_params(self) -> McpServerParams:
return self._server_params
async def list_tools(self) -> List[ToolSchema]:
if not self._actor:
await self.start() # fallback to start the actor if not initialized instead of raising an error
# Why? Because when deserializing the workbench, the actor might not be initialized yet.
# raise RuntimeError("Actor is not initialized. Call start() first.")
if self._actor is None:
raise RuntimeError("Actor is not initialized. Please check the server connection.")
result_future = await self._actor.call("list_tools", None)
list_tool_result = await result_future
assert isinstance(
list_tool_result, ListToolsResult
), f"list_tools must return a CallToolResult, instead of : {str(type(list_tool_result))}"
schema: List[ToolSchema] = []
for tool in list_tool_result.tools:
original_name = tool.name
name = original_name
description = tool.description or ""
# Apply overrides if they exist for this tool
if original_name in self._tool_overrides:
override = self._tool_overrides[original_name]
if override.name is not None:
name = override.name
if override.description is not None:
description = override.description
parameters = ParametersSchema(
type="object",
properties=tool.inputSchema.get("properties", {}),
required=tool.inputSchema.get("required", []),
additionalProperties=tool.inputSchema.get("additionalProperties", False),
)
tool_schema = ToolSchema(
name=name,
description=description,
parameters=parameters,
)
schema.append(tool_schema)
return schema
async def call_tool(
self,
name: str,
arguments: Mapping[str, Any] | None = None,
cancellation_token: CancellationToken | None = None,
call_id: str | None = None,
) -> ToolResult:
if not self._actor:
await self.start() # fallback to start the actor if not initialized instead of raising an error
# Why? Because when deserializing the workbench, the actor might not be initialized yet.
# raise RuntimeError("Actor is not initialized. Call start() first.")
if self._actor is None:
raise RuntimeError("Actor is not initialized. Please check the server connection.")
if not cancellation_token:
cancellation_token = CancellationToken()
if not arguments:
arguments = {}
# Check if the name is an override name and map it back to the original
original_name = self._override_name_to_original.get(name, name)
with trace_tool_span(
tool_name=name, # Use the requested name for tracing
tool_call_id=call_id,
):
try:
result_future = await self._actor.call("call_tool", {"name": original_name, "kargs": arguments})
cancellation_token.link_future(result_future)
result = await result_future
assert isinstance(
result, CallToolResult
), f"call_tool must return a CallToolResult, instead of : {str(type(result))}"
result_parts: List[TextResultContent | ImageResultContent] = []
is_error = result.isError
for content in result.content:
if isinstance(content, TextContent):
result_parts.append(TextResultContent(content=content.text))
elif isinstance(content, ImageContent):
result_parts.append(ImageResultContent(content=Image.from_base64(content.data)))
elif isinstance(content, EmbeddedResource):
# TODO: how to handle embedded resources?
# For now we just use text representation.
result_parts.append(TextResultContent(content=content.model_dump_json()))
else:
raise ValueError(f"Unknown content type from server: {type(content)}")
except Exception as e:
error_message = self._format_errors(e)
is_error = True
result_parts = [TextResultContent(content=error_message)]
return ToolResult(name=name, result=result_parts, is_error=is_error) # Return the requested name
@property
def initialize_result(self) -> Any:
if self._actor:
return self._actor.initialize_result
return None
async def list_prompts(self) -> ListPromptsResult:
"""List available prompts from the MCP server."""
if not self._actor:
await self.start()
if self._actor is None:
raise RuntimeError("Actor is not initialized. Please check the server connection.")
result_future = await self._actor.call("list_prompts", None)
list_prompts_result = await result_future
assert isinstance(
list_prompts_result, ListPromptsResult
), f"list_prompts must return a ListPromptsResult, instead of: {str(type(list_prompts_result))}"
return list_prompts_result
async def list_resources(self) -> ListResourcesResult:
"""List available resources from the MCP server."""
if not self._actor:
await self.start()
if self._actor is None:
raise RuntimeError("Actor is not initialized. Please check the server connection.")
result_future = await self._actor.call("list_resources", None)
list_resources_result = await result_future
assert isinstance(
list_resources_result, ListResourcesResult
), f"list_resources must return a ListResourcesResult, instead of: {str(type(list_resources_result))}"
return list_resources_result
async def list_resource_templates(self) -> ListResourceTemplatesResult:
"""List available resource templates from the MCP server."""
if not self._actor:
await self.start()
if self._actor is None:
raise RuntimeError("Actor is not initialized. Please check the server connection.")
result_future = await self._actor.call("list_resource_templates", None)
list_templates_result = await result_future
assert isinstance(
list_templates_result, ListResourceTemplatesResult
), f"list_resource_templates must return a ListResourceTemplatesResult, instead of: {str(type(list_templates_result))}"
return list_templates_result
async def read_resource(self, uri: str) -> ReadResourceResult:
"""Read a resource from the MCP server."""
if not self._actor:
await self.start()
if self._actor is None:
raise RuntimeError("Actor is not initialized. Please check the server connection.")
result_future = await self._actor.call("read_resource", {"name": None, "kargs": {"uri": uri}})
read_resource_result = await result_future
assert isinstance(
read_resource_result, ReadResourceResult
), f"read_resource must return a ReadResourceResult, instead of: {str(type(read_resource_result))}"
return read_resource_result
async def get_prompt(self, name: str, arguments: Optional[Dict[str, str]] = None) -> GetPromptResult:
"""Get a prompt from the MCP server."""
if not self._actor:
await self.start()
if self._actor is None:
raise RuntimeError("Actor is not initialized. Please check the server connection.")
result_future = await self._actor.call("get_prompt", {"name": name, "kargs": {"arguments": arguments}})
get_prompt_result = await result_future
assert isinstance(
get_prompt_result, GetPromptResult
), f"get_prompt must return a GetPromptResult, instead of: {str(type(get_prompt_result))}"
return get_prompt_result
def _format_errors(self, error: Exception) -> str:
"""Recursively format errors into a string."""
error_message = ""
if hasattr(builtins, "ExceptionGroup") and isinstance(error, builtins.ExceptionGroup):
# ExceptionGroup is available in Python 3.11+.
# TODO: how to make this compatible with Python 3.10?
for sub_exception in error.exceptions: # type: ignore
error_message += self._format_errors(sub_exception) # type: ignore
else:
error_message += f"{str(error)}\n"
return error_message
async def start(self) -> None:
if self._actor:
warnings.warn(
"McpWorkbench is already started. No need to start again.",
UserWarning,
stacklevel=2,
)
return # Already initialized, no need to start again
if isinstance(self._server_params, (StdioServerParams, SseServerParams, StreamableHttpServerParams)):
self._actor = McpSessionActor(self._server_params, model_client=self._model_client)
await self._actor.initialize()
self._actor_loop = asyncio.get_event_loop()
else:
raise ValueError(f"Unsupported server params type: {type(self._server_params)}")
async def stop(self) -> None:
if self._actor:
# Close the actor
await self._actor.close()
self._actor = None
else:
raise RuntimeError("McpWorkbench is not started. Call start() first.")
async def reset(self) -> None:
pass
async def save_state(self) -> Mapping[str, Any]:
return McpWorkbenchState().model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
pass
def _to_config(self) -> McpWorkbenchConfig:
model_client_config = None
if self._model_client is not None:
model_client_config = self._model_client.dump_component()
return McpWorkbenchConfig(
server_params=self._server_params, tool_overrides=self._tool_overrides, model_client=model_client_config
)
@classmethod
def _from_config(cls, config: McpWorkbenchConfig) -> Self:
model_client = None
if config.model_client is not None:
model_client = ChatCompletionClient.load_component(config.model_client)
return cls(server_params=config.server_params, tool_overrides=config.tool_overrides, model_client=model_client)
def __del__(self) -> None:
# Ensure the actor is stopped when the workbench is deleted
# Use getattr to safely handle cases where attributes may not be set (e.g., if __init__ failed)
actor = getattr(self, "_actor", None)
actor_loop = getattr(self, "_actor_loop", None)
if actor and actor_loop:
if actor_loop.is_running() and not actor_loop.is_closed():
actor_loop.call_soon_threadsafe(lambda: asyncio.create_task(self.stop()))
else:
msg = "Cannot safely stop actor at [McpWorkbench.__del__]: loop is closed or not running"
warnings.warn(msg, RuntimeWarning, stacklevel=2)