From bf7d2eedf1d60e18685dbd0214e238a367c9bdc5 Mon Sep 17 00:00:00 2001 From: fzfzzfzzzfzzzz Date: Fri, 11 Sep 2026 00:35:53 +0800 Subject: [PATCH] Python: handle concurrent checkpoint deletion --- .../agent_framework/_workflows/_checkpoint.py | 9 +++--- .../core/tests/workflow/test_checkpoint.py | 32 +++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index d9b30d83919..31d5bf59d47 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -414,11 +414,12 @@ async def delete(self, checkpoint_id: CheckpointID) -> bool: file_path = self._validate_file_path(checkpoint_id) def _delete() -> bool: - if file_path.exists(): + try: file_path.unlink() - logger.info(f"Deleted checkpoint {checkpoint_id} from {file_path}") - return True - return False + except FileNotFoundError: + return False + logger.info(f"Deleted checkpoint {checkpoint_id} from {file_path}") + return True return await asyncio.to_thread(_delete) diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index be8de7c13b6..be8ada37648 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1,11 +1,14 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import json import tempfile +import threading from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any +from unittest.mock import patch import pytest @@ -1258,6 +1261,35 @@ async def test_file_checkpoint_storage_delete(): assert result is False +async def test_file_checkpoint_storage_concurrent_delete(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + checkpoint = WorkflowCheckpoint( + workflow_name="test-workflow", + graph_signature_hash="test-hash", + checkpoint_id="same", + ) + await storage.save(checkpoint) + + file_path = (Path(temp_dir) / "same.json").resolve() + original_exists = Path.exists + barrier = threading.Barrier(2) + + def synchronized_exists(path: Path) -> bool: + result = original_exists(path) + if path.resolve() == file_path: + barrier.wait(timeout=5) + return result + + with patch.object(Path, "exists", synchronized_exists): + results = await asyncio.gather( + storage.delete(checkpoint.checkpoint_id), + storage.delete(checkpoint.checkpoint_id), + ) + + assert sorted(results) == [False, True] + + async def test_file_checkpoint_storage_directory_creation(): with tempfile.TemporaryDirectory() as temp_dir: nested_path = Path(temp_dir) / "nested" / "checkpoint" / "storage"