diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 65dfefb933b..fdc053666e8 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -128,6 +128,18 @@ async def _tracked_coro() -> Any: task = asyncio.Task(_tracked_coro(), loop=loop, **kwargs) else: task = self.original_factory(loop, _tracked_coro(), **kwargs) + + def _close_unstarted_coro(completed: asyncio.Future[Any]) -> None: + # Cancellation can prevent the tracking wrapper from ever entering + # the coroutine it owns, so the wrapper cannot close it itself. + if ( + completed.cancelled() + and inspect.iscoroutine(coro) + and inspect.getcoroutinestate(coro) == inspect.CORO_CREATED + ): + coro.close() + + task.add_done_callback(_close_unstarted_coro) _register(task) return task @@ -696,6 +708,12 @@ def __init__( self._workflow_tasks: set[asyncio.Future[Any]] = set() # Tasks explicitly cancelled by framework cleanup. self._framework_cancelled_tasks: set[asyncio.Future[Any]] = set() + # In-flight per-step checkpoint saves. Cleanup waits for these instead of + # cancelling them, so a completed step's checkpoint lands even when the run + # is torn down (e.g. a sibling branch fails). + self._checkpoint_save_tasks: set[asyncio.Future[str]] = set() + # Caller cancellation also applies to steps finishing during teardown. + self._caller_cancelled = False # Parent task for each tracked workflow task. self._workflow_task_parents: dict[asyncio.Future[Any], asyncio.Task[Any] | None] = {} # Tasks that have entered at least one @step wrapper. @@ -1065,13 +1083,44 @@ def _get_new_descendant_interruption( current = self._workflow_task_parents.get(current) return None + def _cancel_checkpoint_saves(self) -> None: + """Stop existing and future per-step saves after caller cancellation.""" + self._caller_cancelled = True + for save_task in list(self._checkpoint_save_tasks): + save_task.cancel() + async def _cancel_workflow_tasks(self) -> None: """Cancel and drain all tracked workflow tasks.""" + cancellation: asyncio.CancelledError | None = None + while True: + try: + await self._drain_workflow_tasks() + except asyncio.CancelledError as exc: + # Caller cancellation can arrive after sibling-failure cleanup + # has started. Keep it effective throughout the remaining drain. + cancellation = exc + self._cancel_checkpoint_saves() + else: + break + if cancellation is not None: + raise cancellation + + async def _drain_workflow_tasks(self) -> None: + """Drain tracked tasks, protecting saves during framework-owned cleanup.""" current_task = asyncio.current_task() while pending := [task for task in self._workflow_tasks if task is not current_task and not task.done()]: + # Storage adapters may create tasks of their own (including via + # gather). Protect the whole in-flight save subtree, not just its + # root, or cancelling a child still aborts the shielded save. + save_tasks: set[asyncio.Future[Any]] = set(self._checkpoint_save_tasks) + while True: + previous_count = len(save_tasks) + save_tasks.update(task for task, parent in self._workflow_task_parents.items() if parent in save_tasks) + if len(save_tasks) == previous_count: + break waiting_for_existing_cancellation = False for task in pending: - if task in self._framework_cancelled_tasks: + if task in self._framework_cancelled_tasks or task in save_tasks: continue if self._has_pending_cancellation(task): waiting_for_existing_cancellation = True @@ -1080,10 +1129,15 @@ async def _cancel_workflow_tasks(self) -> None: if waiting_for_existing_cancellation: await asyncio.sleep(0) continue - await asyncio.gather(*pending, return_exceptions=True) + # Checkpoint saves are awaited, not cancelled: a completed step's + # checkpoint must land before the run is torn down, whatever storage + # is configured. + # Observe caller cancellation before a step can suppress it while + # returning from its own cancellation cleanup. + await asyncio.shield(asyncio.gather(*pending, return_exceptions=True)) tasks = [task for task in self._workflow_tasks if task is not current_task] if tasks: - await asyncio.gather(*tasks, return_exceptions=True) + await asyncio.shield(asyncio.gather(*tasks, return_exceptions=True)) @staticmethod def _has_pending_cancellation(task: asyncio.Future[Any]) -> bool: @@ -1107,6 +1161,7 @@ async def _settle_workflow_tasks( current_task = asyncio.current_task() settled_tasks = self._pause_owned_tasks inspected_tasks: set[asyncio.Future[Any]] = set() + cancellation_from_settled_task = False def _pending_owned_tasks() -> list[asyncio.Future[Any]]: self._expand_pause_owned_tasks() @@ -1120,6 +1175,7 @@ def _pending_owned_tasks() -> list[asyncio.Future[Any]]: return pending def _raise_completed_failure() -> None: + nonlocal cancellation_from_settled_task for task in settled_tasks: if task in inspected_tasks or not task.done(): continue @@ -1150,6 +1206,7 @@ def _raise_completed_failure() -> None: and self._get_new_descendant_interruption(cast(asyncio.Task[Any], task), {}) is interruption ): continue + cancellation_from_settled_task = True raise if task in self._pause_failure_boundaries: if self._pause_boundary_was_handled(task): @@ -1182,7 +1239,13 @@ def _raise_completed_failure() -> None: await self._cancel_workflow_tasks() if propagate_failures: _raise_completed_failure() - except BaseException: + except BaseException as exc: + if ( + isinstance(exc, asyncio.CancelledError) + and not isinstance(exc, WorkflowInterrupted) + and not cancellation_from_settled_task + ): + self._cancel_checkpoint_saves() await self._cancel_workflow_tasks() raise @@ -2150,7 +2213,36 @@ async def _run_core( if storage is not None: async def _on_step_completed() -> None: - ckpt_chain[0] = await self._save_checkpoint(ctx, storage, ckpt_chain[0]) + if ctx._caller_cancelled: + return + # Build the checkpoint before yielding to the event loop: a + # ready sibling could mutate ctx state before a task scheduled + # here gets to run, snapshotting state from work that did not + # complete. A save that has started must still finish if the + # step is cancelled because a sibling branch fails the run: a + # storage whose save awaits I/O before writing would otherwise + # lose the completed step's result, and a resumed run would + # re-run the step and its side effects. The save runs as its own + # task that cleanup waits for instead of cancelling; both awaits + # use asyncio.shield so a cancellation of this step task cannot + # reach the save through the await. + checkpoint = self._build_checkpoint(ctx, ckpt_chain[0]) + save_task = asyncio.ensure_future(storage.save(checkpoint)) + ctx._checkpoint_save_tasks.add(save_task) + save_task.add_done_callback(ctx._checkpoint_save_tasks.discard) + try: + ckpt_chain[0] = await asyncio.shield(save_task) + except asyncio.CancelledError: + try: + ckpt_chain[0] = await asyncio.shield(save_task) + except asyncio.CancelledError: + raise + except BaseException: + logger.warning( + "Checkpoint save failed while draining after cancellation", + exc_info=True, + ) + raise ctx._on_step_completed = _on_step_completed @@ -2324,8 +2416,8 @@ async def _invoke_workflow() -> Any: try: await asyncio.wait({workflow_task}) except BaseException: + ctx._cancel_checkpoint_saves() ctx._cancel_workflow_task(workflow_task) - await asyncio.gather(workflow_task, return_exceptions=True) await ctx._cancel_workflow_tasks() raise @@ -2376,24 +2468,31 @@ async def _invoke_workflow() -> Any: # Checkpoint helpers # ------------------------------------------------------------------ - async def _save_checkpoint( + def _build_checkpoint( self, ctx: RunContext, - storage: CheckpointStorage, previous_checkpoint_id: str | None = None, - ) -> str: - state = dict(ctx._state) + ) -> WorkflowCheckpoint: + state = _snapshot_workflow_state(ctx._state) state["_step_cache"] = ctx._export_step_cache() state["_step_cache_auto_request_info_counts"] = ctx._export_step_cache_auto_request_info_counts() state["_original_message"] = self._last_message - checkpoint = WorkflowCheckpoint( + return WorkflowCheckpoint( workflow_name=self.name, graph_signature_hash=self.graph_signature_hash, previous_checkpoint_id=previous_checkpoint_id, state=state, pending_request_info_events=dict(ctx._pending_requests), ) + + async def _save_checkpoint( + self, + ctx: RunContext, + storage: CheckpointStorage, + previous_checkpoint_id: str | None = None, + ) -> str: + checkpoint = self._build_checkpoint(ctx, previous_checkpoint_id) return await storage.save(checkpoint) def _compute_signature_hash( diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 8eeb8d28c1c..08b02ea84a8 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 +from collections import Counter from collections.abc import Awaitable, Callable, Generator from contextlib import contextmanager, suppress from dataclasses import dataclass @@ -18,6 +19,7 @@ from agent_framework import ( AgentResponseUpdate, + CheckpointID, CheckpointStorage, Content, ExperimentalFeature, @@ -28,6 +30,7 @@ RunContext, StepWrapper, SupportsAgentRun, + WorkflowCheckpoint, WorkflowEvent, WorkflowEventSource, WorkflowRunResult, @@ -930,6 +933,467 @@ async def wf(x: int) -> int: checkpoints = await storage.list_checkpoints(workflow_name="wf") assert len(checkpoints) == 3 # 2 from first run + 1 final from restore + @pytest.mark.parametrize("save_mode", ["inline", "task", "gather"]) + async def test_completed_parallel_step_is_checkpointed_when_sibling_fails(self, save_mode: str): + """A step that completed before a sibling failed must not run again + after resume, whatever the checkpoint storage's save timing.""" + runs: Counter[str] = Counter() + save_started = asyncio.Event() + + class TransientFailure(Exception): + pass + + class YieldingSaveStorage(InMemoryCheckpointStorage): + """A storage whose save awaits before writing, as a database-backed one does.""" + + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: # type: ignore[override] + async def write() -> CheckpointID: + if runs["done"]: + save_started.set() + await asyncio.sleep(0.01) + return await super(YieldingSaveStorage, self).save(checkpoint) + + if save_mode == "task": + return await asyncio.create_task(write()) + if save_mode == "gather": + return (await asyncio.gather(write()))[0] + return await write() + + storage = YieldingSaveStorage() + + @step(replay_key=lambda x: "prepare") + async def prepare(x: int) -> int: + runs["prepare"] += 1 + return x + + @step(replay_key=lambda x: "done") + async def done(x: int) -> int: + runs["done"] += 1 + return x + 1 + + @step(replay_key=lambda x: "flaky") + async def flaky(x: int) -> int: + runs["flaky"] += 1 + if runs["flaky"] == 1: + await save_started.wait() + raise TransientFailure("sibling failure") + return x + 2 + + @built_workflow(checkpoint_storage=storage) + async def parallel_wf(x: int) -> list[int]: + a = await prepare(x) + return list(await asyncio.gather(done(a), flaky(a))) + + with pytest.raises(TransientFailure): + await parallel_wf.run(0) + + latest = await storage.get_latest(workflow_name="parallel_wf") + assert latest is not None + result = await parallel_wf.run(checkpoint_id=latest.checkpoint_id) + + assert result.get_outputs() == [[1, 2]] + assert runs == Counter({"prepare": 1, "done": 1, "flaky": 2}) + + @pytest.mark.parametrize("mutation", ["top_level", "dict", "list"]) + async def test_step_checkpoint_snapshots_state_before_scheduling_save(self, mutation: str) -> None: + saved_states: list[dict[str, Any]] = [] + done_ready = asyncio.Event() + + class SiblingFailure(Exception): + pass + + class CapturingStorage(InMemoryCheckpointStorage): + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: # type: ignore[override] + saved_states.append(dict(checkpoint.state)) + return await super().save(checkpoint) + + storage = CapturingStorage() + + @step(replay_key=lambda x: "done") + async def done(x: int, ctx: RunContext) -> int: + ctx.set_state("done", True) + done_ready.set() + return x + 1 + + @step(replay_key=lambda x: "mutator") + async def mutator(x: int, ctx: RunContext) -> int: + await done_ready.wait() + if mutation == "dict": + ctx.get_state("nested")["details"]["unfinished"] = True + elif mutation == "list": + ctx.get_state("nested")["items"].append("unfinished") + else: + ctx.set_state("unfinished", True) + raise SiblingFailure("sibling failure") + + @built_workflow(checkpoint_storage=storage) + async def parallel_wf(x: int, ctx: RunContext) -> list[int]: + ctx.set_state("nested", {"details": {}, "items": []}) + return list(await asyncio.gather(done(x), mutator(x))) + + with pytest.raises(SiblingFailure): + await parallel_wf.run(0) + + done_checkpoints = [state for state in saved_states if state.get("done")] + assert done_checkpoints + assert all("unfinished" not in state for state in done_checkpoints) + assert all(state["nested"] == {"details": {}, "items": []} for state in done_checkpoints) + persisted = await storage.list_checkpoints(workflow_name="parallel_wf") + assert persisted + assert all(checkpoint.state["nested"] == {"details": {}, "items": []} for checkpoint in persisted) + + @pytest.mark.parametrize("cancel_mode", ["cancel", "timeout"]) + @pytest.mark.parametrize("phase", ["step", "failure", "pause"]) + async def test_caller_cancellation_skips_checkpoint_save_from_step_finishing_during_teardown( + self, cancel_mode: str, phase: str + ): + step_started = asyncio.Event() + step_finished = asyncio.Event() + workflow_finished = asyncio.Event() + release_step = asyncio.Event() + release_save = asyncio.Event() + + class BlockingSaveStorage(InMemoryCheckpointStorage): + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + await release_save.wait() + return await super().save(checkpoint) + + storage = BlockingSaveStorage() + + @step(replay_key=lambda: "complete") + async def complete() -> int: + step_started.set() + try: + await release_step.wait() + except asyncio.CancelledError: + # A step may finish cooperative cleanup and return a result. + step_finished.set() + return 1 + + @step(replay_key=lambda: "interrupt") + async def interrupt(ctx: RunContext) -> int: + await step_started.wait() + if phase == "failure": + raise RuntimeError("sibling failure") + return await ctx.request_info("continue?", int, request_id="continue") + + @built_workflow(checkpoint_storage=storage) + async def wf(_: int) -> int | list[int]: + try: + if phase == "step": + return await complete() + return list(await asyncio.gather(complete(), interrupt())) + finally: + workflow_finished.set() + + run_task = asyncio.ensure_future(wf.run(0)) + cancellation_task = run_task + expected_error: type[BaseException] + try: + await asyncio.wait_for(step_started.wait(), timeout=5) + if phase != "step": + await asyncio.wait_for(workflow_finished.wait(), timeout=5) + if cancel_mode == "timeout": + cancellation_task = asyncio.ensure_future(asyncio.wait_for(run_task, timeout=0.01)) + expected_error = asyncio.TimeoutError + else: + run_task.cancel() + expected_error = asyncio.CancelledError + + finished, _ = await asyncio.wait({cancellation_task}, timeout=1) + assert cancellation_task in finished, "A checkpoint save created during teardown blocked cancellation" + with pytest.raises(expected_error): + await cancellation_task + assert run_task.cancelled() + assert step_finished.is_set() + assert await storage.list_checkpoints(workflow_name="wf") == [] + finally: + release_step.set() + release_save.set() + await asyncio.gather(cancellation_task, run_task, return_exceptions=True) + + @pytest.mark.parametrize("save_mode", ["inline", "task", "gather"]) + @pytest.mark.parametrize("cancel_mode", ["cancel", "timeout"]) + @pytest.mark.parametrize("phase", ["step", "failure", "pause"]) + async def test_caller_can_cancel_in_flight_checkpoint_save(self, save_mode: str, cancel_mode: str, phase: str): + save_started = asyncio.Event() + release_save = asyncio.Event() + save_finished = asyncio.Event() + workflow_finished = asyncio.Event() + + class BlockingSaveStorage(InMemoryCheckpointStorage): + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + async def write() -> CheckpointID: + save_started.set() + try: + await release_save.wait() + return await super(BlockingSaveStorage, self).save(checkpoint) + finally: + save_finished.set() + + if save_mode == "task": + return await asyncio.create_task(write()) + if save_mode == "gather": + return (await asyncio.gather(write()))[0] + return await write() + + storage = BlockingSaveStorage() + + @step(replay_key=lambda value: "complete") + async def complete(value: int) -> int: + return value + 1 + + @step(replay_key=lambda: "interrupt") + async def interrupt(ctx: RunContext) -> int: + await save_started.wait() + if phase == "failure": + raise RuntimeError("sibling failure") + return await ctx.request_info("continue?", int, request_id="continue") + + @built_workflow(checkpoint_storage=storage) + async def wf(value: int) -> int | list[int]: + try: + if phase == "step": + return await complete(value) + return list(await asyncio.gather(complete(value), interrupt())) + finally: + workflow_finished.set() + + run_task = asyncio.ensure_future(wf.run(0)) + cancellation_task = run_task + expected_error: type[BaseException] + try: + await asyncio.wait_for(save_started.wait(), timeout=5) + if phase != "step": + await asyncio.wait_for(workflow_finished.wait(), timeout=5) + await asyncio.sleep(0) + await asyncio.sleep(0) + if cancel_mode == "timeout": + cancellation_task = asyncio.ensure_future(asyncio.wait_for(run_task, timeout=0.01)) + expected_error = asyncio.TimeoutError + else: + run_task.cancel() + expected_error = asyncio.CancelledError + + finished, _ = await asyncio.wait({cancellation_task}, timeout=1) + assert cancellation_task in finished, "Caller cancellation is blocked by the checkpoint save" + with pytest.raises(expected_error): + await cancellation_task + + assert run_task.cancelled() + assert save_finished.is_set() + assert await storage.list_checkpoints(workflow_name="wf") == [] + finally: + release_save.set() + await asyncio.gather(cancellation_task, run_task, return_exceptions=True) + + async def test_caller_cancellation_before_save_starts_closes_storage_coroutine( + self, recwarn: pytest.WarningsRecorder + ): + storage = InMemoryCheckpointStorage() + run_task: asyncio.Future[Any] | None = None + + @step(replay_key=lambda: "complete") + async def complete() -> int: + assert run_task is not None + run_task.cancel() + return 1 + + @built_workflow(checkpoint_storage=storage) + async def wf(_: int) -> int: + return await complete() + + run_task = asyncio.ensure_future(wf.run(0)) + with pytest.raises(asyncio.CancelledError): + await run_task + run_task = None + # Let task completion callbacks run before checking coroutine ownership. + await asyncio.sleep(0) + await asyncio.sleep(0) + gc.collect() + assert not [warning for warning in recwarn if "was never awaited" in str(warning.message)] + assert await storage.list_checkpoints(workflow_name="wf") == [] + + async def test_double_caller_cancellation_drains_checkpoint_save_and_background_tasks(self): + save_started = asyncio.Event() + release_save = asyncio.Event() + save_finished = asyncio.Event() + save_cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + background_started = asyncio.Event() + background_finished = asyncio.Event() + background_tasks: list[asyncio.Task[None]] = [] + + class BlockingSaveStorage(InMemoryCheckpointStorage): + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + save_started.set() + try: + await release_save.wait() + return await super().save(checkpoint) + finally: + save_cleanup_started.set() + try: + await release_cleanup.wait() + finally: + save_finished.set() + + storage = BlockingSaveStorage() + + @step(replay_key=lambda value: "complete") + async def complete(value: int) -> int: + return value + 1 + + async def background() -> None: + background_started.set() + try: + await asyncio.Event().wait() + finally: + background_finished.set() + + @built_workflow(checkpoint_storage=storage) + async def wf(value: int) -> int: + background_tasks.append(asyncio.create_task(background())) + await background_started.wait() + return await complete(value) + + run_task = asyncio.ensure_future(wf.run(0)) + try: + await asyncio.wait_for(save_started.wait(), timeout=5) + run_task.cancel() + await asyncio.wait_for(save_cleanup_started.wait(), timeout=5) + run_task.cancel() + release_cleanup.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=5) + await asyncio.wait_for(save_finished.wait(), timeout=5) + assert background_finished.is_set() + assert all(task.done() for task in background_tasks) + assert await storage.list_checkpoints(workflow_name="wf") == [] + finally: + release_save.set() + release_cleanup.set() + for task in background_tasks: + task.cancel() + await asyncio.gather(run_task, *background_tasks, return_exceptions=True) + + async def test_pause_sibling_cancellation_preserves_in_flight_checkpoint(self): + save_started = asyncio.Event() + release_save = asyncio.Event() + pause_requested = asyncio.Event() + cancellation_raised = asyncio.Event() + save_cancelled = False + + class BlockingSaveStorage(InMemoryCheckpointStorage): + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + nonlocal save_cancelled + save_started.set() + try: + await release_save.wait() + except asyncio.CancelledError: + save_cancelled = True + raise + return await super().save(checkpoint) + + storage = BlockingSaveStorage() + + @step(replay_key=lambda: "complete") + async def complete() -> int: + return 1 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await save_started.wait() + pause_requested.set() + return await ctx.request_info("continue?", int, request_id="continue") + + @step(replay_key=lambda: "cancel") + async def cancel() -> None: + await pause_requested.wait() + await asyncio.sleep(0) + cancellation_raised.set() + raise asyncio.CancelledError("sibling cancellation") + + @built_workflow(checkpoint_storage=storage) + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(complete(), ask(), cancel())) + + async def release_after_settlement_starts() -> None: + await cancellation_raised.wait() + await asyncio.sleep(0) + release_save.set() + + release_task = asyncio.create_task(release_after_settlement_starts()) + try: + # Cancellation messages are not preserved consistently across supported Python versions. + with pytest.raises(asyncio.CancelledError): + await wf.run(0) + finally: + release_save.set() + await release_task + + assert not save_cancelled + assert len(await storage.list_checkpoints(workflow_name="wf")) == 1 + + async def test_second_cancellation_during_pause_settlement_finishes_save(self): + save_started = asyncio.Event() + release_save = asyncio.Event() + pause_requested = asyncio.Event() + cancellation_raised = asyncio.Event() + save_finished = asyncio.Event() + save_cancelled = False + + class BlockingSaveStorage(InMemoryCheckpointStorage): + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + nonlocal save_cancelled + save_started.set() + try: + await release_save.wait() + except asyncio.CancelledError: + save_cancelled = True + raise + finally: + save_finished.set() + return await super().save(checkpoint) + + storage = BlockingSaveStorage() + + @step(replay_key=lambda: "complete") + async def complete() -> int: + return 1 + + @step(replay_key=lambda: "ask") + async def ask(ctx: RunContext) -> int: + await save_started.wait() + pause_requested.set() + return await ctx.request_info("continue?", int, request_id="continue") + + @step(replay_key=lambda: "cancel") + async def cancel() -> None: + await pause_requested.wait() + await asyncio.sleep(0) + cancellation_raised.set() + raise asyncio.CancelledError("sibling cancellation") + + @built_workflow(checkpoint_storage=storage) + async def wf(_: int) -> list[object]: + return list(await asyncio.gather(complete(), ask(), cancel())) + + run_task = asyncio.ensure_future(wf.run(0)) + try: + await asyncio.wait_for(cancellation_raised.wait(), timeout=5) + await asyncio.sleep(0) + await asyncio.sleep(0) + run_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=5) + await asyncio.wait_for(save_finished.wait(), timeout=2) + assert save_cancelled + finally: + release_save.set() + await asyncio.gather(run_task, return_exceptions=True) + # --------------------------------------------------------------------------- # Step replay identity