From be0266829db26d3097c7ad428a86ad3361628706 Mon Sep 17 00:00:00 2001
From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com>
Date: Fri, 18 Sep 2026 17:29:43 +0530
Subject: [PATCH] Python: break get_latest timestamp ties with the checkpoint
 lineage chain

Identical timestamps at a superstep boundary made max() pick whichever
checkpoint the storage iterated first, so get_latest could return a
stale checkpoint depending on save order. Both storages now first drop
every checkpoint that another one supersedes via previous_checkpoint_id
and only then compare timestamps.

Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com>
---
 .../agent_framework/_workflows/_checkpoint.py | 14 +++++++++--
 .../core/tests/workflow/test_checkpoint.py    | 25 +++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py
index a0eaed0797b..e3b5a4d5e7b 100644
--- a/python/packages/core/agent_framework/_workflows/_checkpoint.py
+++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py
@@ -243,7 +243,13 @@ 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))
+        # Timestamps tie at the same superstep boundary (see iteration_count
+        # note above); break ties by the previous_checkpoint_id lineage chain
+        # so resume picks the checkpoint no other one supersedes, not whichever
+        # the dict happens to iterate first.
+        superseded = {cp.previous_checkpoint_id for cp in checkpoints if cp.previous_checkpoint_id}
+        candidates = [cp for cp in checkpoints if cp.checkpoint_id not in superseded] or checkpoints
+        latest_checkpoint = max(candidates, key=lambda cp: datetime.fromisoformat(cp.timestamp))
         logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}")
         return copy.deepcopy(latest_checkpoint)
 
@@ -854,7 +860,11 @@ 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))
+        # Same lineage tiebreak as InMemoryCheckpointStorage: identical
+        # timestamps at a superstep boundary must not pick an arbitrary one.
+        superseded = {cp.previous_checkpoint_id for cp in checkpoints if cp.previous_checkpoint_id}
+        candidates = [cp for cp in checkpoints if cp.checkpoint_id not in superseded] or checkpoints
+        latest_checkpoint = max(candidates, key=lambda cp: datetime.fromisoformat(cp.timestamp))
         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..f38d7b200e0 100644
--- a/python/packages/core/tests/workflow/test_checkpoint.py
+++ b/python/packages/core/tests/workflow/test_checkpoint.py
@@ -3449,3 +3449,28 @@ 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_memory_get_latest_ties_broken_by_lineage_not_save_order():
+    """Identical timestamps must resolve via the lineage chain, not dict order."""
+    from datetime import datetime, timezone
+
+    ts = datetime(2026, 9, 18, 10, 0, 0, tzinfo=timezone.utc).isoformat()
+    parent = WorkflowCheckpoint(
+        workflow_name="w", graph_signature_hash="h", checkpoint_id="parent", timestamp=ts, iteration_count=1
+    )
+    child = WorkflowCheckpoint(
+        workflow_name="w",
+        graph_signature_hash="h",
+        checkpoint_id="child",
+        timestamp=ts,
+        iteration_count=1,
+        previous_checkpoint_id="parent",
+    )
+
+    for save_order in ([parent, child], [child, parent]):
+        storage = InMemoryCheckpointStorage()
+        for cp in save_order:
+            await storage.save(cp)
+        latest = await storage.get_latest(workflow_name="w")
+        assert latest.checkpoint_id == "child", (
+            f"save order {[cp.checkpoint_id for cp in save_order]} picked {latest.checkpoint_id}"
+        )