From f789e98e441aa1da7bbcaa356b26b72130dbff42 Mon Sep 17 00:00:00 2001 From: ptimizeroracle Date: Thu, 10 Sep 2026 10:59:44 +0200 Subject: [PATCH 1/2] fix: unique temporary filenames for concurrent FileCheckpointStorage saves (#8182) Signed-off-by: ptimizeroracle --- .../agent_framework/_workflows/_checkpoint.py | 6 +++- .../core/tests/workflow/test_checkpoint.py | 30 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index d9b30d83919..58330aa1e4b 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -328,7 +328,11 @@ 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") + # 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") with open(tmp_path, "w") as f: json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False) os.replace(tmp_path, file_path) diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index be8de7c13b6..bf103b54d47 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/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,34 @@ 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_corrupted_file(): with tempfile.TemporaryDirectory() as temp_dir: storage = FileCheckpointStorage(temp_dir) From 6b7311d42c0b75ed9f5894fc381ac939dcd09ee0 Mon Sep 17 00:00:00 2001 From: ptimizeroracle Date: Thu, 10 Sep 2026 11:29:06 +0200 Subject: [PATCH 2/2] fix: clean up temp file when a checkpoint write fails (PR review) Signed-off-by: ptimizeroracle --- .../agent_framework/_workflows/_checkpoint.py | 14 +++++++++++--- .../core/tests/workflow/test_checkpoint.py | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 58330aa1e4b..b3681774c13 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 @@ -333,9 +334,16 @@ def _write_atomic() -> None: # `.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") - with open(tmp_path, "w") as f: - json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False) - os.replace(tmp_path, file_path) + 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) diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index bf103b54d47..9717c13c6da 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1305,6 +1305,24 @@ async def test_file_checkpoint_storage_concurrent_saves_same_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)