Python: fix: unique temporary filenames for concurrent FileCheckpointStorage saves (#8182)

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
@@ -328,10 +329,21 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
         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)
+            # Use a unique temporary filename per writer: a fixed name makes
+            # concurrent saves of the same checkpoint id race on the shared
+            # `.tmp` path, where one writer's `os.replace` consumes the file
+            # another writer is about to replace (#8182).
+            tmp_path = file_path.with_suffix(f".json.{os.getpid()}.{uuid.uuid4().hex}.tmp")
+            try:
+                with open(tmp_path, "w") as f:
+                    json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False)
+                os.replace(tmp_path, file_path)
+            except BaseException:
+                # A failed write must not leave its temp file behind, or
+                # repeated failures accumulate unbounded (PR review).
+                with contextlib.suppress(OSError):
+                    tmp_path.unlink(missing_ok=True)
+                raise
 
         await asyncio.to_thread(_write_atomic)
 

python/packages/core/tests/workflow/test_checkpoint.py

@@ -1,7 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import asyncio
 import json
 import tempfile
+import uuid
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
@@ -1275,6 +1277,52 @@ async def test_file_checkpoint_storage_directory_creation():
         assert file_path.exists()
 
 
+async def test_file_checkpoint_storage_concurrent_saves_same_id():
+    """Concurrent saves of the same checkpoint id must all succeed (#8182).
+
+    A fixed temporary filename made writers race on the shared `.tmp` path:
+    one writer's `os.replace` consumed the file another writer was about to
+    replace, failing that save with a raw `FileNotFoundError`.
+    """
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint_id = str(uuid.uuid4())
+        checkpoints = [
+            WorkflowCheckpoint(
+                workflow_name="concurrent-test",
+                graph_signature_hash="sig",
+                checkpoint_id=checkpoint_id,
+                state={"i": i},
+            )
+            for i in range(10)
+        ]
+
+        results = await asyncio.gather(*[storage.save(cp) for cp in checkpoints], return_exceptions=True)
+        errors = [r for r in results if isinstance(r, Exception)]
+        assert not errors, f"concurrent same-id saves failed: {errors}"
+
+        loaded = await storage.load(checkpoint_id)
+        assert isinstance(loaded.state["i"], int)
+
+
+async def test_file_checkpoint_storage_failed_save_cleans_up_temp_file(monkeypatch):
+    """A save that fails mid-write must not leave its temp file behind (PR review)."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(workflow_name="wf", graph_signature_hash="sig", state={"x": 1})
+
+        def failing_replace(src: str, dst: str) -> None:
+            raise OSError("simulated replace failure")
+
+        monkeypatch.setattr("agent_framework._workflows._checkpoint.os.replace", failing_replace)
+        with pytest.raises(OSError, match="simulated replace failure"):
+            await storage.save(checkpoint)
+        monkeypatch.undo()
+
+        leftovers = list(Path(temp_dir).glob("*.tmp"))
+        assert not leftovers, f"temp files leaked from failed save: {leftovers}"
+
+
 async def test_file_checkpoint_storage_corrupted_file():
     with tempfile.TemporaryDirectory() as temp_dir:
         storage = FileCheckpointStorage(temp_dir)