Python: Resolve checkpoint get_latest timestamp ties through lineage

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.

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,

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