ref: Support outgoing trace propagation in span first (18)

sentry_sdk/_span_batcher.py

@@ -116,9 +116,7 @@ def _flush(self) -> None:
             envelopes = []
             for trace_id, spans in self._span_buffer.items():
                 if spans:
-                    # TODO[span-first]
-                    # dsc = spans[0].dynamic_sampling_context()
-                    dsc = None
+                    dsc = spans[0]._dynamic_sampling_context()
 
                     # Max per envelope is 1000, so if we happen to have more than
                     # 1000 spans in one bucket, we'll need to separate them.

sentry_sdk/ai/utils.py

@@ -14,6 +14,8 @@
 
 import sentry_sdk
 from sentry_sdk.utils import logger
+from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 MAX_GEN_AI_MESSAGE_BYTES = 20_000  # 20KB
 # Maximum characters when only a single message is left after bytes truncation
@@ -525,7 +527,14 @@
 
 
 def get_start_span_function() -> "Callable[..., Any]":
+    if has_span_streaming_enabled(sentry_sdk.get_client().options):
+        return sentry_sdk.traces.start_span
+
     current_span = sentry_sdk.get_current_span()
+    if isinstance(current_span, StreamedSpan):
+        # mypy
+        return sentry_sdk.traces.start_span
+
     transaction_exists = (
         current_span is not None and current_span.containing_transaction is not None
     )

sentry_sdk/api.py

@@ -6,6 +6,7 @@
 from sentry_sdk._init_implementation import init
 from sentry_sdk.consts import INSTRUMENTER
 from sentry_sdk.scope import Scope, _ScopeManager, new_scope, isolation_scope
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import NoOpSpan, Transaction, trace
 from sentry_sdk.crons import monitor
 
@@ -37,6 +38,7 @@
         LogLevelStr,
         SamplingContext,
     )
+    from sentry_sdk.traces import StreamedSpan
     from sentry_sdk.tracing import Span, TransactionKwargs
 
     T = TypeVar("T")
@@ -409,7 +411,9 @@ def set_measurement(name: str, value: float, unit: "MeasurementUnit" = "") -> No
         transaction.set_measurement(name, value, unit)
 
 
-def get_current_span(scope: "Optional[Scope]" = None) -> "Optional[Span]":
+def get_current_span(
+    scope: "Optional[Scope]" = None,
+) -> "Optional[Union[Span, StreamedSpan]]":
     """
     Returns the currently active span if there is one running, otherwise `None`
     """
@@ -525,6 +529,16 @@ def update_current_span(
     if current_span is None:
         return
 
+    if isinstance(current_span, StreamedSpan):
+        warnings.warn(
+            "The `update_current_span` API isn't available in streaming mode. "
+            "Retrieve the current span with get_current_span() and use its API "
+            "directly.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return
+
     if op is not None:
         current_span.op = op
 

sentry_sdk/feature_flags.py

@@ -1,6 +1,7 @@
 import copy
 import sentry_sdk
 from sentry_sdk._lru_cache import LRUCache
+from sentry_sdk.tracing import Span
 from threading import Lock
 
 from typing import TYPE_CHECKING, Any
@@ -61,5 +62,5 @@ def add_feature_flag(flag: str, result: bool) -> None:
     flags.set(flag, result)
 
     span = sentry_sdk.get_current_span()
-    if span:
+    if span and isinstance(span, Span):
         span.set_flag(f"flag.evaluation.{flag}", result)

sentry_sdk/integrations/celery/__init__.py

@@ -14,7 +14,7 @@
 )
 from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch
 from sentry_sdk.integrations.logging import ignore_logger
-from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, TransactionSource
+from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span, TransactionSource
 from sentry_sdk.tracing_utils import Baggage
 from sentry_sdk.utils import (
     capture_internal_exceptions,
@@ -34,7 +34,6 @@
     from typing import Union
 
     from sentry_sdk._types import EventProcessor, Event, Hint, ExcInfo
-    from sentry_sdk.tracing import Span
 
     F = TypeVar("F", bound=Callable[..., Any])
 
@@ -100,7 +99,10 @@ def _set_status(status: str) -> None:
     with capture_internal_exceptions():
         scope = sentry_sdk.get_current_scope()
         if scope.span is not None:
-            scope.span.set_status(status)
+            if isinstance(scope.span, Span):
+                scope.span.set_status(status)
+            else:
+                scope.span.status = "ok" if status == "ok" else "error"
 
 
 def _capture_exception(task: "Any", exc_info: "ExcInfo") -> None:

sentry_sdk/integrations/openai_agents/utils.py

@@ -11,6 +11,7 @@
 from sentry_sdk.consts import SPANDATA, SPANSTATUS, OP
 from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing import Span
 from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import event_from_exception, safe_serialize
 from sentry_sdk.ai._openai_completions_api import _transform_system_instructions
@@ -22,10 +23,10 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Union
     from agents import Usage, TResponseInputItem
 
-    from sentry_sdk.tracing import Span
+    from sentry_sdk.traces import StreamedSpan
     from sentry_sdk._types import TextPart
 
 try:
@@ -46,8 +47,15 @@ def _capture_exception(exc: "Any") -> None:
     sentry_sdk.capture_event(event, hint=hint)
 
 
-def _record_exception_on_span(span: "Span", error: Exception) -> "Any":
+def _record_exception_on_span(
+    span: "Union[Span, StreamedSpan]", error: Exception
+) -> "Any":
     set_span_errored(span)
+
+    if not isinstance(span, Span):
+        # TODO[span-first]: make this work with streamedspans
+        return
+
     span.set_data("span.status", "error")
 
     # Optionally capture the error details if we have them

sentry_sdk/scope.py

@@ -583,8 +583,12 @@
         client = self.get_client()
 
         # If we have an active span, return traceparent from there
-        if has_tracing_enabled(client.options) and self.span is not None:
-            return self.span.to_traceparent()
+        if (
+            has_tracing_enabled(client.options)
+            and self.span is not None
+            and not isinstance(self.span, NoOpStreamedSpan)
+        ):
+            return self.span._to_traceparent()
 
         # else return traceparent from the propagation context
         return self.get_active_propagation_context().to_traceparent()
@@ -597,8 +601,12 @@
         client = self.get_client()
 
         # If we have an active span, return baggage from there
-        if has_tracing_enabled(client.options) and self.span is not None:
-            return self.span.to_baggage()
+        if (
+            has_tracing_enabled(client.options)
+            and self.span is not None
+            and not isinstance(self.span, NoOpStreamedSpan)
+        ):
+            return self.span._to_baggage()
 
         # else return baggage from the propagation context
         return self.get_active_propagation_context().get_baggage()
@@ -607,8 +615,12 @@
         """
         Returns the Sentry "trace" context from the Propagation Context.
         """
-        if has_tracing_enabled(self.get_client().options) and self._span is not None:
-            return self._span.get_trace_context()
+        if (
+            has_tracing_enabled(self.get_client().options)
+            and self._span is not None
+            and not isinstance(self.span, NoOpStreamedSpan)
+        ):
+            return self._span._get_trace_context()
 
         # if we are tracing externally (otel), those values take precedence
         external_propagation_context = get_external_propagation_context()
@@ -673,7 +685,7 @@
         span = span or self.span
 
         if has_tracing_enabled(client.options) and span is not None:
-            for header in span.iter_headers():
+            for header in span._iter_headers():
                 yield header
         elif has_external_propagation_context():
             # when we have an external_propagation_context (otlp)
@@ -720,7 +732,7 @@
         self.clear_breadcrumbs()
         self._should_capture: bool = True
 
-        self._span: "Optional[Span]" = None
+        self._span: "Optional[Union[Span, StreamedSpan]]" = None
         self._session: "Optional[Session]" = None
         self._force_auto_session_tracking: "Optional[bool]" = None
 
@@ -774,6 +786,14 @@
         if self._span is None:
             return None
 
+        if isinstance(self._span, StreamedSpan):
+            warnings.warn(
+                "Scope.transaction is not available in streaming mode.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return None
+
         # there is an orphan span on the scope
         if self._span.containing_transaction is None:
             return None
@@ -803,17 +823,36 @@
             "Assigning to scope.transaction directly is deprecated: use scope.set_transaction_name() instead."
         )
         self._transaction = value
-        if self._span and self._span.containing_transaction:
-            self._span.containing_transaction.name = value
+        if self._span:
+            if isinstance(self._span, StreamedSpan):
+                warnings.warn(
+                    "Scope.transaction is not available in streaming mode.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                return None
+
+            if self._span.containing_transaction:
+                self._span.containing_transaction.name = value
 
     def set_transaction_name(self, name: str, source: "Optional[str]" = None) -> None:
         """Set the transaction name and optionally the transaction source."""
         self._transaction = name
-
-        if self._span and self._span.containing_transaction:
-            self._span.containing_transaction.name = name
-            if source:
-                self._span.containing_transaction.source = source
+        if self._span:
+            if isinstance(self._span, NoOpStreamedSpan):
+                return
+
+            elif isinstance(self._span, StreamedSpan):
+                self._span._segment.name = name
+                if source:
+                    self._span._segment.set_attribute(
+                        "sentry.span.source", getattr(source, "value", source)
+                    )
+
+            elif self._span.containing_transaction:
+                self._span.containing_transaction.name = name
+                if source:
+                    self._span.containing_transaction.source = source
 
         if source:
             self._transaction_info["source"] = source
@@ -836,12 +875,12 @@
             session.update(user=value)
 
     @property
-    def span(self) -> "Optional[Span]":
+    def span(self) -> "Optional[Union[Span, StreamedSpan]]":
         """Get/set current tracing span or transaction."""
         return self._span
 
     @span.setter
-    def span(self, span: "Optional[Span]") -> None:
+    def span(self, span: "Optional[Union[Span, StreamedSpan]]") -> None:
         self._span = span
         # XXX: this differs from the implementation in JS, there Scope.setSpan
         # does not set Scope._transactionName.
@@ -1150,6 +1189,15 @@
         be removed in the next major version. Going forward, it should only
         be used by the SDK itself.
         """
+        client = sentry_sdk.get_client()
+        if has_span_streaming_enabled(client.options):
+            warnings.warn(
+                "Scope.start_span is not available in streaming mode.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return NoOpSpan()
+
         if kwargs.get("description") is not None:
             warnings.warn(
                 "The `description` parameter is deprecated. Please use `name` instead.",
@@ -1169,6 +1217,9 @@
 
             # get current span or transaction
             span = self.span or self.get_isolation_scope().span
+            if isinstance(span, StreamedSpan):
+                # make mypy happy
+                return NoOpSpan()
 
             if span is None:
                 # New spans get the `trace_id` from the scope