From 24eb1132a181c76af863d2b04f18f0d7ac78ccbc Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 19 Aug 2026 13:40:19 +0530 Subject: [PATCH 1/2] Python: fix: make concurrent FileCheckpointStorage saves not race on a shared temp path FileCheckpointStorage.save() wrote to and then renamed a fixed ".json.tmp" path, so concurrent saves of the same checkpoint ID raced over the shared temp file. Whichever save renamed first removed the temp file still being written by another save, which then failed with FileNotFoundError / PermissionError in os.replace. Create a unique temp file per save in the destination directory (so os.replace remains atomic), serialize same-ID writes with a per-ID lock, and retry the atomic move briefly to absorb the transient Windows background-handle PermissionError that surfaces even for fully serialized replaces. Fixes #7748 --- .../agent_framework/_workflows/_checkpoint.py | 111 +++++++++++++++++- .../core/tests/workflow/test_checkpoint.py | 54 +++++++++ 2 files changed, 160 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 3de9460c86d..5b8a073b671 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -7,7 +7,10 @@ import json import logging import os +import threading +import time import uuid +import weakref from collections.abc import Mapping from dataclasses import dataclass, field, fields from datetime import datetime, timezone @@ -246,6 +249,20 @@ async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID] return [cp.checkpoint_id for cp in self._checkpoints.values() if cp.workflow_name == workflow_name] +class _SaveLockRef: + """A reference-counted asyncio.Lock entry for per-checkpoint-ID save serialization. + + ``refs`` counts the holder plus any in-flight acquirers, so the registry entry is + only deleted once the final user releases it. + """ + + __slots__ = ("lock", "refs") + + def __init__(self) -> None: + self.lock = asyncio.Lock() + self.refs = 0 + + class FileCheckpointStorage: """File-based checkpoint storage for persistence. @@ -288,8 +305,45 @@ def __init__( self.storage_path = Path(storage_path) self.storage_path.mkdir(parents=True, exist_ok=True) self._allowed_types: frozenset[str] = frozenset(allowed_checkpoint_types or []) + # Serialize writes per checkpoint ID within each event loop. Entries are + # reference-counted: a save takes a reference before awaiting the lock and + # releases it after the write finishes, and the entry is removed once the + # last user exits. This keeps held/waited-on locks undisturbed (no eviction + # of active entries), bounds map growth to the number of in-flight saves, + # and lets the weak loop key be collected once its entries are gone. + # Locks are keyed per loop because asyncio.Lock is loop-bound. + self._save_locks_by_loop: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, dict[CheckpointID, _SaveLockRef] + ] = weakref.WeakKeyDictionary() + self._save_locks_guard = threading.Lock() logger.info(f"Initialized file checkpoint storage at {self.storage_path}") + def _acquire_save_lock_ref(self, checkpoint_id: CheckpointID) -> tuple[asyncio.AbstractEventLoop, _SaveLockRef]: + """Take a reference on the per-loop, per-checkpoint-ID lock, creating it on first use.""" + loop = asyncio.get_running_loop() + with self._save_locks_guard: + locks = self._save_locks_by_loop.get(loop) + if locks is None: + locks = {} + self._save_locks_by_loop[loop] = locks + entry = locks.get(checkpoint_id) + if entry is None: + entry = _SaveLockRef() + locks[checkpoint_id] = entry + entry.refs += 1 + return loop, entry + + def _release_save_lock_ref( + self, loop: asyncio.AbstractEventLoop, checkpoint_id: CheckpointID, entry: _SaveLockRef + ) -> None: + """Release a reference; the entry is removed once no holder or waiter remains.""" + with self._save_locks_guard: + entry.refs -= 1 + if entry.refs == 0: + locks = self._save_locks_by_loop.get(loop) + if locks is not None and locks.get(checkpoint_id) is entry: + del locks[checkpoint_id] + def _validate_file_path(self, checkpoint_id: CheckpointID) -> Path: """Validate that a checkpoint ID resolves to a path within the storage directory. @@ -325,13 +379,60 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: checkpoint_dict = checkpoint.to_dict() encoded_checkpoint = encode_checkpoint_value(checkpoint_dict) + # Take a lock reference before awaiting so waiters count toward it; the entry + # is removed by _release_save_lock_ref only when no holder or waiter remains. + loop, save_lock_ref = self._acquire_save_lock_ref(checkpoint.checkpoint_id) + + def _replace_with_retry(tmp_path: Path) -> None: + # On Windows, os.replace can transiently fail with PermissionError when a + # background indexer or AV scan briefly holds a handle to the destination + # file. The per-ID lock serializes concurrent save() calls, but the OS + # callback is still external to the process and can trip a transient + # error even when only one replace is in flight. Retry briefly to absorb it. + for attempt in range(5): + try: + os.replace(tmp_path, file_path) + return + except PermissionError: + if attempt == 4: + raise + time.sleep(0.001 * (2**attempt)) + def _write_atomic() -> None: - tmp_path = file_path.with_suffix(".json.tmp") - with open(tmp_path, "w") as f: - json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False) - os.replace(tmp_path, file_path) + # Use a unique temp file per save in the destination directory so + # concurrent saves of the same checkpoint ID never race on a shared + # temporary path, and os.replace remains atomic (same filesystem). + # A short, ID-independent name keeps every destination name accepted by + # _validate_file_path saveable regardless of checkpoint-ID length or + # filesystem limits. + tmp_path: Path | None = None + try: + # O_CREAT | O_EXCL | O_WRONLY with an explicit 0o666 mode, so the file + # is created with the process umask exactly like the previous + # open(..., "w") path was (NamedTemporaryFile would hard-code 0o600 + # and downgrade modes on POSIX after an os.replace over an existing + # checkpoint). + tmp_name = f".maf-ckpt-{uuid.uuid4().hex}.tmp" + tmp_path = file_path.parent / tmp_name + fd = os.open(tmp_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o666) + with os.fdopen(fd, "w") as f: + json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False) + _replace_with_retry(tmp_path) + tmp_path = None + finally: + if tmp_path is not None and tmp_path.exists(): + try: + tmp_path.unlink() + except OSError: + # Best-effort cleanup only; leaking a temp file is harmless + # compared to masking the original exception. + logger.debug(f"Failed to remove checkpoint temp file {tmp_path}", exc_info=True) - await asyncio.to_thread(_write_atomic) + try: + async with save_lock_ref.lock: + await asyncio.to_thread(_write_atomic) + finally: + self._release_save_lock_ref(loop, checkpoint.checkpoint_id, save_lock_ref) logger.info(f"Saved checkpoint {checkpoint.checkpoint_id} to {file_path}") return checkpoint.checkpoint_id diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 5f3da78cd1d..834770a5bd4 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import json import tempfile from dataclasses import dataclass @@ -1115,6 +1116,59 @@ async def test_file_checkpoint_storage_save_and_load(): assert loaded_checkpoint.pending_request_info_events == checkpoint.pending_request_info_events +async def test_file_checkpoint_storage_concurrent_saves_same_id(): + """Concurrent saves of the same checkpoint ID must not fail on a shared temp path. + + Regression for https://github.com/microsoft/agent-framework/issues/7748: + FileCheckpointStorage.save() used a fixed `.json.tmp` temp path, so concurrent + saves raced on it (one rename removed it before another's rename). Uses enough + concurrent saves to reliably trip the race on the unfixed code. + """ + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + checkpoint = WorkflowCheckpoint( + workflow_name="test-workflow", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) + + results = await asyncio.gather(*(storage.save(checkpoint) for _ in range(50)), return_exceptions=True) + + errors = [r for r in results if isinstance(r, BaseException)] + assert not errors, f"concurrent saves raised internal filesystem errors: {errors[:1]!r}" + assert all(r == "shared-id" for r in results) + # One of the saves won; the destination is intact and parseable, not corrupted or truncated. + assert (Path(temp_dir) / "shared-id.json").exists() + loaded = await storage.load("shared-id") + assert loaded.checkpoint_id == checkpoint.checkpoint_id + assert loaded.workflow_name == checkpoint.workflow_name + assert loaded.graph_signature_hash == checkpoint.graph_signature_hash + + +async def test_file_checkpoint_storage_save_lock_registry_bounded(): + """Save-lock registry entries must be released after each save completes. + + The reference-counted bookkeeping in FileCheckpointStorage must remove the entry + for a checkpoint ID once the final holder exits, so successive saves of many + distinct IDs do not accumulate entries in the per-loop registry. + """ + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + for i in range(50): + checkpoint = WorkflowCheckpoint( + workflow_name="test-workflow", + graph_signature_hash="test-hash", + checkpoint_id=f"checkpoint-{i}", + ) + await storage.save(checkpoint) + + # Each save takes and releases its lock reference inside save() itself, so + # by the time we observe the registry here every entry should be gone. + loop = asyncio.get_running_loop() + locks = storage._save_locks_by_loop.get(loop) # pyright: ignore[reportPrivateUsage] + assert not locks + + async def test_file_checkpoint_storage_load_nonexistent(): with tempfile.TemporaryDirectory() as temp_dir: storage = FileCheckpointStorage(temp_dir) From 723c17bb70bccc3f9e661bfd868b562fcc3de5b8 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Thu, 20 Aug 2026 13:01:39 +0530 Subject: [PATCH 2/2] Python: fix: lock file writes at destination path, shield workers from cancellation Address two post-merge review concerns on FileCheckpointStorage.save(): - The per-event-loop, per-instance asyncio.Lock registry could not serialize two FileCheckpointStorage instances pointed at the same directory, nor a single instance driven from two event loops. Replace it with a process-wide, per-canonical-path threading.Lock registry (lazily populated, bounded by distinct destinations actually written). Because asyncio.to_thread runs the actual file I/O on a worker thread, the threading lock serializes correctly across coroutines, loops, and instances and spans the entire open + write + os.replace window. - Caller-side cancellation previously released the write lock early: a CancelledError delivered inside 'await asyncio.to_thread(...)' exited the 'async with' and ran _release_save_lock_ref while the OS thread kept writing. The next save for the same checkpoint_id could then reach os.replace concurrently, reintroducing the PermissionError race and possibly landing stale data over a newer write. Shield the worker so cancellation propagates only after the in-flight write completes. Also replace the refcount-aware registry test with: - registry-bounded test for the new invariant (one lock per destination) - cross-instance concurrent save test (two storages, same directory) - gated os.replace cancel-race test (deterministically parks A's worker mid-publish, confirms save B cannot reach os.replace until A completes) Fixes #7748 --- .../agent_framework/_workflows/_checkpoint.py | 168 ++++++++---------- .../core/tests/workflow/test_checkpoint.py | 157 ++++++++++++++-- 2 files changed, 223 insertions(+), 102 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 5b8a073b671..81124186a4d 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -10,7 +10,6 @@ import threading import time import uuid -import weakref from collections.abc import Mapping from dataclasses import dataclass, field, fields from datetime import datetime, timezone @@ -249,18 +248,31 @@ async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID] return [cp.checkpoint_id for cp in self._checkpoints.values() if cp.workflow_name == workflow_name] -class _SaveLockRef: - """A reference-counted asyncio.Lock entry for per-checkpoint-ID save serialization. - - ``refs`` counts the holder plus any in-flight acquirers, so the registry entry is - only deleted once the final user releases it. - """ - - __slots__ = ("lock", "refs") - - def __init__(self) -> None: - self.lock = asyncio.Lock() - self.refs = 0 +# Process-wide serialization of os.replace() per destination file. +# +# asyncio.Lock is loop-bound, so a per-(loop, checkpoint-id) registry (the previous +# design) could not serialize two FileCheckpointStorage instances pointed at the +# same directory, nor one instance driven from two event loops. A threading.Lock +# keyed by the canonical destination path *does* span coroutines, loops, and +# instances because asyncio.to_thread runs the actual file write on a worker +# thread, and threading primitives serialize across those. +# +# Locks are created lazily on first save and never removed. The registry grows +# by at most one entry per *distinct* checkpoint file ever written; that is +# bounded by the number of files actually present under any FileCheckpointStorage +# directory the process touches — a working-set bound, not unbounded. +_file_locks: dict[Path, threading.Lock] = {} +_file_locks_guard = threading.Lock() + + +def _get_file_lock(file_path: Path) -> threading.Lock: + """Return the process-wide lock guarding *file_path*, creating it on first use.""" + with _file_locks_guard: + lock = _file_locks.get(file_path) + if lock is None: + lock = threading.Lock() + _file_locks[file_path] = lock + return lock class FileCheckpointStorage: @@ -305,45 +317,8 @@ def __init__( self.storage_path = Path(storage_path) self.storage_path.mkdir(parents=True, exist_ok=True) self._allowed_types: frozenset[str] = frozenset(allowed_checkpoint_types or []) - # Serialize writes per checkpoint ID within each event loop. Entries are - # reference-counted: a save takes a reference before awaiting the lock and - # releases it after the write finishes, and the entry is removed once the - # last user exits. This keeps held/waited-on locks undisturbed (no eviction - # of active entries), bounds map growth to the number of in-flight saves, - # and lets the weak loop key be collected once its entries are gone. - # Locks are keyed per loop because asyncio.Lock is loop-bound. - self._save_locks_by_loop: weakref.WeakKeyDictionary[ - asyncio.AbstractEventLoop, dict[CheckpointID, _SaveLockRef] - ] = weakref.WeakKeyDictionary() - self._save_locks_guard = threading.Lock() logger.info(f"Initialized file checkpoint storage at {self.storage_path}") - def _acquire_save_lock_ref(self, checkpoint_id: CheckpointID) -> tuple[asyncio.AbstractEventLoop, _SaveLockRef]: - """Take a reference on the per-loop, per-checkpoint-ID lock, creating it on first use.""" - loop = asyncio.get_running_loop() - with self._save_locks_guard: - locks = self._save_locks_by_loop.get(loop) - if locks is None: - locks = {} - self._save_locks_by_loop[loop] = locks - entry = locks.get(checkpoint_id) - if entry is None: - entry = _SaveLockRef() - locks[checkpoint_id] = entry - entry.refs += 1 - return loop, entry - - def _release_save_lock_ref( - self, loop: asyncio.AbstractEventLoop, checkpoint_id: CheckpointID, entry: _SaveLockRef - ) -> None: - """Release a reference; the entry is removed once no holder or waiter remains.""" - with self._save_locks_guard: - entry.refs -= 1 - if entry.refs == 0: - locks = self._save_locks_by_loop.get(loop) - if locks is not None and locks.get(checkpoint_id) is entry: - del locks[checkpoint_id] - def _validate_file_path(self, checkpoint_id: CheckpointID) -> Path: """Validate that a checkpoint ID resolves to a path within the storage directory. @@ -379,16 +354,13 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: checkpoint_dict = checkpoint.to_dict() encoded_checkpoint = encode_checkpoint_value(checkpoint_dict) - # Take a lock reference before awaiting so waiters count toward it; the entry - # is removed by _release_save_lock_ref only when no holder or waiter remains. - loop, save_lock_ref = self._acquire_save_lock_ref(checkpoint.checkpoint_id) - def _replace_with_retry(tmp_path: Path) -> None: # On Windows, os.replace can transiently fail with PermissionError when a # background indexer or AV scan briefly holds a handle to the destination - # file. The per-ID lock serializes concurrent save() calls, but the OS - # callback is still external to the process and can trip a transient - # error even when only one replace is in flight. Retry briefly to absorb it. + # file. The process-wide per-path lock serializes concurrent save() calls + # to the same destination, but the OS callback is still external to the + # process and can trip a transient error even when only one replace is in + # flight. Retry briefly to absorb it. for attempt in range(5): try: os.replace(tmp_path, file_path) @@ -399,40 +371,54 @@ def _replace_with_retry(tmp_path: Path) -> None: time.sleep(0.001 * (2**attempt)) def _write_atomic() -> None: - # Use a unique temp file per save in the destination directory so - # concurrent saves of the same checkpoint ID never race on a shared - # temporary path, and os.replace remains atomic (same filesystem). - # A short, ID-independent name keeps every destination name accepted by - # _validate_file_path saveable regardless of checkpoint-ID length or - # filesystem limits. - tmp_path: Path | None = None - try: - # O_CREAT | O_EXCL | O_WRONLY with an explicit 0o666 mode, so the file - # is created with the process umask exactly like the previous - # open(..., "w") path was (NamedTemporaryFile would hard-code 0o600 - # and downgrade modes on POSIX after an os.replace over an existing - # checkpoint). - tmp_name = f".maf-ckpt-{uuid.uuid4().hex}.tmp" - tmp_path = file_path.parent / tmp_name - fd = os.open(tmp_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o666) - with os.fdopen(fd, "w") as f: - json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False) - _replace_with_retry(tmp_path) - tmp_path = None - finally: - if tmp_path is not None and tmp_path.exists(): - try: - tmp_path.unlink() - except OSError: - # Best-effort cleanup only; leaking a temp file is harmless - # compared to masking the original exception. - logger.debug(f"Failed to remove checkpoint temp file {tmp_path}", exc_info=True) - - try: - async with save_lock_ref.lock: - await asyncio.to_thread(_write_atomic) - finally: - self._release_save_lock_ref(loop, checkpoint.checkpoint_id, save_lock_ref) + # The threading lock here is the heartbeat of cross-instance/cross-loop + # safety: a same-directory save racing through a different + # FileCheckpointStorage instance — or from another event loop in the + # same process — contends on the same canonical destination path and + # therefore on the same lock. Holding it across the entire open + write + # + replace keeps no window where a second writer can briefly see a + # half-published temp file or reach os.replace concurrently. + with _get_file_lock(file_path): + # Use a unique temp file per save in the destination directory so + # concurrent saves of distinct checkpoint IDs never contend on a + # shared temporary path, and os.replace remains atomic (same + # filesystem). A short, ID-independent name keeps every destination + # name accepted by _validate_file_path saveable regardless of + # checkpoint-ID length or filesystem limits. + tmp_path: Path | None = None + try: + # O_CREAT | O_EXCL | O_WRONLY with an explicit 0o666 mode, so the + # file is created with the process umask exactly like the previous + # open(..., "w") path was (NamedTemporaryFile would hard-code 0o600 + # and downgrade modes on POSIX after an os.replace over an existing + # checkpoint). + tmp_name = f".maf-ckpt-{uuid.uuid4().hex}.tmp" + tmp_path = file_path.parent / tmp_name + fd = os.open(tmp_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o666) + with os.fdopen(fd, "w") as f: + json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False) + _replace_with_retry(tmp_path) + tmp_path = None + finally: + if tmp_path is not None and tmp_path.exists(): + try: + tmp_path.unlink() + except OSError: + # Best-effort cleanup only; leaking a temp file is harmless + # compared to masking the original exception. + logger.debug(f"Failed to remove checkpoint temp file {tmp_path}", exc_info=True) + + # Shield the worker from caller-side cancellation: without the shield, a + # cancellation delivered while the coroutine is suspended inside + # asyncio.to_thread exits the await but leaves the OS thread running, so + # its os.replace can still come in *after* the caller has been cancelled + # and a subsequent save for the same checkpoint ID has started — on + # Windows that reintroduces the PermissionError race this path exists to + # avoid, and it can also overwrite a newer checkpoint with stale data. + # Shielding guarantees the in-flight write completes (or fails) before + # the caller observes a result, so the order seen on disk matches the + # order callers observed. + await asyncio.shield(asyncio.to_thread(_write_atomic)) logger.info(f"Saved checkpoint {checkpoint.checkpoint_id} to {file_path}") return checkpoint.checkpoint_id diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 834770a5bd4..2a52e8a833b 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -2,6 +2,7 @@ import asyncio import json +import os import tempfile from dataclasses import dataclass from datetime import datetime, timezone @@ -1146,27 +1147,161 @@ async def test_file_checkpoint_storage_concurrent_saves_same_id(): async def test_file_checkpoint_storage_save_lock_registry_bounded(): - """Save-lock registry entries must be released after each save completes. + """Process-wide file-lock registry must not grow with sequential saves of the same ID. - The reference-counted bookkeeping in FileCheckpointStorage must remove the entry - for a checkpoint ID once the final holder exits, so successive saves of many - distinct IDs do not accumulate entries in the per-loop registry. + Regression guard for the post-#7748 design: `FileCheckpointStorage` uses a + module-level `_file_locks` dict keyed by canonical destination path, with one + lazy `threading.Lock` created on first save. Saving the same checkpoint ID + repeatedly must reuse that one lock, never append new entries. """ with tempfile.TemporaryDirectory() as temp_dir: storage = FileCheckpointStorage(temp_dir) - for i in range(50): + canonical = (Path(temp_dir) / "shared-id.json").resolve() + + from agent_framework._workflows import _checkpoint as checkpoint_module + + initial_size = len(checkpoint_module._file_locks) # pyright: ignore[reportPrivateUsage] + + for _ in range(50): checkpoint = WorkflowCheckpoint( workflow_name="test-workflow", graph_signature_hash="test-hash", - checkpoint_id=f"checkpoint-{i}", + checkpoint_id="shared-id", ) await storage.save(checkpoint) - # Each save takes and releases its lock reference inside save() itself, so - # by the time we observe the registry here every entry should be gone. - loop = asyncio.get_running_loop() - locks = storage._save_locks_by_loop.get(loop) # pyright: ignore[reportPrivateUsage] - assert not locks + # Saving the same ID repeatedly creates at most one lock registry entry. + assert len(checkpoint_module._file_locks) <= initial_size + 1 # pyright: ignore[reportPrivateUsage] + assert canonical in checkpoint_module._file_locks # pyright: ignore[reportPrivateUsage] + + +async def test_file_checkpoint_storage_concurrent_saves_across_instances(): + """Two FileCheckpointStorage instances to the same directory must serialize same-ID saves. + + Companion regression for a reviewer concern raised while fixing #7748: the + previous per-instance, per-event-loop lock registry did not span a second + FileCheckpointStorage instance pointed at the same directory, so concurrent + saves could still reach os.replace together and trip the Windows PermissionError + race. The fix switched to a process-wide, per-destination threading.Lock + keyed by canonical path. Both instances must complete all saves without + surfacing filesystem errors. + """ + with tempfile.TemporaryDirectory() as temp_dir: + storage_a = FileCheckpointStorage(temp_dir) + storage_b = FileCheckpointStorage(temp_dir) + checkpoint = WorkflowCheckpoint( + workflow_name="test-workflow", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) + + results = await asyncio.gather( + *[storage_a.save(checkpoint) for _ in range(25)], + *[storage_b.save(checkpoint) for _ in range(25)], + return_exceptions=True, + ) + + errors = [r for r in results if isinstance(r, BaseException)] + assert not errors, f"cross-instance concurrent saves raised: {errors[:1]!r}" + assert all(r == "shared-id" for r in results) + assert (Path(temp_dir) / "shared-id.json").exists() + + loaded = await storage_a.load("shared-id") + assert loaded.checkpoint_id == checkpoint.checkpoint_id + assert loaded.workflow_name == checkpoint.workflow_name + + +async def test_file_checkpoint_storage_cancel_does_not_expose_race(monkeypatch): + """Cancelling save() mid-write must not let a later save race the in-flight worker. + + Addressing a reviewer concern raised while fixing #7748: the fix holds the + destination file's threading lock *inside* the worker thread (around the + open + write + os.replace) and shields the worker, so a caller-side + cancellation cannot release the lock early or interrupt the write. The + regression gate monkeypatches ``os.replace`` inside the checkpoint module to + make save A's worker deterministically block at the publish step; save B is + then issued after A's caller was cancelled. Correct behavior requires that + B's worker cannot reach os.replace until A's worker completes: otherwise the + two replaces run concurrently on Windows (PermissionError race) or A's stale + data lands after B. + """ + import threading + + from agent_framework._workflows import _checkpoint as checkpoint_module + + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + real_replace = checkpoint_module.os.replace + replace_started = threading.Event() + release_first_replace = threading.Event() + replace_calls: list[tuple[str, str]] = [] + calls_guard = threading.Lock() + first_call_blocked = threading.Event() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + with calls_guard: + ordinal = len(replace_calls) + 1 + replace_calls.append((os.path.basename(str(src)), os.path.basename(str(dst)))) + if ordinal == 1 and not first_call_blocked.is_set(): + first_call_blocked.set() + replace_started.set() + assert release_first_replace.wait(timeout=10) + return real_replace(src, dst) + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + checkpoint_a = WorkflowCheckpoint( + workflow_name="workflow-a", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) + checkpoint_b = WorkflowCheckpoint( + workflow_name="workflow-b", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) + + task_a = asyncio.create_task(storage.save(checkpoint_a)) + # Wait until A's worker is parked inside os.replace (write in flight). + started = await asyncio.to_thread(replace_started.wait, 10) + assert started, "save A's worker never reached os.replace" + + # Cancel A's caller while its worker is mid-publish. Shielding keeps the + # worker alive; the caller observes CancelledError. + task_a.cancel() + with pytest.raises(asyncio.CancelledError): + await task_a + + # Issue save B and give its worker a real chance to reach os.replace. + # If cancellation had freed the write lock, B's replace would start + # while A is still parked and replace_calls would grow to 2. + task_b_started = asyncio.Event() + task_b_done = asyncio.Event() + + async def run_b() -> None: + task_b_started.set() + await storage.save(checkpoint_b) + task_b_done.set() + + task_b = asyncio.create_task(run_b()) + await task_b_started.wait() + # Brief, bounded: B must remain blocked on the destination lock. A short + # poll window is sufficient because lock handoff is synchronous. + for _ in range(50): + if len(replace_calls) >= 2: + pytest.fail("save B reached os.replace while save A's worker was still mid-write") + await asyncio.sleep(0.01) + assert not task_b_done.is_set() + + # Let A's worker finish; B follows, serialized. Final state is B's data. + release_first_replace.set() + await asyncio.wait_for(task_b, timeout=10) + assert task_b_done.is_set() + assert len(replace_calls) == 2 + + result = await storage.load("shared-id") + assert result.workflow_name == "workflow-b" async def test_file_checkpoint_storage_load_nonexistent():