Python: Normalize naive timestamps to UTC in checkpoint selection

python/packages/core/agent_framework/_workflows/_checkpoint.py

@@ -205,6 +205,14 @@ async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]
         ...
 
 
+def _parse_checkpoint_timestamp(timestamp: str) -> datetime:
+    """Parse an ISO 8601 timestamp string into a timezone-aware UTC datetime for comparison."""
+    dt = datetime.fromisoformat(timestamp)
+    if dt.tzinfo is None:
+        return dt.replace(tzinfo=timezone.utc)
+    return dt
+
+
 def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> WorkflowCheckpoint:
     """Return the latest checkpoint, resolving timestamp ties through the lineage chain.
 
@@ -214,8 +222,8 @@ def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> Workflow
     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]
+    newest_timestamp = max(_parse_checkpoint_timestamp(cp.timestamp) for cp in checkpoints)
+    candidates = [cp for cp in checkpoints if _parse_checkpoint_timestamp(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)

python/packages/core/tests/workflow/test_checkpoint.py

@@ -349,6 +349,46 @@ async def test_get_latest_breaks_ties_between_unsuperseded_checkpoints_by_id(
     assert latest.checkpoint_id == "leaf-b"
 
 
+@pytest.mark.parametrize("storage_kind", ["memory", "file"])
+async def test_get_latest_with_mixed_naive_and_aware_timestamps(tmp_path: Path, storage_kind: str) -> None:
+    storage = InMemoryCheckpointStorage() if storage_kind == "memory" else FileCheckpointStorage(tmp_path)
+    # Checkpoints with naive and timezone-aware timestamps
+    await storage.save(
+        WorkflowCheckpoint(
+            workflow_name="w",
+            graph_signature_hash="h",
+            checkpoint_id="older-naive",
+            timestamp="2026-01-01T00:00:00",
+        )
+    )
+    await storage.save(
+        WorkflowCheckpoint(
+            workflow_name="w",
+            graph_signature_hash="h",
+            checkpoint_id="newer-aware",
+            timestamp="2026-01-02T00:00:00+00:00",
+        )
+    )
+
+    latest = await storage.get_latest(workflow_name="w")
+    assert latest is not None
+    assert latest.checkpoint_id == "newer-aware"
+
+    # Save a newer naive checkpoint to verify naive-as-UTC selection
+    await storage.save(
+        WorkflowCheckpoint(
+            workflow_name="w",
+            graph_signature_hash="h",
+            checkpoint_id="newest-naive",
+            timestamp="2026-01-03T00:00:00",
+        )
+    )
+
+    latest = await storage.get_latest(workflow_name="w")
+    assert latest is not None
+    assert latest.checkpoint_id == "newest-naive"
+
+
 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