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

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],
-        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,