diff --git a/.github/workflows/test-integrations-ai.yml b/.github/workflows/test-integrations-ai.yml index 940e0bed12..b362716cb0 100644 --- a/.github/workflows/test-integrations-ai.yml +++ b/.github/workflows/test-integrations-ai.yml @@ -82,6 +82,10 @@ jobs: run: | set -x # print commands that are executed ./scripts/runtox.sh "py${{ matrix.python-version }}-litellm" + - name: Test mcp + run: | + set -x # print commands that are executed + ./scripts/runtox.sh "py${{ matrix.python-version }}-mcp" - name: Test openai-base run: | set -x # print commands that are executed diff --git a/scripts/populate_tox/config.py b/scripts/populate_tox/config.py index b7f6d9efe7..7d83c0c9dd 100644 --- a/scripts/populate_tox/config.py +++ b/scripts/populate_tox/config.py @@ -242,6 +242,12 @@ "package": "loguru", "num_versions": 2, }, + "mcp": { + "package": "mcp", + "deps": { + "*": ["pytest-asyncio"], + }, + }, "openai-base": { "package": "openai", "integration_name": "openai", diff --git a/scripts/populate_tox/releases.jsonl b/scripts/populate_tox/releases.jsonl index edef42967d..60bbeb8dc2 100644 --- a/scripts/populate_tox/releases.jsonl +++ b/scripts/populate_tox/releases.jsonl @@ -98,7 +98,7 @@ {"info": {"classifiers": ["Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Scientific/Engineering :: Artificial Intelligence"], "name": "huggingface-hub", "requires_python": ">=3.8.0", "version": "0.24.7", "yanked": false}} {"info": {"classifiers": ["Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Scientific/Engineering :: Artificial Intelligence"], "name": "huggingface-hub", "requires_python": ">=3.8.0", "version": "0.28.1", "yanked": false}} {"info": {"classifiers": ["Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Scientific/Engineering :: Artificial Intelligence"], "name": "huggingface-hub", "requires_python": ">=3.8.0", "version": "0.32.6", "yanked": false}} -{"info": {"classifiers": ["Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Scientific/Engineering :: Artificial Intelligence"], "name": "huggingface-hub", "requires_python": ">=3.8.0", "version": "0.35.3", "yanked": false}} +{"info": {"classifiers": ["Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Scientific/Engineering :: Artificial Intelligence"], "name": "huggingface-hub", "requires_python": ">=3.8.0", "version": "0.36.0", "yanked": false}} {"info": {"classifiers": ["Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Science/Research", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.9", "Topic :: Scientific/Engineering :: Artificial Intelligence"], "name": "huggingface-hub", "requires_python": ">=3.9.0", "version": "1.0.0rc7", "yanked": false}} {"info": {"classifiers": ["License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.9"], "name": "langchain", "requires_python": "<4.0,>=3.8.1", "version": "0.1.20", "yanked": false}} {"info": {"classifiers": [], "name": "langchain", "requires_python": "<4.0,>=3.9", "version": "0.3.27", "yanked": false}} @@ -114,6 +114,10 @@ {"info": {"classifiers": ["Development Status :: 5 - Production/Stable", "Environment :: Web Environment", "Intended Audience :: Developers", "Intended Audience :: Information Technology", "Intended Audience :: System Administrators", "License :: OSI Approved :: MIT License", "Natural Language :: English", "Operating System :: OS Independent", "Programming Language :: Python", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Internet", "Topic :: Internet :: WWW/HTTP", "Topic :: Internet :: WWW/HTTP :: HTTP Servers", "Topic :: Software Development", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Application Frameworks", "Topic :: Software Development :: Libraries :: Python Modules", "Typing :: Typed"], "name": "litestar", "requires_python": "<4.0,>=3.8", "version": "2.18.0", "yanked": false}} {"info": {"classifiers": ["Development Status :: 5 - Production/Stable", "Environment :: Web Environment", "Intended Audience :: Developers", "Intended Audience :: Information Technology", "Intended Audience :: System Administrators", "License :: OSI Approved :: MIT License", "Natural Language :: English", "Operating System :: OS Independent", "Programming Language :: Python", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Internet", "Topic :: Internet :: WWW/HTTP", "Topic :: Internet :: WWW/HTTP :: HTTP Servers", "Topic :: Software Development", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Application Frameworks", "Topic :: Software Development :: Libraries :: Python Modules", "Typing :: Typed"], "name": "litestar", "requires_python": "<4.0,>=3.8", "version": "2.6.4", "yanked": false}} {"info": {"classifiers": ["Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Natural Language :: English", "Operating System :: OS Independent", "Programming Language :: Python", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.5", "Programming Language :: Python :: 3.6", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", "Topic :: System :: Logging"], "name": "loguru", "requires_python": "<4.0,>=3.5", "version": "0.7.3", "yanked": false}} +{"info": {"classifiers": ["Development Status :: 4 - Beta", "Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13"], "name": "mcp", "requires_python": ">=3.10", "version": "1.15.0", "yanked": false}} +{"info": {"classifiers": ["Development Status :: 4 - Beta", "Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13"], "name": "mcp", "requires_python": ">=3.10", "version": "1.16.0", "yanked": false}} +{"info": {"classifiers": ["Development Status :: 4 - Beta", "Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13"], "name": "mcp", "requires_python": ">=3.10", "version": "1.17.0", "yanked": false}} +{"info": {"classifiers": ["Development Status :: 4 - Beta", "Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13"], "name": "mcp", "requires_python": ">=3.10", "version": "1.18.0", "yanked": false}} {"info": {"classifiers": ["Intended Audience :: Developers", "Operating System :: MacOS", "Operating System :: Microsoft :: Windows", "Operating System :: OS Independent", "Operating System :: POSIX", "Operating System :: POSIX :: Linux", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Software Development :: Libraries :: Python Modules", "Typing :: Typed"], "name": "openai", "requires_python": ">=3.7.1", "version": "1.0.1", "yanked": false}} {"info": {"classifiers": ["Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Operating System :: MacOS", "Operating System :: Microsoft :: Windows", "Operating System :: OS Independent", "Operating System :: POSIX", "Operating System :: POSIX :: Linux", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Software Development :: Libraries :: Python Modules", "Typing :: Typed"], "name": "openai", "requires_python": ">=3.8", "version": "1.103.0", "yanked": false}} {"info": {"classifiers": ["Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Operating System :: MacOS", "Operating System :: Microsoft :: Windows", "Operating System :: OS Independent", "Operating System :: POSIX", "Operating System :: POSIX :: Linux", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Topic :: Software Development :: Libraries :: Python Modules", "Typing :: Typed"], "name": "openai", "requires_python": ">=3.8", "version": "1.109.1", "yanked": false}} diff --git a/scripts/split_tox_gh_actions/split_tox_gh_actions.py b/scripts/split_tox_gh_actions/split_tox_gh_actions.py index 12222eff1b..50892ad66b 100755 --- a/scripts/split_tox_gh_actions/split_tox_gh_actions.py +++ b/scripts/split_tox_gh_actions/split_tox_gh_actions.py @@ -78,6 +78,7 @@ "langchain-notiktoken", "langgraph", "litellm", + "mcp", "openai-base", "openai-notiktoken", "openai_agents", diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index 4b7d219245..b041e0626b 100644 --- a/sentry_sdk/consts.py +++ b/sentry_sdk/consts.py @@ -749,6 +749,90 @@ class SPANDATA: Example: "MainThread" """ + MCP_TOOL_NAME = "mcp.tool.name" + """ + The name of the MCP tool being called. + Example: "get_weather" + """ + + MCP_PROMPT_NAME = "mcp.prompt.name" + """ + The name of the MCP prompt being retrieved. + Example: "code_review" + """ + + MCP_RESOURCE_URI = "mcp.resource.uri" + """ + The URI of the MCP resource being accessed. + Example: "file:///path/to/resource" + """ + + MCP_METHOD_NAME = "mcp.method.name" + """ + The MCP protocol method name being called. + Example: "tools/call", "prompts/get", "resources/read" + """ + + MCP_REQUEST_ID = "mcp.request.id" + """ + The unique identifier for the MCP request. + Example: "req_123abc" + """ + + MCP_TOOL_RESULT_CONTENT = "mcp.tool.result.content" + """ + The result/output content from an MCP tool execution. + Example: "The weather is sunny" + """ + + MCP_TOOL_RESULT_CONTENT_COUNT = "mcp.tool.result.content_count" + """ + The number of items/keys in the MCP tool result. + Example: 5 + """ + + MCP_TOOL_RESULT_IS_ERROR = "mcp.tool.result.is_error" + """ + Whether the MCP tool execution resulted in an error. + Example: True + """ + + MCP_PROMPT_RESULT_MESSAGE_CONTENT = "mcp.prompt.result.message_content" + """ + The message content from an MCP prompt retrieval. + Example: "Review the following code..." + """ + + MCP_PROMPT_RESULT_MESSAGE_ROLE = "mcp.prompt.result.message_role" + """ + The role of the message in an MCP prompt retrieval (only set for single-message prompts). + Example: "user", "assistant", "system" + """ + + MCP_PROMPT_RESULT_MESSAGE_COUNT = "mcp.prompt.result.message_count" + """ + The number of messages in an MCP prompt result. + Example: 1, 3 + """ + + MCP_RESOURCE_PROTOCOL = "mcp.resource.protocol" + """ + The protocol/scheme of the MCP resource URI. + Example: "file", "http", "https" + """ + + MCP_TRANSPORT = "mcp.transport" + """ + The transport method used for MCP communication. + Example: "pipe" (stdio), "tcp" (HTTP/WebSocket/SSE) + """ + + MCP_SESSION_ID = "mcp.session.id" + """ + The session identifier for the MCP connection. + Example: "a1b2c3d4e5f6" + """ + class SPANSTATUS: """ @@ -845,6 +929,7 @@ class OP: WEBSOCKET_SERVER = "websocket.server" SOCKET_CONNECTION = "socket.connection" SOCKET_DNS = "socket.dns" + MCP_SERVER = "mcp.server" # This type exists to trick mypy and PyCharm into thinking `init` and `Client` diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py index 9e279b8345..b4dd6fc96b 100644 --- a/sentry_sdk/integrations/__init__.py +++ b/sentry_sdk/integrations/__init__.py @@ -149,6 +149,7 @@ def iter_default_integrations(with_auto_enabling_integrations): "launchdarkly": (9, 8, 0), "litellm": (1, 77, 5), "loguru": (0, 7, 0), + "mcp": (1, 15, 0), "openai": (1, 0, 0), "openai_agents": (0, 0, 19), "openfeature": (0, 7, 1), diff --git a/sentry_sdk/integrations/mcp.py b/sentry_sdk/integrations/mcp.py new file mode 100644 index 0000000000..2a2d440616 --- /dev/null +++ b/sentry_sdk/integrations/mcp.py @@ -0,0 +1,552 @@ +""" +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 get_start_span_function +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations import Integration, DidNotEnable +from sentry_sdk.utils import safe_serialize +from sentry_sdk.scope import should_send_default_pii + +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] +except ImportError: + raise DidNotEnable("MCP SDK not installed") + + +if TYPE_CHECKING: + from typing import Any, Callable, Optional + + +class MCPIntegration(Integration): + identifier = "mcp" + origin = "auto.ai.mcp" + + def __init__(self, include_prompts=True): + # type: (bool) -> 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(): + # type: () -> None + """ + Patches MCP server classes to instrument handler execution. + """ + _patch_lowlevel_server() + + +def _get_request_context_data(): + # type: () -> tuple[Optional[str], Optional[str], str] + """ + Extract request ID, session ID, and transport type from the MCP request context. + + Returns: + Tuple of (request_id, session_id, transport). + - request_id: May be None if not available + - session_id: May be None if not available + - transport: "tcp" for HTTP-based, "pipe" for stdio + """ + request_id = None # type: Optional[str] + session_id = None # type: Optional[str] + transport = "pipe" # type: str + + try: + ctx = request_ctx.get() + + if ctx is not None: + request_id = ctx.request_id + if hasattr(ctx, "request") and ctx.request is not None: + transport = "tcp" + request = ctx.request + if hasattr(request, "headers"): + session_id = request.headers.get("mcp-session-id") + + except LookupError: + # No request context available - default to pipe + pass + + return request_id, session_id, transport + + +def _get_span_config(handler_type, item_name): + # type: (str, 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, + handler_name, + span_data_key, + mcp_method_name, + arguments, + request_id, + session_id, + transport, +): + # type: (Any, str, str, str, dict[str, Any], Optional[str], Optional[str], str) -> None + """Set input span data for MCP handlers.""" + # Set handler identifier + span.set_data(span_data_key, handler_name) + span.set_data(SPANDATA.MCP_METHOD_NAME, mcp_method_name) + + # Set transport type + span.set_data(SPANDATA.MCP_TRANSPORT, transport) + + # 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)) + + +def _extract_tool_result_content(result): + # type: (Any) -> Any + """ + Extract meaningful content from MCP tool result. + + Tool handlers can return: + - tuple (UnstructuredContent, StructuredContent): Return the structured content (dict) + - dict (StructuredContent): Return as-is + - Iterable (UnstructuredContent): Extract text from content blocks + """ + if result is None: + return None + + # 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 + # Try to extract text content + if hasattr(result, "__iter__") and not isinstance(result, (str, bytes, dict)): + texts = [] + try: + for item in result: + # Try to get text attribute from ContentBlock objects + if hasattr(item, "text"): + texts.append(item.text) + elif isinstance(item, dict) and "text" in item: + texts.append(item["text"]) + except Exception: + # If extraction fails, return the original + return result + return " ".join(texts) if texts else result + + return result + + +def _set_span_output_data(span, result, result_data_key, handler_type): + # type: (Any, Any, Optional[str], 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: + span.set_data(result_data_key, safe_serialize(extracted)) + # Set content count if result is a dict + if isinstance(extracted, dict): + span.set_data(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 = None # type: Optional[list[str]] + 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: + span.set_data(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: + span.set_data(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: + span.set_data(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 _prepare_handler_data(handler_type, original_args): + # type: (str, tuple[Any, ...]) -> tuple[str, dict[str, Any], str, str, str, Optional[str]] + """ + Prepare common handler data for both async and sync wrappers. + + Returns: + Tuple of (handler_name, arguments, span_data_key, span_name, mcp_method_name, result_data_key) + """ + # Extract handler-specific data based on handler type + if handler_type == "tool": + handler_name = original_args[0] # tool_name + arguments = original_args[1] if len(original_args) > 1 else {} + elif handler_type == "prompt": + handler_name = original_args[0] # name + arguments = original_args[1] if len(original_args) > 1 else {} + # Include name in arguments dict for span data + arguments = {"name": handler_name, **(arguments or {})} + else: # resource + uri = original_args[0] + handler_name = str(uri) if uri else "unknown" + arguments = {} + + # Get span configuration + 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 _async_handler_wrapper(handler_type, func, original_args): + # type: (str, Callable[..., Any], tuple[Any, ...]) -> Any + """ + Async wrapper for MCP handlers. + + Args: + handler_type: "tool", "prompt", or "resource" + func: The async handler function to wrap + original_args: Original arguments passed to the handler + """ + ( + handler_name, + arguments, + span_data_key, + span_name, + mcp_method_name, + result_data_key, + ) = _prepare_handler_data(handler_type, original_args) + + # 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, 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, + transport, + ) + + # For resources, extract and set protocol + if handler_type == "resource": + uri = original_args[0] + 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 + result = await func(*original_args) + 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) + return result + + +def _sync_handler_wrapper(handler_type, func, original_args): + # type: (str, Callable[..., Any], tuple[Any, ...]) -> Any + """ + Sync wrapper for MCP handlers. + + Args: + handler_type: "tool", "prompt", or "resource" + func: The sync handler function to wrap + original_args: Original arguments passed to the handler + """ + ( + handler_name, + arguments, + span_data_key, + span_name, + mcp_method_name, + result_data_key, + ) = _prepare_handler_data(handler_type, original_args) + + # 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, 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, + transport, + ) + + # For resources, extract and set protocol + if handler_type == "resource": + uri = original_args[0] + 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 sync handler + result = func(*original_args) + 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) + return result + + +def _create_instrumented_handler(handler_type, func): + # type: (str, Callable[..., Any]) -> Callable[..., Any] + """ + Create an instrumented version of a handler function (async or sync). + + This function wraps the user's handler with a runtime wrapper that will create + Sentry spans and capture metrics when the handler is actually called. + + The wrapper preserves the async/sync nature of the original function, which is + critical for Python's async/await to work correctly. + + Args: + handler_type: "tool", "prompt", or "resource" - determines span configuration + func: The handler function to instrument (async or sync) + + Returns: + A wrapped version of func that creates Sentry spans on execution + """ + if inspect.iscoroutinefunction(func): + + @wraps(func) + async def async_wrapper(*args): + # type: (*Any) -> Any + return await _async_handler_wrapper(handler_type, func, args) + + return async_wrapper + else: + + @wraps(func) + def sync_wrapper(*args): + # type: (*Any) -> Any + return _sync_handler_wrapper(handler_type, func, args) + + return sync_wrapper + + +def _create_instrumented_decorator( + original_decorator, handler_type, *decorator_args, **decorator_kwargs +): + # type: (Callable[..., Any], str, *Any, **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. Wrap it with instrumentation via _create_instrumented_handler + 3. 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): + # type: (Callable[..., Any]) -> Callable[..., Any] + # First wrap the handler with instrumentation + instrumented_func = _create_instrumented_handler(handler_type, func) + # Then register it with the original MCP decorator + return original_decorator(*decorator_args, **decorator_kwargs)( + instrumented_func + ) + + return instrumented_decorator + + +def _patch_lowlevel_server(): + # type: () -> None + """ + Patches the mcp.server.lowlevel.Server class to instrument handler execution. + """ + # Patch call_tool decorator + original_call_tool = Server.call_tool + + def patched_call_tool(self, **kwargs): + # type: (Server, **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 + + # Patch get_prompt decorator + original_get_prompt = Server.get_prompt + + def patched_get_prompt(self): + # type: (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 + + # Patch read_resource decorator + original_read_resource = Server.read_resource + + def patched_read_resource(self): + # type: (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 diff --git a/setup.py b/setup.py index cf66f30f23..d91f99a1ee 100644 --- a/setup.py +++ b/setup.py @@ -68,6 +68,7 @@ def get_file_text(file_name): "litellm": ["litellm>=1.77.5"], "litestar": ["litestar>=2.0.0"], "loguru": ["loguru>=0.5"], + "mcp": ["mcp>=1.15.0"], "openai": ["openai>=1.0.0", "tiktoken>=0.3.0"], "openfeature": ["openfeature-sdk>=0.7.1"], "opentelemetry": ["opentelemetry-distro>=0.35b0"], diff --git a/tests/integrations/mcp/__init__.py b/tests/integrations/mcp/__init__.py new file mode 100644 index 0000000000..01ef442500 --- /dev/null +++ b/tests/integrations/mcp/__init__.py @@ -0,0 +1,3 @@ +import pytest + +pytest.importorskip("mcp") diff --git a/tests/integrations/mcp/test_mcp.py b/tests/integrations/mcp/test_mcp.py new file mode 100644 index 0000000000..738fdedf48 --- /dev/null +++ b/tests/integrations/mcp/test_mcp.py @@ -0,0 +1,911 @@ +""" +Unit tests for the MCP (Model Context Protocol) integration. + +This test suite covers: +- Tool handlers (sync and async) +- Prompt handlers (sync and async) +- Resource handlers (sync and async) +- Error handling for each handler type +- Request context data extraction (request_id, session_id, transport) +- Tool result content extraction (various formats) +- Span data validation +- Origin tracking + +The tests mock the MCP server components and request context to verify +that the integration properly instruments MCP handlers with Sentry spans. +""" + +import pytest +import json +from unittest import mock + +try: + from unittest.mock import AsyncMock +except ImportError: + + class AsyncMock(mock.MagicMock): + async def __call__(self, *args, **kwargs): + return super(AsyncMock, self).__call__(*args, **kwargs) + + +from mcp.server.lowlevel import Server +from mcp.server.lowlevel.server import request_ctx + +try: + from mcp.server.lowlevel.server import request_ctx +except ImportError: + request_ctx = None + +from sentry_sdk import start_transaction +from sentry_sdk.consts import SPANDATA, OP +from sentry_sdk.integrations.mcp import MCPIntegration + + +@pytest.fixture(autouse=True) +def reset_request_ctx(): + """Reset request context before and after each test""" + if request_ctx is not None: + try: + if request_ctx.get() is not None: + request_ctx.set(None) + except LookupError: + pass + + yield + + if request_ctx is not None: + try: + request_ctx.set(None) + except LookupError: + pass + + +# Mock MCP types and structures +class MockURI: + """Mock URI object for resource testing""" + + def __init__(self, uri_string): + self.scheme = uri_string.split("://")[0] if "://" in uri_string else "" + self.path = uri_string.split("://")[1] if "://" in uri_string else uri_string + self._uri_string = uri_string + + def __str__(self): + return self._uri_string + + +class MockRequestContext: + """Mock MCP request context""" + + def __init__(self, request_id=None, session_id=None, transport="pipe"): + self.request_id = request_id + if transport == "tcp": + self.request = MockHTTPRequest(session_id) + else: + self.request = None + + +class MockHTTPRequest: + """Mock HTTP request for SSE/WebSocket transport""" + + def __init__(self, session_id=None): + self.headers = {} + if session_id: + self.headers["mcp-session-id"] = session_id + + +class MockTextContent: + """Mock TextContent object""" + + def __init__(self, text): + self.text = text + + +class MockPromptMessage: + """Mock PromptMessage object""" + + def __init__(self, role, content_text): + self.role = role + self.content = MockTextContent(content_text) + + +class MockGetPromptResult: + """Mock GetPromptResult object""" + + def __init__(self, messages): + self.messages = messages + + +def test_integration_patches_server(sentry_init): + """Test that MCPIntegration patches the Server class""" + # Get original methods before integration + original_call_tool = Server.call_tool + original_get_prompt = Server.get_prompt + original_read_resource = Server.read_resource + + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + + # After initialization, the methods should be patched + assert Server.call_tool is not original_call_tool + assert Server.get_prompt is not original_get_prompt + assert Server.read_resource is not original_read_resource + + +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +def test_tool_handler_sync( + sentry_init, capture_events, send_default_pii, include_prompts +): + """Test that synchronous tool handlers create proper spans""" + sentry_init( + integrations=[MCPIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-123", transport="pipe") + request_ctx.set(mock_ctx) + + @server.call_tool() + def test_tool(tool_name, arguments): + return {"result": "success", "value": 42} + + with start_transaction(name="mcp tx"): + # Call the tool handler + result = test_tool("calculate", {"x": 10, "y": 5}) + + assert result == {"result": "success", "value": 42} + + (tx,) = events + assert tx["type"] == "transaction" + assert len(tx["spans"]) == 1 + + span = tx["spans"][0] + assert span["op"] == OP.MCP_SERVER + assert span["description"] == "tools/call calculate" + assert span["origin"] == "auto.ai.mcp" + + # Check span data + assert span["data"][SPANDATA.MCP_TOOL_NAME] == "calculate" + assert span["data"][SPANDATA.MCP_METHOD_NAME] == "tools/call" + assert span["data"][SPANDATA.MCP_TRANSPORT] == "pipe" + assert span["data"][SPANDATA.MCP_REQUEST_ID] == "req-123" + assert span["data"]["mcp.request.argument.x"] == "10" + assert span["data"]["mcp.request.argument.y"] == "5" + + # Check PII-sensitive data is only present when both flags are True + if send_default_pii and include_prompts: + assert span["data"][SPANDATA.MCP_TOOL_RESULT_CONTENT] == json.dumps( + { + "result": "success", + "value": 42, + } + ) + assert span["data"][SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT] == 2 + else: + assert SPANDATA.MCP_TOOL_RESULT_CONTENT not in span["data"] + assert SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT not in span["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +async def test_tool_handler_async( + sentry_init, capture_events, send_default_pii, include_prompts +): + """Test that async tool handlers create proper spans""" + sentry_init( + integrations=[MCPIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext( + request_id="req-456", session_id="session-789", transport="tcp" + ) + request_ctx.set(mock_ctx) + + @server.call_tool() + async def test_tool_async(tool_name, arguments): + return {"status": "completed"} + + with start_transaction(name="mcp tx"): + result = await test_tool_async("process", {"data": "test"}) + + assert result == {"status": "completed"} + + (tx,) = events + assert tx["type"] == "transaction" + assert len(tx["spans"]) == 1 + + span = tx["spans"][0] + assert span["op"] == OP.MCP_SERVER + assert span["description"] == "tools/call process" + assert span["origin"] == "auto.ai.mcp" + + # Check span data + assert span["data"][SPANDATA.MCP_TOOL_NAME] == "process" + assert span["data"][SPANDATA.MCP_METHOD_NAME] == "tools/call" + assert span["data"][SPANDATA.MCP_TRANSPORT] == "tcp" + assert span["data"][SPANDATA.MCP_REQUEST_ID] == "req-456" + assert span["data"][SPANDATA.MCP_SESSION_ID] == "session-789" + assert span["data"]["mcp.request.argument.data"] == '"test"' + + # Check PII-sensitive data + if send_default_pii and include_prompts: + assert span["data"][SPANDATA.MCP_TOOL_RESULT_CONTENT] == json.dumps( + {"status": "completed"} + ) + else: + assert SPANDATA.MCP_TOOL_RESULT_CONTENT not in span["data"] + + +def test_tool_handler_with_error(sentry_init, capture_events): + """Test that tool handler errors are captured properly""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-error", transport="pipe") + request_ctx.set(mock_ctx) + + @server.call_tool() + def failing_tool(tool_name, arguments): + raise ValueError("Tool execution failed") + + with start_transaction(name="mcp tx"): + with pytest.raises(ValueError): + failing_tool("bad_tool", {}) + + # Should have error event and transaction + assert len(events) == 2 + error_event, tx = events + + # Check error event + assert error_event["level"] == "error" + assert error_event["exception"]["values"][0]["type"] == "ValueError" + assert error_event["exception"]["values"][0]["value"] == "Tool execution failed" + + # Check transaction and span + assert tx["type"] == "transaction" + assert len(tx["spans"]) == 1 + span = tx["spans"][0] + + # Error flag should be set for tools + assert span["data"][SPANDATA.MCP_TOOL_RESULT_IS_ERROR] is True + assert span["tags"]["status"] == "internal_error" + + +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +def test_prompt_handler_sync( + sentry_init, capture_events, send_default_pii, include_prompts +): + """Test that synchronous prompt handlers create proper spans""" + sentry_init( + integrations=[MCPIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-prompt", transport="pipe") + request_ctx.set(mock_ctx) + + @server.get_prompt() + def test_prompt(name, arguments): + return MockGetPromptResult([MockPromptMessage("user", "Tell me about Python")]) + + with start_transaction(name="mcp tx"): + result = test_prompt("code_help", {"language": "python"}) + + assert result.messages[0].role == "user" + assert result.messages[0].content.text == "Tell me about Python" + + (tx,) = events + assert tx["type"] == "transaction" + assert len(tx["spans"]) == 1 + + span = tx["spans"][0] + assert span["op"] == OP.MCP_SERVER + assert span["description"] == "prompts/get code_help" + assert span["origin"] == "auto.ai.mcp" + + # Check span data + assert span["data"][SPANDATA.MCP_PROMPT_NAME] == "code_help" + assert span["data"][SPANDATA.MCP_METHOD_NAME] == "prompts/get" + assert span["data"][SPANDATA.MCP_TRANSPORT] == "pipe" + assert span["data"][SPANDATA.MCP_REQUEST_ID] == "req-prompt" + assert span["data"]["mcp.request.argument.name"] == '"code_help"' + assert span["data"]["mcp.request.argument.language"] == '"python"' + + # Message count is always captured + assert span["data"][SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT] == 1 + + # For single message prompts, role and content should be captured only with PII + if send_default_pii and include_prompts: + assert span["data"][SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE] == "user" + assert ( + span["data"][SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT] + == "Tell me about Python" + ) + else: + assert SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE not in span["data"] + assert SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT not in span["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +async def test_prompt_handler_async( + sentry_init, capture_events, send_default_pii, include_prompts +): + """Test that async prompt handlers create proper spans""" + sentry_init( + integrations=[MCPIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext( + request_id="req-async-prompt", session_id="session-abc", transport="tcp" + ) + request_ctx.set(mock_ctx) + + @server.get_prompt() + async def test_prompt_async(name, arguments): + return MockGetPromptResult( + [ + MockPromptMessage("system", "You are a helpful assistant"), + MockPromptMessage("user", "What is MCP?"), + ] + ) + + with start_transaction(name="mcp tx"): + result = await test_prompt_async("mcp_info", {}) + + assert len(result.messages) == 2 + + (tx,) = events + assert tx["type"] == "transaction" + assert len(tx["spans"]) == 1 + + span = tx["spans"][0] + assert span["op"] == OP.MCP_SERVER + assert span["description"] == "prompts/get mcp_info" + + # For multi-message prompts, count is always captured + assert span["data"][SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT] == 2 + # Role/content are never captured for multi-message prompts (even with PII) + assert SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE not in span["data"] + assert SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT not in span["data"] + + +def test_prompt_handler_with_error(sentry_init, capture_events): + """Test that prompt handler errors are captured""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-error-prompt", transport="pipe") + request_ctx.set(mock_ctx) + + @server.get_prompt() + def failing_prompt(name, arguments): + raise RuntimeError("Prompt not found") + + with start_transaction(name="mcp tx"): + with pytest.raises(RuntimeError): + failing_prompt("missing_prompt", {}) + + # Should have error event and transaction + assert len(events) == 2 + error_event, tx = events + + assert error_event["level"] == "error" + assert error_event["exception"]["values"][0]["type"] == "RuntimeError" + + +def test_resource_handler_sync(sentry_init, capture_events): + """Test that synchronous resource handlers create proper spans""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-resource", transport="pipe") + request_ctx.set(mock_ctx) + + @server.read_resource() + def test_resource(uri): + return {"content": "file contents", "mime_type": "text/plain"} + + with start_transaction(name="mcp tx"): + uri = MockURI("file:///path/to/file.txt") + result = test_resource(uri) + + assert result["content"] == "file contents" + + (tx,) = events + assert tx["type"] == "transaction" + assert len(tx["spans"]) == 1 + + span = tx["spans"][0] + assert span["op"] == OP.MCP_SERVER + assert span["description"] == "resources/read file:///path/to/file.txt" + assert span["origin"] == "auto.ai.mcp" + + # Check span data + assert span["data"][SPANDATA.MCP_RESOURCE_URI] == "file:///path/to/file.txt" + assert span["data"][SPANDATA.MCP_METHOD_NAME] == "resources/read" + assert span["data"][SPANDATA.MCP_TRANSPORT] == "pipe" + assert span["data"][SPANDATA.MCP_REQUEST_ID] == "req-resource" + assert span["data"][SPANDATA.MCP_RESOURCE_PROTOCOL] == "file" + # Resources don't capture result content + assert SPANDATA.MCP_TOOL_RESULT_CONTENT not in span["data"] + + +@pytest.mark.asyncio +async def test_resource_handler_async(sentry_init, capture_events): + """Test that async resource handlers create proper spans""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext( + request_id="req-async-resource", session_id="session-res", transport="tcp" + ) + request_ctx.set(mock_ctx) + + @server.read_resource() + async def test_resource_async(uri): + return {"data": "resource data"} + + with start_transaction(name="mcp tx"): + uri = MockURI("https://example.com/resource") + result = await test_resource_async(uri) + + assert result["data"] == "resource data" + + (tx,) = events + assert tx["type"] == "transaction" + assert len(tx["spans"]) == 1 + + span = tx["spans"][0] + assert span["op"] == OP.MCP_SERVER + assert span["description"] == "resources/read https://example.com/resource" + + assert span["data"][SPANDATA.MCP_RESOURCE_URI] == "https://example.com/resource" + assert span["data"][SPANDATA.MCP_RESOURCE_PROTOCOL] == "https" + assert span["data"][SPANDATA.MCP_SESSION_ID] == "session-res" + + +def test_resource_handler_with_error(sentry_init, capture_events): + """Test that resource handler errors are captured""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-error-resource", transport="pipe") + request_ctx.set(mock_ctx) + + @server.read_resource() + def failing_resource(uri): + raise FileNotFoundError("Resource not found") + + with start_transaction(name="mcp tx"): + with pytest.raises(FileNotFoundError): + uri = MockURI("file:///missing.txt") + failing_resource(uri) + + # Should have error event and transaction + assert len(events) == 2 + error_event, tx = events + + assert error_event["level"] == "error" + assert error_event["exception"]["values"][0]["type"] == "FileNotFoundError" + + +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (False, False)], +) +def test_tool_result_extraction_tuple( + sentry_init, capture_events, send_default_pii, include_prompts +): + """Test extraction of tool results from tuple format (UnstructuredContent, StructuredContent)""" + sentry_init( + integrations=[MCPIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-tuple", transport="pipe") + request_ctx.set(mock_ctx) + + @server.call_tool() + def test_tool_tuple(tool_name, arguments): + # Return CombinationContent: (UnstructuredContent, StructuredContent) + unstructured = [MockTextContent("Result text")] + structured = {"key": "value", "count": 5} + return (unstructured, structured) + + with start_transaction(name="mcp tx"): + test_tool_tuple("combo_tool", {}) + + (tx,) = events + span = tx["spans"][0] + + # Should extract the structured content (second element of tuple) only with PII + if send_default_pii and include_prompts: + assert span["data"][SPANDATA.MCP_TOOL_RESULT_CONTENT] == json.dumps( + { + "key": "value", + "count": 5, + } + ) + assert span["data"][SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT] == 2 + else: + assert SPANDATA.MCP_TOOL_RESULT_CONTENT not in span["data"] + assert SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT not in span["data"] + + +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (False, False)], +) +def test_tool_result_extraction_unstructured( + sentry_init, capture_events, send_default_pii, include_prompts +): + """Test extraction of tool results from UnstructuredContent (list of content blocks)""" + sentry_init( + integrations=[MCPIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-unstructured", transport="pipe") + request_ctx.set(mock_ctx) + + @server.call_tool() + def test_tool_unstructured(tool_name, arguments): + # Return UnstructuredContent as list of content blocks + return [ + MockTextContent("First part"), + MockTextContent("Second part"), + ] + + with start_transaction(name="mcp tx"): + test_tool_unstructured("text_tool", {}) + + (tx,) = events + span = tx["spans"][0] + + # Should extract and join text from content blocks only with PII + if send_default_pii and include_prompts: + assert ( + span["data"][SPANDATA.MCP_TOOL_RESULT_CONTENT] == '"First part Second part"' + ) + else: + assert SPANDATA.MCP_TOOL_RESULT_CONTENT not in span["data"] + + +def test_request_context_no_context(sentry_init, capture_events): + """Test handling when no request context is available""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Clear request context (simulating no context available) + # This will cause a LookupError when trying to get context + request_ctx.set(None) + + @server.call_tool() + def test_tool_no_ctx(tool_name, arguments): + return {"result": "ok"} + + with start_transaction(name="mcp tx"): + # This should work even without request context + try: + test_tool_no_ctx("tool", {}) + except LookupError: + # If it raises LookupError, that's expected when context is truly missing + pass + + # Should still create span even if context is missing + (tx,) = events + span = tx["spans"][0] + + # Transport defaults to "pipe" when no context + assert span["data"][SPANDATA.MCP_TRANSPORT] == "pipe" + # Request ID and Session ID should not be present + assert SPANDATA.MCP_REQUEST_ID not in span["data"] + assert SPANDATA.MCP_SESSION_ID not in span["data"] + + +def test_span_origin(sentry_init, capture_events): + """Test that span origin is set correctly""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-origin", transport="pipe") + request_ctx.set(mock_ctx) + + @server.call_tool() + def test_tool(tool_name, arguments): + return {"result": "test"} + + with start_transaction(name="mcp tx"): + test_tool("origin_test", {}) + + (tx,) = events + + assert tx["contexts"]["trace"]["origin"] == "manual" + assert tx["spans"][0]["origin"] == "auto.ai.mcp" + + +def test_multiple_handlers(sentry_init, capture_events): + """Test that multiple handler calls create multiple spans""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-multi", transport="pipe") + request_ctx.set(mock_ctx) + + @server.call_tool() + def tool1(tool_name, arguments): + return {"result": "tool1"} + + @server.call_tool() + def tool2(tool_name, arguments): + return {"result": "tool2"} + + @server.get_prompt() + def prompt1(name, arguments): + return MockGetPromptResult([MockPromptMessage("user", "Test prompt")]) + + with start_transaction(name="mcp tx"): + tool1("tool_a", {}) + tool2("tool_b", {}) + prompt1("prompt_a", {}) + + (tx,) = events + assert tx["type"] == "transaction" + assert len(tx["spans"]) == 3 + + # Check that we have different span types + span_ops = [span["op"] for span in tx["spans"]] + assert all(op == OP.MCP_SERVER for op in span_ops) + + span_descriptions = [span["description"] for span in tx["spans"]] + assert "tools/call tool_a" in span_descriptions + assert "tools/call tool_b" in span_descriptions + assert "prompts/get prompt_a" in span_descriptions + + +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (False, False)], +) +def test_prompt_with_dict_result( + sentry_init, capture_events, send_default_pii, include_prompts +): + """Test prompt handler with dict result instead of GetPromptResult object""" + sentry_init( + integrations=[MCPIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-dict-prompt", transport="pipe") + request_ctx.set(mock_ctx) + + @server.get_prompt() + def test_prompt_dict(name, arguments): + # Return dict format instead of GetPromptResult object + return { + "messages": [ + {"role": "user", "content": {"text": "Hello from dict"}}, + ] + } + + with start_transaction(name="mcp tx"): + test_prompt_dict("dict_prompt", {}) + + (tx,) = events + span = tx["spans"][0] + + # Message count is always captured + assert span["data"][SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT] == 1 + + # Role and content only captured with PII + if send_default_pii and include_prompts: + assert span["data"][SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE] == "user" + assert ( + span["data"][SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT] + == "Hello from dict" + ) + else: + assert SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE not in span["data"] + assert SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT not in span["data"] + + +def test_resource_without_protocol(sentry_init, capture_events): + """Test resource handler with URI without protocol scheme""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-no-proto", transport="pipe") + request_ctx.set(mock_ctx) + + @server.read_resource() + def test_resource(uri): + return {"data": "test"} + + with start_transaction(name="mcp tx"): + # URI without protocol + test_resource("simple-path") + + (tx,) = events + span = tx["spans"][0] + + assert span["data"][SPANDATA.MCP_RESOURCE_URI] == "simple-path" + # No protocol should be set + assert SPANDATA.MCP_RESOURCE_PROTOCOL not in span["data"] + + +def test_tool_with_complex_arguments(sentry_init, capture_events): + """Test tool handler with complex nested arguments""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-complex", transport="pipe") + request_ctx.set(mock_ctx) + + @server.call_tool() + def test_tool_complex(tool_name, arguments): + return {"processed": True} + + with start_transaction(name="mcp tx"): + complex_args = { + "nested": {"key": "value", "list": [1, 2, 3]}, + "string": "test", + "number": 42, + } + test_tool_complex("complex_tool", complex_args) + + (tx,) = events + span = tx["spans"][0] + + # Complex arguments should be serialized + assert span["data"]["mcp.request.argument.nested"] == json.dumps( + {"key": "value", "list": [1, 2, 3]} + ) + assert span["data"]["mcp.request.argument.string"] == '"test"' + assert span["data"]["mcp.request.argument.number"] == "42" + + +@pytest.mark.asyncio +async def test_async_handlers_mixed(sentry_init, capture_events): + """Test mixing sync and async handlers in the same transaction""" + sentry_init( + integrations=[MCPIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + server = Server("test-server") + + # Set up mock request context + mock_ctx = MockRequestContext(request_id="req-mixed", transport="pipe") + request_ctx.set(mock_ctx) + + @server.call_tool() + def sync_tool(tool_name, arguments): + return {"type": "sync"} + + @server.call_tool() + async def async_tool(tool_name, arguments): + return {"type": "async"} + + with start_transaction(name="mcp tx"): + sync_result = sync_tool("sync", {}) + async_result = await async_tool("async", {}) + + assert sync_result["type"] == "sync" + assert async_result["type"] == "async" + + (tx,) = events + assert len(tx["spans"]) == 2 + + # Both should be instrumented correctly + assert all(span["op"] == OP.MCP_SERVER for span in tx["spans"]) diff --git a/tox.ini b/tox.ini index 05273d9e82..302ff63a73 100644 --- a/tox.ini +++ b/tox.ini @@ -65,7 +65,7 @@ envlist = {py3.8,py3.10,py3.11}-huggingface_hub-v0.24.7 {py3.8,py3.12,py3.13}-huggingface_hub-v0.28.1 {py3.8,py3.12,py3.13}-huggingface_hub-v0.32.6 - {py3.8,py3.12,py3.13}-huggingface_hub-v0.35.3 + {py3.8,py3.12,py3.13}-huggingface_hub-v0.36.0 {py3.9,py3.12,py3.13}-huggingface_hub-v1.0.0rc7 {py3.9,py3.11,py3.12}-langchain-base-v0.1.20 @@ -82,6 +82,11 @@ envlist = {py3.9,py3.12,py3.13}-litellm-v1.77.7 {py3.9,py3.12,py3.13}-litellm-v1.78.7 + {py3.10,py3.12,py3.13}-mcp-v1.15.0 + {py3.10,py3.12,py3.13}-mcp-v1.16.0 + {py3.10,py3.12,py3.13}-mcp-v1.17.0 + {py3.10,py3.12,py3.13}-mcp-v1.18.0 + {py3.8,py3.11,py3.12}-openai-base-v1.0.1 {py3.8,py3.12,py3.13}-openai-base-v1.109.1 {py3.9,py3.12,py3.13}-openai-base-v2.6.0 @@ -363,7 +368,7 @@ deps = huggingface_hub-v0.24.7: huggingface_hub==0.24.7 huggingface_hub-v0.28.1: huggingface_hub==0.28.1 huggingface_hub-v0.32.6: huggingface_hub==0.32.6 - huggingface_hub-v0.35.3: huggingface_hub==0.35.3 + huggingface_hub-v0.36.0: huggingface_hub==0.36.0 huggingface_hub-v1.0.0rc7: huggingface_hub==1.0.0rc7 huggingface_hub: responses huggingface_hub: pytest-httpx @@ -393,6 +398,12 @@ deps = litellm-v1.77.7: litellm==1.77.7 litellm-v1.78.7: litellm==1.78.7 + mcp-v1.15.0: mcp==1.15.0 + mcp-v1.16.0: mcp==1.16.0 + mcp-v1.17.0: mcp==1.17.0 + mcp-v1.18.0: mcp==1.18.0 + mcp: pytest-asyncio + openai-base-v1.0.1: openai==1.0.1 openai-base-v1.109.1: openai==1.109.1 openai-base-v2.6.0: openai==2.6.0 @@ -782,6 +793,7 @@ setenv = litellm: TESTPATH=tests/integrations/litellm litestar: TESTPATH=tests/integrations/litestar loguru: TESTPATH=tests/integrations/loguru + mcp: TESTPATH=tests/integrations/mcp openai-base: TESTPATH=tests/integrations/openai openai-notiktoken: TESTPATH=tests/integrations/openai openai_agents: TESTPATH=tests/integrations/openai_agents