diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 3ddd0e0c73f..b4158621d7a 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py @@ -5,6 +5,7 @@ from __future__ import annotations import logging +from datetime import datetime from typing import Any, TypedDict from agent_framework._settings import SecretString, load_settings @@ -24,6 +25,28 @@ logger = logging.getLogger(__name__) +# Page size for get_latest's newest-first query. It is normally stopped after the first row +# or a small group of rows sharing the newest timestamp, so a small page avoids reading many +# full checkpoint documents that are then discarded. +_LATEST_CHECKPOINT_PAGE_SIZE = 10 + + +def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> WorkflowCheckpoint: + """Return the latest checkpoint, resolving timestamp ties through the lineage chain. + + Mirrors the tie-breaker used by the core checkpoint storages. Checkpoints taken at the + same superstep boundary can share a timestamp. Among the checkpoints with the newest + timestamp, one that another of them names as its ``previous_checkpoint_id`` is + superseded, so the result does not depend on the order the backend returns rows in. + If several unsuperseded candidates remain, the one with the greatest ``checkpoint_id`` + is returned, so the choice among them is stable too. + """ + newest_timestamp = max(datetime.fromisoformat(cp.timestamp) for cp in checkpoints) + candidates = [cp for cp in checkpoints if datetime.fromisoformat(cp.timestamp) == newest_timestamp] + superseded_ids = {cp.previous_checkpoint_id for cp in candidates} + leaves = [cp for cp in candidates if cp.checkpoint_id not in superseded_ids] + return max(leaves or candidates, key=lambda cp: cp.checkpoint_id) + class AzureCosmosCheckpointSettings(TypedDict, total=False): """Settings for CosmosCheckpointStorage resolved from args and environment.""" @@ -353,7 +376,14 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: """ await self._ensure_container_proxy() - query = "SELECT * FROM c WHERE c.workflow_name = @workflow_name ORDER BY c.timestamp DESC OFFSET 0 LIMIT 1" + # Checkpoints taken at the same superstep boundary can share a timestamp, so collect + # every checkpoint at the newest timestamp and resolve the tie through lineage + # instead of letting the backend pick one row. The newest timestamp and its rows come + # from one query ordered newest first rather than a MAX query followed by a second + # query, so a checkpoint saved or deleted in between cannot leave the result pinned to + # a stale timestamp, or empty while older checkpoints remain. Iteration stops at the + # first older row, so only the pages holding the tied group are read. + query = "SELECT * FROM c WHERE c.workflow_name = @workflow_name ORDER BY c.timestamp DESC" parameters: list[dict[str, object]] = [ {"name": "@workflow_name", "value": workflow_name}, ] @@ -362,18 +392,26 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: query=query, parameters=parameters, partition_key=workflow_name, + max_item_count=_LATEST_CHECKPOINT_PAGE_SIZE, ) + newest_documents: list[dict[str, Any]] = [] async for item in items: - checkpoint = self._document_to_checkpoint(item) - logger.debug( - "Latest checkpoint for workflow %s is %s", - workflow_name, - checkpoint.checkpoint_id, - ) - return checkpoint - - return None + if newest_documents and item.get("timestamp") != newest_documents[0].get("timestamp"): + break + newest_documents.append(item) + + if not newest_documents: + return None + + candidates = [self._document_to_checkpoint(document) for document in newest_documents] + checkpoint = _select_latest_checkpoint(candidates) + logger.debug( + "Latest checkpoint for workflow %s is %s", + workflow_name, + 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/azure-cosmos/tests/test_cosmos_checkpoint_storage.py b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py index e4fd305171c..87a20fe65ff 100644 --- a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py +++ b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py @@ -7,7 +7,7 @@ from collections.abc import AsyncIterator from contextlib import suppress from dataclasses import dataclass -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -81,6 +81,37 @@ def _checkpoint_to_cosmos_document(checkpoint: WorkflowCheckpoint) -> dict[str, return doc +def _mock_get_latest_queries(container: MagicMock, docs: list[dict[str, Any]]) -> None: + """Answer get_latest's newest-first query the way Cosmos orders it, keeping tied rows in the given order.""" + newest_first = sorted(docs, key=lambda doc: doc["timestamp"], reverse=True) + container.query_items.return_value = _to_async_iter(newest_first) + + +class _ConcurrentlyWrittenContainer: + """A Cosmos container whose partition another writer changes right after each query is issued. + + Each query reads a snapshot of the partition, newest first, as Cosmos serves one result page. + """ + + def __init__(self, documents: list[dict[str, Any]], concurrent_write: Any) -> None: + self.documents = list(documents) + self._concurrent_write = concurrent_write + self.queries: list[str] = [] + + def query_items(self, *, query: str, partition_key: str, **kwargs: Any) -> AsyncIterator[Any]: + assert "ORDER BY c.timestamp DESC" in query + self.queries.append(query) + snapshot = sorted( + (doc for doc in self.documents if doc["workflow_name"] == partition_key), + key=lambda doc: doc["timestamp"], + reverse=True, + ) + write, self._concurrent_write = self._concurrent_write, None + if write is not None: + write(self.documents) + return _to_async_iter(snapshot) + + @pytest.fixture def mock_container() -> MagicMock: container = MagicMock() @@ -351,9 +382,7 @@ async def test_delete_cosmos_not_found_returns_false(mock_container: MagicMock) async def test_get_latest_returns_latest_checkpoint(mock_container: MagicMock) -> None: cp = _make_checkpoint(checkpoint_id="cp-latest", timestamp="2025-06-01T00:00:00+00:00") - mock_container.query_items.return_value = _to_async_iter([ - _checkpoint_to_cosmos_document(cp), - ]) + _mock_get_latest_queries(mock_container, [_checkpoint_to_cosmos_document(cp)]) storage = CosmosCheckpointStorage(container_client=mock_container) result = await storage.get_latest(workflow_name="test-workflow") @@ -371,15 +400,134 @@ async def test_get_latest_returns_none_when_empty(mock_container: MagicMock) -> assert result is None -async def test_get_latest_uses_order_by_desc_with_limit(mock_container: MagicMock) -> None: - mock_container.query_items.return_value = _to_async_iter([]) +async def test_get_latest_reads_the_newest_rows_with_one_ordered_query(mock_container: MagicMock) -> None: + cp = _make_checkpoint(checkpoint_id="cp-latest", timestamp="2025-06-01T00:00:00+00:00") + _mock_get_latest_queries(mock_container, [_checkpoint_to_cosmos_document(cp)]) storage = CosmosCheckpointStorage(container_client=mock_container) await storage.get_latest(workflow_name="test-workflow") - kwargs = mock_container.query_items.call_args.kwargs - assert "ORDER BY c.timestamp DESC" in kwargs["query"] - assert "OFFSET 0 LIMIT 1" in kwargs["query"] + (query_call,) = mock_container.query_items.call_args_list + assert "ORDER BY c.timestamp DESC" in query_call.kwargs["query"] + assert "MAX(" not in query_call.kwargs["query"] + assert "LIMIT" not in query_call.kwargs["query"] + assert query_call.kwargs["parameters"] == [{"name": "@workflow_name", "value": "test-workflow"}] + assert query_call.kwargs["partition_key"] == "test-workflow" + assert query_call.kwargs["max_item_count"] > 0 + + +async def test_get_latest_stops_reading_at_the_first_older_checkpoint(mock_container: MagicMock) -> None: + timestamp = "2025-06-01T00:00:00+00:00" + tied = [ + _checkpoint_to_cosmos_document(_make_checkpoint(checkpoint_id="parent", timestamp=timestamp)), + _checkpoint_to_cosmos_document( + _make_checkpoint(checkpoint_id="child", previous_checkpoint_id="parent", timestamp=timestamp) + ), + ] + older = _checkpoint_to_cosmos_document( + _make_checkpoint(checkpoint_id="older", timestamp="2025-05-01T00:00:00+00:00") + ) + + async def _rows() -> AsyncIterator[Any]: + for doc in [*tied, older]: + yield doc + raise AssertionError("get_latest read past the first older checkpoint") + + mock_container.query_items.return_value = _rows() + + storage = CosmosCheckpointStorage(container_client=mock_container) + result = await storage.get_latest(workflow_name="test-workflow") + + assert result is not None + assert result.checkpoint_id == "child" + + +async def test_get_latest_is_not_pinned_to_a_stale_timestamp_by_a_concurrent_save() -> None: + older = _make_checkpoint(checkpoint_id="cp-1", timestamp="2025-05-01T00:00:00+00:00") + newest = _make_checkpoint( + checkpoint_id="cp-2", previous_checkpoint_id="cp-1", timestamp="2025-06-01T00:00:00+00:00" + ) + saved = _make_checkpoint(checkpoint_id="cp-3", previous_checkpoint_id="cp-2", timestamp="2025-07-01T00:00:00+00:00") + container = _ConcurrentlyWrittenContainer( + [_checkpoint_to_cosmos_document(older), _checkpoint_to_cosmos_document(newest)], + lambda documents: documents.append(_checkpoint_to_cosmos_document(saved)), + ) + storage = CosmosCheckpointStorage(container_client=cast(Any, container)) + + during_save = await storage.get_latest(workflow_name="test-workflow") + after_save = await storage.get_latest(workflow_name="test-workflow") + + # The first call answers from the one read it made, before the save landed. + assert during_save is not None + assert during_save.checkpoint_id == "cp-2" + assert len(container.queries) == 2 + assert after_save is not None + assert after_save.checkpoint_id == "cp-3" + + +async def test_get_latest_falls_back_to_older_checkpoints_when_the_newest_is_deleted_concurrently() -> None: + older = _make_checkpoint(checkpoint_id="cp-1", timestamp="2025-05-01T00:00:00+00:00") + newest = _make_checkpoint( + checkpoint_id="cp-2", previous_checkpoint_id="cp-1", timestamp="2025-06-01T00:00:00+00:00" + ) + container = _ConcurrentlyWrittenContainer( + [_checkpoint_to_cosmos_document(older), _checkpoint_to_cosmos_document(newest)], + lambda documents: documents.remove(next(doc for doc in documents if doc["checkpoint_id"] == "cp-2")), + ) + storage = CosmosCheckpointStorage(container_client=cast(Any, container)) + + during_delete = await storage.get_latest(workflow_name="test-workflow") + after_delete = await storage.get_latest(workflow_name="test-workflow") + + # The first call returns the checkpoint its one read saw, never None while checkpoints remain. + assert during_delete is not None + assert during_delete.checkpoint_id == "cp-2" + assert len(container.queries) == 2 + assert after_delete is not None + assert after_delete.checkpoint_id == "cp-1" + + +@pytest.mark.parametrize("row_order", [("parent", "child", "grandchild"), ("grandchild", "child", "parent")]) +async def test_get_latest_resolves_timestamp_ties_through_lineage( + mock_container: MagicMock, row_order: tuple[str, ...] +) -> None: + timestamp = "2025-06-01T00:00:00+00:00" + checkpoints = { + "parent": _make_checkpoint(checkpoint_id="parent", timestamp=timestamp), + "child": _make_checkpoint(checkpoint_id="child", previous_checkpoint_id="parent", timestamp=timestamp), + "grandchild": _make_checkpoint(checkpoint_id="grandchild", previous_checkpoint_id="child", timestamp=timestamp), + } + _mock_get_latest_queries( + mock_container, [_checkpoint_to_cosmos_document(checkpoints[checkpoint_id]) for checkpoint_id in row_order] + ) + + storage = CosmosCheckpointStorage(container_client=mock_container) + result = await storage.get_latest(workflow_name="test-workflow") + + assert result is not None + assert result.checkpoint_id == "grandchild" + + +@pytest.mark.parametrize("row_order", [("leaf-a", "leaf-b"), ("leaf-b", "leaf-a")]) +async def test_get_latest_breaks_ties_between_unsuperseded_checkpoints_by_id( + mock_container: MagicMock, row_order: tuple[str, ...] +) -> None: + timestamp = "2025-06-01T00:00:00+00:00" + _mock_get_latest_queries( + mock_container, + [ + _checkpoint_to_cosmos_document( + _make_checkpoint(checkpoint_id=checkpoint_id, previous_checkpoint_id="parent", timestamp=timestamp) + ) + for checkpoint_id in row_order + ], + ) + + storage = CosmosCheckpointStorage(container_client=mock_container) + result = await storage.get_latest(workflow_name="test-workflow") + + assert result is not None + assert result.checkpoint_id == "leaf-b" # --- Tests for list_checkpoint_ids --- @@ -740,7 +888,7 @@ async def test_get_latest_blocks_unlisted_app_type(mock_container: MagicMock) -> """get_latest raises when the checkpoint contains an unlisted application type.""" checkpoint = _make_checkpoint_with_state({"data": _AppState(label="x", count=1)}) doc = _checkpoint_to_cosmos_document(checkpoint) - mock_container.query_items.return_value = _to_async_iter([doc]) + _mock_get_latest_queries(mock_container, [doc]) storage = CosmosCheckpointStorage(container_client=mock_container) @@ -752,7 +900,7 @@ async def test_get_latest_allows_listed_app_type(mock_container: MagicMock) -> N """get_latest decodes checkpoints with listed application types.""" checkpoint = _make_checkpoint_with_state({"data": _AppState(label="latest", count=9)}) doc = _checkpoint_to_cosmos_document(checkpoint) - mock_container.query_items.return_value = _to_async_iter([doc]) + _mock_get_latest_queries(mock_container, [doc]) storage = CosmosCheckpointStorage( container_client=mock_container, diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index cd78725d6a1..cbca7ef52e9 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -205,6 +205,22 @@ async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID] ... +def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> WorkflowCheckpoint: + """Return the latest checkpoint, resolving timestamp ties through the lineage chain. + + Checkpoints taken at the same superstep boundary can share a timestamp. Among the + checkpoints with the newest timestamp, one that another of them names as its + ``previous_checkpoint_id`` is superseded, so the result does not depend on save order. + If several unsuperseded candidates remain, the one with the greatest ``checkpoint_id`` + is returned, so the choice among them is stable too. + """ + newest_timestamp = max(datetime.fromisoformat(cp.timestamp) for cp in checkpoints) + candidates = [cp for cp in checkpoints if datetime.fromisoformat(cp.timestamp) == newest_timestamp] + superseded_ids = {cp.previous_checkpoint_id for cp in candidates} + leaves = [cp for cp in candidates if cp.checkpoint_id not in superseded_ids] + return max(leaves or candidates, key=lambda cp: cp.checkpoint_id) + + class InMemoryCheckpointStorage: """In-memory checkpoint storage for testing and development.""" @@ -243,7 +259,7 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: checkpoints = [cp for cp in self._checkpoints.values() if cp.workflow_name == workflow_name] if not checkpoints: return None - latest_checkpoint = max(checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp)) + latest_checkpoint = _select_latest_checkpoint(checkpoints) logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}") return copy.deepcopy(latest_checkpoint) @@ -854,7 +870,7 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: 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)) + latest_checkpoint = _select_latest_checkpoint(checkpoints) logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}") return latest_checkpoint diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 6e13a414f8b..e1c6a77f6ea 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -290,6 +290,65 @@ async def test_memory_checkpoint_storage_get_latest(): assert latest_none is None +@pytest.mark.parametrize("storage_kind", ["memory", "file"]) +@pytest.mark.parametrize("save_order", [("parent", "child", "grandchild"), ("grandchild", "child", "parent")]) +async def test_get_latest_resolves_timestamp_ties_through_lineage( + tmp_path: Path, storage_kind: str, save_order: tuple[str, ...] +) -> None: + storage = InMemoryCheckpointStorage() if storage_kind == "memory" else FileCheckpointStorage(tmp_path) + timestamp = "2026-10-11T12:00:00+00:00" + checkpoints = { + "parent": WorkflowCheckpoint( + workflow_name="w", graph_signature_hash="h", checkpoint_id="parent", timestamp=timestamp + ), + "child": WorkflowCheckpoint( + workflow_name="w", + graph_signature_hash="h", + checkpoint_id="child", + previous_checkpoint_id="parent", + timestamp=timestamp, + ), + "grandchild": WorkflowCheckpoint( + workflow_name="w", + graph_signature_hash="h", + checkpoint_id="grandchild", + previous_checkpoint_id="child", + timestamp=timestamp, + ), + } + for checkpoint_id in save_order: + await storage.save(checkpoints[checkpoint_id]) + + latest = await storage.get_latest(workflow_name="w") + + assert latest is not None + assert latest.checkpoint_id == "grandchild" + + +@pytest.mark.parametrize("storage_kind", ["memory", "file"]) +@pytest.mark.parametrize("save_order", [("leaf-a", "leaf-b"), ("leaf-b", "leaf-a")]) +async def test_get_latest_breaks_ties_between_unsuperseded_checkpoints_by_id( + tmp_path: Path, storage_kind: str, save_order: tuple[str, ...] +) -> None: + storage = InMemoryCheckpointStorage() if storage_kind == "memory" else FileCheckpointStorage(tmp_path) + timestamp = "2026-10-11T12:00:00+00:00" + for checkpoint_id in save_order: + await storage.save( + WorkflowCheckpoint( + workflow_name="w", + graph_signature_hash="h", + checkpoint_id=checkpoint_id, + previous_checkpoint_id="parent", + timestamp=timestamp, + ) + ) + + latest = await storage.get_latest(workflow_name="w") + + assert latest is not None + assert latest.checkpoint_id == "leaf-b" + + async def test_workflow_checkpoint_chaining_via_previous_checkpoint_id(): """Test that consecutive checkpoints created by a workflow are properly chained via previous_checkpoint_id.""" from typing_extensions import Never diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py index 487d6882cd0..539287adc74 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py @@ -34,6 +34,23 @@ def _encoded_checkpoint_hash(value: Any) -> str: return hashlib.sha256(payload).hexdigest() +def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> WorkflowCheckpoint: + """Return the latest checkpoint, resolving timestamp ties through the lineage chain. + + Mirrors the tie-breaker used by the core checkpoint storages. Checkpoints taken at the + same superstep boundary can share a timestamp. Among the checkpoints with the newest + timestamp, one that another of them names as its ``previous_checkpoint_id`` is + superseded, so the result does not depend on the order the store lists keys in. + If several unsuperseded candidates remain, the one with the greatest ``checkpoint_id`` + is returned, so the choice among them is stable too. + """ + newest_timestamp = max(datetime.fromisoformat(cp.timestamp) for cp in checkpoints) + candidates = [cp for cp in checkpoints if datetime.fromisoformat(cp.timestamp) == newest_timestamp] + superseded_ids = {cp.previous_checkpoint_id for cp in candidates} + leaves = [cp for cp in candidates if cp.checkpoint_id not in superseded_ids] + return max(leaves or candidates, key=lambda cp: cp.checkpoint_id) + + def _store_scope(config: AgentConfig, platform_context: FoundryAgentRequestContext) -> FoundryRequestScope | None: if not config.is_hosted: return None @@ -227,7 +244,7 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: checkpoints = await self.list_checkpoints(workflow_name=workflow_name) if not checkpoints: return None - return max(checkpoints, key=lambda checkpoint: datetime.fromisoformat(checkpoint.timestamp)) + return _select_latest_checkpoint(checkpoints) async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]: checkpoints = await self.list_checkpoints(workflow_name=workflow_name) diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 6ed3b84fdce..5c622ca1e75 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -30,12 +30,17 @@ class _NotAllowed: def _checkpoint( - checkpoint_id: str, *, workflow_name: str = "workflow", timestamp: str = "2026-01-01T00:00:00+00:00" + checkpoint_id: str, + *, + workflow_name: str = "workflow", + timestamp: str = "2026-01-01T00:00:00+00:00", + previous_checkpoint_id: str | None = None, ) -> WorkflowCheckpoint: return WorkflowCheckpoint( workflow_name=workflow_name, graph_signature_hash="graph-hash", checkpoint_id=checkpoint_id, + previous_checkpoint_id=previous_checkpoint_id, timestamp=timestamp, ) @@ -301,6 +306,49 @@ async def test_get_latest_uses_timestamp_and_list_ids_filters(monkeypatch: pytes assert await storage.list_checkpoint_ids(workflow_name="workflow") == ["newer", "older"] +@pytest.mark.parametrize("list_order", [("parent", "child", "grandchild"), ("grandchild", "child", "parent")]) +async def test_get_latest_resolves_timestamp_ties_through_lineage( + monkeypatch: pytest.MonkeyPatch, list_order: tuple[str, ...] +) -> None: + storage = FoundryCheckpointStore("context-1", _platform_context()) + timestamp = "2026-01-02T00:00:00+00:00" + checkpoints = { + "older": _checkpoint("older", timestamp="2026-01-01T00:00:00+00:00"), + "parent": _checkpoint("parent", timestamp=timestamp, previous_checkpoint_id="older"), + "child": _checkpoint("child", timestamp=timestamp, previous_checkpoint_id="parent"), + "grandchild": _checkpoint("grandchild", timestamp=timestamp, previous_checkpoint_id="child"), + } + listed = [checkpoints["older"], *(checkpoints[checkpoint_id] for checkpoint_id in list_order)] + monkeypatch.setattr(storage, "list_checkpoints", AsyncMock(return_value=listed)) + + latest = await storage.get_latest(workflow_name="workflow") + + assert latest is not None + assert latest.checkpoint_id == "grandchild" + + +@pytest.mark.parametrize("list_order", [("leaf-a", "leaf-b"), ("leaf-b", "leaf-a")]) +async def test_get_latest_breaks_ties_between_unsuperseded_checkpoints_by_id( + monkeypatch: pytest.MonkeyPatch, list_order: tuple[str, ...] +) -> None: + storage = FoundryCheckpointStore("context-1", _platform_context()) + timestamp = "2026-01-02T00:00:00+00:00" + checkpoints = { + "leaf-a": _checkpoint("leaf-a", timestamp=timestamp, previous_checkpoint_id="older"), + "leaf-b": _checkpoint("leaf-b", timestamp=timestamp, previous_checkpoint_id="older"), + } + listed = [ + _checkpoint("older", timestamp="2026-01-01T00:00:00+00:00"), + *(checkpoints[checkpoint_id] for checkpoint_id in list_order), + ] + monkeypatch.setattr(storage, "list_checkpoints", AsyncMock(return_value=listed)) + + latest = await storage.get_latest(workflow_name="workflow") + + assert latest is not None + assert latest.checkpoint_id == "leaf-b" + + async def test_get_latest_returns_none_when_no_checkpoints_exist(monkeypatch: pytest.MonkeyPatch) -> None: storage = FoundryCheckpointStore("context-1", _platform_context()) monkeypatch.setattr(storage, "list_checkpoints", AsyncMock(return_value=[]))