diff --git a/sentry_sdk/integrations/mcp.py b/sentry_sdk/integrations/mcp.py index 47fda272b7..983aa041ca 100644 --- a/sentry_sdk/integrations/mcp.py +++ b/sentry_sdk/integrations/mcp.py @@ -352,19 +352,20 @@ def _prepare_handler_data( ) -async def _async_handler_wrapper( +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": """ - Async wrapper for MCP handlers. + Wrapper for MCP handlers. Args: handler_type: "tool", "prompt", or "resource" - func: The async handler function to wrap + 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 @@ -421,73 +422,11 @@ async def _async_handler_wrapper( # Execute the async handler if self is not None: original_args = (self, *original_args) - result = await func(*original_args, **original_kwargs) - except Exception as e: - # Set error flag for tools - if handler_type == "tool": - span.set_data(SPANDATA.MCP_TOOL_RESULT_IS_ERROR, True) - sentry_sdk.capture_exception(e) - raise - - _set_span_output_data(span, result, result_data_key, handler_type) - return result - -def _sync_handler_wrapper( - handler_type: str, func: "Callable[..., Any]", original_args: "tuple[Any, ...]" -) -> "Any": - """ - Sync wrapper for MCP handlers. + result = func(*original_args, **original_kwargs) + if force_await or inspect.isawaitable(result): + result = await result - 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, mcp_transport = _get_request_context_data() - - # Set input span data - _set_span_input_data( - span, - handler_name, - span_data_key, - mcp_method_name, - arguments, - request_id, - session_id, - mcp_transport, - ) - - # 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": @@ -499,41 +438,6 @@ def _sync_handler_wrapper( return result -def _create_instrumented_handler( - handler_type: str, func: "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: "Any") -> "Any": - return await _async_handler_wrapper(handler_type, func, args) - - return async_wrapper - else: - - @wraps(func) - def sync_wrapper(*args: "Any") -> "Any": - return _sync_handler_wrapper(handler_type, func, args) - - return sync_wrapper - - def _create_instrumented_decorator( original_decorator: "Callable[..., Any]", handler_type: str, @@ -547,8 +451,7 @@ def _create_instrumented_decorator( 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 + 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. @@ -564,12 +467,12 @@ def _create_instrumented_decorator( """ def instrumented_decorator(func: "Callable[..., Any]") -> "Callable[..., Any]": - # First wrap the handler with instrumentation - instrumented_func = _create_instrumented_handler(handler_type, func) + @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)( - instrumented_func - ) + return original_decorator(*decorator_args, **decorator_kwargs)(wrapper) return instrumented_decorator @@ -634,7 +537,7 @@ def _patch_fastmcp() -> None: async def patched_get_prompt_mcp( self: "Any", *args: "Any", **kwargs: "Any" ) -> "Any": - return await _async_handler_wrapper( + return await _handler_wrapper( "prompt", original_get_prompt_mcp, args, @@ -651,7 +554,7 @@ async def patched_get_prompt_mcp( async def patched_read_resource_mcp( self: "Any", *args: "Any", **kwargs: "Any" ) -> "Any": - return await _async_handler_wrapper( + return await _handler_wrapper( "resource", original_read_resource_mcp, args,