diff --git a/sentry_sdk/integrations/mcp.py b/sentry_sdk/integrations/mcp.py index 837f896443..2f00b65fd7 100644 --- a/sentry_sdk/integrations/mcp.py +++ b/sentry_sdk/integrations/mcp.py @@ -57,6 +57,7 @@ if TYPE_CHECKING: from typing import Any, Awaitable, Callable, Optional, Union + from mcp.server.context import CallNext, HandlerResult from mcp_types import ( CallToolResult, GetPromptResult, @@ -268,41 +269,6 @@ def _extract_text_from_content_blocks(content_blocks: "Any") -> "Any": return " ".join(texts) if texts else content_blocks -# Handler data preparation and wrapping - - -def _is_v2_context(original_args: "tuple[Any, ...]") -> bool: - """Check if original_args contains a v2 ServerRequestContext as the first element.""" - return ( - ServerRequestContext is not None - and bool(original_args) - and isinstance(original_args[0], ServerRequestContext) - ) - - -def _extract_handler_data_from_params( - handler_type: str, - params: "Any", -) -> "tuple[str, dict[str, Any]]": - """ - Extract handler name and arguments from a v2 typed params object. - - In MCP SDK v2, handlers receive (ctx, params) where params is a typed - Pydantic model (CallToolRequestParams, GetPromptRequestParams, etc.). - """ - if handler_type == "tool": - handler_name = getattr(params, "name", "unknown") - arguments = getattr(params, "arguments", None) or {} - elif handler_type == "prompt": - handler_name = getattr(params, "name", "unknown") - arguments = getattr(params, "arguments", None) or {} - else: # resource - handler_name = str(getattr(params, "uri", "unknown")) - arguments = {} - - return handler_name, arguments - - def _extract_handler_data_from_args( handler_type: str, original_args: "tuple[Any, ...]", @@ -376,33 +342,17 @@ async def _tool_handler_wrapper( if original_kwargs is None: original_kwargs = {} - # Detect v1 vs v2: MCP SDK v2 passes (ServerRequestContext, params) to handlers - ctx: "Optional[Any]" = None - params: "Optional[Any]" = None - if ( - ServerRequestContext is not None - and original_args - and isinstance(original_args[0], ServerRequestContext) - ): - ctx = original_args[0] - if len(original_args) > 1: - params = original_args[1] - handler_name, arguments = _extract_handler_data_from_params("tool", params) - else: - handler_name = "unknown" - arguments = {} - else: - handler_name, arguments = _extract_handler_data_from_args( - "tool", original_args, original_kwargs - ) + handler_name, arguments = _extract_handler_data_from_args( + "tool", original_args, original_kwargs + ) # Get request ID, session ID, and transport from context - request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx) + request_id, session_id, mcp_transport = _get_request_context_data() span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) # Start span and execute - with _active_http_scopes(ctx=ctx): + with _active_http_scopes(): span_mgr: "Union[Span, StreamedSpan]" if span_streaming: span_mgr = sentry_sdk.traces.start_span( @@ -474,6 +424,103 @@ async def _tool_handler_wrapper( return result +async def _instrument_v2_tool_call( + ctx: "ServerRequestContext[Any, Any]", + call_next: "CallNext", +) -> "HandlerResult": + """ + Instrument a tool call as observed by the MCP Server middleware. + Creates and manages the MCP span and attaches all attributes on the span. + """ + if ctx.params is None or ctx.params.get("name") is None: + return await call_next(ctx) + + handler_name = ctx.params["name"] + arguments = ctx.params.get("arguments") + if arguments is None: + arguments = {} + + # Get request ID, session ID, and transport from context + request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx) + + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + + # Start span and execute + with _active_http_scopes(ctx=ctx): + span_mgr: "Union[Span, StreamedSpan]" + if span_streaming: + span_mgr = sentry_sdk.traces.start_span( + name=f"tools/call {handler_name}", + attributes={ + "sentry.op": OP.MCP_SERVER, + "sentry.origin": MCPIntegration.origin, + }, + ) + else: + span_mgr = get_start_span_function()( + op=OP.MCP_SERVER, + name=f"tools/call {handler_name}", + origin=MCPIntegration.origin, + ) + + with span_mgr as span: + # Set input span data + _set_span_input_data( + span, + handler_name, + SPANDATA.MCP_TOOL_NAME, + "tools/call", + arguments, + request_id, + session_id, + mcp_transport, + ) + + try: + result = await call_next(ctx) + + except Exception as e: + # Set error flag for tools + _set_span_data_attribute(span, SPANDATA.MCP_TOOL_RESULT_IS_ERROR, True) + sentry_sdk.capture_exception(e) + raise + + if result is None: + return result + + # Get integration to check PII settings + integration = sentry_sdk.get_client().get_integration(MCPIntegration) + if integration is None: + return result + + # Check if we should include sensitive data + should_include_data = ( + should_send_default_pii() and integration.include_prompts + ) + + result_content = result + if isinstance(result, dict) and "structuredContent" in result: + result_content = result["structuredContent"] + elif isinstance(result, dict) and isinstance(result.get("content"), list): + result_content = _extract_text_from_content_blocks(result["content"]) + + if result_content is not None and should_include_data: + _set_span_data_attribute( + span, + SPANDATA.MCP_TOOL_RESULT_CONTENT, + safe_serialize(result_content), + ) + # Set content count if result is a dict + if isinstance(result_content, dict): + _set_span_data_attribute( + span, + SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT, + len(result_content), + ) + + return result + + async def _prompt_handler_wrapper( func: "Callable[..., Awaitable[Union[GetPromptResult, InputRequiredResult]]]", original_args: "tuple[Any, ...]", @@ -494,33 +541,17 @@ async def _prompt_handler_wrapper( if original_kwargs is None: original_kwargs = {} - # Detect v1 vs v2: MCP SDK v2 passes (ServerRequestContext, params) to handlers - ctx: "Optional[Any]" = None - if ( - ServerRequestContext is not None - and original_args - and isinstance(original_args[0], ServerRequestContext) - ): - ctx = original_args[0] - if len(original_args) > 1: - handler_name, arguments = _extract_handler_data_from_params( - "prompt", original_args[1] - ) - else: - handler_name = "unknown" - arguments = {} - else: - handler_name, arguments = _extract_handler_data_from_args( - "prompt", original_args, original_kwargs - ) + handler_name, arguments = _extract_handler_data_from_args( + "prompt", original_args, original_kwargs + ) # Get request ID, session ID, and transport from context - request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx) + request_id, session_id, mcp_transport = _get_request_context_data() span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) # Start span and execute - with _active_http_scopes(ctx=ctx): + with _active_http_scopes(): span_mgr: "Union[Span, StreamedSpan]" if span_streaming: span_mgr = sentry_sdk.traces.start_span( @@ -642,6 +673,143 @@ async def _prompt_handler_wrapper( return result +async def _instrument_v2_prompt_get( + ctx: "ServerRequestContext[Any, Any]", + call_next: "CallNext", +) -> "HandlerResult": + """ + Instrument a prompt retrieval as observed by the MCP Server middleware. + Creates and manages the MCP span and attaches all attributes on the span. + """ + if ctx.params is None or ctx.params.get("name") is None: + return await call_next(ctx) + + handler_name = ctx.params["name"] + arguments = ctx.params.get("arguments") + if arguments is None: + arguments = {} + + # Get request ID, session ID, and transport from context + request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx) + + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + + # Start span and execute + with _active_http_scopes(ctx=ctx): + span_mgr: "Union[Span, StreamedSpan]" + if span_streaming: + span_mgr = sentry_sdk.traces.start_span( + name=f"prompts/get {handler_name}", + attributes={ + "sentry.op": OP.MCP_SERVER, + "sentry.origin": MCPIntegration.origin, + }, + ) + else: + span_mgr = get_start_span_function()( + op=OP.MCP_SERVER, + name=f"prompts/get {handler_name}", + origin=MCPIntegration.origin, + ) + + with span_mgr as span: + # Set input span data + _set_span_input_data( + span, + handler_name, + SPANDATA.MCP_PROMPT_NAME, + "prompts/get", + arguments, + request_id, + session_id, + mcp_transport, + ) + + try: + result = await call_next(ctx) + except Exception as e: + sentry_sdk.capture_exception(e) + raise + + if result is None: + return result + + # Get integration to check PII settings + integration = sentry_sdk.get_client().get_integration(MCPIntegration) + if integration is None: + return result + + # Check if we should include sensitive data + should_include_data = ( + should_send_default_pii() and integration.include_prompts + ) + + # For prompts, count messages and set role/content only for single-message prompts + try: + messages: "Optional[list[str]]" = None + message_count = 0 + + # Check if result has messages attribute (GetPromptResult) + if hasattr(result, "messages") and result.messages: + messages = result.messages + message_count = len(messages) + # Also check if result is a dict with messages + elif isinstance(result, dict) and result.get("messages"): + messages = result["messages"] + message_count = len(messages) + + # Always set message count if we found messages + if message_count > 0: + _set_span_data_attribute( + span, SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT, message_count + ) + + # Only set role and content for single-message prompts if PII is allowed + if message_count == 1 and should_include_data and messages: + first_message = messages[0] + # Extract role + role = None + if hasattr(first_message, "role"): + role = first_message.role + elif isinstance(first_message, dict) and "role" in first_message: + role = first_message["role"] + + if role: + _set_span_data_attribute( + span, SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE, role + ) + + # Extract content text + content_text = None + if hasattr(first_message, "content"): + msg_content = first_message.content + # Content can be a TextContent object or similar + if hasattr(msg_content, "text"): + content_text = msg_content.text + elif isinstance(msg_content, dict) and "text" in msg_content: + content_text = msg_content["text"] + elif isinstance(msg_content, str): + content_text = msg_content + elif isinstance(first_message, dict) and "content" in first_message: + msg_content = first_message["content"] + if isinstance(msg_content, dict) and "text" in msg_content: + content_text = msg_content["text"] + elif isinstance(msg_content, str): + content_text = msg_content + + if content_text: + _set_span_data_attribute( + span, + SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT, + content_text, + ) + except Exception: + # Silently ignore if we can't extract message info + pass + + return result + + async def _resource_handler_wrapper( func: "Callable[..., Awaitable[Union[ReadResourceResult, InputRequiredResult]]]", original_args: "tuple[Any, ...]", @@ -662,35 +830,17 @@ async def _resource_handler_wrapper( if original_kwargs is None: original_kwargs = {} - # Detect v1 vs v2: MCP SDK v2 passes (ServerRequestContext, params) to handlers - ctx: "Optional[Any]" = None - params: "Optional[Any]" = None - if ( - ServerRequestContext is not None - and original_args - and isinstance(original_args[0], ServerRequestContext) - ): - ctx = original_args[0] - if len(original_args) > 1: - params = original_args[1] - handler_name, arguments = _extract_handler_data_from_params( - "resource", params - ) - else: - handler_name = "unknown" - arguments = {} - else: - handler_name, arguments = _extract_handler_data_from_args( - "resource", original_args, original_kwargs - ) + handler_name, arguments = _extract_handler_data_from_args( + "resource", original_args, original_kwargs + ) # Get request ID, session ID, and transport from context - request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx) + request_id, session_id, mcp_transport = _get_request_context_data() span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) # Start span and execute - with _active_http_scopes(ctx=ctx): + with _active_http_scopes(): span_mgr: "Union[Span, StreamedSpan]" if span_streaming: span_mgr = sentry_sdk.traces.start_span( @@ -720,16 +870,10 @@ async def _resource_handler_wrapper( mcp_transport, ) - uri = None - if params is not None: - uri = getattr(params, "uri", None) - - # v1 scenario - if ServerRequestContext is None: - if original_args: - uri = original_args[0] - else: - uri = original_kwargs.get("uri") + if original_args: + uri = original_args[0] + else: + uri = original_kwargs.get("uri") protocol = None if uri is not None and hasattr(uri, "scheme"): @@ -755,6 +899,77 @@ async def _resource_handler_wrapper( return result +async def _instrument_v2_resource_read( + ctx: "ServerRequestContext[Any, Any]", + call_next: "CallNext", +) -> "HandlerResult": + """ + Instrument getting a resource as observed by the MCP Server middleware. + Creates and manages the MCP span and attaches all attributes on the span. + """ + if ctx.params is None or ctx.params.get("uri") is None: + return await call_next(ctx) + + handler_name = ctx.params["uri"] + + # Get request ID, session ID, and transport from context + request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx) + + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + + # Start span and execute + with _active_http_scopes(ctx=ctx): + span_mgr: "Union[Span, StreamedSpan]" + if span_streaming: + span_mgr = sentry_sdk.traces.start_span( + name=f"resources/read {handler_name}", + attributes={ + "sentry.op": OP.MCP_SERVER, + "sentry.origin": MCPIntegration.origin, + }, + ) + else: + span_mgr = get_start_span_function()( + op=OP.MCP_SERVER, + name=f"resources/read {handler_name}", + origin=MCPIntegration.origin, + ) + + with span_mgr as span: + # Set input span data + _set_span_input_data( + span, + handler_name, + SPANDATA.MCP_RESOURCE_URI, + "resources/read", + {}, + request_id, + session_id, + mcp_transport, + ) + + uri = None + if ctx.params is not None: + uri = getattr(ctx.params, "uri", None) + + protocol = None + if uri is not None and hasattr(uri, "scheme"): + protocol = uri.scheme + elif handler_name and "://" in handler_name: + protocol = handler_name.split("://")[0] + if protocol: + _set_span_data_attribute(span, SPANDATA.MCP_RESOURCE_PROTOCOL, protocol) + + try: + result = await call_next(ctx) + + except Exception as e: + sentry_sdk.capture_exception(e) + raise + + return result + + def _patch_lowlevel_server() -> None: """ Patches the mcp.server.lowlevel.Server class to instrument handler execution. @@ -828,6 +1043,21 @@ async def wrapper(*args: "Any") -> "Any": Server.read_resource = patched_read_resource # type: ignore[attr-defined] +async def _sentry_middleware( + ctx: "ServerRequestContext[Any, Any]", call_next: "CallNext" +) -> "HandlerResult": + if ctx.method == "tools/call": + return await _instrument_v2_tool_call(ctx, call_next) + + if ctx.method == "prompts/get": + return await _instrument_v2_prompt_get(ctx, call_next) + + if ctx.method == "resources/read": + return await _instrument_v2_resource_read(ctx, call_next) + + return await call_next(ctx) + + def _patch_lowlevel_server_v2() -> None: """Patches the v2 Server to wrap tool/prompt/resource handlers. @@ -840,113 +1070,11 @@ def _patch_lowlevel_server_v2() -> None: @wraps(original_init) def patched_init(self: "Server", *args: "Any", **kwargs: "Any") -> None: - on_tool_call = kwargs.get("on_call_tool") - if on_tool_call is not None and not getattr( - on_tool_call, "__sentry_mcp_wrapped__", False - ): - - @wraps(on_tool_call) - async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": - return await _tool_handler_wrapper( - on_tool_call, args, kwargs, force_await=False - ) - - wrapper.__sentry_mcp_wrapped__ = True # type: ignore[attr-defined] - kwargs["on_call_tool"] = wrapper - - on_get_prompt = kwargs.get("on_get_prompt") - if on_get_prompt is not None and not getattr( - on_get_prompt, "__sentry_mcp_wrapped__", False - ): - - @wraps(on_get_prompt) - async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": - return await _prompt_handler_wrapper( - on_get_prompt, args, kwargs, force_await=False - ) - - wrapper.__sentry_mcp_wrapped__ = True # type: ignore[attr-defined] - kwargs["on_get_prompt"] = wrapper - - on_read_resource = kwargs.get("on_read_resource") - if on_read_resource is not None and not getattr( - on_read_resource, "__sentry_mcp_wrapped__", False - ): - - @wraps(on_read_resource) - async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": - return await _resource_handler_wrapper( - on_read_resource, args, kwargs, force_await=False - ) - - wrapper.__sentry_mcp_wrapped__ = True # type: ignore[attr-defined] - kwargs["on_read_resource"] = wrapper - original_init(self, *args, **kwargs) + self.middleware.append(_sentry_middleware) Server.__init__ = patched_init # type: ignore[method-assign] - original_add_request_handler = Server.add_request_handler - - def patched_add_request_handler( - self: "Server", - method: str, - params_type: "Any", - handler: "Callable[..., Any]", - *args: "Any", - **kwargs: "Any", - ) -> None: - if getattr(handler, "__sentry_mcp_wrapped__", False): - return original_add_request_handler( - self, method, params_type, handler, *args, **kwargs - ) - - if method == "tools/call": - - @wraps(handler) - async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": - return await _tool_handler_wrapper( - handler, args, kwargs, force_await=False - ) - - wrapper.__sentry_mcp_wrapped__ = True # type: ignore[attr-defined] - - return original_add_request_handler( - self, method, params_type, wrapper, *args, **kwargs - ) - if method == "prompts/get": - - @wraps(handler) - async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": - return await _prompt_handler_wrapper( - handler, args, kwargs, force_await=False - ) - - wrapper.__sentry_mcp_wrapped__ = True # type: ignore[attr-defined] - - return original_add_request_handler( - self, method, params_type, wrapper, *args, **kwargs - ) - if method == "resources/read": - - @wraps(handler) - async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": - return await _resource_handler_wrapper( - handler, args, kwargs, force_await=False - ) - - wrapper.__sentry_mcp_wrapped__ = True # type: ignore[attr-defined] - - return original_add_request_handler( - self, method, params_type, wrapper, *args, **kwargs - ) - - original_add_request_handler( - self, method, params_type, handler, *args, **kwargs - ) - - Server.add_request_handler = patched_add_request_handler # type: ignore[method-assign] - def _patch_handle_request() -> None: original_handle_request = StreamableHTTPServerTransport.handle_request diff --git a/tests/integrations/mcp/test_mcp.py b/tests/integrations/mcp/test_mcp.py index 8fe82ee0be..3eb4155c15 100644 --- a/tests/integrations/mcp/test_mcp.py +++ b/tests/integrations/mcp/test_mcp.py @@ -134,18 +134,7 @@ def __init__(self, text): def test_integration_patches_server(sentry_init): """Test that MCPIntegration patches the Server class""" - if IS_MCP_V2: - original_add_request_handler = Server.add_request_handler - original_init = Server.__init__ - - sentry_init( - integrations=[MCPIntegration()], - traces_sample_rate=1.0, - ) - - assert Server.add_request_handler is not original_add_request_handler - assert Server.__init__ is not original_init - else: + if not IS_MCP_V2: original_call_tool = Server.call_tool original_get_prompt = Server.get_prompt original_read_resource = Server.read_resource