Skip to content
152 changes: 141 additions & 11 deletions python/packages/core/agent_framework/_workflows/_checkpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,14 @@
import os
import threading
import time
import unicodedata
import uuid
from collections.abc import Callable, Mapping
from concurrent.futures import Future
from dataclasses import dataclass, field, fields
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeAlias
from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeAlias, cast

from ..exceptions import WorkflowCheckpointException

Expand Down Expand Up @@ -339,7 +340,7 @@ def _encode_checkpoint_for_storage(checkpoint: WorkflowCheckpoint, allowed_types
# checkpoint ID from different processes still rely on ``os.replace`` being atomic on
# the underlying filesystem for the file not to be seen half-written; which write
# survives is undefined.
_destination_queues: dict[Path, _DestinationQueue] = {}
_destination_queues: dict[tuple[Path, str], _DestinationQueue] = {}
# A completed predecessor invokes callbacks synchronously. If cancellation lands while
# a save is entering the queue, its deferred release can therefore run re-entrantly on
# the same thread before the enqueue frame has left this guard. A plain Lock deadlocks
Expand Down Expand Up @@ -375,14 +376,37 @@ class _WriteTicket:
error: BaseException | None = None


def _fs_key(file_name: str) -> str:
"""The filesystem's folded spelling of a filename.

Case-insensitive volumes (the default on Windows and macOS) merge ``a`` and
``A``; APFS and HFS+ additionally merge Unicode normalization forms. Both
folds are lossy, so two names sharing a key cannot coexist there.
``str.casefold()`` alone does not normalize.
"""
return unicodedata.normalize("NFC", file_name).casefold()


def _destination_key(file_path: Path) -> tuple[Path, str]:
"""The queue key for *file_path*: one entry per filesystem-folded name.

Variant spellings of one checkpoint (``Round-A.json`` and ``round-a.json``)
resolve to a single file on a folding volume, so their writes must share one
queue or the second replace can land first and overwrite the newer one
(#9466 review).
"""
return (file_path.parent, _fs_key(file_path.name))


def _enqueue_write(file_path: Path) -> _WriteTicket:
"""Take a place in *file_path*'s queue without waiting for it."""
completion: Future[None] = Future()
key = _destination_key(file_path)
with _destination_queues_guard:
queue = _destination_queues.get(file_path)
queue = _destination_queues.get(key)
if queue is None:
queue = _DestinationQueue()
_destination_queues[file_path] = queue
_destination_queues[key] = queue
predecessor = queue.tail
queue.tail = completion
queue.pending += 1
Expand Down Expand Up @@ -543,11 +567,11 @@ def _release_one_write(ticket: _WriteTicket) -> None:
if ticket.released:
return
ticket.released = True
queue = _destination_queues.get(ticket.path)
queue = _destination_queues.get(_destination_key(ticket.path))
if queue is not None:
queue.pending -= 1
if queue.pending <= 0:
del _destination_queues[ticket.path]
del _destination_queues[_destination_key(ticket.path)]
# Outside the guard: waking the successor must not happen while holding the lock its
# own release will need.
if not ticket.completion.done():
Expand Down Expand Up @@ -623,6 +647,87 @@ def _validate_file_path(self, checkpoint_id: CheckpointID) -> Path:
raise WorkflowCheckpointException(f"Invalid checkpoint ID: {checkpoint_id}")
return file_path

@staticmethod
def _read_stored_checkpoint_id(file_path: Path) -> CheckpointID:
"""Read the persisted identity used to distinguish filesystem aliases."""
try:
with open(file_path, encoding="utf-8") as f:
encoded_checkpoint = json.load(f)
except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
raise WorkflowCheckpointException(f"Cannot verify the checkpoint identity stored in {file_path}.") from exc
if not isinstance(encoded_checkpoint, dict):
raise WorkflowCheckpointException(f"Cannot verify the checkpoint identity stored in {file_path}.")
checkpoint_data = cast("dict[str, Any]", encoded_checkpoint)
stored_id = checkpoint_data.get("checkpoint_id")
if not isinstance(stored_id, str):
raise WorkflowCheckpointException(f"Cannot verify the checkpoint identity stored in {file_path}.")
return stored_id

@staticmethod
def _assert_no_variant_collision(file_path: Path, checkpoint_id: CheckpointID) -> None:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please enforce the same exact checkpoint-ID identity on load() and delete(). On a case/normalization-folding filesystem, after saving Round-A, _validate_file_path("round-a") resolves to the same existing file: load("round-a") returns the other checkpoint and delete("round-a") deletes Round-A. Under the fail-closed policy, a variant ID should not read or remove a checkpoint whose stored ID/spelling differs; add regressions for both operations.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both operations now go through _assert_stored_spelling, which compares the directory's stored spelling against the requested one and refuses when only a case/normalization fold brings them together: load("round-a") raises rather than returning Round-A's state, and delete("round-a") raises rather than unlinking it. Same-ID access and genuinely distinct IDs on a case-sensitive volume are unchanged. Regressions are test_file_load_rejects_case_variant_checkpoint_id and test_file_delete_rejects_case_variant_checkpoint_id in the conformance file.

"""Reject an ID that the filesystem cannot store as an independent file."""
if not file_path.exists():
return
# Use the requested spelling because resolve() can return the on-disk casing on Windows.
expected_name = f"{checkpoint_id}.json"
try:
siblings = [entry.name for entry in file_path.parent.iterdir()]
except OSError as exc:
raise WorkflowCheckpointException(
f"Cannot verify checkpoint ID {checkpoint_id!r} against existing checkpoints: "
"the storage directory could not be listed."
) from exc
expected_key = _fs_key(expected_name)
variant = next((name for name in siblings if _fs_key(name) == expected_key), None)
if variant is None:
return
# Decomposing filesystems can return an exact spelling that still belongs to another ID.
if FileCheckpointStorage._read_stored_checkpoint_id(file_path) == checkpoint_id:
return
raise WorkflowCheckpointException(
f"Checkpoint ID {checkpoint_id!r} collides with the existing checkpoint "
f"{variant.removesuffix('.json')!r}: on this filesystem both names "
"resolve to one file, and saving would replace that checkpoint's data. "
"Choose a checkpoint ID that differs by more than letter case or "
"Unicode normalization form."
)

@staticmethod
def _assert_stored_spelling(
file_path: Path,
checkpoint_id: CheckpointID,
encoded_checkpoint: dict[str, Any] | None = None,
) -> None:
"""Refuse to act on a file that belongs to another checkpoint ID.

Filesystems can change the spelling returned by directory enumeration,
so the persisted checkpoint ID, rather than the filename alone, decides identity.
"""
expected_name = f"{checkpoint_id}.json"
try:
siblings = [entry.name for entry in file_path.parent.iterdir()]
except OSError as exc:
raise WorkflowCheckpointException(
f"Cannot verify checkpoint ID {checkpoint_id!r} against existing checkpoints: "
"the storage directory could not be listed."
) from exc
variant = next((name for name in siblings if _fs_key(name) == _fs_key(expected_name)), None)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yufeng He (@he-yufeng), this rejects a folded sibling without checking whether the requested path actually exists. On a case-sensitive filesystem, after saving Round-A, delete("round-a") now raises WorkflowCheckpointException, even though round-a.json is absent and the documented result for a missing ID is False. These are independent IDs on that filesystem, so there is no alias to prevent. Please skip the spelling rejection when the requested destination does not exist and add a case-sensitive regression confirming that deleting the absent variant returns False without touching Round-A.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Implemented in 8126b01dd99ea9977eea188fe09d9949dbacc46f.

Delete now checks whether the requested destination exists before doing any variant verification, all inside the worker under the deletion lock. An absent independent round-a.json therefore returns False without touching Round-A.json.

The case-sensitive regression is included in the checkpoint conformance file. It is the single skipped test in this APFS run, so I am not claiming a local case-sensitive execution; the remaining 111 tests in the two-file checkpoint sweep passed.

if variant is None:
return
if encoded_checkpoint is None:
stored_id = FileCheckpointStorage._read_stored_checkpoint_id(file_path)
elif isinstance(encoded_checkpoint, dict) and isinstance(encoded_checkpoint.get("checkpoint_id"), str):
stored_id = encoded_checkpoint["checkpoint_id"]
else:
raise WorkflowCheckpointException(f"Cannot verify the checkpoint identity stored in {file_path}.")
if stored_id == checkpoint_id:
return
raise WorkflowCheckpointException(
f"No checkpoint found with ID {checkpoint_id}: the stored checkpoint is spelled "
f"{variant.removesuffix('.json')!r}, and on this filesystem both names resolve to "
"one file. Refusing to read or delete another checkpoint's data."
)

@classmethod
def _delete_lock(cls, file_path: Path) -> threading.Lock:
"""Return the process-local deletion lock for a checkpoint file."""
Expand Down Expand Up @@ -672,6 +777,15 @@ def _write_atomic() -> None:
# filesystem). A short, ID-independent name keeps every destination
# name accepted by _validate_file_path saveable regardless of
# checkpoint-ID length or filesystem limits.
#
# The collision guard is repeated here, under destination ownership,
# because the pre-enqueue check cannot see a variant sibling another
# save is about to create: two first-time variant saves both observed
# "no destination" and the second replace would win (#9466 review).
# Variant spellings share one queue, so this re-check runs after the
# predecessor has actually landed. Cross-process races stay a
# documented limitation.
self._assert_no_variant_collision(file_path, checkpoint.checkpoint_id)
tmp_path: Path | None = None
try:
# O_CREAT | O_EXCL | O_WRONLY with an explicit 0o666 mode, so the
Expand Down Expand Up @@ -796,19 +910,27 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:

Raises:
WorkflowCheckpointException: If no checkpoint with the given ID exists,
or if checkpoint decoding fails.
or if checkpoint decoding fails. Also when the ID resolves to a
differently spelled checkpoint only through a filesystem
case/normalization fold: under fail-closed, that ID names no
checkpoint and the other one's data is not read.
"""
file_path = self._validate_file_path(checkpoint_id)

if not file_path.exists():
raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}")

def _read() -> dict[str, Any]:
if not file_path.exists():
raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}")
with open(file_path, encoding="utf-8") as f:
return json.load(f)
encoded_checkpoint = json.load(f)
self._assert_stored_spelling(file_path, checkpoint_id, encoded_checkpoint)
return encoded_checkpoint

try:
encoded_checkpoint = await asyncio.to_thread(_read)
except WorkflowCheckpointException:
raise
except FileNotFoundError as ex:
raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}") from ex
except UnicodeDecodeError as ex:
raise WorkflowCheckpointException(
f"Checkpoint file for {checkpoint_id} is not valid UTF-8 and cannot be loaded."
Expand Down Expand Up @@ -866,12 +988,20 @@ async def delete(self, checkpoint_id: CheckpointID) -> bool:

Returns:
True if the checkpoint was successfully deleted, False if no checkpoint with the given ID exists.

Raises:
WorkflowCheckpointException: If the ID resolves to a differently spelled
checkpoint only through a filesystem case/normalization fold, or the
storage directory cannot be listed to rule that out.
"""
file_path = self._validate_file_path(checkpoint_id)
file_lock = self._delete_lock(file_path)

def _delete() -> bool:
with file_lock:
if not file_path.exists():
return False
self._assert_stored_spelling(file_path, checkpoint_id)
try:
file_path.unlink()
except FileNotFoundError:
Expand Down
23 changes: 15 additions & 8 deletions python/packages/core/tests/workflow/test_checkpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -2336,7 +2336,8 @@ def make(name: str) -> WorkflowCheckpoint:
# queued behind the holder before cancelling, or this would be testing a
# cancellation that raced the enqueue instead.
registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage]
canonical = (Path(temp_dir) / "shared-id.json").resolve()
# The destination queue is keyed by the filesystem-folded name (#9466).
canonical = checkpoint_module._destination_key((Path(temp_dir) / "shared-id.json").resolve())
for _ in range(400):
entry = registry.get(canonical)
if entry is not None and entry.pending == 3:
Expand Down Expand Up @@ -2566,7 +2567,7 @@ def test_release_write_tolerates_an_already_dropped_entry():
# No entry was ever created for this path.
checkpoint_module._release_write(ticket) # pyright: ignore[reportPrivateUsage]
assert ticket.completion.done()
assert path not in checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage]
assert checkpoint_module._destination_key(path) not in checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage]


def test_file_checkpoint_storage_abandoned_loop_does_not_own_destination_forever(monkeypatch, tmp_path):
Expand Down Expand Up @@ -2734,7 +2735,8 @@ def make(name: str) -> WorkflowCheckpoint:
third = asyncio.create_task(storage.save(make("third")))

registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage]
canonical = (Path(temp_dir) / "shared-id.json").resolve()
# The destination queue is keyed by the filesystem-folded name (#9466).
canonical = checkpoint_module._destination_key((Path(temp_dir) / "shared-id.json").resolve())
for _ in range(400):
entry = registry.get(canonical)
if entry is not None and entry.pending == 3:
Expand Down Expand Up @@ -2793,7 +2795,8 @@ def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim

monkeypatch.setattr(checkpoint_module.os, "replace", gated_replace)

canonical = (tmp_path / "shared-id.json").resolve()
# The destination queue is keyed by the filesystem-folded name (#9466).
canonical = checkpoint_module._destination_key((tmp_path / "shared-id.json").resolve())
observed: dict[str, bool] = {}

async def main() -> None:
Expand Down Expand Up @@ -2849,7 +2852,8 @@ def test_file_checkpoint_storage_shutdown_before_the_write_starts_does_not_hang(

from agent_framework._workflows import _checkpoint as checkpoint_module

canonical = (tmp_path / "shared-id.json").resolve()
# The destination queue is keyed by the filesystem-folded name (#9466).
canonical = checkpoint_module._destination_key((tmp_path / "shared-id.json").resolve())
outcome: dict[str, object] = {}

# Hold the submitted write at its first syscall. Without this the count below races the
Expand Down Expand Up @@ -3134,7 +3138,8 @@ def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim
def make(name: str) -> WorkflowCheckpoint:
return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id")

canonical = (tmp_path / "shared-id.json").resolve()
# The destination queue is keyed by the filesystem-folded name (#9466).
canonical = checkpoint_module._destination_key((tmp_path / "shared-id.json").resolve())
registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage]

holder_loop = asyncio.new_event_loop()
Expand Down Expand Up @@ -3231,7 +3236,8 @@ def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim
def make(name: str) -> WorkflowCheckpoint:
return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id")

canonical = (tmp_path / "shared-id.json").resolve()
# The destination queue is keyed by the filesystem-folded name (#9466).
canonical = checkpoint_module._destination_key((tmp_path / "shared-id.json").resolve())
registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage]

holder_loop = asyncio.new_event_loop()
Expand Down Expand Up @@ -3324,7 +3330,8 @@ def gated_replace(src, dst): # noqa: ANN001, ANN202 - test shim
def make(name: str) -> WorkflowCheckpoint:
return WorkflowCheckpoint(workflow_name=name, graph_signature_hash="test-hash", checkpoint_id="shared-id")

canonical = (tmp_path / "shared-id.json").resolve()
# The destination queue is keyed by the filesystem-folded name (#9466).
canonical = checkpoint_module._destination_key((tmp_path / "shared-id.json").resolve())
registry = checkpoint_module._destination_queues # pyright: ignore[reportPrivateUsage]

def wait_for_pending(count: int, timeout: float = 20.0) -> None:
Expand Down
Loading
Loading