Python: fix(workflow): keep a sibling pause when a nested step failure was handled

python/packages/core/agent_framework/_workflows/_functional.py

@@ -1008,6 +1008,51 @@ def _pause_boundary_was_handled(self, task: asyncio.Future[Any]) -> bool:
             return False
         return waiter.exception() is None
 
+    def _failure_was_handled_by_awaiter(self, task: asyncio.Future[Any]) -> bool:
+        """Return whether the code awaiting this task's boundary caught the failure.
+
+        A step that runs as a task inside another step fails inside the inner
+        concurrency boundary, so the boundary's waiter keeps the exception even
+        after the awaiting code caught it. The waiter failing with that same
+        exception shows the failure reached the boundary, and that task
+        completing without raising it shows the code handled it.
+
+        The failure also has to have been delivered: a future's exception is
+        retrieved exactly when the awaiting code receives it, and retrieving it
+        clears the future's traceback flag. A waiter that still carries the
+        flag means nobody received the failure, because the awaiting code was
+        cancelled before it could resume, or because the boundary was never
+        awaited at all.
+        """
+        waiter = self._pause_boundary_waiters.get(task)
+        if waiter is None or not waiter.done() or waiter.cancelled():
+            return False
+        error = task.exception()
+        if error is None:
+            return False
+        if getattr(waiter, "_log_traceback", False):
+            # Nobody retrieved the waiter's exception, so it was never delivered.
+            return False
+        try:
+            delivered = waiter.exception()
+        except Exception:
+            return False
+        if delivered is not error:
+            # The waiter failed for another reason, a pause interruption for
+            # example, so this failure never reached the awaiting code.
+            return False
+        awaiter = self._workflow_task_parents.get(task)
+        if awaiter is None or not awaiter.done() or awaiter.cancelled():
+            return False
+        cancelling = getattr(awaiter, "cancelling", None)
+        cancelling_count = cast(Callable[[], int], cancelling)() if callable(cancelling) else 0
+        if cancelling_count > 0:
+            # The awaiting task was cancelled, so it may never have received
+            # the failure. Cancelling a waiter clears its traceback flag, so the
+            # flag alone cannot rule that out.
+            return False
+        return awaiter.exception() is None
+
     def _pause_root_waiter_was_handled(self) -> bool:
         waiter = self._pause_root_waiter
         if waiter is None or not waiter.done() or waiter.cancelled():
@@ -1152,6 +1197,11 @@ def _raise_completed_failure() -> None:
                             continue
                         raise
                 if task in self._pause_failure_boundaries:
+                    # The delivery check reads the waiter's traceback flag, so it
+                    # has to run before anything retrieves the waiter's exception.
+                    if self._failure_was_handled_by_awaiter(task):
+                        inspected_tasks.add(task)
+                        continue
                     if self._pause_boundary_was_handled(task):
                         inspected_tasks.add(task)
                         continue

python/packages/core/tests/workflow/test_functional_workflow.py

@@ -8,6 +8,7 @@
 import gc
 import json
 import logging
+import sys
 from collections.abc import Awaitable, Callable, Generator
 from contextlib import contextmanager, suppress
 from dataclasses import dataclass
@@ -99,6 +100,65 @@ async def failing_step(x: int) -> int:
     raise ValueError(f"step failed with {x}")
 
 
+class NestedStepError(Exception):
+    """Raised by a nested step whose failure the workflow's own code handles."""
+
+
+@step
+async def nested_boom_step() -> int:
+    raise NestedStepError("boom")
+
+
+@step
+async def gather_failing_step() -> int:
+    (value,) = await asyncio.gather(nested_boom_step())
+    return value
+
+
+@step
+async def ask_for_review_step(name: str, ctx: RunContext) -> int:
+    return await ctx.request_info(f"{name}?", int, request_id=name)
+
+
+@step
+async def detached_gather_step(ctx: RunContext) -> str:
+    """Start a gather with a failing child and a pausing child, and never await it."""
+
+    async def failing() -> int:
+        raise NestedStepError("boom")
+
+    async def pausing() -> int:
+        return await ctx.request_info("review?", int, request_id="review")
+
+    asyncio.gather(failing(), pausing())
+    return "detached"
+
+
+@step
+async def cancelled_awaiter_step(ctx: RunContext) -> int:
+    """Get cancelled after the gather finished but before it can deliver its failure."""
+
+    async def failing() -> int:
+        raise NestedStepError("boom")
+
+    async def pausing() -> int:
+        return await ctx.request_info("review?", int, request_id="review")
+
+    async def helper() -> int:
+        task = asyncio.current_task()
+        gather = asyncio.gather(failing(), pausing())
+        if task is not None:
+            # Registered before the await, so it runs before the task resumes with
+            # the gather and the gather's exception is never retrieved.
+            gather.add_done_callback(lambda _: task.cancel())
+        try:
+            return await gather
+        except asyncio.CancelledError:
+            return -1
+
+    return await asyncio.create_task(helper())
+
+
 # ---------------------------------------------------------------------------
 # Basic execution
 # ---------------------------------------------------------------------------
@@ -456,6 +516,64 @@ async def auto_id_wf(x: int, ctx: RunContext) -> None:
         assert len(events) == 1
         assert events[0].request_id  # should be a non-empty uuid string
 
+    async def test_a_gather_nobody_awaited_still_surfaces_its_failure(self):
+        """A failure nobody received must not be taken for a handled one.
+
+        The gather that holds the failing child is never awaited, so nothing
+        ever received the exception, and the task that created it finishes
+        normally. The workflow still has to report the failure.
+        """
+
+        @built_workflow
+        async def flow(x: int) -> str:
+            return await detached_gather_step()
+
+        with pytest.raises(NestedStepError):
+            await flow.run(0)
+
+    @pytest.mark.skipif(sys.version_info < (3, 11), reason="needs Task.cancelling()")
+    async def test_a_cancellation_before_the_failure_is_delivered_still_surfaces_it(self):
+        """A cancellation that interrupts the delivery must not hide the failure.
+
+        The awaiting task is cancelled after the gather finished but before it
+        resumes, so it never receives the gather's exception and only sees the
+        cancellation. The failure was never handled and has to surface.
+        """
+
+        @built_workflow
+        async def flow(x: int) -> int:
+            return await cancelled_awaiter_step()
+
+        with pytest.raises(NestedStepError):
+            await flow.run(0)
+
+    async def test_nested_handled_step_failure_does_not_replace_sibling_pause(self):
+        """A handled step failure must not replace a sibling branch's pending request.
+
+        A step that runs as a task inside another step fails inside the inner
+        concurrency boundary, so the direct waiter keeps carrying the exception
+        even after the surrounding branch caught it. That branch then completes
+        normally, so the workflow must keep the pending request instead of
+        re-raising the failure the workflow's own code already handled.
+        """
+
+        async def left() -> int:
+            try:
+                return await gather_failing_step()
+            except NestedStepError:
+                return -1
+
+        @built_workflow
+        async def flow(x: int) -> list:
+            return list(await asyncio.gather(left(), ask_for_review_step("review")))
+
+        result = await flow.run(0)
+        pending = result.get_request_info_events()
+        assert [event.request_id for event in pending] == ["review"]
+
+        result = await flow.run(responses={event.request_id: 1 for event in pending})
+        assert result.get_outputs() == [[-1, 1]]
+
 
 # ---------------------------------------------------------------------------
 # Error handling