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
@@ -182,6 +183,7 @@ def __init__(
             Awaitable[None],
         ]
         | None = None,
+        tracer_provider: TracerProvider | None = None,
     ):
         self.name = name
         self.version = version
@@ -197,6 +199,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
@@ -378,6 +381,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.networks import AnyUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from starlette.applications import Starlette
@@ -144,6 +145,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,
@@ -176,6 +178,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,126 @@
+import contextlib
+from collections.abc import Iterator
+from typing import Any
+
+from opentelemetry.trace import Span, SpanKind, StatusCode, Tracer
+
+from mcp import types
+from mcp.shared.exceptions import MCPError
+
+# OTel Semantic Conventions for MCP and GenAI
+# See: https://github.com/open-telemetry/semantic-conventions/blob/v1.40.0/docs/gen-ai/mcp.md
+MCP_METHOD_NAME = "mcp.method.name"
+MCP_RESOURCE_URI = "mcp.resource.uri"
+JSONRPC_REQUEST_ID = "jsonrpc.request.id"
+
+GEN_AI_TOOL_NAME = "gen_ai.tool.name"
+GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
+GEN_AI_PROMPT_NAME = "gen_ai.prompt.name"
+
+ERROR_TYPE = "error.type"
+RPC_RESPONSE_STATUS_CODE = "rpc.response.status_code"
+
+
+def _get_span_name(
+    request: types.ClientRequest | types.ServerRequest | types.ClientNotification | types.ServerNotification,
+) -> str:
+    """Computes the span name based on the request type and parameters."""
+    target = None
+    match request:
+        case types.CallToolRequest():
+            target = request.params.name
+        case types.GetPromptRequest():
+            target = request.params.name
+        case _:
+            pass
+
+    if target:
+        return f"{request.method} {target}"
+    return request.method
+
+
+def _get_common_attributes(
+    request: types.ClientRequest | types.ServerRequest | types.ClientNotification | types.ServerNotification,
+    *,
+    json_rpc_request_id: int | str | None = None,
+) -> dict[str, Any]:
+    """Computes common attributes for both client and server spans."""
+    attributes = {MCP_METHOD_NAME: request.method}
+
+    if json_rpc_request_id is not None:
+        attributes[JSONRPC_REQUEST_ID] = str(json_rpc_request_id)
+
+    match request:
+        case types.CallToolRequest():
+            attributes[GEN_AI_TOOL_NAME] = request.params.name
+            attributes[GEN_AI_OPERATION_NAME] = "execute_tool"
+        case types.GetPromptRequest():
+            attributes[GEN_AI_PROMPT_NAME] = request.params.name
+        case (
+            types.ReadResourceRequest()
+            | types.SubscribeRequest()
+            | types.UnsubscribeRequest()
+            | types.ResourceUpdatedNotification()
+        ):
+            attributes[MCP_RESOURCE_URI] = request.params.uri
+        case _:
+            pass
+    return attributes
+
+
+_ERROR_NAMES = {
+    types.INVALID_PARAMS: "invalid_params",
+    types.METHOD_NOT_FOUND: "method_not_found",
+    types.CONNECTION_CLOSED: "connection_closed",
+    types.REQUEST_TIMEOUT: "timeout",
+    types.PARSE_ERROR: "parse_error",
+    types.INTERNAL_ERROR: "internal_error",
+    types.INVALID_REQUEST: "invalid_request",
+    types.URL_ELICITATION_REQUIRED: "url_elicitation_required",
+}
+
+
+def _record_error_data(span: Span, e: types.ErrorData, record_status: bool = True) -> None:
+    """Record an MCP protocol error on the span set status
+
+    https://github.com/open-telemetry/semantic-conventions/blob/v1.40.0/docs/general/recording-errors.md
+    """
+    if not span.is_recording():
+        return
+
+    span.set_attribute(ERROR_TYPE, _ERROR_NAMES.get(e.code, str(e.code)))
+    span.set_attribute(RPC_RESPONSE_STATUS_CODE, str(e.code))
+    span.set_status(status=StatusCode.ERROR, description=e.message)
+
+
+@contextlib.contextmanager
+def mcp_client_span(
+    tracer: Tracer,
+    request: types.ClientRequest | types.ServerRequest | types.ClientNotification | types.ServerNotification,
+    *,
+    json_rpc_request_id: int | str | 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
+    """
+    span_name = _get_span_name(request)
+    attributes = _get_common_attributes(request, json_rpc_request_id=json_rpc_request_id)
+
+    reraise_exc = None
+    with tracer.start_as_current_span(
+        span_name,
+        kind=SpanKind.CLIENT,
+        attributes=attributes,
+        set_status_on_exception=False,
+    ) as span:
+        try:
+            yield span
+        except MCPError as mcp_error:
+            _record_error_data(span, mcp_error.error)
+            span.record_exception(mcp_error)
+            # re-raise outside of with block to avoid overwriting span status
+            reraise_exc = mcp_error
+
+    if reraise_exc:
+        raise reraise_exc

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.
@@ -268,10 +273,10 @@ async def send_request(
             # Store the callback for this request
             self._progress_callbacks[request_id] = progress_callback
 
-        # Propagate opentelemetry trace context
-        self._inject_otel_context(request_data)
+        async def make_request():
+            # 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`.