-
Notifications
You must be signed in to change notification settings - Fork 572
fix(mcp): Nest MCP spans under HTTP transactions #5292
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
4c49c1a
a34ce9d
c994d0f
f4c30c2
3f13802
98d9929
3e016e9
1b1cec5
31d32b7
784f57a
0e95d7f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -21,6 +21,7 @@ | |
| 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] | ||
| from mcp.server.streamable_http import StreamableHTTPServerTransport # type: ignore[import-not-found] | ||
| except ImportError: | ||
| raise DidNotEnable("MCP SDK not installed") | ||
|
|
||
|
|
@@ -31,7 +32,9 @@ | |
|
|
||
|
|
||
| if TYPE_CHECKING: | ||
| from typing import Any, Callable, Optional | ||
| from typing import Any, Callable, Optional, Tuple | ||
|
|
||
| from starlette.types import Receive, Scope, Send # type: ignore[import-not-found] | ||
|
|
||
|
|
||
| class MCPIntegration(Integration): | ||
|
|
@@ -54,11 +57,34 @@ def setup_once() -> None: | |
| Patches MCP server classes to instrument handler execution. | ||
| """ | ||
| _patch_lowlevel_server() | ||
| _patch_handle_request() | ||
|
|
||
| if FastMCP is not None: | ||
| _patch_fastmcp() | ||
|
|
||
|
|
||
| def _get_active_http_scopes() -> ( | ||
| "Optional[Tuple[Optional[sentry_sdk.Scope], Optional[sentry_sdk.Scope]]]" | ||
| ): | ||
| try: | ||
| ctx = request_ctx.get() | ||
| except LookupError: | ||
| return None | ||
|
|
||
| if ( | ||
| ctx is None | ||
| or not hasattr(ctx, "request") | ||
| or ctx.request is None | ||
| or "state" not in ctx.request.scope | ||
alexander-alderman-webb marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| ): | ||
| return None | ||
|
|
||
| return ( | ||
| ctx.request.scope["state"].get("sentry_sdk.current_scope"), | ||
| ctx.request.scope["state"].get("sentry_sdk.isolation_scope"), | ||
| ) | ||
|
|
||
|
|
||
| def _get_request_context_data() -> "tuple[Optional[str], Optional[str], str]": | ||
| """ | ||
| Extract request ID, session ID, and MCP transport type from the request context. | ||
|
|
@@ -381,56 +407,67 @@ async def _async_handler_wrapper( | |
| result_data_key, | ||
| ) = _prepare_handler_data(handler_type, original_args, original_kwargs) | ||
|
|
||
| # 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, | ||
| ) | ||
| scopes = _get_active_http_scopes() | ||
|
|
||
| # For resources, extract and set protocol | ||
| if handler_type == "resource": | ||
| if original_args: | ||
| uri = original_args[0] | ||
| else: | ||
| uri = original_kwargs.get("uri") | ||
| if scopes is None: | ||
| current_scope = None | ||
| isolation_scope = None | ||
| else: | ||
| current_scope, isolation_scope = scopes | ||
|
|
||
| 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) | ||
| # Get request ID, session ID, and transport from context | ||
| request_id, session_id, mcp_transport = _get_request_context_data() | ||
|
|
||
| try: | ||
| # 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 | ||
| # Start span and execute | ||
| with sentry_sdk.scope.use_isolation_scope(isolation_scope): | ||
| with sentry_sdk.scope.use_scope(current_scope): | ||
| with get_start_span_function()( | ||
| op=OP.MCP_SERVER, | ||
| name=span_name, | ||
| origin=MCPIntegration.origin, | ||
| ) as span: | ||
| # 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": | ||
| if original_args: | ||
| uri = original_args[0] | ||
| else: | ||
| uri = original_kwargs.get("uri") | ||
|
|
||
| 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 | ||
| 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) | ||
|
|
||
| _set_span_output_data(span, result, result_data_key, handler_type) | ||
| return result | ||
| return result | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Sync handlers missing HTTP scope nestingMedium Severity The Additional Locations (1)
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The |
||
|
|
||
|
|
||
| def _sync_handler_wrapper( | ||
|
|
@@ -618,6 +655,25 @@ def patched_read_resource( | |
| Server.read_resource = patched_read_resource | ||
|
|
||
|
|
||
| def _patch_handle_request() -> None: | ||
| original_handle_request = StreamableHTTPServerTransport.handle_request | ||
|
|
||
| @wraps(original_handle_request) | ||
| async def patched_handle_request( | ||
| self: "StreamableHTTPServerTransport", | ||
| scope: "Scope", | ||
| receive: "Receive", | ||
| send: "Send", | ||
| ) -> None: | ||
| scope.setdefault("state", {})["sentry_sdk.isolation_scope"] = ( | ||
| sentry_sdk.get_isolation_scope() | ||
| ) | ||
| scope["state"]["sentry_sdk.current_scope"] = sentry_sdk.get_current_scope() | ||
| await original_handle_request(self, scope, receive, send) | ||
|
|
||
| StreamableHTTPServerTransport.handle_request = patched_handle_request | ||
|
|
||
|
|
||
| def _patch_fastmcp() -> None: | ||
| """ | ||
| Patches the standalone fastmcp package's FastMCP class. | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.