diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index cd78725d6a1..79155df7531 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -8,6 +8,7 @@ import json import logging import os +import stat import threading import time import uuid @@ -29,6 +30,17 @@ # Type alias for checkpoint IDs in case we want to change the # underlying type in the future (e.g., to UUID or a custom class) CheckpointID: TypeAlias = str +_MAX_GET_LATEST_ATTEMPTS = 3 +_GET_LATEST_RETRY_DELAY_SECONDS = 0.05 + + +class _CheckpointScanInvalidated(WorkflowCheckpointException): + """Internal signal that an optimistic scan was invalidated by concurrent changes. + + Raised only inside the worker thread used by ``get_latest`` and always consumed + by its retry loop. Callers see a plain :class:`WorkflowCheckpointException` + only when every attempt is invalidated. + """ @dataclass(slots=True) @@ -779,16 +791,23 @@ def _read() -> dict[str, Any]: f"Checkpoint file for {checkpoint_id} is not valid JSON and cannot be loaded." ) from ex - from ._checkpoint_encoding import decode_checkpoint_value - - try: - decoded_checkpoint_dict = decode_checkpoint_value(encoded_checkpoint, allowed_types=self._allowed_types) - except WorkflowCheckpointException: - raise - checkpoint = WorkflowCheckpoint.from_dict(decoded_checkpoint_dict) + checkpoint = self._decode_checkpoint(encoded_checkpoint) logger.info(f"Loaded checkpoint {checkpoint_id} from {file_path}") return checkpoint + def _decode_checkpoint(self, encoded_checkpoint: dict[str, Any]) -> WorkflowCheckpoint: + """Decode an already-read checkpoint using this reader's deserialization policy. + + Shared by :meth:`load` and :meth:`get_latest` so both apply the same + ``allowed_checkpoint_types`` policy. Callers must have already validated the + on-disk state if snapshot consistency matters; this method only handles + deserialization. + """ + from ._checkpoint_encoding import decode_checkpoint_value + + decoded = decode_checkpoint_value(encoded_checkpoint, allowed_types=self._allowed_types) + return WorkflowCheckpoint.from_dict(decoded) + async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: """List checkpoint objects for a given workflow name. @@ -845,18 +864,208 @@ def _delete() -> bool: async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: """Get the latest checkpoint for a given workflow name. + Select from stored metadata and decode only the selected checkpoint, using the same + data that was read during selection. Files whose owning workflow cannot be + identified (unparseable JSON, a missing or non-string ``workflow_name``, or + duplicate ``workflow_name`` keys) are quarantined: they are skipped with a + warning so a single malformed or legacy file cannot permanently break recovery + for every workflow sharing this storage root. Invalid content in a file that IS + identified as belonging to the requested workflow still raises, and invalid + metadata for a known different workflow is ignored. + File versions and directory membership are checked before accepting the scan. + Concurrent changes under the storage root, including an atomic save for a + different workflow, invalidate the optimistic scan; the scan is retried + internally a bounded number of times before raising, so ordinary concurrent + checkpointing does not fail unrelated recovery. Writes after validation belong + to a subsequent snapshot. This relies on the filesystem reporting file changes + and writers using atomic replacement, as ``save`` does. + Checkpoints with equal timestamps retain directory iteration order. + Nested directories are included, with in-tree directory aliases visited only once. + Reading and decoding run in a worker thread so decoding does not block the event + loop. + Args: workflow_name: The name of the workflow to get the latest checkpoint for. Returns: - The latest WorkflowCheckpoint object for the specified workflow name, or None if no checkpoints exist. + The latest WorkflowCheckpoint from the validated scan, or None if no matching + checkpoints exist. + + Raises: + WorkflowCheckpointException: If a checkpoint owned by ``workflow_name`` cannot + be read or decoded, the directory cannot be read, matching metadata is + invalid, or the scan changes on every retry attempt. """ - checkpoints = await self.list_checkpoints(workflow_name=workflow_name) - if not checkpoints: - return None - latest_checkpoint = max(checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp)) - logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}") - return latest_checkpoint + + def _version(info: os.stat_result) -> tuple[int, int, int]: + # On Windows, Path.stat() and os.fstat() use different Win32 APIs and can + # return different values for st_dev (volume serial number) and st_ino + # (file ID). Only st_size + st_mtime_ns + st_ctime_ns are reliable across + # both APIs, which is sufficient to detect both content changes (mtime, + # size) and file replacement (ctime on Windows is creation time, which + # updates on os.replace). Symlink escapes are caught separately by + # _validate_file_path and the is_relative_to(root) check in _checkpoint_paths. + return info.st_size, info.st_mtime_ns, info.st_ctime_ns + + def _checkpoint_paths(root: Path) -> tuple[list[Path], dict[Path, tuple[int, int, int]]]: + paths: list[Path] = [] + directories: dict[Path, tuple[int, int, int]] = {} + pending = [root] + while pending: + directory = pending.pop().resolve() + if not directory.is_relative_to(root): + raise WorkflowCheckpointException( + f"Checkpoint directory {directory} is outside storage root {root}" + ) + if directory in directories: + continue + directories[directory] = _version(directory.stat()) + with os.scandir(directory) as entries: + for entry in entries: + if entry.is_dir(): + pending.append(Path(entry.path)) + elif os.path.normcase(entry.name).endswith(".json"): + paths.append(Path(entry.path)) + return paths, directories + + def _make_duplicate_capturing_hook() -> tuple[ + Callable[[list[tuple[str, Any]]], dict[str, Any]], + list[str], + ]: + """Return a JSON ``object_pairs_hook`` that records duplicate keys instead of raising. + + The hook is paired with a list that collects duplicate key names as they are + encountered during parsing. The caller decides when to raise based on workflow + ownership, so a foreign-workflow file with duplicate keys does not block + recovery of healthy checkpoints for the requested workflow (per the documented + policy that invalid metadata for a known different workflow is ignored). + """ + duplicate_keys: list[str] = [] + + def hook(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + duplicate_keys.append(key) + else: + result[key] = value + return result + + return hook, duplicate_keys + + def _read_latest() -> WorkflowCheckpoint | None: + root = self.storage_path.resolve() + paths, directories = _checkpoint_paths(root) + versions: dict[Path, tuple[int, int, int]] = {} + latest: tuple[datetime, dict[str, Any]] | None = None + for file_path in paths: + try: + self._validate_file_path(str(file_path.relative_to(root))[:-5]) + info = file_path.stat() + if not stat.S_ISREG(info.st_mode): + raise ValueError("checkpoint path is not a regular file") + version = _version(info) + hook, duplicate_keys = _make_duplicate_capturing_hook() + with open(file_path, encoding="utf-8") as f: + stored: dict[str, Any] = json.load(f, object_pairs_hook=hook) + versions[file_path] = version + owner: str | None = None + if ( + isinstance(stored, dict) + and "workflow_name" not in duplicate_keys + and isinstance(stored.get("workflow_name"), str) + ): + owner = stored["workflow_name"] + if owner is None: + logger.warning( + "Quarantining unidentifiable checkpoint file %s while finding " + "the latest checkpoint for workflow '%s'", + file_path, + workflow_name, + ) + continue + if owner != workflow_name: + continue + if "__pickled__" in stored or "__type__" in stored: + raise ValueError("Top-level pickle markers are not checkpoint metadata") + if duplicate_keys: + raise ValueError(f"Duplicate checkpoint JSON key: {duplicate_keys[0]}") + timestamp = datetime.fromisoformat(stored["timestamp"]) + checkpoint_id = stored["checkpoint_id"] + if not isinstance(checkpoint_id, str): + raise ValueError("checkpoint_id must be a string") + if self._validate_file_path(checkpoint_id) != file_path.resolve(): + raise ValueError("checkpoint_id does not match its checkpoint file") + if latest is None or timestamp > latest[0]: + latest = (timestamp, stored) + except (json.JSONDecodeError, UnicodeDecodeError) as ex: + logger.warning( + "Quarantining checkpoint file %s that is not valid JSON while " + "finding the latest checkpoint for workflow '%s': %s", + file_path, + workflow_name, + ex, + ) + continue + except Exception as e: + raise WorkflowCheckpointException( + f"Failed to read checkpoint metadata for workflow {workflow_name} from {file_path}: {e}" + ) from e + for file_path, version in versions.items(): + try: + current = _version(file_path.stat()) + except FileNotFoundError as ex: + raise _CheckpointScanInvalidated( + f"Checkpoint {file_path} was removed during the scan; retry get_latest" + ) from ex + if current != version: + raise _CheckpointScanInvalidated( + f"Checkpoint {file_path} changed during the scan; retry get_latest" + ) + current_paths, current_directories = _checkpoint_paths(root) + if ( + set(current_paths) != set(paths) + or current_directories != directories + or any(_version(path.stat()) != version for path, version in directories.items()) + or self.storage_path.resolve() != root + ): + raise _CheckpointScanInvalidated( + f"Checkpoint directory {self.storage_path} changed during the scan; retry get_latest" + ) + return self._decode_checkpoint(latest[1]) if latest else None + + attempt = 0 + while True: + attempt += 1 + try: + checkpoint = await asyncio.to_thread(_read_latest) + except _CheckpointScanInvalidated as exc: + if attempt >= _MAX_GET_LATEST_ATTEMPTS: + raise WorkflowCheckpointException( + f"Checkpoint directory {self.storage_path} kept changing during " + f"get_latest for workflow '{workflow_name}' after " + f"{_MAX_GET_LATEST_ATTEMPTS} attempts; retry once concurrent " + "checkpoint writes settle" + ) from exc + logger.debug( + "Checkpoint directory changed while finding the latest checkpoint for " + "workflow '%s' (attempt %d of %d); retrying", + workflow_name, + attempt, + _MAX_GET_LATEST_ATTEMPTS, + ) + await asyncio.sleep(_GET_LATEST_RETRY_DELAY_SECONDS * attempt) + continue + except WorkflowCheckpointException: + raise + except Exception as e: + raise WorkflowCheckpointException( + f"Failed to get latest checkpoint for workflow {workflow_name} from {self.storage_path}: {e}" + ) from e + if checkpoint is None: + return None + logger.debug(f"Latest checkpoint for workflow {workflow_name} is {checkpoint.checkpoint_id}") + return checkpoint async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]: """List checkpoint IDs for a given workflow name. diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 6e13a414f8b..417c539cda5 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -24,6 +24,7 @@ WorkflowCheckpointException, WorkflowEvent, ) +from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value from agent_framework._workflows._runner_context import WorkflowMessage @@ -62,6 +63,37 @@ class _TestCustomData: tags: list[str] +@dataclass +class _UnlistedState: + """A state type outside the deserialization allow list.""" + + stage: str + + +def _stamped_checkpoint(state: dict[str, Any], timestamp: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint( + workflow_name="recovery-workflow", + graph_signature_hash="test-hash", + state=state, + timestamp=timestamp, + ) + + +def _write_checkpoint_file(storage_path: Path, checkpoint: WorkflowCheckpoint) -> None: + """Write a checkpoint JSON file directly, bypassing save()'s decode validation. + + Simulates the real-world scenario the bug is about: a file on disk that this + reader's allowed_checkpoint_types cannot decode (written by an older framework + version, a different storage config, or a type since unregistered). + """ + encoded = encode_checkpoint_value(checkpoint.to_dict()) + file_path = storage_path / f"{checkpoint.checkpoint_id}.json" + file_path.write_text( + json.dumps(encoded, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + + # region test WorkflowCheckpoint @@ -3449,3 +3481,431 @@ async def test_await_signal_through_cancellation_returns_immediately_for_a_resol checkpoint_module._await_signal_through_cancellation(resolved), # pyright: ignore[reportPrivateUsage] timeout=5, ) + + +async def test_file_checkpoint_storage_get_latest_does_not_silently_return_stale_state(): + """A newest checkpoint that cannot be decoded must raise, not fall back to an older one.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + _write_checkpoint_file( + Path(temp_dir), + _stamped_checkpoint({"app": _UnlistedState(stage="second")}, "2026-01-01T11:00:00+00:00"), + ) + with pytest.raises(WorkflowCheckpointException, match="allowed_checkpoint_types"): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_rejects_missing_directory(tmp_path: Path) -> None: + storage = FileCheckpointStorage(tmp_path / "checkpoints") + storage.storage_path.rmdir() + with pytest.raises(WorkflowCheckpointException) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + assert isinstance(exc.value.__cause__, FileNotFoundError) + + +async def test_file_checkpoint_storage_get_latest_returns_the_newest(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + newest_id = await storage.save(_stamped_checkpoint({"stage": "second"}, "2026-01-01T11:00:00+00:00")) + latest = await storage.get_latest(workflow_name="recovery-workflow") + assert latest is not None + assert latest.checkpoint_id == newest_id + + +async def test_file_checkpoint_storage_get_latest_without_checkpoints_is_none(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + assert await storage.get_latest(workflow_name="recovery-workflow") is None + + +@pytest.mark.parametrize( + ("older_timestamp", "newer_timestamp"), + [ + ("2026-01-01T12:00:00+03:00", "2026-01-01T10:00:00+00:00"), # tz-aware comparison + ("2026-01-01T09:00:00", "2026-01-01T10:00:00"), # naive comparison + ], +) +async def test_file_checkpoint_storage_get_latest_only_decodes_selected_checkpoint( + tmp_path: Path, older_timestamp: str, newer_timestamp: str +) -> None: + storage = FileCheckpointStorage(tmp_path) + _write_checkpoint_file( + tmp_path, + _stamped_checkpoint({"app": _UnlistedState(stage="first")}, older_timestamp), + ) + newest_id = await storage.save(_stamped_checkpoint({"stage": "second"}, newer_timestamp)) + latest = await storage.get_latest(workflow_name="recovery-workflow") + assert latest is not None + assert latest.checkpoint_id == newest_id + assert latest.state == {"stage": "second"} + + +@pytest.mark.parametrize( + ("field_name", "value", "remove"), + [ + ("timestamp", None, True), + ("timestamp", None, False), + ("timestamp", 42, False), + ("timestamp", "", False), + ("checkpoint_id", None, True), + ("checkpoint_id", None, False), + ("checkpoint_id", 42, False), + ("checkpoint_id", "../outside", False), + ], +) +async def test_file_checkpoint_storage_get_latest_rejects_invalid_metadata( + tmp_path: Path, field_name: str, value: Any, remove: bool +) -> None: + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + newest_id = await storage.save(_stamped_checkpoint({"stage": "second"}, "2026-01-01T11:00:00+00:00")) + file_path = tmp_path / f"{newest_id}.json" + data = json.loads(file_path.read_text()) + if remove: + del data[field_name] + else: + data[field_name] = value + file_path.write_text(json.dumps(data)) + with pytest.raises(WorkflowCheckpointException, match=newest_id) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + assert exc.value.__cause__ is not None + + +@pytest.mark.parametrize( + "contents", + ["{ invalid json }", "[]", "null", "{}", '{"workflow_name": null}', '{"workflow_name": 1}'], +) +@pytest.mark.parametrize("has_healthy_checkpoint", [False, True]) +async def test_file_checkpoint_storage_get_latest_quarantines_unidentifiable_files( + tmp_path: Path, contents: str, has_healthy_checkpoint: bool +) -> None: + """A file whose owner cannot be identified is skipped, not fatal. + + One malformed or legacy file must not poison get_latest() for every workflow + sharing the storage root, so it is quarantined (skipped with a warning) and + left in place for inspection. Healthy checkpoints for the requested workflow + are still returned. + """ + storage = FileCheckpointStorage(tmp_path) + healthy_id: str | None = None + if has_healthy_checkpoint: + healthy_id = await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + corrupt = tmp_path / "unidentifiable.json" + corrupt.write_text(contents) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + if has_healthy_checkpoint: + assert latest is not None + assert latest.checkpoint_id == healthy_id + else: + assert latest is None + assert corrupt.read_text() == contents + + +async def test_file_checkpoint_storage_get_latest_shared_root_is_not_poisoned_by_foreign_files( + tmp_path: Path, +) -> None: + """Regression: corruption owned by one workflow must not break another's recovery. + + Two workflows share one storage root. 'recovery-workflow' is healthy, while the + root also holds an unparseable file and a valid-but-foreign checkpoint with a + newer timestamp. get_latest() must still return the healthy newest checkpoint. + """ + storage = FileCheckpointStorage(tmp_path) + ours_id = await storage.save(_stamped_checkpoint({"stage": "ours"}, "2026-01-01T10:00:00+00:00")) + (tmp_path / "legacy-broken.json").write_text("{ not valid json") + foreign = _stamped_checkpoint({"stage": "theirs"}, "2026-01-01T11:00:00+00:00") + foreign_dict = foreign.to_dict() + foreign_dict["workflow_name"] = "other-workflow" + (tmp_path / f"{foreign.checkpoint_id}.json").write_text(json.dumps(encode_checkpoint_value(foreign_dict))) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == ours_id + assert latest.state == {"stage": "ours"} + + +@pytest.mark.parametrize("deny_directory", [False, True]) +async def test_file_checkpoint_storage_get_latest_surfaces_permission_errors( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, deny_directory: bool +) -> None: + import builtins + + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({}, "2026-01-01T11:00:00+00:00")) + + def deny(*args: Any, **kwargs: Any) -> Any: + raise PermissionError("checkpoint storage read denied") + + monkeypatch.setattr(os if deny_directory else builtins, "scandir" if deny_directory else "open", deny) + with pytest.raises(WorkflowCheckpointException) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + assert isinstance(exc.value.__cause__, PermissionError) + + +@pytest.mark.parametrize("change", ["create", "replace", "overwrite"]) +async def test_file_checkpoint_storage_get_latest_returns_fresh_snapshot_after_concurrent_change( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + """A concurrent change invalidates the scan once; the internal retry returns fresh state. + + The first attempt detects the mutation and retries instead of failing; the + second attempt observes the change and returns the new, consistent snapshot. + """ + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00") + checkpoint_id = await storage.save(checkpoint) + path = tmp_path / f"{checkpoint_id}.json" + original_load = json.load + mutated = False + + def mutate_once(stream: Any, **kwargs: Any) -> Any: + nonlocal mutated + data = original_load(stream, **kwargs) + # Release the read handle before deleting or replacing the file on Windows: + # an open file handle blocks unlink() and os.replace() with PermissionError, + # so the mutation below would not land while the scan holds the handle. + stream.close() + if not mutated: + mutated = True + if change == "create": + newer = _stamped_checkpoint({}, "2026-01-01T11:00:00+00:00") + encoded_newer = encode_checkpoint_value(newer.to_dict()) + (tmp_path / f"{newer.checkpoint_id}.json").write_text( + json.dumps(encoded_newer, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + else: + replacement = checkpoint.to_dict() + replacement["state"] = {"stage": "other-different-content"} + encoded_replacement = encode_checkpoint_value(replacement) + if change == "replace": + temporary = tmp_path / "replacement.tmp" + temporary.write_text( + json.dumps(encoded_replacement, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + os.replace(temporary, path) + else: + path.write_text( + json.dumps(encoded_replacement, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + return data + + monkeypatch.setattr(json, "load", mutate_once) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + if change == "create": + assert latest.state == {} + else: + assert latest.state == {"stage": "other-different-content"} + + +async def test_file_checkpoint_storage_get_latest_returns_none_when_file_removed_during_read( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """After a concurrent deletion the fresh snapshot honestly contains no checkpoints. + + The removal invalidates the first attempt; the retry rescans, finds no + checkpoint files, and returns None instead of stale data or an error. + """ + storage = FileCheckpointStorage(tmp_path) + checkpoint_id = await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + path = tmp_path / f"{checkpoint_id}.json" + original_load = json.load + mutated = False + + def mutate_once(stream: Any, **kwargs: Any) -> Any: + nonlocal mutated + data = original_load(stream, **kwargs) + # Release the read handle before unlinking on Windows. + stream.close() + if not mutated: + mutated = True + path.unlink() + return data + + monkeypatch.setattr(json, "load", mutate_once) + + assert await storage.get_latest(workflow_name="recovery-workflow") is None + + +async def test_file_checkpoint_storage_get_latest_raises_after_exhausting_retries( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """When the directory changes on every attempt, get_latest raises a clear error. + + Each attempt is invalidated by a new foreign checkpoint file appearing during + the scan. After the bounded number of internal retries is exhausted, the error + tells the caller the storage root is under heavy concurrent writes. + """ + storage = FileCheckpointStorage(tmp_path) + checkpoint_id = await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + counter = 0 + original_load = json.load + + def churn_on_every_read(stream: Any, **kwargs: Any) -> Any: + nonlocal counter + data = original_load(stream, **kwargs) + stream.close() + churn = _stamped_checkpoint({}, f"2026-01-01T12:0{counter}:00+00:00") + churn_dict = churn.to_dict() + churn_dict["workflow_name"] = "churn-workflow" + (tmp_path / f"{churn.checkpoint_id}.json").write_text(json.dumps(encode_checkpoint_value(churn_dict))) + counter += 1 + return data + + monkeypatch.setattr(json, "load", churn_on_every_read) + + with pytest.raises(WorkflowCheckpointException, match="kept changing"): + await storage.get_latest(workflow_name="recovery-workflow") + + assert (tmp_path / f"{checkpoint_id}.json").exists() + + +async def test_file_checkpoint_storage_get_latest_survives_concurrent_foreign_save( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An atomic save for a different workflow during the scan must not fail recovery. + + Regression for the shared-root invalidation comment: the optimistic scan is + retried internally, the second attempt skips the foreign checkpoint, and the + requested workflow's checkpoint is returned instead of a 'changed during the + scan' error. + """ + storage = FileCheckpointStorage(tmp_path) + checkpoint_id = await storage.save(_stamped_checkpoint({"stage": "ours"}, "2026-01-01T10:00:00+00:00")) + original_load = json.load + injected = False + + def inject_foreign_save(stream: Any, **kwargs: Any) -> Any: + nonlocal injected + data = original_load(stream, **kwargs) + stream.close() + if not injected: + injected = True + foreign = _stamped_checkpoint({"stage": "theirs"}, "2026-01-01T11:00:00+00:00") + foreign_dict = foreign.to_dict() + foreign_dict["workflow_name"] = "another-workflow" + (tmp_path / f"{foreign.checkpoint_id}.json").write_text( + json.dumps(encode_checkpoint_value(foreign_dict)), + encoding="utf-8", + ) + return data + + monkeypatch.setattr(json, "load", inject_foreign_save) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == checkpoint_id + assert latest.state == {"stage": "ours"} + + +async def test_file_checkpoint_storage_list_checkpoints_still_skips_undecodable_entries(): + """list_checkpoints keeps returning what it can read, so one bad file cannot break listing.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + readable_id = await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + _write_checkpoint_file( + Path(temp_dir), + _stamped_checkpoint({"app": _UnlistedState(stage="second")}, "2026-01-01T11:00:00+00:00"), + ) + listed = await storage.list_checkpoints(workflow_name="recovery-workflow") + assert [cp.checkpoint_id for cp in listed] == [readable_id] + + +async def test_file_checkpoint_storage_get_latest_honors_reader_allowed_types(tmp_path: Path) -> None: + checkpoint = _stamped_checkpoint({"app": _UnlistedState(stage="second")}, "2026-01-01T11:00:00+00:00") + _write_checkpoint_file(tmp_path, checkpoint) + reader = FileCheckpointStorage( + tmp_path, allowed_checkpoint_types=[f"{_UnlistedState.__module__}:{_UnlistedState.__qualname__}"] + ) + assert await reader.get_latest(workflow_name="recovery-workflow") == checkpoint + + +async def test_file_checkpoint_storage_get_latest_reads_selected_file_once( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00") + await storage.save(checkpoint) + original_load = json.load + reads = 0 + + def count_reads(stream: Any, **kwargs: Any) -> Any: + nonlocal reads + reads += 1 + return original_load(stream, **kwargs) + + monkeypatch.setattr(json, "load", count_reads) + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + assert reads == 1 + + +async def test_file_checkpoint_storage_get_latest_quarantines_duplicate_workflow_name_keys( + tmp_path: Path, +) -> None: + """A file with a duplicated workflow_name key has ambiguous ownership. + + It might belong to any workflow sharing the root, so it is quarantined (skipped + with a warning) instead of poisoning get_latest() for all of them. Ownership is + never trusted from a duplicated key. + """ + storage = FileCheckpointStorage(tmp_path) + saved_id = await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + (tmp_path / "ambiguous.json").write_text( + '{"workflow_name": "recovery-workflow", "workflow_name": "other-workflow"}' + ) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == saved_id + + +async def test_file_checkpoint_storage_get_latest_rejects_duplicate_metadata_keys_in_matching_workflow( + tmp_path: Path, +) -> None: + """A matching-workflow file with duplicate non-ownership keys must raise. + + Once ownership is established and matches the requested workflow, key + uniqueness is enforced across the whole file. + """ + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + (tmp_path / "duplicate-timestamp.json").write_text( + '{"workflow_name": "recovery-workflow", ' + '"timestamp": "2026-01-01T10:00:00+00:00", ' + '"timestamp": "2026-01-01T10:00:00+00:00"}' + ) + with pytest.raises(WorkflowCheckpointException, match="Duplicate checkpoint JSON key"): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_ignores_duplicate_keys_in_foreign_workflow( + tmp_path: Path, +) -> None: + """Duplicate keys in a file belonging to a different workflow must not block recovery. + + Per the documented policy, invalid metadata for a known different workflow is + ignored. A foreign file with duplicate keys must be skipped after ownership is + determined, not rejected during parsing. + """ + storage = FileCheckpointStorage(tmp_path) + newest_id = await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T11:00:00+00:00")) + (tmp_path / "foreign.json").write_text( + '{"workflow_name": "other-workflow", ' + '"timestamp": "2026-01-01T10:00:00+00:00", ' + '"timestamp": "2026-01-01T10:00:00+00:00"}' + ) + latest = await storage.get_latest(workflow_name="recovery-workflow") + assert latest is not None + assert latest.checkpoint_id == newest_id