Extract OTel context from _meta in incoming requests

This works for both incoming client->server requests and server->client requests like MCP Sampling.

Author: aabmass

src/mcp/client/sse.py

@@ -132,7 +132,8 @@ async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
                     async def post_writer(endpoint_url: str):
                         try:
                             async with write_stream_reader:
-                                async for session_message in write_stream_reader:
+
+                                async def handle_message(session_message: SessionMessage) -> None:
                                     logger.debug(f"Sending client message: {session_message}")
                                     response = await client.post(
                                         endpoint_url,
@@ -144,6 +145,13 @@ async def post_writer(endpoint_url: str):
                                     )
                                     response.raise_for_status()
                                     logger.debug(f"Client message sent successfully: {response.status_code}")
+
+                                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:  # pragma: lax no cover
                             logger.exception("Error in post_writer")
                         finally:

src/mcp/client/streamable_http.py

@@ -441,7 +441,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
@@ -478,8 +479,12 @@ 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()

src/mcp/server/lowlevel/server.py

@@ -393,7 +393,13 @@ 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:
+                        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 typing import Any
 
 from mcp.types import JSONRPCMessage, RequestId
@@ -49,4 +50,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
@@ -8,7 +9,8 @@
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from opentelemetry.propagate import inject
+from opentelemetry import context as otel_context
+from opentelemetry.propagate import extract, inject
 from pydantic import BaseModel, TypeAdapter
 from typing_extensions import Self
 
@@ -80,11 +82,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()
@@ -363,10 +367,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),
@@ -379,6 +382,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)
@@ -427,15 +431,35 @@ async def _receive_loop(self) -> None:
                                             logging.exception("Progress callback raised an exception")
                                 await self._received_notification(notification)
                                 await self._handle_incoming(notification)
-                        except Exception:
+                        except Exception:  # pragma: lax no cover
                             # For other validation errors, log and continue
-                            logging.warning(  # pragma: no cover
+                            logging.warning(
                                 f"Failed to validate notification:. Message was: {message.message}",
                                 exc_info=True,
                             )
                     else:  # Response or error
                         await self._handle_response(message)
 
+                async def _handle_message_with_otel(message: SessionMessage) -> None:
+                    meta = None
+                    if isinstance(message.message, (JSONRPCRequest | JSONRPCNotification)) and message.message.params:
+                        meta = message.message.params.get("_meta")
+
+                    extracted_ctx = extract(meta) if meta else None
+                    otel_token = otel_context.attach(extracted_ctx) if extracted_ctx else None
+                    try:
+                        await handle_message(message)
+                    finally:
+                        if otel_token:
+                            otel_context.detach(otel_token)
+
+                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_with_otel, message)
+
             except anyio.ClosedResourceError:
                 # This is expected when the client disconnects abruptly.
                 # Without this handler, the exception would propagate up and

tests/shared/test_otel_context_meta.py

@@ -79,6 +79,11 @@ async def tool_with_sampling(topic: str, ctx: Context[ServerSession, None]) -> s
             )
         return "ran sampling"
 
+    @mcp.tool()
+    async def tool_that_checks_trace_context() -> str:
+        """Returns current span details to verify parent propagation."""
+        return trace.format_trace_id(trace.get_current_span().get_span_context().trace_id)
+
     return mcp
 
 
@@ -106,8 +111,9 @@ async def patched_client(server: MCPServer, monkeypatch: pytest.MonkeyPatch):
     async def sampling_callback(
         context: RequestContext[ClientSession], params: types.CreateMessageRequestParams
     ) -> types.CreateMessageResult:
+        current_trace_id = trace.format_trace_id(trace.get_current_span().get_span_context().trace_id)
         return types.CreateMessageResult(
-            role="assistant", content=TextContent(type="text", text="hello"), model="foomodel"
+            role="assistant", content=TextContent(type="text", text=current_trace_id), model="foomodel"
         )
 
     async with create_client_server_memory_streams() as (client_streams, server_streams):
@@ -231,6 +237,17 @@ async def test_with_existing_meta(
     assert patched_client.client_to_server_messages == expect_client_to_server
 
 
+@pytest.mark.anyio
+async def test_trace_context_extraction(patched_client: PatchedClient):
+    """Test that OTEL context is successfully extracted on the receiving end."""
+
+    with trace.use_span(SPAN_IN_CLIENT):
+        result = await patched_client.session.call_tool("tool_that_checks_trace_context")
+
+    # Verify that SPAN_IN_CLIENT was extracted and made it through to the handler
+    assert result.content[0] == snapshot(TextContent(text="00000000000000000000000000000123"))
+
+
 @pytest.mark.anyio
 async def test_list_tools_with_span(patched_client: PatchedClient):
     """Test that OTEL context is injected into the _meta field of a tools/list request."""
@@ -316,7 +333,11 @@ async def test_server_side_sampling_propagates_to_client(patched_client: Patched
             JSONRPCResponse(
                 jsonrpc="2.0",
                 id=0,
-                result={"role": "assistant", "content": {"type": "text", "text": "hello"}, "model": "foomodel"},
+                result={
+                    "role": "assistant",
+                    "content": {"type": "text", "text": "00000000000000000000000000000456"},
+                    "model": "foomodel",
+                },
             ),
         ]
     )