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

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

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 `<id>.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)