From 24eb1132a181c76af863d2b04f18f0d7ac78ccbc Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 19 Aug 2026 13:40:19 +0530 Subject: [PATCH 1/6] 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/6] 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(): From e4f30b8858893b6c47f2d3a62c30315c99b66c2f Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Tue, 8 Sep 2026 23:06:45 +0530 Subject: [PATCH 3/6] Python: take checkpoint destination ownership before submitting the write Implements the revised A shape chosen in review: process-wide coordination keyed by canonical destination, asynchronous waiters, and registry entries released after the final queued or running operation. Three review findings shared one cause -- ownership was established inside the executor thread rather than before submission. Acquiring the destination lock inside the worker meant every same-path save occupied an `asyncio.to_thread` worker while merely waiting, so a burst could fill the default pool and stall unrelated work including checkpoint loads, and could deadlock once the write that had to finish first was queued behind those waiters. Ownership now comes from a queue keyed by the canonical path and is taken before the write is submitted, so only the active write holds a worker. Waiters suspend on `asyncio.wrap_future` over a `concurrent.futures.Future`, which is not bound to a loop, so one chain also orders writers enqueued from different event loops -- something a per-path lock could not do. The registry was a dict that never removed an entry. Because `WorkflowCheckpoint` generates a fresh UUID by default, every save retained one entry for the life of the process. Entries are now reference-counted by queued-or-running operations and dropped by the last release. Cancellation is handled explicitly, both ways the review allowed. A save cancelled while still waiting for the destination has nothing in the executor, so it is simply removed -- and it still hands ownership on, or every later save for that path would wait forever. A save cancelled mid-write drains its worker before propagating, because releasing first would let the next save start while this replace is still in flight. That drain absorbs re-delivered cancellations: once a task has a cancellation pending its next await raises immediately, so a single shielded await returns with the worker still running and the write can still land after ownership is released. 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 no write was submitted. Releasing only from the coroutine regressed on the design it replaced: a `threading.Lock` was released by the worker thread, which outlives the loop that submitted it, whereas a coroutine that never resumes -- its loop closed with the task still pending -- would leave the destination owned for the life of the process and hang every later save to it. The two releasers are idempotent, so exactly one of them accounts for each ticket. Tests are deterministic and cover each case: registry release for repeated, distinct-ID and concurrent saves; queued saves not occupying executor threads, gated by a two-worker executor with three saves in flight; serialization across two real event loops; cancellation before submission writing nothing and not stalling the follower; repeated cancellation still draining; a write failing while draining being reported rather than lost; a failed save releasing its destination and leaving the path usable; recovery of a destination whose loop was closed mid-write; and the defensive release path. --- .../agent_framework/_workflows/_checkpoint.py | 269 +++++--- .../core/tests/workflow/test_checkpoint.py | 633 +++++++++++++++--- 2 files changed, 748 insertions(+), 154 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 8dc558b04b4..3d64840b12b 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -11,6 +11,7 @@ import time import uuid from collections.abc import Mapping +from concurrent.futures import Future from dataclasses import dataclass, field, fields from datetime import datetime, timezone from pathlib import Path @@ -249,31 +250,128 @@ 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] -# Process-wide serialization of os.replace() per destination file. +# Process-wide serialization of writes 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. +# asyncio.Lock is loop-bound, so a per-(loop, checkpoint-id) registry cannot +# serialize two FileCheckpointStorage instances pointed at the same directory, nor one +# instance driven from two event loops. Ownership is therefore handed out by a +# process-wide queue keyed by the canonical destination path, and taken *before* the +# write is submitted to a worker thread rather than inside it. # -# 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 +# Waiting on a ``concurrent.futures.Future`` through ``asyncio.wrap_future`` is what +# makes that work in both directions: it is not bound to a loop, so a single chain +# orders writers regardless of which loop enqueued them, and waiters suspend on the +# event loop instead of occupying an ``asyncio.to_thread`` worker. Blocking on a +# ``threading.Lock`` inside the worker instead -- the previous design -- let a burst of +# same-path saves fill the default executor and stall unrelated ``to_thread`` work, +# including checkpoint loads, and could deadlock once the write that had to finish +# first was queued behind those waiters. +# +# Entries are reference-counted by queued-or-running operations and dropped when the +# last one releases, so a process that saves many distinct checkpoint IDs -- the +# default, since ``WorkflowCheckpoint`` generates a fresh UUID -- does not retain an +# entry per ID for its lifetime. +# +# The scope really is one process. Several replicas sharing a checkpoint directory -- a +# mounted volume, a network share -- get no serialization from this, because there is no +# shared state between them to coordinate through. Concurrent saves of the same +# checkpoint ID from different processes still rely on ``os.replace`` being atomic on +# the underlying filesystem for the file not to be seen half-written; which write +# survives is undefined. +_destination_queues: dict[Path, _DestinationQueue] = {} +_destination_queues_guard = threading.Lock() + + +@dataclass +class _DestinationQueue: + """FIFO ownership hand-off for one destination path.""" + + #: Completion signal of the most recently enqueued operation, awaited by the next. + tail: Future[None] | None = None + #: Operations queued or running for this path; the entry is dropped at zero. + pending: int = 0 + + +@dataclass +class _WriteTicket: + """One operation's place in a destination's queue.""" + + path: Path + #: Signal to await before taking ownership; ``None`` when the queue was empty. + predecessor: Future[None] | None + #: Signal this operation resolves to hand ownership to its successor. + completion: Future[None] + #: Set by whichever of the worker thread or the coroutine releases first. + released: bool = False + + +def _enqueue_write(file_path: Path) -> _WriteTicket: + """Take a place in *file_path*'s queue without waiting for it.""" + completion: Future[None] = Future() + with _destination_queues_guard: + queue = _destination_queues.get(file_path) + if queue is None: + queue = _DestinationQueue() + _destination_queues[file_path] = queue + predecessor = queue.tail + queue.tail = completion + queue.pending += 1 + return _WriteTicket(path=file_path, predecessor=predecessor, completion=completion) + + +async def _drain_cancelled_write(worker: asyncio.Future[None], file_path: Path) -> None: + """Wait for *worker* to finish while this task is being cancelled. + + A single ``await asyncio.shield(worker)`` is not enough here. Once the enclosing + task has a cancellation pending, the next await raises ``CancelledError`` again + immediately, so the shield returns without the worker having finished and the write + can still land after ownership is released. Absorbing the re-delivered + cancellations until the worker is genuinely done is what makes the drain real. + """ + while not worker.done(): + try: + await asyncio.shield(worker) + except asyncio.CancelledError: + continue + except Exception: + break + if worker.done() and not worker.cancelled(): + exception = worker.exception() + if exception is not None: + # The caller is being cancelled and will not see this, but a write that + # failed while draining should not vanish silently. + logger.warning(f"Checkpoint write to {file_path} failed while draining after cancellation: {exception!r}") + + +def _release_write(ticket: _WriteTicket) -> None: + """Hand ownership to the next waiter and drop the entry once nothing is queued. + + Called from two places, whichever gets there first, and idempotent so both can: + + * the worker thread, as the last act of the write itself. This is what keeps a + destination usable if the event loop that submitted the write goes away before the + coroutine can resume -- the thread finishes and releases regardless, where a + release that only happened in the coroutine would leave the path owned forever and + hang every later save to it. + * the coroutine, on any exit path. This is the only releaser when the write was + never submitted, which is the case for a save cancelled while still queued. + + Resolving the completion signal is what keeps the chain moving, so an operation that + is cancelled or fails must still release rather than stall every later writer. + """ + with _destination_queues_guard: + if ticket.released: + return + ticket.released = True + queue = _destination_queues.get(ticket.path) + if queue is not None: + queue.pending -= 1 + if queue.pending <= 0: + del _destination_queues[ticket.path] + # Outside the guard: waking the successor must not happen while holding the lock its + # own release will need. + if not ticket.completion.done(): + ticket.completion.set_result(None) class FileCheckpointStorage: @@ -358,10 +456,10 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: 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 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. + # file. The destination queue serializes concurrent save() calls to the same + # path, 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) @@ -372,54 +470,75 @@ def _replace_with_retry(tmp_path: Path) -> None: time.sleep(0.001 * (2**attempt)) def _write_atomic() -> None: - # 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)) + # No lock here: ownership of the destination is already held by the caller + # of this function, taken from the process-wide queue before the write was + # submitted. Acquiring it on the worker instead is what let same-path saves + # pile up inside the executor. + # + # 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) + + def _write_atomic_and_release(ticket: _WriteTicket) -> None: + # Release on the worker thread, as the write's last act. The coroutine's + # `finally` would otherwise be the only releaser, and it never runs if the + # loop that submitted this write is gone -- leaving the destination owned + # forever and hanging every later save to it. + try: + _write_atomic() + finally: + _release_write(ticket) + + ticket = _enqueue_write(file_path) + try: + if ticket.predecessor is not None: + # Suspends on the event loop, not on an executor worker, and orders this + # save behind every earlier one for the same destination even when they + # were enqueued from a different loop. A cancellation delivered here + # propagates with nothing submitted and nothing written, which is the + # cheapest correct outcome: `finally` hands ownership straight to the + # next waiter. + await asyncio.wrap_future(ticket.predecessor) + + worker = asyncio.ensure_future(asyncio.to_thread(_write_atomic_and_release, ticket)) + try: + # Shield so a cancellation arriving mid-write cannot leave the worker + # running past this frame: its os.replace would otherwise land after the + # caller returned, overwriting whatever a later save had published. + await asyncio.shield(worker) + except asyncio.CancelledError: + # Ownership is still held, so drain before propagating. Releasing first + # would let the next save begin while this replace is still in flight, + # which is the race the queue exists to prevent. + await _drain_cancelled_write(worker, file_path) + raise + finally: + _release_write(ticket) 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 2a52e8a833b..488937a4c27 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -2,8 +2,10 @@ import asyncio import json +import logging import os import tempfile +from concurrent.futures import Future as ConcurrentFuture from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -1146,33 +1148,54 @@ async def test_file_checkpoint_storage_concurrent_saves_same_id(): assert loaded.graph_signature_hash == checkpoint.graph_signature_hash -async def test_file_checkpoint_storage_save_lock_registry_bounded(): - """Process-wide file-lock registry must not grow with sequential saves of the same ID. +async def test_file_checkpoint_storage_destination_queue_registry_released(): + """The destination registry must be empty again once nothing is queued or running. - 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. + Reviewer concern on #7757: the previous design kept a process-wide dict of + `threading.Lock` keyed by destination path and never removed an entry. Because + `WorkflowCheckpoint` generates a fresh UUID by default, every save retained one + entry for the lifetime of the process. Entries are now reference-counted by + queued-or-running operations and dropped by the last release, so repeated saves of + one ID and saves of many distinct IDs both settle back to nothing held. """ with tempfile.TemporaryDirectory() as temp_dir: storage = FileCheckpointStorage(temp_dir) - 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] + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + baseline = len(registry) - for _ in range(50): - checkpoint = WorkflowCheckpoint( - workflow_name="test-workflow", - graph_signature_hash="test-hash", - checkpoint_id="shared-id", + for _ in range(20): + await storage.save( + WorkflowCheckpoint( + workflow_name="test-workflow", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) ) - await storage.save(checkpoint) - - # 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] + assert len(registry) == baseline, "same-ID saves left an entry behind" + + # The default: every checkpoint gets its own UUID, so each save is a distinct + # destination. This is the case that grew without bound before. + for _ in range(20): + await storage.save(WorkflowCheckpoint(workflow_name="test-workflow", graph_signature_hash="test-hash")) + assert len(registry) == baseline, "distinct-ID saves leaked registry entries" + + # Concurrent same-path saves share one entry while queued, and release it. + await asyncio.gather( + *( + storage.save( + WorkflowCheckpoint( + workflow_name="test-workflow", + graph_signature_hash="test-hash", + checkpoint_id="concurrent-id", + ) + ) + for _ in range(8) + ) + ) + assert len(registry) == baseline, "concurrent saves left an entry behind" async def test_file_checkpoint_storage_concurrent_saves_across_instances(): @@ -1211,19 +1234,15 @@ async def test_file_checkpoint_storage_concurrent_saves_across_instances(): 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. +async def test_file_checkpoint_storage_cancel_drains_before_releasing(monkeypatch): + """A cancelled save must not release the destination while its write is in flight. - 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. + Reviewer concern on #7757: shielding alone let the caller observe `CancelledError` + while the worker kept running, so a later save on another loop could take the + destination, complete, and then be overwritten when the cancelled worker finally + ran its `os.replace`. The cancellation path now drains the worker before + propagating, which means the cancelled caller does not return until its own write + has finished -- so nothing it wrote can land after a later save. """ import threading @@ -1235,73 +1254,56 @@ async def test_file_checkpoint_storage_cancel_does_not_expose_race(monkeypatch): real_replace = checkpoint_module.os.replace replace_started = threading.Event() release_first_replace = threading.Event() - replace_calls: list[tuple[str, str]] = [] + # `_replace_with_retry` may call os.replace more than once for a single write, so + # count completed writes rather than attempts. + attempts = 0 + replace_calls: list[str] = [] calls_guard = threading.Lock() - first_call_blocked = threading.Event() def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + nonlocal attempts 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() + attempts += 1 + first = attempts == 1 + if first: replace_started.set() assert release_first_replace.wait(timeout=10) - return real_replace(src, dst) + real_replace(src, dst) + with calls_guard: + replace_calls.append(os.path.basename(str(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", - ) + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, 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" + task_a = asyncio.create_task(storage.save(make("workflow-a"))) + assert await asyncio.to_thread(replace_started.wait, 10), "save A 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 + # The drain is the point: A must still be running, holding the destination, + # because its own worker has not finished. + for _ in range(20): + await asyncio.sleep(0.01) + if task_a.done(): + break + assert not task_a.done(), "cancelled save returned while its write was still in flight" - # 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") + # A later save cannot take the destination while A is draining. + task_b = asyncio.create_task(storage.save(make("workflow-b"))) + for _ in range(20): await asyncio.sleep(0.01) - assert not task_b_done.is_set() + assert attempts == 1, "save B submitted a write while A was still draining" + assert not task_b.done() - # Let A's worker finish; B follows, serialized. Final state is B's data. release_first_replace.set() + with pytest.raises(asyncio.CancelledError): + await task_a 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" + assert len(replace_calls) == 2 + # B ran second, so B's data is what survives. + assert (await storage.load("shared-id")).workflow_name == "workflow-b" async def test_file_checkpoint_storage_load_nonexistent(): @@ -1955,3 +1957,476 @@ async def test_file_checkpoint_storage_roundtrip_empty_collections(): # endregion + + +async def test_file_checkpoint_storage_queued_saves_do_not_occupy_executor_threads(monkeypatch): + """Queued same-path saves must wait on the event loop, not inside the executor. + + Reviewer concern on #7757: the previous design acquired the destination lock inside + the worker, so every same-path save occupied an `asyncio.to_thread` worker while + merely waiting. A burst could fill the default pool and stall unrelated work -- + including checkpoint loads -- and deadlock once the write that had to finish first + was queued behind those waiters. Ownership is now taken before submission, so only + the active write holds a worker. + + The gate is a deliberately small executor: with three same-path saves in flight and + only two workers, an unrelated `to_thread` call still has to get a thread. + """ + import threading + from concurrent.futures import ThreadPoolExecutor + + 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 + first_replace_reached = threading.Event() + release_first_replace = threading.Event() + inside_worker = 0 + max_inside_worker = 0 + counter_guard = threading.Lock() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + nonlocal inside_worker, max_inside_worker + with counter_guard: + inside_worker += 1 + max_inside_worker = max(max_inside_worker, inside_worker) + first = inside_worker == 1 and not first_replace_reached.is_set() + try: + if first: + first_replace_reached.set() + assert release_first_replace.wait(timeout=10) + real_replace(src, dst) + finally: + with counter_guard: + inside_worker -= 1 + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="ckpt-test") + # `set_default_executor(None)` is rejected, so the original has to be put back + # by attribute. Routed through an untyped local rather than ignore comments. + loop: Any = asyncio.get_running_loop() + previous_executor = loop._default_executor + loop.set_default_executor(executor) + try: + saves = [ + asyncio.create_task( + storage.save( + WorkflowCheckpoint( + workflow_name=f"w{index}", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) + ) + ) + for index in range(3) + ] + assert await asyncio.to_thread(first_replace_reached.wait, 10) + + # The decisive assertion: two saves are queued behind the parked one, and an + # unrelated to_thread call must still be scheduled. Under the old design the + # waiters held both workers and this timed out. + marker = await asyncio.wait_for(asyncio.to_thread(lambda: "scheduled"), timeout=5) + assert marker == "scheduled" + + release_first_replace.set() + await asyncio.wait_for(asyncio.gather(*saves), timeout=10) + assert max_inside_worker == 1, f"{max_inside_worker} writes ran concurrently for one destination" + finally: + loop._default_executor = previous_executor + executor.shutdown(wait=True) + + +async def test_file_checkpoint_storage_cancel_before_submission_writes_nothing(monkeypatch): + """A save cancelled while queued must never submit its write, and must not stall the queue. + + This is the other half of the cancellation contract: the reviewer asked that a + cancelled write either be removed before submission or drained. A save cancelled + while still waiting for the destination has nothing in the executor, so it is simply + removed -- and it still has to hand ownership on, or every later save for that path + would wait forever. + """ + 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 + first_replace_reached = threading.Event() + release_first_replace = threading.Event() + # `_replace_with_retry` may call os.replace more than once for a single write, so + # count completed writes rather than attempts. + attempts = 0 + written: list[str] = [] + guard = threading.Lock() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + nonlocal attempts + with guard: + attempts += 1 + first = attempts == 1 + if first: + first_replace_reached.set() + assert release_first_replace.wait(timeout=10) + real_replace(src, dst) + with guard: + written.append(os.path.basename(str(dst))) + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id") + + holder = asyncio.create_task(storage.save(make("holder"))) + assert await asyncio.to_thread(first_replace_reached.wait, 10) + + queued = asyncio.create_task(storage.save(make("cancelled"))) + follower = asyncio.create_task(storage.save(make("follower"))) + + # Wait on observable state rather than a sleep: both saves must actually be + # queued behind the holder before cancelling, or this would be testing a + # cancellation that raced the enqueue instead. + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + canonical = (Path(temp_dir) / "shared-id.json").resolve() + for _ in range(400): + entry = registry.get(canonical) + if entry is not None and entry.pending == 3: + break + await asyncio.sleep(0.005) + entry = registry.get(canonical) + assert entry is not None and entry.pending == 3, ( + f"expected holder + two queued saves, got {entry.pending if entry else 'no entry'}" + ) + + queued.cancel() + with pytest.raises(asyncio.CancelledError): + await queued + # Nothing was submitted for the cancelled save: the only write in flight is the + # holder's, still parked before its replace. + assert attempts == 1, "the cancelled save submitted a write" + assert written == [] + + release_first_replace.set() + await asyncio.wait_for(asyncio.gather(holder, follower), timeout=10) + + # Two writes total -- the holder and the follower. The cancelled save contributed + # none, and crucially did not stall the follower behind it. + assert len(written) == 2 + assert (await storage.load("shared-id")).workflow_name == "follower" + assert not registry, "a cancelled save left its destination entry behind" + + +def test_file_checkpoint_storage_serializes_across_event_loops(monkeypatch, tmp_path): + """Saves driven from separate event loops must still serialize per destination. + + Reviewer concern on #7757: ownership has to be established by something that is not + bound to a loop, or two workers queued from different loops can reach `os.replace` + 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. + + Deliberately not an async test: it needs two real loops running at once. + """ + import threading + import time + + from agent_framework._workflows import _checkpoint as checkpoint_module + + real_replace = checkpoint_module.os.replace + inside_replace = 0 + max_inside_replace = 0 + completed: list[str] = [] + guard = threading.Lock() + both_enqueued = threading.Barrier(2, timeout=10) + + def observing_replace(src, dst): # noqa: ANN001, ANN202 - test shim + nonlocal inside_replace, max_inside_replace + with guard: + inside_replace += 1 + max_inside_replace = max(max_inside_replace, inside_replace) + try: + # Widen the window a real overlap would land in. + time.sleep(0.05) + real_replace(src, dst) + finally: + with guard: + inside_replace -= 1 + completed.append(os.path.basename(str(dst))) + + monkeypatch.setattr(checkpoint_module.os, "replace", observing_replace) + + errors: list[BaseException] = [] + + def run_loop(name: str) -> None: + async def main() -> None: + storage = FileCheckpointStorage(str(tmp_path)) + # Make both loops reach save() at the same time so the queue, not timing, + # is what orders them. + await asyncio.to_thread(both_enqueued.wait) + await storage.save( + WorkflowCheckpoint( + workflow_name=name, + graph_signature_hash="test-hash", + checkpoint_id="cross-loop-id", + ) + ) + + try: + asyncio.run(main()) + except BaseException as exc: # noqa: BLE001 - surfaced to the test below + errors.append(exc) + + threads = [threading.Thread(target=run_loop, args=(f"loop-{index}",)) for index in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + assert not thread.is_alive(), "a save driven from its own event loop never finished" + + assert not errors, f"saves raised: {errors!r}" + assert len(completed) == 2 + assert max_inside_replace == 1, f"{max_inside_replace} writes reached os.replace together across event loops" + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + assert not registry, "cross-loop saves left a destination entry behind" + + +async def test_file_checkpoint_storage_repeated_cancellation_still_drains(monkeypatch): + """Cancelling again while a save is draining must not abandon the in-flight write. + + Once a task has a cancellation pending, its next await raises `CancelledError` + immediately -- so a drain built on a single `await asyncio.shield(worker)` returns + with the worker still running and the write can land after ownership is released. + A task group or supervisor that cancels more than once reaches exactly that path, + so the drain absorbs re-delivered cancellations until the worker is genuinely done. + """ + 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_replace = threading.Event() + finished_writes: list[str] = [] + guard = threading.Lock() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + replace_started.set() + assert release_replace.wait(timeout=10) + real_replace(src, dst) + with guard: + finished_writes.append(os.path.basename(str(dst))) + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + task = asyncio.create_task( + storage.save( + WorkflowCheckpoint( + workflow_name="drained", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) + ) + ) + assert await asyncio.to_thread(replace_started.wait, 10) + + # Cancel repeatedly while the worker is parked. Each one is re-delivered into + # the drain, which must keep waiting rather than return early. + for _ in range(5): + task.cancel() + await asyncio.sleep(0.01) + assert not task.done(), "drain gave up while the write was still in flight" + assert finished_writes == [] + + release_replace.set() + with pytest.raises(asyncio.CancelledError): + await task + + # The write it owned completed before it propagated, so nothing lands later. + assert finished_writes == ["shared-id.json"] + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + assert not registry + + +async def test_file_checkpoint_storage_write_failing_during_drain_is_reported(monkeypatch, caplog): + """A write that fails while draining must surface in the log, not vanish. + + The cancelled caller never sees the write's exception -- it receives + `CancelledError` -- so the only place a failure can be noticed is the log. + """ + import threading + + from agent_framework._workflows import _checkpoint as checkpoint_module + + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + replace_started = threading.Event() + release_replace = threading.Event() + + def failing_replace(src, dst): # noqa: ANN001, ANN202 - test shim + replace_started.set() + assert release_replace.wait(timeout=10) + raise OSError("disk went away mid-publish") + + monkeypatch.setattr(checkpoint_module.os, "replace", failing_replace) + + task = asyncio.create_task( + storage.save( + WorkflowCheckpoint( + workflow_name="failing", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) + ) + ) + assert await asyncio.to_thread(replace_started.wait, 10) + + task.cancel() + await asyncio.sleep(0.01) + with caplog.at_level(logging.WARNING, logger=checkpoint_module.logger.name): + release_replace.set() + with pytest.raises(asyncio.CancelledError): + await task + + assert any("failed while draining after cancellation" in record.message for record in caplog.records), ( + f"the drained write's failure was not reported: {[r.message for r in caplog.records]}" + ) + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + assert not registry, "a failed drained write left its destination entry behind" + + +def test_release_write_tolerates_an_already_dropped_entry(): + """Releasing a ticket whose entry is gone must be a no-op, not a KeyError. + + Defensive: the registry entry is dropped by whichever operation releases last, so a + release must never assume its entry is still present. + """ + from agent_framework._workflows import _checkpoint as checkpoint_module + + path = Path("/nonexistent/never-enqueued.json") + ticket = checkpoint_module._WriteTicket( # pyright: ignore[reportPrivateUsage] + path=path, + predecessor=None, + completion=ConcurrentFuture(), + ) + # No entry was ever created for this path. + checkpoint_module._release_write(ticket) # pyright: ignore[reportPrivateUsage] + assert ticket.completion.done() + assert path not in checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + + +def test_file_checkpoint_storage_abandoned_loop_does_not_own_destination_forever(monkeypatch, tmp_path): + """A loop that goes away mid-write must not leave its destination owned forever. + + Ownership is released as the write's last act on the worker thread, not only in the + coroutine's `finally`. A thread outlives the loop that submitted it, so the write + finishes and hands the destination on; releasing solely from the coroutine would + leave the path owned for the life of the process and hang every later save to it. + + Deliberately not an async test: it has to close a loop out from under a pending task. + """ + import threading + import time + + from agent_framework._workflows import _checkpoint as checkpoint_module + + real_replace = checkpoint_module.os.replace + parked = threading.Event() + release_parked = threading.Event() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + parked.set() + assert release_parked.wait(timeout=20) + real_replace(src, dst) + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="abandoned-id") + + abandoned_loop = asyncio.new_event_loop() + # Closing a loop with a pending task makes asyncio report "Task was destroyed but it + # is pending!" through the loop's exception handler. That is exactly the situation + # under test, so silence it rather than leaving the noise in CI output. + abandoned_loop.set_exception_handler(lambda loop, context: None) + try: + storage = FileCheckpointStorage(str(tmp_path)) + + async def start_and_park() -> asyncio.Task[str]: + task = abandoned_loop.create_task(storage.save(make("abandoned"))) + await asyncio.to_thread(parked.wait, 20) + return task + + pending = abandoned_loop.run_until_complete(start_and_park()) + assert not pending.done() + finally: + # Abrupt: no cancellation, so the coroutine's `finally` never runs. + abandoned_loop.close() + + # Let the orphaned worker finish. Its release happens on the thread. + release_parked.set() + deadline = time.monotonic() + 10 + while checkpoint_module._destination_queues and time.monotonic() < deadline: # pyright: ignore[reportPrivateUsage] + time.sleep(0.05) + assert not checkpoint_module._destination_queues, ( # pyright: ignore[reportPrivateUsage] + "the abandoned loop's destination entry was never released" + ) + + # A fresh loop must be able to save to the same destination. + outcome: dict[str, bool] = {} + + def run_later_save() -> None: + async def main() -> None: + later_storage = FileCheckpointStorage(str(tmp_path)) + try: + await asyncio.wait_for(later_storage.save(make("later")), timeout=10) + outcome["saved"] = True + except asyncio.TimeoutError: + outcome["saved"] = False + + asyncio.run(main()) + + thread = threading.Thread(target=run_later_save) + thread.start() + thread.join(timeout=30) + assert not thread.is_alive() + assert outcome.get("saved") is True, "a later save to the abandoned destination hung" + + +async def test_file_checkpoint_storage_failed_save_releases_the_destination(monkeypatch): + """A save that raises must not leave its destination owned. + + The failure path matters as much as cancellation: if a raising write kept ownership, + one transient disk error would hang every later save to that checkpoint for the life + of the process. + """ + from agent_framework._workflows import _checkpoint as checkpoint_module + + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + def exploding_replace(src, dst): # noqa: ANN001, ANN202 - test shim + raise OSError("simulated disk failure") + + monkeypatch.setattr(checkpoint_module.os, "replace", exploding_replace) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id") + + with pytest.raises(OSError, match="simulated disk failure"): + await storage.save(make("fails")) + + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + assert not registry, "a failed save kept its destination" + + # And the path is still usable once the failure clears. + monkeypatch.undo() + await asyncio.wait_for(storage.save(make("after-failure")), timeout=10) + assert (await storage.load("shared-id")).workflow_name == "after-failure" + assert not registry From 30a92df288e9770837a0f1179e2670f28d5867c1 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 9 Sep 2026 19:56:01 +0530 Subject: [PATCH 4/6] Python: tie checkpoint write ownership to the write, not to asyncio state Follow-up review on #7757 found two ways a later save could still overtake an earlier one. Both had the same cause: asyncio objects were standing in for the state of work that lives outside asyncio, and cancellation can mark those done while the work continues. `asyncio.wrap_future` chains cancellation into the future it wraps. Cancelling the middle of three queued saves therefore cancelled the hand-off signal the first save still had to resolve, and the unconditional release then let the third reach `os.replace` beside the first -- after which the first write could land last and overwrite the newer checkpoint. Waiting now goes through `_wait_for_signal`, which feeds a fresh per-wait future from a done-callback and leaves the shared signal untouched no matter what happens to the waiter. A save cancelled while queued no longer resolves its own signal either; it defers the hand-off until its predecessor has actually finished, so the queue stays ordered while still guaranteeing nothing waits behind it forever. Draining on the task wrapping `asyncio.to_thread` was the second: a cancelled task reports `done()` while its function is still inside `os.replace`, so loop shutdown released ownership with the write in flight. The drain now waits on the signal the worker thread resolves, which cancellation cannot mark done. Fixing those exposed two more failures of the same kind, both found by probing rather than by review. `ensure_future(asyncio.to_thread(...))` makes the write a Task, and shutdown cancels every task -- possibly before it reaches the executor. Nothing then released the destination, and the drain waited for a signal that could never be resolved, so the save never completed at all. The write is now submitted with `run_in_executor`, which submits synchronously and returns a plain Future that `asyncio.all_tasks()` does not include, plus a done-callback that releases on success, failure and cancellation alike. Deferring a cancelled ticket's hand-off onto its predecessor made release recursive: resolving one signal ran the next ticket's callback synchronously, so a run of cancelled queued saves longer than the recursion limit raised `RecursionError` on the worker thread partway through and left the destination owned for good. Releases are now collected per thread and drained in a loop, so the depth is flat however long the run. Two smaller points from the same review. A write that fails while draining is recorded on the ticket rather than read back from the worker future, because a cancelled task hides its own exception and the cancelled caller never sees it. And the drain registers one done-callback pointing at whichever waiter is current, rather than one per wait, so a caller cancelling in a loop no longer grows that list in step with it. Tests: cancelling the middle of three queued saves cannot let the third overtake; a cancelled worker task does not release while the thread writes; shutdown before the write starts neither hangs nor strands the destination, asserted structurally on the write not being a task so the regression fails cleanly instead of hanging CI; a deferred-release run of 2000 links does not recurse; submission against a dead executor releases; and resolving a signal after its loop closed does not raise. --- .../agent_framework/_workflows/_checkpoint.py | 243 ++++++++--- .../core/tests/workflow/test_checkpoint.py | 378 ++++++++++++++++++ 2 files changed, 569 insertions(+), 52 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 3d64840b12b..04cc505df32 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import contextlib import copy import json import logging @@ -258,14 +259,21 @@ async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID] # process-wide queue keyed by the canonical destination path, and taken *before* the # write is submitted to a worker thread rather than inside it. # -# Waiting on a ``concurrent.futures.Future`` through ``asyncio.wrap_future`` is what -# makes that work in both directions: it is not bound to a loop, so a single chain -# orders writers regardless of which loop enqueued them, and waiters suspend on the -# event loop instead of occupying an ``asyncio.to_thread`` worker. Blocking on a -# ``threading.Lock`` inside the worker instead -- the previous design -- let a burst of -# same-path saves fill the default executor and stall unrelated ``to_thread`` work, -# including checkpoint loads, and could deadlock once the write that had to finish -# first was queued behind those waiters. +# Hand-off runs over ``concurrent.futures.Future`` signals, which are not bound to a +# loop, so a single chain orders writers regardless of which loop enqueued them, and +# waiters suspend on the event loop instead of occupying an ``asyncio.to_thread`` +# worker. Blocking on a ``threading.Lock`` inside the worker instead -- the original +# design -- let a burst of same-path saves fill the default executor and stall unrelated +# ``to_thread`` work, including checkpoint loads, and could deadlock once the write that +# had to finish first was queued behind those waiters. +# +# Those signals are awaited through ``_wait_for_signal``, never ``asyncio.wrap_future``. +# Cancellation is the reason: ``wrap_future`` chains it into the future it wraps, so one +# cancelled waiter would cancel the signal an earlier writer still has to resolve, and +# the ``asyncio`` task wrapping a submitted write reports ``done()`` while its function +# is still running on the executor thread. Both let a later save start its ``os.replace`` +# beside an earlier one, after which the earlier write can land last and overwrite the +# newer checkpoint. # # Entries are reference-counted by queued-or-running operations and dropped when the # last one releases, so a process that saves many distinct checkpoint IDs -- the @@ -303,6 +311,10 @@ class _WriteTicket: completion: Future[None] #: Set by whichever of the worker thread or the coroutine releases first. released: bool = False + #: Failure raised by the write, recorded on the worker thread. A cancelled caller + #: never sees it, and a cancelled task hides its own exception, so the ticket is the + #: only place it survives. + error: BaseException | None = None def _enqueue_write(file_path: Path) -> _WriteTicket: @@ -319,46 +331,134 @@ def _enqueue_write(file_path: Path) -> _WriteTicket: return _WriteTicket(path=file_path, predecessor=predecessor, completion=completion) -async def _drain_cancelled_write(worker: asyncio.Future[None], file_path: Path) -> None: - """Wait for *worker* to finish while this task is being cancelled. +def _wait_for_signal(source: Future[None]) -> asyncio.Future[None]: + """Return a fresh awaitable that completes when *source* does. + + Deliberately not ``asyncio.wrap_future``: that chains cancellation into the future it + wraps, so a waiter cancelled here would cancel the hand-off signal an earlier writer + still has to resolve, and every later writer keyed to it. A per-wait future fed by a + done-callback leaves *source* untouched no matter what happens to the waiter. + """ + loop = asyncio.get_running_loop() + waiter: asyncio.Future[None] = loop.create_future() + + def _resolve(_completed: Future[None]) -> None: + def _set() -> None: + if not waiter.done(): + waiter.set_result(None) + + # A closed loop means nobody is left to wake. + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(_set) + + source.add_done_callback(_resolve) + return waiter + + +async def _await_signal_through_cancellation(source: Future[None]) -> None: + """Wait until *source* resolves, absorbing cancellations delivered meanwhile. - A single ``await asyncio.shield(worker)`` is not enough here. Once the enclosing - task has a cancellation pending, the next await raises ``CancelledError`` again - immediately, so the shield returns without the worker having finished and the write - can still land after ownership is released. Absorbing the re-delivered - cancellations until the worker is genuinely done is what makes the drain real. + Used to wait out a write this coroutine still owns. Waiting on the ``asyncio`` task + wrapping the write is not equivalent: a cancelled task reports ``done()`` while its + function is still running on the executor thread, so ownership would be released + with the write still in flight. Only the signal the worker thread resolves tracks + the write itself, and cancellation cannot mark it done. + + The done-callback is registered once and points at whichever waiter is current, + rather than one callback per wait. Each cancellation would otherwise leave another + callback on the signal, so a caller cancelling in a loop would grow that list in + step with it and pay for the whole list when the write finally lands. """ - while not worker.done(): + if source.done(): + # Purely to avoid registering a callback that would fire straight back; the loop + # condition below already short-circuits, so this is not load-bearing. + return + + loop = asyncio.get_running_loop() + current: dict[str, asyncio.Future[None] | None] = {"waiter": None} + + def _resolve(_completed: Future[None]) -> None: + def _set() -> None: + waiter = current["waiter"] + if waiter is not None and not waiter.done(): + waiter.set_result(None) + + # A closed loop means nobody is left to wake. + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(_set) + + source.add_done_callback(_resolve) + while not source.done(): + waiter: asyncio.Future[None] = loop.create_future() + current["waiter"] = waiter try: - await asyncio.shield(worker) + await waiter except asyncio.CancelledError: + # Re-delivered cancellation. The signal may have resolved against the waiter + # we just abandoned, so the loop condition is what decides whether to stop. continue - except Exception: - break - if worker.done() and not worker.cancelled(): - exception = worker.exception() - if exception is not None: - # The caller is being cancelled and will not see this, but a write that - # failed while draining should not vanish silently. - logger.warning(f"Checkpoint write to {file_path} failed while draining after cancellation: {exception!r}") + + +def _release_write_after(ticket: _WriteTicket, predecessor: Future[None]) -> None: + """Hand *ticket*'s ownership on only once *predecessor* has actually finished. + + A ticket cancelled while still queued must not resolve its own signal immediately: + an earlier writer still owns the destination, and the next waiter would start its + write alongside that one. Deferring keeps the queue in order while still guaranteeing + the hand-off happens, so nothing waits forever behind a cancelled save. + """ + predecessor.add_done_callback(lambda _completed: _release_write(ticket)) + + +class _ReleaseTrampoline(threading.local): + """Releases waiting for the one currently unwinding on this thread. + + Resolving one ticket's signal runs the next ticket's deferred-release callback + synchronously, so a run of cancelled queued saves would otherwise nest one stack + frame per link and raise ``RecursionError`` on the worker thread partway through -- + leaving the destination owned for good. Collecting them here and draining in a loop + keeps the depth flat however long the run is. + """ + + queued: list[_WriteTicket] | None = None + + +_release_trampoline = _ReleaseTrampoline() def _release_write(ticket: _WriteTicket) -> None: """Hand ownership to the next waiter and drop the entry once nothing is queued. - Called from two places, whichever gets there first, and idempotent so both can: + Called from three places, whichever gets there first, and idempotent so all can: * the worker thread, as the last act of the write itself. This is what keeps a destination usable if the event loop that submitted the write goes away before the - coroutine can resume -- the thread finishes and releases regardless, where a - release that only happened in the coroutine would leave the path owned forever and - hang every later save to it. - * the coroutine, on any exit path. This is the only releaser when the write was - never submitted, which is the case for a save cancelled while still queued. + coroutine can resume -- the thread finishes and releases regardless. + * a done-callback on the submitted write, which fires on success, failure and + cancellation alike, so the executor dropping the work still releases. + * the coroutine, when the write was never submitted at all. Resolving the completion signal is what keeps the chain moving, so an operation that is cancelled or fails must still release rather than stall every later writer. """ + queued = _release_trampoline.queued + if queued is not None: + # A release is already unwinding on this thread; let it drain this one. + queued.append(ticket) + return + + draining: list[_WriteTicket] = [] + _release_trampoline.queued = draining + try: + _release_one_write(ticket) + while draining: + _release_one_write(draining.pop(0)) + finally: + _release_trampoline.queued = None + + +def _release_one_write(ticket: _WriteTicket) -> None: + """Release exactly one ticket. Only ``_release_write`` should call this.""" with _destination_queues_guard: if ticket.released: return @@ -511,34 +611,73 @@ def _write_atomic_and_release(ticket: _WriteTicket) -> None: # forever and hanging every later save to it. try: _write_atomic() + except BaseException as exc: + ticket.error = exc + raise finally: _release_write(ticket) ticket = _enqueue_write(file_path) - try: - if ticket.predecessor is not None: - # Suspends on the event loop, not on an executor worker, and orders this - # save behind every earlier one for the same destination even when they - # were enqueued from a different loop. A cancellation delivered here - # propagates with nothing submitted and nothing written, which is the - # cheapest correct outcome: `finally` hands ownership straight to the - # next waiter. - await asyncio.wrap_future(ticket.predecessor) - - worker = asyncio.ensure_future(asyncio.to_thread(_write_atomic_and_release, ticket)) + if ticket.predecessor is not None: + # Suspends on the event loop, not on an executor worker, and orders this save + # behind every earlier one for the same destination even when they were + # enqueued from a different loop. try: - # Shield so a cancellation arriving mid-write cannot leave the worker - # running past this frame: its os.replace would otherwise land after the - # caller returned, overwriting whatever a later save had published. - await asyncio.shield(worker) - except asyncio.CancelledError: - # Ownership is still held, so drain before propagating. Releasing first - # would let the next save begin while this replace is still in flight, - # which is the race the queue exists to prevent. - await _drain_cancelled_write(worker, file_path) + await _wait_for_signal(ticket.predecessor) + except BaseException: + # Nothing was submitted and nothing written, but the hand-off cannot + # happen yet: an earlier writer still owns the destination, and resolving + # this ticket's signal now would let the next save run its os.replace + # alongside that one -- after which the earlier write can land last and + # overwrite the newer checkpoint. Defer the hand-off until the + # predecessor has actually finished. + _release_write_after(ticket, ticket.predecessor) raise - finally: + + # Ownership held from here. Submit through `run_in_executor` rather than + # `ensure_future(asyncio.to_thread(...))`: the latter creates a Task, and loop + # shutdown cancels every task, so the write could be cancelled before it ever + # reached the executor -- leaving nobody to release the destination and this + # coroutine waiting on a signal that could never be resolved. `run_in_executor` + # returns a plain Future that `asyncio.all_tasks()` does not include, and it + # submits synchronously, so returning from it means the write really is queued. + loop = asyncio.get_running_loop() + try: + write_future = loop.run_in_executor(None, _write_atomic_and_release, ticket) + except BaseException: + # Never reached the executor, so no worker will release on our behalf. _release_write(ticket) + raise + # The callback fires on success, failure and cancellation alike, so the + # destination is released even if the executor drops the work before the function + # runs. Idempotent with the worker thread's own release, which means no flag is + # needed to decide which of the two owns it. + write_future.add_done_callback(lambda _completed: _release_write(ticket)) + + try: + # Shield so a cancellation arriving mid-write cannot leave the write running + # past this frame: its os.replace would otherwise land after the caller + # returned, overwriting whatever a later save had published. + await asyncio.shield(write_future) + except asyncio.CancelledError: + # Wait out the write itself rather than the task wrapping it. A cancelled + # task reports done() while its function is still running on the executor + # thread, so waiting on `worker` would return with the replace still in + # flight and ownership about to be released. + await _await_signal_through_cancellation(ticket.completion) + if ticket.error is not None: + # The caller is receiving CancelledError and will never see this, and a + # cancelled task hides its own exception, so the log is the only place a + # write that failed while draining can surface. + logger.warning( + f"Checkpoint write to {file_path} failed while draining after cancellation: {ticket.error!r}" + ) + raise + except BaseException: + # The write failed. The worker already released in its own `finally`; this is + # idempotent cover for a submission that never got that far. + _release_write(ticket) + raise 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 488937a4c27..91923c0d60c 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -4,7 +4,9 @@ import json import logging import os +import sys import tempfile +import time from concurrent.futures import Future as ConcurrentFuture from dataclasses import dataclass from datetime import datetime, timezone @@ -2430,3 +2432,379 @@ def make(name: str) -> WorkflowCheckpoint: await asyncio.wait_for(storage.save(make("after-failure")), timeout=10) assert (await storage.load("shared-id")).workflow_name == "after-failure" assert not registry + + +async def test_file_checkpoint_storage_cancelling_a_queued_save_does_not_let_the_next_overtake(monkeypatch): + """A save cancelled while queued must not hand off before its predecessor finishes. + + Reviewer concern on #7757: `asyncio.wrap_future` chains cancellation into the future + it wraps, so cancelling the middle of three queued saves cancelled the shared + hand-off signal, and the unconditional release then let the third save reach + `os.replace` while the first still owned the destination -- after which the first + write could land last and overwrite the newer checkpoint. The hand-off for a + cancelled ticket is now deferred until its predecessor actually completes. + """ + 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 + holder_parked = threading.Event() + release_holder = threading.Event() + inside_replace = 0 + max_inside_replace = 0 + completed: list[str] = [] + guard = threading.Lock() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + nonlocal inside_replace, max_inside_replace + with guard: + inside_replace += 1 + max_inside_replace = max(max_inside_replace, inside_replace) + park = inside_replace == 1 and not holder_parked.is_set() + try: + if park: + holder_parked.set() + assert release_holder.wait(timeout=15) + real_replace(src, dst) + with guard: + completed.append(os.path.basename(str(dst))) + finally: + with guard: + inside_replace -= 1 + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id") + + holder = asyncio.create_task(storage.save(make("holder"))) + assert await asyncio.to_thread(holder_parked.wait, 15) + + queued = asyncio.create_task(storage.save(make("cancelled"))) + third = asyncio.create_task(storage.save(make("third"))) + + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + canonical = (Path(temp_dir) / "shared-id.json").resolve() + for _ in range(400): + entry = registry.get(canonical) + if entry is not None and entry.pending == 3: + break + await asyncio.sleep(0.005) + entry = registry.get(canonical) + assert entry is not None and entry.pending == 3 + + queued.cancel() + with pytest.raises(asyncio.CancelledError): + await queued + + # The holder is still parked inside its replace. The third save must not have + # taken the destination on the back of the cancellation. + for _ in range(30): + await asyncio.sleep(0.01) + with guard: + assert max_inside_replace == 1, "a second write started while the holder still owned the path" + assert completed == [], "a write completed while the holder was still parked" + assert not third.done() + + release_holder.set() + await asyncio.wait_for(asyncio.gather(holder, third), timeout=15) + + with guard: + assert max_inside_replace == 1, "writes overlapped for one destination" + assert len(completed) == 2, f"expected holder + third, got {completed}" + assert (await storage.load("shared-id")).workflow_name == "third" + assert not registry, "a cancelled queued save leaked its destination entry" + + +def test_file_checkpoint_storage_cancelled_worker_task_does_not_release_early(monkeypatch, tmp_path): + """Cancelling the worker task must not release ownership while the thread writes. + + Reviewer concern on #7757: loop shutdown cancels `save()` and the task wrapping + `asyncio.to_thread`, which makes `worker.done()` true while the function is still + blocked in `os.replace`. Draining on that task therefore returned immediately and + ownership was released with the write in flight. The drain now waits on the signal + the worker thread resolves, which cancellation cannot mark done. + + Deliberately not an async test: it cancels every task the way shutdown does. + """ + import threading + import time + + from agent_framework._workflows import _checkpoint as checkpoint_module + + real_replace = checkpoint_module.os.replace + parked = threading.Event() + release_parked = threading.Event() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + parked.set() + assert release_parked.wait(timeout=15) + real_replace(src, dst) + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + canonical = (tmp_path / "shared-id.json").resolve() + observed: dict[str, bool] = {} + + async def main() -> None: + storage = FileCheckpointStorage(str(tmp_path)) + task = asyncio.create_task( + storage.save( + WorkflowCheckpoint(workflow_name="holder", graph_signature_hash="test-hash", checkpoint_id="shared-id") + ) + ) + await asyncio.to_thread(parked.wait, 15) + + # What loop shutdown does: cancel every remaining task, including the one + # wrapping asyncio.to_thread. + for pending in [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]: + pending.cancel() + + for _ in range(40): + await asyncio.sleep(0.01) + if canonical not in checkpoint_module._destination_queues: # pyright: ignore[reportPrivateUsage] + break + # The write is still parked, so the destination must still be owned. + observed["released_early"] = canonical not in checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + observed["still_writing"] = not release_parked.is_set() + + release_parked.set() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(main()) + + assert observed["still_writing"], "the test never held the write open" + assert not observed["released_early"], "ownership was released while the write was still in flight" + # The worker thread releases once its write lands, so nothing is left owned. + deadline = time.monotonic() + 10 + while checkpoint_module._destination_queues and time.monotonic() < deadline: # pyright: ignore[reportPrivateUsage] + time.sleep(0.05) + assert not checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + + +def test_file_checkpoint_storage_shutdown_before_the_write_starts_does_not_hang(monkeypatch, tmp_path): + """Loop shutdown before the write reaches the executor must not strand the save. + + Submitting through `asyncio.ensure_future(asyncio.to_thread(...))` makes the write a + Task, and shutdown cancels every task -- potentially before it has reached the + executor. Nothing would then release the destination, and a coroutine waiting for the + write's completion signal would wait for a signal that could never be resolved, so + the save never finished at all. Submission now goes through `run_in_executor`, which + returns a plain Future that `asyncio.all_tasks()` does not include, and a done + callback releases even if the executor drops the work. + + Deliberately not an async test: it has to cancel every task the way shutdown does. + """ + from agent_framework._workflows import _checkpoint as checkpoint_module + + canonical = (tmp_path / "shared-id.json").resolve() + outcome: dict[str, object] = {} + + async def main() -> None: + storage = FileCheckpointStorage(str(tmp_path)) + asyncio.create_task( + storage.save( + WorkflowCheckpoint(workflow_name="victim", graph_signature_hash="test-hash", checkpoint_id="shared-id") + ) + ) + # One tick, so the save coroutine actually runs and submits its write. Without + # it the write has not been created yet and cancelling only the save takes the + # clean cancelled-before-submission path, which was never the broken case. + await asyncio.sleep(0) + victims = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()] + outcome["task_count"] = len(victims) + + # Assert the structural property *before* cancelling anything. If the write is a + # task, shutdown sweeps it and the save can end up waiting for a completion + # signal nothing will ever resolve -- which hangs `asyncio.run`'s own shutdown, + # outside any `wait_for` this test could wrap around it. Failing here keeps the + # regression a clean assertion rather than a hung CI job. + assert len(victims) == 1, ( + f"the write must not be a task, or loop shutdown cancels it: found {len(victims)} tasks" + ) + + for task in victims: + task.cancel() + await asyncio.wait_for(asyncio.gather(*victims, return_exceptions=True), timeout=10) + + asyncio.run(main()) + assert outcome["task_count"] == 1 + + deadline = time.monotonic() + 10 + while canonical in checkpoint_module._destination_queues and time.monotonic() < deadline: # pyright: ignore[reportPrivateUsage] + time.sleep(0.05) + assert canonical not in checkpoint_module._destination_queues, ( # pyright: ignore[reportPrivateUsage] + "shutdown before the write started left the destination owned forever" + ) + + +async def test_file_checkpoint_storage_executor_shutdown_at_submission_releases(): + """If the write cannot be submitted at all, the coroutine must release the destination. + + `run_in_executor` raises synchronously against a shut-down executor, so no worker + exists to release on our behalf. This is the one path where the coroutine, not the + worker thread, owns the release. + """ + from concurrent.futures import ThreadPoolExecutor + + from agent_framework._workflows import _checkpoint as checkpoint_module + + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + dead_executor = ThreadPoolExecutor(max_workers=1) + dead_executor.shutdown(wait=True) + loop: Any = asyncio.get_running_loop() + previous_executor = loop._default_executor + loop.set_default_executor(dead_executor) + try: + with pytest.raises(RuntimeError): + await storage.save( + WorkflowCheckpoint( + workflow_name="no-executor", + graph_signature_hash="test-hash", + checkpoint_id="shared-id", + ) + ) + finally: + loop._default_executor = previous_executor + + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + assert not registry, "a save that could not be submitted kept its destination" + + # The path is still usable once an executor is available again. + await asyncio.wait_for( + storage.save( + WorkflowCheckpoint(workflow_name="after", graph_signature_hash="test-hash", checkpoint_id="shared-id") + ), + timeout=10, + ) + assert (await storage.load("shared-id")).workflow_name == "after" + + +def test_wait_for_signal_survives_its_loop_being_closed(): + """Resolving a signal after its waiter's loop closed must not raise. + + The worker thread outlives the loop that submitted the write, so it can resolve a + hand-off signal that a now-dead loop was waiting on. `call_soon_threadsafe` raises + `RuntimeError` on a closed loop, and that runs inside a `concurrent.futures` + done-callback -- where an exception would be swallowed into the callback machinery + rather than surfacing usefully. + """ + import threading + from concurrent.futures import Future as ConcurrentFuture + + from agent_framework._workflows import _checkpoint as checkpoint_module + + source: ConcurrentFuture[None] = ConcurrentFuture() + loop = asyncio.new_event_loop() + try: + + async def register() -> None: + checkpoint_module._wait_for_signal(source) # pyright: ignore[reportPrivateUsage] + + loop.run_until_complete(register()) + finally: + loop.close() + + # The worker thread resolves it after the loop is gone. + errors: list[BaseException] = [] + + def resolve() -> None: + try: + source.set_result(None) + except BaseException as exc: # noqa: BLE001 - surfaced below + errors.append(exc) + + thread = threading.Thread(target=resolve) + thread.start() + thread.join(timeout=10) + assert not thread.is_alive() + assert not errors, f"resolving after loop close raised: {errors!r}" + assert source.done() + + +def test_await_signal_through_cancellation_survives_its_loop_being_closed(): + """Resolving the signal after the draining loop closed must not raise. + + The worker thread outlives the loop that submitted the write, so it can resolve a + signal that a now-dead loop was draining on. `call_soon_threadsafe` raises + `RuntimeError` against a closed loop, inside a `concurrent.futures` done-callback + where it would be swallowed rather than surface. + """ + import threading + from concurrent.futures import Future as ConcurrentFuture + + from agent_framework._workflows import _checkpoint as checkpoint_module + + source: ConcurrentFuture[None] = ConcurrentFuture() + loop = asyncio.new_event_loop() + loop.set_exception_handler(lambda active_loop, context: None) + try: + # Start the drain, let it subscribe and suspend, then abandon the loop. + drain = loop.create_task( + checkpoint_module._await_signal_through_cancellation(source) # pyright: ignore[reportPrivateUsage] + ) + loop.run_until_complete(asyncio.sleep(0)) + assert not drain.done() + finally: + loop.close() + + errors: list[BaseException] = [] + + def resolve() -> None: + try: + source.set_result(None) + except BaseException as exc: # noqa: BLE001 - surfaced below + errors.append(exc) + + thread = threading.Thread(target=resolve) + thread.start() + thread.join(timeout=10) + assert not thread.is_alive() + assert not errors, f"resolving after the draining loop closed raised: {errors!r}" + assert source.done() + + +def test_release_write_does_not_recurse_per_deferred_link(): + """A long run of deferred releases must not nest one stack frame per link. + + A save cancelled while queued defers its hand-off onto its predecessor, so resolving + the first ticket runs the second's callback, which resolves the third, and so on. Done + naively that is synchronous recursion: a run longer than the recursion limit raised + `RecursionError` on the worker thread partway through and left the destination owned + for the life of the process. The chain is drained in a loop instead. + + Built from tickets directly so the depth can exceed the recursion limit without + thousands of real saves. + """ + from concurrent.futures import Future as ConcurrentFuture + + from agent_framework._workflows import _checkpoint as checkpoint_module + + depth = 2000 + assert depth > sys.getrecursionlimit(), "the run has to be longer than the recursion limit to prove anything" + + path = Path("/nonexistent/deferred-chain.json") + tickets = [ + checkpoint_module._WriteTicket( # pyright: ignore[reportPrivateUsage] + path=path, + predecessor=None, + completion=ConcurrentFuture(), + ) + for _ in range(depth) + ] + # Each link releases only once the one before it has. + for earlier, later in zip(tickets, tickets[1:]): + checkpoint_module._release_write_after(later, earlier.completion) # pyright: ignore[reportPrivateUsage] + + # Releasing the head must unwind the whole run. + checkpoint_module._release_write(tickets[0]) # pyright: ignore[reportPrivateUsage] + + unresolved = [index for index, ticket in enumerate(tickets) if not ticket.completion.done()] + assert not unresolved, f"{len(unresolved)} links never released, first at index {unresolved[0]}" + assert all(ticket.released for ticket in tickets) From 6783a30e531c733b70842143a7df3e1244effc82 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Thu, 10 Sep 2026 06:25:02 +0530 Subject: [PATCH 5/6] Python: release a queued checkpoint save whose event loop is abandoned Follow-up review on #7757. A submitted write is released by its worker thread, which outlives the loop that submitted it, but a ticket still waiting for its predecessor has no worker to fall back on. If that waiter's loop closes, `call_soon_threadsafe` raises `RuntimeError`, the suppression swallowed it, and the coroutine never resumed -- so nothing resolved the ticket's hand-off signal and every later save for that destination waited forever. Reproduced: the entry stays at `pending=1` and a later save from a fresh loop times out. The predecessor's callback now owns that case. `_wait_for_signal` takes an `on_abandoned` hook, invoked when the waiter's loop has closed, and the queued wait passes its own release -- performed on the thread that resolved the predecessor rather than on a loop that will never run again. The drain keeps suppressing, and the two guards now differ deliberately: by the time anything drains, the write has been submitted, so the worker thread and the submitted future's callback both release without needing that loop. The comment says so, since the asymmetry is otherwise unexplained. Two windows this cannot close, both stated in the comment rather than implied: a loop that closes after `call_soon_threadsafe` succeeded but before the callback runs, and one abandoned without being closed at all. Reaching either needs a loop closed with tasks still pending, which asyncio already reports as an error. Graceful shutdown is unaffected -- `asyncio.run` cancels pending tasks first, which drives a queued save through its cancellation path and defers the hand-off onto its predecessor. That path now has its own test, since it is the guarantee callers actually depend on. --- .../agent_framework/_workflows/_checkpoint.py | 41 ++- .../core/tests/workflow/test_checkpoint.py | 277 ++++++++++++++++++ 2 files changed, 312 insertions(+), 6 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 04cc505df32..8f3f8d95722 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -11,7 +11,7 @@ import threading import time import uuid -from collections.abc import Mapping +from collections.abc import Callable, Mapping from concurrent.futures import Future from dataclasses import dataclass, field, fields from datetime import datetime, timezone @@ -331,7 +331,11 @@ def _enqueue_write(file_path: Path) -> _WriteTicket: return _WriteTicket(path=file_path, predecessor=predecessor, completion=completion) -def _wait_for_signal(source: Future[None]) -> asyncio.Future[None]: +def _wait_for_signal( + source: Future[None], + *, + on_abandoned: Callable[[], None] | None = None, +) -> asyncio.Future[None]: """Return a fresh awaitable that completes when *source* does. Deliberately not ``asyncio.wrap_future``: that chains cancellation into the future it @@ -347,9 +351,24 @@ def _set() -> None: if not waiter.done(): waiter.set_result(None) - # A closed loop means nobody is left to wake. - with contextlib.suppress(RuntimeError): + try: loop.call_soon_threadsafe(_set) + except RuntimeError: + # The waiter's loop has closed, so the coroutine suspended here will never + # resume and cannot do anything on its way out. Whatever it owed -- releasing + # its place in the queue, above all -- has to happen here instead, on the + # thread that resolved the signal. + # + # This catches a loop that is already closed when the signal resolves. It + # cannot catch one that closes in the window after `call_soon_threadsafe` + # succeeded but before the callback runs, nor one abandoned without being + # closed at all: both leave a queued ticket unreleased. Reaching either needs + # a loop closed with tasks still pending, which asyncio already reports as an + # error, and a graceful shutdown is unaffected -- `asyncio.run` cancels + # pending tasks first, which drives the queued save through its cancellation + # path and defers the hand-off onto the predecessor. + if on_abandoned is not None: + on_abandoned() source.add_done_callback(_resolve) return waiter @@ -383,7 +402,10 @@ def _set() -> None: if waiter is not None and not waiter.done(): waiter.set_result(None) - # A closed loop means nobody is left to wake. + # Safe to swallow here, unlike the queued wait in `_wait_for_signal`. By the + # time anything drains, the write has been submitted, so the worker thread and + # the submitted future's callback both release the ticket without needing this + # loop; a closed loop only means there is nobody left to wake. with contextlib.suppress(RuntimeError): loop.call_soon_threadsafe(_set) @@ -623,7 +645,14 @@ def _write_atomic_and_release(ticket: _WriteTicket) -> None: # behind every earlier one for the same destination even when they were # enqueued from a different loop. try: - await _wait_for_signal(ticket.predecessor) + await _wait_for_signal( + ticket.predecessor, + # If this loop dies while we are queued, nothing else would release + # this ticket: the write was never submitted, so there is no worker + # thread to fall back on, and every later save for the destination + # would wait on a signal nobody resolves. + on_abandoned=lambda: _release_write(ticket), + ) except BaseException: # Nothing was submitted and nothing written, but the hand-off cannot # happen yet: an earlier writer still owns the destination, and resolving diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 91923c0d60c..d9bee711bb0 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -2808,3 +2808,280 @@ def test_release_write_does_not_recurse_per_deferred_link(): unresolved = [index for index, ticket in enumerate(tickets) if not ticket.completion.done()] assert not unresolved, f"{len(unresolved)} links never released, first at index {unresolved[0]}" assert all(ticket.released for ticket in tickets) + + +def test_file_checkpoint_storage_abandoned_loop_while_queued_releases_the_destination(monkeypatch, tmp_path): + """A save whose loop dies while it is still queued must not strand the destination. + + Reviewer follow-up on #7757: the submitted write is released by its worker thread, but + a ticket still *waiting* for its predecessor has no worker. If that waiter's loop + closes, `call_soon_threadsafe` raises and the coroutine never resumes, so nothing + resolves its hand-off signal and every later save for the destination waits forever. + The predecessor's callback now performs the release itself in that case. + + Deliberately not an async test: it needs two loops and has to close one of them. + """ + import threading + import time + + from agent_framework._workflows import _checkpoint as checkpoint_module + + real_replace = checkpoint_module.os.replace + parked = threading.Event() + release_parked = threading.Event() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + parked.set() + assert release_parked.wait(timeout=20) + real_replace(src, dst) + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id") + + canonical = (tmp_path / "shared-id.json").resolve() + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + + holder_loop = asyncio.new_event_loop() + queued_loop = asyncio.new_event_loop() + queued_loop.set_exception_handler(lambda loop, context: None) + try: + holder_storage = FileCheckpointStorage(str(tmp_path)) + + async def start_holder() -> asyncio.Task[str]: + task = holder_loop.create_task(holder_storage.save(make("holder"))) + await asyncio.to_thread(parked.wait, 20) + return task + + holder = holder_loop.run_until_complete(start_holder()) + + # A second save queues behind the parked holder on its own loop. + queued_storage = FileCheckpointStorage(str(tmp_path)) + + async def start_queued() -> None: + queued_loop.create_task(queued_storage.save(make("queued"))) + for _ in range(400): + entry = registry.get(canonical) + if entry is not None and entry.pending == 2: + return + await asyncio.sleep(0.005) + raise AssertionError("the second save never queued behind the holder") + + queued_loop.run_until_complete(start_queued()) + finally: + # Abandoned while its ticket is still queued and nothing has been submitted. + queued_loop.close() + + release_parked.set() + holder_loop.run_until_complete(asyncio.gather(holder, return_exceptions=True)) + holder_loop.close() + + deadline = time.monotonic() + 10 + while canonical in registry and time.monotonic() < deadline: + time.sleep(0.05) + assert canonical not in registry, "the abandoned queued save left the destination owned" + + # And the destination is usable again from a fresh loop. + outcome: dict[str, bool] = {} + + def later_save() -> None: + async def main() -> None: + storage = FileCheckpointStorage(str(tmp_path)) + try: + await asyncio.wait_for(storage.save(make("later")), timeout=10) + outcome["saved"] = True + except asyncio.TimeoutError: + outcome["saved"] = False + + asyncio.run(main()) + + thread = threading.Thread(target=later_save) + thread.start() + thread.join(timeout=30) + assert not thread.is_alive() + assert outcome.get("saved") is True, "a later save to the abandoned destination hung" + + +def test_file_checkpoint_storage_graceful_shutdown_releases_a_queued_save(tmp_path, monkeypatch): + """A queued save released through the normal shutdown path, which is the guarantee. + + `asyncio.run` cancels pending tasks before closing, so a save still waiting for its + predecessor takes its cancellation path and defers the hand-off onto that + predecessor. Pinning it because the `on_abandoned` hook only covers a loop that is + already closed when the signal resolves -- this is the path that has to stay safe + without it. + """ + import threading + import time + + from agent_framework._workflows import _checkpoint as checkpoint_module + + real_replace = checkpoint_module.os.replace + parked = threading.Event() + release_parked = threading.Event() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + parked.set() + assert release_parked.wait(timeout=20) + real_replace(src, dst) + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id") + + canonical = (tmp_path / "shared-id.json").resolve() + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + + holder_loop = asyncio.new_event_loop() + try: + holder_storage = FileCheckpointStorage(str(tmp_path)) + + async def start_holder() -> asyncio.Task[str]: + task = holder_loop.create_task(holder_storage.save(make("holder"))) + await asyncio.to_thread(parked.wait, 20) + return task + + holder = holder_loop.run_until_complete(start_holder()) + + # A second save queues behind it, on a loop that exits the normal way. + def run_queued() -> None: + async def main() -> None: + storage = FileCheckpointStorage(str(tmp_path)) + asyncio.create_task(storage.save(make("queued"))) + for _ in range(400): + entry = registry.get(canonical) + if entry is not None and entry.pending == 2: + return + await asyncio.sleep(0.005) + raise AssertionError("the second save never queued behind the holder") + + asyncio.run(main()) + + thread = threading.Thread(target=run_queued) + thread.start() + thread.join(timeout=30) + assert not thread.is_alive() + + release_parked.set() + holder_loop.run_until_complete(asyncio.gather(holder, return_exceptions=True)) + finally: + holder_loop.close() + + deadline = time.monotonic() + 10 + while canonical in registry and time.monotonic() < deadline: + time.sleep(0.05) + assert canonical not in registry, "a gracefully cancelled queued save left the destination owned" + + +def test_file_checkpoint_storage_abandonment_mid_chain_keeps_the_queue_ordered(tmp_path, monkeypatch): + """Releasing an abandoned queued ticket must hand off to its successor, in order. + + The two-ticket case only shows the entry clearing. With a live save queued *behind* + the abandoned one, the release has to propagate through that link and still serialize + the writes -- if it handed off early, the successor's `os.replace` would run beside + the holder's. + """ + import threading + import time + + from agent_framework._workflows import _checkpoint as checkpoint_module + + real_replace = checkpoint_module.os.replace + parked = threading.Event() + release_parked = threading.Event() + inside_replace = 0 + peak_inside_replace = 0 + completed: list[str] = [] + guard = threading.Lock() + + def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim + nonlocal inside_replace, peak_inside_replace + with guard: + inside_replace += 1 + peak_inside_replace = max(peak_inside_replace, inside_replace) + park = inside_replace == 1 and not parked.is_set() + try: + if park: + parked.set() + assert release_parked.wait(timeout=25) + real_replace(src, dst) + with guard: + completed.append(os.path.basename(str(dst))) + finally: + with guard: + inside_replace -= 1 + + monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id") + + canonical = (tmp_path / "shared-id.json").resolve() + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + + def wait_for_pending(count: int, timeout: float = 20.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + entry = registry.get(canonical) + if entry is not None and entry.pending == count: + return + time.sleep(0.005) + entry = registry.get(canonical) + raise AssertionError(f"expected {count} queued, saw {entry.pending if entry else 0}") + + holder_loop = asyncio.new_event_loop() + abandoned_loop = asyncio.new_event_loop() + abandoned_loop.set_exception_handler(lambda loop, context: None) + successor: dict[str, bool] = {} + try: + holder_storage = FileCheckpointStorage(str(tmp_path)) + + async def start_holder() -> asyncio.Task[str]: + task = holder_loop.create_task(holder_storage.save(make("holder"))) + await asyncio.to_thread(parked.wait, 25) + return task + + holder = holder_loop.run_until_complete(start_holder()) + + async def start_abandoned() -> None: + storage = FileCheckpointStorage(str(tmp_path)) + abandoned_loop.create_task(storage.save(make("abandoned"))) + await asyncio.sleep(0) + + abandoned_loop.run_until_complete(start_abandoned()) + wait_for_pending(2) + + def run_successor() -> None: + async def main() -> None: + storage = FileCheckpointStorage(str(tmp_path)) + try: + await asyncio.wait_for(storage.save(make("successor")), timeout=25) + successor["saved"] = True + except asyncio.TimeoutError: + successor["saved"] = False + + asyncio.run(main()) + + thread = threading.Thread(target=run_successor) + thread.start() + wait_for_pending(3) + + abandoned_loop.close() + + release_parked.set() + holder_loop.run_until_complete(asyncio.gather(holder, return_exceptions=True)) + thread.join(timeout=30) + assert not thread.is_alive() + finally: + holder_loop.close() + + assert successor.get("saved") is True, "the save queued behind the abandoned one never ran" + with guard: + assert peak_inside_replace == 1, "writes overlapped across the abandoned link" + assert len(completed) == 2, f"expected holder + successor, got {completed}" + deadline = time.monotonic() + 10 + while canonical in registry and time.monotonic() < deadline: + time.sleep(0.05) + assert canonical not in registry From a86e94698311b6cdb4bca96b8c6397c59ee5bc26 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Thu, 10 Sep 2026 08:33:54 +0530 Subject: [PATCH 6/6] Python: cover the checkpoint write error paths this PR introduced Three branches added by this PR had no test. Coverage of `_checkpoint.py` was reporting 97% with the gaps sitting exactly on the error handling: * the bounded `os.replace` retry giving up after five attempts. The suite only ever reached the surrounding lines by accident -- the concurrency tests trip a real `PermissionError` on Windows now and then, which is environment-dependent and never happens on Linux CI. The new test pins both halves that matter: the retry is bounded, and giving up still releases the destination rather than wedging every later save to it. * the temp-file cleanup swallowing an `OSError`. It runs in a `finally` while a write exception is propagating, so letting one out would substitute a misleading "cannot remove temp file" for the real disk failure. Reverting the handler makes the new test fail with exactly that substitution. * the fast path in `_await_signal_through_cancellation` for a signal that has already resolved. Exercised directly, since arranging the race through `save()` is not deterministic. No source change. Every line this PR adds is now covered; the three that remain uncovered in the file all predate it. --- .../core/tests/workflow/test_checkpoint.py | 113 ++++++++++++++++++ 1 file changed, 113 insertions(+) diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index d9bee711bb0..d560305d8b3 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -3085,3 +3085,116 @@ async def main() -> None: while canonical in registry and time.monotonic() < deadline: time.sleep(0.05) assert canonical not in registry + + +async def test_file_checkpoint_storage_replace_retry_gives_up_and_releases(monkeypatch): + """A destination that never becomes replaceable must fail loudly, not silently or forever. + + The retry around ``os.replace`` exists for a Windows-specific transient: an indexer or + AV scan briefly holding a handle to the destination. It is bounded on purpose -- a + handle held for good has to surface as an error rather than a hang -- and giving up + must still release the destination, or one stuck file would wedge every later save to + it for the life of the process. + + Both halves are asserted here because the suite only ever reached this code by + accident: the concurrency tests occasionally trip a real ``PermissionError`` on + Windows, which is environment-dependent and absent on Linux CI. + """ + from agent_framework._workflows import _checkpoint as checkpoint_module + + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + attempts: list[int] = [] + real_replace = checkpoint_module.os.replace + + def always_locked(src, dst): # noqa: ANN001, ANN202 - test shim + attempts.append(1) + raise PermissionError("simulated handle held by another process") + + monkeypatch.setattr(checkpoint_module.os, "replace", always_locked) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="locked-id") + + with pytest.raises(PermissionError, match="simulated handle held"): + await storage.save(make("never-lands")) + + assert len(attempts) == 5, f"expected the bounded retry to stop at 5 attempts, got {len(attempts)}" + + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + assert not registry, "giving up on the replace kept the destination owned" + + # No temp file survives the failure, and the path works once the lock clears. + leftovers = await asyncio.to_thread(lambda: list(Path(temp_dir).glob(".maf-ckpt-*.tmp"))) + assert not leftovers, f"the abandoned write leaked temp files: {leftovers}" + + monkeypatch.setattr(checkpoint_module.os, "replace", real_replace) + await asyncio.wait_for(storage.save(make("after-the-lock")), timeout=10) + assert (await storage.load("locked-id")).workflow_name == "after-the-lock" + assert not registry + + +async def test_file_checkpoint_storage_temp_cleanup_failure_does_not_mask_the_write_error(monkeypatch, caplog): + """A temp file that cannot be removed must not replace the error that stranded it. + + The cleanup runs in a ``finally`` while a write exception is propagating. Letting an + ``OSError`` out of it there would substitute a misleading "cannot remove temp file" + for the real disk failure the caller needs to see, and would skip the release that + keeps the destination usable. + """ + from agent_framework._workflows import _checkpoint as checkpoint_module + + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + def exploding_replace(src, dst): # noqa: ANN001, ANN202 - test shim + raise OSError("the real disk failure") + + real_unlink = Path.unlink + + def refuse_temp_unlink(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003, ANN202 - test shim + if self.name.startswith(".maf-ckpt-"): + raise OSError("temp file is not removable either") + return real_unlink(self, *args, **kwargs) + + monkeypatch.setattr(checkpoint_module.os, "replace", exploding_replace) + monkeypatch.setattr(Path, "unlink", refuse_temp_unlink) + + def make(name: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="masked-id") + + with ( + caplog.at_level(logging.DEBUG, logger=checkpoint_module.logger.name), + pytest.raises(OSError, match="the real disk failure"), + ): + await storage.save(make("fails")) + + assert any("Failed to remove checkpoint temp file" in record.message for record in caplog.records), ( + "the swallowed cleanup failure left no diagnostic behind" + ) + + registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage] + assert not registry, "a failed cleanup kept the destination owned" + + monkeypatch.undo() + await asyncio.wait_for(storage.save(make("after-failure")), timeout=10) + assert (await storage.load("masked-id")).workflow_name == "after-failure" + + +async def test_await_signal_through_cancellation_returns_immediately_for_a_resolved_signal(): + """The drain must not suspend on a write that already finished. + + Exercised directly because arranging the race -- a cancellation arriving in the + window after the worker resolved the signal but before the coroutine resumes -- is + not deterministic through ``save()``. The fast path only skips registering a callback + that would fire straight back; this pins the behaviour it is allowed to have. + """ + from agent_framework._workflows import _checkpoint as checkpoint_module + + resolved: ConcurrentFuture[None] = ConcurrentFuture() + resolved.set_result(None) + + await asyncio.wait_for( + checkpoint_module._await_signal_through_cancellation(resolved), # pyright: ignore[reportPrivateUsage] + timeout=5, + )