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.

Author: Alex Wang

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:

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:

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: