refactor: update span lifecycle to use sdk over setting context manually

Author: keith-decker

util/opentelemetry-util-genai/src/opentelemetry/util/genai/handler.py

@@ -29,7 +29,7 @@
     handler = get_telemetry_handler()
 
     # Create an invocation object with your request data
-    # The span and context_token attributes are set by the TelemetryHandler, and
+    # The span and span_scope attributes are set by the TelemetryHandler, and
     # managed by the TelemetryHandler during the lifecycle of the span.
 
     # Use the context manager to manage the lifecycle of an LLM invocation.
@@ -63,7 +63,6 @@
 from contextlib import contextmanager
 from typing import Iterator, Optional
 
-from opentelemetry import context as otel_context
 from opentelemetry.semconv._incubating.attributes import (
     gen_ai_attributes as GenAI,
 )
@@ -72,13 +71,16 @@
     SpanKind,
     TracerProvider,
     get_tracer,
-    set_span_in_context,
+    use_span,
 )
 from opentelemetry.util.genai.span_utils import (
     _apply_error_attributes,
     _apply_finish_attributes,
 )
-from opentelemetry.util.genai.types import Error, LLMInvocation
+from opentelemetry.util.genai.types import (
+    Error,
+    LLMInvocation,
+)
 from opentelemetry.util.genai.version import __version__
 
 
@@ -101,41 +103,56 @@ def start_llm(
         invocation: LLMInvocation,
     ) -> LLMInvocation:
         """Start an LLM invocation and create a pending span entry."""
-        # Create a span and attach it as current; keep the token to detach later
+        # Create a span and activate it using the OpenTelemetry helper scope
         span = self._tracer.start_span(
             name=f"{GenAI.GenAiOperationNameValues.CHAT.value} {invocation.request_model}",
             kind=SpanKind.CLIENT,
         )
         invocation.span = span
-        invocation.context_token = otel_context.attach(
-            set_span_in_context(span)
+        scope = use_span(
+            span,
+            end_on_exit=False,
+            record_exception=False,
+            set_status_on_exception=False,
         )
+        scope.__enter__()
+        invocation.span_scope = scope
         return invocation
 
     def stop_llm(self, invocation: LLMInvocation) -> LLMInvocation:  # pylint: disable=no-self-use
         """Finalize an LLM invocation successfully and end its span."""
-        if invocation.context_token is None or invocation.span is None:
+        if invocation.span_scope is None or invocation.span is None:
             # TODO: Provide feedback that this invocation was not started
             return invocation
 
-        _apply_finish_attributes(invocation.span, invocation)
-        # Detach context and end span
-        otel_context.detach(invocation.context_token)
-        invocation.span.end()
+        scope = invocation.span_scope
+        span = invocation.span
+        try:
+            _apply_finish_attributes(span, invocation)
+        finally:
+            scope.__exit__(None, None, None)
+            span.end()
+            invocation.span_scope = None
+            invocation.span = None
         return invocation
 
     def fail_llm(  # pylint: disable=no-self-use
         self, invocation: LLMInvocation, error: Error
     ) -> LLMInvocation:
         """Fail an LLM invocation and end its span with error status."""
-        if invocation.context_token is None or invocation.span is None:
+        if invocation.span_scope is None or invocation.span is None:
             # TODO: Provide feedback that this invocation was not started
             return invocation
 
-        _apply_error_attributes(invocation.span, error)
-        # Detach context and end span
-        otel_context.detach(invocation.context_token)
-        invocation.span.end()
+        scope = invocation.span_scope
+        span = invocation.span
+        try:
+            _apply_error_attributes(span, error)
+        finally:
+            scope.__exit__(None, None, None)
+            span.end()
+            invocation.span_scope = None
+            invocation.span = None
         return invocation
 
     @contextmanager
@@ -144,21 +161,22 @@ def llm(
     ) -> Iterator[LLMInvocation]:
         """Context manager for LLM invocations.
 
-        Only set data attributes on the invocation object, do not modify the span or context.
-
-        Starts the span on entry. On normal exit, finalizes the invocation and ends the span.
-        If an exception occurs inside the context, marks the span as error, ends it, and
-        re-raises the original exception.
+        Only set data attributes on the invocation object, do not modify the span or
+        context. Starts the span on entry. On normal exit, finalizes the invocation and
+        ends the span. If an exception occurs inside the context, marks the span as
+        error, ends it, and re-raises the original exception.
         """
         if invocation is None:
-            invocation = LLMInvocation(
-                request_model="",
-            )
+            invocation = LLMInvocation(request_model="")
+
         self.start_llm(invocation)
         try:
             yield invocation
         except Exception as exc:
-            self.fail_llm(invocation, Error(message=str(exc), type=type(exc)))
+            self.fail_llm(
+                invocation,
+                Error(message=str(exc), type=type(exc)),
+            )
             raise
         self.stop_llm(invocation)
 

util/opentelemetry-util-genai/src/opentelemetry/util/genai/span_utils.py

@@ -124,9 +124,14 @@ def _apply_finish_attributes(span: Span, invocation: LLMInvocation) -> None:
 
 
 def _apply_error_attributes(span: Span, error: Error) -> None:
-    """Apply status and error attributes common to error() paths."""
-    span.set_status(Status(StatusCode.ERROR, error.message))
+    """Apply status and error attributes common to error() paths.
+
+    The helper stays separate from the context manager's exception handling so we can
+    attach the GenAI-specific error attributes defined in the semantic conventions
+    (see https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-agent-spans.md).
+    """
     if span.is_recording():
+        span.set_status(Status(StatusCode.ERROR, error.message))
         span.set_attribute(ErrorAttributes.ERROR_TYPE, error.type.__qualname__)
 
 

util/opentelemetry-util-genai/src/opentelemetry/util/genai/types.py

@@ -13,18 +13,21 @@
 # limitations under the License.
 
 
-from contextvars import Token
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Type, Union
+from typing import (
+    Any,
+    ContextManager,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Type,
+    Union,
+)
 
-from typing_extensions import TypeAlias
-
-from opentelemetry.context import Context
 from opentelemetry.trace import Span
 
-ContextToken: TypeAlias = Token[Context]
-
 
 class ContentCapturingMode(Enum):
     # Do not capture content (default).
@@ -95,13 +98,16 @@ def _new_str_any_dict() -> Dict[str, Any]:
 class LLMInvocation:
     """
     Represents a single LLM call invocation. When creating an LLMInvocation object,
-    only update the data attributes. The span and context_token attributes are
-    set by the TelemetryHandler.
+    only update the data attributes. The span and span_scope attributes are set by the
+    TelemetryHandler.
     """
 
     request_model: str
-    context_token: Optional[ContextToken] = None
     span: Optional[Span] = None
+    # Internal handle returned by opentelemetry.trace.use_span to keep the span active.
+    span_scope: Optional[ContextManager[Span]] = field(
+        default=None, compare=False, repr=False
+    )
     input_messages: List[InputMessage] = field(
         default_factory=_new_input_messages
     )