sentry-sdk 2.42.1__py2.py3-none-any.whl → 2.43.0__py2.py3-none-any.whl

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

sentry_sdk/integrations/openai_agents/__init__.py

@@ -5,6 +5,7 @@ from .patches import (
     _create_get_all_tools_wrapper,
     _create_run_wrapper,
     _patch_agent_run,
+    _patch_error_tracing,
 )
 
 try:
@@ -48,6 +49,7 @@ class OpenAIAgentsIntegration(Integration):
     @staticmethod
     def setup_once():
         # type: () -> None
+        _patch_error_tracing()
         _patch_tools()
         _patch_model()
         _patch_runner()

sentry_sdk/integrations/openai_agents/patches/__init__.py

@@ -2,3 +2,4 @@ from .models import _create_get_model_wrapper  # noqa: F401
 from .tools import _create_get_all_tools_wrapper  # noqa: F401
 from .runner import _create_run_wrapper  # noqa: F401
 from .agent_run import _patch_agent_run  # noqa: F401
+from .error_tracing import _patch_error_tracing  # noqa: F401

sentry_sdk/integrations/openai_agents/patches/error_tracing.py

@@ -0,0 +1,77 @@
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk.consts import SPANSTATUS
+from sentry_sdk.tracing_utils import set_span_errored
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Optional
+
+
+def _patch_error_tracing():
+    # type: () -> None
+    """
+    Patches agents error tracing function to inject our span error logic
+    when a tool execution fails.
+
+    In newer versions, the function is at: agents.util._error_tracing.attach_error_to_current_span
+    In older versions, it was at: agents._utils.attach_error_to_current_span
+
+    This works even when the module or function doesn't exist.
+    """
+    error_tracing_module = None
+
+    # Try newer location first (agents.util._error_tracing)
+    try:
+        from agents.util import _error_tracing
+
+        error_tracing_module = _error_tracing
+    except (ImportError, AttributeError):
+        pass
+
+    # Try older location (agents._utils)
+    if error_tracing_module is None:
+        try:
+            import agents._utils
+
+            error_tracing_module = agents._utils
+        except (ImportError, AttributeError):
+            # Module doesn't exist in either location, nothing to patch
+            return
+
+    # Check if the function exists
+    if not hasattr(error_tracing_module, "attach_error_to_current_span"):
+        return
+
+    original_attach_error = error_tracing_module.attach_error_to_current_span
+
+    @wraps(original_attach_error)
+    def sentry_attach_error_to_current_span(error, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        """
+        Wraps agents' error attachment to also set Sentry span status to error.
+        This allows us to properly track tool execution errors even though
+        the agents library swallows exceptions.
+        """
+        # Set the current Sentry span to errored
+        current_span = sentry_sdk.get_current_span()
+        if current_span is not None:
+            set_span_errored(current_span)
+            current_span.set_data("span.status", "error")
+
+            # Optionally capture the error details if we have them
+            if hasattr(error, "__class__"):
+                current_span.set_data("error.type", error.__class__.__name__)
+            if hasattr(error, "__str__"):
+                error_message = str(error)
+                if error_message:
+                    current_span.set_data("error.message", error_message)
+
+        # Call the original function
+        return original_attach_error(error, *args, **kwargs)
+
+    error_tracing_module.attach_error_to_current_span = (
+        sentry_attach_error_to_current_span
+    )

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -0,0 +1,47 @@
+from sentry_sdk.integrations import DidNotEnable, Integration
+
+
+try:
+    import pydantic_ai  # type: ignore
+except ImportError:
+    raise DidNotEnable("pydantic-ai not installed")
+
+
+from .patches import (
+    _patch_agent_run,
+    _patch_graph_nodes,
+    _patch_model_request,
+    _patch_tool_execution,
+)
+
+
+class PydanticAIIntegration(Integration):
+    identifier = "pydantic_ai"
+    origin = f"auto.ai.{identifier}"
+
+    def __init__(self, include_prompts=True):
+        # type: (bool) -> None
+        """
+        Initialize the Pydantic AI integration.
+
+        Args:
+            include_prompts: Whether to include prompts and messages in span data.
+                Requires send_default_pii=True. Defaults to True.
+        """
+        self.include_prompts = include_prompts
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        """
+        Set up the pydantic-ai integration.
+
+        This patches the key methods in pydantic-ai to create Sentry spans for:
+        - Agent invocations (Agent.run methods)
+        - Model requests (AI client calls)
+        - Tool executions
+        """
+        _patch_agent_run()
+        _patch_graph_nodes()
+        _patch_model_request()
+        _patch_tool_execution()

sentry_sdk/integrations/pydantic_ai/consts.py

@@ -0,0 +1 @@
+SPAN_ORIGIN = "auto.ai.pydantic_ai"