diff --git a/python/packages/autogen-ext/pyproject.toml b/python/packages/autogen-ext/pyproject.toml index ba4cd4bff..db82cfa29 100644 --- a/python/packages/autogen-ext/pyproject.toml +++ b/python/packages/autogen-ext/pyproject.toml @@ -144,7 +144,7 @@ semantic-kernel-all = [ rich = ["rich>=13.9.4"] -mcp = ["mcp>=1.6.0"] +mcp = ["mcp>=1.8.1"] canvas = [ "unidiff>=0.7.5", ] diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py index bde36794c..fcf414874 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py @@ -1,9 +1,10 @@ from ._actor import McpSessionActor -from ._config import McpServerParams, SseServerParams, StdioServerParams +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__ = [ @@ -13,6 +14,8 @@ __all__ = [ "StdioServerParams", "SseMcpToolAdapter", "SseServerParams", + "StreamableHttpMcpToolAdapter", + "StreamableHttpServerParams", "McpServerParams", "mcp_server_tools", "McpWorkbench", diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py index 215102c5e..6e4f51275 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py @@ -1,3 +1,4 @@ +from datetime import timedelta from typing import Any, Literal from mcp import StdioServerParameters @@ -18,10 +19,24 @@ class SseServerParams(BaseModel): type: Literal["SseServerParams"] = "SseServerParams" - url: str - headers: dict[str, Any] | None = None - timeout: float = 5 - sse_read_timeout: float = 60 * 5 + 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. -McpServerParams = Annotated[StdioServerParams | SseServerParams, Field(discriminator="type")] +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: timedelta = timedelta(seconds=30) # HTTP timeout for regular operations. + sse_read_timeout: timedelta = timedelta(seconds=60 * 5) # Timeout for SSE read operations. + terminate_on_close: bool = True + + +McpServerParams = Annotated[ + StdioServerParams | SseServerParams | StreamableHttpServerParams, Field(discriminator="type") +] diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py index 5e23cff92..37af9ef4f 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py @@ -1,15 +1,16 @@ from mcp import ClientSession -from ._config import McpServerParams, SseServerParams, StdioServerParams +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]: +) -> list[StdioMcpToolAdapter | SseMcpToolAdapter | StreamableHttpMcpToolAdapter]: """Creates a list of MCP tool adapters that can be used with AutoGen agents. This factory function connects to an MCP server and returns adapters for all available tools. @@ -26,14 +27,14 @@ async def mcp_server_tools( Args: server_params (McpServerParams): Connection parameters for the MCP server. Can be either StdioServerParams for command-line tools or - SseServerParams for HTTP/SSE services. + 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]: A list of tool adapters ready to use - with AutoGen agents. + list[StdioMcpToolAdapter | SseMcpToolAdapter | StreamableHttpMcpToolAdapter]: + A list of tool adapters ready to use with AutoGen agents. Examples: @@ -200,4 +201,9 @@ async def mcp_server_tools( 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)}") diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py index d4d696769..590f209c3 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py @@ -5,8 +5,9 @@ from typing import AsyncGenerator from mcp import ClientSession 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 +from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams @asynccontextmanager @@ -24,5 +25,22 @@ async def create_mcp_server_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) as session: + async with ClientSession( + read_stream=read, + write_stream=write, + read_timeout_seconds=timedelta(seconds=server_params.sse_read_timeout), + ) as session: + yield session + elif isinstance(server_params, StreamableHttpServerParams): + async with streamablehttp_client(**server_params.model_dump(exclude={"type"})) 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=server_params.sse_read_timeout, + ) as session: yield session diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_streamable_http.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_streamable_http.py new file mode 100644 index 000000000..169300b81 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_streamable_http.py @@ -0,0 +1,121 @@ +from autogen_core import Component +from mcp import ClientSession, Tool +from pydantic import BaseModel +from typing_extensions import Self + +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 "autogen-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 datetime import timedelta + from autogen_ext.models.openai import OpenAIChatCompletionClient + from autogen_ext.tools.mcp import StreamableHttpMcpToolAdapter, StreamableHttpServerParams + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.ui import Console + from autogen_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=timedelta(seconds=30), + sse_read_timeout=timedelta(seconds=60 * 5), + 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 = "autogen_ext.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) diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py index fd43ff852..1dfe633c6 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py @@ -17,7 +17,7 @@ from pydantic import BaseModel from typing_extensions import Self from ._actor import McpSessionActor -from ._config import McpServerParams, SseServerParams, StdioServerParams +from ._config import McpServerParams, SseServerParams, StdioServerParams, StreamableHttpServerParams class McpWorkbenchConfig(BaseModel): @@ -252,7 +252,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]): ) return # Already initialized, no need to start again - if isinstance(self._server_params, (StdioServerParams, SseServerParams)): + if isinstance(self._server_params, (StdioServerParams, SseServerParams, StreamableHttpServerParams)): self._actor = McpSessionActor(self._server_params) await self._actor.initialize() self._actor_loop = asyncio.get_event_loop() diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py index 0a31fd128..4f7dfe11b 100644 --- a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py +++ b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py @@ -17,6 +17,8 @@ from autogen_ext.tools.mcp import ( SseServerParams, StdioMcpToolAdapter, StdioServerParams, + StreamableHttpMcpToolAdapter, + StreamableHttpServerParams, create_mcp_server_session, mcp_server_tools, ) @@ -62,6 +64,19 @@ def sample_sse_tool() -> Tool: ) +@pytest.fixture +def sample_streamable_http_tool() -> Tool: + return Tool( + name="test_streamable_http_tool", + description="A test StreamableHttp tool", + inputSchema={ + "type": "object", + "properties": {"test_param": {"type": "string"}}, + "required": ["test_param"], + }, + ) + + @pytest.fixture def mock_sse_session() -> AsyncMock: session = AsyncMock(spec=ClientSession) @@ -71,6 +86,15 @@ def mock_sse_session() -> AsyncMock: return session +@pytest.fixture +def mock_streamable_http_session() -> AsyncMock: + session = AsyncMock(spec=ClientSession) + session.initialize = AsyncMock() + session.call_tool = AsyncMock() + session.list_tools = AsyncMock() + return session + + @pytest.fixture def mock_session() -> AsyncMock: session = AsyncMock(spec=ClientSession) @@ -411,6 +435,122 @@ async def test_sse_adapter_from_server_params( ) +@pytest.mark.asyncio +async def test_streamable_http_adapter_config_serialization(sample_streamable_http_tool: Tool) -> None: + """Test that StreamableHttp adapter can be saved to and loaded from config.""" + params = StreamableHttpServerParams(url="http://test-url") + original_adapter = StreamableHttpMcpToolAdapter(server_params=params, tool=sample_streamable_http_tool) + config = original_adapter.dump_component() + loaded_adapter = StreamableHttpMcpToolAdapter.load_component(config) + + # Test that the loaded adapter has the same properties + assert loaded_adapter.name == "test_streamable_http_tool" + assert loaded_adapter.description == "A test StreamableHttp tool" + + # Verify schema structure + schema = loaded_adapter.schema + assert "parameters" in schema, "Schema must have parameters" + params_schema = schema["parameters"] + assert isinstance(params_schema, dict), "Parameters must be a dict" + assert "type" in params_schema, "Parameters must have type" + assert "required" in params_schema, "Parameters must have required fields" + assert "properties" in params_schema, "Parameters must have properties" + + # Compare schema content + assert params_schema["type"] == sample_streamable_http_tool.inputSchema["type"] + assert params_schema["required"] == sample_streamable_http_tool.inputSchema["required"] + assert ( + params_schema["properties"]["test_param"]["type"] + == sample_streamable_http_tool.inputSchema["properties"]["test_param"]["type"] + ) + + +@pytest.mark.asyncio +async def test_streamable_http_tool_execution( + sample_streamable_http_tool: Tool, + mock_streamable_http_session: AsyncMock, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test that StreamableHttp adapter properly executes tools through ClientSession.""" + params = StreamableHttpServerParams(url="http://test-url") + mock_context = AsyncMock() + mock_context.__aenter__.return_value = mock_streamable_http_session + + mock_streamable_http_session.call_tool.return_value = MagicMock( + isError=False, + content=[ + TextContent( + text="test_output", + type="text", + annotations=Annotations(audience=["user", "assistant"], priority=0.7), + ), + ], + ) + + monkeypatch.setattr( + "autogen_ext.tools.mcp._base.create_mcp_server_session", + lambda *args, **kwargs: mock_context, # type: ignore + ) + + with caplog.at_level(logging.INFO): + adapter = StreamableHttpMcpToolAdapter(server_params=params, tool=sample_streamable_http_tool) + result = await adapter.run_json( + args=schema_to_pydantic_model(sample_streamable_http_tool.inputSchema)( + **{"test_param": "test"} + ).model_dump(), + cancellation_token=CancellationToken(), + ) + + assert result == mock_streamable_http_session.call_tool.return_value.content + mock_streamable_http_session.initialize.assert_called_once() + mock_streamable_http_session.call_tool.assert_called_once() + + # Check log. + assert "test_output" in caplog.text + + +@pytest.mark.asyncio +async def test_streamable_http_adapter_from_server_params( + sample_streamable_http_tool: Tool, + mock_streamable_http_session: AsyncMock, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that StreamableHttp adapter can be created from server parameters.""" + params = StreamableHttpServerParams(url="http://test-url") + mock_context = AsyncMock() + mock_context.__aenter__.return_value = mock_streamable_http_session + monkeypatch.setattr( + "autogen_ext.tools.mcp._base.create_mcp_server_session", + lambda *args, **kwargs: mock_context, # type: ignore + ) + + mock_streamable_http_session.list_tools.return_value.tools = [sample_streamable_http_tool] + + adapter = await StreamableHttpMcpToolAdapter.from_server_params(params, "test_streamable_http_tool") + + assert isinstance(adapter, StreamableHttpMcpToolAdapter) + assert adapter.name == "test_streamable_http_tool" + assert adapter.description == "A test StreamableHttp tool" + + # Verify schema structure + schema = adapter.schema + assert "parameters" in schema, "Schema must have parameters" + params_schema = schema["parameters"] + assert isinstance(params_schema, dict), "Parameters must be a dict" + assert "type" in params_schema, "Parameters must have type" + assert "required" in params_schema, "Parameters must have required fields" + assert "properties" in params_schema, "Parameters must have properties" + + # Compare schema content + assert params_schema["type"] == sample_streamable_http_tool.inputSchema["type"] + assert params_schema["required"] == sample_streamable_http_tool.inputSchema["required"] + assert ( + params_schema["properties"]["test_param"]["type"] + == sample_streamable_http_tool.inputSchema["properties"]["test_param"]["type"] + ) + + @pytest.mark.asyncio async def test_mcp_server_fetch() -> None: params = StdioServerParams( diff --git a/python/uv.lock b/python/uv.lock index 953cbc6d9..69d8e5e82 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -771,7 +771,7 @@ requires-dist = [ { name = "markitdown", extras = ["all"], marker = "extra == 'file-surfer'", specifier = "~=0.1.0a3" }, { name = "markitdown", extras = ["all"], marker = "extra == 'magentic-one'", specifier = "~=0.1.0a3" }, { name = "markitdown", extras = ["all"], marker = "extra == 'web-surfer'", specifier = "~=0.1.0a3" }, - { name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.6.0" }, + { name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.8.1" }, { name = "nbclient", marker = "extra == 'jupyter-executor'", specifier = ">=0.10.2" }, { name = "ollama", marker = "extra == 'ollama'", specifier = ">=0.4.7" }, { name = "openai", marker = "extra == 'openai'", specifier = ">=1.66.5" }, @@ -4209,7 +4209,7 @@ wheels = [ [[package]] name = "mcp" -version = "1.6.0" +version = "1.9.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -4217,13 +4217,14 @@ dependencies = [ { name = "httpx-sse" }, { name = "pydantic" }, { name = "pydantic-settings" }, + { name = "python-multipart" }, { name = "sse-starlette" }, { name = "starlette" }, - { name = "uvicorn" }, + { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/95/d2/f587cb965a56e992634bebc8611c5b579af912b74e04eb9164bd49527d21/mcp-1.6.0.tar.gz", hash = "sha256:d9324876de2c5637369f43161cd71eebfd803df5a95e46225cab8d280e366723", size = 200031 } +sdist = { url = "https://files.pythonhosted.org/packages/ea/03/77c49cce3ace96e6787af624611b627b2828f0dca0f8df6f330a10eea51e/mcp-1.9.2.tar.gz", hash = "sha256:3c7651c053d635fd235990a12e84509fe32780cd359a5bbef352e20d4d963c05", size = 333066 } wheels = [ - { url = "https://files.pythonhosted.org/packages/10/30/20a7f33b0b884a9d14dd3aa94ff1ac9da1479fe2ad66dd9e2736075d2506/mcp-1.6.0-py3-none-any.whl", hash = "sha256:7bd24c6ea042dbec44c754f100984d186620d8b841ec30f1b19eda9b93a634d0", size = 76077 }, + { url = "https://files.pythonhosted.org/packages/5d/a6/8f5ee9da9f67c0fd8933f63d6105f02eabdac8a8c0926728368ffbb6744d/mcp-1.9.2-py3-none-any.whl", hash = "sha256:bc29f7fd67d157fef378f89a4210384f5fecf1168d0feb12d22929818723f978", size = 131083 }, ] [[package]]