Add logfire.instrument_mcp() method (#966)

Author: alexmojaki

logfire-api/logfire_api/__init__.py

@@ -147,6 +147,8 @@ def instrument_aiohttp_client(self, *args, **kwargs) -> None: ...
 
             def instrument_system_metrics(self, *args, **kwargs) -> None: ...
 
+            def instrument_mcp(self, *args, **kwargs) -> None: ...
+
             def shutdown(self, *args, **kwargs) -> None: ...
 
         DEFAULT_LOGFIRE_INSTANCE = Logfire()
@@ -191,6 +193,7 @@ def shutdown(self, *args, **kwargs) -> None: ...
         instrument_pymongo = DEFAULT_LOGFIRE_INSTANCE.instrument_pymongo
         instrument_mysql = DEFAULT_LOGFIRE_INSTANCE.instrument_mysql
         instrument_system_metrics = DEFAULT_LOGFIRE_INSTANCE.instrument_system_metrics
+        instrument_mcp = DEFAULT_LOGFIRE_INSTANCE.instrument_mcp
         shutdown = DEFAULT_LOGFIRE_INSTANCE.shutdown
         suppress_scopes = DEFAULT_LOGFIRE_INSTANCE.suppress_scopes
 

logfire/__init__.py

@@ -49,6 +49,7 @@
 instrument_pymongo = DEFAULT_LOGFIRE_INSTANCE.instrument_pymongo
 instrument_mysql = DEFAULT_LOGFIRE_INSTANCE.instrument_mysql
 instrument_system_metrics = DEFAULT_LOGFIRE_INSTANCE.instrument_system_metrics
+instrument_mcp = DEFAULT_LOGFIRE_INSTANCE.instrument_mcp
 suppress_scopes = DEFAULT_LOGFIRE_INSTANCE.suppress_scopes
 shutdown = DEFAULT_LOGFIRE_INSTANCE.shutdown
 with_tags = DEFAULT_LOGFIRE_INSTANCE.with_tags
@@ -138,6 +139,7 @@ def loguru_handler() -> Any:
     'instrument_pymongo',
     'instrument_mysql',
     'instrument_system_metrics',
+    'instrument_mcp',
     'AutoTraceModule',
     'with_tags',
     'with_settings',

logfire/_internal/integrations/mcp.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING, Any
+
+from mcp.shared.session import BaseSession, SendRequestT
+from mcp.types import CallToolRequest
+
+if TYPE_CHECKING:
+    from logfire import Logfire
+
+
+def instrument_mcp(logfire_instance: Logfire):
+    logfire_instance = logfire_instance.with_settings(custom_scope_suffix='mcp')
+
+    original = BaseSession.send_request  # type: ignore
+
+    @functools.wraps(original)  # type: ignore
+    async def send_request(self, request: SendRequestT, *args, **kwargs: Any):  # type: ignore
+        attributes: dict[str, Any] = {
+            'request': request,
+            # 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
+            span_name += f': {method}'
+            attributes['rpc.method'] = method
+            if isinstance(root, CallToolRequest):
+                span_name += f' {root.params.name}'
+
+        with logfire_instance.span(span_name, **attributes) as span:
+            result = await original(self, request, *args, **kwargs)  # type: ignore
+            span.set_attribute('response', result)
+            return result  # type: ignore
+
+    BaseSession.send_request = send_request

logfire/_internal/main.py

@@ -899,6 +899,13 @@ 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."""
+        from .integrations.mcp import instrument_mcp
+
+        self._warn_if_not_initialized_for_instrumentation()
+        instrument_mcp(self)
+
     def instrument_pydantic(
         self,
         record: PydanticPluginRecordValues = 'all',