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.

Author: Alex Wang

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(

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,
         )
 

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,
         )
 

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.
 

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()