From 8cbe2a6de5273bce472ec3dbf9f90b20fc6d6e94 Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Mon, 2 Nov 2026 00:49:03 +0530 Subject: [PATCH 1/2] fix(workflows): support naive and timezone-aware timestamps in _select_latest_checkpoint Normalize ISO 8601 timestamps without timezone offsets to UTC when selecting the latest checkpoint, preventing TypeError when comparing offset-naive and offset-aware datetimes across saved checkpoints. Signed-off-by: Hardik Kaurani --- .../core/agent_framework/_workflows/_checkpoint.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) 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) From f935ace732cbb17602313fcb4cb7f280d3d0dafa Mon Sep 17 00:00:00 2001 From: Hardik Kaurani Date: Mon, 2 Nov 2026 00:49:37 +0530 Subject: [PATCH 2/2] test(workflows): add regression tests for get_latest with naive and aware timestamps Add parameterized regression tests for InMemoryCheckpointStorage and FileCheckpointStorage verifying that get_latest correctly selects the newest checkpoint when storages contain a mixture of naive and timezone-aware timestamps without raising TypeError. Signed-off-by: Hardik Kaurani --- .../core/tests/workflow/test_checkpoint.py | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) 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