Merge branch 'master' into feat/span-first-2

Author: sentrivana

sentry_sdk/ai/utils.py

@@ -16,7 +16,6 @@
 from sentry_sdk.utils import logger
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
-from sentry_sdk.consts import SPANDATA
 
 MAX_GEN_AI_MESSAGE_BYTES = 20_000  # 20KB
 # Maximum characters when only a single message is left after bytes truncation
@@ -714,8 +713,6 @@ def truncate_and_annotate_messages(
     if len(messages) > 1:
         scope._gen_ai_original_message_count[span.span_id] = len(messages)
 
-    span.set_data(SPANDATA.META_GEN_AI_ORIGINAL_INPUT_MESSAGES_LENGTH, len(messages))
-
     return [truncated_message]
 
 

sentry_sdk/consts.py

@@ -879,14 +879,6 @@ class SPANDATA:
     Example: "a1b2c3d4e5f6"
     """
 
-    META_GEN_AI_ORIGINAL_INPUT_MESSAGES_LENGTH = (
-        "sentry.sdk_meta.gen_ai.input.messages.original_length"
-    )
-    """
-    The original number of input non-system instruction messages, before SDK trimming.
-    Example: 4
-    """
-
 
 class SPANSTATUS:
     """

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,

tests/integrations/anthropic/test_anthropic.py

@@ -1012,7 +1012,6 @@ def test_anthropic_message_truncation(sentry_init, capture_events):
     assert len(parsed_messages) == 1
     assert "small message 5" in str(parsed_messages[0])
 
-    assert chat_span["data"][SPANDATA.META_GEN_AI_ORIGINAL_INPUT_MESSAGES_LENGTH] == 5
     assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5
 
 
@@ -1064,7 +1063,6 @@ async def test_anthropic_message_truncation_async(sentry_init, capture_events):
     assert len(parsed_messages) == 1
     assert "small message 5" in str(parsed_messages[0])
 
-    assert chat_span["data"][SPANDATA.META_GEN_AI_ORIGINAL_INPUT_MESSAGES_LENGTH] == 5
     assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5
 
 

tests/integrations/google_genai/test_google_genai.py

@@ -983,7 +983,6 @@ def test_google_genai_message_truncation(
     assert parsed_messages[0]["role"] == "user"
     assert small_content in parsed_messages[0]["content"]
 
-    assert invoke_span["data"][SPANDATA.META_GEN_AI_ORIGINAL_INPUT_MESSAGES_LENGTH] == 2
     assert (
         event["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 2
     )

tests/integrations/langchain/test_langchain.py

@@ -1070,8 +1070,6 @@ def test_langchain_message_truncation(sentry_init, capture_events):
     assert isinstance(parsed_messages, list)
     assert len(parsed_messages) == 1
     assert "small message 5" in str(parsed_messages[0])
-
-    assert llm_span["data"][SPANDATA.META_GEN_AI_ORIGINAL_INPUT_MESSAGES_LENGTH] == 5
     assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5
 
 

tests/integrations/langgraph/test_langgraph.py

@@ -1384,6 +1384,4 @@ def original_invoke(self, *args, **kwargs):
     assert isinstance(parsed_messages, list)
     assert len(parsed_messages) == 1
     assert "small message 5" in str(parsed_messages[0])
-
-    assert invoke_span["data"][SPANDATA.META_GEN_AI_ORIGINAL_INPUT_MESSAGES_LENGTH] == 5
     assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5

tests/integrations/litellm/test_litellm.py

@@ -754,8 +754,6 @@ def test_litellm_message_truncation(sentry_init, capture_events):
     assert isinstance(parsed_messages, list)
     assert len(parsed_messages) == 1
     assert "small message 5" in str(parsed_messages[0])
-
-    assert chat_span["data"][SPANDATA.META_GEN_AI_ORIGINAL_INPUT_MESSAGES_LENGTH] == 5
     assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5
 
 

tests/integrations/openai/test_openai.py

@@ -1607,7 +1607,6 @@ def test_ai_client_span_responses_api(
         "gen_ai.usage.total_tokens": 30,
         "gen_ai.request.model": "gpt-4o",
         "gen_ai.response.text": "the model response",
-        "sentry.sdk_meta.gen_ai.input.messages.original_length": 1,
         "thread.id": mock.ANY,
         "thread.name": mock.ANY,
     }
@@ -1911,7 +1910,6 @@ async def test_ai_client_span_responses_async_api(
         "gen_ai.usage.output_tokens.reasoning": 8,
         "gen_ai.usage.total_tokens": 30,
         "gen_ai.response.text": "the model response",
-        "sentry.sdk_meta.gen_ai.input.messages.original_length": 1,
         "thread.id": mock.ANY,
         "thread.name": mock.ANY,
     }
@@ -2179,7 +2177,6 @@ async def test_ai_client_span_streaming_responses_async_api(
         "gen_ai.usage.total_tokens": 30,
         "gen_ai.request.model": "gpt-4o",
         "gen_ai.response.text": "the model response",
-        "sentry.sdk_meta.gen_ai.input.messages.original_length": 1,
         "thread.id": mock.ANY,
         "thread.name": mock.ANY,
     }