fix(otel): release unreleased scope on operation re-entry

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

@@ -174,8 +174,24 @@ def _attempt_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
     # Context scope helpers
     # ------------------------------------------------------------------
     def _attach_context(self, key: str, new_context: Context) -> None:
-        """Attach a context and remember its token under ``key``."""
+        """Attach a context and remember its token under ``key``.
+
+        A token already stored under ``key`` means the previous scope for that
+        operation was never released: its user function did not reach
+        ``on_user_function_end``, because it suspended and a timed in-process
+        resume re-entered the same operation. Release it before attaching, so
+        the new token does not bury one that can never be detached -- otherwise
+        releasing the new scope would restore the abandoned span and leave it
+        current for later work on this thread.
+        """
         with self._lock:
+            if key in self._context_tokens:
+                logger.debug(
+                    "Releasing an unreleased context scope for %s before "
+                    "re-attaching; its user function did not report an end.",
+                    key,
+                )
+                self._detach_context(key)
             self._context_tokens[key] = (
                 threading.get_ident(),
                 otel_context.attach(new_context),

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

@@ -176,8 +176,24 @@ def _attempt_span_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
     # Context scope helpers
     # ------------------------------------------------------------------
     def _attach_context(self, key: str, new_context: Context) -> None:
-        """Attach a context and remember its token under ``key``."""
+        """Attach a context and remember its token under ``key``.
+
+        A token already stored under ``key`` means the previous scope for that
+        operation was never released: its user function did not reach
+        ``on_user_function_end``, because it suspended and a timed in-process
+        resume re-entered the same operation. Release it before attaching, so
+        the new token does not bury one that can never be detached -- otherwise
+        releasing the new scope would restore the abandoned span and leave it
+        current for later work on this thread.
+        """
         with self._operation_spans_lock:
+            if key in self._context_tokens:
+                logger.debug(
+                    "Releasing an unreleased context scope for %s before "
+                    "re-attaching; its user function did not report an end.",
+                    key,
+                )
+                self._detach_context(key)
             self._context_tokens[key] = (
                 threading.get_ident(),
                 context.attach(new_context),

packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import threading
 from concurrent.futures import ThreadPoolExecutor
 from datetime import UTC, datetime
 from types import SimpleNamespace
@@ -594,13 +595,15 @@ def _step_end_info(
     )
 
 
-def _context_start_info(operation_id: str) -> UserFunctionStartInfo:
+def _context_start_info(
+    operation_id: str, parent_id: str | None = None
+) -> UserFunctionStartInfo:
     return UserFunctionStartInfo(
         operation_id=operation_id,
         operation_type=OperationType.CONTEXT,
         sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
         name=operation_id,
-        parent_id=None,
+        parent_id=parent_id,
         start_time=START_TIME,
         is_replayed=False,
         status=OperationStatus.STARTED,
@@ -609,13 +612,15 @@ def _context_start_info(operation_id: str) -> UserFunctionStartInfo:
     )
 
 
-def _context_end_info(operation_id: str) -> UserFunctionEndInfo:
+def _context_end_info(
+    operation_id: str, parent_id: str | None = None
+) -> UserFunctionEndInfo:
     return UserFunctionEndInfo(
         operation_id=operation_id,
         operation_type=OperationType.CONTEXT,
         sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
         name=operation_id,
-        parent_id=None,
+        parent_id=parent_id,
         start_time=START_TIME,
         is_replayed=False,
         status=OperationStatus.STARTED,
@@ -803,3 +808,164 @@ def test_detach_ignores_token_attached_on_another_thread():
     assert otel_context.get_current() == before_context
 
     plugin.on_invocation_end(_invocation_end_info())
+
+
+# ---------------------------------------------------------------------------
+# Re-entering an operation whose scope was never released
+# ---------------------------------------------------------------------------
+def test_reentered_child_context_does_not_leave_abandoned_span_current():
+    """Verify a timed in-process resume unwinds the abandoned scope.
+
+    A suspended child context never reaches on_user_function_end, and the
+    map/parallel coordinator can resume that branch in-process, re-entering the
+    same operation ID on the same thread. Without releasing the first scope, the
+    second scope's detach would restore the abandoned span and leave it current.
+    """
+    plugin, _ = _create_plugin()
+    plugin.on_invocation_start(_invocation_start_info())
+    before_context = otel_context.get_current()
+    context_id = "ctx-1"
+
+    # First run: the child context suspends, so no end hook fires.
+    plugin.on_user_function_start(_context_start_info(context_id))
+    suspended_span = plugin._get_span(context_id)
+    assert suspended_span is not None
+
+    # Timed in-process resume re-enters the same operation.
+    plugin.on_user_function_start(_context_start_info(context_id))
+    assert len([key for key in plugin._context_tokens if key == context_id]) == 1
+
+    plugin.on_user_function_end(_context_end_info(context_id))
+
+    assert otel_context.get_current() == before_context
+    assert context_id not in plugin._context_tokens
+    assert (
+        trace.get_current_span().get_span_context().span_id
+        != suspended_span.get_span_context().span_id
+    )
+
+    plugin.on_invocation_end(_invocation_end_info())
+
+
+def test_reentered_step_attempt_releases_the_previous_scope():
+    """Verify re-entering the same attempt key unwinds the previous scope."""
+    plugin, _ = _create_plugin()
+    plugin.on_invocation_start(_invocation_start_info())
+    before_context = otel_context.get_current()
+
+    plugin.on_user_function_start(_step_start_info("step-1"))
+    plugin.on_user_function_start(_step_start_info("step-1"))
+    plugin.on_user_function_end(_step_end_info("step-1"))
+
+    assert otel_context.get_current() == before_context
+    assert set(plugin._context_tokens) == {"__invocation_context__"}
+
+    plugin.on_invocation_end(_invocation_end_info())
+    assert plugin._context_tokens == {}
+
+
+def test_reentry_on_another_thread_leaves_the_originating_worker_dirty():
+    """Pin what re-entry can and cannot clean up across threads.
+
+    A resumed branch can land on a different pool thread than the one that
+    suspended. Re-entry drops the foreign token instead of resetting it, because
+    a context token can only be reset on its own thread, and it unwinds cleanly
+    on the thread that re-entered. The worker that suspended keeps the abandoned
+    span current: releasing it needs a hook invoked on that thread when the user
+    function fails to complete, which the SDK does not provide. The worker is
+    kept alive here so this limitation is asserted rather than hidden by pool
+    shutdown; the assertion flips once such a hook exists.
+    """
+    plugin, _ = _create_plugin()
+    plugin.on_invocation_start(_invocation_start_info())
+    before_context = otel_context.get_current()
+    span_key = "step-1:attempt:1"
+
+    with ThreadPoolExecutor(max_workers=1) as worker:
+        # The suspending run happens on the worker and never reports an end.
+        worker.submit(
+            plugin.on_user_function_start, _step_start_info("step-1")
+        ).result()
+        abandoned_span = plugin._get_span(span_key)
+        assert abandoned_span is not None
+        foreign_thread_ident, _foreign_token = plugin._context_tokens[span_key]
+        assert foreign_thread_ident != threading.get_ident()
+
+        # The timed resume lands on this thread instead.
+        plugin.on_user_function_start(_step_start_info("step-1"))
+        assert plugin._context_tokens[span_key][0] == threading.get_ident()
+        plugin.on_user_function_end(_step_end_info("step-1"))
+
+        # This thread unwound to where it started.
+        assert otel_context.get_current() == before_context
+
+        # The originating worker is still carrying the abandoned span.
+        worker_span_id = worker.submit(
+            lambda: trace.get_current_span().get_span_context().span_id
+        ).result()
+        assert worker_span_id == abandoned_span.get_span_context().span_id
+
+    plugin.on_invocation_end(_invocation_end_info())
+
+
+def test_nested_reentry_restores_the_abandoned_outer_scope():
+    """Pin nested re-entry: correct ids, but the abandoned outer span object.
+
+    When an outer child context and an inner one both suspend, re-entry releases
+    each scope in the order the operations are replayed, which is not the reverse
+    of the order they were attached. Ending the inner operation therefore
+    restores the scope captured for the abandoned outer span rather than the
+    resumed one. Deterministic CONTEXT span ids make the two indistinguishable
+    downstream -- same trace id and span id, so parenting and log correlation are
+    unaffected -- but the current span object is one that is never exported, so
+    anything an instrumentation library records on it is lost. Reverse-order
+    unwinding needs the SDK to report the suspension; this test documents the
+    current behaviour and flips when that lands.
+    """
+    plugin, _ = _create_plugin()
+    plugin.on_invocation_start(_invocation_start_info())
+    before_context = otel_context.get_current()
+
+    # Both contexts suspend, so neither reports an end.
+    plugin.on_user_function_start(_context_start_info("ctx-outer"))
+    abandoned_outer = plugin._get_span("ctx-outer")
+    plugin.on_user_function_start(
+        _context_start_info("ctx-inner", parent_id="ctx-outer")
+    )
+    assert abandoned_outer is not None
+
+    # The timed in-process resume replays both contexts, outer first.
+    plugin.on_user_function_start(_context_start_info("ctx-outer"))
+    resumed_outer = plugin._get_span("ctx-outer")
+    plugin.on_user_function_start(
+        _context_start_info("ctx-inner", parent_id="ctx-outer")
+    )
+    resumed_inner = plugin._get_span("ctx-inner")
+    assert resumed_outer is not None
+    assert resumed_inner is not None
+    assert resumed_outer is not abandoned_outer
+
+    # Resumed inner code runs under the resumed inner span.
+    assert trace.get_current_span() is resumed_inner
+
+    plugin.on_user_function_end(_context_end_info("ctx-inner", parent_id="ctx-outer"))
+
+    # The restored scope carries the abandoned outer span, whose ids match the
+    # resumed one because CONTEXT span ids are derived from the operation id.
+    assert trace.get_current_span() is abandoned_outer
+    assert (
+        abandoned_outer.get_span_context().span_id
+        == resumed_outer.get_span_context().span_id
+    )
+    assert (
+        abandoned_outer.get_span_context().trace_id
+        == resumed_outer.get_span_context().trace_id
+    )
+
+    # Leaving the outer context still unwinds to where the invocation started.
+    plugin.on_user_function_end(_context_end_info("ctx-outer"))
+    assert otel_context.get_current() == before_context
+    assert set(plugin._context_tokens) == {"__invocation_context__"}
+
+    plugin.on_invocation_end(_invocation_end_info())
+    assert plugin._context_tokens == {}