MCP logging (#974)

Co-authored-by: Copilot <[email protected]>

Author: alexmojaki

logfire/_internal/integrations/mcp.py

@@ -3,20 +3,23 @@
 import functools
 from typing import TYPE_CHECKING, Any
 
-from mcp.shared.session import BaseSession, SendRequestT
-from mcp.types import CallToolRequest
+from mcp.client.session import ClientSession
+from mcp.shared.session import BaseSession
+from mcp.types import CallToolRequest, LoggingMessageNotification
+
+from logfire._internal.utils import handle_internal_errors
 
 if TYPE_CHECKING:
-    from logfire import Logfire
+    from logfire import LevelName, Logfire
 
 
 def instrument_mcp(logfire_instance: Logfire):
     logfire_instance = logfire_instance.with_settings(custom_scope_suffix='mcp')
 
-    original = BaseSession.send_request  # type: ignore
+    original_send_request = BaseSession.send_request  # type: ignore
 
-    @functools.wraps(original)  # type: ignore
-    async def send_request(self, request: SendRequestT, *args, **kwargs: Any):  # type: ignore
+    @functools.wraps(original_send_request)  # type: ignore
+    async def send_request(self: Any, request: Any, *args: Any, **kwargs: Any):
         attributes: dict[str, Any] = {
             'request': request,
             # https://opentelemetry.io/docs/specs/semconv/rpc/json-rpc/
@@ -35,8 +38,28 @@ async def send_request(self, request: SendRequestT, *args, **kwargs: Any):  # ty
                 span_name += f' {root.params.name}'
 
         with logfire_instance.span(span_name, **attributes) as span:
-            result = await original(self, request, *args, **kwargs)  # type: ignore
+            result = await original_send_request(self, request, *args, **kwargs)
             span.set_attribute('response', result)
-            return result  # type: ignore
+            return result
 
     BaseSession.send_request = send_request
+
+    original_received_notification = ClientSession._received_notification  # type: ignore
+
+    @functools.wraps(original_received_notification)
+    async def _received_notification(self: Any, notification: Any, *args: Any, **kwargs: Any):
+        with handle_internal_errors:
+            if isinstance(notification.root, LoggingMessageNotification):  # pragma: no branch
+                params = notification.root.params
+                level: LevelName
+                if params.level in ('critical', 'alert', 'emergency'):
+                    level = 'fatal'
+                else:
+                    level = params.level
+                span_name = 'MCP server log'
+                if params.logger:
+                    span_name += f' from {params.logger}'
+                logfire_instance.log(level, span_name, attributes=dict(data=params.data))
+        await original_received_notification(self, notification, *args, **kwargs)
+
+    ClientSession._received_notification = _received_notification  # type: ignore