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

src/mcp/shared/_context_streams.py

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

src/mcp/shared/_stream_protocols.py

@@ -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: ...

src/mcp/shared/memory.py

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

src/mcp/shared/session.py

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