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

@@ -45,7 +45,9 @@
 LAST_EVENT_ID = "last-event-id"
 
 # Reconnection defaults
-DEFAULT_RECONNECTION_DELAY_MS = 1000  # 1 second fallback when server doesn't provide retry
+DEFAULT_RECONNECTION_DELAY_MS = (
+    1000  # 1 second fallback when server doesn't provide retry
+)
 MAX_RECONNECTION_ATTEMPTS = 2  # Max retry attempts before giving up
 
 
@@ -104,7 +106,10 @@ def _is_initialization_request(self, message: JSONRPCMessage) -> bool:
 
     def _is_initialized_notification(self, message: JSONRPCMessage) -> bool:
         """Check if the message is an initialized notification."""
-        return isinstance(message, JSONRPCNotification) and message.method == "notifications/initialized"
+        return (
+            isinstance(message, JSONRPCNotification)
+            and message.method == "notifications/initialized"
+        )
 
     def _maybe_extract_session_id_from_response(self, response: httpx.Response) -> None:
         """Extract and store session ID from response headers."""
@@ -113,16 +118,23 @@ def _maybe_extract_session_id_from_response(self, response: httpx.Response) -> N
             self.session_id = new_session_id
             logger.info(f"Received session ID: {self.session_id}")
 
-    def _maybe_extract_protocol_version_from_message(self, message: JSONRPCMessage) -> None:
+    def _maybe_extract_protocol_version_from_message(
+        self, message: JSONRPCMessage
+    ) -> None:
         """Extract protocol version from initialization response message."""
         if isinstance(message, JSONRPCResponse) and message.result:  # pragma: no branch
             try:
                 # Parse the result as InitializeResult for type safety
-                init_result = InitializeResult.model_validate(message.result, by_name=False)
+                init_result = InitializeResult.model_validate(
+                    message.result, by_name=False
+                )
                 self.protocol_version = str(init_result.protocol_version)
                 logger.info(f"Negotiated protocol version: {self.protocol_version}")
             except Exception:  # pragma: no cover
-                logger.warning("Failed to parse initialization response as InitializeResult", exc_info=True)
+                logger.warning(
+                    "Failed to parse initialization response as InitializeResult",
+                    exc_info=True,
+                )
                 logger.warning(f"Raw result: {message.result}")
 
     async def _handle_sse_event(
@@ -150,7 +162,9 @@ async def _handle_sse_event(
                     self._maybe_extract_protocol_version_from_message(message)
 
                 # If this is a response and we have original_request_id, replace it
-                if original_request_id is not None and isinstance(message, JSONRPCResponse | JSONRPCError):
+                if original_request_id is not None and isinstance(
+                    message, JSONRPCResponse | JSONRPCError
+                ):
                     message.id = original_request_id
 
                 session_message = SessionMessage(message)
@@ -167,8 +181,14 @@ async def _handle_sse_event(
             except Exception as exc:  # pragma: no cover
                 logger.exception("Error parsing SSE message")
                 if original_request_id is not None:
-                    error_data = ErrorData(code=PARSE_ERROR, message=f"Failed to parse SSE message: {exc}")
-                    error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=original_request_id, error=error_data))
+                    error_data = ErrorData(
+                        code=PARSE_ERROR, message=f"Failed to parse SSE message: {exc}"
+                    )
+                    error_msg = SessionMessage(
+                        JSONRPCError(
+                            jsonrpc="2.0", id=original_request_id, error=error_data
+                        )
+                    )
                     await read_stream_writer.send(error_msg)
                     return True
                 await read_stream_writer.send(exc)
@@ -177,7 +197,9 @@ async def _handle_sse_event(
             logger.warning(f"Unknown SSE event: {sse.event}")
             return False
 
-    async def handle_get_stream(self, client: httpx.AsyncClient, read_stream_writer: StreamWriter) -> None:
+    async def handle_get_stream(
+        self, client: httpx.AsyncClient, read_stream_writer: StreamWriter
+    ) -> None:
         """Handle GET stream for server-initiated messages with auto-reconnect."""
         last_event_id: str | None = None
         retry_interval_ms: int | None = None
@@ -192,7 +214,9 @@ async def handle_get_stream(self, client: httpx.AsyncClient, read_stream_writer:
                 if last_event_id:
                     headers[LAST_EVENT_ID] = last_event_id
 
-                async with aconnect_sse(client, "GET", self.url, headers=headers) as event_source:
+                async with aconnect_sse(
+                    client, "GET", self.url, headers=headers
+                ) as event_source:
                     event_source.response.raise_for_status()
                     logger.debug("GET SSE connection established")
 
@@ -214,11 +238,17 @@ async def handle_get_stream(self, client: httpx.AsyncClient, read_stream_writer:
                 attempt += 1
 
             if attempt >= MAX_RECONNECTION_ATTEMPTS:  # pragma: no cover
-                logger.debug(f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded")
+                logger.debug(
+                    f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded"
+                )
                 return
 
             # Wait before reconnecting
-            delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS
+            delay_ms = (
+                retry_interval_ms
+                if retry_interval_ms is not None
+                else DEFAULT_RECONNECTION_DELAY_MS
+            )
             logger.info(f"GET stream disconnected, reconnecting in {delay_ms}ms...")
             await anyio.sleep(delay_ms / 1000.0)
 
@@ -228,14 +258,18 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
         if ctx.metadata and ctx.metadata.resumption_token:
             headers[LAST_EVENT_ID] = ctx.metadata.resumption_token
         else:
-            raise ResumptionError("Resumption request requires a resumption token")  # pragma: no cover
+            raise ResumptionError(
+                "Resumption request requires a resumption token"
+            )  # pragma: no cover
 
         # Extract original request ID to map responses
         original_request_id = None
         if isinstance(ctx.session_message.message, JSONRPCRequest):  # pragma: no branch
             original_request_id = ctx.session_message.message.id
 
-        async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source:
+        async with aconnect_sse(
+            ctx.client, "GET", self.url, headers=headers
+        ) as event_source:
             event_source.response.raise_for_status()
             logger.debug("Resumption GET SSE connection established")
 
@@ -268,8 +302,12 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
 
             if response.status_code == 404:  # pragma: no branch
                 if isinstance(message, JSONRPCRequest):  # pragma: no branch
-                    error_data = ErrorData(code=INVALID_REQUEST, message="Session terminated")
-                    session_message = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data))
+                    error_data = ErrorData(
+                        code=INVALID_REQUEST, message="Session terminated"
+                    )
+                    session_message = SessionMessage(
+                        JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data)
+                    )
                     await ctx.read_stream_writer.send(session_message)
                 return
 
@@ -283,14 +321,22 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                 content_type = response.headers.get("content-type", "").lower()
                 if content_type.startswith("application/json"):
                     await self._handle_json_response(
-                        response, ctx.read_stream_writer, is_initialization, request_id=message.id
+                        response,
+                        ctx.read_stream_writer,
+                        is_initialization,
+                        request_id=message.id,
                     )
                 elif content_type.startswith("text/event-stream"):
                     await self._handle_sse_response(response, ctx, is_initialization)
                 else:
                     logger.error(f"Unexpected content type: {content_type}")
-                    error_data = ErrorData(code=INVALID_REQUEST, message=f"Unexpected content type: {content_type}")
-                    error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data))
+                    error_data = ErrorData(
+                        code=INVALID_REQUEST,
+                        message=f"Unexpected content type: {content_type}",
+                    )
+                    error_msg = SessionMessage(
+                        JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data)
+                    )
                     await ctx.read_stream_writer.send(error_msg)
 
     async def _handle_json_response(
@@ -314,8 +360,12 @@ async def _handle_json_response(
             await read_stream_writer.send(session_message)
         except (httpx.StreamError, ValidationError) as exc:
             logger.exception("Error parsing JSON response")
-            error_data = ErrorData(code=PARSE_ERROR, message=f"Failed to parse JSON response: {exc}")
-            error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=request_id, error=error_data))
+            error_data = ErrorData(
+                code=PARSE_ERROR, message=f"Failed to parse JSON response: {exc}"
+            )
+            error_msg = SessionMessage(
+                JSONRPCError(jsonrpc="2.0", id=request_id, error=error_data)
+            )
             await read_stream_writer.send(error_msg)
 
     async def _handle_sse_response(
@@ -348,7 +398,11 @@ async def _handle_sse_response(
                     sse,
                     ctx.read_stream_writer,
                     original_request_id=original_request_id,
-                    resumption_callback=(ctx.metadata.on_resumption_token_update if ctx.metadata else None),
+                    resumption_callback=(
+                        ctx.metadata.on_resumption_token_update
+                        if ctx.metadata
+                        else None
+                    ),
                     is_initialization=is_initialization,
                 )
                 # If the SSE event indicates completion, like returning respose/error
@@ -374,11 +428,17 @@ async def _handle_reconnection(
         """Reconnect with Last-Event-ID to resume stream after server disconnect."""
         # Bail if max retries exceeded
         if attempt >= MAX_RECONNECTION_ATTEMPTS:  # pragma: no cover
-            logger.debug(f"Max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded")
+            logger.debug(
+                f"Max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded"
+            )
             return
 
         # Always wait - use server value or default
-        delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS
+        delay_ms = (
+            retry_interval_ms
+            if retry_interval_ms is not None
+            else DEFAULT_RECONNECTION_DELAY_MS
+        )
         await anyio.sleep(delay_ms / 1000.0)
 
         headers = self._prepare_headers()
@@ -390,7 +450,9 @@ async def _handle_reconnection(
             original_request_id = ctx.session_message.message.id
 
         try:
-            async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source:
+            async with aconnect_sse(
+                ctx.client, "GET", self.url, headers=headers
+            ) as event_source:
                 event_source.response.raise_for_status()
                 logger.info("Reconnected to SSE stream")
 
@@ -408,19 +470,25 @@ async def _handle_reconnection(
                         sse,
                         ctx.read_stream_writer,
                         original_request_id,
-                        ctx.metadata.on_resumption_token_update if ctx.metadata else None,
+                        ctx.metadata.on_resumption_token_update
+                        if ctx.metadata
+                        else None,
                     )
                     if is_complete:
                         await event_source.response.aclose()
                         return
 
                 # Stream ended again without response - reconnect again (reset attempt counter)
                 logger.info("SSE stream disconnected, reconnecting...")
-                await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, 0)
+                await self._handle_reconnection(
+                    ctx, reconnect_last_event_id, reconnect_retry_ms, 0
+                )
         except Exception as e:  # pragma: no cover
             logger.debug(f"Reconnection failed: {e}")
             # Try to reconnect again if we still have an event ID
-            await self._handle_reconnection(ctx, last_event_id, retry_interval_ms, attempt + 1)
+            await self._handle_reconnection(
+                ctx, last_event_id, retry_interval_ms, attempt + 1
+            )
 
     async def post_writer(
         self,
@@ -434,7 +502,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
@@ -471,8 +540,14 @@ async def handle_request_async():
                     else:
                         await handle_request_async()
 
-        except Exception:  # pragma: lax no cover
-            logger.exception("Error in post_writer")
+                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:
             await read_stream_writer.aclose()
             await write_stream.aclose()
@@ -526,8 +601,12 @@ async def streamable_http_client(
     Example:
         See examples/snippets/clients/ for usage patterns.
     """
-    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 = anyio.create_memory_object_stream[
+        SessionMessage | Exception
+    ](0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream[
+        SessionMessage
+    ](0)
 
     # Determine if we need to create and manage the client
     client_provided = http_client is not None
@@ -549,7 +628,9 @@ async def streamable_http_client(
                     await stack.enter_async_context(client)
 
                 def start_get_stream() -> None:
-                    tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
+                    tg.start_soon(
+                        transport.handle_get_stream, client, read_stream_writer
+                    )
 
                 tg.start_soon(
                     transport.post_writer,

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