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

examples/clients/simple-auth-client/mcp_simple_auth_client/main.py

@@ -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],
+        read_stream: ReadStream[SessionMessage | Exception],
-        write_stream: MemoryObjectSendStream[SessionMessage],
+        write_stream: WriteStream[SessionMessage],
     ):
         """Run the MCP session with the given streams."""
         print("🤝 Initializing MCP session...")

pyproject.toml

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

src/mcp/client/__main__.py

@@ -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],
+    read_stream: ReadStream[SessionMessage | Exception],
-    write_stream: MemoryObjectSendStream[SessionMessage],
+    write_stream: WriteStream[SessionMessage],
     client_info: types.Implementation | None = None,
 ):
     async with ClientSession(

src/mcp/client/_transport.py

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

src/mcp/client/session.py

@@ -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],
+        read_stream: ReadStream[SessionMessage | Exception],
-        write_stream: MemoryObjectSendStream[SessionMessage],
+        write_stream: WriteStream[SessionMessage],
         read_timeout_seconds: float | None = None,
         sampling_callback: SamplingFnT | None = None,
         elicitation_callback: ElicitationFnT | None = None,

src/mcp/client/sse.py

@@ -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)
+            read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
-            write_stream, write_stream_reader = anyio.create_memory_object_stream(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")
 

src/mcp/client/streamable_http.py

@@ -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]
+StreamWriter = ContextSendStream[SessionMessageOrError]
-StreamReader = MemoryObjectReceiveStream[SessionMessage]
+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)
+        read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
-        write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
+        write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
 
         async with (
             read_stream_writer,

src/mcp/server/lowlevel/server.py

@@ -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],
+        read_stream: ReadStream[SessionMessage | Exception],
-        write_stream: MemoryObjectSendStream[SessionMessage],
+        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,

src/mcp/server/session.py

@@ -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],
+        read_stream: ReadStream[SessionMessage | Exception],
-        write_stream: MemoryObjectSendStream[SessionMessage],
+        write_stream: WriteStream[SessionMessage],
         init_options: InitializationOptions,
         stateless: bool = False,
     ) -> None:

src/mcp/server/sse.py

@@ -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]
+        read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
-        write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
+        write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
-
-        read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
-        write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
 
         session_id = uuid4()
         self._read_stream_writers[session_id] = read_stream_writer

src/mcp/server/stdio.py

@@ -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, read_stream = create_context_streams[SessionMessage | Exception](0)
-    read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
+    write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
-
-    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)
 
     async def stdin_reader():
         try:

src/mcp/server/streamable_http.py

@@ -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_writer: ContextSendStream[SessionMessage | Exception] | None = None
-    _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception] | None = None
+    _read_stream: ContextReceiveStream[SessionMessage | Exception] | None = None
-    _write_stream: MemoryObjectSendStream[SessionMessage] | None = None
+    _write_stream: ContextSendStream[SessionMessage] | None = None
-    _write_stream_reader: MemoryObjectReceiveStream[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],
+            ReadStream[SessionMessage | Exception],
-            MemoryObjectSendStream[SessionMessage],
+            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)
+        read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
-        write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
+        write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
 
         # Store the streams
         self._read_stream_writer = read_stream_writer

src/mcp/server/websocket.py

@@ -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, read_stream = create_context_streams[SessionMessage | Exception](0)
-    read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
+    write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
-
-    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)
 
     async def ws_reader():
         try:

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 mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
-from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-
 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)
+    server_to_client_send, server_to_client_receive = create_context_streams[SessionMessage | Exception](1)
-    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[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],
+        read_stream: ReadStream[SessionMessage | Exception],
-        write_stream: MemoryObjectSendStream[SessionMessage],
+        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):
+                async def _handle_session_message(message: SessionMessage) -> None:
-                        await self._handle_incoming(message)
+                    sender_context: contextvars.Context | None = getattr(self._read_stream, "last_context", None)
-                    elif isinstance(message.message, JSONRPCRequest):
+                    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

tests/client/conftest.py

@@ -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] = []
 

tests/client/test_client.py

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

tests/shared/test_streamable_http.py

@@ -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
+    # Create a context-aware stream writer (matches StreamWriter type alias)
-    write_stream, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](1)
+    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)
+        with pytest.raises(TimeoutError):
-        assert write_stream.statistics().current_buffer_used == 0
+            with anyio.fail_after(0):
+                await read_stream.receive()
     finally:
         await write_stream.aclose()
         await read_stream.aclose()