Python: fix(workflow): finish in-flight checkpoint saves when a sibling step fails

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.
@@ -1069,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:
+                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,6 +1093,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 +2166,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
 

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,67 @@ 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})
+
 
 # ---------------------------------------------------------------------------
 # Step replay identity