From 12e80eebf88d273237239131e96bded23afa6d67 Mon Sep 17 00:00:00 2001
From: Max Claus <maxclaus@posteo.com>
Date: Wed, 31 Dec 2025 00:07:17 -0300
Subject: [PATCH] fix(ai): Support distributed tracing for MCP tool calls

MCP tool calls now properly continue distributed traces from parent services
instead of creating new transactions with new trace IDs.

The issue occurred because when get_start_span_function() executes during MCP
tool handler wrapping, the HTTP transaction isn't accessible due to async
context isolation. The function only checked the current scope, which often
contained unrelated spans with different trace IDs.

This fix adds checking for sentry-trace headers in the MCP request context
and uses continue_trace() to create transactions that inherit the trace_id
and parent_span_id from incoming requests.

Benefits:
- MCP tools join distributed traces correctly
- Transactions still appear in MCP Insights dashboard
- Fully backward compatible
- No new dependencies
---
 sentry_sdk/ai/utils.py | 74 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 74 insertions(+)

diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py
index 1d2b4483c9..ef9ea9c8e2 100644
--- a/sentry_sdk/ai/utils.py
+++ b/sentry_sdk/ai/utils.py
@@ -100,6 +100,33 @@ def normalize_message_roles(messages: "list[dict[str, Any]]") -> "list[dict[str,
 
 
 def get_start_span_function() -> "Callable[..., Any]":
+    """
+    Determine whether to create a span or transaction for AI operations.
+
+    Checks in priority order:
+    1. MCP request context has sentry-trace headers: create transaction continuing the trace
+    2. Transaction exists in current scope: create child span
+    3. Otherwise: create new transaction
+
+    This ensures MCP tool calls properly continue distributed traces while maintaining
+    correct span hierarchies for local operations.
+    """
+    # Check for distributed trace headers in MCP request context for cross-service tracing
+    try:
+        from mcp.server.lowlevel.server import request_ctx  # type: ignore[import-not-found]
+
+        ctx = request_ctx.get()
+        if ctx and hasattr(ctx, "request") and ctx.request is not None:
+            request = ctx.request
+            if hasattr(request, "headers"):
+                sentry_trace = request.headers.get("sentry-trace")
+                if sentry_trace:
+                    return _create_transaction_from_mcp_headers
+    except (ImportError, LookupError):
+        # MCP not installed or no request context, fall through to normal logic
+        pass
+
+    # Normal logic: create span if transaction exists, otherwise create transaction
     current_span = sentry_sdk.get_current_span()
     transaction_exists = (
         current_span is not None and current_span.containing_transaction is not None
@@ -107,6 +134,53 @@ def get_start_span_function() -> "Callable[..., Any]":
     return sentry_sdk.start_span if transaction_exists else sentry_sdk.start_transaction
 
 
+def _create_transaction_from_mcp_headers(
+    op: "Optional[str]" = None,
+    name: "Optional[str]" = None,
+    origin: str = "manual",
+    **kwargs: "Any",
+) -> "Any":
+    """
+    Create transaction continuing distributed trace from MCP request headers.
+
+    Extracts sentry-trace and baggage headers from MCP request context and uses
+    continue_trace() to create a transaction inheriting trace_id and parent_span_id.
+    Ensures MCP tool executions join the distributed trace while remaining transactions
+    for MCP Insights dashboard compatibility.
+    """
+    try:
+        from mcp.server.lowlevel.server import request_ctx  # type: ignore[import-not-found]
+
+        ctx = request_ctx.get()
+        if ctx and hasattr(ctx, "request") and ctx.request is not None:
+            request = ctx.request
+            if hasattr(request, "headers"):
+                headers = {}
+
+                # Extract trace propagation headers
+                sentry_trace = request.headers.get("sentry-trace")
+                baggage = request.headers.get("baggage")
+
+                if sentry_trace:
+                    headers["sentry-trace"] = sentry_trace
+                    if baggage:
+                        headers["baggage"] = baggage
+
+                    # Use continue_trace to create transaction inheriting trace context
+                    isolation_scope = sentry_sdk.get_isolation_scope()
+                    transaction = isolation_scope.continue_trace(
+                        environ_or_headers=headers, op=op, name=name, origin=origin
+                    )
+
+                    # Start transaction on scope so it becomes active
+                    return sentry_sdk.start_transaction(transaction=transaction)
+    except Exception as e:
+        logger.debug("Could not create transaction from MCP headers: %s", e)
+
+    # Fallback: create new transaction if headers unavailable
+    return sentry_sdk.start_transaction(op=op, name=name, origin=origin, **kwargs)
+
+
 def _truncate_single_message_content_if_present(
     message: "Dict[str, Any]", max_chars: int
 ) -> "Dict[str, Any]":