fix(otel): balance context attach and detach across plugin lifecycles

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

The worst consequence was cross-execution trace pollution. The invocation-start
attach ran on the Lambda handler thread, which is reused across warm
invocations, so the next execution's context extractor and ambient-parent
lookup adopted the previous execution's ended Workflow span -- merging two
unrelated executions into a single trace.

Changes:

- Add context_scope, a thread-confined LIFO stack of attach tokens. Detaches
  unwind downwards so the underlying ContextVar is always reset in order:
  unlike Java's Scope.close(), ContextVar.reset() writes back its captured
  value unconditionally and would otherwise revive a stale context. The stack
  is module level so the two plugins, which ship as separate entry points and
  can be enabled together, still unwind in true LIFO order. An epoch check
  discards scopes a suspended operation left behind, since the SDK re-raises
  SuspendExecution without calling on_user_function_end.
- Pair every user-function attach with a detach on the same thread, replacing
  the re-attach that previously stood in for restoring the enclosing context.
- Unwind any remaining scopes at invocation end.
- Drop the invocation-start attach entirely. User code runs on a separate
  worker and ThreadPoolExecutor does not copy contextvars, so that attach
  never reached the code it was meant to parent; it only leaked. The Workflow
  and Invocation spans are used as explicit parents instead, matching the Java
  plugins, which never make either span current.

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. Adds coverage for nested contexts, sequential steps, failures, suspension,
worker-thread hooks, both plugins on one thread, and warm invocation reuse
keeping two executions in separate traces.

Ambient spans emitted outside any operation are no longer parented to the
Invocation span; the README documents this, and log correlation is unchanged
because the logging filter resolves through the plugin's span registry.

Author: Alex Wang

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

@@ -235,6 +235,23 @@ 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.
 
+### Active Span Scope
+
+The plugin makes a span current only while your step or child-context function
+runs, and detaches it when that function returns. Two consequences are worth
+knowing:
+
+- Auto-instrumented calls (botocore, urllib3, and similar) made **inside** a step
+  or child context become children of that operation's span.
+- Auto-instrumented calls made **outside** any operation -- for example directly
+  in the handler between two steps -- are not parented to the durable spans. In
+  an ADOT deployment they attach to the ambient Lambda invocation span instead.
+  Put such work in a step if you need it inside the durable trace.
+
+Log correlation is unaffected either way: the logging 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/context_scope.py

@@ -0,0 +1,179 @@
+"""Balanced ``opentelemetry.context`` attach/detach bookkeeping for the plugins.
+
+The OpenTelemetry Context specification requires every ``context.attach()`` to
+have a corresponding ``context.detach(token)``. Detaching is only possible with
+the token that ``attach`` returned, so the token has to survive from the hook
+that attached to the hook that pops it -- the plugin hooks are separate calls,
+so the idiomatic ``with tracer.start_as_current_span(...)`` form is unavailable.
+
+Two properties of the runtime shape the design:
+
+* **Tokens are thread-confined.** The plugin hooks run on several threads: the
+  invocation hooks on the Lambda handler thread, the user-function hooks on the
+  ``dex-handler`` worker that runs user code, and on a branch worker for each
+  ``map``/``parallel`` branch. ``ContextVar.reset()`` only accepts a token
+  created in the same ``contextvars.Context``, so each thread keeps its own
+  stack and only ever detaches its own tokens.
+* **Detach order matters.** Unlike OpenTelemetry Java's ``Scope.close()`` --
+  which ignores a close that does not represent the current context --
+  ``ContextVar.reset()`` unconditionally writes back the token's captured value.
+  Detaching out of order therefore *revives* a stale context instead of failing
+  safe. The stack is module level rather than per plugin instance so that two
+  plugins attaching on the same thread (both ship as separate entry points and
+  can be enabled together) still unwind in true LIFO order.
+
+Scopes are keyed by ``(owner, key)`` so a plugin instance can pop the exact
+scope it pushed, while :func:`exit_scope` still unwinds anything stacked above
+it. Nothing here raises: a plugin must never break an execution over
+observability bookkeeping.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from contextvars import Token
+
+    from opentelemetry.context import Context
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class _Entry:
+    """One attached scope: who pushed it, under what key, and its token."""
+
+    owner_id: int
+    key: str
+    epoch: int
+    token: Token[Context]
+
+
+class _ThreadState(threading.local):
+    """Per-thread LIFO stack of attached scopes."""
+
+    def __init__(self) -> None:
+        self.entries: list[_Entry] = []
+
+
+_state = _ThreadState()
+
+
+def _detach(entry: _Entry) -> None:
+    """Detach one entry, swallowing any failure."""
+    from opentelemetry import context as otel_context
+
+    try:
+        otel_context.detach(entry.token)
+    except Exception:  # noqa: BLE001 - observability must not break execution
+        logger.debug("Failed to detach OTel context scope %s", entry.key, exc_info=True)
+
+
+def enter_scope(owner: Any, key: str, context: Context, epoch: int = 0) -> None:
+    """Attach ``context`` on this thread and remember how to restore it.
+
+    Any scope still on this thread's stack from an earlier ``epoch`` is unwound
+    first. That covers the paths where a paired pop never runs: the SDK re-raises
+    ``SuspendExecution`` without calling ``on_user_function_end``, so a suspended
+    operation leaves its scope attached, and a reused thread would otherwise
+    inherit it.
+
+    Args:
+        owner: The plugin instance pushing the scope.
+        key: Registry key for the scope, unique per owner (operation or attempt).
+        context: The context to attach.
+        epoch: The owner's invocation counter; scopes from older epochs are
+            discarded before the new scope is pushed.
+    """
+    from opentelemetry import context as otel_context
+
+    owner_id = id(owner)
+    _discard_stale(owner_id, epoch)
+    try:
+        token = otel_context.attach(context)
+    except Exception:  # noqa: BLE001
+        logger.debug("Failed to attach OTel context scope %s", key, exc_info=True)
+        return
+    _state.entries.append(_Entry(owner_id=owner_id, key=key, epoch=epoch, token=token))
+
+
+def exit_scope(owner: Any, key: str) -> None:
+    """Detach the scope ``owner`` pushed under ``key``, restoring what preceded it.
+
+    Scopes stacked above the target are detached first so the underlying
+    ``ContextVar`` is always reset in LIFO order. A key this thread never pushed
+    is a no-op -- the scope belongs to another thread (or was already unwound),
+    and detaching someone else's token would corrupt the context.
+    """
+    owner_id = id(owner)
+    index = _find_last(owner_id, key)
+    if index is None:
+        return
+    for entry in reversed(_state.entries[index:]):
+        _detach(entry)
+    del _state.entries[index:]
+
+
+def unwind(owner: Any) -> None:
+    """Detach every scope ``owner`` still holds on this thread, newest first.
+
+    Called at invocation end so the handler thread is left exactly as the plugin
+    found it. Scopes this owner pushed on *other* threads cannot be detached from
+    here; those threads are created per invocation and their context dies with
+    them.
+    """
+    owner_id = id(owner)
+    index = _find_first(owner_id)
+    if index is None:
+        return
+    for entry in reversed(_state.entries[index:]):
+        _detach(entry)
+    del _state.entries[index:]
+
+
+def depth(owner: Any | None = None) -> int:
+    """Return the number of scopes attached on this thread (for tests)."""
+    if owner is None:
+        return len(_state.entries)
+    owner_id = id(owner)
+    return sum(1 for entry in _state.entries if entry.owner_id == owner_id)
+
+
+def _discard_stale(owner_id: int, epoch: int) -> None:
+    """Unwind this owner's scopes left over from a previous epoch."""
+    index = next(
+        (
+            position
+            for position, entry in enumerate(_state.entries)
+            if entry.owner_id == owner_id and entry.epoch != epoch
+        ),
+        None,
+    )
+    if index is None:
+        return
+    for entry in reversed(_state.entries[index:]):
+        _detach(entry)
+    del _state.entries[index:]
+
+
+def _find_last(owner_id: int, key: str) -> int | None:
+    """Index of this owner's most recent scope for ``key``, if any."""
+    for position in range(len(_state.entries) - 1, -1, -1):
+        entry = _state.entries[position]
+        if entry.owner_id == owner_id and entry.key == key:
+            return position
+    return None
+
+
+def _find_first(owner_id: int) -> int | None:
+    """Index of this owner's oldest scope, if any."""
+    for position, entry in enumerate(_state.entries):
+        if entry.owner_id == owner_id:
+            return position
+    return None

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

@@ -51,6 +51,7 @@
     Tracer,
 )
 
+from aws_durable_execution_sdk_python_otel import context_scope
 from aws_durable_execution_sdk_python_otel.context_extractors import (
     ContextExtractor,
     xray_context_extractor,
@@ -143,6 +144,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
         self._invocation_span: Span | None = None
         self._operation_spans: dict[str, Span] = {}
         self._lock = threading.RLock()
+        # Bumped every invocation. context_scope uses it to discard scopes a
+        # previous invocation left attached on a reused thread.
+        self._epoch = 0
 
         if self._config.enrich_logger:
             install_log_filter(self)
@@ -168,6 +172,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 pushed 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()
@@ -204,6 +219,7 @@ def _resolve_parent(self, parent_id: str | None) -> Span | None:
     # ------------------------------------------------------------------
     def on_invocation_start(self, info: InvocationStartInfo) -> None:
         logger.debug("Durable invocation started: %s", info)
+        self._epoch += 1
         self._execution_arn = info.execution_arn or ""
         self._extracted_context = self._context_extractor(info)
         self._id_generator.set_trace_id(self._execution_arn, info.execution_start_time)
@@ -213,12 +229,16 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
         # is parented to the ambient Lambda invocation span.
         self._start_invocation_span(info)
 
-        # Make the Workflow span the active span so auto-instrumented spans
-        # created during the invocation become its children.
-        if self._workflow_span is not None:
-            otel_context.attach(
-                trace.set_span_in_context(self._workflow_span, self._extracted_context)
-            )
+        # No context is attached here. Nothing on this thread needs it: user code
+        # runs on a separate worker (ThreadPoolExecutor does not copy
+        # contextvars), so an attach here would never reach it, while the Lambda
+        # handler thread is reused across warm invocations -- an unpaired attach
+        # would leak an ended span into the next execution, whose context
+        # extractor and ambient-parent lookup would then adopt it and merge two
+        # executions into one trace. The Workflow and Invocation spans are used
+        # as explicit parents instead (see _resolve_parent), matching the Java
+        # plugins, which never make either span current. Log correlation for this
+        # thread resolves through get_current_span_context().
 
     def _start_workflow_span(self, info: InvocationStartInfo) -> None:
         if not self._execution_arn:
@@ -329,6 +349,13 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
                 logger.exception("force_flush failed at invocation end")
 
     def _reset_state(self) -> None:
+        # Detach anything this plugin still holds on this thread so the handler
+        # thread is left exactly as it was found. Scopes attached on the
+        # per-invocation worker threads cannot be detached from here (a token is
+        # only resettable in the context that created it); those threads are
+        # destroyed with the invocation, and any scope a suspended operation left
+        # behind is discarded by the epoch check on the next enter_scope.
+        context_scope.unwind(self)
         self._execution_arn = ""
         self._extracted_context = None
         self._workflow_span = None
@@ -455,14 +482,30 @@ 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))
+        # Attach on this worker thread so auto-instrumented calls made by the
+        # user function become children of this span. The scope is pushed onto
+        # whatever is already current (rather than replacing it with
+        # _extracted_context) so an ambient context on this thread survives; the
+        # span's own parent was chosen explicitly in _start_span.
+        context_scope.enter_scope(
+            self,
+            self._scope_key(info),
+            trace.set_span_in_context(span, otel_context.get_current()),
+            epoch=self._epoch,
+        )
 
     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"
             )
+        # Detach first, on the same thread that attached, so the context this
+        # operation was entered from is restored exactly. Detaching (rather than
+        # attaching the enclosing span again) is what keeps the scopes balanced:
+        # a nested operation lands back on its parent's still-attached scope, and
+        # a top-level one lands back on the thread's ambient context.
+        context_scope.exit_scope(self, self._scope_key(info))
         key = (
             self._attempt_key(info)
             if info.operation_type is OperationType.STEP
@@ -496,17 +539,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
     # ------------------------------------------------------------------