Skip to content
150 changes: 103 additions & 47 deletions sentry_sdk/integrations/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
try:
from mcp.server.lowlevel import Server # type: ignore[import-not-found]
from mcp.server.lowlevel.server import request_ctx # type: ignore[import-not-found]
from mcp.server.streamable_http import StreamableHTTPServerTransport # type: ignore[import-not-found]
except ImportError:
raise DidNotEnable("MCP SDK not installed")

Expand All @@ -31,7 +32,9 @@


if TYPE_CHECKING:
from typing import Any, Callable, Optional
from typing import Any, Callable, Optional, Tuple

from starlette.types import Receive, Scope, Send # type: ignore[import-not-found]


class MCPIntegration(Integration):
Expand All @@ -54,11 +57,34 @@ def setup_once() -> None:
Patches MCP server classes to instrument handler execution.
"""
_patch_lowlevel_server()
_patch_handle_request()

if FastMCP is not None:
_patch_fastmcp()


def _get_active_http_scopes() -> (
"Optional[Tuple[Optional[sentry_sdk.Scope], Optional[sentry_sdk.Scope]]]"
):
try:
ctx = request_ctx.get()
except LookupError:
return None

if (
ctx is None
or not hasattr(ctx, "request")
or ctx.request is None
or "state" not in ctx.request.scope
):
return None

return (
ctx.request.scope["state"].get("sentry_sdk.current_scope"),
ctx.request.scope["state"].get("sentry_sdk.isolation_scope"),
)


def _get_request_context_data() -> "tuple[Optional[str], Optional[str], str]":
"""
Extract request ID, session ID, and MCP transport type from the request context.
Expand Down Expand Up @@ -381,56 +407,67 @@ async def _async_handler_wrapper(
result_data_key,
) = _prepare_handler_data(handler_type, original_args, original_kwargs)

# Start span and execute
with get_start_span_function()(
op=OP.MCP_SERVER,
name=span_name,
origin=MCPIntegration.origin,
) as span:
# Get request ID, session ID, and transport from context
request_id, session_id, mcp_transport = _get_request_context_data()

# Set input span data
_set_span_input_data(
span,
handler_name,
span_data_key,
mcp_method_name,
arguments,
request_id,
session_id,
mcp_transport,
)
scopes = _get_active_http_scopes()

# For resources, extract and set protocol
if handler_type == "resource":
if original_args:
uri = original_args[0]
else:
uri = original_kwargs.get("uri")
if scopes is None:
current_scope = None
isolation_scope = None
else:
current_scope, isolation_scope = scopes

protocol = None
if hasattr(uri, "scheme"):
protocol = uri.scheme
elif handler_name and "://" in handler_name:
protocol = handler_name.split("://")[0]
if protocol:
span.set_data(SPANDATA.MCP_RESOURCE_PROTOCOL, protocol)
# Get request ID, session ID, and transport from context
request_id, session_id, mcp_transport = _get_request_context_data()

try:
# Execute the async handler
if self is not None:
original_args = (self, *original_args)
result = await func(*original_args, **original_kwargs)
except Exception as e:
# Set error flag for tools
if handler_type == "tool":
span.set_data(SPANDATA.MCP_TOOL_RESULT_IS_ERROR, True)
sentry_sdk.capture_exception(e)
raise
# Start span and execute
with sentry_sdk.scope.use_isolation_scope(isolation_scope):
with sentry_sdk.scope.use_scope(current_scope):
with get_start_span_function()(
op=OP.MCP_SERVER,
name=span_name,
origin=MCPIntegration.origin,
) as span:
# Set input span data
_set_span_input_data(
span,
handler_name,
span_data_key,
mcp_method_name,
arguments,
request_id,
session_id,
mcp_transport,
)

# For resources, extract and set protocol
if handler_type == "resource":
if original_args:
uri = original_args[0]
else:
uri = original_kwargs.get("uri")

protocol = None
if hasattr(uri, "scheme"):
protocol = uri.scheme
elif handler_name and "://" in handler_name:
protocol = handler_name.split("://")[0]
if protocol:
span.set_data(SPANDATA.MCP_RESOURCE_PROTOCOL, protocol)

try:
# Execute the async handler
if self is not None:
original_args = (self, *original_args)
result = await func(*original_args, **original_kwargs)
except Exception as e:
# Set error flag for tools
if handler_type == "tool":
span.set_data(SPANDATA.MCP_TOOL_RESULT_IS_ERROR, True)
sentry_sdk.capture_exception(e)
raise

_set_span_output_data(span, result, result_data_key, handler_type)

_set_span_output_data(span, result, result_data_key, handler_type)
return result
return result
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sync handlers missing HTTP scope nesting

Medium Severity

The _async_handler_wrapper was updated to use use_isolation_scope and use_scope context managers for nesting MCP spans under HTTP transactions, but _sync_handler_wrapper was not similarly updated. Sync MCP handlers using HTTP-based transports (like SSE) will not have their spans properly nested under the HTTP transaction, while async handlers will. This creates inconsistent behavior depending on whether a handler is sync or async, even though both can be used with HTTP transport.

Additional Locations (1)

Fix in Cursor Fix in Web

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The _sync_handler_wrapper probably shouldn't exist in the first place. The decorators we are wrapping are for async functions only.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it a lot of extra work to update the sync wrapper as well? We should investigate if we need it at some point (not in this PR), but as long as we're not sure if it's actually unused, (and if it's not too much extra effort,) we might want to update it, too.



def _sync_handler_wrapper(
Expand Down Expand Up @@ -618,6 +655,25 @@ def patched_read_resource(
Server.read_resource = patched_read_resource


def _patch_handle_request() -> None:
original_handle_request = StreamableHTTPServerTransport.handle_request

@wraps(original_handle_request)
async def patched_handle_request(
self: "StreamableHTTPServerTransport",
scope: "Scope",
receive: "Receive",
send: "Send",
) -> None:
scope.setdefault("state", {})["sentry_sdk.isolation_scope"] = (
sentry_sdk.get_isolation_scope()
)
scope["state"]["sentry_sdk.current_scope"] = sentry_sdk.get_current_scope()
await original_handle_request(self, scope, receive, send)

StreamableHTTPServerTransport.handle_request = patched_handle_request


def _patch_fastmcp() -> None:
"""
Patches the standalone fastmcp package's FastMCP class.
Expand Down
13 changes: 11 additions & 2 deletions sentry_sdk/scope.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@

F = TypeVar("F", bound=Callable[..., Any])
T = TypeVar("T")
S = TypeVar("S", bound=Optional["Scope"])


# Holds data that will be added to **all** events sent by this process.
Expand Down Expand Up @@ -1786,7 +1787,7 @@ def new_scope() -> "Generator[Scope, None, None]":


@contextmanager
def use_scope(scope: "Scope") -> "Generator[Scope, None, None]":
def use_scope(scope: "S") -> "Generator[S, None, None]":
"""
.. versionadded:: 2.0.0

Expand All @@ -1808,6 +1809,10 @@ def use_scope(scope: "Scope") -> "Generator[Scope, None, None]":
sentry_sdk.capture_message("hello, again") # will NOT include `color` tag.

"""
if scope is None:
yield scope
return

# set given scope as current scope
token = _current_scope.set(scope)

Expand Down Expand Up @@ -1871,7 +1876,7 @@ def isolation_scope() -> "Generator[Scope, None, None]":


@contextmanager
def use_isolation_scope(isolation_scope: "Scope") -> "Generator[Scope, None, None]":
def use_isolation_scope(isolation_scope: "S") -> "Generator[S, None, None]":
"""
.. versionadded:: 2.0.0

Expand All @@ -1892,6 +1897,10 @@ def use_isolation_scope(isolation_scope: "Scope") -> "Generator[Scope, None, Non
sentry_sdk.capture_message("hello, again") # will NOT include `color` tag.

"""
if isolation_scope is None:
yield isolation_scope
return

# fork current scope
current_scope = Scope.get_current_scope()
forked_current_scope = current_scope.fork()
Expand Down
1 change: 1 addition & 0 deletions tests/integrations/fastmcp/test_fastmcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,7 @@ class MockHTTPRequest:
def __init__(self, session_id=None, transport="http"):
self.headers = {}
self.query_params = {}
self.scope = {}

if transport == "sse":
# SSE transport uses query parameter
Expand Down
Loading
Loading