Propagate contextvars through anyio streams

TODO:
- Update a recipe to show it working
- Consider adding an integration test of some kind

Author: aabmass

src/mcp/client/streamable_http.py

@@ -430,7 +430,8 @@ async def post_writer(
         """Handle writing requests to the server."""
         try:
             async with write_stream_reader:
-                async for session_message in write_stream_reader:
+
+                async def handle_message(session_message: SessionMessage) -> None:
                     message = session_message.message
                     metadata = (
                         session_message.metadata
@@ -467,6 +468,10 @@ async def handle_request_async():
                     else:
                         await handle_request_async()
 
+                async for session_message in write_stream_reader:
+                    async with anyio.create_task_group() as tg_local:
+                        session_message.context.run(tg_local.start_soon, handle_message, session_message)
+
         except Exception:
             logger.exception("Error in post_writer")  # pragma: no cover
         finally:

src/mcp/server/lowlevel/server.py

@@ -683,7 +683,14 @@ async def run(
                 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:
+                        logger.debug("Got a context to propagate, %s", message.context)
+                        context = message.context
+                    else:
+                        context = contextvars.copy_context()
+
+                    context.run(
+                        tg.start_soon,
                         self._handle_message,
                         message,
                         session,

src/mcp/shared/message.py

@@ -4,8 +4,9 @@
 to support transport-specific features like resumability.
 """
 
+import contextvars
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from mcp.types import JSONRPCMessage, RequestId
 
@@ -46,4 +47,5 @@ class SessionMessage:
     """A message with specific metadata for transport-specific features."""
 
     message: JSONRPCMessage
+    context: contextvars.Context = field(default_factory=contextvars.copy_context)
     metadata: MessageMetadata = None

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
@@ -77,11 +78,13 @@ def __init__(
         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()
@@ -330,10 +333,9 @@ def _receive_notification_adapter(self) -> TypeAdapter[ReceiveNotificationT]:
     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):  # pragma: no cover
-                        await self._handle_incoming(message)
-                    elif isinstance(message.message, JSONRPCRequest):
+
+                async def handle_message(message: SessionMessage) -> 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),
@@ -346,6 +348,7 @@ async def _receive_loop(self) -> None:
                                 session=self,
                                 on_complete=lambda r: self._in_flight.pop(r.request_id, None),
                                 message_metadata=message.metadata,
+                                context=message.context,
                             )
                             self._in_flight[responder.request_id] = responder
                             await self._received_request(responder)
@@ -403,6 +406,13 @@ async def _receive_loop(self) -> None:
                     else:  # Response or error
                         await self._handle_response(message)
 
+                async for message in self._read_stream:
+                    if isinstance(message, Exception):  # pragma: no cover
+                        await self._handle_incoming(message)
+                    else:
+                        async with anyio.create_task_group() as tg:
+                            message.context.run(tg.start_soon, handle_message, message)
+
             except anyio.ClosedResourceError:
                 # This is expected when the client disconnects abruptly.
                 # Without this handler, the exception would propagate up and