Python: serialize per-step checkpoint saves in functional workflows

[Ruiming Zhao (uuzzrm)]

Motivation & Context

Per-step checkpoints in functional workflows assume a single linear history: each new checkpoint links to the previous one via previous_checkpoint_id. That assumption breaks as soon as two @step calls complete around the same time (e.g. asyncio.gather). Both completion callbacks read the same chain head before either writes the updated value back, so both create checkpoints that point at the same parent - two sibling roots instead of one chain. A later final/approval checkpoint continues from only one of them, and the other becomes unreachable when traversing the history from the latest checkpoint.

This is the same invariant #4588 / #6695 established for restoration-time ancestry, but on the concurrent save path inside a single run. The final result is unaffected; checkpoint ordering and history traversal are what fork.

Description & Review Guide

• What are the major changes?
  The _on_step_completed closure in FunctionalWorkflow._run_core now guards the read-save-update of the chain head with an asyncio.Lock, so concurrent step completions serialize on the checkpoint chain instead of racing on it. The lock is scoped to one run, so it does not couple independent workflows sharing the same storage.
• What is the impact of these changes?
  Per-step checkpointing behaves the same for sequential workflows (the lock is uncontended). For concurrent steps, each checkpoint now always chains to the one saved just before it, keeping exactly one root and a fully reachable history. The regression test uses an in-memory storage whose save() yields to the event loop, which reproduces the file/database interleaving deterministically - with the fix it sees one root and a complete chain, without it two sibling roots.
• What do you want reviewers to focus on?
  Whether guarding _on_step_completed covers every path that can write per-step checkpoints, and whether the lock placement (per-run, next to the chain head it protects) reads naturally.

Related Issue

Fixes #7647

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) - a workflow keeps the label and title prefix in sync automatically.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Fixes a concurrency issue where parallel step completions could fork the checkpoint lineage by serializing checkpoint saves, and adds a regression test that deterministically reproduces the interleaving seen with real checkpoint backends.

Changes:

• Serialize per-step checkpoint save/update with an asyncio.Lock to prevent sibling-root checkpoints under parallel step completion.
• Add a regression test using a yielding in-memory checkpoint storage to force interleaving and validate a single linear checkpoint chain.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
python/packages/core/agent_framework/_workflows/_functional.py	Adds a lock around step-completion checkpoint saving to prevent lineage forks under concurrent step completion.
python/packages/core/tests/workflow/test_functional_workflow.py	Adds a deterministic regression test (and storage shim) that exercises the parallel-step checkpoint lineage scenario.

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

[Ruiming Zhao (uuzzrm)]

Good catch, fixed in f68cb2b - removed the duplicated header block so the section boundary is clean.