merge master

Author: alexander-alderman-webb

sentry_sdk/integrations/mcp.py

@@ -20,7 +20,12 @@
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
-from sentry_sdk.utils import nullcontext, package_version, safe_serialize
+from sentry_sdk.utils import (
+    has_data_collection_enabled,
+    nullcontext,
+    package_version,
+    safe_serialize,
+)
 
 try:
     from mcp.server.lowlevel import Server
@@ -322,13 +327,21 @@ async def _tool_handler_wrapper(
         original_kwargs: Original keyword arguments passed to the handler
         self: Optional instance for bound methods
     """
+    client = sentry_sdk.get_client()
+
     if original_kwargs is None:
         original_kwargs = {}
 
     handler_name, arguments = _extract_handler_data_from_args(
         "tool", original_args, original_kwargs
     )
 
+    if has_data_collection_enabled(client.options):
+        if not client.options["data_collection"]["gen_ai"]["inputs"]:
+            # Arguments can contain sensitive data and shouldn't be added to the span
+            # if the user has opted out via this config.
+            arguments = {}
+
     ctx = None
     try:
         ctx = request_ctx.get()
@@ -388,14 +401,17 @@ async def _tool_handler_wrapper(
                 return result
 
             # Get integration to check PII settings
-            integration = sentry_sdk.get_client().get_integration(MCPIntegration)
+            integration = client.get_integration(MCPIntegration)
             if integration is None:
                 return result
 
             # Check if we should include sensitive data
-            should_include_data = (
-                should_send_default_pii() and integration.include_prompts
-            )
+            should_include_data = False
+            if has_data_collection_enabled(client.options):
+                if client.options["data_collection"]["gen_ai"]["outputs"]:
+                    should_include_data = True
+            elif should_send_default_pii() and integration.include_prompts:
+                should_include_data = True
 
             extracted = _extract_tool_result_content(result)
             if extracted is not None and should_include_data:
@@ -422,15 +438,22 @@ async def _instrument_v2_tool_call(
     if ctx.params is None or ctx.params.get("name") is None:
         return await call_next(ctx)
 
+    client = sentry_sdk.get_client()
     handler_name = ctx.params["name"]
     arguments = ctx.params.get("arguments")
     if arguments is None:
         arguments = {}
 
+    if has_data_collection_enabled(client.options):
+        if not client.options["data_collection"]["gen_ai"]["inputs"]:
+            # Arguments can contain sensitive data and shouldn't be added to the span
+            # if the user has opted out via this config.
+            arguments = {}
+
     # Get request ID, session ID, and transport from context
     request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx)
 
-    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    span_streaming = has_span_streaming_enabled(client.options)
 
     # Start span and execute
     with _active_http_scopes(ctx=ctx):
@@ -474,22 +497,25 @@ async def _instrument_v2_tool_call(
                 return result
 
             # Get integration to check PII settings
-            integration = sentry_sdk.get_client().get_integration(MCPIntegration)
+            integration = client.get_integration(MCPIntegration)
             if integration is None:
                 return result
 
             # Check if we should include sensitive data
-            should_include_data = (
-                should_send_default_pii() and integration.include_prompts
-            )
+            should_include_result_data = False
+            if has_data_collection_enabled(client.options):
+                if client.options["data_collection"]["gen_ai"]["outputs"]:
+                    should_include_result_data = True
+            elif should_send_default_pii() and integration.include_prompts:
+                should_include_result_data = True
 
             result_content = result
             if "structuredContent" in result:
                 result_content = result["structuredContent"]
             elif isinstance(result.get("content"), list):
                 result_content = _extract_text_from_content_blocks(result["content"])
 
-            if result_content is not None and should_include_data:
+            if result_content is not None and should_include_result_data:
                 _set_span_data_attribute(
                     span,
                     SPANDATA.MCP_TOOL_RESULT_CONTENT,
@@ -526,10 +552,17 @@ async def _prompt_handler_wrapper(
     if original_kwargs is None:
         original_kwargs = {}
 
+    client = sentry_sdk.get_client()
     handler_name, arguments = _extract_handler_data_from_args(
         "prompt", original_args, original_kwargs
     )
 
+    if has_data_collection_enabled(client.options):
+        if not client.options["data_collection"]["gen_ai"]["inputs"]:
+            # Arguments can contain sensitive data and shouldn't be added to the span
+            # if the user has opted out via this config.
+            arguments = {}
+
     ctx = None
     try:
         ctx = request_ctx.get()
@@ -539,7 +572,7 @@ async def _prompt_handler_wrapper(
     # Get request ID, session ID, and transport from context
     request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx)
 
-    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    span_streaming = has_span_streaming_enabled(client.options)
 
     # Start span and execute
     with _active_http_scopes(ctx=ctx):
@@ -589,14 +622,17 @@ async def _prompt_handler_wrapper(
                 return result
 
             # Get integration to check PII settings
-            integration = sentry_sdk.get_client().get_integration(MCPIntegration)
+            integration = client.get_integration(MCPIntegration)
             if integration is None:
                 return result
 
             # Check if we should include sensitive data
-            should_include_data = (
-                should_send_default_pii() and integration.include_prompts
-            )
+            should_include_result_data = False
+            if has_data_collection_enabled(client.options):
+                if client.options["data_collection"]["gen_ai"]["inputs"]:
+                    should_include_result_data = True
+            elif should_send_default_pii() and integration.include_prompts:
+                should_include_result_data = True
 
             # For prompts, count messages and set role/content only for single-message prompts
             try:
@@ -619,7 +655,7 @@ async def _prompt_handler_wrapper(
                     )
 
                 # Only set role and content for single-message prompts if PII is allowed
-                if message_count == 1 and should_include_data and messages:
+                if message_count == 1 and should_include_result_data and messages:
                     first_message = messages[0]
                     # Extract role
                     role = None
@@ -675,15 +711,24 @@ async def _instrument_v2_prompt_get(
     if ctx.params is None or ctx.params.get("name") is None:
         return await call_next(ctx)
 
+    client = sentry_sdk.get_client()
     handler_name = ctx.params["name"]
+
     arguments = ctx.params.get("arguments")
+
     if arguments is None:
         arguments = {}
 
+    if has_data_collection_enabled(client.options):
+        if not client.options["data_collection"]["gen_ai"]["inputs"]:
+            # Arguments can contain sensitive data and shouldn't be added to the span
+            # if the user has opted out via this config.
+            arguments = {}
+
     # Get request ID, session ID, and transport from context
     request_id, session_id, mcp_transport = _get_request_context_data(ctx=ctx)
 
-    span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+    span_streaming = has_span_streaming_enabled(client.options)
 
     # Start span and execute
     with _active_http_scopes(ctx=ctx):
@@ -726,14 +771,17 @@ async def _instrument_v2_prompt_get(
                 return result
 
             # Get integration to check PII settings
-            integration = sentry_sdk.get_client().get_integration(MCPIntegration)
+            integration = client.get_integration(MCPIntegration)
             if integration is None:
                 return result
 
             # Check if we should include sensitive data
-            should_include_data = (
-                should_send_default_pii() and integration.include_prompts
-            )
+            should_include_result_data = False
+            if has_data_collection_enabled(client.options):
+                if client.options["data_collection"]["gen_ai"]["inputs"]:
+                    should_include_result_data = True
+            elif should_send_default_pii() and integration.include_prompts:
+                should_include_result_data = True
 
             # For prompts, count messages and set role/content only for single-message prompts
             try:
@@ -751,7 +799,7 @@ async def _instrument_v2_prompt_get(
                     )
 
                 # Only set role and content for single-message prompts if PII is allowed
-                if message_count == 1 and should_include_data and messages:
+                if message_count == 1 and should_include_result_data and messages:
                     first_message = messages[0]
                     # Extract role
                     role = None

sentry_sdk/scope.py

@@ -861,10 +861,7 @@ def set_transaction_name(self, name: str, source: "Optional[str]" = None) -> Non
         """Set the transaction name and optionally the transaction source."""
         self._transaction = name
         if self._span:
-            if isinstance(self._span, NoOpStreamedSpan):
-                return
-
-            elif isinstance(self._span, StreamedSpan):
+            if isinstance(self._span, StreamedSpan):
                 self._span._segment.name = name
                 if source:
                     self._span._segment.set_attribute(
@@ -925,12 +922,24 @@ def streamed_span(self, span: "Optional[StreamedSpan]") -> None:
 
         # Also set _transaction and _transaction_info in streaming mode as this
         # is used for populating events and linking them to segments
-        if type(span) is StreamedSpan and span._is_segment():
+        if not isinstance(span, StreamedSpan) or not span._is_segment():
+            return
+
+        if type(span) is StreamedSpan:
             self._transaction = span.name
             if span._attributes.get("sentry.segment.name.source"):
                 self._transaction_info["source"] = str(
                     span._attributes["sentry.segment.name.source"]
                 )
+            return
+
+        if type(span) is NoOpStreamedSpan:
+            if span._name is not None:
+                self._transaction = span.name
+            if span._attributes.get("sentry.segment.name.source"):
+                self._transaction_info["source"] = str(
+                    span._attributes["sentry.segment.name.source"]
+                )
 
     @property
     def profile(self) -> "Optional[Profile]":
@@ -1309,6 +1318,8 @@ def start_streamed_span(
 
             if is_ignored_span(name, attributes):
                 return NoOpStreamedSpan(
+                    name=name,
+                    attributes=attributes,
                     scope=self,
                     segment=None,
                     trace_id=propagation_context.trace_id,
@@ -1329,6 +1340,8 @@ def start_streamed_span(
 
             if sampled is False or sampled is None:
                 return NoOpStreamedSpan(
+                    name=name,
+                    attributes=attributes,
                     scope=self,
                     segment=None,
                     trace_id=propagation_context.trace_id,
@@ -1359,6 +1372,8 @@ def start_streamed_span(
         with new_scope():
             if is_ignored_span(name, attributes):
                 return NoOpStreamedSpan(
+                    name=name,
+                    attributes=attributes,
                     segment=parent_span._segment,
                     trace_id=parent_span.trace_id,
                     parent_span_id=parent_span.span_id,
@@ -1368,6 +1383,8 @@ def start_streamed_span(
 
             if isinstance(parent_span, NoOpStreamedSpan):
                 return NoOpStreamedSpan(
+                    name=name,
+                    attributes=attributes,
                     segment=parent_span._segment,
                     trace_id=parent_span.trace_id,
                     parent_span_id=parent_span.span_id,

sentry_sdk/traces.py

@@ -632,6 +632,8 @@ class NoOpStreamedSpan(StreamedSpan):
 
     def __init__(
         self,
+        name: "Optional[str]" = None,
+        attributes: "Optional[Attributes]" = None,
         segment: "Optional[StreamedSpan]" = None,
         trace_id: "Optional[str]" = None,
         parent_span_id: "Optional[str]" = None,
@@ -643,6 +645,13 @@ def __init__(
         sample_rand: "Optional[float]" = None,
         sample_rate: "Optional[float]" = None,
     ) -> None:
+        self._name = name  # type: ignore[assignment]
+        self._attributes = {}
+        if attributes is not None and "sentry.segment.name.source" in attributes:
+            self.set_attribute(
+                "sentry.segment.name.source", attributes["sentry.segment.name.source"]
+            )
+
         self._span_id: "Optional[str]" = None
 
         self._sampled = sampled
@@ -721,10 +730,17 @@ def get_attributes(self) -> "Attributes":
         return {}
 
     def set_attribute(self, key: str, value: "AttributeValue") -> None:
-        pass
+        if key != "sentry.segment.name.source":
+            return
+
+        super().set_attribute("sentry.segment.name.source", value)
 
     def set_attributes(self, attributes: "Attributes") -> None:
-        pass
+        for key, value in attributes.items():
+            if key != "sentry.segment.name.source":
+                continue
+
+            self.set_attribute("sentry.segment.name.source", value)
 
     def remove_attribute(self, key: str) -> None:
         pass
@@ -739,11 +755,11 @@ def status(self, status: "Union[SpanStatus, str]") -> None:
 
     @property
     def name(self) -> str:
-        return ""
+        return self._name or ""
 
     @name.setter
-    def name(self, value: str) -> None:
-        pass
+    def name(self, name: str) -> None:
+        self._name = name
 
     @property
     def active(self) -> bool: