fix(otel): balance otel context attach and detach

- Track every context.attach() token in the plugin and release it at
  the matching lifecycle end: user-function scopes in
  on_user_function_end, the invocation scope during invocation cleanup
- Replace the re-attach of the enclosing span in on_user_function_end
  with a detach, so nested operations manage a balanced scope instead
  of stacking a new one
- Release scopes still open at invocation end (a suspended user
  function or an aborted invocation) so nothing survives into a warm
  invocation
- Skip tokens recorded on another thread, which OpenTelemetry cannot
  reset
- Replace the autouse OTel context reset in the plugin tests with an
  assertion that each test leaves the context as it found it, and
  cover nested child contexts, sequential steps, failures, suspension,
  and warm invocation reuse

Resolves #643

Author: Alex Wang

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py

@@ -17,6 +17,11 @@
 - context is attached inside the synchronous ``on_user_function_*`` hooks and
 log correlation is handled by :mod:`log_filter`), the hook wiring mirrors the
 existing :class:`~aws_durable_execution_sdk_python_otel.invocation_plugin.InvocationOtelPlugin`.
+
+Every context the plugin attaches is tracked by its token and detached at the
+matching lifecycle end: a user-function scope is released in
+``on_user_function_end`` and the invocation scope in ``on_invocation_end``, so
+the plugin never leaves an ended or suspended span current.
 """
 
 from __future__ import annotations
@@ -75,6 +80,9 @@
 # Registry key for the invocation span (operations use their operation_id).
 _INVOCATION_KEY = "__invocation__"
 
+# Token key for the invocation-level context scope attached at invocation start.
+_INVOCATION_CONTEXT_KEY = "__invocation_context__"
+
 
 def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None:
     """Convert a datetime to an OTel timestamp (ns since epoch), or None."""
@@ -114,6 +122,11 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
         self._workflow_span: Span | None = None
         self._invocation_span: Span | None = None
         self._operation_spans: dict[str, Span] = {}
+        # Tokens returned by context.attach(), keyed by the span registry key,
+        # paired with the thread that attached them. Every attach the plugin
+        # owns is released through _detach_context so the plugin never leaves a
+        # scope on the context stack.
+        self._context_tokens: dict[str, tuple[int, object]] = {}
         self._lock = threading.RLock()
         self._tracing_enabled = False
 
@@ -157,6 +170,46 @@ def _pop_span(self, key: str) -> Span | None:
     def _attempt_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
         return f"{info.operation_id}:attempt:{info.attempt or 1}"
 
+    # ------------------------------------------------------------------
+    # Context scope helpers
+    # ------------------------------------------------------------------
+    def _attach_context(self, key: str, new_context: Context) -> None:
+        """Attach a context and remember its token under ``key``."""
+        with self._lock:
+            self._context_tokens[key] = (
+                threading.get_ident(),
+                otel_context.attach(new_context),
+            )
+
+    def _detach_context(self, key: str) -> None:
+        """Detach the context attached under ``key``, restoring its predecessor.
+
+        A context token can only be reset on the thread that created it, so a
+        token recorded on another thread is dropped instead of detached (OTel
+        logs an error for a cross-thread reset). In practice the pairs always
+        line up: invocation hooks run on the Lambda handler thread and
+        user-function hooks run on the thread executing user code.
+        """
+        with self._lock:
+            entry = self._context_tokens.pop(key, None)
+        if entry is None:
+            return
+        thread_ident, token = entry
+        if thread_ident == threading.get_ident():
+            otel_context.detach(token)  # type: ignore[arg-type]
+
+    def _detach_remaining_contexts(self) -> None:
+        """Release scopes still open, newest first, so nothing outlives the plugin.
+
+        Reached when a lifecycle end hook never fires -- for example a user
+        function that suspends, or a warm invocation that starts before the
+        previous one was cleaned up.
+        """
+        with self._lock:
+            keys = list(reversed(self._context_tokens))
+        for key in keys:
+            self._detach_context(key)
+
     def get_current_span_context(self) -> SpanContext | None:
         """Return the active span context for log correlation (see log_filter)."""
         span_context = trace.get_current_span().get_span_context()
@@ -225,10 +278,13 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
         self._start_invocation_span(info)
 
         # Make the Workflow span the active span so auto-instrumented spans
-        # created during the invocation become its children.
+        # created during the invocation become its children. The token is
+        # released in _reset_state at invocation end, restoring the context that
+        # was active before the invocation started.
         if self._workflow_span is not None:
-            otel_context.attach(
-                trace.set_span_in_context(self._workflow_span, self._extracted_context)
+            self._attach_context(
+                _INVOCATION_CONTEXT_KEY,
+                trace.set_span_in_context(self._workflow_span, self._extracted_context),
             )
 
     def _start_workflow_span(self, info: InvocationStartInfo) -> None:
@@ -324,6 +380,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
                 logger.exception("force_flush failed at invocation end")
 
     def _reset_state(self) -> None:
+        self._detach_remaining_contexts()
         self._execution_arn = ""
         self._execution_trace_id = None
         self._extracted_context = None
@@ -439,25 +496,29 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
                 info.parent_id
             )
             name = f"{info.name or info.operation_id} attempt {info.attempt or 1}"
+            key = self._attempt_key(info)
             span = self._start_span(
                 operation_id=info.operation_id,
                 name=name,
                 info=info,
                 parent=parent,
                 start_time=info.start_time,
-                span_key=self._attempt_key(info),
+                span_key=key,
                 deterministic=False,
             )
         else:  # CONTEXT
             parent = self._resolve_parent(info.parent_id)
+            key = info.operation_id
             span = self._start_span(
                 operation_id=info.operation_id,
                 name=info.name or info.operation_id,
                 info=info,
                 parent=parent,
                 start_time=info.start_time,
             )
-        otel_context.attach(trace.set_span_in_context(span, self._extracted_context))
+        self._attach_context(
+            key, trace.set_span_in_context(span, self._extracted_context)
+        )
 
     def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
         logger.debug("Durable user function ended: %s", info)
@@ -500,16 +561,11 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
             if popped is not None:
                 popped.end(end_time=_to_otel_timestamp(end_time))
 
-        # Restore the enclosing span as active (parent op, else invocation/workflow).
-        enclosing = (
-            self._get_span(info.parent_id)
-            or self._invocation_span
-            or self._workflow_span
-        )
-        if enclosing is not None:
-            otel_context.attach(
-                trace.set_span_in_context(enclosing, self._extracted_context)
-            )
+        # Restore the enclosing context by releasing the scope this user
+        # function attached, so the parent operation (or, at the top level, the
+        # context that was active before the operation) becomes current again
+        # without stacking another scope.
+        self._detach_context(key)
 
     # ------------------------------------------------------------------
     # Attributes

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py

@@ -123,6 +123,11 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
         self._workflow_span: Span | None = None
         # Maps operation ID (None for root) to the active span.
         self._operation_spans: dict[str | None, Span] = {}
+        # Tokens returned by context.attach(), keyed by the span registry key,
+        # paired with the thread that attached them. Every attach the plugin
+        # owns is released through _detach_context so the plugin never leaves a
+        # scope on the context stack.
+        self._context_tokens: dict[str, tuple[int, object]] = {}
         self._operation_spans_lock = threading.RLock()
         self._tracing_enabled = False
 
@@ -167,19 +172,62 @@ def _attempt_span_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
         """Return the registry key for a STEP attempt span."""
         return f"{info.operation_id}:attempt:{info.attempt or 1}"
 
+    # ------------------------------------------------------------------
+    # Context scope helpers
+    # ------------------------------------------------------------------
+    def _attach_context(self, key: str, new_context: Context) -> None:
+        """Attach a context and remember its token under ``key``."""
+        with self._operation_spans_lock:
+            self._context_tokens[key] = (
+                threading.get_ident(),
+                context.attach(new_context),
+            )
+
+    def _detach_context(self, key: str) -> None:
+        """Detach the context attached under ``key``, restoring its predecessor.
+
+        A context token can only be reset on the thread that created it, so a
+        token recorded on another thread is dropped instead of detached (OTel
+        logs an error for a cross-thread reset). In practice the pairs always
+        line up: user-function hooks run on the thread executing user code, and
+        both the start and end hook for one attempt run on that same thread.
+        """
+        with self._operation_spans_lock:
+            entry = self._context_tokens.pop(key, None)
+        if entry is None:
+            return
+        thread_ident, token = entry
+        if thread_ident == threading.get_ident():
+            context.detach(token)  # type: ignore[arg-type]
+
+    def _detach_remaining_contexts(self) -> None:
+        """Release scopes still open, newest first, so nothing outlives the plugin.
+
+        Reached when a lifecycle end hook never fires -- for example a user
+        function that suspends, or a warm invocation that starts before the
+        previous one was cleaned up.
+        """
+        with self._operation_spans_lock:
+            keys = list(reversed(self._context_tokens))
+        for key in keys:
+            self._detach_context(key)
+
     def get_current_span_context(self) -> SpanContext | None:
         """Return the span context to use for log correlation.
 
         Resolution order:
         1. The span attached to the OTel thread-local context. Inside a step
            this is the active attempt span, and inside a child context this is
            the active context span (attached in
-           on_user_function_start), and between operations it is the enclosing
-           operation span (restored in on_user_function_end).
+           on_user_function_start), and between the steps of a child context it
+           is the enclosing context span, restored when on_user_function_end
+           detaches the inner scope.
         2. The invocation span from the plugin registry. This is the path used
            for top-level handler code: the invocation span is never attached to
            the worker thread's context, so the registry is the only way to
-           resolve it.
+           resolve it. It also covers code between top-level operations, where
+           detaching the operation scope restores a context with no durable
+           span.
 
         Returns:
             A valid SpanContext, or None if no span is active.
@@ -445,6 +493,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
 
     def _reset_state(self) -> None:
         """Clear per-invocation state for warm Lambda environment reuse."""
+        self._detach_remaining_contexts()
         self._execution_arn = ""
         self._execution_trace_id = None
         self._extracted_context = None
@@ -557,7 +606,9 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
             span_key=span_key,
             deterministic_span_id=info.operation_type is not OperationType.STEP,
         )
-        context.attach(trace.set_span_in_context(span, self._extracted_context))
+        self._attach_context(
+            span_key, trace.set_span_in_context(span, self._extracted_context)
+        )
 
     def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
         """Called when a context or step operation finishes user code.
@@ -607,16 +658,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
             if end_timestamp is not None and end_timestamp == info.start_time:
                 end_timestamp += datetime.timedelta(microseconds=1)
             self._end_span(span_key, end_timestamp)
-        # Restore the enclosing operation span as current so code that runs
-        # after this operation (e.g. between steps in a child context)
-        # correlates to its enclosing operation, not the operation that just
-        # ended. For a top-level operation (parent_id is None) this is the
-        # invocation span; for a nested operation it is the parent context span.
-        parent_span = self._get_span(info.parent_id) or self._get_span(None)
-        if parent_span:
-            context.attach(
-                trace.set_span_in_context(parent_span, self._extracted_context)
-            )
+        # Restore the enclosing context by releasing the scope this user
+        # function attached. Code that runs after this operation (e.g. between
+        # steps in a child context) correlates to its enclosing operation
+        # again -- the parent context span for a nested operation, and the
+        # context active before the operation for a top-level one, where
+        # get_current_span_context falls back to the invocation span.
+        self._detach_context(span_key)
 
     def _extract_attributes(self, info: Any) -> _SpanAttributes:
         """Extract durable execution fields as OpenTelemetry span attributes.