diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 3de9460c86..81124186a4 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -7,6 +7,8 @@ import json import logging import os +import threading +import time import uuid from collections.abc import Mapping from dataclasses import dataclass, field, fields @@ -246,6 +248,33 @@ 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. +# +# 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: """File-based checkpoint storage for persistence. @@ -325,13 +354,71 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: checkpoint_dict = checkpoint.to_dict() encoded_checkpoint = encode_checkpoint_value(checkpoint_dict) - 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) + 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. + 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)) - await asyncio.to_thread(_write_atomic) + 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)) 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 5f3da78cd1..2a52e8a833 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import json +import os import tempfile from dataclasses import dataclass from datetime import datetime, timezone @@ -1115,6 +1117,193 @@ 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(): + """Process-wide file-lock registry must not grow with sequential saves of the same ID. + + 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) + 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="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] + + +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(): with tempfile.TemporaryDirectory() as temp_dir: storage = FileCheckpointStorage(temp_dir)