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
@@ -25,6 +26,23 @@
 logger = logging.getLogger(__name__)
 
 
+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,9 +371,33 @@ 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 fetch
+        # every checkpoint at the newest timestamp and resolve the tie through lineage
+        # instead of letting the backend pick one row.
+        max_query = "SELECT VALUE MAX(c.timestamp) FROM c WHERE c.workflow_name = @workflow_name"
+        max_parameters: list[dict[str, object]] = [
+            {"name": "@workflow_name", "value": workflow_name},
+        ]
+
+        max_items = self._container_proxy.query_items(  # type: ignore[union-attr]
+            query=max_query,
+            parameters=max_parameters,
+            partition_key=workflow_name,
+        )
+
+        newest_timestamp: str | None = None
+        async for value in max_items:
+            if isinstance(value, str):
+                newest_timestamp = value
+            break
+
+        if newest_timestamp is None:
+            return None
+
+        query = "SELECT * FROM c WHERE c.workflow_name = @workflow_name AND c.timestamp = @timestamp"
         parameters: list[dict[str, object]] = [
             {"name": "@workflow_name", "value": workflow_name},
+            {"name": "@timestamp", "value": newest_timestamp},
         ]
 
         items = self._container_proxy.query_items(  # type: ignore[union-attr]
@@ -364,16 +406,17 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None:
             partition_key=workflow_name,
         )
 
-        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
+        candidates = [self._document_to_checkpoint(item) async for item in items]
+        if not candidates:
+            return None
 
-        return None
+        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

@@ -81,6 +81,12 @@ def _checkpoint_to_cosmos_document(checkpoint: WorkflowCheckpoint) -> dict[str,
     return doc
 
 
+def _mock_get_latest_queries(container: MagicMock, newest_docs: list[dict[str, Any]]) -> None:
+    """Answer get_latest's newest-timestamp query, then its query for the rows at that timestamp."""
+    newest_timestamps = [max(doc["timestamp"] for doc in newest_docs)] if newest_docs else []
+    container.query_items.side_effect = [_to_async_iter(newest_timestamps), _to_async_iter(newest_docs)]
+
+
 @pytest.fixture
 def mock_container() -> MagicMock:
     container = MagicMock()
@@ -351,9 +357,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 +375,78 @@ 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_fetches_every_checkpoint_at_the_newest_timestamp(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"]
+    max_call, candidates_call = mock_container.query_items.call_args_list
+    assert "MAX(c.timestamp)" in max_call.kwargs["query"]
+    assert max_call.kwargs["parameters"] == [{"name": "@workflow_name", "value": "test-workflow"}]
+    assert max_call.kwargs["partition_key"] == "test-workflow"
+    assert "c.timestamp = @timestamp" in candidates_call.kwargs["query"]
+    assert "LIMIT" not in candidates_call.kwargs["query"]
+    assert candidates_call.kwargs["parameters"] == [
+        {"name": "@workflow_name", "value": "test-workflow"},
+        {"name": "@timestamp", "value": "2025-06-01T00:00:00+00:00"},
+    ]
+    assert candidates_call.kwargs["partition_key"] == "test-workflow"
+
+
+async def test_get_latest_returns_none_when_newest_rows_disappear(mock_container: MagicMock) -> None:
+    mock_container.query_items.side_effect = [
+        _to_async_iter(["2025-06-01T00:00:00+00:00"]),
+        _to_async_iter([]),
+    ]
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+
+    assert await storage.get_latest(workflow_name="test-workflow") is None
+
+
+@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 +807,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 +819,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
 

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

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)