fix: flatten React Router structure — no nested Routes, each page is a top-level route
This commit is contained in:
enzo
@@ -0,0 +1,880 @@
|
||||
"""
|
||||
Sentry integration for MCP (Model Context Protocol) servers.
|
||||
|
||||
This integration instruments MCP servers to create spans for tool, prompt,
|
||||
and resource handler execution, and captures errors that occur during execution.
|
||||
|
||||
Supports the low-level `mcp.server.lowlevel.Server` API.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from functools import wraps
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import sentry_sdk
|
||||
from sentry_sdk.ai.utils import _set_span_data_attribute, get_start_span_function
|
||||
from sentry_sdk.consts import OP, SPANDATA
|
||||
from sentry_sdk.integrations import DidNotEnable, Integration
|
||||
from sentry_sdk.scope import should_send_default_pii
|
||||
from sentry_sdk.traces import StreamedSpan
|
||||
from sentry_sdk.tracing_utils import has_span_streaming_enabled
|
||||
from sentry_sdk.utils import nullcontext, package_version, safe_serialize
|
||||
|
||||
MCP_PACKAGE_VERSION = package_version("mcp")
|
||||
|
||||
try:
|
||||
from mcp.server.lowlevel import Server
|
||||
from mcp.server.streamable_http import (
|
||||
StreamableHTTPServerTransport,
|
||||
)
|
||||
|
||||
if MCP_PACKAGE_VERSION and MCP_PACKAGE_VERSION < (2, 0, 0):
|
||||
from mcp.server.lowlevel.server import ( # type: ignore[attr-defined]
|
||||
request_ctx,
|
||||
)
|
||||
except ImportError:
|
||||
raise DidNotEnable("MCP SDK not installed")
|
||||
|
||||
try:
|
||||
from fastmcp import FastMCP # type: ignore[import-not-found]
|
||||
except ImportError:
|
||||
FastMCP = None
|
||||
|
||||
if MCP_PACKAGE_VERSION and MCP_PACKAGE_VERSION >= (2, 0, 0):
|
||||
try:
|
||||
from mcp.server.context import (
|
||||
ServerRequestContext,
|
||||
)
|
||||
except ImportError:
|
||||
ServerRequestContext = None # type: ignore[assignment,misc]
|
||||
else:
|
||||
ServerRequestContext = None # type: ignore[assignment,misc]
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any, Callable, ContextManager, Optional, Tuple, Union
|
||||
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from sentry_sdk.traces import StreamedSpan
|
||||
from sentry_sdk.tracing import Span
|
||||
|
||||
|
||||
class MCPIntegration(Integration):
|
||||
identifier = "mcp"
|
||||
origin = "auto.ai.mcp"
|
||||
|
||||
def __init__(self, include_prompts: bool = True) -> None:
|
||||
"""
|
||||
Initialize the MCP integration.
|
||||
|
||||
Args:
|
||||
include_prompts: Whether to include prompts (tool results and prompt content)
|
||||
in span data. Requires send_default_pii=True. Default is True.
|
||||
"""
|
||||
self.include_prompts = include_prompts
|
||||
|
||||
@staticmethod
|
||||
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(
|
||||
ctx: "Optional[Any]" = None,
|
||||
) -> "Optional[Tuple[Optional[sentry_sdk.Scope], Optional[sentry_sdk.Scope]]]":
|
||||
if MCP_PACKAGE_VERSION and MCP_PACKAGE_VERSION < (2, 0, 0):
|
||||
if ctx is None:
|
||||
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.isolation_scope"),
|
||||
ctx.request.scope["state"].get("sentry_sdk.current_scope"),
|
||||
)
|
||||
|
||||
|
||||
def _get_request_context_data(
|
||||
ctx: "Optional[Any]" = None,
|
||||
) -> "tuple[Optional[str], Optional[str], str]":
|
||||
"""
|
||||
Extract request ID, session ID, and MCP transport type from the request context.
|
||||
|
||||
Returns:
|
||||
Tuple of (request_id, session_id, mcp_transport).
|
||||
- request_id: May be None if not available
|
||||
- session_id: May be None if not available
|
||||
- mcp_transport: "http", "sse", "stdio"
|
||||
"""
|
||||
request_id: "Optional[str]" = None
|
||||
session_id: "Optional[str]" = None
|
||||
mcp_transport: str = "stdio"
|
||||
if MCP_PACKAGE_VERSION and MCP_PACKAGE_VERSION < (2, 0, 0):
|
||||
if ctx is None:
|
||||
try:
|
||||
ctx = request_ctx.get()
|
||||
except LookupError:
|
||||
return request_id, session_id, mcp_transport
|
||||
|
||||
if ctx is not None:
|
||||
request_id = ctx.request_id
|
||||
if hasattr(ctx, "request") and ctx.request is not None:
|
||||
request = ctx.request
|
||||
# Detect transport type by checking request characteristics
|
||||
if hasattr(request, "query_params") and request.query_params.get(
|
||||
"session_id"
|
||||
):
|
||||
# SSE transport uses query parameter
|
||||
mcp_transport = "sse"
|
||||
session_id = request.query_params.get("session_id")
|
||||
elif hasattr(request, "headers") and request.headers.get("mcp-session-id"):
|
||||
# StreamableHTTP transport uses header
|
||||
mcp_transport = "http"
|
||||
session_id = request.headers.get("mcp-session-id")
|
||||
|
||||
return request_id, session_id, mcp_transport
|
||||
|
||||
|
||||
def _get_span_config(
|
||||
handler_type: str, item_name: str
|
||||
) -> "tuple[str, str, str, Optional[str]]":
|
||||
"""
|
||||
Get span configuration based on handler type.
|
||||
|
||||
Returns:
|
||||
Tuple of (span_data_key, span_name, mcp_method_name, result_data_key)
|
||||
Note: result_data_key is None for resources
|
||||
"""
|
||||
if handler_type == "tool":
|
||||
span_data_key = SPANDATA.MCP_TOOL_NAME
|
||||
mcp_method_name = "tools/call"
|
||||
result_data_key = SPANDATA.MCP_TOOL_RESULT_CONTENT
|
||||
elif handler_type == "prompt":
|
||||
span_data_key = SPANDATA.MCP_PROMPT_NAME
|
||||
mcp_method_name = "prompts/get"
|
||||
result_data_key = SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT
|
||||
else: # resource
|
||||
span_data_key = SPANDATA.MCP_RESOURCE_URI
|
||||
mcp_method_name = "resources/read"
|
||||
result_data_key = None # Resources don't capture result content
|
||||
|
||||
span_name = f"{mcp_method_name} {item_name}"
|
||||
return span_data_key, span_name, mcp_method_name, result_data_key
|
||||
|
||||
|
||||
def _set_span_input_data(
|
||||
span: "Union[StreamedSpan, Span]",
|
||||
handler_name: str,
|
||||
span_data_key: str,
|
||||
mcp_method_name: str,
|
||||
arguments: "dict[str, Any]",
|
||||
request_id: "Optional[str]",
|
||||
session_id: "Optional[str]",
|
||||
mcp_transport: str,
|
||||
) -> None:
|
||||
"""Set input span data for MCP handlers."""
|
||||
|
||||
# Set handler identifier
|
||||
_set_span_data_attribute(span, span_data_key, handler_name)
|
||||
_set_span_data_attribute(span, SPANDATA.MCP_METHOD_NAME, mcp_method_name)
|
||||
|
||||
# Set transport/MCP transport type
|
||||
_set_span_data_attribute(
|
||||
span,
|
||||
SPANDATA.NETWORK_TRANSPORT,
|
||||
"pipe" if mcp_transport == "stdio" else "tcp",
|
||||
)
|
||||
_set_span_data_attribute(span, SPANDATA.MCP_TRANSPORT, mcp_transport)
|
||||
|
||||
# Set request_id if provided
|
||||
if request_id:
|
||||
_set_span_data_attribute(span, SPANDATA.MCP_REQUEST_ID, request_id)
|
||||
|
||||
# Set session_id if provided
|
||||
if session_id:
|
||||
_set_span_data_attribute(span, SPANDATA.MCP_SESSION_ID, session_id)
|
||||
|
||||
# Set request arguments (excluding common request context objects)
|
||||
for k, v in arguments.items():
|
||||
_set_span_data_attribute(span, f"mcp.request.argument.{k}", safe_serialize(v))
|
||||
|
||||
|
||||
def _extract_tool_result_content(result: "Any") -> "Any":
|
||||
"""
|
||||
Extract meaningful content from MCP tool result.
|
||||
|
||||
Tool handlers can return:
|
||||
- CallToolResult (mcp v2+): Has .content list and optional .structured_content
|
||||
- tuple (UnstructuredContent, StructuredContent): Return the structured content (dict)
|
||||
- dict (StructuredContent): Return as-is
|
||||
- list/Iterable (UnstructuredContent): Extract text from content blocks
|
||||
"""
|
||||
if result is None:
|
||||
return None
|
||||
|
||||
# Handle v2 CallToolResult-like objects (has .content list attribute)
|
||||
if hasattr(result, "content") and isinstance(
|
||||
getattr(result, "content", None), list
|
||||
):
|
||||
# This is only present when a tool declares an output_schema
|
||||
structured = getattr(result, "structured_content", None)
|
||||
if structured is not None:
|
||||
return structured
|
||||
return _extract_text_from_content_blocks(result.content)
|
||||
|
||||
# Handle CombinationContent: tuple of (UnstructuredContent, StructuredContent)
|
||||
if isinstance(result, tuple) and len(result) == 2:
|
||||
# Return the structured content (2nd element)
|
||||
return result[1]
|
||||
|
||||
# Handle StructuredContent: dict
|
||||
if isinstance(result, dict):
|
||||
return result
|
||||
|
||||
# Handle UnstructuredContent: iterable of ContentBlock objects
|
||||
if hasattr(result, "__iter__") and not isinstance(result, (str, bytes, dict)):
|
||||
return _extract_text_from_content_blocks(result)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _extract_text_from_content_blocks(content_blocks: "Any") -> "Any":
|
||||
texts = []
|
||||
try:
|
||||
for item in content_blocks:
|
||||
if hasattr(item, "text"):
|
||||
texts.append(item.text)
|
||||
elif isinstance(item, dict) and "text" in item:
|
||||
texts.append(item["text"])
|
||||
except Exception:
|
||||
return content_blocks
|
||||
return " ".join(texts) if texts else content_blocks
|
||||
|
||||
|
||||
def _set_span_output_data(
|
||||
span: "Union[StreamedSpan, Span]",
|
||||
result: "Any",
|
||||
result_data_key: "Optional[str]",
|
||||
handler_type: str,
|
||||
) -> None:
|
||||
"""Set output span data for MCP handlers."""
|
||||
if result is None:
|
||||
return
|
||||
|
||||
# Get integration to check PII settings
|
||||
integration = sentry_sdk.get_client().get_integration(MCPIntegration)
|
||||
if integration is None:
|
||||
return
|
||||
|
||||
# Check if we should include sensitive data
|
||||
should_include_data = should_send_default_pii() and integration.include_prompts
|
||||
|
||||
# For tools, extract the meaningful content
|
||||
if handler_type == "tool":
|
||||
extracted = _extract_tool_result_content(result)
|
||||
if (
|
||||
extracted is not None
|
||||
and should_include_data
|
||||
and result_data_key is not None
|
||||
):
|
||||
_set_span_data_attribute(span, result_data_key, safe_serialize(extracted))
|
||||
# Set content count if result is a dict
|
||||
if isinstance(extracted, dict):
|
||||
_set_span_data_attribute(
|
||||
span, SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT, len(extracted)
|
||||
)
|
||||
elif handler_type == "prompt":
|
||||
# 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 and result_data_key is not None:
|
||||
_set_span_data_attribute(span, result_data_key, content_text)
|
||||
except Exception:
|
||||
# Silently ignore if we can't extract message info
|
||||
pass
|
||||
# Resources don't capture result content (result_data_key is None)
|
||||
|
||||
|
||||
# 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 {}
|
||||
arguments = {"name": handler_name, **arguments}
|
||||
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, ...]",
|
||||
original_kwargs: "Optional[dict[str, Any]]" = None,
|
||||
) -> "tuple[str, dict[str, Any]]":
|
||||
"""
|
||||
Extract handler name and arguments from v1 positional args.
|
||||
|
||||
In MCP SDK v1, handlers receive positional args:
|
||||
- Tool: (tool_name, arguments)
|
||||
- Prompt: (name, arguments)
|
||||
- Resource: (uri,)
|
||||
"""
|
||||
original_kwargs = original_kwargs or {}
|
||||
|
||||
if handler_type == "tool":
|
||||
if original_args:
|
||||
handler_name = original_args[0]
|
||||
elif original_kwargs.get("name"):
|
||||
handler_name = original_kwargs["name"]
|
||||
|
||||
arguments = {}
|
||||
if len(original_args) > 1:
|
||||
arguments = original_args[1]
|
||||
elif original_kwargs.get("arguments"):
|
||||
arguments = original_kwargs["arguments"]
|
||||
|
||||
elif handler_type == "prompt":
|
||||
if original_args:
|
||||
handler_name = original_args[0]
|
||||
elif original_kwargs.get("name"):
|
||||
handler_name = original_kwargs["name"]
|
||||
|
||||
arguments = {}
|
||||
if len(original_args) > 1:
|
||||
arguments = original_args[1]
|
||||
elif original_kwargs.get("arguments"):
|
||||
arguments = original_kwargs["arguments"]
|
||||
|
||||
arguments = {"name": handler_name, **(arguments or {})}
|
||||
|
||||
else: # resource
|
||||
handler_name = "unknown"
|
||||
if original_args:
|
||||
handler_name = str(original_args[0])
|
||||
elif original_kwargs.get("uri"):
|
||||
handler_name = str(original_kwargs["uri"])
|
||||
|
||||
arguments = {}
|
||||
|
||||
return handler_name, arguments
|
||||
|
||||
|
||||
def _prepare_handler_data(
|
||||
handler_type: str,
|
||||
original_args: "tuple[Any, ...]",
|
||||
original_kwargs: "Optional[dict[str, Any]]" = None,
|
||||
params: "Optional[Any]" = None,
|
||||
) -> "tuple[str, dict[str, Any], str, str, str, Optional[str]]":
|
||||
"""
|
||||
Prepare common handler data for both v1 and v2 MCP SDK.
|
||||
|
||||
Args:
|
||||
handler_type: "tool", "prompt", or "resource"
|
||||
original_args: Original positional args (v1 path)
|
||||
original_kwargs: Original keyword args (v1 path)
|
||||
params: Typed params object from v2 ServerRequestContext path
|
||||
|
||||
Returns:
|
||||
Tuple of (handler_name, arguments, span_data_key, span_name, mcp_method_name, result_data_key)
|
||||
"""
|
||||
if params is not None:
|
||||
handler_name, arguments = _extract_handler_data_from_params(
|
||||
handler_type, params
|
||||
)
|
||||
elif _is_v2_context(original_args):
|
||||
handler_name = "unknown"
|
||||
arguments = {}
|
||||
else:
|
||||
handler_name, arguments = _extract_handler_data_from_args(
|
||||
handler_type, original_args, original_kwargs
|
||||
)
|
||||
|
||||
span_data_key, span_name, mcp_method_name, result_data_key = _get_span_config(
|
||||
handler_type, handler_name
|
||||
)
|
||||
|
||||
return (
|
||||
handler_name,
|
||||
arguments,
|
||||
span_data_key,
|
||||
span_name,
|
||||
mcp_method_name,
|
||||
result_data_key,
|
||||
)
|
||||
|
||||
|
||||
async def _handler_wrapper(
|
||||
handler_type: str,
|
||||
func: "Callable[..., Any]",
|
||||
original_args: "tuple[Any, ...]",
|
||||
original_kwargs: "Optional[dict[str, Any]]" = None,
|
||||
self: "Optional[Any]" = None,
|
||||
force_await: bool = True,
|
||||
) -> "Any":
|
||||
"""
|
||||
Wrapper for MCP handlers.
|
||||
|
||||
Args:
|
||||
handler_type: "tool", "prompt", or "resource"
|
||||
func: The handler function to wrap
|
||||
original_args: Original arguments passed to the handler
|
||||
original_kwargs: Original keyword arguments passed to the handler
|
||||
self: Optional instance for bound methods
|
||||
"""
|
||||
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]
|
||||
params = original_args[1] if len(original_args) > 1 else None
|
||||
|
||||
(
|
||||
handler_name,
|
||||
arguments,
|
||||
span_data_key,
|
||||
span_name,
|
||||
mcp_method_name,
|
||||
result_data_key,
|
||||
) = _prepare_handler_data(
|
||||
handler_type, original_args, original_kwargs, params=params
|
||||
)
|
||||
|
||||
scopes = _get_active_http_scopes(ctx=ctx)
|
||||
|
||||
isolation_scope_context: "ContextManager[Any]"
|
||||
current_scope_context: "ContextManager[Any]"
|
||||
|
||||
if scopes is None:
|
||||
isolation_scope_context = nullcontext()
|
||||
current_scope_context = nullcontext()
|
||||
else:
|
||||
isolation_scope, current_scope = scopes
|
||||
|
||||
isolation_scope_context = (
|
||||
nullcontext()
|
||||
if isolation_scope is None
|
||||
else sentry_sdk.scope.use_isolation_scope(isolation_scope)
|
||||
)
|
||||
current_scope_context = (
|
||||
nullcontext()
|
||||
if current_scope is None
|
||||
else sentry_sdk.scope.use_scope(current_scope)
|
||||
)
|
||||
|
||||
# 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 isolation_scope_context:
|
||||
with current_scope_context:
|
||||
span_mgr: "Union[Span, StreamedSpan, ContextManager[None]]"
|
||||
if span_streaming:
|
||||
if sentry_sdk.traces.get_current_span() is not None:
|
||||
span_mgr = sentry_sdk.traces.start_span(
|
||||
name=span_name,
|
||||
attributes={
|
||||
"sentry.op": OP.MCP_SERVER,
|
||||
"sentry.origin": MCPIntegration.origin,
|
||||
},
|
||||
)
|
||||
else:
|
||||
span_mgr = nullcontext()
|
||||
else:
|
||||
span_mgr = get_start_span_function()(
|
||||
op=OP.MCP_SERVER,
|
||||
name=span_name,
|
||||
origin=MCPIntegration.origin,
|
||||
)
|
||||
|
||||
with span_mgr as span:
|
||||
if span is not None:
|
||||
# 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":
|
||||
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")
|
||||
|
||||
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:
|
||||
# Execute the async handler
|
||||
if self is not None:
|
||||
original_args = (self, *original_args)
|
||||
|
||||
result = func(*original_args, **original_kwargs)
|
||||
if force_await or inspect.isawaitable(result):
|
||||
result = await result
|
||||
|
||||
except Exception as e:
|
||||
# Set error flag for tools
|
||||
if span is not None and handler_type == "tool":
|
||||
_set_span_data_attribute(
|
||||
span, SPANDATA.MCP_TOOL_RESULT_IS_ERROR, True
|
||||
)
|
||||
sentry_sdk.capture_exception(e)
|
||||
raise
|
||||
|
||||
if span is not None:
|
||||
_set_span_output_data(span, result, result_data_key, handler_type)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _create_instrumented_decorator(
|
||||
original_decorator: "Callable[..., Any]",
|
||||
handler_type: str,
|
||||
*decorator_args: "Any",
|
||||
**decorator_kwargs: "Any",
|
||||
) -> "Callable[..., Any]":
|
||||
"""
|
||||
Create an instrumented version of an MCP decorator.
|
||||
|
||||
This function intercepts MCP decorators (like @server.call_tool()) and injects
|
||||
Sentry instrumentation into the handler registration flow. The returned decorator
|
||||
will:
|
||||
1. Receive the user's handler function
|
||||
2. Pass the instrumented version to the original MCP decorator
|
||||
|
||||
This ensures that when the handler is called at runtime, it's already wrapped
|
||||
with Sentry spans and metrics collection.
|
||||
|
||||
Args:
|
||||
original_decorator: The original MCP decorator method (e.g., Server.call_tool)
|
||||
handler_type: "tool", "prompt", or "resource" - determines span configuration
|
||||
decorator_args: Positional arguments to pass to the original decorator (e.g., self)
|
||||
decorator_kwargs: Keyword arguments to pass to the original decorator
|
||||
|
||||
Returns:
|
||||
A decorator function that instruments handlers before registering them
|
||||
"""
|
||||
|
||||
def instrumented_decorator(func: "Callable[..., Any]") -> "Callable[..., Any]":
|
||||
@wraps(func)
|
||||
async def wrapper(*args: "Any") -> "Any":
|
||||
return await _handler_wrapper(handler_type, func, args, force_await=False)
|
||||
|
||||
# Then register it with the original MCP decorator
|
||||
return original_decorator(*decorator_args, **decorator_kwargs)(wrapper)
|
||||
|
||||
return instrumented_decorator
|
||||
|
||||
|
||||
_METHOD_TO_HANDLER_TYPE = {
|
||||
"tools/call": "tool",
|
||||
"prompts/get": "prompt",
|
||||
"resources/read": "resource",
|
||||
}
|
||||
|
||||
# In MCP SDK v2, tool/prompt/resource handlers are most commonly registered via
|
||||
# the Server(...) constructor kwargs rather than add_request_handler. The in-tree
|
||||
# high-level MCPServer also wires its handlers through these kwargs.
|
||||
_KWARG_TO_HANDLER_TYPE = {
|
||||
"on_call_tool": "tool",
|
||||
"on_get_prompt": "prompt",
|
||||
"on_read_resource": "resource",
|
||||
}
|
||||
|
||||
|
||||
def _patch_lowlevel_server() -> None:
|
||||
"""
|
||||
Patches the mcp.server.lowlevel.Server class to instrument handler execution.
|
||||
"""
|
||||
if MCP_PACKAGE_VERSION and MCP_PACKAGE_VERSION >= (2, 0, 0):
|
||||
_patch_lowlevel_server_v2()
|
||||
else:
|
||||
_patch_lowlevel_server_v1()
|
||||
|
||||
|
||||
def _patch_lowlevel_server_v1() -> None:
|
||||
"""Patches v1 Server decorator methods (call_tool, get_prompt, read_resource)."""
|
||||
# Patch call_tool decorator
|
||||
original_call_tool = Server.call_tool # type: ignore[attr-defined]
|
||||
|
||||
def patched_call_tool(
|
||||
self: "Server", **kwargs: "Any"
|
||||
) -> "Callable[[Callable[..., Any]], Callable[..., Any]]":
|
||||
"""Patched version of Server.call_tool that adds Sentry instrumentation."""
|
||||
return lambda func: _create_instrumented_decorator(
|
||||
original_call_tool, "tool", self, **kwargs
|
||||
)(func)
|
||||
|
||||
Server.call_tool = patched_call_tool # type: ignore[attr-defined]
|
||||
|
||||
# Patch get_prompt decorator
|
||||
original_get_prompt = Server.get_prompt # type: ignore[attr-defined]
|
||||
|
||||
def patched_get_prompt(
|
||||
self: "Server",
|
||||
) -> "Callable[[Callable[..., Any]], Callable[..., Any]]":
|
||||
"""Patched version of Server.get_prompt that adds Sentry instrumentation."""
|
||||
return lambda func: _create_instrumented_decorator(
|
||||
original_get_prompt, "prompt", self
|
||||
)(func)
|
||||
|
||||
Server.get_prompt = patched_get_prompt # type: ignore[attr-defined]
|
||||
|
||||
# Patch read_resource decorator
|
||||
original_read_resource = Server.read_resource # type: ignore[attr-defined]
|
||||
|
||||
def patched_read_resource(
|
||||
self: "Server",
|
||||
) -> "Callable[[Callable[..., Any]], Callable[..., Any]]":
|
||||
"""Patched version of Server.read_resource that adds Sentry instrumentation."""
|
||||
return lambda func: _create_instrumented_decorator(
|
||||
original_read_resource, "resource", self
|
||||
)(func)
|
||||
|
||||
Server.read_resource = patched_read_resource # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _wrap_v2_handler(
|
||||
handler_type: str, handler: "Callable[..., Any]"
|
||||
) -> "Callable[..., Any]":
|
||||
"""Wrap a v2 (ctx, params) handler with Sentry instrumentation.
|
||||
|
||||
Idempotent: an already-wrapped handler is returned unchanged so handlers
|
||||
registered through more than one path (e.g. MCPServer building a Server)
|
||||
are not double-wrapped.
|
||||
"""
|
||||
if getattr(handler, "__sentry_mcp_wrapped__", False):
|
||||
return handler
|
||||
|
||||
@wraps(handler)
|
||||
async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
|
||||
return await _handler_wrapper(
|
||||
handler_type, handler, args, kwargs, force_await=False
|
||||
)
|
||||
|
||||
wrapper.__sentry_mcp_wrapped__ = True # type: ignore[attr-defined]
|
||||
return wrapper
|
||||
|
||||
|
||||
def _patch_lowlevel_server_v2() -> None:
|
||||
"""Patches the v2 Server to wrap tool/prompt/resource handlers.
|
||||
|
||||
Handlers can be registered either via the Server(...) constructor kwargs
|
||||
(on_call_tool/on_get_prompt/on_read_resource) — the path the in-tree
|
||||
MCPServer and most lowlevel examples use — or via add_request_handler.
|
||||
Both are patched.
|
||||
"""
|
||||
original_init = Server.__init__
|
||||
|
||||
@wraps(original_init)
|
||||
def patched_init(self: "Server", *args: "Any", **kwargs: "Any") -> None:
|
||||
for kwarg, handler_type in _KWARG_TO_HANDLER_TYPE.items():
|
||||
handler = kwargs.get(kwarg)
|
||||
if handler is not None:
|
||||
kwargs[kwarg] = _wrap_v2_handler(handler_type, handler)
|
||||
original_init(self, *args, **kwargs)
|
||||
|
||||
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:
|
||||
handler_type = _METHOD_TO_HANDLER_TYPE.get(method)
|
||||
if handler_type is not None:
|
||||
handler = _wrap_v2_handler(handler_type, handler)
|
||||
|
||||
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
|
||||
|
||||
@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 # type: ignore[method-assign]
|
||||
|
||||
|
||||
def _patch_fastmcp() -> None:
|
||||
"""
|
||||
Patches the standalone fastmcp package's FastMCP class.
|
||||
|
||||
The standalone fastmcp package (v2.14.0+) registers its own handlers for
|
||||
prompts and resources directly, bypassing the Server decorators we patch.
|
||||
This function patches the _get_prompt_mcp and _read_resource_mcp methods
|
||||
to add instrumentation for those handlers.
|
||||
"""
|
||||
if FastMCP is not None and hasattr(FastMCP, "_get_prompt_mcp"):
|
||||
original_get_prompt_mcp = FastMCP._get_prompt_mcp
|
||||
|
||||
@wraps(original_get_prompt_mcp)
|
||||
async def patched_get_prompt_mcp(
|
||||
self: "Any", *args: "Any", **kwargs: "Any"
|
||||
) -> "Any":
|
||||
return await _handler_wrapper(
|
||||
"prompt",
|
||||
original_get_prompt_mcp,
|
||||
args,
|
||||
kwargs,
|
||||
self,
|
||||
)
|
||||
|
||||
FastMCP._get_prompt_mcp = patched_get_prompt_mcp
|
||||
|
||||
if FastMCP is not None and hasattr(FastMCP, "_read_resource_mcp"):
|
||||
original_read_resource_mcp = FastMCP._read_resource_mcp
|
||||
|
||||
@wraps(original_read_resource_mcp)
|
||||
async def patched_read_resource_mcp(
|
||||
self: "Any", *args: "Any", **kwargs: "Any"
|
||||
) -> "Any":
|
||||
return await _handler_wrapper(
|
||||
"resource",
|
||||
original_read_resource_mcp,
|
||||
args,
|
||||
kwargs,
|
||||
self,
|
||||
)
|
||||
|
||||
FastMCP._read_resource_mcp = patched_read_resource_mcp
|
||||
Reference in New Issue
Block a user