diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index cd78725d6a1..0f89e2cf165 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -10,6 +10,7 @@ import os import threading import time +import unicodedata import uuid from collections.abc import Callable, Mapping from concurrent.futures import Future @@ -287,7 +288,7 @@ async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID] # 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 @@ -323,14 +324,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 @@ -491,11 +515,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(): @@ -571,6 +595,88 @@ def _validate_file_path(self, checkpoint_id: CheckpointID) -> Path: raise WorkflowCheckpointException(f"Invalid checkpoint ID: {checkpoint_id}") return file_path + @staticmethod + def _assert_no_variant_collision(file_path: Path, checkpoint_id: CheckpointID) -> None: + """Reject an ID that the filesystem cannot store as an independent file. + + On case-insensitive filesystems (the default on Windows and macOS), + ``Round-A.json`` and ``round-a.json`` name the same file, so a save under + the variant spelling would silently replace the other checkpoint and its + reads would then return the wrong workflow's state (#9444). APFS and HFS+ + additionally fold Unicode normalization forms, so ``café-1.json`` spelled + NFC and NFD is also one file there (#9466 review). On filesystems without + those folds the variant path simply does not exist, so the directory scan + below is skipped entirely there. + """ + if not file_path.exists(): + return + # Derive the expected spelling from the requested ID, not the resolved + # path: on Windows Path.resolve() canonicalizes an existing file to its + # on-disk casing, so file_path.name would already be the colliding + # name and the check below would accept the overwrite (#9444). + expected_name = f"{checkpoint_id}.json" + try: + siblings = [entry.name for entry in file_path.parent.iterdir()] + except OSError as exc: + # Permissions can allow lookup/creation/replacement while denying + # enumeration, so skipping the scan here would let a variant + # save proceed unchecked. Fail closed instead. + raise WorkflowCheckpointException( + f"Cannot verify checkpoint ID {checkpoint_id!r} against existing checkpoints: " + "the storage directory could not be listed." + ) from exc + # The filesystem key is the NFC-normalized, case-folded spelling: both + # folds are lossy on the volumes that perform them, so two IDs sharing + # a key cannot coexist there. str.casefold() alone does not normalize. + expected_key = _fs_key(expected_name) + variant = next( + (name for name in siblings if name != expected_name and _fs_key(name) == expected_key), + None, + ) + if variant is None or expected_name in siblings: + # No filesystem-folded variant exists, or the exact spelling is + # already stored (a same-ID update on a volume where distinct + # spellings legitimately coexist). + 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) -> None: + """Refuse to act on a file that only a filesystem fold resolved to. + + On a case- or normalization-folding volume, ``load("round-a")`` resolves + to the stored ``Round-A.json`` and would read another checkpoint's data; + ``delete("round-a")`` would remove it. Under the fail-closed policy a + variant ID names no checkpoint at all, so both operations must stop + where the stored spelling differs (#9466 review). + """ + expected_name = f"{checkpoint_id}.json" + try: + siblings = [entry.name for entry in file_path.parent.iterdir()] + except OSError as exc: + # Enumeration can fail where lookup succeeds; acting on a fold we + # could not inspect is exactly the aliasing this guard prevents. + raise WorkflowCheckpointException( + f"Cannot verify checkpoint ID {checkpoint_id!r} against existing checkpoints: " + "the storage directory could not be listed." + ) from exc + if expected_name in siblings: + return + variant = next((name for name in siblings if _fs_key(name) == _fs_key(expected_name)), None) + if variant is None: + 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.""" @@ -592,6 +698,7 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value file_path = self._validate_file_path(checkpoint.checkpoint_id) + self._assert_no_variant_collision(file_path, checkpoint.checkpoint_id) checkpoint_dict = checkpoint.to_dict() # Fail at save time if encoding or restore validation fails (#8181). try: @@ -633,6 +740,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 @@ -757,9 +873,13 @@ 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) + self._assert_stored_spelling(file_path, checkpoint_id) if not file_path.exists(): raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}") @@ -827,8 +947,14 @@ 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) + self._assert_stored_spelling(file_path, checkpoint_id) file_lock = self._delete_lock(file_path) def _delete() -> bool: diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 6e13a414f8b..db69850af17 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -2277,7 +2277,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: @@ -2507,7 +2508,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): @@ -2675,7 +2676,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: @@ -2734,7 +2736,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: @@ -2790,7 +2793,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 @@ -3075,7 +3079,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() @@ -3172,7 +3177,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() @@ -3265,7 +3271,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: diff --git a/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py b/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py index e1ce83d81f8..866ee04069a 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py +++ b/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py @@ -4,6 +4,7 @@ from __future__ import annotations +import asyncio from pathlib import Path from typing import Any, cast @@ -152,3 +153,190 @@ async def test_load_missing_id_raises(conformance_storage: CheckpointStorage) -> async def test_save_returns_id(conformance_storage: CheckpointStorage) -> None: checkpoint = _conformance_checkpoint() assert await conformance_storage.save(checkpoint) == checkpoint.checkpoint_id + + +def _filesystem_is_case_insensitive(path: Path) -> bool: + probe = path / "CkAsE.probe" + probe.write_text("x") + try: + return (path / "ckase.probe").exists() + finally: + probe.unlink() + + +async def test_file_save_rejects_case_variant_checkpoint_id(tmp_path: Path) -> None: + """A case-only variant ID must not silently replace the other checkpoint (#9444).""" + # Probe inside the test on tmp_path: a module-level probe races under + # pytest-xdist workers sharing the source directory, and tmp_path is the + # volume that actually holds the checkpoints. + if not _filesystem_is_case_insensitive(tmp_path): + pytest.skip("case-variant collision only occurs on case-insensitive filesystems") + storage = FileCheckpointStorage(tmp_path / "checkpoints") + first = WorkflowCheckpoint( + workflow_name="first-workflow", + graph_signature_hash="topology-a", + checkpoint_id="Round-A", + state={"marker": "first"}, + ) + second = WorkflowCheckpoint( + workflow_name="second-workflow", + graph_signature_hash="topology-b", + checkpoint_id="round-a", + state={"marker": "second"}, + ) + + first_id = await storage.save(first) + with pytest.raises(WorkflowCheckpointException, match="collides"): + await storage.save(second) + + # The first checkpoint survives untouched: its data and listing are intact. + restored = await storage.load(first_id) + assert restored.checkpoint_id == "Round-A" + assert restored.workflow_name == "first-workflow" + assert restored.state["marker"] == "first" + listed = await storage.list_checkpoint_ids(workflow_name="first-workflow") + assert listed == ["Round-A"] + + # Re-saving the exact same ID remains a normal update, not a collision. + first.state["marker"] = "first-updated" + assert await storage.save(first) == "Round-A" + assert (await storage.load("Round-A")).state["marker"] == "first-updated" + + + +async def test_file_load_rejects_case_variant_checkpoint_id(tmp_path: Path) -> None: + """A variant ID must not read another checkpoint's data (#9466 review). + + On a folding filesystem ``load("round-a")`` resolves to the stored + ``Round-A.json``; under fail-closed it names no checkpoint at all. + """ + if not _filesystem_is_case_insensitive(tmp_path): + pytest.skip("case-variant collision only occurs on case-insensitive filesystems") + storage = FileCheckpointStorage(tmp_path / "checkpoints") + first = WorkflowCheckpoint( + workflow_name="first-workflow", + graph_signature_hash="topology-a", + checkpoint_id="Round-A", + state={"marker": "first"}, + ) + await storage.save(first) + + with pytest.raises(WorkflowCheckpointException, match="spelled 'Round-A'"): + await storage.load("round-a") + + # The stored checkpoint is untouched and still reachable under its own ID. + restored = await storage.load("Round-A") + assert restored.state["marker"] == "first" + + +async def test_file_delete_rejects_case_variant_checkpoint_id(tmp_path: Path) -> None: + """A variant ID must not delete another checkpoint (#9466 review).""" + if not _filesystem_is_case_insensitive(tmp_path): + pytest.skip("case-variant collision only occurs on case-insensitive filesystems") + storage = FileCheckpointStorage(tmp_path / "checkpoints") + first = WorkflowCheckpoint( + workflow_name="first-workflow", + graph_signature_hash="topology-a", + checkpoint_id="Round-A", + state={"marker": "first"}, + ) + await storage.save(first) + + with pytest.raises(WorkflowCheckpointException, match="spelled 'Round-A'"): + await storage.delete("round-a") + + assert (await storage.load("Round-A")).state["marker"] == "first" + assert await storage.list_checkpoint_ids(workflow_name="first-workflow") == ["Round-A"] + + +async def test_concurrent_variant_saves_do_not_overwrite(tmp_path: Path) -> None: + """Two first-time variant saves must not both pass the guard (#9466 review). + + The collision check used to run before the write was queued, so both saves + observed "no destination" and the second atomic replace silently won. + """ + if not _filesystem_is_case_insensitive(tmp_path): + pytest.skip("case-variant collision only occurs on case-insensitive filesystems") + storage = FileCheckpointStorage(tmp_path / "checkpoints") + upper = WorkflowCheckpoint( + workflow_name="upper-workflow", + graph_signature_hash="topology-a", + checkpoint_id="Round-B", + state={"marker": "upper"}, + ) + lower = WorkflowCheckpoint( + workflow_name="lower-workflow", + graph_signature_hash="topology-b", + checkpoint_id="round-b", + state={"marker": "lower"}, + ) + + results = await asyncio.gather(storage.save(upper), storage.save(lower), return_exceptions=True) + errors = [r for r in results if isinstance(r, WorkflowCheckpointException)] + saved_ids = [r for r in results if isinstance(r, str)] + assert len(errors) == 1, f"expected exactly one save to be rejected, got {results!r}" + assert "collides" in str(errors[0]) + assert len(saved_ids) == 1 + + # Whichever spelling won keeps its own data; nothing is silently replaced. + stored = await storage.list_checkpoint_ids(workflow_name=f"{'upper' if saved_ids[0] == 'Round-B' else 'lower'}-workflow") + assert stored == [saved_ids[0]] + assert (await storage.load(saved_ids[0])).state["marker"] == ("upper" if saved_ids[0] == "Round-B" else "lower") + + +def _filesystem_folds_unicode_normalization(path: Path) -> bool: + """True when the volume treats NFC and NFD spellings of a name as one file.""" + import unicodedata + + nfc = unicodedata.normalize("NFC", "café-norm.probe") + nfd = unicodedata.normalize("NFD", "café-norm.probe") + probe = path / nfc + probe.write_text("x") + try: + return (path / nfd).exists() + finally: + probe.unlink() + + +@pytest.mark.parametrize("mixed_case", [False, True], ids=["nfc-nfd", "case-nfc-nfd"]) +async def test_file_save_rejects_normalization_variant_checkpoint_id(tmp_path: Path, mixed_case: bool) -> None: + """An NFC/NFD spelling variant must not silently replace a checkpoint (#9466). + + APFS and HFS+ fold Unicode normalization forms, so two IDs that differ only + in normalization name one file there; str.casefold() alone does not catch + it. On volumes without that fold the two files coexist and there is nothing + to reject. + """ + import unicodedata + + if not _filesystem_folds_unicode_normalization(tmp_path): + pytest.skip("normalization-variant collision only occurs on normalization-insensitive filesystems") + first_spelling = unicodedata.normalize("NFC", "Café-1" if mixed_case else "café-1") + second_spelling = unicodedata.normalize("NFD", "café-1") + assert first_spelling != second_spelling # the test would be vacuous otherwise + + storage = FileCheckpointStorage(tmp_path / "checkpoints") + first = WorkflowCheckpoint( + workflow_name="first-workflow", + graph_signature_hash="topology-a", + checkpoint_id=first_spelling, + state={"marker": "first"}, + ) + second = WorkflowCheckpoint( + workflow_name="second-workflow", + graph_signature_hash="topology-b", + checkpoint_id=second_spelling, + state={"marker": "second"}, + ) + + first_id = await storage.save(first) + with pytest.raises(WorkflowCheckpointException, match="collides"): + await storage.save(second) + + restored = await storage.load(first_id) + assert restored.checkpoint_id == first_spelling + assert restored.state["marker"] == "first" + + # Re-saving the exact same spelling remains a normal update. + first.state["marker"] = "first-updated" + assert await storage.save(first) == first_spelling