Restore correct context propagation for async handlers

ctx.run(async_fn) only sets context during coroutine creation, not
execution. For session._receive_loop, pass sender_ctx explicitly to
the handler which stores it on RequestResponder.context. For client
post_writers (sse/streamable_http), restore ctx.run(tg.start_soon, ...)
so the httpx calls actually run in the sender's context.

Author: Kludex

src/mcp/client/sse.py

@@ -143,7 +143,8 @@ async def _send_message(session_message: SessionMessage) -> None:
                                 async for session_message in write_stream_reader:
                                     sender_ctx = write_stream_reader.last_context
                                     if sender_ctx is not None:
-                                        await sender_ctx.run(_send_message, session_message)
+                                        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

src/mcp/client/streamable_http.py

@@ -482,7 +482,8 @@ async def handle_request_async():
                 async for session_message in write_stream_reader:
                     sender_ctx = write_stream_reader.last_context
                     if sender_ctx is not None:
-                        await sender_ctx.run(_handle_message, session_message)
+                        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
 

src/mcp/shared/session.py

@@ -338,7 +338,10 @@ async def _receive_loop(self) -> None:
         async with self._read_stream, self._write_stream:
             try:
 
-                async def _handle_session_message(message: SessionMessage) -> None:
+                async def _handle_session_message(
+                    message: SessionMessage,
+                    sender_context: contextvars.Context | None = None,
+                ) -> None:
                     if isinstance(message.message, JSONRPCRequest):
                         try:
                             validated_request = self._receive_request_adapter.validate_python(
@@ -352,7 +355,7 @@ async def _handle_session_message(message: SessionMessage) -> None:
                                 session=self,
                                 on_complete=lambda r: self._in_flight.pop(r.request_id, None),
                                 message_metadata=message.metadata,
-                                context=contextvars.copy_context(),
+                                context=sender_context,
                             )
                             self._in_flight[responder.request_id] = responder
                             await self._received_request(responder)
@@ -416,11 +419,7 @@ async def _handle_session_message(message: SessionMessage) -> None:
                         continue
 
                     sender_ctx: contextvars.Context | None = getattr(self._read_stream, "last_context", None)
-                    if sender_ctx is not None:
-                        coro = sender_ctx.run(_handle_session_message, message)
-                        await coro
-                    else:
-                        await _handle_session_message(message)
+                    await _handle_session_message(message, sender_context=sender_ctx)
 
             except anyio.ClosedResourceError:
                 # This is expected when the client disconnects abruptly.