modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
This commit is contained in:
Felix Weinberger 2026-01-16 15:49:26 +00:00 • committed by GitHub
commit df039bf97c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 1298 additions and 526 deletions
Download patch file Download diff file Expand all files Collapse all files

57
src/mcp/shared/memory.py
View file

@ -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()