add further trace propagation for responses

Author: liustve

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/instrumentation/mcp/instrumentation.py

@@ -1,24 +1,23 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-from dataclasses import dataclass
 import json
-from typing import Any, AsyncGenerator, Callable, Collection, Dict, Optional, Tuple, Union, cast
+from typing import Any, Callable, Collection, Dict, Optional, Tuple, Union
 
 from wrapt import register_post_import_hook, wrap_function_wrapper
 
 from opentelemetry import trace
-from opentelemetry.trace import SpanKind, Status, StatusCode
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
-from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.propagate import get_global_textmap
-
-from .version import __version__
+from opentelemetry.semconv.attributes.network_attributes import NetworkTransportValues
+from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace import INVALID_SPAN, SpanKind, Status, StatusCode
 
 from .semconv import (
-    MCPSpanAttributes,
     MCPMethodValue,
+    MCPSpanAttributes,
 )
+from .version import __version__
 
 
 class McpInstrumentor(BaseInstrumentor):
@@ -29,6 +28,7 @@ class McpInstrumentor(BaseInstrumentor):
     _DEFAULT_CLIENT_SPAN_NAME = "span.mcp.client"
     _DEFAULT_SERVER_SPAN_NAME = "span.mcp.server"
     _MCP_SESSION_ID_HEADER = "mcp-session-id"
+    _MCP_META_FIELD = "_meta"
 
     def __init__(self, **kwargs):
         super().__init__()
@@ -55,6 +55,14 @@ def _instrument(self, **kwargs: Any) -> None:
             ),
             "mcp.shared.session",
         )
+        register_post_import_hook(
+            lambda _: wrap_function_wrapper(
+                "mcp.shared.session",
+                "BaseSession._send_response",
+                self._wrap_send_response,
+            ),
+            "mcp.shared.session",
+        )
         register_post_import_hook(
             lambda _: wrap_function_wrapper(
                 "mcp.server.lowlevel.server",
@@ -75,6 +83,7 @@ def _instrument(self, **kwargs: Any) -> None:
     def _uninstrument(self, **kwargs: Any) -> None:
         unwrap("mcp.shared.session", "BaseSession.send_request")
         unwrap("mcp.shared.session", "BaseSession.send_notification")
+        unwrap("mcp.shared.session", "BaseSession._send_response")
         unwrap("mcp.server.lowlevel.server", "Server._handle_request")
         unwrap("mcp.server.lowlevel.server", "Server._handle_notification")
 
@@ -98,7 +107,7 @@ def _wrap_session_send(
              args: Positional arguments passed to the original send_request/send_notification method
              kwargs: Keyword arguments passed to the original send_request/send_notification method
         """
-        from mcp.types import ClientRequest, ClientNotification, ServerRequest, ServerNotification
+        from mcp.types import ClientNotification, ClientRequest, ServerNotification, ServerRequest
 
         async def async_wrapper():
             message: Optional[Union[ClientRequest, ClientNotification, ServerRequest, ServerNotification]] = (
@@ -120,14 +129,18 @@ async def async_wrapper():
 
             if "params" not in message_json:
                 message_json["params"] = {}
-            if "_meta" not in message_json["params"]:
-                message_json["params"]["_meta"] = {}
+            if self._MCP_META_FIELD not in message_json["params"]:
+                message_json["params"][self._MCP_META_FIELD] = {}
+
+            parent_ctx = None
+            if message_json["params"][self._MCP_META_FIELD]:
+                parent_ctx = self.propagators.extract(message_json["params"][self._MCP_META_FIELD])
 
-            with self.tracer.start_as_current_span(name=span_name, kind=span_kind) as span:
+            with self.tracer.start_as_current_span(name=span_name, kind=span_kind, context=parent_ctx) as span:
                 ctx = trace.set_span_in_context(span)
                 carrier = {}
                 self.propagators.inject(carrier=carrier, context=ctx)
-                message_json["params"]["_meta"].update(carrier)
+                message_json["params"][self._MCP_META_FIELD].update(carrier)
 
                 McpInstrumentor._generate_mcp_message_attrs(span, message, request_id)
 
@@ -227,11 +240,14 @@ async def _wrap_server_message_handler(
         with self.tracer.start_as_current_span(
             self._DEFAULT_SERVER_SPAN_NAME, kind=SpanKind.SERVER, context=parent_ctx
         ) as server_span:
-            
-            # Extract session ID if available
+
+            # session_id only exits if the transport protocol is Streamable HTTP
             session_id = self._extract_session_id(args)
             if session_id:
                 server_span.set_attribute(MCPSpanAttributes.MCP_SESSION_ID, session_id)
+                server_span.set_attribute(SpanAttributes.NETWORK_TRANSPORT, NetworkTransportValues.PIPE.value)
+            else:
+                server_span.set_attribute(SpanAttributes.NETWORK_TRANSPORT, NetworkTransportValues.TCP.value)
 
             self._generate_mcp_message_attrs(server_span, incoming_msg, request_id)
 
@@ -243,27 +259,60 @@ async def _wrap_server_message_handler(
                 server_span.set_status(Status(StatusCode.ERROR, str(e)))
                 server_span.record_exception(e)
                 raise
-    
+
+    async def _wrap_send_response(
+        self, wrapped: Callable, instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> Any:
+        """
+        Instruments BaseSession._send_response to propagate trace context into response.
+
+        Note: we do not need to generate another span for the reponse as it falls under
+              the _wrap_server_message_handler
+        """
+        response = args[1] if len(args) > 1 else kwargs.get("response", None)
+
+        if not response:
+            return await wrapped(*args, **kwargs)
+
+        current_span = trace.get_current_span()
+        if current_span is not INVALID_SPAN:
+            # Inject trace context into response
+            carrier = {}
+            self.propagators.inject(carrier=carrier, context=trace.set_span_in_context(current_span))
+
+            response_json = response.model_dump(by_alias=True, mode="json", exclude_none=True)
+
+            if self._MCP_META_FIELD not in response_json:
+                response_json[self._MCP_META_FIELD] = {}
+            response_json[self._MCP_META_FIELD].update(carrier)
+
+            modified_response = response.model_validate(response_json)
+
+            if len(args) > 1:
+                args = args[:1] + (modified_response,) + args[2:]
+            else:
+                kwargs["response"] = modified_response
+
+        return await wrapped(*args, **kwargs)
+
     def _extract_session_id(self, args: Tuple[Any, ...]) -> Optional[str]:
         """
         Extract session ID from server method arguments.
         """
         try:
-            from mcp.shared.session import RequestResponder  # pylint: disable=import-outside-toplevel
-            from mcp.shared.message import ServerMessageMetadata  # pylint: disable=import-outside-toplevel
-            
             message = args[0]
-            if isinstance(message, RequestResponder):
-                if message.message_metadata and isinstance(message.message_metadata, ServerMessageMetadata):
-                    request_context = message.message_metadata.request_context
-                    if request_context:
-                        headers = getattr(request_context, 'headers', None)
+            if hasattr(message, "message_metadata"):
+                message_metadata = message.message_metadata
+                if message.message_metadata and hasattr(message_metadata, "request_context"):
+                    request_context = message_metadata.request_context
+                    if request_context and hasattr(request_context, "headers"):
+                        headers = request_context.headers
                         if headers:
                             return headers.get(self._MCP_SESSION_ID_HEADER)
             return None
         except Exception:
             return None
-        
+
     @staticmethod
     def _generate_mcp_message_attrs(span: trace.Span, message, request_id: Optional[int]) -> None:
         import mcp.types as types  # pylint: disable=import-outside-toplevel,consider-using-from-import
@@ -326,4 +375,4 @@ def serialize(args: dict[str, Any]) -> str:
         try:
             return json.dumps(args)
         except Exception:
-            return ""
+            return "unknown_args"