fix: retrieval of session ID during a request

Author: constantinius

sentry_sdk/integrations/mcp/__init__.py

@@ -27,15 +27,9 @@ def setup_once():
         Patches MCP server classes to instrument handler execution.
         """
         from sentry_sdk.integrations.mcp.lowlevel import patch_lowlevel_server
-        from sentry_sdk.integrations.mcp.transport import (
-            patch_streamable_http_transport,
-        )
 
         # Patch server classes to instrument handlers
         patch_lowlevel_server()
 
-        # Patch HTTP transport to track session IDs
-        patch_streamable_http_transport()
-
 
 __all__ = ["MCPIntegration"]

sentry_sdk/integrations/mcp/lowlevel.py

@@ -12,7 +12,7 @@
 from sentry_sdk.integrations.mcp import MCPIntegration
 from sentry_sdk.integrations.mcp.transport import (
     detect_mcp_transport_from_context,
-    mcp_session_id_ctx,
+    get_session_id_from_context,
 )
 from sentry_sdk.utils import safe_serialize
 
@@ -24,6 +24,28 @@
     from typing import Any, Callable
 
 
+def _get_request_context_data():
+    # type: () -> tuple[str | None, str | None]
+    """
+    Extract request ID and session ID from the MCP request context.
+
+    Returns:
+        Tuple of (request_id, session_id). Either value may be None if not available.
+    """
+    request_id = None  # type: str | None
+    session_id = None  # type: str | None
+
+    try:
+        ctx = request_ctx.get()
+        request_id = ctx.request_id
+        session_id = get_session_id_from_context(ctx)
+    except LookupError:
+        # No request context available
+        pass
+
+    return request_id, session_id
+
+
 def _get_span_config(handler_type, item_name):
     # type: (str, str) -> tuple[str, str, str, str | None]
     """
@@ -51,9 +73,15 @@ def _get_span_config(handler_type, item_name):
 
 
 def _set_span_input_data(
-    span, handler_name, span_data_key, mcp_method_name, arguments, request_id=None
+    span,
+    handler_name,
+    span_data_key,
+    mcp_method_name,
+    arguments,
+    request_id=None,
+    session_id=None,
 ):
-    # type: (Any, str, str, str, dict[str, Any], str | None) -> None
+    # type: (Any, str, str, str, dict[str, Any], str | None, str | None) -> None
     """Set input span data for MCP handlers."""
     # Set handler identifier
     span.set_data(span_data_key, handler_name)
@@ -69,19 +97,14 @@ def _set_span_input_data(
         # No request context available - likely stdio
         span.set_data(SPANDATA.MCP_TRANSPORT, "pipe")
 
-    # Extract session ID from context variable if available (HTTP transport)
-    try:
-        session_id = mcp_session_id_ctx.get()
-        if session_id:
-            span.set_data(SPANDATA.MCP_SESSION_ID, session_id)
-    except Exception:
-        # Session ID not available or transport module not imported
-        pass
-
     # Set request_id if provided
     if request_id:
         span.set_data(SPANDATA.MCP_REQUEST_ID, request_id)
 
+    # Set session_id if provided
+    if session_id:
+        span.set_data(SPANDATA.MCP_SESSION_ID, session_id)
+
     # Set request arguments (excluding common request context objects)
     for k, v in arguments.items():
         span.set_data(f"mcp.request.argument.{k}", safe_serialize(v))
@@ -262,13 +285,8 @@ async def _async_handler_wrapper(handler_type, func, original_args):
         name=span_name,
         origin=MCPIntegration.origin,
     ) as span:
-        # Get request ID from context
-        request_id = None
-        try:
-            ctx = request_ctx.get()
-            request_id = ctx.request_id
-        except LookupError:
-            pass
+        # Get request ID and session ID from context
+        request_id, session_id = _get_request_context_data()
 
         # Set input span data
         _set_span_input_data(
@@ -278,6 +296,7 @@ async def _async_handler_wrapper(handler_type, func, original_args):
             mcp_method_name,
             arguments,
             request_id,
+            session_id,
         )
 
         # For resources, extract and set protocol
@@ -329,13 +348,8 @@ def _sync_handler_wrapper(handler_type, func, original_args):
         name=span_name,
         origin=MCPIntegration.origin,
     ) as span:
-        # Get request ID from context
-        request_id = None
-        try:
-            ctx = request_ctx.get()
-            request_id = ctx.request_id
-        except LookupError:
-            pass
+        # Get request ID and session ID from context
+        request_id, session_id = _get_request_context_data()
 
         # Set input span data
         _set_span_input_data(
@@ -345,6 +359,7 @@ def _sync_handler_wrapper(handler_type, func, original_args):
             mcp_method_name,
             arguments,
             request_id,
+            session_id,
         )
 
         # For resources, extract and set protocol

sentry_sdk/integrations/mcp/transport.py

@@ -9,20 +9,15 @@
 - If there's an HTTP request context (SSE/WebSocket), transport is "tcp"
 - If there's no request context (stdio), transport is "pipe"
 
-Session ID tracking is done by patching the StreamableHTTPServerTransport to store
-the session ID in a context variable that can be accessed during handler execution.
+Session ID tracking is done by patching Server._run_request_handler to store a
+reference to the server instance in the request context. The session ID can then
+be retrieved from the server's transport when needed during handler execution.
 """
 
-import contextvars
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Optional, Any
-    from starlette.types import Receive, Scope, Send
-
-
-# Context variable to store the current MCP session ID
-mcp_session_id_ctx = contextvars.ContextVar("mcp_session_id", default=None)  # type: contextvars.ContextVar[Optional[str]]
 
 
 def detect_mcp_transport_from_context(request_ctx):
@@ -58,36 +53,32 @@ def detect_mcp_transport_from_context(request_ctx):
     return None
 
 
-def patch_streamable_http_transport():
-    # type: () -> None
+def get_session_id_from_context(request_ctx):
+    # type: (Any) -> Optional[str]
     """
-    Patches the StreamableHTTPServerTransport to store session IDs in context.
+    Extract session ID from the request context.
+
+    The session ID is sent by the client in the MCP-Session-Id header and is
+    available in the Starlette Request object stored in ctx.request.
 
-    This allows handler code to access the session ID via the mcp_session_id_ctx
-    context variable, regardless of whether it's the first request (where the
-    session ID hasn't been sent to the client yet) or a subsequent request.
+    Args:
+        request_ctx: The MCP request context object
+
+    Returns:
+        Session ID string if available, None otherwise
     """
     try:
-        from mcp.server.streamable_http import StreamableHTTPServerTransport
-    except ImportError:
-        # StreamableHTTP transport not available
-        return
-
-    original_handle_request = StreamableHTTPServerTransport.handle_request
-
-    async def patched_handle_request(self, scope, receive, send):
-        # type: (Any, Scope, Receive, Send) -> None
-        """Wrap handle_request to set session ID in context."""
-        # Store session ID in context variable before handling request
-        token = None
-        if hasattr(self, "mcp_session_id") and self.mcp_session_id:
-            token = mcp_session_id_ctx.set(self.mcp_session_id)
-
-        try:
-            await original_handle_request(self, scope, receive, send)
-        finally:
-            # Reset context after request
-            if token is not None:
-                mcp_session_id_ctx.reset(token)
-
-    StreamableHTTPServerTransport.handle_request = patched_handle_request  # type: ignore
+        # The Starlette Request object is stored in ctx.request
+        if hasattr(request_ctx, "request") and request_ctx.request is not None:
+            request = request_ctx.request
+
+            # Check if it's a Starlette Request with headers
+            if hasattr(request, "headers"):
+                # The session ID is sent in the mcp-session-id header
+                # MCP_SESSION_ID_HEADER = "mcp-session-id"
+                return request.headers.get("mcp-session-id")
+
+    except Exception:
+        pass
+
+    return None