Python: fix: make concurrent FileCheckpointStorage saves not race on a shared temp path

[Manjunath Janardhan (manjunathshiva)]

Motivation & Context

Concurrent calls to FileCheckpointStorage.save() for the same checkpoint ID raced over a shared temporary file (<checkpoint-id>.json.tmp). Each save wrote to and then os.replaced that fixed path, so whichever save renamed first removed the temp file still being written by a competing save, which then failed — FileNotFoundError on POSIX (per the issue) and PermissionError [WinError 5] on Windows (where the concurrent-replace also trips the destination-lock check). In the issue's 100-concurrent-save repro, between 49 and 71 saves failed.

Description & Review Guide

Three cooperating parts in _checkpoint.py, the last two reshaped by review.

• Unique temp file per save. Each save writes to its own .maf-ckpt-<uuid>.tmp in the
  destination directory, created with os.open(O_CREAT | O_EXCL | O_WRONLY, 0o666) so the file
  keeps the process umask exactly as the previous open(..., "w") did. Same directory keeps
  os.replace atomic, and per-save temp names remove the shared-path collision that Python: [Bug]: Concurrent FileCheckpointStorage saves fail due to shared temporary path #7748 reported.

• Ownership taken before the write is submitted. A process-wide queue keyed by the canonical
  destination path hands out ownership; only then is the write submitted to the executor. Waiters
  suspend on a concurrent.futures.Future signal, awaited through a per-wait future fed by a
  done-callback rather than through asyncio.wrap_future, which chains cancellation into the future
  it wraps.

  Two properties follow from that choice. A queued same-path save occupies no executor worker, where
  previously each one held a worker purely to block on threading.Lock.acquire() — a burst could
  fill the default pool, stall unrelated to_thread work including checkpoint loads, and deadlock
  once the write that had to finish first was queued behind those waiters. And because a
  concurrent.futures.Future is not bound to a loop, one chain orders writers enqueued from
  different event loops, which a per-path asyncio primitive cannot do.

  Ownership is released as the write's last act on the worker thread, with the coroutine's finally
  as a backstop for the case where nothing was submitted. Releasing only from the coroutine would be
  a regression against the threading.Lock it replaces: a lock is released by the worker thread,
  which outlives the loop that submitted it, whereas a coroutine whose loop is closed while its task
  is pending never releases -- leaving the destination owned for the life of the process.

  Registry entries are reference-counted by queued-or-running operations and dropped by the last
  release. The previous dict never removed an entry, and since WorkflowCheckpoint generates a fresh
  UUID by default, ordinary saves each retained one for the life of the process.

• Cancellation, handled explicitly in both directions. Cancelled while still waiting for the
  destination: the write is never submitted, so there is nothing left to land later — and ownership is
  still handed on, or every later save for that path would wait forever. Cancelled once the write is
  running: the worker is drained before the cancellation propagates and before ownership is released,
  since releasing first would let the next save begin while this os.replace is in flight.

  The drain absorbs re-delivered cancellations. A single await asyncio.shield(worker) is not
  enough: once a task has a cancellation pending its next await raises immediately, so the shield
  returns with the worker still running.

• Ownership survives the loop that took it. A submitted write is released by its worker thread,
  which outlives the loop that submitted it. A ticket still waiting for its predecessor has no
  worker to fall back on, so the predecessor's callback releases it directly when the waiter's loop
  has closed — on the thread that resolved the predecessor, rather than on a loop that will never
  run again. Without that, a queued save whose loop went away left its hand-off signal unresolved
  and every later save for that destination waited forever.

  Two windows this cannot close: a loop that closes after the wake-up was queued but before it runs,
  and one abandoned without being closed at all. Both need a loop closed with tasks still pending,
  which asyncio already reports as an error. A graceful shutdown is unaffected, because
  asyncio.run cancels pending tasks first and that drives a queued save through its cancellation
  path — pinned by its own test, since that is the path callers actually depend on.

• Bounded replace retry. Retained from the original fix. On Windows os.replace can transiently
  raise PermissionError even fully serialized, because a background indexer or AV scan briefly
  holds a handle to the destination (19 failures in 200 purely sequential replaces locally). Five
  attempts with 1→2→4→8→16 ms backoff absorb it.

• What are the major changes? Destination ownership moved out of the worker thread and in front of
  submission, reference-counted registry entries, and an explicit cancellation contract. No public
  API change.

• What is the impact of these changes? Concurrent saves of one checkpoint ID no longer fail, and
  last-writer-wins now follows the order callers observed rather than the order the executor happened
  to schedule. Queued saves no longer consume executor capacity, so a save burst cannot starve
  checkpoint loads.

  One observable behaviour change: a cancelled save() returns after its own write finishes, where
  before it returned promptly and the write continued behind it. That is the fix — a cancelled write
  landing later could overwrite a newer checkpoint — but it does mean a caller that cancels no longer
  proceeds immediately.

  Deliberately out of scope: delete() does not join the queue, so a delete racing a save to the
  same path remains unserialized. Python: [Bug]: Concurrent FileCheckpointStorage saves fail due to shared temporary path #7748 is about concurrent saves; happy to add it if wanted.

• What do you want reviewers to focus on? The cancellation contract, since it is the part with a
  caller-visible consequence, and whether draining should be labelled a breaking change.

  Also the two loop-abandonment windows described above. I chose to state them rather than engineer
  around them: closing them means tracking each ticket's loop and reaping, which is real complexity
  in this path for a failure only reachable through a loop closed with tasks still pending. Say if
  you would rather have the machinery.

  Second, the ownership hand-off. The signal is a concurrent.futures.Future so any loop can await
  it, and nothing asyncio-owned may mark it done — cancelling a waiter must not reach it, and the
  submitted write must not be a Task, or loop shutdown can cancel it before it reaches the executor.
  An earlier revision of this description claimed wrap_future was safe here on the strength of a
  check that read the wrapped future one loop iteration too early; it is not, and the mechanism has
  been replaced.

  Validation: core 5354 passing with _checkpoint.py at 98% line coverage, poe syntax and
  poe typing clean across all five checkers, and ag-ui, declarative and foundry_hosting green as
  the downstream users of checkpoint storage. The three lines still uncovered in the file all
  predate this PR, apart from two that arrived with Python: FileCheckpointStorage save/load symmetry (#8181) #8214's own save-time validation when it
  merged into this branch.

  Python: FileCheckpointStorage save/load symmetry (#8181) #8214 merged while this was open and rewrote save() too. Resolved keeping its encoding="utf-8" and its save-time validation rather than letting "keep ours" discard them -- the encoding fix had landed inside the conflicting hunk, so it needed restoring by hand. That is a real bug: ensure_ascii=False into a cp1252 stream raised UnicodeEncodeError on a checkpoint containing CJK text or emoji, and fixing only the write side would have traded a loud write error for silent read corruption. Its validation raises before _enqueue_write, so a rejected checkpoint never takes a destination ticket.

  This PR adds or rewrites 25 tests. I re-ran all 25 against main's _checkpoint.py rather than
  restating an earlier claim, and all 25 fail there. An earlier revision of this description said
  "ten", which had been true several rounds ago and was never updated.

  End to end, the issue's own repro -- 100 concurrent save() calls for one checkpoint ID -- fails
  59/100 then 85/100 with PermissionError on main (matching the 49-71 reported) and 0/100 here
  across three runs. The surviving checkpoint is also deterministic now: main leaves a different
  winner each run because the last os.replace to land is whichever the executor scheduled last,
  where this branch leaves the last save enqueued, every run.

Related Issue

Fixes #7748

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.

[Copilot]

Pull request overview

Note

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

This PR addresses a concurrency race in FileCheckpointStorage.save() when multiple saves target the same checkpoint ID, and adds a regression test to ensure concurrent saves don’t fail.

Changes:

• Added a concurrency regression test covering concurrent saves with the same checkpoint ID.
• Updated FileCheckpointStorage.save() to use unique temp files and added per-checkpoint-ID serialization plus a retry loop around os.replace() for Windows.
• Implemented best-effort cleanup for temp files when failures occur.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
python/packages/core/tests/workflow/test_checkpoint.py	Adds a regression test that exercises concurrent save() calls for the same checkpoint ID.
python/packages/core/agent_framework/_workflows/_checkpoint.py	Makes save() more robust under concurrency by changing temp-file strategy, adding per-ID locking, and retrying replace on Windows.

---

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

[Evan Mattson (moonbox3)]

/review

[github-actions]

MAF Automated Review — Iteration 1

Result: Findings reported
Scope: full PR (1 commit(s)): fd0a15ab45a6
Model: gpt-5.6-sol

Overview

The change removes the shared temporary-path race by giving each save a unique same-directory file, publishing it atomically, and serializing same-ID saves at the coroutine level. The regression test verifies concurrent saves complete and leave a parseable checkpoint. Residual compatibility and lifecycle issues remain around active-lock eviction, maximum-length checkpoint IDs, event-loop retention, and changed file permissions.

Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
4 verified findings remained after source verification (4 medium) across 1 file. Details are attached to the affected lines below.

Affected areas: python/packages/core/agent_framework/_workflows/_checkpoint.py

[Copilot]

🔵 Needs a closer look

Cross-loop synchronization, cancellation, and executor-shutdown behavior warrant final human review.

Review details

Suppressed comments (1)

python/packages/core/tests/workflow/test_checkpoint.py:2209

• This says the cross-loop signal is awaited with asyncio.wrap_future, but _wait_for_signal deliberately avoids that API because cancellation propagates into the wrapped future. Please describe the per-loop waiter callback actually under test.

    together. The queue hands ownership over a `concurrent.futures.Future`, which any
    loop can await through `asyncio.wrap_future`, so a single chain orders both.

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Balanced

[Evan Mattson (moonbox3)]

Please fix the failing unit tests.

[Manjunath Janardhan (manjunathshiva)]

Fixed in c10466c94 — and thanks for updating the branch, I have merged your two merge commits
rather than pushing over them.

The failure was not an assertion. The worker process died, with no Python traceback:

worker 'gw2' crashed while running
'test_file_checkpoint_storage_abandonment_mid_chain_keeps_the_queue_ordered'

Cause: an unreferenced pending task. That test did
abandoned_loop.create_task(storage.save(...)) without keeping the result, so the task became
collectable the moment the call returned — and the loop it belongs to is deliberately closed while
the task is still pending. That left the task's finalizer, which touches the closed loop, to run at
whatever point the garbage collector chose, including during worker teardown.

The shape of the failure supports that. It is not a version boundary: 3.11, 3.12 and 3.14 failed on
Windows while 3.10 and 3.13 passed, and every Linux job passed. Non-monotonic scatter like that
is a timing flake rather than an API difference, and GC timing is the thing that varies between
runs.

Three sibling tests had the same unreferenced-task pattern, all of them in tests that abandon or
shut down a loop on purpose, so all four now hold their task for the duration of the test. What made
the inconsistency obvious is that ..._abandoned_loop_while_queued_releases_the_destination already
did this correctly — the crashing test was the copy that did not. Holding the reference does not
weaken what any of them prove: the abandonment under test is of the loop, not of the task object,
and the task stays pending either way. It only moves the finalizer to a deterministic point.

One caveat I would rather state than gloss over: I could not reproduce the crash locally, so CI
is the real verdict here, not my machine. Ten attempts on a 3.12 interpreter — the test alone, the
whole file under xdist, the workflow directory, and the full core suite under xdist six times — all
passed. After the fix, five more full core runs under xdist on 3.12 are clean, and poe check -P core
is green at 5197 passing with _checkpoint.py at 99%. So I am confident in the mechanism and that it
matches the symptom exactly, but not in a way I can demonstrate by reproduction.

The Python workflows on this head are sitting at action_required and cannot run until someone
approves them — could you release them when you get a chance? Given I could not reproduce the crash
locally, your CI is the only thing that can actually confirm this.

[Eduard van Valkenburg (eavanvalkenburg)]

Thanks for the update. The review thread from the previous request is now resolved, but the current head is still not ready: Python - Tests and Merge Gatekeeper are failing, several Python matrix jobs were cancelled, and the branch conflicts with main. Could you please get the checks passing and resolve the merge conflicts, then re-request review? Thanks!

[Manjunath Janardhan (manjunathshiva)]

Fixed in edb83345d, and main merged in 03c612f4f so the conflict is cleared.

The failure was found 0 tasks in
test_file_checkpoint_storage_shutdown_before_the_write_starts_does_not_hang, and it was my test's
timing assumption rather than the code. It created the save task, yielded once with
await asyncio.sleep(0), then asserted exactly one pending task to prove the write is not itself a
task. But a single yield drains every ready callback, so on a fast tmpfs the executor write could
finish and resolve the shield inside that same batch -- the save completed and asyncio.all_tasks()
came back empty. Windows passed and Linux did not for exactly that reason.

The write is now held at its first syscall, so the save is deterministically suspended when the
count is taken. Three details worth naming:

• The gate covers only this save's own temp file. checkpoint_module.os is the real os module, so
  the patch is process-wide, and blocking every os.open would have stalled pytest's own I/O and
  coverage writes on the same event and deadlocked the run. I had it wrong that way first.
• The assertions are split. An empty list now reports that the gate did not hold, which is a problem
  with the test; two tasks reports the write being a task, which is the regression. One message could
  not honestly cover both, and 0 was indistinguishable from the flake it replaced.
• A failing assertion releases the held worker before propagating, so the real error surfaces
  immediately rather than after the gate's timeout.

Verified rather than asserted: reverting the source to ensure_future(asyncio.to_thread(...)) still
fails the test, now with found 2 tasks and the pending to_thread() task in the output. Stable
over ten runs of the test and five shuffled xdist runs of the whole file.

The merge had two import conflicts, both resolved keeping what this branch needs -- import time
for _replace_with_retry's backoff, which main does not have, and main's new module-level
import threading alongside ours. No test was lost on either side (74 here plus main's, 75 after),
and #8214's encoding="utf-8" and save-time validation are both still present.

Validation: core 5354 passing, _checkpoint.py at 98% -- the two uncovered lines are #8214's own
exception branch -- with poe syntax and poe typing clean across all five checkers, plus
orchestrations, ag-ui, declarative and foundry_hosting green.

The workflows on this head are queued at action_required again; could you release them?