feat(plugin): report incomplete user function outcomes (#661)

* feat(plugin): report incomplete user function outcomes


---------

Co-authored-by: Alex Wang <wangyb@amazon.com>

Author: wangyb-A

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

@@ -167,9 +167,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,
+    ) -> str:
         return f"{info.operation_id}:attempt:{info.attempt or 1}"
 
+    @classmethod
+    def _user_function_key(
+        cls,
+        info: UserFunctionStartInfo | UserFunctionEndInfo,
+    ) -> 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
     # ------------------------------------------------------------------
@@ -511,12 +527,12 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
             raise RuntimeError(
                 "on_user_function_start only supports CONTEXT and STEP operations"
             )
+        key = self._user_function_key(info)
         if info.operation_type is OperationType.STEP:
             parent = self._get_span(info.operation_id) or self._resolve_parent(
                 info.parent_id
             )
             name = f"{info.name or info.operation_id} attempt {info.attempt or 1}"
-            key = self._attempt_key(info)
             span = self._start_span(
                 operation_id=info.operation_id,
                 name=name,
@@ -528,7 +544,6 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
             )
         else:  # CONTEXT
             parent = self._resolve_parent(info.parent_id)
-            key = info.operation_id
             span = self._start_span(
                 operation_id=info.operation_id,
                 name=info.name or info.operation_id,
@@ -548,17 +563,16 @@ 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(
                 "on_user_function_end without matching on_user_function_start"
             )
-        if info.operation_type is OperationType.STEP:
+        if (
+            info.operation_type is OperationType.STEP
+            and info.outcome is not UserFunctionOutcome.INCOMPLETE
+        ):
             span.set_attributes(self._operation_attributes(info))
             if info.outcome is UserFunctionOutcome.FAILED:
                 span.set_status(

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

@@ -168,10 +168,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,
+    ) -> 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,
+    ) -> 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
     # ------------------------------------------------------------------
@@ -613,11 +629,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
         span_name = info.name or info.operation_id
         if info.operation_type is OperationType.STEP:
             span_name = f"{span_name} attempt {info.attempt or 1}"
-        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._start_span(
             operation_id=info.operation_id,
             name=span_name,
@@ -648,19 +660,17 @@ 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(
                 "on_user_function_end called without matching on_user_function_start"
             )
 
-        if info.operation_type is OperationType.STEP:
+        if (
+            info.operation_type is OperationType.STEP
+            and info.outcome is not UserFunctionOutcome.INCOMPLETE
+        ):
             span.set_attributes(self._extract_attributes(info))
             if info.outcome is UserFunctionOutcome.FAILED:
                 span.set_status(

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

@@ -613,6 +613,48 @@ def _step_end_info(
     )
 
 
+def _step_incomplete_info(
+    operation_id: str,
+    parent_id: str | None = None,
+    attempt: int = 1,
+) -> UserFunctionEndInfo:
+    return UserFunctionEndInfo(
+        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,
+        outcome=UserFunctionOutcome.INCOMPLETE,
+        end_time=END_TIME,
+        error=None,
+    )
+
+
+def _context_incomplete_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=parent_id,
+        start_time=START_TIME,
+        is_replayed=False,
+        status=OperationStatus.STARTED,
+        is_replay_children=False,
+        attempt=1,
+        outcome=UserFunctionOutcome.INCOMPLETE,
+        end_time=END_TIME,
+        error=None,
+    )
+
+
 def _context_start_info(
     operation_id: str, parent_id: str | None = None
 ) -> UserFunctionStartInfo:
@@ -882,75 +924,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_end with 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_end(_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 end callback 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_end(
+        _context_incomplete_info("ctx-inner", parent_id="ctx-outer")
+    )
+    assert trace.get_current_span() is suspended_outer
+
+    plugin.on_user_function_end(_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"))
@@ -960,27 +1012,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__"}