Implement MCP client opentelemetry tracing

finish client spans

Inject tracer_provider into servers

Author: aabmass

src/mcp/client/client.py

@@ -6,6 +6,8 @@
 from dataclasses import KW_ONLY, dataclass, field
 from typing import Any
 
+from opentelemetry.trace import TracerProvider
+
 from mcp.client._memory import InMemoryTransport
 from mcp.client._transport import Transport
 from mcp.client.session import ClientSession, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
@@ -95,6 +97,8 @@ async def main():
     elicitation_callback: ElicitationFnT | None = None
     """Callback for handling elicitation requests."""
 
+    tracer_provider: TracerProvider | None = None
+
     _session: ClientSession | None = field(init=False, default=None)
     _exit_stack: AsyncExitStack | None = field(init=False, default=None)
     _transport: Transport = field(init=False)
@@ -126,6 +130,7 @@ async def __aenter__(self) -> Client:
                     message_handler=self.message_handler,
                     client_info=self.client_info,
                     elicitation_callback=self.elicitation_callback,
+                    tracer_provider=self.tracer_provider,
                 )
             )
 

src/mcp/client/session.py

@@ -5,6 +5,7 @@
 
 import anyio.lowlevel
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from opentelemetry.trace import TracerProvider
 from pydantic import TypeAdapter
 
 from mcp import types
@@ -121,8 +122,11 @@ def __init__(
         *,
         sampling_capabilities: types.SamplingCapability | None = None,
         experimental_task_handlers: ExperimentalTaskHandlers | None = None,
+        tracer_provider: TracerProvider | None = None,
     ) -> None:
-        super().__init__(read_stream, write_stream, read_timeout_seconds=read_timeout_seconds)
+        super().__init__(
+            read_stream, write_stream, read_timeout_seconds=read_timeout_seconds, tracer_provider=tracer_provider
+        )
         self._client_info = client_info or DEFAULT_CLIENT_INFO
         self._sampling_callback = sampling_callback or _default_sampling_callback
         self._sampling_capabilities = sampling_capabilities

src/mcp/server/lowlevel/server.py

@@ -46,6 +46,7 @@ async def main():
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from opentelemetry.trace import TracerProvider
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
@@ -184,6 +185,7 @@ def __init__(
             Awaitable[None],
         ]
         | None = None,
+        tracer_provider: TracerProvider | None = None,
     ):
         self.name = name
         self.version = version
@@ -199,6 +201,7 @@ def __init__(
         ] = {}
         self._experimental_handlers: ExperimentalHandlers[LifespanResultT] | None = None
         self._session_manager: StreamableHTTPSessionManager | None = None
+        self._tracer_provider = tracer_provider
         logger.debug("Initializing server %r", name)
 
         # Populate internal handler dicts from on_* kwargs
@@ -380,6 +383,7 @@ async def run(
                     write_stream,
                     initialization_options,
                     stateless=stateless,
+                    tracer_provider=self._tracer_provider,
                 )
             )
 

src/mcp/server/mcpserver/server.py

@@ -12,6 +12,7 @@
 
 import anyio
 import pydantic_core
+from opentelemetry.trace import TracerProvider
 from pydantic import BaseModel
 from pydantic.networks import AnyUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -146,6 +147,7 @@ def __init__(
         warn_on_duplicate_prompts: bool = True,
         lifespan: Callable[[MCPServer[LifespanResultT]], AbstractAsyncContextManager[LifespanResultT]] | None = None,
         auth: AuthSettings | None = None,
+        tracer_provider: TracerProvider | None = None,
     ):
         self.settings = Settings(
             debug=debug,
@@ -178,6 +180,7 @@ def __init__(
             # TODO(Marcelo): It seems there's a type mismatch between the lifespan type from an MCPServer and Server.
             # We need to create a Lifespan type that is a generic on the server type, like Starlette does.
             lifespan=(lifespan_wrapper(self, self.settings.lifespan) if self.settings.lifespan else default_lifespan),  # type: ignore
+            tracer_provider=tracer_provider,
         )
         # Validate auth configuration
         if self.settings.auth is not None:

src/mcp/server/session.py

@@ -34,6 +34,7 @@ async def handle_list_prompts(ctx: RequestContext, params) -> ListPromptsResult:
 import anyio
 import anyio.lowlevel
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from opentelemetry.trace import TracerProvider
 from pydantic import AnyUrl, TypeAdapter
 
 from mcp import types
@@ -83,8 +84,9 @@ def __init__(
         write_stream: MemoryObjectSendStream[SessionMessage],
         init_options: InitializationOptions,
         stateless: bool = False,
+        tracer_provider: TracerProvider | None = None,
     ) -> None:
-        super().__init__(read_stream, write_stream)
+        super().__init__(read_stream, write_stream, tracer_provider=tracer_provider)
         self._stateless = stateless
         self._initialization_state = (
             InitializationState.Initialized if stateless else InitializationState.NotInitialized

src/mcp/shared/_otel_utils.py

@@ -0,0 +1,63 @@
+import contextlib
+from collections.abc import Iterator
+
+from opentelemetry.trace import Span, SpanKind, Tracer
+
+from mcp import types
+from mcp.shared.exceptions import MCPError
+
+
+@contextlib.contextmanager
+def mcp_client_span(
+    tracer: Tracer,
+    request: types.ClientRequest | types.ServerRequest | types.ClientNotification | types.ServerNotification,
+    *,
+    json_rpc_request_id: int | None = None,
+) -> Iterator[Span]:
+    """Starts an MCP client span as current span
+
+    https://github.com/open-telemetry/semantic-conventions/blob/v1.40.0/docs/gen-ai/mcp.md#client
+    """
+    attributes = {"mcp.method.name": request.method}
+
+    # When omitted, the request is treated as a notification. Instrumentations SHOULD NOT
+    # capture this attribute when the id is null or omitted.
+    if json_rpc_request_id is not None:
+        attributes["jsonrpc.request.id"] = str(json_rpc_request_id)
+
+    target = None
+
+    match request:
+        case types.CallToolRequest():
+            target = request.params.name
+            attributes["gen_ai.tool.name"] = target
+            attributes["gen_ai.operation.name"] = "execute_tool"
+        case types.GetPromptRequest():
+            target = request.params.name
+            attributes["gen_ai.prompt.name"] = target
+        case (
+            types.ReadResourceRequest()
+            | types.SubscribeRequest()
+            | types.UnsubscribeRequest()
+            | types.ResourceUpdatedNotification()
+        ):
+            attributes["mcp.resource.uri"] = request.params.uri
+        case _:
+            pass
+
+    if target:
+        span_name = f"{request.method} {target}"
+    else:
+        span_name = request.method
+
+    with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT, attributes=attributes) as span:
+        try:
+            yield span
+        except MCPError as e:
+            if span.is_recording():
+                if e.code == types.REQUEST_TIMEOUT:
+                    span.set_attribute("error.type", "timeout")
+                else:
+                    span.set_attribute("error.type", str(e.code))
+                    span.set_attribute("rpc.response.status_code", str(e.code))
+            raise

src/mcp/shared/session.py

@@ -11,9 +11,11 @@
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from opentelemetry import context as otel_context
 from opentelemetry.propagate import extract, inject
+from opentelemetry.trace import TracerProvider, get_tracer
 from pydantic import BaseModel, TypeAdapter
 from typing_extensions import Self
 
+from mcp.shared._otel_utils import mcp_client_span
 from mcp.shared.exceptions import MCPError
 from mcp.shared.message import MessageMetadata, ServerMessageMetadata, SessionMessage
 from mcp.shared.response_router import ResponseRouter
@@ -190,6 +192,8 @@ def __init__(
         write_stream: MemoryObjectSendStream[SessionMessage],
         # If none, reading will never time out
         read_timeout_seconds: float | None = None,
+        *,
+        tracer_provider: TracerProvider | None = None,
     ) -> None:
         self._read_stream = read_stream
         self._write_stream = write_stream
@@ -200,6 +204,7 @@ def __init__(
         self._progress_callbacks = {}
         self._response_routers = []
         self._exit_stack = AsyncExitStack()
+        self._tracer = get_tracer("mcp", tracer_provider=tracer_provider)
 
     def add_response_router(self, router: ResponseRouter) -> None:
         """Register a response router to handle responses for non-standard requests.
@@ -256,22 +261,22 @@ async def send_request(
         response_stream, response_stream_reader = anyio.create_memory_object_stream[JSONRPCResponse | JSONRPCError](1)
         self._response_streams[request_id] = response_stream
 
-        # Set up progress token if progress callback is provided
-        request_data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
-        if progress_callback is not None:
-            # Use request_id as progress token
-            if "params" not in request_data:  # pragma: lax no cover
-                request_data["params"] = {}
-            if "_meta" not in request_data["params"]:  # pragma: lax no cover
-                request_data["params"]["_meta"] = {}
-            request_data["params"]["_meta"]["progressToken"] = request_id
-            # Store the callback for this request
-            self._progress_callbacks[request_id] = progress_callback
+        async def make_request() -> ReceiveResultT:
+            # Set up progress token if progress callback is provided
+            request_data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
+            if progress_callback is not None:
+                # Use request_id as progress token
+                if "params" not in request_data:  # pragma: lax no cover
+                    request_data["params"] = {}
+                if "_meta" not in request_data["params"]:  # pragma: lax no cover
+                    request_data["params"]["_meta"] = {}
+                request_data["params"]["_meta"]["progressToken"] = request_id
+                # Store the callback for this request
+                self._progress_callbacks[request_id] = progress_callback
+
+            # Propagate opentelemetry trace context
+            self._inject_otel_context(request_data)
 
-        # Propagate opentelemetry trace context
-        self._inject_otel_context(request_data)
-
-        try:
             jsonrpc_request = JSONRPCRequest(jsonrpc="2.0", id=request_id, **request_data)
             await self._write_stream.send(SessionMessage(message=jsonrpc_request, metadata=metadata))
 
@@ -291,6 +296,9 @@ async def send_request(
             else:
                 return result_type.model_validate(response_or_error.result, by_name=False)
 
+        try:
+            with mcp_client_span(self._tracer, request, json_rpc_request_id=request_id):
+                return await make_request()
         finally:
             self._response_streams.pop(request_id, None)
             self._progress_callbacks.pop(request_id, None)
@@ -315,7 +323,9 @@ async def send_notification(
             message=jsonrpc_notification,
             metadata=ServerMessageMetadata(related_request_id=related_request_id) if related_request_id else None,
         )
-        await self._write_stream.send(session_message)
+
+        with mcp_client_span(self._tracer, notification):
+            await self._write_stream.send(session_message)
 
     def _inject_otel_context(self, request: dict[str, Any]) -> None:
         """Propagate OpenTelemetry context in `_meta`.