ref(mcp): Use middleware when available

Author: alexander-alderman-webb

sentry_sdk/integrations/mcp.py

@@ -55,6 +55,7 @@
 if TYPE_CHECKING:
     from typing import Any, Awaitable, Callable, ContextManager, Optional, Tuple, Union
 
+    from mcp.server.context import CallNext, HandlerResult
     from mcp_types import (
         CallToolResult,
         GetPromptResult,
@@ -943,6 +944,21 @@ async def wrapper(*args: "Any") -> "Any":
     Server.read_resource = patched_read_resource  # type: ignore[attr-defined]
 
 
+async def _sentry_middleware(
+    ctx: "ServerRequestContext[Any, Any]", call_next: "CallNext"
+) -> "HandlerResult":
+    if ctx.method == "tools/call":
+        return await _tool_handler_wrapper(call_next, ctx)
+
+    if ctx.method == "prompts/get":
+        return await _prompt_handler_wrapper(call_next, ctx)
+
+    if ctx.method == "resources/read":
+        return await _resource_handler_wrapper(call_next, ctx)
+
+    return await call_next(ctx)
+
+
 def _patch_lowlevel_server_v2() -> None:
     """Patches the v2 Server to wrap tool/prompt/resource handlers.
 
@@ -955,113 +971,11 @@ def _patch_lowlevel_server_v2() -> None:
 
     @wraps(original_init)
     def patched_init(self: "Server", *args: "Any", **kwargs: "Any") -> None:
-        on_tool_call = kwargs.get("on_call_tool")
-        if on_tool_call is not None and not getattr(
-            on_tool_call, "__sentry_mcp_wrapped__", False
-        ):
-
-            @wraps(on_tool_call)
-            async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
-                return await _tool_handler_wrapper(
-                    on_tool_call, args, kwargs, force_await=False
-                )
-
-            wrapper.__sentry_mcp_wrapped__ = True  # type: ignore[attr-defined]
-            kwargs["on_call_tool"] = wrapper
-
-        on_get_prompt = kwargs.get("on_get_prompt")
-        if on_get_prompt is not None and not getattr(
-            on_get_prompt, "__sentry_mcp_wrapped__", False
-        ):
-
-            @wraps(on_get_prompt)
-            async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
-                return await _prompt_handler_wrapper(
-                    on_get_prompt, args, kwargs, force_await=False
-                )
-
-            wrapper.__sentry_mcp_wrapped__ = True  # type: ignore[attr-defined]
-            kwargs["on_get_prompt"] = wrapper
-
-        on_read_resource = kwargs.get("on_read_resource")
-        if on_read_resource is not None and not getattr(
-            on_read_resource, "__sentry_mcp_wrapped__", False
-        ):
-
-            @wraps(on_read_resource)
-            async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
-                return await _resource_handler_wrapper(
-                    on_read_resource, args, kwargs, force_await=False
-                )
-
-            wrapper.__sentry_mcp_wrapped__ = True  # type: ignore[attr-defined]
-            kwargs["on_read_resource"] = wrapper
-
         original_init(self, *args, **kwargs)
+        self.middleware.append(_sentry_middleware)
 
     Server.__init__ = patched_init  # type: ignore[method-assign]
 
-    original_add_request_handler = Server.add_request_handler
-
-    def patched_add_request_handler(
-        self: "Server",
-        method: str,
-        params_type: "Any",
-        handler: "Callable[..., Any]",
-        *args: "Any",
-        **kwargs: "Any",
-    ) -> None:
-        if getattr(handler, "__sentry_mcp_wrapped__", False):
-            return original_add_request_handler(
-                self, method, params_type, handler, *args, **kwargs
-            )
-
-        if method == "tools/call":
-
-            @wraps(handler)
-            async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
-                return await _tool_handler_wrapper(
-                    handler, args, kwargs, force_await=False
-                )
-
-            wrapper.__sentry_mcp_wrapped__ = True  # type: ignore[attr-defined]
-
-            return original_add_request_handler(
-                self, method, params_type, wrapper, *args, **kwargs
-            )
-        if method == "prompts/get":
-
-            @wraps(handler)
-            async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
-                return await _prompt_handler_wrapper(
-                    handler, args, kwargs, force_await=False
-                )
-
-            wrapper.__sentry_mcp_wrapped__ = True  # type: ignore[attr-defined]
-
-            return original_add_request_handler(
-                self, method, params_type, wrapper, *args, **kwargs
-            )
-        if method == "resources/read":
-
-            @wraps(handler)
-            async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
-                return await _resource_handler_wrapper(
-                    handler, args, kwargs, force_await=False
-                )
-
-            wrapper.__sentry_mcp_wrapped__ = True  # type: ignore[attr-defined]
-
-            return original_add_request_handler(
-                self, method, params_type, wrapper, *args, **kwargs
-            )
-
-        original_add_request_handler(
-            self, method, params_type, handler, *args, **kwargs
-        )
-
-    Server.add_request_handler = patched_add_request_handler  # type: ignore[method-assign]
-
 
 def _patch_handle_request() -> None:
     original_handle_request = StreamableHTTPServerTransport.handle_request