Python: fix get_latest silently returning a stale checkpoint on decode failure

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.