Make shared session possible for MCP tool

python/packages/autogen-core/src/autogen_core/utils/_json_to_pydantic.py

@@ -229,11 +229,14 @@
                 item_type = self.get_ref(item_schema["$ref"].split("/")[-1])
             else:
                 item_type_name = item_schema.get("type")
-                if item_type_name not in TYPE_MAPPING:
+                if item_type_name is None:
+                    item_type = List[str]
+                elif item_type_name not in TYPE_MAPPING:
                     raise UnsupportedKeywordError(
                         f"Unsupported or missing item type `{item_type_name}` for array field `{key}` in `{model_name}`"
                     )
-                item_type = TYPE_MAPPING[item_type_name]
+                else:
+                    item_type = TYPE_MAPPING[item_type_name]
 
             base_type = conlist(item_type, **constraints) if constraints else List[item_type]  # type: ignore[valid-type]
 

python/packages/autogen-ext/pyproject.toml

@@ -135,10 +135,7 @@ semantic-kernel-all = [
 
 rich = ["rich>=13.9.4"]
 
-mcp = [
-    "mcp>=1.6.0",
-    "json-schema-to-pydantic>=0.2.4"
-]
+mcp = ["mcp>=1.6.0"]
 
 [tool.hatch.build.targets.wheel]
 packages = ["src/autogen_ext"]

python/packages/autogen-ext/src/autogen_ext/experimental/task_centric_memory/__init__.py

@@ -1,4 +1,4 @@
-from .memory_controller import MemoryController, MemoryControllerConfig
 from ._memory_bank import MemoryBankConfig
+from .memory_controller import MemoryController, MemoryControllerConfig
 
 __all__ = ["MemoryController", "MemoryControllerConfig", "MemoryBankConfig"]

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

@@ -1,9 +1,11 @@
 from ._config import McpServerParams, SseServerParams, StdioServerParams
 from ._factory import mcp_server_tools
+from ._session import create_mcp_server_session
 from ._sse import SseMcpToolAdapter
 from ._stdio import StdioMcpToolAdapter
 
 __all__ = [
+    "create_mcp_server_session",
     "StdioMcpToolAdapter",
     "StdioServerParams",
     "SseMcpToolAdapter",

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

@@ -1,12 +1,12 @@
 import asyncio
 import builtins
 from abc import ABC
-from typing import Any, Generic, Type, TypeVar
+from typing import Any, Dict, Generic, Type, TypeVar
 
 from autogen_core import CancellationToken
 from autogen_core.tools import BaseTool
-from json_schema_to_pydantic import create_model
-from mcp import Tool
+from autogen_core.utils import schema_to_pydantic_model
+from mcp import ClientSession, Tool
 from pydantic import BaseModel
 
 from ._config import McpServerParams
@@ -26,16 +26,17 @@
 
     component_type = "tool"
 
-    def __init__(self, server_params: TServerParams, tool: Tool) -> None:
+    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 = create_model(tool.inputSchema, allow_undefined_array_items=True)
+        input_model = schema_to_pydantic_model(tool.inputSchema)
 
         # Use Any as return type since MCP tool returns can vary
         return_type: Type[Any] = object
@@ -61,20 +62,27 @@
         # for many servers.
         kwargs = args.model_dump(exclude_unset=True)
 
-        try:
-            async with create_mcp_server_session(self._server_params) as session:
-                await session.initialize()
+        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)
 
-                if cancellation_token.is_cancelled():
-                    raise Exception("Operation cancelled")
+    async def _run(self, args: Dict[str, Any], cancellation_token: CancellationToken, session: ClientSession) -> Any:
+        try:
+            if cancellation_token.is_cancelled():
+                raise Exception("Operation cancelled")
 
-                result_future = asyncio.ensure_future(session.call_tool(name=self._tool.name, arguments=kwargs))
-                cancellation_token.link_future(result_future)
-                result = await result_future
+            result_future = asyncio.ensure_future(session.call_tool(name=self._tool.name, arguments=args))
+            cancellation_token.link_future(result_future)
+            result = await result_future
 
-                if result.isError:
-                    raise Exception(f"MCP tool execution failed: {result.content}")
-                return result.content
+            if result.isError:
+                raise Exception(f"MCP tool execution failed: {result.content}")
+            return result.content
         except Exception as e:
             error_message = self._format_errors(e)
             raise Exception(error_message) from e

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

@@ -1,3 +1,5 @@
+from mcp import ClientSession
+
 from ._config import McpServerParams, SseServerParams, StdioServerParams
 from ._session import create_mcp_server_session
 from ._sse import SseMcpToolAdapter
@@ -6,6 +8,7 @@
 
 async def mcp_server_tools(
     server_params: McpServerParams,
+    session: ClientSession | None = None,
 ) -> list[StdioMcpToolAdapter | SseMcpToolAdapter]:
     """Creates a list of MCP tool adapters that can be used with AutoGen agents.
 
@@ -24,6 +27,9 @@
         server_params (McpServerParams): Connection parameters for the MCP server.
             Can be either StdioServerParams for command-line tools or
             SseServerParams 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
@@ -110,6 +116,58 @@
 
             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 autogen_agentchat.agents import AssistantAgent
+            from autogen_agentchat.conditions import TextMentionTermination
+            from autogen_agentchat.teams import RoundRobinGroupChat
+            from autogen_agentchat.ui import Console
+            from autogen_ext.models.openai import OpenAIChatCompletionClient
+            from autogen_ext.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:**
 
@@ -130,13 +188,16 @@
 
     For more examples and detailed usage, see the samples directory in the package repository.
     """
-    async with create_mcp_server_session(server_params) as session:
-        await session.initialize()
+    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) for tool in tools.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) for tool in tools.tools]
+        return [SseMcpToolAdapter(server_params=server_params, tool=tool, session=session) for tool in tools.tools]
     raise ValueError(f"Unsupported server params type: {type(server_params)}")

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

@@ -1,5 +1,5 @@
 from autogen_core import Component
-from mcp import Tool
+from mcp import ClientSession, Tool
 from pydantic import BaseModel
 from typing_extensions import Self
 
@@ -35,8 +35,11 @@ class SseMcpToolAdapter(
 
     Args:
         server_params (SseServerParameters): Parameters for the MCP server connection,
-            including URL, headers, and timeouts
-        tool (Tool): The MCP tool to wrap
+            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
@@ -86,8 +89,8 @@ async def main() -> None:
     component_config_schema = SseMcpToolAdapterConfig
     component_provider_override = "autogen_ext.tools.mcp.SseMcpToolAdapter"
 
-    def __init__(self, server_params: SseServerParams, tool: Tool) -> None:
-        super().__init__(server_params=server_params, tool=tool)
+    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:
         """

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

@@ -1,5 +1,5 @@
 from autogen_core import Component
-from mcp import Tool
+from mcp import ClientSession, Tool
 from pydantic import BaseModel
 from typing_extensions import Self
 
@@ -37,15 +37,18 @@ class StdioMcpToolAdapter(
         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:`~autogen_ext.tools.mcp.mcp_server_tools` for examples.
     """
 
     component_config_schema = StdioMcpToolAdapterConfig
     component_provider_override = "autogen_ext.tools.mcp.StdioMcpToolAdapter"
 
-    def __init__(self, server_params: StdioServerParams, tool: Tool) -> None:
-        super().__init__(server_params=server_params, tool=tool)
+    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:
         """