fix(otel): balance context attach and detach

Both plugins called opentelemetry.context.attach() without keeping
the returned token, and "restored" the enclosing span by attaching
another context rather than detaching. Every operation pushed two
context layers and popped none, leaving an ended span current after
its scope had finished.

The worst consequence: the invocation-start attach runs on the
Lambda handler thread, which is reused across warm invocations, so
the next execution's context extractor and GLOBAL-mode
ambient-parent lookup adopted the previous execution's ended
Workflow span. A valid parent overrides the deterministic ID
generator's trace ID, so two unrelated durable executions merged
into a single trace.

Each attach now keeps its token, keyed the same way as the span
registry, and the hook that pairs with it pops the scope:

- on_invocation_start attaches; on_invocation_end detaches.
- on_user_function_start attaches; on_user_function_end detaches,
  replacing the re-attach that stood in for restoring the
  enclosing context.
- Anything still held when the invocation ends is swept, since the
  SDK re-raises SuspendExecution without calling
  on_user_function_end.

A scope is only detached while it is still the current one. That
mirrors OpenTelemetry Java's ScopeImpl.close(), which ignores a
close that does not represent the current context, and it matters
more in Python: ContextVar.reset writes back its captured value
unconditionally, so an out-of-order or wrong-thread detach would
revive a stale context instead of failing safe. A skipped detach
keeps its entry so the owning thread can still undo it.

Tests no longer reset the OTel context to isolate themselves; an
autouse fixture asserts instead that every test leaves the context
exactly as it found it, and the tests that drove hooks without
completing the lifecycle now complete it. Adds coverage for warm
invocation reuse keeping two executions in separate traces, a
suspended operation's scope being swept, sequential steps not
accumulating layers, exact restore on success and failure, nested
child contexts, worker-thread confinement, and the identity guard
skipping both out-of-order and cross-thread detaches.

Author: Alex Wang

packages/aws-durable-execution-sdk-python-otel/README.md

@@ -235,6 +235,12 @@ context onto every emitted log record using these attributes:
 These attributes are only set when a valid span context is active, so any log
 formatter or schema must treat the fields as optional.
 
+Between two operations the plugin holds no span current: the scope a step or child
+context attached is detached when that function returns. Log correlation is
+unaffected -- the filter resolves the trace context from the plugin's own span
+registry, so records emitted between operations still carry the invocation's
+`traceId` and `spanId`.
+
 ## Verification
 
 After deploying your function with the plugin configured:

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

@@ -23,6 +23,7 @@
 import datetime
 import logging
 import threading
+from contextvars import Token
 from typing import Any
 
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -142,11 +143,70 @@ 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] = {}
+        # Contexts this plugin has attached, keyed the same way as the span
+        # registry, so each attach can be undone by the hook that pairs with it.
+        self._scopes: dict[str, tuple[Token[Context], Context]] = {}
         self._lock = threading.RLock()
 
         if self._config.enrich_logger:
             install_log_filter(self)
 
+    # ------------------------------------------------------------------
+    # Context scopes
+    # ------------------------------------------------------------------
+    def _enter_scope(self, key: str, context: Context) -> None:
+        """Attach ``context`` and remember what is needed to restore it.
+
+        ``otel_context.attach`` returns a token that is the only way to undo it,
+        and the hook that attaches is not the hook that pops, so the token has to
+        be kept. The context is kept alongside it for the identity check in
+        :meth:`_exit_scope`.
+        """
+        with self._lock:
+            self._scopes[key] = (otel_context.attach(context), context)
+
+    def _exit_scope(self, key: str) -> None:
+        """Restore the context that preceded the scope attached under ``key``.
+
+        Only detaches when the scope being popped is still the current one. This
+        mirrors OpenTelemetry Java's ``ScopeImpl.close()``, which ignores a close
+        that does not represent the current context, and it matters more here:
+        ``ContextVar.reset`` writes back its captured value unconditionally, so an
+        out-of-order or wrong-thread detach would *revive* a stale context instead
+        of failing safe. Skipping leaves the layer attached, which is inert.
+        """
+        with self._lock:
+            entry = self._scopes.get(key)
+            if entry is None:
+                return
+            token, context = entry
+            if otel_context.get_current() is not context:
+                # Not ours to pop right now: another scope is stacked above it, or
+                # this is not the thread that attached it. The entry is left in
+                # place so it can still be undone later -- discarding the token
+                # here would strand the context permanently.
+                logger.debug("Skipping out-of-scope OTel context detach for %s", key)
+                return
+            del self._scopes[key]
+        try:
+            otel_context.detach(token)
+        except Exception:  # noqa: BLE001 - observability must not break execution
+            logger.debug("Failed to detach OTel context for %s", key, exc_info=True)
+
+    def _exit_all_scopes(self) -> None:
+        """Pop every scope this plugin still holds, newest first.
+
+        Reached when a hook that would have popped a scope never ran: the SDK
+        re-raises ``SuspendExecution`` without calling ``on_user_function_end``.
+        Scopes attached on another thread fail the identity check and are dropped
+        without detaching; those threads are per-invocation and their context dies
+        with them.
+        """
+        with self._lock:
+            keys = list(reversed(self._scopes))
+        for key in keys:
+            self._exit_scope(key)
+
     # ------------------------------------------------------------------
     # Span registry helpers
     # ------------------------------------------------------------------
@@ -168,6 +228,17 @@ 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}"
 
+    @classmethod
+    def _scope_key(cls, info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
+        """Return the context-scope key for a user-function hook pair.
+
+        Mirrors the span registry key so the scope attached by
+        ``on_user_function_start`` is the one ``on_user_function_end`` pops.
+        """
+        if info.operation_type is OperationType.STEP:
+            return cls._attempt_key(info)
+        return info.operation_id
+
     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()
@@ -214,10 +285,16 @@ 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. Paired with the
+        # _exit_scope in on_invocation_end: this thread is the Lambda handler
+        # thread, which is reused across warm invocations, so leaving it attached
+        # let the next execution's context extractor and ambient-parent lookup
+        # adopt this execution's ended Workflow span -- merging two unrelated
+        # executions into one trace.
         if self._workflow_span is not None:
-            otel_context.attach(
-                trace.set_span_in_context(self._workflow_span, self._extracted_context)
+            self._enter_scope(
+                _INVOCATION_KEY,
+                trace.set_span_in_context(self._workflow_span, self._extracted_context),
             )
 
     def _start_workflow_span(self, info: InvocationStartInfo) -> None:
@@ -329,6 +406,9 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
                 logger.exception("force_flush failed at invocation end")
 
     def _reset_state(self) -> None:
+        # Undo the invocation scope, and anything a suspended operation left
+        # behind, so the handler thread is returned to the state it was found in.
+        self._exit_all_scopes()
         self._execution_arn = ""
         self._extracted_context = None
         self._workflow_span = None
@@ -455,14 +535,22 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
                 parent=parent,
                 start_time=info.start_time,
             )
-        otel_context.attach(trace.set_span_in_context(span, self._extracted_context))
+        self._enter_scope(
+            self._scope_key(info),
+            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)
         if info.operation_type not in (OperationType.CONTEXT, OperationType.STEP):
             raise RuntimeError(
                 "on_user_function_end only supports CONTEXT and STEP operations"
             )
+        # Pop the scope this operation attached, restoring exactly what preceded
+        # it. Detaching rather than attaching the enclosing span again is what
+        # keeps this balanced: the previous code pushed a second context here, so
+        # every operation added a layer and removed none.
+        self._exit_scope(self._scope_key(info))
         key = (
             self._attempt_key(info)
             if info.operation_type is OperationType.STEP
@@ -496,17 +584,6 @@ 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)
-            )
-
     # ------------------------------------------------------------------
     # Attributes
     # ------------------------------------------------------------------

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

@@ -5,6 +5,7 @@
 import datetime
 import logging
 import threading
+from contextvars import Token
 from typing import Any
 
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -167,6 +168,9 @@ 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] = {}
+        # Contexts this plugin has attached, keyed the same way as the span
+        # registry, so each attach can be undone by the hook that pairs with it.
+        self._scopes: dict[str, tuple[Token[Context], Context]] = {}
         self._operation_spans_lock = threading.RLock()
 
         if self._enrich_logger:
@@ -176,6 +180,59 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
             # plugin is constructed), so the handlers are available here.
             install_log_filter(self)
 
+    def _enter_scope(self, key: str, context_to_attach: Context) -> None:
+        """Attach a context and remember what is needed to restore it.
+
+        ``context.attach`` returns a token that is the only way to undo it, and
+        the hook that attaches is not the hook that pops, so the token has to be
+        kept. The context is kept alongside it for the identity check in
+        :meth:`_exit_scope`.
+        """
+        with self._operation_spans_lock:
+            self._scopes[key] = (context.attach(context_to_attach), context_to_attach)
+
+    def _exit_scope(self, key: str) -> None:
+        """Restore the context that preceded the scope attached under ``key``.
+
+        Only detaches when the scope being popped is still the current one. This
+        mirrors OpenTelemetry Java's ``ScopeImpl.close()``, which ignores a close
+        that does not represent the current context, and it matters more here:
+        ``ContextVar.reset`` writes back its captured value unconditionally, so an
+        out-of-order or wrong-thread detach would *revive* a stale context instead
+        of failing safe. Skipping leaves the layer attached, which is inert.
+        """
+        with self._operation_spans_lock:
+            entry = self._scopes.get(key)
+            if entry is None:
+                return
+            token, attached = entry
+            if context.get_current() is not attached:
+                # Not ours to pop right now: another scope is stacked above it, or
+                # this is not the thread that attached it. The entry is left in
+                # place so it can still be undone later -- discarding the token
+                # here would strand the context permanently.
+                logger.debug("Skipping out-of-scope OTel context detach for %s", key)
+                return
+            del self._scopes[key]
+        try:
+            context.detach(token)
+        except Exception:  # noqa: BLE001 - observability must not break execution
+            logger.debug("Failed to detach OTel context for %s", key, exc_info=True)
+
+    def _exit_all_scopes(self) -> None:
+        """Pop every scope this plugin still holds, newest first.
+
+        Reached when a hook that would have popped a scope never ran: the SDK
+        re-raises ``SuspendExecution`` without calling ``on_user_function_end``.
+        Scopes attached on another thread fail the identity check and are dropped
+        without detaching; those threads are per-invocation and their context dies
+        with them.
+        """
+        with self._operation_spans_lock:
+            keys = list(reversed(self._scopes))
+        for key in keys:
+            self._exit_scope(key)
+
     def _set_span(self, operation_id: str | None, span: Span) -> None:
         """Register the active span for an operation ID."""
         with self._operation_spans_lock:
@@ -196,6 +253,17 @@ 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}"
 
+    @classmethod
+    def _scope_key(cls, info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
+        """Return the context-scope key for a user-function hook pair.
+
+        Mirrors the span registry key so the scope attached by
+        ``on_user_function_start`` is the one ``on_user_function_end`` pops.
+        """
+        if info.operation_type is OperationType.STEP:
+            return cls._attempt_span_key(info)
+        return info.operation_id
+
     def get_current_span_context(self) -> SpanContext | None:
         """Return the span context to use for log correlation.
 
@@ -453,6 +521,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
                     self._workflow_span.set_status(StatusCode.OK)
                 self._workflow_span.end()
 
+        # Undo anything a suspended operation left attached, so no scope outlives
+        # the invocation that created it.
+        self._exit_all_scopes()
+
         # Clear all per-invocation state to prevent leaks across warm Lambda reuses
         self._execution_arn = ""
         self._extracted_context = None
@@ -562,7 +634,10 @@ 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._enter_scope(
+            self._scope_key(info),
+            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.
@@ -578,6 +653,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
             raise RuntimeError(
                 "on_user_function_end should only be called for CONTEXT and STEP operations"
             )
+        # Pop the scope this operation attached, restoring exactly what preceded
+        # it. Detaching rather than attaching the enclosing span again is what
+        # keeps this balanced: the previous code pushed a second context here, so
+        # every operation added a layer and removed none. Between operations the
+        # log filter resolves through the span registry (see
+        # get_current_span_context), so correlation is unaffected.
+        self._exit_scope(self._scope_key(info))
         # key = f"{info.operation_id}-{int(info.start_time.timestamp())}"
         span_key = (
             self._attempt_span_key(info)
@@ -610,16 +692,6 @@ 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)
-            )
 
     def _extract_attributes(self, info: Any) -> _SpanAttributes:
         """Extract durable execution fields as OpenTelemetry span attributes.