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

Propagate contextvars.Context through anyio streams without modifying SessionMessage (#2298)

Browse source
This commit is contained in:
Marcelo Trylesinski 2026-03-31 12:49:38 -04:00 • committed by GitHub
commit e6235d1667
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 310 additions and 94 deletions
Download patch file Download diff file Expand all files Collapse all files

6
examples/clients/simple-auth-client/mcp_simple_auth_client/main.py
View file

@ -18,7 +18,7 @@ from typing import Any
from urllib.parse import parse_qs, urlparse
import httpx
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.client._transport import ReadStream, WriteStream
from mcp.client.auth import OAuthClientProvider, TokenStorage
from mcp.client.session import ClientSession
from mcp.client.sse import sse_client
@ -241,8 +241,8 @@ class SimpleAuthClient:
async def _run_session(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
):
"""Run the MCP session with the given streams."""
print("🤝 Initializing MCP session...")

5
pyproject.toml
View file

@ -219,13 +219,10 @@ skip_covered = true
show_missing = true
ignore_errors = true
precision = 2
exclude_lines = [
"pragma: no cover",
exclude_also = [
"pragma: lax no cover",
"if TYPE_CHECKING:",
"@overload",
"raise NotImplementedError",
"^\\s*\\.\\.\\.\\s*$",
]
# https://coverage.readthedocs.io/en/latest/config.html#paths

6
src/mcp/client/__main__.py
View file

@ -6,9 +6,9 @@ from functools import partial
from urllib.parse import urlparse
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp import types
from mcp.client._transport import ReadStream, WriteStream
from mcp.client.session import ClientSession
from mcp.client.sse import sse_client
from mcp.client.stdio import StdioServerParameters, stdio_client
@ -33,8 +33,8 @@ async def message_handler(
async def run_session(
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
client_info: types.Implementation | None = None,
):
async with ClientSession(

7
src/mcp/client/_transport.py
View file

@ -5,11 +5,12 @@ from __future__ import annotations
from contextlib import AbstractAsyncContextManager
from typing import Protocol
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.message import SessionMessage
TransportStreams = tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
__all__ = ["ReadStream", "WriteStream", "Transport", "TransportStreams"]
TransportStreams = tuple[ReadStream[SessionMessage | Exception], WriteStream[SessionMessage]]
class Transport(AbstractAsyncContextManager[TransportStreams], Protocol):

6
src/mcp/client/session.py
View file

@ -4,10 +4,10 @@ import logging
from typing import Any, Protocol
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import TypeAdapter
from mcp import types
from mcp.client._transport import ReadStream, WriteStream
from mcp.client.experimental import ExperimentalClientFeatures
from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
from mcp.shared._context import RequestContext
@ -109,8 +109,8 @@ class ClientSession(
):
def __init__(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
read_timeout_seconds: float | None = None,
sampling_callback: SamplingFnT | None = None,
elicitation_callback: ElicitationFnT | None = None,

23
src/mcp/client/sse.py
View file

@ -7,11 +7,11 @@ from urllib.parse import parse_qs, urljoin, urlparse
import anyio
import httpx
from anyio.abc import TaskStatus
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from httpx_sse import aconnect_sse
from httpx_sse._exceptions import SSEError
from mcp import types
from mcp.shared._context_streams import create_context_streams
from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
from mcp.shared.message import SessionMessage
@ -51,12 +51,6 @@ async def sse_client(
auth: Optional HTTPX authentication handler.
on_session_created: Optional callback invoked with the session ID when received.
"""
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}")
async with httpx_client_factory(
headers=headers, auth=auth, timeout=httpx.Timeout(timeout, read=sse_read_timeout)
@ -65,8 +59,8 @@ async def sse_client(
event_source.response.raise_for_status()
logger.debug("SSE connection established")
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
try:
@ -124,7 +118,8 @@ async def sse_client(
async def post_writer(endpoint_url: str):
try:
async with write_stream_reader, write_stream:
async for session_message in write_stream_reader:
async def _send_message(session_message: SessionMessage) -> None:
logger.debug(f"Sending client message: {session_message}")
response = await client.post(
endpoint_url,
@ -136,6 +131,14 @@ async def sse_client(
)
response.raise_for_status()
logger.debug(f"Client message sent successfully: {response.status_code}")
async for session_message in write_stream_reader:
sender_ctx = write_stream_reader.last_context
if sender_ctx is not None:
async with anyio.create_task_group() as tg:
sender_ctx.run(tg.start_soon, _send_message, session_message)
else:
await _send_message(session_message) # pragma: no cover
except Exception: # pragma: lax no cover
logger.exception("Error in post_writer")

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

@ -11,11 +11,11 @@ from dataclasses import dataclass
import anyio
import httpx
from anyio.abc import TaskGroup
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
from pydantic import ValidationError
from mcp.client._transport import TransportStreams
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
from mcp.shared._httpx_utils import create_mcp_http_client
from mcp.shared.message import ClientMessageMetadata, SessionMessage
from mcp.types import (
@ -38,8 +38,8 @@ 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]
StreamWriter = ContextSendStream[SessionMessageOrError]
StreamReader = ContextReceiveStream[SessionMessage]
MCP_SESSION_ID = "mcp-session-id"
MCP_PROTOCOL_VERSION = "mcp-protocol-version"
@ -434,14 +434,15 @@ class StreamableHTTPTransport:
client: httpx.AsyncClient,
write_stream_reader: StreamReader,
read_stream_writer: StreamWriter,
write_stream: MemoryObjectSendStream[SessionMessage],
write_stream: ContextSendStream[SessionMessage],
start_get_stream: Callable[[], None],
tg: TaskGroup,
) -> None:
"""Handle writing requests to the server."""
try:
async with write_stream_reader, read_stream_writer, write_stream:
async for session_message in write_stream_reader:
async def _handle_message(session_message: SessionMessage) -> None:
message = session_message.message
metadata = (
session_message.metadata
@ -478,6 +479,14 @@ class StreamableHTTPTransport:
else:
await handle_request_async()
async for session_message in write_stream_reader:
sender_ctx = write_stream_reader.last_context
if sender_ctx is not None:
async with anyio.create_task_group() as tg_local:
sender_ctx.run(tg_local.start_soon, _handle_message, session_message)
else:
await _handle_message(session_message) # pragma: no cover
except Exception: # pragma: lax no cover
logger.exception("Error in post_writer")
@ -547,8 +556,8 @@ async def streamable_http_client(
if not client_provided:
await stack.enter_async_context(client)
read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
async with (
read_stream_writer,

15
src/mcp/server/lowlevel/server.py
View file

@ -36,6 +36,7 @@ handler callables by method string.
from __future__ import annotations
import contextvars
import logging
import warnings
from collections.abc import AsyncIterator, Awaitable, Callable
@ -44,7 +45,6 @@ from importlib.metadata import version as importlib_version
from typing import Any, Generic
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.authentication import AuthenticationMiddleware
@ -65,6 +65,7 @@ from mcp.server.session import ServerSession
from mcp.server.streamable_http import EventStore
from mcp.server.streamable_http_manager import StreamableHTTPASGIApp, StreamableHTTPSessionManager
from mcp.server.transport_security import TransportSecuritySettings
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.exceptions import MCPError
from mcp.shared.message import ServerMessageMetadata, SessionMessage
from mcp.shared.session import RequestResponder
@ -355,8 +356,8 @@ class Server(Generic[LifespanResultT]):
async def run(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
initialization_options: InitializationOptions,
# When False, exceptions are returned as messages to the client.
# When True, exceptions are raised, which will cause the server to shut down
@ -391,7 +392,13 @@ class Server(Generic[LifespanResultT]):
async for message in session.incoming_messages:
logger.debug("Received message: %s", message)
tg.start_soon(
if isinstance(message, RequestResponder) and message.context is not None:
context = message.context
else:
context = contextvars.copy_context()
context.run(
tg.start_soon,
self._handle_message,
message,
session,

7
src/mcp/server/session.py
View file

@ -33,13 +33,14 @@ from typing import Any, TypeVar, overload
import anyio
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from anyio.streams.memory import MemoryObjectReceiveStream
from pydantic import AnyUrl, TypeAdapter
from mcp import types
from mcp.server.experimental.session_features import ExperimentalServerSessionFeatures
from mcp.server.models import InitializationOptions
from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.exceptions import StatelessModeNotSupported
from mcp.shared.experimental.tasks.capabilities import check_tasks_capability
from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY
@ -79,8 +80,8 @@ class ServerSession(
def __init__(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
init_options: InitializationOptions,
stateless: bool = False,
) -> None:

13
src/mcp/server/sse.py
View file

@ -43,7 +43,6 @@ from urllib.parse import quote
from uuid import UUID, uuid4
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import ValidationError
from sse_starlette import EventSourceResponse
from starlette.requests import Request
@ -55,6 +54,7 @@ from mcp.server.transport_security import (
TransportSecurityMiddleware,
TransportSecuritySettings,
)
from mcp.shared._context_streams import ContextSendStream, create_context_streams
from mcp.shared.message import ServerMessageMetadata, SessionMessage
logger = logging.getLogger(__name__)
@ -72,7 +72,7 @@ class SseServerTransport:
"""
_endpoint: str
_read_stream_writers: dict[UUID, MemoryObjectSendStream[SessionMessage | Exception]]
_read_stream_writers: dict[UUID, ContextSendStream[SessionMessage | Exception]]
_security: TransportSecurityMiddleware
def __init__(self, endpoint: str, security_settings: TransportSecuritySettings | None = None) -> None:
@ -129,14 +129,9 @@ class SseServerTransport:
raise ValueError("Request validation failed")
logger.debug("Setting up SSE connection")
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
session_id = uuid4()
self._read_stream_writers[session_id] = read_stream_writer

12
src/mcp/server/stdio.py
View file

@ -23,9 +23,9 @@ from io import TextIOWrapper
import anyio
import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp import types
from mcp.shared._context_streams import create_context_streams
from mcp.shared.message import SessionMessage
@ -43,14 +43,8 @@ async def stdio_server(stdin: anyio.AsyncFile[str] | None = None, stdout: anyio.
if not stdout:
stdout = anyio.wrap_file(TextIOWrapper(sys.stdout.buffer, encoding="utf-8"))
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
async def stdin_reader():
try:

18
src/mcp/server/streamable_http.py
View file

@ -25,6 +25,8 @@ from starlette.responses import Response
from starlette.types import Receive, Scope, Send
from mcp.server.transport_security import TransportSecurityMiddleware, TransportSecuritySettings
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.message import ServerMessageMetadata, SessionMessage
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
from mcp.types import (
@ -119,10 +121,10 @@ class StreamableHTTPServerTransport:
"""
# Server notification streams for POST requests as well as standalone SSE stream
_read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception] | None = None
_read_stream: MemoryObjectReceiveStream[SessionMessage | Exception] | None = None
_write_stream: MemoryObjectSendStream[SessionMessage] | None = None
_write_stream_reader: MemoryObjectReceiveStream[SessionMessage] | None = None
_read_stream_writer: ContextSendStream[SessionMessage | Exception] | None = None
_read_stream: ContextReceiveStream[SessionMessage | Exception] | None = None
_write_stream: ContextSendStream[SessionMessage] | None = None
_write_stream_reader: ContextReceiveStream[SessionMessage] | None = None
_security: TransportSecurityMiddleware
def __init__(
@ -954,8 +956,8 @@ class StreamableHTTPServerTransport:
self,
) -> AsyncGenerator[
tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
ReadStream[SessionMessage | Exception],
WriteStream[SessionMessage],
],
None,
]:
@ -967,8 +969,8 @@ class StreamableHTTPServerTransport:
# Create the memory streams for this connection
read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
# Store the streams
self._read_stream_writer = read_stream_writer

12
src/mcp/server/websocket.py
View file

@ -1,12 +1,12 @@
from contextlib import asynccontextmanager
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic_core import ValidationError
from starlette.types import Receive, Scope, Send
from starlette.websockets import WebSocket
from mcp import types
from mcp.shared._context_streams import create_context_streams
from mcp.shared.message import SessionMessage
@ -19,14 +19,8 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
websocket = WebSocket(scope, receive, send)
await websocket.accept(subprotocol="mcp")
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
write_stream: MemoryObjectSendStream[SessionMessage]
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
async def ws_reader():
try:

119
src/mcp/shared/_context_streams.py Normal file
View file

@ -0,0 +1,119 @@
"""Context-aware memory stream wrappers.
anyio memory streams do not propagate ``contextvars.Context`` across task
boundaries. These thin wrappers capture the sender's context at ``send()``
time and expose it on the receive side via ``last_context``, so consumers
can restore it with ``ctx.run(handler, item)``.
The iteration interface is unchanged (yields ``T``, not tuples), keeping
these wrappers duck-type compatible with plain ``MemoryObjectSendStream``
and ``MemoryObjectReceiveStream``.
"""
from __future__ import annotations
import contextvars
from types import TracebackType
from typing import Any, Generic, TypeVar
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
T = TypeVar("T")
# Internal payload carried through the underlying raw stream.
_Envelope = tuple[contextvars.Context, T]
class ContextSendStream(Generic[T]):
"""Send-side wrapper that snapshots ``contextvars.copy_context()`` on every ``send()``."""
__slots__ = ("_inner",)
def __init__(self, inner: MemoryObjectSendStream[_Envelope[T]]) -> None:
self._inner = inner
async def send(self, item: T) -> None:
await self._inner.send((contextvars.copy_context(), item))
def close(self) -> None:
self._inner.close()
async def aclose(self) -> None:
await self._inner.aclose()
def clone(self) -> ContextSendStream[T]: # pragma: no cover
return ContextSendStream(self._inner.clone())
async def __aenter__(self) -> ContextSendStream[T]:
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> bool | None:
await self.aclose()
return None
class ContextReceiveStream(Generic[T]):
"""Receive-side wrapper that yields ``T`` and stores the sender's context in ``last_context``."""
__slots__ = ("_inner", "last_context")
def __init__(self, inner: MemoryObjectReceiveStream[_Envelope[T]]) -> None:
self._inner = inner
self.last_context: contextvars.Context | None = None
async def receive(self) -> T:
ctx, item = await self._inner.receive()
self.last_context = ctx
return item
def close(self) -> None:
self._inner.close()
async def aclose(self) -> None:
await self._inner.aclose()
def clone(self) -> ContextReceiveStream[T]: # pragma: no cover
return ContextReceiveStream(self._inner.clone())
def __aiter__(self) -> ContextReceiveStream[T]:
return self
async def __anext__(self) -> T:
try:
return await self.receive()
except anyio.EndOfStream:
raise StopAsyncIteration
async def __aenter__(self) -> ContextReceiveStream[T]:
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> bool | None:
await self.aclose()
return None
class create_context_streams(
tuple[ContextSendStream[T], ContextReceiveStream[T]],
):
"""Create context-aware memory object streams.
Supports ``create_context_streams[T](n)`` bracket syntax,
matching anyio's ``create_memory_object_stream`` API style.
"""
def __new__(cls, max_buffer_size: float = 0) -> tuple[ContextSendStream[T], ContextReceiveStream[T]]: # type: ignore[type-var]
raw_send: MemoryObjectSendStream[Any]
raw_receive: MemoryObjectReceiveStream[Any]
raw_send, raw_receive = anyio.create_memory_object_stream(max_buffer_size)
return (ContextSendStream(raw_send), ContextReceiveStream(raw_receive))

49
src/mcp/shared/_stream_protocols.py Normal file
View file

@ -0,0 +1,49 @@
"""Stream protocols for MCP transports.
These are general-purpose protocols satisfied by both ``MemoryObjectSendStream``/
``MemoryObjectReceiveStream`` and the context-aware wrappers in ``_context_streams``.
"""
from __future__ import annotations
from types import TracebackType
from typing import Protocol, TypeVar
from typing_extensions import Self
T_co = TypeVar("T_co", covariant=True)
T_contra = TypeVar("T_contra", contravariant=True)
class ReadStream(Protocol[T_co]):
"""Protocol for reading items from a stream.
Consumers that need the sender's context should use
``getattr(stream, 'last_context', None)``.
"""
async def receive(self) -> T_co: ...
async def aclose(self) -> None: ...
def __aiter__(self) -> ReadStream[T_co]: ...
async def __anext__(self) -> T_co: ...
async def __aenter__(self) -> Self: ...
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> bool | None: ...
class WriteStream(Protocol[T_contra]):
"""Protocol for writing items to a stream."""
async def send(self, item: T_contra, /) -> None: ...
async def aclose(self) -> None: ...
async def __aenter__(self) -> Self: ...
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> bool | None: ...

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

@ -5,12 +5,10 @@ from __future__ import annotations
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
from mcp.shared.message import SessionMessage
MessageStream = tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
MessageStream = tuple[ContextReceiveStream[SessionMessage | Exception], ContextSendStream[SessionMessage | Exception]]
@asynccontextmanager
@ -22,8 +20,8 @@ async def create_client_server_memory_streams() -> AsyncGenerator[tuple[MessageS
(read_stream, write_stream)
"""
# Create streams for both directions
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
server_to_client_send, server_to_client_receive = create_context_streams[SessionMessage | Exception](1)
client_to_server_send, client_to_server_receive = create_context_streams[SessionMessage | Exception](1)
client_streams = (server_to_client_receive, client_to_server_send)
server_streams = (client_to_server_receive, server_to_client_send)

26
src/mcp/shared/session.py
View file

@ -1,5 +1,6 @@
from __future__ import annotations
import contextvars
import logging
from collections.abc import Callable
from contextlib import AsyncExitStack
@ -7,10 +8,11 @@ from types import TracebackType
from typing import Any, Generic, Protocol, TypeVar
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from anyio.streams.memory import MemoryObjectSendStream
from pydantic import BaseModel, TypeAdapter
from typing_extensions import Self
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.exceptions import MCPError
from mcp.shared.message import MessageMetadata, ServerMessageMetadata, SessionMessage
from mcp.shared.response_router import ResponseRouter
@ -79,11 +81,13 @@ class RequestResponder(Generic[ReceiveRequestT, SendResultT]):
session: BaseSession[SendRequestT, SendNotificationT, SendResultT, ReceiveRequestT, ReceiveNotificationT],
on_complete: Callable[[RequestResponder[ReceiveRequestT, SendResultT]], Any],
message_metadata: MessageMetadata = None,
context: contextvars.Context | None = None,
) -> None:
self.request_id = request_id
self.request_meta = request_meta
self.request = request
self.message_metadata = message_metadata
self.context = context
self._session = session
self._completed = False
self._cancel_scope = anyio.CancelScope()
@ -181,8 +185,8 @@ class BaseSession(
def __init__(
self,
read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
write_stream: MemoryObjectSendStream[SessionMessage],
read_stream: ReadStream[SessionMessage | Exception],
write_stream: WriteStream[SessionMessage],
# If none, reading will never time out
read_timeout_seconds: float | None = None,
) -> None:
@ -333,10 +337,10 @@ class BaseSession(
async def _receive_loop(self) -> None:
async with self._read_stream, self._write_stream:
try:
async for message in self._read_stream:
if isinstance(message, Exception):
await self._handle_incoming(message)
elif isinstance(message.message, JSONRPCRequest):
async def _handle_session_message(message: SessionMessage) -> None:
sender_context: contextvars.Context | None = getattr(self._read_stream, "last_context", None)
if isinstance(message.message, JSONRPCRequest):
try:
validated_request = self._receive_request_adapter.validate_python(
message.message.model_dump(by_alias=True, mode="json", exclude_none=True),
@ -349,6 +353,7 @@ class BaseSession(
session=self,
on_complete=lambda r: self._in_flight.pop(r.request_id, None),
message_metadata=message.metadata,
context=sender_context,
)
self._in_flight[responder.request_id] = responder
await self._received_request(responder)
@ -406,6 +411,13 @@ class BaseSession(
else: # Response or error
await self._handle_response(message)
async for message in self._read_stream:
if isinstance(message, Exception):
await self._handle_incoming(message)
continue
await _handle_session_message(message)
except anyio.ClosedResourceError:
# This is expected when the client disconnects abruptly.
# Without this handler, the exception would propagate up and

4
tests/client/conftest.py
View file

@ -4,15 +4,15 @@ from typing import Any
from unittest.mock import patch
import pytest
from anyio.streams.memory import MemoryObjectSendStream
import mcp.shared.memory
from mcp.client._transport import WriteStream
from mcp.shared.message import SessionMessage
from mcp.types import JSONRPCNotification, JSONRPCRequest
class SpyMemoryObjectSendStream:
def __init__(self, original_stream: MemoryObjectSendStream[SessionMessage]):
def __init__(self, original_stream: WriteStream[SessionMessage]):
self.original_stream = original_stream
self.sent_messages: list[SessionMessage] = []

33
tests/client/test_client.py
View file

@ -2,6 +2,9 @@
from __future__ import annotations
import contextvars
from collections.abc import Iterator
from contextlib import contextmanager
from unittest.mock import patch
import anyio
@ -320,3 +323,33 @@ async def test_client_uses_transport_directly(app: MCPServer):
structured_content={"result": "Hello, Transport!"},
)
)
_TEST_CONTEXTVAR = contextvars.ContextVar("test_var", default="initial")
@contextmanager
def _set_test_contextvar(value: str) -> Iterator[None]:
token = _TEST_CONTEXTVAR.set(value)
try:
yield
finally:
_TEST_CONTEXTVAR.reset(token)
async def test_context_propagation():
"""Sender's contextvars.Context is propagated to the server handler."""
server = MCPServer("test")
@server.tool()
async def check_context() -> str:
"""Return the contextvar value visible to the handler."""
return _TEST_CONTEXTVAR.get()
async with Client(server) as client:
with _set_test_contextvar("client_value"):
result = await client.call_tool("check_context", {})
assert result.content[0].text == "client_value", ( # type: ignore[union-attr]
"Server handler did not see the sender's contextvars.Context"
)

10
tests/shared/test_streamable_http.py
View file

@ -45,6 +45,7 @@ from mcp.server.streamable_http import (
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.server.transport_security import TransportSecuritySettings
from mcp.shared._context import RequestContext
from mcp.shared._context_streams import create_context_streams
from mcp.shared._httpx_utils import (
MCP_DEFAULT_SSE_READ_TIMEOUT,
MCP_DEFAULT_TIMEOUT,
@ -1783,8 +1784,8 @@ async def test_handle_sse_event_skips_empty_data():
# Create a mock SSE event with empty data (keep-alive ping)
mock_sse = ServerSentEvent(event="message", data="", id=None, retry=None)
# Create a mock stream writer
write_stream, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](1)
# Create a context-aware stream writer (matches StreamWriter type alias)
write_stream, read_stream = create_context_streams[SessionMessage | Exception](1)
try:
# Call _handle_sse_event with empty data - should return False and not raise
@ -1794,8 +1795,9 @@ async def test_handle_sse_event_skips_empty_data():
assert result is False
# Nothing should have been written to the stream
# Check buffer is empty (statistics().current_buffer_used returns buffer size)
assert write_stream.statistics().current_buffer_used == 0
with pytest.raises(TimeoutError):
with anyio.fail_after(0):
await read_stream.receive()
finally:
await write_stream.aclose()
await read_stream.aclose()