From 72b554867bbc7545528d60f04767b227abb32d92 Mon Sep 17 00:00:00 2001 From: ShiYuan Xiong <2757609992@qq.com> Date: Fri, 23 Oct 2026 20:17:26 +0800 Subject: [PATCH 1/2] fix(workflow): finish in-flight checkpoint saves when a sibling step fails A step that completed and emitted executor_completed before a sibling failed could have its per-step checkpoint save cancelled by run teardown. asyncio.gather cancels the completed step's task as soon as the sibling raises, and the failure-path task sweep cancelled the save task itself (the loop task factory tracks every task), so a storage whose save awaits I/O before writing lost the checkpoint. The latest checkpoint then held no trace of the finished step and a resumed run re-ran it and its side effects. Run the save as its own task, shield both awaits so a cancellation of the step task cannot reach it, and have the task sweep wait for in-flight saves instead of cancelling them. A completed step is now restored from the checkpoint whatever storage is used. --- .../agent_framework/_workflows/_functional.py | 34 +++++++++++- .../workflow/test_functional_workflow.py | 52 +++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 65dfefb933b..ee17ebd88a0 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -696,6 +696,10 @@ 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() # 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. @@ -1071,7 +1075,7 @@ async def _cancel_workflow_tasks(self) -> None: while pending := [task for task in self._workflow_tasks if task is not current_task and not task.done()]: 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 self._checkpoint_save_tasks: continue if self._has_pending_cancellation(task): waiting_for_existing_cancellation = True @@ -1080,6 +1084,9 @@ async def _cancel_workflow_tasks(self) -> None: if waiting_for_existing_cancellation: await asyncio.sleep(0) continue + # Checkpoint saves are awaited, not cancelled: a completed step's + # checkpoint must land before the run is torn down, whatever storage + # is configured. await asyncio.gather(*pending, return_exceptions=True) tasks = [task for task in self._workflow_tasks if task is not current_task] if tasks: @@ -2150,7 +2157,30 @@ 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]) + # A checkpoint save that has started must finish even if the step + # is cancelled (e.g. 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. + save_task = asyncio.ensure_future(self._save_checkpoint(ctx, storage, ckpt_chain[0])) + 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 diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 8eeb8d28c1c..d609512ab42 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,55 @@ 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 + async def test_completed_parallel_step_is_checkpointed_when_sibling_fails(self): + """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() + + 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] + await asyncio.sleep(0.01) + return await super().save(checkpoint) + + 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: + 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}) + # --------------------------------------------------------------------------- # Step replay identity From ba8eed8895b8ed69e21cf5560bdb8c54564bb2d4 Mon Sep 17 00:00:00 2001 From: hehe764 <243333382+hehe764@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:14:40 +0800 Subject: [PATCH 2/2] fix(workflow): preserve checkpoint save descendant tasks --- .../agent_framework/_workflows/_functional.py | 11 ++++++++++- .../tests/workflow/test_functional_workflow.py | 18 +++++++++++++++--- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index ee17ebd88a0..8eb3fbfd467 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -1073,9 +1073,18 @@ async def _cancel_workflow_tasks(self) -> None: """Cancel and drain all tracked workflow tasks.""" 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 or task in self._checkpoint_save_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 diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index d609512ab42..b7d1b5f7592 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -933,10 +933,12 @@ 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 - async def test_completed_parallel_step_is_checkpointed_when_sibling_fails(self): + @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 @@ -945,8 +947,17 @@ 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] - await asyncio.sleep(0.01) - return await super().save(checkpoint) + 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() @@ -964,6 +975,7 @@ async def done(x: int) -> int: async def flaky(x: int) -> int: runs["flaky"] += 1 if runs["flaky"] == 1: + await save_started.wait() raise TransientFailure("sibling failure") return x + 2