diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 245950e0a93..ea8786d3edb 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/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) diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index e1c6a77f6ea..48373f118be 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/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