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

Support different transports in Client (#1972)

Browse source
This commit is contained in:
Marcelo Trylesinski 2026-01-30 13:11:27 +01:00 • committed by GitHub
commit 21822053df
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 544 additions and 631 deletions
Download patch file Download diff file Expand all files Collapse all files

3
src/mcp/client/__init__.py
View file

@ -1,6 +1,7 @@
"""MCP Client module."""
from mcp.client._transport import Transport
from mcp.client.client import Client
from mcp.client.session import ClientSession
__all__ = ["Client", "ClientSession"]
__all__ = ["Client", "ClientSession", "Transport"]

55
src/mcp/client/_memory.py
View file

@ -2,17 +2,17 @@
from __future__ import annotations
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from collections.abc import AsyncIterator
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from types import TracebackType
from typing import Any
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.client._transport import TransportStreams
from mcp.server import Server
from mcp.server.mcpserver import MCPServer
from mcp.shared.memory import create_client_server_memory_streams
from mcp.shared.message import SessionMessage
class InMemoryTransport:
@ -21,19 +21,6 @@ class InMemoryTransport:
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 = MCPServer("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] | MCPServer, *, raise_exceptions: bool = False) -> None:
@ -45,26 +32,15 @@ class InMemoryTransport:
"""
self._server = server
self._raise_exceptions = raise_exceptions
self._cm: AbstractAsyncContextManager[TransportStreams] | None = None
@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
"""
async def _connect(self) -> AsyncIterator[TransportStreams]:
"""Connect to the server and yield streams for communication."""
# Unwrap MCPServer to get underlying Server
actual_server: Server[Any]
if isinstance(self._server, MCPServer):
actual_server = self._server._lowlevel_server # type: ignore[reportPrivateUsage]
# TODO(Marcelo): Make `lowlevel_server` public.
actual_server: Server[Any] = self._server._lowlevel_server # type: ignore[reportPrivateUsage]
else:
actual_server = self._server
@ -87,3 +63,16 @@ class InMemoryTransport:
yield client_read, client_write
finally:
tg.cancel_scope.cancel()
async def __aenter__(self) -> TransportStreams:
"""Connect to the server and return streams for communication."""
self._cm = self._connect()
return await self._cm.__aenter__()
async def __aexit__(
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
) -> None:
"""Close the transport and stop the server."""
if self._cm is not None: # pragma: no branch
await self._cm.__aexit__(exc_type, exc_val, exc_tb)
self._cm = None

20
src/mcp/client/_transport.py Normal file
View file

@ -0,0 +1,20 @@
"""Transport protocol for MCP clients."""
from __future__ import annotations
from contextlib import AbstractAsyncContextManager
from typing import Protocol
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.shared.message import SessionMessage
TransportStreams = tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
class Transport(AbstractAsyncContextManager[TransportStreams], Protocol):
"""Protocol for MCP transports.
A transport is an async context manager that yields read and write streams
for bidirectional communication with an MCP server.
"""

111
src/mcp/client/client.py
View file

@ -3,10 +3,13 @@
from __future__ import annotations
from contextlib import AsyncExitStack
from dataclasses import KW_ONLY, dataclass, field
from typing import Any
from mcp.client._memory import InMemoryTransport
from mcp.client._transport import Transport
from mcp.client.session import ClientSession, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
from mcp.client.streamable_http import streamable_http_client
from mcp.server import Server
from mcp.server.mcpserver import MCPServer
from mcp.shared.session import ProgressFnT
@ -30,6 +33,7 @@ from mcp.types import (
)
@dataclass
class Client:
"""A high-level MCP client for connecting to MCP servers.
@ -55,52 +59,53 @@ class Client:
```
"""
# TODO(felixweinberger): Expand to support all transport types:
# - 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)
server: Server[Any] | MCPServer | Transport | str
"""The MCP server to connect to.
def __init__(
self,
server: Server[Any] | MCPServer,
*,
# TODO(Marcelo): When do `raise_exceptions=True` actually raises?
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: Implementation | None = None,
elicitation_callback: ElicitationFnT | None = None,
) -> None:
"""Initialize the client with a server.
If the server is a `Server` or `MCPServer` instance, it will be wrapped in an `InMemoryTransport`.
If the server is a URL string, it will be used as the URL for a `streamable_http_client` transport.
If the server is a `Transport` instance, it will be used directly.
"""
Args:
server: The MCP server to connect to (Server or MCPServer 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
_: KW_ONLY
self._session: ClientSession | None = None
self._exit_stack: AsyncExitStack | None = None
# TODO(Marcelo): When do `raise_exceptions=True` actually raises?
raise_exceptions: bool = False
"""Whether to raise exceptions from the server."""
read_timeout_seconds: float | None = None
"""Timeout for read operations."""
sampling_callback: SamplingFnT | None = None
"""Callback for handling sampling requests."""
list_roots_callback: ListRootsFnT | None = None
"""Callback for handling list roots requests."""
logging_callback: LoggingFnT | None = None
"""Callback for handling logging notifications."""
# TODO(Marcelo): Why do we have both "callback" and "handler"?
message_handler: MessageHandlerFnT | None = None
"""Callback for handling raw messages."""
client_info: Implementation | None = None
"""Client implementation info to send to server."""
elicitation_callback: ElicitationFnT | None = None
"""Callback for handling elicitation requests."""
_session: ClientSession | None = field(init=False, default=None)
_exit_stack: AsyncExitStack | None = field(init=False, default=None)
_transport: Transport = field(init=False)
def __post_init__(self) -> None:
if isinstance(self.server, Server | MCPServer):
self._transport = InMemoryTransport(self.server, raise_exceptions=self.raise_exceptions)
elif isinstance(self.server, str):
self._transport = streamable_http_client(self.server)
else:
self._transport = self.server
async def __aenter__(self) -> Client:
"""Enter the async context manager."""
@ -108,26 +113,22 @@ class Client:
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())
read_stream, write_stream = await exit_stack.enter_async_context(self._transport)
# 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,
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

2
src/mcp/client/session_group.py
View file

@ -296,7 +296,7 @@ class ClientSessionGroup:
http_client=httpx_client,
terminate_on_close=server_params.terminate_on_close,
)
read, write, _ = await session_stack.enter_async_context(client)
read, write = await session_stack.enter_async_context(client)
session = await session_stack.enter_async_context(
mcp.ClientSession(

25
src/mcp/client/streamable_http.py
View file

@ -14,6 +14,7 @@ from anyio.abc import TaskGroup
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
from mcp.client._transport import TransportStreams
from mcp.shared._httpx_utils import create_mcp_http_client
from mcp.shared.message import ClientMessageMetadata, SessionMessage
from mcp.types import (
@ -31,10 +32,10 @@ from mcp.types import (
logger = logging.getLogger(__name__)
# TODO(Marcelo): Put the TransportStreams in a module under shared, so we can import here.
SessionMessageOrError = SessionMessage | Exception
StreamWriter = MemoryObjectSendStream[SessionMessageOrError]
StreamReader = MemoryObjectReceiveStream[SessionMessage]
GetSessionIdCallback = Callable[[], str | None]
MCP_SESSION_ID = "mcp-session-id"
MCP_PROTOCOL_VERSION = "mcp-protocol-version"
@ -200,8 +201,8 @@ class StreamableHTTPTransport:
# Stream ended normally (server closed) - reset attempt counter
attempt = 0
except Exception as exc: # pragma: lax no cover
logger.debug(f"GET stream error: {exc}")
except Exception: # pragma: lax no cover
logger.debug("GET stream error", exc_info=True)
attempt += 1
if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover
@ -488,25 +489,22 @@ class StreamableHTTPTransport:
except Exception as exc: # pragma: no cover
logger.warning(f"Session termination failed: {exc}")
# TODO(Marcelo): Check the TODO below, and cover this with tests if necessary.
def get_session_id(self) -> str | None:
"""Get the current session ID."""
return self.session_id
return self.session_id # pragma: no cover
# TODO(Marcelo): I've dropped the `get_session_id` callback because it breaks the Transport protocol. Is that needed?
# It's a completely wrong abstraction, so removal is a good idea. But if we need the client to find the session ID,
# we should think about a better way to do it. I believe we can achieve it with other means.
@asynccontextmanager
async def streamable_http_client(
url: str,
*,
http_client: httpx.AsyncClient | None = None,
terminate_on_close: bool = True,
) -> AsyncGenerator[
tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
GetSessionIdCallback,
],
None,
]:
) -> AsyncGenerator[TransportStreams, None]:
"""Client transport for StreamableHTTP.
Args:
@ -520,7 +518,6 @@ async def streamable_http_client(
Tuple containing:
- read_stream: Stream for reading messages from the server
- write_stream: Stream for sending messages to the server
- get_session_id_callback: Function to retrieve the current session ID
Example:
See examples/snippets/clients/ for usage patterns.
@ -561,7 +558,7 @@ async def streamable_http_client(
)
try:
yield (read_stream, write_stream, transport.get_session_id)
yield read_stream, write_stream
finally:
if transport.session_id and terminate_on_close:
await transport.terminate_session(client)

4
src/mcp/shared/_httpx_utils.py
View file

@ -66,9 +66,7 @@ def create_mcp_http_client(
response = await client.get("/protected-endpoint")
"""
# Set MCP defaults
kwargs: dict[str, Any] = {
"follow_redirects": True,
}
kwargs: dict[str, Any] = {"follow_redirects": True}
# Handle timeout
if timeout is None: