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

The fix has two cooperating parts in _checkpoint.py:

• Unique temp file per save. Replace the fixed with_suffix(".json.tmp") path with tempfile.NamedTemporaryFile(dir=<dest-dir>, prefix=<id>., suffix=.json.tmp, delete=False). Creating the temp file in the destination directory keeps os.replace atomic (same filesystem), and each save now owns its own temp file, eliminating the shared-path collision entirely.

• Per-ID write serialization + bounded replace retry. Add a per-checkpoint-ID threading.Lock map so concurrent saves of the same ID execute their os.replace serially. On Windows, os.replace can still transiently raise PermissionError even when fully serialized — a background indexer/AV scan can briefly hold a handle to the just-replaced destination file (I reproduced 19 failures out of 200 purely single-threaded sequential os.replace calls on Windows locally). A short bounded retry (5 attempts, 1→2→4→8→16 ms backoff) absorbs that transient.

• Temp files are cleaned up on failure (finally + unlink(missing_ok=True)), best-effort.

• Regression test test_file_checkpoint_storage_concurrent_saves_same_id runs 50 concurrent saves of a shared checkpoint ID and asserts zero raise internal filesystem exceptions and the final file exists.

• What are the major changes? Unique temp files + per-ID lock + replace retry in FileCheckpointStorage.save() (packages/core/_workflows/_checkpoint.py); one regression test, no public API change.

• What is the impact of these changes? Concurrent saves of the same checkpoint ID no longer fail; each save's last writer wins, which is the documented FileCheckpointStorage semantics. Not a breaking change.

• What do you want reviewers to focus on? Whether the retry constant (5 attempts, exponential backoff up to 16 ms) and the bounded set of retried exceptions (PermissionError only) are the right shape, or whether a broader/narrower policy is preferred.

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

[github-actions]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/core/agent_framework/_workflows
_checkpoint.py	189	9	95%	339, 368–371, 404–406, 409
TOTAL	47369	4377	90%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
9647	36 💤	0 ❌	0 🔥	2m 9s ⏱️

[Evan Mattson (moonbox3)]

Following up on the earlier cancellation concern: this commit moves serialization into _write_atomic() and protects the asyncio.to_thread() task with asyncio.shield(), but a save cancelled before its worker starts remains queued after its caller receives CancelledError. A later save on another event loop can acquire _get_file_lock(), return successfully, and then be overwritten when the cancelled worker eventually runs. Could the cancellation path either remove a queued write or drain it before propagating cancellation?

[Evan Mattson (moonbox3)]

Following up on the earlier thread-pool comment: could waiting happen before work enters the default executor? This commit moves _get_file_lock() inside _write_atomic(), so every same-path save() now occupies a worker while waiting; a burst can fill the pool and delay unrelated asyncio.to_thread() work, including checkpoint loads and independent saves. A path-keyed queue or dedicated write executor could preserve cross-instance serialization while delegating only the active write.

[Manjunath Janardhan (manjunathshiva)]

Thanks — reading these together, all three land on the same root: _get_file_lock inside the worker, no refcounting, no thought given to where the wait happens. Wanted to put the candidate shapes in front of you before committing, since they're not equivalent:

Verified against the current tree

• _get_file_lock(file_path) sits inside _write_atomic, which runs via asyncio.to_thread on the default executor. A burst of N saves on one ID = N default-pool workers all blocked in threading.Lock.acquire(), so asyncio.to_thread for unrelated work (loads, independent saves to other paths) queues until a worker frees. Default max_workers = min(32, cpu + 4), so even ~16 racing saves can saturate on a smaller host. Real.
• _file_locks is unbounded — every distinct checkpoint ID ever saved (and not remove()d) leaves a Path → threading.Lock entry. UUID-keyed IDs make this accumulate in long-running orchestrators. Real.
• Cancel-before-start: the shielded inner Task wraps loop.run_in_executor. If the caller's Task.cancel() lands before the executor submission completes, the wrapped Task is cancelled and the executor does not run the function (CPython 3.11+: to_thread does loop.create_task(coro) whose first step is the submission; cancel before that step → no worker ever picked it up). Empirically hard to hit but present; the post-cancel stale-write concern from the first round is already mitigated by the per-path lock (next save serializes, so final state is correct even if a stale write briefly runs).

Fix shapes on the table

• (A) Refcounted registry, lock acquired in the coroutine body — not inside the worker.

  entry=await_acquire_path_lock(file_path) # bumps holders under guard; does threading.Lock.acquire() via asyncio.to_thread for the brief acquire onlytry:
  awaitasyncio.shield(asyncio.to_thread(_write_atomic)) # worker does pure I/O; lock already held by the coroutinefinally:
  _release_path_lock(file_path, entry) # drops entry when holders hit 0

  Addresses (2) by moving the wait out of the long-lived worker (a burst of N same-ID saves spawns N short-lived to_thread acquisitions, not N held workers) and (3) via refcount. (1) still bounded by the CPython window above, but the lock guarantees final-state ordering.

• (B) Per-path single-worker ThreadPoolExecutor, refcounted.

  w=_writer_for(file_path) # single-thread executor per pathtry:
  awaitloop.run_in_executor(w.executor, _write_atomic) # serialization from queueing on that executor, no threading.Lock neededfinally:
  _release_writer(file_path, w) # shutdown wait=True when refcount hits 0

  Strongest response to (2): same-path saves queue on their own dedicated thread, never touch the default pool even while waiting; concurrent saves to different IDs get parallel writers. (3) and the cancel-before-start window unchanged. Cost: one OS thread per in-flight path, plus shutdown pacing at the end of each path's lifetime.

• (C) Keep the current epoch, accept (2) as a documented limit and (3) via refcounting only. Smallest possible diff (~25 lines): _PathLock wrapper + refcount, no other restructuring. Doesn't address the pool-starvation concern at all.

• (D) Push back on (1). The window is a CPython asyncio.to_thread submission race (cancel delivered between loop.create_task and the Task's first __step). Every realistic cancel I can trigger lands mid-run_in_executor (worker already started) or after the worker finished (no-op), not in the narrow pre-submit step. The post-cancel stale-write concern is already moot under the process-wide lock: even if the stale write runs, the next save will re-publish after the lock serializes it, so final state is consistent. A regression test can enforce that.

My lean is (A) — cleanest middle ground, keeps the existing shape, addresses 2+3 directly and explains 1 as a runtime edge. (B) is cleanest architecturally but adds executor lifetime management. Want me to proceed with (A), or would you rather take (B)?

[Evan Mattson (moonbox3)]

Following up on the earlier registry-growth discussion: could _file_locks release entries after the final holder or waiter exits? This commit replaces the cleaned-up per-loop entries with a process-wide dictionary that never removes a Path; because WorkflowCheckpoint creates a new UUID by default, normal saves and later delete() calls retain one lock per checkpoint forever. A reference-counted entry would preserve process-wide coordination while tying lock lifetime to active work.

[Evan Mattson (moonbox3)]

Manjunath Janardhan (@manjunathshiva) please handle all comments per guidelines.

[Manjunath Janardhan (manjunathshiva)]

Manjunath Janardhan (@manjunathshiva) please handle all comments per guidelines.

Yes, working on it