Python: [Bug]: Functional workflow drops a finished @step's checkpoint when a sibling step fails during the save, so resuming runs the step again #9615

Observed Behavior

A functional workflow with checkpoint storage runs prepare, then two @step calls with asyncio.gather: done returns at once, and flaky fails the first time it runs. The failed run is then resumed from the latest checkpoint in a fresh workflow.

done finishes before flaky fails. It emits executor_completed, and its result goes into the run's step cache. Then the step awaits its per-step checkpoint save (ctx._on_step_completed()). If the storage's save yields to the event loop before it writes, as a database-backed storage does, flaky fails during that await. The workflow then cancels its remaining tasks (_cancel_workflow_tasks), and that cancels done's save. The failure path of _run_core saves no checkpoint. So the latest checkpoint holds only prepare's result, and the resumed run runs done again, with its side effects.

With InMemoryCheckpointStorage the save finishes without yielding, so done is not run again. FileCheckpointStorage shields its write and finishes it even when cancelled, so it is not affected either. Whether a finished step runs again therefore depends on how the storage's save awaits.

Expected Behavior

A step that finished, and emitted executor_completed, before the run failed is restored from the checkpoint and not run again, whatever storage is used. The functional-workflow sample steps_and_checkpointing.py says "Each step saves a checkpoint after it completes" and "On crash recovery from a checkpoint, earlier step results are restored". StepWrapper documents "a checkpoint is saved after each live execution when checkpoint storage is configured". With either storage the script should print runs {'prepare': 1, 'done': 1, 'flaky': 2}.

Steps to Reproduce

1. pip install agent-framework-core==1.23.0
2. Run the script below. RemoteStorage is InMemoryCheckpointStorage whose save first awaits 10 ms, like a network round trip (for example CosmosCheckpointStorage.save, which awaits the container before it upserts).
3. With RemoteStorage, done runs twice.

Minimal Reproduction

import asyncio
from collections import Counter
from importlib.metadata import version

from agent_framework import InMemoryCheckpointStorage, step, workflow

runs = Counter()


class Boom(Exception):
    pass


class RemoteStorage(InMemoryCheckpointStorage):
    """A storage whose save awaits I/O, as a database or a remote store does."""

    async def save(self, checkpoint):
        await asyncio.sleep(0.01)
        return await super().save(checkpoint)


@step
async def prepare(x: int) -> int:
    runs["prepare"] += 1
    return x


@step
async def done(x: int) -> int:
    runs["done"] += 1
    return x + 1


@step
async def flaky(x: int) -> int:
    runs["flaky"] += 1
    if runs["flaky"] == 1:
        raise Boom("transient failure")
    return x + 2


@workflow
async def flow(x: int) -> list[int]:
    a = await prepare(x)
    return list(await asyncio.gather(done(a), flaky(a)))


async def main():
    for storage in (InMemoryCheckpointStorage(), RemoteStorage()):
        runs.clear()
        app = flow.build(checkpoint_storage=storage)
        completed = []
        try:
            async for event in app.run(0, stream=True):
                if event.type == "executor_completed":
                    completed.append(event.executor_id)
        except Boom:
            pass
        latest = await storage.get_latest(workflow_name=flow.name)
        result = await flow.build(checkpoint_storage=storage).run(checkpoint_id=latest.checkpoint_id)
        print(f"{type(storage).__name__}: completed before the failure {completed}; "
              f"resumed {result.get_outputs()}; runs {dict(runs)}")
    print(f"agent-framework-core {version('agent-framework-core')}")


asyncio.run(main())

Error Messages and Stack Traces

No exception. Output:

InMemoryCheckpointStorage: completed before the failure ['prepare', 'done']; resumed [[1, 2]]; runs {'prepare': 1, 'done': 1, 'flaky': 2}
RemoteStorage: completed before the failure ['prepare', 'done']; resumed [[1, 2]]; runs {'prepare': 1, 'done': 2, 'flaky': 2}
agent-framework-core 1.23.0

Where `done`'s save is cancelled (a `traceback.print_stack()` in `RemoteStorage.save` when it receives `CancelledError`):

  File "agent_framework/_workflows/_functional.py", line 125, in _tracked_coro
  File "agent_framework/_workflows/_functional.py", line 1601, in __call__
    await ctx._on_step_completed()
  File "agent_framework/_workflows/_functional.py", line 2153, in _on_step_completed
    ckpt_chain[0] = await self._save_checkpoint(ctx, storage, ckpt_chain[0])
  File "agent_framework/_workflows/_functional.py", line 2397, in _save_checkpoint
    return await storage.save(checkpoint)
  File "repro.py", line 19, in save

Package Versions

agent-framework-core: 1.23.0 (also 1.22.0, and main at a8e284c)

Python Version

Python 3.13.9

Operating System

Linux

Regression

No

Additional Context

• StepWrapper.__call__ caches the result (ctx._set_cached_result), emits executor_completed, and then awaits ctx._on_step_completed() (_functional.py lines 1597-1601). When a sibling fails, _execute calls ctx._cancel_workflow_tasks() (line 2352), which cancels the step's task while that save is pending. The failure branch of _run_core (line 2231) then yields the events and re-raises without saving a checkpoint. The WorkflowInterrupted branch does save one.
• In the reproduction no checkpoint after prepare reaches the storage, so get_latest returns prepare's checkpoint.
• Possible direction: on the failure path, save a checkpoint of the step cache before re-raising, as the interruption path does. Locally, with that change (3 lines after for event in ctx._get_events(): yield event in the except Exception branch), the reproduction prints runs {'prepare': 1, 'done': 1, 'flaky': 2} for both storages, and tests/workflow/test_functional_workflow.py passes (209 tests). Another option is to let a step's pending save finish when the task is cancelled, as FileCheckpointStorage.save does for its own write.
• Related but different: Python: [Bug]: Parallel functional steps fork the checkpoint lineage #7647 (concurrent per-step saves fork the checkpoint lineage, without losing results).

Acknowledgements

• I searched existing issues and did not find a duplicate.
• I personally verified this behavior and the reproduction details are authentic.
• I will wait for explicit maintainer agreement before starting implementation of a non-trivial change.