Add ergonomic Client class for testing MCP servers (#1870)

src/mcp/__init__.py

@@ -1,3 +1,4 @@
+from .client.client import Client
 from .client.session import ClientSession
 from .client.session_group import ClientSessionGroup
 from .client.stdio import StdioServerParameters, stdio_client
@@ -66,6 +67,7 @@ from .types import Role as SamplingRole
 
 __all__ = [
     "CallToolRequest",
+    "Client",
     "ClientCapabilities",
     "ClientNotification",
     "ClientRequest",

src/mcp/client/__init__.py

@@ -0,0 +1,9 @@
+"""MCP Client module."""
+
+from mcp.client.client import Client
+from mcp.client.session import ClientSession
+
+__all__ = [
+    "Client",
+    "ClientSession",
+]

src/mcp/client/_memory.py

@@ -0,0 +1,97 @@
+"""In-memory transport for testing MCP servers without network overhead."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import Any
+
+import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
+from mcp.server import Server
+from mcp.server.fastmcp import FastMCP
+from mcp.shared.memory import create_client_server_memory_streams
+from mcp.shared.message import SessionMessage
+
+
+class InMemoryTransport:
+    """
+    In-memory transport for testing MCP servers without network overhead.
+
+    This transport starts the server in a background task and provides
+    streams for client-side communication. The server is automatically
+    stopped when the context manager exits.
+
+    Example:
+        server = FastMCP("test")
+        transport = InMemoryTransport(server)
+
+        async with transport.connect() as (read_stream, write_stream):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+                # Use the session...
+
+    Or more commonly, use with Client:
+        async with Client(server) as client:
+            result = await client.call_tool("my_tool", {...})
+    """
+
+    def __init__(
+        self,
+        server: Server[Any] | FastMCP,
+        *,
+        raise_exceptions: bool = False,
+    ) -> None:
+        """
+        Initialize the in-memory transport.
+
+        Args:
+            server: The MCP server to connect to (Server or FastMCP instance)
+            raise_exceptions: Whether to raise exceptions from the server
+        """
+        self._server = server
+        self._raise_exceptions = raise_exceptions
+
+    @asynccontextmanager
+    async def connect(
+        self,
+    ) -> AsyncGenerator[
+        tuple[
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
+        ],
+        None,
+    ]:
+        """
+        Connect to the server and return streams for communication.
+
+        Yields:
+            A tuple of (read_stream, write_stream) for bidirectional communication
+        """
+        # Unwrap FastMCP to get underlying Server
+        actual_server: Server[Any]
+        if isinstance(self._server, FastMCP):
+            actual_server = self._server._mcp_server  # type: ignore[reportPrivateUsage]
+        else:
+            actual_server = self._server
+
+        async with create_client_server_memory_streams() as (client_streams, server_streams):
+            client_read, client_write = client_streams
+            server_read, server_write = server_streams
+
+            async with anyio.create_task_group() as tg:
+                # Start server in background
+                tg.start_soon(
+                    lambda: actual_server.run(
+                        server_read,
+                        server_write,
+                        actual_server.create_initialization_options(),
+                        raise_exceptions=self._raise_exceptions,
+                    )
+                )
+
+                try:
+                    yield client_read, client_write
+                finally:
+                    tg.cancel_scope.cancel()

src/mcp/client/client.py

@@ -0,0 +1,302 @@
+"""Unified MCP Client that wraps ClientSession with transport management."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import AsyncExitStack
+from typing import Any
+
+from pydantic import AnyUrl
+
+import mcp.types as types
+from mcp.client._memory import InMemoryTransport
+from mcp.client.session import (
+    ClientSession,
+    ElicitationFnT,
+    ListRootsFnT,
+    LoggingFnT,
+    MessageHandlerFnT,
+    SamplingFnT,
+)
+from mcp.server import Server
+from mcp.server.fastmcp import FastMCP
+from mcp.shared.session import ProgressFnT
+
+logger = logging.getLogger(__name__)
+
+
+class Client:
+    """A high-level MCP client for connecting to MCP servers.
+
+    Currently supports in-memory transport for testing. Pass a Server or
+    FastMCP instance directly to the constructor.
+
+    Example:
+        ```python
+        from mcp.client import Client
+        from mcp.server.fastmcp import FastMCP
+
+        server = FastMCP("test")
+
+        @server.tool()
+        def add(a: int, b: int) -> int:
+            return a + b
+
+        async with Client(server) as client:
+            result = await client.call_tool("add", {"a": 1, "b": 2})
+        ```
+    """
+
+    # TODO(felixweinberger): Expand to support all transport types (like FastMCP 2):
+    # - Add ClientTransport base class with connect_session() method
+    # - Add StreamableHttpTransport, SSETransport, StdioTransport
+    # - Add infer_transport() to auto-detect transport from input type
+    # - Accept URL strings, Path objects, config dicts in constructor
+    # - Add auth support (OAuth, bearer tokens)
+
+    def __init__(
+        self,
+        server: Server[Any] | FastMCP,
+        *,
+        raise_exceptions: bool = False,
+        read_timeout_seconds: float | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        list_roots_callback: ListRootsFnT | None = None,
+        logging_callback: LoggingFnT | None = None,
+        message_handler: MessageHandlerFnT | None = None,
+        client_info: types.Implementation | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
+    ) -> None:
+        """
+        Initialize the client with a server.
+
+        Args:
+            server: The MCP server to connect to (Server or FastMCP instance)
+            raise_exceptions: Whether to raise exceptions from the server
+            read_timeout_seconds: Timeout for read operations
+            sampling_callback: Callback for handling sampling requests
+            list_roots_callback: Callback for handling list roots requests
+            logging_callback: Callback for handling logging notifications
+            message_handler: Callback for handling raw messages
+            client_info: Client implementation info to send to server
+            elicitation_callback: Callback for handling elicitation requests
+        """
+        self._server = server
+        self._raise_exceptions = raise_exceptions
+        self._read_timeout_seconds = read_timeout_seconds
+        self._sampling_callback = sampling_callback
+        self._list_roots_callback = list_roots_callback
+        self._logging_callback = logging_callback
+        self._message_handler = message_handler
+        self._client_info = client_info
+        self._elicitation_callback = elicitation_callback
+
+        self._session: ClientSession | None = None
+        self._exit_stack: AsyncExitStack | None = None
+
+    async def __aenter__(self) -> Client:
+        """Enter the async context manager."""
+        if self._session is not None:
+            raise RuntimeError("Client is already entered; cannot reenter")
+
+        async with AsyncExitStack() as exit_stack:
+            # Create transport and connect
+            transport = InMemoryTransport(self._server, raise_exceptions=self._raise_exceptions)
+            read_stream, write_stream = await exit_stack.enter_async_context(transport.connect())
+
+            # Create session
+            self._session = await exit_stack.enter_async_context(
+                ClientSession(
+                    read_stream=read_stream,
+                    write_stream=write_stream,
+                    read_timeout_seconds=self._read_timeout_seconds,
+                    sampling_callback=self._sampling_callback,
+                    list_roots_callback=self._list_roots_callback,
+                    logging_callback=self._logging_callback,
+                    message_handler=self._message_handler,
+                    client_info=self._client_info,
+                    elicitation_callback=self._elicitation_callback,
+                )
+            )
+
+            # Initialize the session
+            await self._session.initialize()
+
+            # Transfer ownership to self for __aexit__ to handle
+            self._exit_stack = exit_stack.pop_all()
+            return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any,
+    ) -> None:
+        """Exit the async context manager."""
+        if self._exit_stack:
+            await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+        self._session = None
+
+    @property
+    def session(self) -> ClientSession:
+        """
+        Get the underlying ClientSession.
+
+        This provides access to the full ClientSession API for advanced use cases.
+
+        Raises:
+            RuntimeError: If accessed before entering the context manager.
+        """
+        if self._session is None:
+            raise RuntimeError("Client must be used within an async context manager")
+        return self._session
+
+    @property
+    def server_capabilities(self) -> types.ServerCapabilities | None:
+        """The server capabilities received during initialization, or None if not yet initialized."""
+        return self.session.get_server_capabilities()
+
+    async def send_ping(self) -> types.EmptyResult:
+        """Send a ping request to the server."""
+        return await self.session.send_ping()
+
+    async def send_progress_notification(
+        self,
+        progress_token: str | int,
+        progress: float,
+        total: float | None = None,
+        message: str | None = None,
+    ) -> None:
+        """Send a progress notification to the server."""
+        await self.session.send_progress_notification(
+            progress_token=progress_token,
+            progress=progress,
+            total=total,
+            message=message,
+        )
+
+    async def set_logging_level(self, level: types.LoggingLevel) -> types.EmptyResult:
+        """Set the logging level on the server."""
+        return await self.session.set_logging_level(level)
+
+    async def list_resources(
+        self,
+        params: types.PaginatedRequestParams | None = None,
+    ) -> types.ListResourcesResult:
+        """List available resources from the server."""
+        return await self.session.list_resources(params=params)
+
+    async def list_resource_templates(
+        self,
+        params: types.PaginatedRequestParams | None = None,
+    ) -> types.ListResourceTemplatesResult:
+        """List available resource templates from the server."""
+        return await self.session.list_resource_templates(params=params)
+
+    async def read_resource(self, uri: str | AnyUrl) -> types.ReadResourceResult:
+        """
+        Read a resource from the server.
+
+        Args:
+            uri: The URI of the resource to read
+
+        Returns:
+            The resource content
+        """
+        return await self.session.read_resource(uri)
+
+    async def subscribe_resource(self, uri: str | AnyUrl) -> types.EmptyResult:
+        """Subscribe to resource updates."""
+        return await self.session.subscribe_resource(uri)
+
+    async def unsubscribe_resource(self, uri: str | AnyUrl) -> types.EmptyResult:
+        """Unsubscribe from resource updates."""
+        return await self.session.unsubscribe_resource(uri)
+
+    async def call_tool(
+        self,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        read_timeout_seconds: float | None = None,
+        progress_callback: ProgressFnT | None = None,
+        *,
+        meta: dict[str, Any] | None = None,
+    ) -> types.CallToolResult:
+        """
+        Call a tool on the server.
+
+        Args:
+            name: The name of the tool to call
+            arguments: Arguments to pass to the tool
+            read_timeout_seconds: Timeout for the tool call
+            progress_callback: Callback for progress updates
+            meta: Additional metadata for the request
+
+        Returns:
+            The tool result
+        """
+        return await self.session.call_tool(
+            name=name,
+            arguments=arguments,
+            read_timeout_seconds=read_timeout_seconds,
+            progress_callback=progress_callback,
+            meta=meta,
+        )
+
+    async def list_prompts(
+        self,
+        params: types.PaginatedRequestParams | None = None,
+    ) -> types.ListPromptsResult:
+        """List available prompts from the server."""
+        return await self.session.list_prompts(params=params)
+
+    async def get_prompt(
+        self,
+        name: str,
+        arguments: dict[str, str] | None = None,
+    ) -> types.GetPromptResult:
+        """
+        Get a prompt from the server.
+
+        Args:
+            name: The name of the prompt
+            arguments: Arguments to pass to the prompt
+
+        Returns:
+            The prompt content
+        """
+        return await self.session.get_prompt(name=name, arguments=arguments)
+
+    async def complete(
+        self,
+        ref: types.ResourceTemplateReference | types.PromptReference,
+        argument: dict[str, str],
+        context_arguments: dict[str, str] | None = None,
+    ) -> types.CompleteResult:
+        """
+        Get completions for a prompt or resource template argument.
+
+        Args:
+            ref: Reference to the prompt or resource template
+            argument: The argument to complete
+            context_arguments: Additional context arguments
+
+        Returns:
+            Completion suggestions
+        """
+        return await self.session.complete(
+            ref=ref,
+            argument=argument,
+            context_arguments=context_arguments,
+        )
+
+    async def list_tools(
+        self,
+        params: types.PaginatedRequestParams | None = None,
+    ) -> types.ListToolsResult:
+        """List available tools from the server."""
+        return await self.session.list_tools(params=params)
+
+    async def send_roots_list_changed(self) -> None:
+        """Send a notification that the roots list has changed."""
+        await self.session.send_roots_list_changed()

src/mcp/client/session.py

@@ -407,7 +407,6 @@ class ClientSession(
         """Send a tools/list request.
 
         Args:
-            cursor: Simple cursor string for pagination (deprecated, use params instead)
             params: Full pagination parameters including cursor and any future fields
         """
         result = await self.send_request(

src/mcp/shared/memory.py

@@ -6,15 +6,10 @@ from __future__ import annotations
 
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import Any
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
-import mcp.types as types
-from mcp.client.session import ClientSession, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
-from mcp.server import Server
-from mcp.server.fastmcp import FastMCP
 from mcp.shared.message import SessionMessage
 
 MessageStream = tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
@@ -43,55 +38,3 @@ async def create_client_server_memory_streams() -> AsyncGenerator[tuple[MessageS
         server_to_client_send,
     ):
         yield client_streams, server_streams
-
-
-@asynccontextmanager
-async def create_connected_server_and_client_session(
-    server: Server[Any] | FastMCP,
-    read_timeout_seconds: float | None = None,
-    sampling_callback: SamplingFnT | None = None,
-    list_roots_callback: ListRootsFnT | None = None,
-    logging_callback: LoggingFnT | None = None,
-    message_handler: MessageHandlerFnT | None = None,
-    client_info: types.Implementation | None = None,
-    raise_exceptions: bool = False,
-    elicitation_callback: ElicitationFnT | None = None,
-) -> AsyncGenerator[ClientSession, None]:
-    """Creates a ClientSession that is connected to a running MCP server."""
-
-    # TODO(Marcelo): we should have a proper `Client` that can use this "in-memory transport",
-    # and we should expose a method in the `FastMCP` so we don't access a private attribute.
-    if isinstance(server, FastMCP):  # pragma: no cover
-        server = server._mcp_server  # type: ignore[reportPrivateUsage]
-
-    async with create_client_server_memory_streams() as (client_streams, server_streams):
-        client_read, client_write = client_streams
-        server_read, server_write = server_streams
-
-        # Create a cancel scope for the server task
-        async with anyio.create_task_group() as tg:
-            tg.start_soon(
-                lambda: server.run(
-                    server_read,
-                    server_write,
-                    server.create_initialization_options(),
-                    raise_exceptions=raise_exceptions,
-                )
-            )
-
-            try:
-                async with ClientSession(
-                    read_stream=client_read,
-                    write_stream=client_write,
-                    read_timeout_seconds=read_timeout_seconds,
-                    sampling_callback=sampling_callback,
-                    list_roots_callback=list_roots_callback,
-                    logging_callback=logging_callback,
-                    message_handler=message_handler,
-                    client_info=client_info,
-                    elicitation_callback=elicitation_callback,
-                ) as client_session:
-                    await client_session.initialize()
-                    yield client_session
-            finally:  # pragma: no cover
-                tg.cancel_scope.cancel()