Merge branch 'master' into toxgen/update-05-21

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/spans/ai_client.py

@@ -31,7 +31,7 @@ def ai_client_span(
 
     span = sentry_sdk.start_span(
         op=OP.GEN_AI_CHAT,
-        description=f"chat {model_name}",
+        name=f"chat {model_name}",
         origin=SPAN_ORIGIN,
     )
     # TODO-anton: remove hardcoded stuff and replace something that also works for embedding and so on

sentry_sdk/integrations/openai_agents/utils.py

@@ -224,7 +224,7 @@ def _create_mcp_execute_tool_spans(
         if output.__class__.__name__ == "McpCall":
             with sentry_sdk.start_span(
                 op=OP.GEN_AI_EXECUTE_TOOL,
-                description=f"execute_tool {output.name}",
+                name=f"execute_tool {output.name}",
                 start_timestamp=span.start_timestamp,
             ) as execute_tool_span:
                 execute_tool_span.set_data(SPANDATA.GEN_AI_TOOL_NAME, output.name)

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py

@@ -19,7 +19,6 @@
 from .utils import (
     _serialize_binary_content_item,
     _serialize_image_url_item,
-    _set_usage_data,
 )
 
 if TYPE_CHECKING:
@@ -150,16 +149,6 @@ def update_invoke_agent_span(span: "sentry_sdk.tracing.Span", result: "Any") ->
             span, SPANDATA.GEN_AI_RESPONSE_TEXT, str(output), unpack=False
         )
 
-    # Set token usage data if available
-    if hasattr(result, "usage") and callable(result.usage):
-        try:
-            usage = result.usage()
-            if usage:
-                _set_usage_data(span, usage)
-        except Exception:
-            # If usage() call fails, continue without setting usage data
-            pass
-
     # Set model name from response if available
     if hasattr(result, "response"):
         try:

sentry_sdk/integrations/strawberry.py

@@ -8,7 +8,9 @@
 from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
 from sentry_sdk.integrations.logging import ignore_logger
 from sentry_sdk.scope import should_send_default_pii
-from sentry_sdk.tracing import TransactionSource
+from sentry_sdk.traces import SegmentSource
+from sentry_sdk.tracing import Span, TransactionSource
+from sentry_sdk.tracing_utils import StreamedSpan, has_span_streaming_enabled
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
@@ -183,50 +185,111 @@ def on_operation(self) -> "Generator[None, None, None]":
         event_processor = _make_request_event_processor(self.execution_context)
         scope.add_event_processor(event_processor)
 
-        graphql_span = sentry_sdk.start_span(
-            op=op,
-            name=description,
-            origin=StrawberryIntegration.origin,
-        )
-        graphql_span.__enter__()
+        client = sentry_sdk.get_client()
+        is_span_streaming_enabled = has_span_streaming_enabled(client.options)
+        if is_span_streaming_enabled:
+            additional_attributes: "dict[str, Any]" = {}
+
+            if should_send_default_pii():
+                additional_attributes["graphql.document"] = self.execution_context.query
+
+            if operation_name:
+                additional_attributes["graphql.operation.name"] = operation_name
+
+            graphql_span = sentry_sdk.traces.start_span(
+                name=description,
+                attributes={
+                    "sentry.origin": StrawberryIntegration.origin,
+                    "sentry.op": op,
+                    "graphql.operation.type": operation_type,
+                    **additional_attributes,
+                },
+            )
+        else:
+            graphql_span = sentry_sdk.start_span(
+                op=op,
+                name=description,
+                origin=StrawberryIntegration.origin,
+            )
+            graphql_span.__enter__()
+
+        if type(graphql_span) is Span:
+            if should_send_default_pii():
+                graphql_span.set_data("graphql.document", self.execution_context.query)
 
-        graphql_span.set_data("graphql.operation.type", operation_type)
-        graphql_span.set_data("graphql.operation.name", operation_name)
-        if should_send_default_pii():
-            graphql_span.set_data("graphql.document", self.execution_context.query)
-        graphql_span.set_data("graphql.resource_name", self._resource_name)
+            graphql_span.set_data("graphql.operation.type", operation_type)
+            graphql_span.set_data("graphql.operation.name", operation_name)
+            # This attribute is being removed in streamed spans
+            graphql_span.set_data("graphql.resource_name", self._resource_name)
 
         yield
 
-        transaction = graphql_span.containing_transaction
-        if transaction and self.execution_context.operation_name:
-            transaction.name = self.execution_context.operation_name
-            transaction.source = TransactionSource.COMPONENT
-            transaction.op = op
+        if type(graphql_span) is StreamedSpan:
+            if self.execution_context.operation_name:
+                segment = graphql_span._segment
+                segment.set_attribute("sentry.span.source", SegmentSource.COMPONENT)
+                segment.set_attribute("sentry.op", op)
+                segment.name = self.execution_context.operation_name
+        elif isinstance(graphql_span, Span):
+            transaction = graphql_span.containing_transaction
+            if transaction and self.execution_context.operation_name:
+                transaction.name = self.execution_context.operation_name
+                transaction.source = TransactionSource.COMPONENT
+                transaction.op = op
 
         graphql_span.__exit__(None, None, None)
 
     def on_validate(self) -> "Generator[None, None, None]":
-        validation_span = sentry_sdk.start_span(
-            op=OP.GRAPHQL_VALIDATE,
-            name="validation",
-            origin=StrawberryIntegration.origin,
-        )
+        client = sentry_sdk.get_client()
+        is_span_streaming_enabled = has_span_streaming_enabled(client.options)
+
+        if is_span_streaming_enabled:
+            validation_span = sentry_sdk.traces.start_span(
+                name="validation",
+                attributes={
+                    "sentry.op": OP.GRAPHQL_VALIDATE,
+                    "sentry.origin": StrawberryIntegration.origin,
+                },
+            )
+        else:
+            validation_span = sentry_sdk.start_span(
+                op=OP.GRAPHQL_VALIDATE,
+                name="validation",
+                origin=StrawberryIntegration.origin,
+            )
 
         yield
 
-        validation_span.finish()
+        if isinstance(validation_span, StreamedSpan):
+            validation_span.end()
+        else:
+            validation_span.finish()
 
     def on_parse(self) -> "Generator[None, None, None]":
-        parsing_span = sentry_sdk.start_span(
-            op=OP.GRAPHQL_PARSE,
-            name="parsing",
-            origin=StrawberryIntegration.origin,
-        )
+        client = sentry_sdk.get_client()
+        is_span_streaming_enabled = has_span_streaming_enabled(client.options)
+
+        if is_span_streaming_enabled:
+            parsing_span = sentry_sdk.traces.start_span(
+                name="parsing",
+                attributes={
+                    "sentry.op": OP.GRAPHQL_PARSE,
+                    "sentry.origin": StrawberryIntegration.origin,
+                },
+            )
+        else:
+            parsing_span = sentry_sdk.start_span(
+                op=OP.GRAPHQL_PARSE,
+                name="parsing",
+                origin=StrawberryIntegration.origin,
+            )
 
         yield
 
-        parsing_span.finish()
+        if isinstance(parsing_span, StreamedSpan):
+            parsing_span.end()
+        else:
+            parsing_span.finish()
 
     def should_skip_tracing(
         self,
@@ -263,6 +326,18 @@ async def resolve(
 
         field_path = "{}.{}".format(info.parent_type, info.field_name)
 
+        client = sentry_sdk.get_client()
+        is_span_streaming_enabled = has_span_streaming_enabled(client.options)
+        if is_span_streaming_enabled:
+            with sentry_sdk.traces.start_span(
+                name=f"resolving {field_path}",
+                attributes={
+                    "sentry.origin": StrawberryIntegration.origin,
+                    "sentry.op": OP.GRAPHQL_RESOLVE,
+                },
+            ):
+                return await self._resolve(_next, root, info, *args, **kwargs)
+
         with sentry_sdk.start_span(
             op=OP.GRAPHQL_RESOLVE,
             name="resolving {}".format(field_path),
@@ -290,6 +365,18 @@ def resolve(
 
         field_path = "{}.{}".format(info.parent_type, info.field_name)
 
+        client = sentry_sdk.get_client()
+        is_span_streaming_enabled = has_span_streaming_enabled(client.options)
+        if is_span_streaming_enabled:
+            with sentry_sdk.traces.start_span(
+                name=f"resolving {field_path}",
+                attributes={
+                    "sentry.origin": StrawberryIntegration.origin,
+                    "sentry.op": OP.GRAPHQL_RESOLVE,
+                },
+            ):
+                return _next(root, info, *args, **kwargs)
+
         with sentry_sdk.start_span(
             op=OP.GRAPHQL_RESOLVE,
             name="resolving {}".format(field_path),