diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 719e21f0073..8e2dfe19e05 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/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 diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 59e4019b5df..8076224c412 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/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