From 6b72aa799ed248420a85a919e96f5aa01bfe0db9 Mon Sep 17 00:00:00 2001
From: Alex Wang <wangyb@amazon.com>
Date: Fri, 14 Aug 2026 19:42:22 +0000
Subject: [PATCH 1/4] fix(otel): balance context attach and detach

Both OTel 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 is not in the issue. The invocation-start
attach ran 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.

Two runtime properties shaped the fix. The hooks run on several
threads -- the invocation hooks on the Lambda handler thread, the
user-function hooks on the worker that runs user code, and on a
branch worker per map/parallel branch -- and ContextVar.reset()
only accepts a token created in the same contextvars.Context.
And unlike Java's ScopeImpl.close(), which ignores a close that
does not represent the current context, ContextVar.reset() writes
back its captured value unconditionally, so detaching out of
order revives a stale context instead of failing safe.

Changes:

- Add context_scope, a thread-confined LIFO stack of attach
  tokens. Detaches unwind downwards so the ContextVar is always
  reset in order. 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. 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.

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.

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.

Fixes #643
---
 .../README.md                                 |  17 +
 .../context_scope.py                          | 179 +++++++++
 .../execution_plugin.py                       |  68 +++-
 .../invocation_plugin.py                      |  65 ++-
 .../tests/test_context_scope.py               | 372 ++++++++++++++++++
 .../tests/test_execution_plugin.py            |  95 ++++-
 .../test_execution_plugin_integration.py      |  18 +-
 .../tests/test_invocation_plugin.py           | 104 +++--
 .../test_invocation_plugin_integration.py     |  18 +-
 .../tests/test_log_filter.py                  |  90 +++++
 10 files changed, 942 insertions(+), 84 deletions(-)
 create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
 create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py

diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md
index 4eb563aa..12201792 100644
--- a/packages/aws-durable-execution-sdk-python-otel/README.md
+++ b/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:
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
new file mode 100644
index 00000000..54c46020
--- /dev/null
+++ b/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
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
index bc829db6..2d84b104 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
+++ b/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,7 +482,17 @@ 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)
@@ -463,6 +500,12 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
             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
     # ------------------------------------------------------------------
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
index 4f42ca32..60988469 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
@@ -34,6 +34,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,
@@ -168,6 +169,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
         # Maps operation ID (None for root) to the active span.
         self._operation_spans: dict[str | None, Span] = {}
         self._operation_spans_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._enrich_logger:
             # Install the root-logger filter so every log record is stamped with
@@ -196,19 +200,32 @@ 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 pushed 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.
 
         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).
+           the active context span (attached in on_user_function_start). After a
+           nested operation ends, its scope is detached and the enclosing child
+           context span -- still attached -- becomes current again.
         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.
+           any thread's context, so the registry is the only way to resolve it.
+           It also covers code that runs between top-level operations, where
+           detaching the operation scope leaves the thread's ambient context
+           current.
 
         Returns:
             A valid SpanContext, or None if no span is active.
@@ -364,6 +381,7 @@ def _end_span(
     def on_invocation_start(self, info: InvocationStartInfo) -> None:
         """Called at the start of each invocation. Creates the invocation span."""
         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)
@@ -454,6 +472,13 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
                 self._workflow_span.end()
 
         # Clear all per-invocation state to prevent leaks across warm Lambda reuses
+        # 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
@@ -562,7 +587,17 @@ 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))
+        # 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, context.get_current()),
+            epoch=self._epoch,
+        )
 
     def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
         """Called when a context or step operation finishes user code.
@@ -578,6 +613,14 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
             raise RuntimeError(
                 "on_user_function_end should only be called for 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:
+        # code after a nested operation lands back on its parent context span,
+        # still attached here, and a top-level operation lands back on the
+        # thread's ambient context. Between-operation log records resolve through
+        # get_current_span_context(), which falls back to the invocation span.
+        context_scope.exit_scope(self, self._scope_key(info))
         # key = f"{info.operation_id}-{int(info.start_time.timestamp())}"
         span_key = (
             self._attempt_span_key(info)
@@ -610,16 +653,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.
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
new file mode 100644
index 00000000..abef60ba
--- /dev/null
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
@@ -0,0 +1,372 @@
+"""Tests for balanced OTel context scopes across the plugin lifecycles.
+
+Covers :mod:`context_scope` directly plus the lifecycle paths the plugins rely
+on it for: a suspended operation whose end hook never runs, hooks that execute
+on a worker thread, and both plugins attaching on the same thread.
+"""
+
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+from datetime import UTC, datetime
+
+import opentelemetry.context as otel_context
+import pytest
+from aws_durable_execution_sdk_python.lambda_service import (
+    InvocationStatus,
+    OperationStatus,
+    OperationSubType,
+    OperationType,
+)
+from aws_durable_execution_sdk_python.plugin import (
+    InvocationEndInfo,
+    InvocationStartInfo,
+    UserFunctionEndInfo,
+    UserFunctionOutcome,
+    UserFunctionStartInfo,
+)
+from opentelemetry import trace
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from aws_durable_execution_sdk_python_otel import context_scope
+from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin
+from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin
+from aws_durable_execution_sdk_python_otel.otel_plugin_config import (
+    OtelPluginConfig,
+    ProviderSource,
+)
+
+
+START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC)
+END_TIME = datetime(2024, 1, 2, 3, 4, 6, tzinfo=UTC)
+EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST"
+
+
+@pytest.fixture(autouse=True)
+def _assert_otel_context_balanced():
+    """Fail any test that leaves an OTel context scope attached."""
+    before = otel_context.get_current()
+    before_depth = context_scope.depth()
+    yield
+    assert context_scope.depth() == before_depth
+    assert otel_context.get_current() is before
+
+
+class _Owner:
+    """Stand-in for a plugin instance; context_scope only uses its identity."""
+
+
+def _span_context(name: str) -> Context:
+    """Return a context carrying a non-recording span named after ``name``."""
+    return trace.set_span_in_context(
+        trace.NonRecordingSpan(
+            trace.SpanContext(
+                trace_id=abs(hash(name)) % (1 << 128) or 1,
+                span_id=abs(hash(name)) % (1 << 64) or 1,
+                is_remote=False,
+                trace_flags=trace.TraceFlags(trace.TraceFlags.SAMPLED),
+            )
+        ),
+        otel_context.get_current(),
+    )
+
+
+# ---------------------------------------------------------------------------
+# context_scope helper
+# ---------------------------------------------------------------------------
+def test_enter_then_exit_restores_the_previous_context():
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    context_scope.enter_scope(owner, "a", _span_context("a"))
+    assert otel_context.get_current() is not before
+    assert context_scope.depth(owner) == 1
+
+    context_scope.exit_scope(owner, "a")
+    assert otel_context.get_current() is before
+    assert context_scope.depth(owner) == 0
+
+
+def test_nested_scopes_pop_in_lifo_order():
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    context_scope.enter_scope(owner, "outer", _span_context("outer"))
+    outer = otel_context.get_current()
+    context_scope.enter_scope(owner, "inner", _span_context("inner"))
+
+    context_scope.exit_scope(owner, "inner")
+    assert otel_context.get_current() is outer
+    context_scope.exit_scope(owner, "outer")
+    assert otel_context.get_current() is before
+
+
+def test_exit_unwinds_scopes_stacked_above_the_target():
+    """Popping an outer scope must also drop anything left above it.
+
+    ``ContextVar.reset`` writes back the token's captured value unconditionally,
+    so resetting out of order would revive a stale context. Unwinding downwards
+    keeps the underlying variable strictly LIFO.
+    """
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    context_scope.enter_scope(owner, "outer", _span_context("outer"))
+    context_scope.enter_scope(owner, "orphan", _span_context("orphan"))
+
+    context_scope.exit_scope(owner, "outer")
+
+    assert otel_context.get_current() is before
+    assert context_scope.depth(owner) == 0
+
+
+def test_exit_with_unknown_key_is_a_noop():
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    context_scope.exit_scope(owner, "never-pushed")
+
+    assert otel_context.get_current() is before
+    assert context_scope.depth(owner) == 0
+
+
+def test_enter_discards_scopes_from_a_previous_epoch():
+    """A scope a suspended operation left behind must not outlive its invocation.
+
+    The SDK re-raises ``SuspendExecution`` without calling
+    ``on_user_function_end``, so the next invocation on a reused thread finds a
+    stale scope; the epoch check drops it before attaching.
+    """
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    context_scope.enter_scope(owner, "suspended", _span_context("suspended"), epoch=1)
+    assert context_scope.depth(owner) == 1
+
+    context_scope.enter_scope(owner, "next", _span_context("next"), epoch=2)
+    assert context_scope.depth(owner) == 1
+
+    context_scope.exit_scope(owner, "next")
+    assert otel_context.get_current() is before
+    assert context_scope.depth(owner) == 0
+
+
+def test_unwind_detaches_every_scope_for_one_owner():
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    context_scope.enter_scope(owner, "a", _span_context("a"))
+    context_scope.enter_scope(owner, "b", _span_context("b"))
+
+    context_scope.unwind(owner)
+
+    assert otel_context.get_current() is before
+    assert context_scope.depth(owner) == 0
+
+
+def test_unwind_is_a_noop_for_an_owner_with_no_scopes():
+    before = otel_context.get_current()
+
+    context_scope.unwind(_Owner())
+
+    assert otel_context.get_current() is before
+
+
+def test_scopes_are_confined_to_the_thread_that_attached_them():
+    """A worker thread's scopes must not appear on, or be poppable from, another.
+
+    Tokens are only resettable in the ``contextvars.Context`` that created them,
+    so the stack is per thread.
+    """
+    owner = _Owner()
+    before = otel_context.get_current()
+    observed: dict[str, object] = {}
+
+    def worker() -> None:
+        context_scope.enter_scope(owner, "worker", _span_context("worker"))
+        observed["worker_depth"] = context_scope.depth(owner)
+        observed["worker_span_valid"] = (
+            trace.get_current_span().get_span_context().is_valid
+        )
+        context_scope.exit_scope(owner, "worker")
+        observed["worker_depth_after"] = context_scope.depth(owner)
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(worker).result()
+
+    assert observed["worker_depth"] == 1
+    assert observed["worker_span_valid"] is True
+    assert observed["worker_depth_after"] == 0
+    # The calling thread never saw the worker's scope.
+    assert context_scope.depth(owner) == 0
+    assert otel_context.get_current() is before
+
+
+# ---------------------------------------------------------------------------
+# Plugin lifecycle paths
+# ---------------------------------------------------------------------------
+def _execution_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]:
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    return (
+        ExecutionOtelPlugin(
+            OtelPluginConfig(
+                provider_source=ProviderSource.EXPLICIT,
+                tracer_provider=provider,
+                context_extractor=lambda _: Context(),
+                enrich_logger=False,
+            )
+        ),
+        exporter,
+    )
+
+
+def _invocation_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]:
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    return (
+        InvocationOtelPlugin(
+            OtelPluginConfig(
+                provider_source=ProviderSource.EXPLICIT,
+                tracer_provider=provider,
+                context_extractor=lambda _: Context(),
+                enrich_logger=False,
+            )
+        ),
+        exporter,
+    )
+
+
+def _invocation_start() -> InvocationStartInfo:
+    return InvocationStartInfo(
+        request_id="request-1",
+        execution_arn=EXECUTION_ARN,
+        execution_start_time=START_TIME,
+        is_first_invocation=True,
+    )
+
+
+def _invocation_end(
+    status: InvocationStatus = InvocationStatus.SUCCEEDED,
+) -> InvocationEndInfo:
+    return InvocationEndInfo(
+        request_id="request-1",
+        execution_arn=EXECUTION_ARN,
+        execution_start_time=START_TIME,
+        is_first_invocation=True,
+        status=status,
+        error=None,
+    )
+
+
+def _step_start(operation_id: str) -> UserFunctionStartInfo:
+    return UserFunctionStartInfo(
+        operation_id=operation_id,
+        operation_type=OperationType.STEP,
+        sub_type=OperationSubType.STEP,
+        name=operation_id,
+        parent_id=None,
+        start_time=START_TIME,
+        is_replayed=False,
+        status=OperationStatus.STARTED,
+        is_replay_children=False,
+        attempt=1,
+    )
+
+
+def _step_end(operation_id: str) -> UserFunctionEndInfo:
+    return UserFunctionEndInfo(
+        operation_id=operation_id,
+        operation_type=OperationType.STEP,
+        sub_type=OperationSubType.STEP,
+        name=operation_id,
+        parent_id=None,
+        start_time=START_TIME,
+        end_time=END_TIME,
+        is_replayed=False,
+        status=OperationStatus.SUCCEEDED,
+        is_replay_children=False,
+        attempt=1,
+        outcome=UserFunctionOutcome.SUCCEEDED,
+        error=None,
+    )
+
+
+@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
+def test_invocation_end_unwinds_a_suspended_operation_scope(factory):
+    """A step that suspends never gets its end hook; invocation end cleans up.
+
+    ``wrap_user_function`` re-raises ``SuspendExecution`` without calling
+    ``on_user_function_end``, so the scope is still attached when the invocation
+    winds down.
+    """
+    plugin, _ = factory()
+    before = otel_context.get_current()
+
+    plugin.on_invocation_start(_invocation_start())
+    plugin.on_user_function_start(_step_start("step-suspends"))
+    assert context_scope.depth(plugin) == 1
+
+    plugin.on_invocation_end(_invocation_end(InvocationStatus.PENDING))
+
+    assert context_scope.depth(plugin) == 0
+    assert otel_context.get_current() is before
+
+
+@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
+def test_user_function_hooks_on_a_worker_thread_leave_the_caller_alone(factory):
+    """Verify hooks running on the user-code thread do not touch the caller.
+
+    User code runs on a worker the SDK owns, and ``ThreadPoolExecutor`` does not
+    copy contextvars, so the plugin's scope must stay on that thread.
+    """
+    plugin, _ = factory()
+    plugin.on_invocation_start(_invocation_start())
+    before = otel_context.get_current()
+    observed: dict[str, object] = {}
+
+    def run_step() -> None:
+        plugin.on_user_function_start(_step_start("step-1"))
+        observed["inside"] = trace.get_current_span().get_span_context().is_valid
+        plugin.on_user_function_end(_step_end("step-1"))
+        observed["after"] = context_scope.depth(plugin)
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(run_step).result()
+
+    assert observed["inside"] is True
+    assert observed["after"] == 0
+    assert otel_context.get_current() is before
+    assert trace.get_current_span().get_span_context().is_valid is False
+
+    plugin.on_invocation_end(_invocation_end())
+
+
+def test_two_plugins_on_one_thread_unwind_in_lifo_order():
+    """Both plugins ship as entry points and can be enabled together.
+
+    Hooks dispatch in registration order, so the second plugin's scope is
+    detached while the first plugin's end hook runs. The shared stack keeps the
+    underlying ContextVar LIFO instead of reviving the first plugin's scope.
+    """
+    first, _ = _execution_plugin()
+    second, _ = _invocation_plugin()
+    before = otel_context.get_current()
+
+    for plugin in (first, second):
+        plugin.on_invocation_start(_invocation_start())
+    for plugin in (first, second):
+        plugin.on_user_function_start(_step_start("step-1"))
+    for plugin in (first, second):
+        plugin.on_user_function_end(_step_end("step-1"))
+
+    assert context_scope.depth() == 0
+    assert otel_context.get_current() is before
+
+    for plugin in (first, second):
+        plugin.on_invocation_end(_invocation_end())
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py
index 75488662..39c42c44 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py
@@ -28,6 +28,7 @@
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
+from aws_durable_execution_sdk_python_otel import context_scope
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     derive_workflow_span_id,
     operation_id_to_span_id,
@@ -45,17 +46,22 @@
 
 
 @pytest.fixture(autouse=True)
-def _reset_otel_context():
-    """Reset the OTel thread-local context around each test.
+def _assert_otel_context_balanced():
+    """Fail any test that leaves an OTel context scope attached.
 
-    The plugin attaches spans via context.attach() without detaching, so state
-    would otherwise leak between tests running on the same thread.
+    The plugins must detach every context they attach, so no reset is needed to
+    isolate tests -- instead this asserts the invariant.
     """
-    token = otel_context.attach(Context())
-    try:
-        yield
-    finally:
-        otel_context.detach(token)
+    before = otel_context.get_current()
+    before_depth = context_scope.depth()
+    yield
+    assert context_scope.depth() == before_depth, (
+        "test left OTel context scopes attached: "
+        f"{context_scope.depth() - before_depth} extra"
+    )
+    assert otel_context.get_current() is before, (
+        "test did not restore the OTel context it started with"
+    )
 
 
 def _create_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]:
@@ -446,6 +452,77 @@ def test_default_mode_invocation_span_parented_to_ambient_span(monkeypatch):
     assert invocation.parent.span_id == ambient.get_span_context().span_id
 
 
+# ----------------------------------------------------------------------
+# Warm invocation reuse: no context leaks from one execution into the next
+# ----------------------------------------------------------------------
+def test_invocation_end_restores_the_pre_invocation_context():
+    """Verify the plugin leaves the handler thread's context untouched.
+
+    The invocation lifecycle attaches nothing on this thread, so the Workflow
+    span must not be current afterwards -- an ended span left current is what
+    previously bled into the next invocation on a warm container.
+    """
+    plugin, _ = _create_plugin()
+    before = otel_context.get_current()
+
+    plugin.on_invocation_start(_invocation_start_info())
+    plugin.on_invocation_end(_invocation_end_info())
+
+    assert otel_context.get_current() is before
+    assert trace.get_current_span().get_span_context().is_valid is False
+
+
+def test_warm_reuse_does_not_share_a_trace_between_executions(monkeypatch):
+    """Verify a reused plugin instance keeps two executions in separate traces.
+
+    In GLOBAL mode the Invocation span is parented to whatever is ambient. When
+    the previous invocation left its Workflow span attached, that span became the
+    parent and its trace ID won, merging two unrelated executions into one trace.
+    """
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider)
+    # The default X-Ray extractor falls back to the ambient context when no
+    # trace header is present, so it observes any leak too.
+    monkeypatch.delenv("_X_AMZN_TRACE_ID", raising=False)
+    plugin = ExecutionOtelPlugin(
+        OtelPluginConfig(
+            provider_source=ProviderSource.GLOBAL,
+            enrich_logger=False,
+        )
+    )
+
+    def run(arn: str) -> None:
+        info = InvocationStartInfo(
+            request_id="request-1",
+            execution_arn=arn,
+            execution_start_time=START_TIME,
+            is_first_invocation=True,
+        )
+        plugin.on_invocation_start(info)
+        plugin.on_invocation_end(
+            InvocationEndInfo(
+                request_id="request-1",
+                execution_arn=arn,
+                execution_start_time=START_TIME,
+                is_first_invocation=True,
+                status=InvocationStatus.SUCCEEDED,
+                error=None,
+            )
+        )
+
+    run(EXECUTION_ARN)
+    run(EXECUTION_ARN + "-second")
+
+    invocations = [s for s in exporter.get_finished_spans() if s.name == "Invocation"]
+    assert len(invocations) == 2
+    first, second = invocations
+    assert first.context.trace_id != second.context.trace_id
+    # The second invocation is a root, not a child of the first execution's span.
+    assert second.parent is None
+
+
 def test_open_operation_span_not_exported_at_invocation_end():
     """A suspended operation (started, not ended) must not be exported.
 
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py
index 98f450c5..d760fed7 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py
@@ -37,6 +37,7 @@
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
+from aws_durable_execution_sdk_python_otel import context_scope
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     derive_workflow_span_id,
     operation_id_to_span_id,
@@ -60,12 +61,17 @@
 
 
 @pytest.fixture(autouse=True)
-def _reset_otel_context():
-    token = otel_context.attach(Context())
-    try:
-        yield
-    finally:
-        otel_context.detach(token)
+def _assert_otel_context_balanced():
+    """Fail any test that leaves an OTel context scope attached."""
+    before = otel_context.get_current()
+    before_depth = context_scope.depth()
+    yield
+    assert context_scope.depth() == before_depth, (
+        "test left OTel context scopes attached"
+    )
+    assert otel_context.get_current() is before, (
+        "test did not restore the OTel context it started with"
+    )
 
 
 def _provider() -> tuple[TracerProvider, InMemorySpanExporter]:
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py
index d6cf063a..64ebcc36 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py
@@ -31,6 +31,7 @@
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from opentelemetry.trace import SpanKind, StatusCode
 
+from aws_durable_execution_sdk_python_otel import context_scope
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     derive_workflow_span_id,
     operation_id_to_span_id,
@@ -48,17 +49,23 @@
 
 
 @pytest.fixture(autouse=True)
-def _reset_otel_context():
-    """Reset the OTel thread-local context before and after each test.
+def _assert_otel_context_balanced():
+    """Fail any test that leaves an OTel context scope attached.
 
-    The plugin attaches spans via context.attach() without ever detaching,
-    so state would otherwise leak between tests running on the same thread.
+    The plugins must detach every context they attach, so no reset is needed to
+    isolate tests -- instead this asserts the invariant. A test that drives hooks
+    the SDK itself leaves unpaired (a suspension) is responsible for unwinding.
     """
-    token = otel_context.attach(Context())
-    try:
-        yield
-    finally:
-        otel_context.detach(token)
+    before = otel_context.get_current()
+    before_depth = context_scope.depth()
+    yield
+    assert context_scope.depth() == before_depth, (
+        "test left OTel context scopes attached: "
+        f"{context_scope.depth() - before_depth} extra"
+    )
+    assert otel_context.get_current() is before, (
+        "test did not restore the OTel context it started with"
+    )
 
 
 def _create_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]:
@@ -791,13 +798,14 @@ def update_span(index: int) -> None:
 
 
 # ----------------------------------------------------------------------
-# on_user_function_end restores the invocation span to the context
+# on_user_function_end restores the context the operation was entered from
 # ----------------------------------------------------------------------
-def test_user_function_end_restores_invocation_span():
-    """Verify the invocation span is current again after a step completes."""
+def test_user_function_end_restores_enclosing_context():
+    """Verify the exact pre-step context is restored after a step completes."""
     plugin, _ = _create_plugin()
     plugin.on_invocation_start(_invocation_start_info())
-    invocation_span_id = plugin._get_span(None).get_span_context().span_id
+    before = otel_context.get_current()
+    before_depth = context_scope.depth(plugin)
 
     operation_id = "step-1"
     plugin.on_user_function_start(_user_function_start_info(operation_id))
@@ -808,18 +816,23 @@ def test_user_function_end_restores_invocation_span():
         trace.get_current_span().get_span_context().span_id
         == active_attempt_span.get_span_context().span_id
     )
+    assert context_scope.depth(plugin) == before_depth + 1
 
     plugin.on_user_function_end(_user_function_end_info(operation_id))
 
-    # After the step, the invocation span is restored.
-    assert trace.get_current_span().get_span_context().span_id == invocation_span_id
+    # The scope is detached, restoring the context byte for byte. The invocation
+    # span is never attached (matching the Java plugins), so between-step log
+    # correlation goes through get_current_span_context() instead.
+    assert otel_context.get_current() is before
+    assert context_scope.depth(plugin) == before_depth
+    plugin.on_invocation_end(_invocation_end_info())
 
 
-def test_user_function_end_restores_invocation_span_on_failure():
-    """Verify the invocation span is restored even when the step fails."""
+def test_user_function_end_restores_enclosing_context_on_failure():
+    """Verify the context is restored even when the step fails."""
     plugin, _ = _create_plugin()
     plugin.on_invocation_start(_invocation_start_info())
-    invocation_span_id = plugin._get_span(None).get_span_context().span_id
+    before = otel_context.get_current()
 
     operation_id = "step-fail"
     plugin.on_user_function_start(_user_function_start_info(operation_id))
@@ -827,21 +840,27 @@ def test_user_function_end_restores_invocation_span_on_failure():
         _user_function_end_info(operation_id, outcome=UserFunctionOutcome.FAILED)
     )
 
-    assert trace.get_current_span().get_span_context().span_id == invocation_span_id
+    assert otel_context.get_current() is before
+    assert context_scope.depth(plugin) == 0
+    plugin.on_invocation_end(_invocation_end_info())
 
 
-def test_user_function_end_restores_invocation_span_across_multiple_steps():
-    """Verify between-step context is the invocation span across many steps."""
+def test_sequential_steps_do_not_accumulate_scopes():
+    """Verify N sequential steps leave no residue: depth returns to 0 each time."""
     plugin, _ = _create_plugin()
     plugin.on_invocation_start(_invocation_start_info())
-    invocation_span_id = plugin._get_span(None).get_span_context().span_id
+    before = otel_context.get_current()
 
     for index in range(3):
         operation_id = f"step-{index}"
         plugin.on_user_function_start(_user_function_start_info(operation_id))
+        assert context_scope.depth(plugin) == 1
         plugin.on_user_function_end(_user_function_end_info(operation_id))
-        # Between each step, the invocation span is the current span.
-        assert trace.get_current_span().get_span_context().span_id == invocation_span_id
+        # Between each step the context is exactly what preceded the step, and
+        # no scope has accumulated.
+        assert otel_context.get_current() is before
+        assert context_scope.depth(plugin) == 0
+    plugin.on_invocation_end(_invocation_end_info())
 
 
 # ----------------------------------------------------------------------
@@ -878,6 +897,11 @@ def test_get_current_span_context_returns_operation_span_inside_step():
     assert active_attempt_span is not None
     assert span_context.span_id == active_attempt_span.get_span_context().span_id
 
+    # Close the lifecycle so the attempt scope is detached; the autouse fixture
+    # asserts no scope outlives the test.
+    plugin.on_user_function_end(_user_function_end_info(operation_id))
+    plugin.on_invocation_end(_invocation_end_info())
+
 
 def test_get_current_span_context_returns_invocation_span_between_steps():
     """Verify between-step code resolves back to the invocation span context."""
@@ -929,25 +953,37 @@ def test_user_function_end_restores_parent_context_span_for_nested_step():
     )
 
     # After the inner step, the enclosing child-context span is current again,
-    # NOT the invocation span.
+    # NOT the invocation span. The inner scope was detached, landing back on the
+    # context scope that is still attached.
     assert trace.get_current_span().get_span_context().span_id == context_span_id
     assert (
         trace.get_current_span().get_span_context().span_id
         != plugin._get_span(None).get_span_context().span_id
     )
 
+    plugin.on_user_function_end(
+        _user_function_end_info(context_id, operation_type=OperationType.CONTEXT)
+    )
+    plugin.on_invocation_end(_invocation_end_info())
+
 
-def test_user_function_end_falls_back_to_invocation_when_parent_missing():
-    """Verify a top-level step (parent_id=None) restores the invocation span."""
+def test_top_level_step_restores_ambient_context():
+    """Verify a top-level step (parent_id=None) restores the ambient context.
+
+    There is no enclosing operation scope to fall back to, so the thread returns
+    to whatever was current before the step. Log correlation for that window is
+    covered by get_current_span_context()'s invocation-span fallback.
+    """
     plugin, _ = _create_plugin()
     plugin.on_invocation_start(_invocation_start_info())
-    invocation_span_id = plugin._get_span(None).get_span_context().span_id
+    before = otel_context.get_current()
 
     operation_id = "step-1"
     plugin.on_user_function_start(_user_function_start_info(operation_id))
     plugin.on_user_function_end(_user_function_end_info(operation_id))
 
-    assert trace.get_current_span().get_span_context().span_id == invocation_span_id
+    assert otel_context.get_current() is before
+    plugin.on_invocation_end(_invocation_end_info())
 
 
 def test_get_current_span_context_returns_context_span_between_nested_steps():
@@ -979,6 +1015,11 @@ def test_get_current_span_context_returns_context_span_between_nested_steps():
     assert span_context.span_id == context_span.get_span_context().span_id
     assert span_context.span_id != plugin._get_span(None).get_span_context().span_id
 
+    plugin.on_user_function_end(
+        _user_function_end_info(context_id, operation_type=OperationType.CONTEXT)
+    )
+    plugin.on_invocation_end(_invocation_end_info())
+
 
 def test_nested_steps_restore_context_span_across_multiple_iterations():
     """Verify each inner step restores the child-context span between iterations."""
@@ -1002,6 +1043,11 @@ def test_nested_steps_restore_context_span_across_multiple_iterations():
         # Between each inner step, the child-context span is current.
         assert trace.get_current_span().get_span_context().span_id == context_span_id
 
+    plugin.on_user_function_end(
+        _user_function_end_info(context_id, operation_type=OperationType.CONTEXT)
+    )
+    plugin.on_invocation_end(_invocation_end_info())
+
 
 @pytest.mark.parametrize(
     ("status", "expected_code"),
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py
index c1c6d4b1..e5344439 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py
@@ -41,6 +41,7 @@
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from opentelemetry.trace import SpanKind
 
+from aws_durable_execution_sdk_python_otel import context_scope
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     operation_id_to_span_id,
 )
@@ -63,12 +64,17 @@
 
 
 @pytest.fixture(autouse=True)
-def _reset_otel_context():
-    token = otel_context.attach(Context())
-    try:
-        yield
-    finally:
-        otel_context.detach(token)
+def _assert_otel_context_balanced():
+    """Fail any test that leaves an OTel context scope attached."""
+    before = otel_context.get_current()
+    before_depth = context_scope.depth()
+    yield
+    assert context_scope.depth() == before_depth, (
+        "test left OTel context scopes attached"
+    )
+    assert otel_context.get_current() is before, (
+        "test did not restore the OTel context it started with"
+    )
 
 
 def _provider() -> tuple[TracerProvider, InMemorySpanExporter]:
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py
index f419f9ad..91f56b1d 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py
@@ -6,18 +6,26 @@
 from datetime import UTC, datetime
 
 from aws_durable_execution_sdk_python.lambda_service import (
+    InvocationStatus,
     OperationStatus,
     OperationType,
 )
 from aws_durable_execution_sdk_python.plugin import (
+    InvocationEndInfo,
     InvocationStartInfo,
+    UserFunctionEndInfo,
+    UserFunctionOutcome,
     UserFunctionStartInfo,
 )
+import opentelemetry.context as otel_context
+import pytest
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
+from aws_durable_execution_sdk_python_otel import context_scope
+from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin
 from aws_durable_execution_sdk_python_otel.log_filter import (
     OtelContextLogFilter,
     install_log_filter,
@@ -33,6 +41,32 @@
 EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST"
 
 
+@pytest.fixture(autouse=True)
+def _assert_otel_context_balanced():
+    """Fail any test that leaves an OTel context scope attached."""
+    before = otel_context.get_current()
+    before_depth = context_scope.depth()
+    yield
+    assert context_scope.depth() == before_depth
+    assert otel_context.get_current() is before
+
+
+def _create_execution_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]:
+    """Create an ExecutionOtelPlugin wired to an in-memory span exporter."""
+    exporter = InMemorySpanExporter()
+    trace_provider = TracerProvider()
+    trace_provider.add_span_processor(SimpleSpanProcessor(exporter))
+    plugin = ExecutionOtelPlugin(
+        OtelPluginConfig(
+            provider_source=ProviderSource.EXPLICIT,
+            tracer_provider=trace_provider,
+            context_extractor=lambda _: Context(),
+            enrich_logger=False,
+        )
+    )
+    return plugin, exporter
+
+
 def _create_plugin(
     enrich_logger: bool = True,
 ) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]:
@@ -77,6 +111,37 @@ def _user_function_start_info(operation_id: str) -> UserFunctionStartInfo:
     )
 
 
+def _invocation_end_info() -> InvocationEndInfo:
+    """Create standard invocation end info for tests."""
+    return InvocationEndInfo(
+        request_id="request-1",
+        execution_arn=EXECUTION_ARN,
+        execution_start_time=START_TIME,
+        is_first_invocation=True,
+        status=InvocationStatus.SUCCEEDED,
+        error=None,
+    )
+
+
+def _user_function_end_info(operation_id: str) -> UserFunctionEndInfo:
+    """Create standard user function end info for tests."""
+    return UserFunctionEndInfo(
+        operation_id=operation_id,
+        operation_type=OperationType.STEP,
+        sub_type=None,
+        name="fetch-user",
+        parent_id=None,
+        start_time=START_TIME,
+        end_time=START_TIME,
+        is_replayed=False,
+        status=OperationStatus.SUCCEEDED,
+        is_replay_children=False,
+        attempt=1,
+        outcome=UserFunctionOutcome.SUCCEEDED,
+        error=None,
+    )
+
+
 def _make_record() -> logging.LogRecord:
     """Create a bare LogRecord for filtering."""
     return logging.LogRecord(
@@ -148,6 +213,31 @@ def test_filter_uses_attempt_span_inside_user_function():
     expected_span_id = format(attempt_span.get_span_context().span_id, "016x")
     assert record.spanId == expected_span_id
 
+    plugin.on_user_function_end(_user_function_end_info(operation_id))
+    plugin.on_invocation_end(_invocation_end_info())
+
+
+def test_execution_plugin_handler_thread_uses_the_invocation_span():
+    """Handler-thread records correlate to the Invocation span, not Workflow.
+
+    ExecutionOtelPlugin attaches nothing on the handler thread, so the filter
+    resolves through the plugin registry, which prefers the Invocation span. The
+    trace ID is shared with the Workflow span either way.
+    """
+    plugin, _ = _create_execution_plugin()
+    plugin.on_invocation_start(_invocation_start_info())
+
+    record = _make_record()
+    OtelContextLogFilter(plugin).filter(record)
+
+    invocation_context = plugin._invocation_span.get_span_context()
+    workflow_context = plugin._workflow_span.get_span_context()
+    assert record.spanId == format(invocation_context.span_id, "016x")
+    assert record.spanId != format(workflow_context.span_id, "016x")
+    assert record.traceId == format(workflow_context.trace_id, "032x")
+
+    plugin.on_invocation_end(_invocation_end_info())
+
 
 def test_install_log_filter_attaches_to_handlers():
     """install_log_filter adds the filter to each handler on the target logger."""

From c79cfdaec2b42f298a7f7f34388020e664e876da Mon Sep 17 00:00:00 2001
From: Alex Wang <wangyb@amazon.com>
Date: Fri, 14 Aug 2026 22:11:06 +0000
Subject: [PATCH 2/4] fix(otel): address review on context scopes

Three review findings, all real:

Same-invocation re-entry. The epoch check only caught a previous
invocation's leftovers, but the same operation key can be entered
twice inside one invocation: a suspended operation is re-entered
when its branch is resubmitted, and its first scope is still
attached because the suspending path has no end hook. The second
enter stacked on the first and the end hook popped one, leaving a
stale layer per re-entry. enter_scope now unwinds an existing
(owner, key) even when the epoch matches.

Extracted context values. Basing every scope on the current
context dropped baggage and other non-span values supplied by the
context extractor, because the worker running user code starts
with an empty context. The outermost scope on a thread is now
layered onto the extracted context; nested scopes keep using the
current one, which already carries it transitively.

Ambient span vs durable span. In GLOBAL mode the ADOT Lambda span
stays current on the handler thread, so get_current_span_context
returned it instead of the Invocation span, contradicting what the
previous commit documented. The current span is now trusted only
while this plugin holds a scope on this thread; otherwise the
registry answers. The earlier test missed this by using an
explicit provider with no ambient span.

Adds the tests each finding asked for: same-key re-entry at the
helper level and through the plugin hooks on a worker thread,
baggage surviving into user code and into a nested scope, and a
GLOBAL-mode ambient Lambda span not displacing the Invocation
span in log records.
---
 .../context_scope.py                          |  27 +++++
 .../execution_plugin.py                       |  46 ++++++--
 .../invocation_plugin.py                      |  57 ++++++----
 .../tests/test_context_scope.py               | 104 +++++++++++++++++-
 .../tests/test_log_filter.py                  |  41 +++++++
 5 files changed, 246 insertions(+), 29 deletions(-)

diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
index 54c46020..52ce9b1f 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
@@ -95,6 +95,7 @@ def enter_scope(owner: Any, key: str, context: Context, epoch: int = 0) -> None:
 
     owner_id = id(owner)
     _discard_stale(owner_id, epoch)
+    _discard_reentered(owner_id, key)
     try:
         token = otel_context.attach(context)
     except Exception:  # noqa: BLE001
@@ -145,6 +146,32 @@ def depth(owner: Any | None = None) -> int:
     return sum(1 for entry in _state.entries if entry.owner_id == owner_id)
 
 
+def _discard_reentered(owner_id: int, key: str) -> None:
+    """Unwind a scope this owner already holds under ``key`` on this thread.
+
+    The epoch check only catches a *previous invocation's* leftovers. The same
+    operation key can also be entered twice inside one invocation: a suspended
+    operation is re-entered when its branch is resubmitted, and its first scope
+    is still attached because the suspending path had no end hook to pop it.
+    Without this, the second enter would stack on the first and the eventual end
+    hook -- which pops one scope -- would leave the original attached, one stale
+    layer per re-entry.
+    """
+    index = next(
+        (
+            position
+            for position, entry in enumerate(_state.entries)
+            if entry.owner_id == owner_id and entry.key == key
+        ),
+        None,
+    )
+    if index is None:
+        return
+    for entry in reversed(_state.entries[index:]):
+        _detach(entry)
+    del _state.entries[index:]
+
+
 def _discard_stale(owner_id: int, epoch: int) -> None:
     """Unwind this owner's scopes left over from a previous epoch."""
     index = next(
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
index 2d84b104..69f5dfbd 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
@@ -183,11 +183,40 @@ def _scope_key(cls, info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
             return cls._attempt_key(info)
         return info.operation_id
 
+    def _scope_base_context(self) -> Context:
+        """Return the context a new operation scope should be layered onto.
+
+        For the outermost durable scope on a thread, the extracted upstream
+        context is the base: user code runs on a worker whose context starts
+        empty, so anything the context extractor supplied -- a remote parent,
+        baggage -- would be dropped by using the current context there, and
+        downstream propagation inside steps would lose it.
+
+        For a nested scope the current context already carries that extracted
+        context transitively, via the enclosing scope, so it is used as the base
+        to preserve whatever ran in between (ambient spans, baggage added by user
+        code).
+        """
+        if context_scope.depth(self) > 0:
+            return otel_context.get_current()
+        return self._extracted_context or otel_context.get_current()
+
     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()
-        if span_context and span_context.is_valid:
-            return span_context
+        """Return the active span context for log correlation (see log_filter).
+
+        The current span is used only while this plugin holds an operation scope
+        on this thread -- inside a step or child context, where the current span
+        is the one this plugin attached. Otherwise the registry is used, so a
+        record emitted between operations, or on the handler thread, correlates to
+        the durable Invocation span rather than to whatever else happens to be
+        current. That distinction matters in GLOBAL (ADOT) mode: the ambient
+        Lambda span is current on the handler thread and would otherwise be
+        reported in place of the durable span.
+        """
+        if context_scope.depth(self) > 0:
+            span_context = trace.get_current_span().get_span_context()
+            if span_context and span_context.is_valid:
+                return span_context
         for candidate in (self._invocation_span, self._workflow_span):
             if candidate is not None:
                 ctx = candidate.get_span_context()
@@ -483,14 +512,13 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
                 start_time=info.start_time,
             )
         # 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.
+        # user function become children of this span. The span's own parent was
+        # chosen explicitly in _start_span; this only sets what is ambient while
+        # the user function runs.
         context_scope.enter_scope(
             self,
             self._scope_key(info),
-            trace.set_span_in_context(span, otel_context.get_current()),
+            trace.set_span_in_context(span, self._scope_base_context()),
             epoch=self._epoch,
         )
 
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
index 60988469..0059a910 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
@@ -211,28 +211,48 @@ def _scope_key(cls, info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
             return cls._attempt_span_key(info)
         return info.operation_id
 
+    def _scope_base_context(self) -> Context:
+        """Return the context a new operation scope should be layered onto.
+
+        For the outermost durable scope on a thread, the extracted upstream
+        context is the base: user code runs on a worker whose context starts
+        empty, so anything the context extractor supplied -- a remote parent,
+        baggage -- would be dropped by using the current context there, and
+        downstream propagation inside steps would lose it.
+
+        For a nested scope the current context already carries that extracted
+        context transitively, via the enclosing scope, so it is used as the base
+        to preserve whatever ran in between (ambient spans, baggage added by user
+        code).
+        """
+        if context_scope.depth(self) > 0:
+            return context.get_current()
+        return self._extracted_context or context.get_current()
+
     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). After a
-           nested operation ends, its scope is detached and the enclosing child
-           context span -- still attached -- becomes current again.
-        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
-           any thread's context, so the registry is the only way to resolve it.
-           It also covers code that runs between top-level operations, where
-           detaching the operation scope leaves the thread's ambient context
-           current.
+        1. The span attached to the OTel thread-local context, but only while this
+           plugin holds an operation scope on this thread. Inside a step that is
+           the active attempt span, and inside a child context the active context
+           span. After a nested operation ends, its scope is detached and the
+           enclosing child context span -- still attached -- becomes current
+           again.
+        2. The invocation span from the plugin registry. This covers top-level
+           handler code (the invocation span is never attached to any thread's
+           context) and code between top-level operations. Gating step 1 on an
+           owned scope matters in GLOBAL (ADOT) mode: the ambient Lambda span is
+           current on the handler thread and would otherwise be reported in place
+           of the durable span.
 
         Returns:
             A valid SpanContext, or None if no span is active.
         """
-        span_context = trace.get_current_span().get_span_context()
-        if span_context and span_context.is_valid:
-            return span_context
+        if context_scope.depth(self) > 0:
+            span_context = trace.get_current_span().get_span_context()
+            if span_context and span_context.is_valid:
+                return span_context
 
         invocation_span = self._get_span(None)
         if invocation_span:
@@ -588,14 +608,13 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
             deterministic_span_id=info.operation_type is not OperationType.STEP,
         )
         # 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.
+        # user function become children of this span. The span's own parent was
+        # chosen explicitly in _start_span; this only sets what is ambient while
+        # the user function runs.
         context_scope.enter_scope(
             self,
             self._scope_key(info),
-            trace.set_span_in_context(span, context.get_current()),
+            trace.set_span_in_context(span, self._scope_base_context()),
             epoch=self._epoch,
         )
 
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
index abef60ba..fd0799e2 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
@@ -25,7 +25,7 @@
     UserFunctionOutcome,
     UserFunctionStartInfo,
 )
-from opentelemetry import trace
+from opentelemetry import baggage, trace
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
@@ -133,6 +133,42 @@ def test_exit_with_unknown_key_is_a_noop():
     assert context_scope.depth(owner) == 0
 
 
+def test_reentering_the_same_key_replaces_the_previous_scope():
+    """Re-entering a key inside one invocation must not stack a second scope.
+
+    A suspended operation is re-entered when its branch is resubmitted, and its
+    first scope is still attached because the suspending path had no end hook to
+    pop it. The epoch is unchanged, so only the same-key guard catches this.
+    """
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    context_scope.enter_scope(owner, "wfc-1", _span_context("poll-1"), epoch=1)
+    context_scope.enter_scope(owner, "wfc-1", _span_context("poll-2"), epoch=1)
+
+    assert context_scope.depth(owner) == 1
+
+    context_scope.exit_scope(owner, "wfc-1")
+    assert otel_context.get_current() is before
+    assert context_scope.depth(owner) == 0
+
+
+def test_reentry_guard_keeps_enclosing_scopes():
+    """Re-entering a nested key must not disturb the scope it is nested in."""
+    owner = _Owner()
+    context_scope.enter_scope(owner, "ctx", _span_context("ctx"), epoch=1)
+    enclosing = otel_context.get_current()
+
+    context_scope.enter_scope(owner, "inner", _span_context("inner-1"), epoch=1)
+    context_scope.enter_scope(owner, "inner", _span_context("inner-2"), epoch=1)
+
+    assert context_scope.depth(owner) == 2
+    context_scope.exit_scope(owner, "inner")
+    assert otel_context.get_current() is enclosing
+
+    context_scope.exit_scope(owner, "ctx")
+
+
 def test_enter_discards_scopes_from_a_previous_epoch():
     """A scope a suspended operation left behind must not outlive its invocation.
 
@@ -347,6 +383,72 @@ def run_step() -> None:
     plugin.on_invocation_end(_invocation_end())
 
 
+@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
+def test_extracted_context_values_reach_user_code_on_a_worker_thread(factory):
+    """Values from the context extractor must be current inside user code.
+
+    The worker running user code starts with an empty context, so the outermost
+    durable scope has to be layered onto the extracted context -- otherwise
+    baggage and any other non-span values the extractor supplied are dropped and
+    downstream instrumentation inside the step cannot propagate them.
+    """
+    plugin, _ = factory()
+    # An extractor that supplies baggage, as a propagator-based one would.
+    plugin._context_extractor = lambda _info: baggage.set_baggage(
+        "tenant", "acme", context=Context()
+    )
+    plugin.on_invocation_start(_invocation_start())
+    observed: dict[str, object] = {}
+
+    def run_step() -> None:
+        plugin.on_user_function_start(_step_start("step-1"))
+        observed["inside"] = baggage.get_baggage("tenant")
+        # A nested scope keeps it too, since it layers onto the current context.
+        plugin.on_user_function_start(_step_start("step-2"))
+        observed["nested"] = baggage.get_baggage("tenant")
+        plugin.on_user_function_end(_step_end("step-2"))
+        plugin.on_user_function_end(_step_end("step-1"))
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(run_step).result()
+
+    assert observed["inside"] == "acme"
+    assert observed["nested"] == "acme"
+
+    plugin.on_invocation_end(_invocation_end())
+
+
+@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
+def test_suspend_then_reenter_then_end_leaves_no_residue(factory):
+    """A suspended operation re-entered in the same invocation stays balanced.
+
+    The first start has no matching end (the SDK re-raises SuspendExecution), so
+    the re-entry must replace that scope rather than stack on it.
+    """
+    plugin, _ = factory()
+    before = otel_context.get_current()
+    plugin.on_invocation_start(_invocation_start())
+    observed: dict[str, object] = {}
+
+    def run_polls() -> None:
+        # Poll 1 suspends: start fires, end never does.
+        plugin.on_user_function_start(_step_start("wfc-1"))
+        # Poll 2 re-enters the same operation and completes.
+        plugin.on_user_function_start(_step_start("wfc-1"))
+        observed["depth_after_reentry"] = context_scope.depth(plugin)
+        plugin.on_user_function_end(_step_end("wfc-1"))
+        observed["depth_after_end"] = context_scope.depth(plugin)
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(run_polls).result()
+
+    assert observed["depth_after_reentry"] == 1
+    assert observed["depth_after_end"] == 0
+    assert otel_context.get_current() is before
+
+    plugin.on_invocation_end(_invocation_end())
+
+
 def test_two_plugins_on_one_thread_unwind_in_lifo_order():
     """Both plugins ship as entry points and can be enabled together.
 
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py
index 91f56b1d..dc9af8a9 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py
@@ -18,6 +18,7 @@
     UserFunctionStartInfo,
 )
 import opentelemetry.context as otel_context
+from opentelemetry import trace
 import pytest
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import TracerProvider
@@ -239,6 +240,46 @@ def test_execution_plugin_handler_thread_uses_the_invocation_span():
     plugin.on_invocation_end(_invocation_end_info())
 
 
+def test_ambient_lambda_span_does_not_displace_the_invocation_span(monkeypatch):
+    """An ambient ADOT span must not be reported in place of the durable span.
+
+    In GLOBAL mode the Lambda invocation span from the ADOT layer stays current on
+    the handler thread. Log records emitted there must still carry the durable
+    Invocation span, so the filter only trusts the current span while the plugin
+    holds an operation scope on that thread.
+    """
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider)
+    plugin = ExecutionOtelPlugin(
+        OtelPluginConfig(
+            provider_source=ProviderSource.GLOBAL,
+            context_extractor=lambda _: Context(),
+            enrich_logger=False,
+        )
+    )
+
+    ambient = provider.get_tracer("ambient").start_span("lambda-invocation")
+    token = otel_context.attach(trace.set_span_in_context(ambient))
+    try:
+        plugin.on_invocation_start(_invocation_start_info())
+        record = _make_record()
+        OtelContextLogFilter(plugin).filter(record)
+
+        ambient_span_id = format(ambient.get_span_context().span_id, "016x")
+        invocation_span_id = format(
+            plugin._invocation_span.get_span_context().span_id, "016x"
+        )
+        assert record.spanId == invocation_span_id
+        assert record.spanId != ambient_span_id
+
+        plugin.on_invocation_end(_invocation_end_info())
+    finally:
+        otel_context.detach(token)
+        ambient.end()
+
+
 def test_install_log_filter_attaches_to_handlers():
     """install_log_filter adds the filter to each handler on the target logger."""
     plugin, _ = _create_plugin()

From 6c0c70f55cec2f8bc132da81a9c89eb3e4b8cc1f Mon Sep 17 00:00:00 2001
From: Alex Wang <wangyb@amazon.com>
Date: Fri, 14 Aug 2026 23:00:10 +0000
Subject: [PATCH 3/4] fix(otel): drop scopes a worker no longer owns

Two more review findings, both real.

Branch workers have no branch affinity. If branch A suspends
without an end hook and its pool worker next runs branch B, A's
scope has the same epoch and a different key, so neither the epoch
check nor the same-key check cleared it. B nested inside A and, on
exit, detached back into it -- later records on that worker
correlated to the wrong branch, and one layer accumulated per
suspended branch.

Replaces the same-key guard with an ancestry check: a scope may
stay attached only while the operation it belongs to is still
running on this thread, so anything above the new scope's parent is
stale, and when the parent is absent -- a root-level operation, or
one whose parent ran elsewhere -- nothing held here can enclose it.
The plugins pass the enclosing operation as parent_key. The normal
nesting path is a no-op, which matters because detaching
necessarily discards entries above the cut, including a second
plugin's.

An empty extracted context was being discarded. Context subclasses
dict, so an empty one is falsy and `extracted or current` silently
inverted the intent of an extractor that returns an empty context
to isolate the operation, inheriting the worker's ambient baggage
and suppression values instead. Now tested with `is not None`.

Adds the tests both findings asked for: a sibling scope dropped
rather than nested into, a nested scope whose parent never ran on
this thread, a plugin-level branch-A-suspends-then-branch-B case
pinned to one worker, and an empty extracted context isolating the
operation from ambient baggage. Two existing helper tests nested
without declaring a parent, which now reads as a root operation;
they pass parent_key like the plugins do.
---
 .../context_scope.py                          |  80 ++++++---
 .../execution_plugin.py                       |   9 +-
 .../invocation_plugin.py                      |   9 +-
 .../tests/test_context_scope.py               | 157 +++++++++++++++++-
 4 files changed, 221 insertions(+), 34 deletions(-)

diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
index 52ce9b1f..fcbcb6a2 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
@@ -75,14 +75,20 @@ def _detach(entry: _Entry) -> None:
         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:
+def enter_scope(
+    owner: Any,
+    key: str,
+    context: Context,
+    epoch: int = 0,
+    parent_key: str | None = None,
+) -> 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.
+    Scopes this owner holds that the new one does not nest inside are unwound
+    first, as are any left over from an earlier ``epoch``. Both cover 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 the worker that ran it goes on to other work.
 
     Args:
         owner: The plugin instance pushing the scope.
@@ -90,12 +96,14 @@ def enter_scope(owner: Any, key: str, context: Context, epoch: int = 0) -> None:
         context: The context to attach.
         epoch: The owner's invocation counter; scopes from older epochs are
             discarded before the new scope is pushed.
+        parent_key: Scope key of the enclosing operation, or None for a
+            root-level one. Used to tell legitimate nesting from a stale scope.
     """
     from opentelemetry import context as otel_context
 
     owner_id = id(owner)
     _discard_stale(owner_id, epoch)
-    _discard_reentered(owner_id, key)
+    _unwind_non_ancestors(owner_id, parent_key)
     try:
         token = otel_context.attach(context)
     except Exception:  # noqa: BLE001
@@ -146,30 +154,48 @@ def depth(owner: Any | None = None) -> int:
     return sum(1 for entry in _state.entries if entry.owner_id == owner_id)
 
 
-def _discard_reentered(owner_id: int, key: str) -> None:
-    """Unwind a scope this owner already holds under ``key`` on this thread.
+def _unwind_non_ancestors(owner_id: int, parent_key: str | None) -> None:
+    """Drop scopes this owner holds on this thread that the new scope is not inside.
 
-    The epoch check only catches a *previous invocation's* leftovers. The same
-    operation key can also be entered twice inside one invocation: a suspended
-    operation is re-entered when its branch is resubmitted, and its first scope
-    is still attached because the suspending path had no end hook to pop it.
-    Without this, the second enter would stack on the first and the eventual end
-    hook -- which pops one scope -- would leave the original attached, one stale
-    layer per re-entry.
+    A scope may stay attached only while the operation it belongs to is still
+    running on this thread, and the only place that can be checked is here: the
+    suspending path has no end hook, so a suspended operation leaves its scope
+    behind. Two cases produce one:
+
+    * A branch-pool worker has no branch affinity. If branch A suspends and its
+      worker next runs branch B, A's scope has the same epoch and a different key,
+      so neither the epoch nor a same-key check clears it. B would nest inside A
+      and, on exit, detach back into it.
+    * The same operation can be re-entered when its branch is resubmitted.
+
+    The new scope nests inside its parent operation, so anything above that parent
+    is stale. When the parent is absent -- a root-level operation, or a parent that
+    never ran on this thread -- nothing this owner holds here can enclose it.
+
+    Detaching necessarily discards entries stacked above the cut, including other
+    owners', because the underlying ``ContextVar`` can only be reset in order. The
+    normal nesting path is a no-op, so that does not disturb a second plugin
+    tracking the same operations.
     """
-    index = next(
-        (
-            position
-            for position, entry in enumerate(_state.entries)
-            if entry.owner_id == owner_id and entry.key == key
-        ),
-        None,
-    )
-    if index is None:
+    positions = [
+        position
+        for position, entry in enumerate(_state.entries)
+        if entry.owner_id == owner_id
+    ]
+    if not positions:
         return
-    for entry in reversed(_state.entries[index:]):
+    if parent_key is not None and _state.entries[positions[-1]].key == parent_key:
+        # Normal nesting: the innermost scope we hold is the new scope's parent.
+        return
+    cut = positions[0]
+    if parent_key is not None:
+        for position in reversed(positions):
+            if _state.entries[position].key == parent_key:
+                cut = position + 1
+                break
+    for entry in reversed(_state.entries[cut:]):
         _detach(entry)
-    del _state.entries[index:]
+    del _state.entries[cut:]
 
 
 def _discard_stale(owner_id: int, epoch: int) -> None:
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
index 69f5dfbd..0ffd5b62 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
@@ -199,7 +199,13 @@ def _scope_base_context(self) -> Context:
         """
         if context_scope.depth(self) > 0:
             return otel_context.get_current()
-        return self._extracted_context or otel_context.get_current()
+        # An extractor may deliberately return an empty Context to isolate the
+        # operation from whatever is ambient. Context subclasses dict, so an empty
+        # one is falsy -- test for None, or that intent silently inverts into
+        # inheriting the worker's ambient baggage and suppression values.
+        if self._extracted_context is not None:
+            return self._extracted_context
+        return otel_context.get_current()
 
     def get_current_span_context(self) -> SpanContext | None:
         """Return the active span context for log correlation (see log_filter).
@@ -520,6 +526,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
             self._scope_key(info),
             trace.set_span_in_context(span, self._scope_base_context()),
             epoch=self._epoch,
+            parent_key=info.parent_id,
         )
 
     def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
index 0059a910..686fbb10 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
@@ -227,7 +227,13 @@ def _scope_base_context(self) -> Context:
         """
         if context_scope.depth(self) > 0:
             return context.get_current()
-        return self._extracted_context or context.get_current()
+        # An extractor may deliberately return an empty Context to isolate the
+        # operation from whatever is ambient. Context subclasses dict, so an empty
+        # one is falsy -- test for None, or that intent silently inverts into
+        # inheriting the worker's ambient baggage and suppression values.
+        if self._extracted_context is not None:
+            return self._extracted_context
+        return context.get_current()
 
     def get_current_span_context(self) -> SpanContext | None:
         """Return the span context to use for log correlation.
@@ -616,6 +622,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
             self._scope_key(info),
             trace.set_span_in_context(span, self._scope_base_context()),
             epoch=self._epoch,
+            parent_key=info.parent_id,
         )
 
     def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
index fd0799e2..978af0e2 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
@@ -96,7 +96,9 @@ def test_nested_scopes_pop_in_lifo_order():
 
     context_scope.enter_scope(owner, "outer", _span_context("outer"))
     outer = otel_context.get_current()
-    context_scope.enter_scope(owner, "inner", _span_context("inner"))
+    context_scope.enter_scope(
+        owner, "inner", _span_context("inner"), parent_key="outer"
+    )
 
     context_scope.exit_scope(owner, "inner")
     assert otel_context.get_current() is outer
@@ -115,7 +117,9 @@ def test_exit_unwinds_scopes_stacked_above_the_target():
     before = otel_context.get_current()
 
     context_scope.enter_scope(owner, "outer", _span_context("outer"))
-    context_scope.enter_scope(owner, "orphan", _span_context("orphan"))
+    context_scope.enter_scope(
+        owner, "orphan", _span_context("orphan"), parent_key="outer"
+    )
 
     context_scope.exit_scope(owner, "outer")
 
@@ -138,7 +142,7 @@ def test_reentering_the_same_key_replaces_the_previous_scope():
 
     A suspended operation is re-entered when its branch is resubmitted, and its
     first scope is still attached because the suspending path had no end hook to
-    pop it. The epoch is unchanged, so only the same-key guard catches this.
+    pop it. The epoch is unchanged, so only the ancestry check catches this.
     """
     owner = _Owner()
     before = otel_context.get_current()
@@ -153,14 +157,58 @@ def test_reentering_the_same_key_replaces_the_previous_scope():
     assert context_scope.depth(owner) == 0
 
 
+def test_a_suspended_sibling_scope_is_dropped_not_nested_into():
+    """A worker that moves to another branch must not inherit the old one's scope.
+
+    Branch-pool workers have no branch affinity. Branch A suspends without an end
+    hook; the same worker then runs branch B. B's scope must replace A's, not
+    stack on it -- otherwise exiting B would detach back into A and correlate
+    later records to the wrong branch.
+    """
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    # Branch A runs and suspends: entered, never exited.
+    context_scope.enter_scope(owner, "branch-a", _span_context("branch-a"), epoch=1)
+
+    # The worker picks up branch B, a sibling: same epoch, different key, no parent.
+    context_scope.enter_scope(owner, "branch-b", _span_context("branch-b"), epoch=1)
+
+    assert context_scope.depth(owner) == 1
+
+    context_scope.exit_scope(owner, "branch-b")
+    assert otel_context.get_current() is before
+    assert context_scope.depth(owner) == 0
+
+
+def test_a_scope_whose_parent_never_ran_here_replaces_the_stale_one():
+    """A nested scope whose parent is absent cannot nest inside what is here."""
+    owner = _Owner()
+    before = otel_context.get_current()
+
+    context_scope.enter_scope(owner, "ctx-a", _span_context("ctx-a"), epoch=1)
+    # An operation nested under a context that ran on another thread.
+    context_scope.enter_scope(
+        owner, "step-b", _span_context("step-b"), epoch=1, parent_key="ctx-b"
+    )
+
+    assert context_scope.depth(owner) == 1
+    context_scope.exit_scope(owner, "step-b")
+    assert otel_context.get_current() is before
+
+
 def test_reentry_guard_keeps_enclosing_scopes():
     """Re-entering a nested key must not disturb the scope it is nested in."""
     owner = _Owner()
     context_scope.enter_scope(owner, "ctx", _span_context("ctx"), epoch=1)
     enclosing = otel_context.get_current()
 
-    context_scope.enter_scope(owner, "inner", _span_context("inner-1"), epoch=1)
-    context_scope.enter_scope(owner, "inner", _span_context("inner-2"), epoch=1)
+    context_scope.enter_scope(
+        owner, "inner", _span_context("inner-1"), epoch=1, parent_key="ctx"
+    )
+    context_scope.enter_scope(
+        owner, "inner", _span_context("inner-2"), epoch=1, parent_key="ctx"
+    )
 
     assert context_scope.depth(owner) == 2
     context_scope.exit_scope(owner, "inner")
@@ -169,6 +217,30 @@ def test_reentry_guard_keeps_enclosing_scopes():
     context_scope.exit_scope(owner, "ctx")
 
 
+def test_legitimate_nesting_leaves_another_owners_scope_alone():
+    """Two plugins tracking the same operations must not evict each other.
+
+    The normal nesting path has to be a no-op, because detaching necessarily
+    discards everything stacked above the cut.
+    """
+    first, second = _Owner(), _Owner()
+    before = otel_context.get_current()
+
+    context_scope.enter_scope(first, "ctx", _span_context("first-ctx"), epoch=1)
+    context_scope.enter_scope(second, "ctx", _span_context("second-ctx"), epoch=1)
+    context_scope.enter_scope(
+        first, "step", _span_context("first-step"), epoch=1, parent_key="ctx"
+    )
+
+    # The second plugin's context scope survived the first plugin's nesting.
+    assert context_scope.depth(second) == 1
+
+    context_scope.exit_scope(first, "step")
+    context_scope.exit_scope(second, "ctx")
+    context_scope.exit_scope(first, "ctx")
+    assert otel_context.get_current() is before
+
+
 def test_enter_discards_scopes_from_a_previous_epoch():
     """A scope a suspended operation left behind must not outlive its invocation.
 
@@ -383,6 +455,81 @@ def run_step() -> None:
     plugin.on_invocation_end(_invocation_end())
 
 
+@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
+def test_worker_reassigned_to_another_branch_drops_the_suspended_scope(factory):
+    """One worker running branch A (suspends) then branch B leaves no residue.
+
+    Branch pools reuse workers with no branch affinity, so the worker that ran a
+    suspended branch goes on to another. B must not nest inside A's abandoned
+    scope, and after B ends the thread must be clean.
+    """
+    plugin, _ = factory()
+    before = otel_context.get_current()
+    plugin.on_invocation_start(_invocation_start())
+    observed: dict[str, object] = {}
+
+    def branch_a() -> None:
+        # Suspends: start fires, end never does.
+        plugin.on_user_function_start(_step_start("branch-a"))
+
+    def branch_b() -> None:
+        plugin.on_user_function_start(_step_start("branch-b"))
+        observed["depth_in_b"] = context_scope.depth(plugin)
+        observed["current_is_b"] = trace.get_current_span().get_span_context().span_id
+        plugin.on_user_function_end(_step_end("branch-b"))
+        observed["depth_after_b"] = context_scope.depth(plugin)
+        observed["span_after_b_valid"] = (
+            trace.get_current_span().get_span_context().is_valid
+        )
+
+    # A single worker, so branch B is guaranteed to land on branch A's thread.
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(branch_a).result()
+        executor.submit(branch_b).result()
+
+    # B replaced A's scope rather than stacking on it.
+    assert observed["depth_in_b"] == 1
+    # Exiting B did not fall back into A.
+    assert observed["depth_after_b"] == 0
+    assert observed["span_after_b_valid"] is False
+    assert otel_context.get_current() is before
+
+    plugin.on_invocation_end(_invocation_end(InvocationStatus.PENDING))
+
+
+@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
+def test_an_empty_extracted_context_isolates_the_operation(factory):
+    """An extractor returning an empty Context must not inherit ambient values.
+
+    Context subclasses dict, so an empty one is falsy; treating it as "no context"
+    would silently invert the extractor's intent and leak the worker's ambient
+    baggage into the operation.
+    """
+    plugin, _ = factory()
+    plugin._context_extractor = lambda _info: Context()
+    plugin.on_invocation_start(_invocation_start())
+    observed: dict[str, object] = {}
+
+    def run_step() -> None:
+        # Something ambient on the worker, as auto-instrumentation might leave.
+        token = otel_context.attach(
+            baggage.set_baggage("ambient", "leaked", context=otel_context.get_current())
+        )
+        try:
+            plugin.on_user_function_start(_step_start("step-1"))
+            observed["inside"] = baggage.get_baggage("ambient")
+            plugin.on_user_function_end(_step_end("step-1"))
+        finally:
+            otel_context.detach(token)
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(run_step).result()
+
+    assert observed["inside"] is None
+
+    plugin.on_invocation_end(_invocation_end())
+
+
 @pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
 def test_extracted_context_values_reach_user_code_on_a_worker_thread(factory):
     """Values from the context extractor must be current inside user code.

From 1e9fa5c0c3566e7b0f800bab59faec24d0be7b95 Mon Sep 17 00:00:00 2001
From: Alex Wang <wangyb@amazon.com>
Date: Fri, 14 Aug 2026 23:45:53 +0000
Subject: [PATCH 4/4] fix(otel): stop inferring nesting from parent_id

Reverts the ancestry check from the previous commit and fixes the
order in which the attached context is built. Both were review
findings; the first one was a regression I introduced.

parent_id is checkpoint hierarchy, not the Python call stack. A
virtual (FLAT) map/parallel branch deliberately reports its inner
operations' parent as the grandparent -- None for a top-level
branch -- while the branch's own context scope is still running
(see DurableContext.is_virtual and create_child_context, where
child_parent_id is the *parent's* parent when is_virtual). Treating
such an inner step as root-level therefore detached the live branch
scope at the first inner step, and work between two inner steps
fell out of the durable trace. That is worse than the abandoned
sibling scope the check was meant to catch, so the narrower
same-key guard is restored.

The gap that leaves -- a scope abandoned by a *different* operation
on the same branch-pool worker -- cannot be closed from the hook
payloads, because a live FLAT branch scope is indistinguishable
from an abandoned sibling. It needs the SDK to report the end of a
suspended user function, which is tracked separately; the docstring
says so rather than implying the helper handles it.

Second finding: the context to attach was built by the caller
before enter_scope ran its cleanup, so it copied baggage and
suppression values out of the very scope about to be detached, and
detaching afterwards could not remove them from an already-built
Context. enter_scope now takes a factory and calls it after
cleanup.

Adds a FLAT-branch test asserting the branch scope stays current
across two inner steps that report no parent, and a test that the
factory observes the post-cleanup context. Drops the three tests
that asserted the reverted rule.
---
 .../context_scope.py                          |  96 ++++---
 .../execution_plugin.py                       |   5 +-
 .../invocation_plugin.py                      |   5 +-
 .../tests/test_context_scope.py               | 243 +++++++++---------
 4 files changed, 176 insertions(+), 173 deletions(-)

diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
index fcbcb6a2..f36b238b 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py
@@ -33,7 +33,7 @@
 import logging
 import threading
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Callable
 
 
 if TYPE_CHECKING:
@@ -78,34 +78,36 @@ def _detach(entry: _Entry) -> None:
 def enter_scope(
     owner: Any,
     key: str,
-    context: Context,
+    context_factory: Callable[[], Context],
     epoch: int = 0,
-    parent_key: str | None = None,
 ) -> None:
-    """Attach ``context`` on this thread and remember how to restore it.
+    """Attach a context on this thread and remember how to restore it.
 
-    Scopes this owner holds that the new one does not nest inside are unwound
-    first, as are any left over from an earlier ``epoch``. Both cover the paths
-    where a paired pop never runs: the SDK re-raises ``SuspendExecution`` without
+    Scopes left over from an earlier ``epoch``, or from an earlier entry under the
+    same ``key``, are unwound first: the SDK re-raises ``SuspendExecution`` without
     calling ``on_user_function_end``, so a suspended operation leaves its scope
-    attached, and the worker that ran it goes on to other work.
+    attached.
+
+    ``context_factory`` is called *after* that cleanup, not before. The context to
+    attach is normally derived from what is current, so building it first would
+    copy values from a scope that is about to be detached -- baggage, suppression
+    flags -- and detaching afterwards cannot remove them from a context that has
+    already been constructed.
 
     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.
+        context_factory: Builds the context to attach, called after cleanup.
         epoch: The owner's invocation counter; scopes from older epochs are
             discarded before the new scope is pushed.
-        parent_key: Scope key of the enclosing operation, or None for a
-            root-level one. Used to tell legitimate nesting from a stale scope.
     """
     from opentelemetry import context as otel_context
 
     owner_id = id(owner)
     _discard_stale(owner_id, epoch)
-    _unwind_non_ancestors(owner_id, parent_key)
+    _discard_reentered(owner_id, key)
     try:
-        token = otel_context.attach(context)
+        token = otel_context.attach(context_factory())
     except Exception:  # noqa: BLE001
         logger.debug("Failed to attach OTel context scope %s", key, exc_info=True)
         return
@@ -154,48 +156,38 @@ def depth(owner: Any | None = None) -> int:
     return sum(1 for entry in _state.entries if entry.owner_id == owner_id)
 
 
-def _unwind_non_ancestors(owner_id: int, parent_key: str | None) -> None:
-    """Drop scopes this owner holds on this thread that the new scope is not inside.
-
-    A scope may stay attached only while the operation it belongs to is still
-    running on this thread, and the only place that can be checked is here: the
-    suspending path has no end hook, so a suspended operation leaves its scope
-    behind. Two cases produce one:
-
-    * A branch-pool worker has no branch affinity. If branch A suspends and its
-      worker next runs branch B, A's scope has the same epoch and a different key,
-      so neither the epoch nor a same-key check clears it. B would nest inside A
-      and, on exit, detach back into it.
-    * The same operation can be re-entered when its branch is resubmitted.
-
-    The new scope nests inside its parent operation, so anything above that parent
-    is stale. When the parent is absent -- a root-level operation, or a parent that
-    never ran on this thread -- nothing this owner holds here can enclose it.
-
-    Detaching necessarily discards entries stacked above the cut, including other
-    owners', because the underlying ``ContextVar`` can only be reset in order. The
-    normal nesting path is a no-op, so that does not disturb a second plugin
-    tracking the same operations.
+def _discard_reentered(owner_id: int, key: str) -> None:
+    """Unwind a scope this owner already holds under ``key`` on this thread.
+
+    The epoch check only catches a *previous invocation's* leftovers. The same
+    operation key can also be entered twice inside one invocation, when a
+    suspended operation is re-entered after its branch is resubmitted, and its
+    first scope is still attached because the suspending path had no end hook to
+    pop it. Without this, the second enter would stack on the first and the
+    eventual end hook -- which pops one scope -- would leave the original
+    attached.
+
+    A scope abandoned by a *different* operation on this thread cannot be
+    detected here. Physical nesting is not derivable from the hook payloads:
+    ``parent_id`` is checkpoint hierarchy, and a FLAT map/parallel branch
+    deliberately reports its inner operations' parent as the grandparent (see
+    ``DurableContext.is_virtual``), so a live branch scope would be
+    indistinguishable from an abandoned sibling. Closing that gap needs the SDK
+    to report the end of a suspended user function, which it does not do today.
     """
-    positions = [
-        position
-        for position, entry in enumerate(_state.entries)
-        if entry.owner_id == owner_id
-    ]
-    if not positions:
-        return
-    if parent_key is not None and _state.entries[positions[-1]].key == parent_key:
-        # Normal nesting: the innermost scope we hold is the new scope's parent.
+    index = next(
+        (
+            position
+            for position, entry in enumerate(_state.entries)
+            if entry.owner_id == owner_id and entry.key == key
+        ),
+        None,
+    )
+    if index is None:
         return
-    cut = positions[0]
-    if parent_key is not None:
-        for position in reversed(positions):
-            if _state.entries[position].key == parent_key:
-                cut = position + 1
-                break
-    for entry in reversed(_state.entries[cut:]):
+    for entry in reversed(_state.entries[index:]):
         _detach(entry)
-    del _state.entries[cut:]
+    del _state.entries[index:]
 
 
 def _discard_stale(owner_id: int, epoch: int) -> None:
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
index 0ffd5b62..6decc311 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
@@ -524,9 +524,10 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
         context_scope.enter_scope(
             self,
             self._scope_key(info),
-            trace.set_span_in_context(span, self._scope_base_context()),
+            # Built after enter_scope's cleanup so it cannot inherit values
+            # from a scope that is about to be detached.
+            lambda: trace.set_span_in_context(span, self._scope_base_context()),
             epoch=self._epoch,
-            parent_key=info.parent_id,
         )
 
     def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
index 686fbb10..890ac49b 100644
--- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
+++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
@@ -620,9 +620,10 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
         context_scope.enter_scope(
             self,
             self._scope_key(info),
-            trace.set_span_in_context(span, self._scope_base_context()),
+            # Built after enter_scope's cleanup so it cannot inherit values
+            # from a scope that is about to be detached.
+            lambda: trace.set_span_in_context(span, self._scope_base_context()),
             epoch=self._epoch,
-            parent_key=info.parent_id,
         )
 
     def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
index 978af0e2..bf787a0b 100644
--- a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
+++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py
@@ -81,7 +81,7 @@ def test_enter_then_exit_restores_the_previous_context():
     owner = _Owner()
     before = otel_context.get_current()
 
-    context_scope.enter_scope(owner, "a", _span_context("a"))
+    context_scope.enter_scope(owner, "a", lambda: _span_context("a"))
     assert otel_context.get_current() is not before
     assert context_scope.depth(owner) == 1
 
@@ -94,11 +94,9 @@ def test_nested_scopes_pop_in_lifo_order():
     owner = _Owner()
     before = otel_context.get_current()
 
-    context_scope.enter_scope(owner, "outer", _span_context("outer"))
+    context_scope.enter_scope(owner, "outer", lambda: _span_context("outer"))
     outer = otel_context.get_current()
-    context_scope.enter_scope(
-        owner, "inner", _span_context("inner"), parent_key="outer"
-    )
+    context_scope.enter_scope(owner, "inner", lambda: _span_context("inner"))
 
     context_scope.exit_scope(owner, "inner")
     assert otel_context.get_current() is outer
@@ -116,10 +114,8 @@ def test_exit_unwinds_scopes_stacked_above_the_target():
     owner = _Owner()
     before = otel_context.get_current()
 
-    context_scope.enter_scope(owner, "outer", _span_context("outer"))
-    context_scope.enter_scope(
-        owner, "orphan", _span_context("orphan"), parent_key="outer"
-    )
+    context_scope.enter_scope(owner, "outer", lambda: _span_context("outer"))
+    context_scope.enter_scope(owner, "orphan", lambda: _span_context("orphan"))
 
     context_scope.exit_scope(owner, "outer")
 
@@ -147,8 +143,8 @@ def test_reentering_the_same_key_replaces_the_previous_scope():
     owner = _Owner()
     before = otel_context.get_current()
 
-    context_scope.enter_scope(owner, "wfc-1", _span_context("poll-1"), epoch=1)
-    context_scope.enter_scope(owner, "wfc-1", _span_context("poll-2"), epoch=1)
+    context_scope.enter_scope(owner, "wfc-1", lambda: _span_context("poll-1"), epoch=1)
+    context_scope.enter_scope(owner, "wfc-1", lambda: _span_context("poll-2"), epoch=1)
 
     assert context_scope.depth(owner) == 1
 
@@ -157,58 +153,14 @@ def test_reentering_the_same_key_replaces_the_previous_scope():
     assert context_scope.depth(owner) == 0
 
 
-def test_a_suspended_sibling_scope_is_dropped_not_nested_into():
-    """A worker that moves to another branch must not inherit the old one's scope.
-
-    Branch-pool workers have no branch affinity. Branch A suspends without an end
-    hook; the same worker then runs branch B. B's scope must replace A's, not
-    stack on it -- otherwise exiting B would detach back into A and correlate
-    later records to the wrong branch.
-    """
-    owner = _Owner()
-    before = otel_context.get_current()
-
-    # Branch A runs and suspends: entered, never exited.
-    context_scope.enter_scope(owner, "branch-a", _span_context("branch-a"), epoch=1)
-
-    # The worker picks up branch B, a sibling: same epoch, different key, no parent.
-    context_scope.enter_scope(owner, "branch-b", _span_context("branch-b"), epoch=1)
-
-    assert context_scope.depth(owner) == 1
-
-    context_scope.exit_scope(owner, "branch-b")
-    assert otel_context.get_current() is before
-    assert context_scope.depth(owner) == 0
-
-
-def test_a_scope_whose_parent_never_ran_here_replaces_the_stale_one():
-    """A nested scope whose parent is absent cannot nest inside what is here."""
-    owner = _Owner()
-    before = otel_context.get_current()
-
-    context_scope.enter_scope(owner, "ctx-a", _span_context("ctx-a"), epoch=1)
-    # An operation nested under a context that ran on another thread.
-    context_scope.enter_scope(
-        owner, "step-b", _span_context("step-b"), epoch=1, parent_key="ctx-b"
-    )
-
-    assert context_scope.depth(owner) == 1
-    context_scope.exit_scope(owner, "step-b")
-    assert otel_context.get_current() is before
-
-
 def test_reentry_guard_keeps_enclosing_scopes():
     """Re-entering a nested key must not disturb the scope it is nested in."""
     owner = _Owner()
-    context_scope.enter_scope(owner, "ctx", _span_context("ctx"), epoch=1)
+    context_scope.enter_scope(owner, "ctx", lambda: _span_context("ctx"), epoch=1)
     enclosing = otel_context.get_current()
 
-    context_scope.enter_scope(
-        owner, "inner", _span_context("inner-1"), epoch=1, parent_key="ctx"
-    )
-    context_scope.enter_scope(
-        owner, "inner", _span_context("inner-2"), epoch=1, parent_key="ctx"
-    )
+    context_scope.enter_scope(owner, "inner", lambda: _span_context("inner-1"), epoch=1)
+    context_scope.enter_scope(owner, "inner", lambda: _span_context("inner-2"), epoch=1)
 
     assert context_scope.depth(owner) == 2
     context_scope.exit_scope(owner, "inner")
@@ -217,30 +169,6 @@ def test_reentry_guard_keeps_enclosing_scopes():
     context_scope.exit_scope(owner, "ctx")
 
 
-def test_legitimate_nesting_leaves_another_owners_scope_alone():
-    """Two plugins tracking the same operations must not evict each other.
-
-    The normal nesting path has to be a no-op, because detaching necessarily
-    discards everything stacked above the cut.
-    """
-    first, second = _Owner(), _Owner()
-    before = otel_context.get_current()
-
-    context_scope.enter_scope(first, "ctx", _span_context("first-ctx"), epoch=1)
-    context_scope.enter_scope(second, "ctx", _span_context("second-ctx"), epoch=1)
-    context_scope.enter_scope(
-        first, "step", _span_context("first-step"), epoch=1, parent_key="ctx"
-    )
-
-    # The second plugin's context scope survived the first plugin's nesting.
-    assert context_scope.depth(second) == 1
-
-    context_scope.exit_scope(first, "step")
-    context_scope.exit_scope(second, "ctx")
-    context_scope.exit_scope(first, "ctx")
-    assert otel_context.get_current() is before
-
-
 def test_enter_discards_scopes_from_a_previous_epoch():
     """A scope a suspended operation left behind must not outlive its invocation.
 
@@ -251,10 +179,12 @@ def test_enter_discards_scopes_from_a_previous_epoch():
     owner = _Owner()
     before = otel_context.get_current()
 
-    context_scope.enter_scope(owner, "suspended", _span_context("suspended"), epoch=1)
+    context_scope.enter_scope(
+        owner, "suspended", lambda: _span_context("suspended"), epoch=1
+    )
     assert context_scope.depth(owner) == 1
 
-    context_scope.enter_scope(owner, "next", _span_context("next"), epoch=2)
+    context_scope.enter_scope(owner, "next", lambda: _span_context("next"), epoch=2)
     assert context_scope.depth(owner) == 1
 
     context_scope.exit_scope(owner, "next")
@@ -266,8 +196,8 @@ def test_unwind_detaches_every_scope_for_one_owner():
     owner = _Owner()
     before = otel_context.get_current()
 
-    context_scope.enter_scope(owner, "a", _span_context("a"))
-    context_scope.enter_scope(owner, "b", _span_context("b"))
+    context_scope.enter_scope(owner, "a", lambda: _span_context("a"))
+    context_scope.enter_scope(owner, "b", lambda: _span_context("b"))
 
     context_scope.unwind(owner)
 
@@ -294,7 +224,7 @@ def test_scopes_are_confined_to_the_thread_that_attached_them():
     observed: dict[str, object] = {}
 
     def worker() -> None:
-        context_scope.enter_scope(owner, "worker", _span_context("worker"))
+        context_scope.enter_scope(owner, "worker", lambda: _span_context("worker"))
         observed["worker_depth"] = context_scope.depth(owner)
         observed["worker_span_valid"] = (
             trace.get_current_span().get_span_context().is_valid
@@ -405,6 +335,41 @@ def _step_end(operation_id: str) -> UserFunctionEndInfo:
     )
 
 
+def _context_start(operation_id: str, parent_id: str | None) -> UserFunctionStartInfo:
+    """Start info for a child-context (including virtual/FLAT branch) body."""
+    return UserFunctionStartInfo(
+        operation_id=operation_id,
+        operation_type=OperationType.CONTEXT,
+        sub_type=None,
+        name=operation_id,
+        parent_id=parent_id,
+        start_time=START_TIME,
+        is_replayed=False,
+        status=OperationStatus.STARTED,
+        is_replay_children=False,
+        attempt=None,
+    )
+
+
+def _context_end(operation_id: str, parent_id: str | None) -> UserFunctionEndInfo:
+    """End info for a child-context body."""
+    return UserFunctionEndInfo(
+        operation_id=operation_id,
+        operation_type=OperationType.CONTEXT,
+        sub_type=None,
+        name=operation_id,
+        parent_id=parent_id,
+        start_time=START_TIME,
+        end_time=END_TIME,
+        is_replayed=False,
+        status=OperationStatus.SUCCEEDED,
+        is_replay_children=False,
+        attempt=None,
+        outcome=UserFunctionOutcome.SUCCEEDED,
+        error=None,
+    )
+
+
 @pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
 def test_invocation_end_unwinds_a_suspended_operation_scope(factory):
     """A step that suspends never gets its end hook; invocation end cleans up.
@@ -455,46 +420,90 @@ def run_step() -> None:
     plugin.on_invocation_end(_invocation_end())
 
 
-@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
-def test_worker_reassigned_to_another_branch_drops_the_suspended_scope(factory):
-    """One worker running branch A (suspends) then branch B leaves no residue.
+def test_the_new_context_is_built_after_stale_scopes_are_dropped():
+    """The attached context must not inherit a scope that is being detached.
 
-    Branch pools reuse workers with no branch affinity, so the worker that ran a
-    suspended branch goes on to another. B must not nest inside A's abandoned
-    scope, and after B ends the thread must be clean.
+    The context is normally derived from what is current, so building it before
+    cleanup would copy baggage or suppression values out of the stale scope, and
+    detaching afterwards cannot remove them from an already-built Context.
     """
-    plugin, _ = factory()
+    owner = _Owner()
     before = otel_context.get_current()
+
+    # A stale scope under the same key, carrying baggage.
+    context_scope.enter_scope(
+        owner,
+        "op-1",
+        lambda: baggage.set_baggage(
+            "poll", "first", context=otel_context.get_current()
+        ),
+        epoch=1,
+    )
+    assert baggage.get_baggage("poll") == "first"
+
+    # Re-entering builds its context from whatever is current at that moment,
+    # which must already be the pre-stale context.
+    observed: dict[str, object] = {}
+
+    def build() -> Context:
+        observed["seen_during_build"] = baggage.get_baggage("poll")
+        return otel_context.get_current()
+
+    context_scope.enter_scope(owner, "op-1", build, epoch=1)
+
+    assert observed["seen_during_build"] is None
+    assert baggage.get_baggage("poll") is None
+
+    context_scope.exit_scope(owner, "op-1")
+    assert otel_context.get_current() is before
+
+
+@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
+def test_flat_branch_scope_survives_its_inner_operations(factory):
+    """A FLAT map/parallel branch scope must stay attached across inner steps.
+
+    A virtual (FLAT) branch reports its inner operations' parent as the
+    grandparent -- None for a top-level branch -- while the branch's own context
+    scope is still running (see DurableContext.is_virtual). Physical nesting is
+    therefore not derivable from parent_id, and treating an inner step as
+    root-level must not detach the live branch scope: work between two inner steps
+    would fall out of the durable trace.
+    """
+    plugin, _ = factory()
     plugin.on_invocation_start(_invocation_start())
     observed: dict[str, object] = {}
 
-    def branch_a() -> None:
-        # Suspends: start fires, end never does.
-        plugin.on_user_function_start(_step_start("branch-a"))
-
-    def branch_b() -> None:
-        plugin.on_user_function_start(_step_start("branch-b"))
-        observed["depth_in_b"] = context_scope.depth(plugin)
-        observed["current_is_b"] = trace.get_current_span().get_span_context().span_id
-        plugin.on_user_function_end(_step_end("branch-b"))
-        observed["depth_after_b"] = context_scope.depth(plugin)
-        observed["span_after_b_valid"] = (
-            trace.get_current_span().get_span_context().is_valid
+    def run_branch() -> None:
+        # The virtual branch body: a CONTEXT operation at the top level.
+        plugin.on_user_function_start(
+            _context_start("flat-branch", parent_id=None),
         )
+        branch_span_id = trace.get_current_span().get_span_context().span_id
+
+        # Inner steps of a FLAT branch report parent_id=None, not the branch.
+        for index in range(2):
+            step_id = f"flat-branch-step-{index}"
+            plugin.on_user_function_start(_step_start(step_id))
+            plugin.on_user_function_end(_step_end(step_id))
+            # Between inner steps the branch scope is still current.
+            observed[f"between-{index}"] = (
+                trace.get_current_span().get_span_context().span_id
+            )
+
+        observed["branch"] = branch_span_id
+        plugin.on_user_function_end(
+            _context_end("flat-branch", parent_id=None),
+        )
+        observed["after_branch_depth"] = context_scope.depth(plugin)
 
-    # A single worker, so branch B is guaranteed to land on branch A's thread.
     with ThreadPoolExecutor(max_workers=1) as executor:
-        executor.submit(branch_a).result()
-        executor.submit(branch_b).result()
-
-    # B replaced A's scope rather than stacking on it.
-    assert observed["depth_in_b"] == 1
-    # Exiting B did not fall back into A.
-    assert observed["depth_after_b"] == 0
-    assert observed["span_after_b_valid"] is False
-    assert otel_context.get_current() is before
+        executor.submit(run_branch).result()
 
-    plugin.on_invocation_end(_invocation_end(InvocationStatus.PENDING))
+    assert observed["between-0"] == observed["branch"]
+    assert observed["between-1"] == observed["branch"]
+    assert observed["after_branch_depth"] == 0
+
+    plugin.on_invocation_end(_invocation_end())
 
 
 @pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])