feat(plugin): report user functions that do not complete

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

@@ -40,6 +40,7 @@
     OperationStartInfo,
     OperationType,
     UserFunctionEndInfo,
+    UserFunctionIncompleteInfo,
     UserFunctionOutcome,
     UserFunctionStartInfo,
 )
@@ -167,9 +168,25 @@ def _pop_span(self, key: str) -> Span | None:
             return self._operation_spans.pop(key, None)
 
     @staticmethod
-    def _attempt_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
+    def _attempt_key(
+        info: UserFunctionStartInfo | UserFunctionEndInfo | UserFunctionIncompleteInfo,
+    ) -> str:
         return f"{info.operation_id}:attempt:{info.attempt or 1}"
 
+    @classmethod
+    def _user_function_key(
+        cls,
+        info: UserFunctionStartInfo | UserFunctionEndInfo | UserFunctionIncompleteInfo,
+    ) -> str:
+        """Return the registry key a user function's span and scope are stored under.
+
+        STEP user functions are attempts, so each attempt gets its own key; a
+        CONTEXT is entered once per invocation and uses the operation id.
+        """
+        if info.operation_type is OperationType.STEP:
+            return cls._attempt_key(info)
+        return info.operation_id
+
     # ------------------------------------------------------------------
     # Context scope helpers
     # ------------------------------------------------------------------
@@ -544,11 +561,7 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
             raise RuntimeError(
                 "on_user_function_end only supports CONTEXT and STEP operations"
             )
-        key = (
-            self._attempt_key(info)
-            if info.operation_type is OperationType.STEP
-            else info.operation_id
-        )
+        key = self._user_function_key(info)
         span = self._get_span(key)
         if span is None:
             raise RuntimeError(
@@ -583,6 +596,24 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
         # without stacking another scope.
         self._detach_context(key)
 
+    def on_user_function_incomplete(self, info: UserFunctionIncompleteInfo) -> None:
+        """Release the scope of a user function that reported no outcome.
+
+        A suspended, orphaned or interrupted user function never reaches
+        ``on_user_function_end``, so its scope is released here instead, on the
+        thread that attached it. Nested scopes unwind in reverse order because
+        this fires as the exception propagates outward: the inner operation is
+        released before its enclosing one.
+
+        The span is left open and registered. The operation is not finished --
+        a suspended one resumes and ends later -- so ending it here would export
+        a span for work that is still in flight.
+        """
+        logger.debug("Durable user function incomplete: %s", info)
+        if not self._tracing_enabled:
+            return
+        self._detach_context(self._user_function_key(info))
+
     # ------------------------------------------------------------------
     # Attributes
     # ------------------------------------------------------------------

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

@@ -16,6 +16,7 @@
     OperationStartInfo,
     OperationType,
     UserFunctionEndInfo,
+    UserFunctionIncompleteInfo,
     UserFunctionOutcome,
     UserFunctionStartInfo,
 )
@@ -168,10 +169,26 @@ def _get_span(self, operation_id: str | None) -> Span | None:
             return self._operation_spans.get(operation_id)
 
     @staticmethod
-    def _attempt_span_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str:
+    def _attempt_span_key(
+        info: UserFunctionStartInfo | UserFunctionEndInfo | UserFunctionIncompleteInfo,
+    ) -> str:
         """Return the registry key for a STEP attempt span."""
         return f"{info.operation_id}:attempt:{info.attempt or 1}"
 
+    @classmethod
+    def _user_function_span_key(
+        cls,
+        info: UserFunctionStartInfo | UserFunctionEndInfo | UserFunctionIncompleteInfo,
+    ) -> str:
+        """Return the registry key a user function's span and scope are stored under.
+
+        STEP user functions are attempts, so each attempt gets its own key; a
+        CONTEXT is entered once per invocation and uses the operation id.
+        """
+        if info.operation_type is OperationType.STEP:
+            return cls._attempt_span_key(info)
+        return info.operation_id
+
     # ------------------------------------------------------------------
     # Context scope helpers
     # ------------------------------------------------------------------
@@ -642,12 +659,7 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
             raise RuntimeError(
                 "on_user_function_end should only be called for CONTEXT and STEP operations"
             )
-        # key = f"{info.operation_id}-{int(info.start_time.timestamp())}"
-        span_key = (
-            self._attempt_span_key(info)
-            if info.operation_type is OperationType.STEP
-            else info.operation_id
-        )
+        span_key = self._user_function_span_key(info)
         span = self._get_span(span_key)
         if not span:
             raise RuntimeError(
@@ -682,6 +694,25 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
         # get_current_span_context falls back to the invocation span.
         self._detach_context(span_key)
 
+    def on_user_function_incomplete(self, info: UserFunctionIncompleteInfo) -> None:
+        """Release the scope of a user function that reported no outcome.
+
+        A suspended, orphaned or interrupted user function never reaches
+        ``on_user_function_end``, so its scope is released here instead, on the
+        thread that attached it. Nested scopes unwind in reverse order because
+        this fires as the exception propagates outward: the inner operation is
+        released before its enclosing one.
+
+        The span is left open and registered. The operation is not finished --
+        a suspended one resumes and ends later -- so ending it here would export
+        a span for work that is still in flight. Spans still open at invocation
+        end are closed by ``on_invocation_end``.
+        """
+        logger.debug("Durable user function incomplete: %s", info)
+        if not self._tracing_enabled:
+            return
+        self._detach_context(self._user_function_span_key(info))
+
     def _extract_attributes(self, info: Any) -> _SpanAttributes:
         """Extract durable execution fields as OpenTelemetry span attributes.
 

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

@@ -22,6 +22,7 @@
     OperationStartInfo,
     OperationType,
     UserFunctionEndInfo,
+    UserFunctionIncompleteInfo,
     UserFunctionOutcome,
     UserFunctionStartInfo,
 )
@@ -595,6 +596,42 @@ def _step_end_info(
     )
 
 
+def _step_incomplete_info(
+    operation_id: str,
+    parent_id: str | None = None,
+    attempt: int = 1,
+) -> UserFunctionIncompleteInfo:
+    return UserFunctionIncompleteInfo(
+        operation_id=operation_id,
+        operation_type=OperationType.STEP,
+        sub_type=OperationSubType.STEP,
+        name=operation_id,
+        parent_id=parent_id,
+        start_time=START_TIME,
+        is_replayed=False,
+        status=OperationStatus.STARTED,
+        is_replay_children=False,
+        attempt=attempt,
+    )
+
+
+def _context_incomplete_info(
+    operation_id: str, parent_id: str | None = None
+) -> UserFunctionIncompleteInfo:
+    return UserFunctionIncompleteInfo(
+        operation_id=operation_id,
+        operation_type=OperationType.CONTEXT,
+        sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
+        name=operation_id,
+        parent_id=parent_id,
+        start_time=START_TIME,
+        is_replayed=False,
+        status=OperationStatus.STARTED,
+        is_replay_children=False,
+        attempt=1,
+    )
+
+
 def _context_start_info(
     operation_id: str, parent_id: str | None = None
 ) -> UserFunctionStartInfo:
@@ -864,75 +901,85 @@ def test_reentered_step_attempt_releases_the_previous_scope():
     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.
+def test_suspension_releases_the_scope_on_the_originating_worker():
+    """Verify the suspending worker releases its own scope.
 
-    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.
+    A suspended user function reports no outcome, so the SDK fires
+    on_user_function_incomplete on the thread that ran it -- the only thread that
+    can reset its context token. The worker is kept alive and probed to prove it
+    is left clean even though the resume lands on a different thread.
     """
-    plugin, _ = _create_plugin()
+    plugin, exporter = _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.
+        def suspend_on_worker() -> tuple[int, bool]:
+            plugin.on_user_function_start(_step_start_info("step-1"))
+            attached_span_id = trace.get_current_span().get_span_context().span_id
+            plugin.on_user_function_incomplete(_step_incomplete_info("step-1"))
+            return (
+                attached_span_id,
+                trace.get_current_span().get_span_context().is_valid,
+            )
+
+        attached_span_id, span_still_current = worker.submit(suspend_on_worker).result()
+        suspended_span = plugin._get_span(span_key)
+
+        # The scope was released on the worker, and its span is left open.
+        assert attached_span_id != 0
+        assert span_still_current is False
+        assert span_key not in plugin._context_tokens
+        assert suspended_span is not None
+        assert not exporter.get_finished_spans()
+
+        # The timed resume lands on this thread, with nothing stale to unwind.
         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
+        # The originating worker is still clean.
+        worker_span_valid = worker.submit(
+            lambda: trace.get_current_span().get_span_context().is_valid
         ).result()
-        assert worker_span_id == abandoned_span.get_span_context().span_id
+        assert worker_span_valid is False
 
     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.
+def test_nested_suspension_unwinds_scopes_in_reverse_order():
+    """Verify nested suspends release inner-first and resume without stale scopes.
 
-    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.
+    The incomplete hook fires as the exception propagates outward, so the inner
+    context's scope is released before its enclosing one. On resume, ending the
+    inner operation restores the resumed outer scope rather than the one captured
+    for the suspended run.
     """
-    plugin, _ = _create_plugin()
+    plugin, exporter = _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")
+    suspended_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
+    assert suspended_outer is not None
+
+    # Both contexts suspend: the inner one unwinds first.
+    plugin.on_user_function_incomplete(
+        _context_incomplete_info("ctx-inner", parent_id="ctx-outer")
+    )
+    assert trace.get_current_span() is suspended_outer
+
+    plugin.on_user_function_incomplete(_context_incomplete_info("ctx-outer"))
+    assert otel_context.get_current() == before_context
+    assert set(plugin._context_tokens) == {"__invocation_context__"}
+    # Neither span is ended: both operations are still in flight.
+    assert not exporter.get_finished_spans()
 
     # The timed in-process resume replays both contexts, outer first.
     plugin.on_user_function_start(_context_start_info("ctx-outer"))
@@ -942,27 +989,14 @@ def test_nested_reentry_restores_the_abandoned_outer_scope():
     )
     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 resumed_outer is not suspended_outer
     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
-    )
+    # The resumed outer scope is restored, not the one from the suspended run.
+    assert trace.get_current_span() is resumed_outer
 
-    # 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__"}