Experimental MCP propagation (#1103)

Co-authored-by: Alex Hall <[email protected]>

Author: samuelcolvin

logfire/_internal/integrations/mcp.py

@@ -1,34 +1,38 @@
 from __future__ import annotations
 
 import functools
+from contextlib import ExitStack, contextmanager
 from typing import TYPE_CHECKING, Any
 
 from mcp.client.session import ClientSession
+from mcp.server import Server
 from mcp.shared.session import BaseSession
-from mcp.types import CallToolRequest, LoggingMessageNotification
+from mcp.types import CallToolRequest, LoggingMessageNotification, RequestParams
+from pydantic import TypeAdapter
 
 from logfire._internal.utils import handle_internal_errors
+from logfire.propagate import attach_context, get_context
 
 if TYPE_CHECKING:
     from logfire import LevelName, Logfire
 
 
-def instrument_mcp(logfire_instance: Logfire):
+def instrument_mcp(logfire_instance: Logfire, propagate_otel_context: bool):
     logfire_instance = logfire_instance.with_settings(custom_scope_suffix='mcp')
 
     original_send_request = BaseSession.send_request  # type: ignore
 
     @functools.wraps(original_send_request)  # type: ignore
     async def send_request(self: Any, request: Any, *args: Any, **kwargs: Any):
+        root = request.root
         attributes: dict[str, Any] = {
-            'request': request,
+            'request': root,
             # https://opentelemetry.io/docs/specs/semconv/rpc/json-rpc/
             'rpc.system': 'jsonrpc',
             'rpc.jsonrpc.version': '2.0',
         }
         span_name = 'MCP request'
 
-        root = request.root
         # method should always exist, but it's had to verify because the request type is a RootModel
         # of a big union, instead of just using a base class with a method attribute.
         if method := getattr(root, 'method', None):  # pragma: no branch
@@ -38,6 +42,20 @@ async def send_request(self: Any, request: Any, *args: Any, **kwargs: Any):
                 span_name += f' {root.params.name}'
 
         with logfire_instance.span(span_name, **attributes) as span:
+            with handle_internal_errors:
+                if propagate_otel_context:  # pragma: no branch
+                    carrier = get_context()
+                    if params := getattr(root, 'params', None):
+                        if meta := getattr(params, 'meta', None):  # pragma: no cover # TODO
+                            dumped_meta = meta.model_dump()
+                        else:
+                            dumped_meta = {}
+                        # Prioritise existing values in meta over the context carrier.
+                        # RequestParams.Meta should allow basically anything, we're being extra careful here.
+                        params.meta = RequestParams.Meta.model_validate({**carrier, **dumped_meta})
+                    else:
+                        root.params = _request_params_type_adapter(type(root)).validate_python({'_meta': carrier})  # type: ignore
+
             result = await original_send_request(self, request, *args, **kwargs)
             span.set_attribute('response', result)
             return result
@@ -63,3 +81,44 @@ async def _received_notification(self: Any, notification: Any, *args: Any, **kwa
         await original_received_notification(self, notification, *args, **kwargs)
 
     ClientSession._received_notification = _received_notification  # type: ignore
+
+    original_handle_client_request = ClientSession._received_request  # type: ignore
+
+    @functools.wraps(original_handle_client_request)
+    async def _received_request_client(self: Any, responder: Any) -> None:  # pragma: no cover
+        request = responder.request.root
+        span_name = 'MCP client handle request'
+        with _handle_request_with_context(request, span_name):
+            await original_handle_client_request(self, responder)
+
+    ClientSession._received_request = _received_request_client  # type: ignore
+
+    original_handle_server_request = Server._handle_request  # type: ignore
+
+    @functools.wraps(original_handle_server_request)  # type: ignore
+    async def _handle_request(self: Any, message: Any, request: Any, *args: Any, **kwargs: Any) -> Any:
+        span_name = 'MCP server handle request'
+        with _handle_request_with_context(request, span_name):
+            return await original_handle_server_request(self, message, request, *args, **kwargs)
+
+    Server._handle_request = _handle_request  # type: ignore
+
+    @contextmanager
+    def _handle_request_with_context(request: Any, span_name: str):
+        with ExitStack() as exit_stack:
+            if (  # pragma: no branch
+                propagate_otel_context
+                and (params := getattr(request, 'params', None))
+                and (meta := getattr(params, 'meta', None))
+            ):
+                exit_stack.enter_context(attach_context(meta.model_dump()))
+            if method := getattr(request, 'method', None):  # pragma: no branch
+                span_name += f': {method}'
+            with logfire_instance.span(span_name, request=request):
+                yield
+
+
+@functools.lru_cache
+def _request_params_type_adapter(root_type: Any):
+    params_type = root_type.model_fields['params'].annotation
+    return TypeAdapter(params_type)

logfire/_internal/main.py

@@ -900,12 +900,17 @@ def install_auto_tracing(
     def _warn_if_not_initialized_for_instrumentation(self):
         self.config.warn_if_not_initialized('Instrumentation will have no effect')
 
-    def instrument_mcp(self) -> None:
-        """Instrument [MCP](https://modelcontextprotocol.io/) requests such as tool calls."""
+    def instrument_mcp(self, *, propagate_otel_context: bool = True) -> None:
+        """Instrument [MCP](https://modelcontextprotocol.io/) requests such as tool calls.
+
+        Args:
+            propagate_otel_context: Whether to enable propagation of the OpenTelemetry context.
+                Set to False to prevent setting extra fields like `traceparent` on the metadata of requests.
+        """
         from .integrations.mcp import instrument_mcp
 
         self._warn_if_not_initialized_for_instrumentation()
-        instrument_mcp(self)
+        instrument_mcp(self, propagate_otel_context)
 
     def instrument_pydantic(
         self,