From 740fca80680c77f514f8db0b09c79e696b1ec9f9 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Sun, 11 Oct 2026 20:59:58 +0530 Subject: [PATCH 1/5] Python: Resolve checkpoint get_latest timestamp ties through lineage InMemoryCheckpointStorage.get_latest and FileCheckpointStorage.get_latest picked max() by timestamp alone. Checkpoints taken at the same superstep boundary can share a timestamp, so the result depended on save order and a resume could restore a superseded checkpoint. Among the checkpoints with the newest timestamp, skip any that another one names as previous_checkpoint_id. Fixes #8513 Co-Authored-By: Claude Opus 5.5 --- .../agent_framework/_workflows/_checkpoint.py | 17 +++++++-- .../core/tests/workflow/test_checkpoint.py | 35 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index cd78725d6a1..3ada216377d 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -205,6 +205,19 @@ 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. + """ + 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} + return next((cp for cp in candidates if cp.checkpoint_id not in superseded_ids), candidates[0]) + + class InMemoryCheckpointStorage: """In-memory checkpoint storage for testing and development.""" @@ -243,7 +256,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 +867,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 diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 6e13a414f8b..d278494865d 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -290,6 +290,41 @@ 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" + + 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 From 83ecc6912f9df9231451ca97beedf875a7686ba3 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:51:27 +0530 Subject: [PATCH 2/5] Python: Apply checkpoint lineage tie-breaker to Cosmos and Foundry stores CosmosCheckpointStorage.get_latest asked Cosmos for a single row ordered only by timestamp, and FoundryCheckpointStore.get_latest used max() over timestamps, so a parent/child pair sharing the newest timestamp could still resume from the superseded parent depending on backend/list order. Cosmos now queries MAX(c.timestamp) for the workflow partition, fetches every checkpoint at that timestamp, and resolves the tie through previous_checkpoint_id. Foundry applies the same tie-breaker to its listed checkpoints. The helper mirrors core's _select_latest_checkpoint because both packages support agent-framework-core releases that predate it. Co-Authored-By: Claude Opus 5.5 --- .../_checkpoint_storage.py | 60 ++++++++++++++--- .../tests/test_cosmos_checkpoint_storage.py | 65 ++++++++++++++++--- .../_state_store.py | 16 ++++- .../foundry_hosting/tests/test_state_store.py | 28 +++++++- 4 files changed, 147 insertions(+), 22 deletions(-) diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 3ddd0e0c73f..64c2526e2a5 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/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,20 @@ 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. + """ + 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} + return next((cp for cp in candidates if cp.checkpoint_id not in superseded_ids), candidates[0]) + + class AzureCosmosCheckpointSettings(TypedDict, total=False): """Settings for CosmosCheckpointStorage resolved from args and environment.""" @@ -353,9 +368,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 +403,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. diff --git a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py index e4fd305171c..f81767b9f82 100644 --- a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py +++ b/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,56 @@ 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" # --- Tests for list_checkpoint_ids --- @@ -740,7 +785,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 +797,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, diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py index 487d6882cd0..2588fabc790 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py @@ -34,6 +34,20 @@ 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. + """ + 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} + return next((cp for cp in candidates if cp.checkpoint_id not in superseded_ids), candidates[0]) + + def _store_scope(config: AgentConfig, platform_context: FoundryAgentRequestContext) -> FoundryRequestScope | None: if not config.is_hosted: return None @@ -227,7 +241,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) diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 6ed3b84fdce..1f5a7cdbd1c 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -30,12 +30,17 @@ class _NotAllowed: def _checkpoint( - checkpoint_id: str, *, workflow_name: str = "workflow", timestamp: str = "2026-01-01T00:00:00+00:00" + checkpoint_id: str, + *, + workflow_name: str = "workflow", + timestamp: str = "2026-01-01T00:00:00+00:00", + previous_checkpoint_id: str | None = None, ) -> WorkflowCheckpoint: return WorkflowCheckpoint( workflow_name=workflow_name, graph_signature_hash="graph-hash", checkpoint_id=checkpoint_id, + previous_checkpoint_id=previous_checkpoint_id, timestamp=timestamp, ) @@ -301,6 +306,27 @@ async def test_get_latest_uses_timestamp_and_list_ids_filters(monkeypatch: pytes assert await storage.list_checkpoint_ids(workflow_name="workflow") == ["newer", "older"] +@pytest.mark.parametrize("list_order", [("parent", "child", "grandchild"), ("grandchild", "child", "parent")]) +async def test_get_latest_resolves_timestamp_ties_through_lineage( + monkeypatch: pytest.MonkeyPatch, list_order: tuple[str, ...] +) -> None: + storage = FoundryCheckpointStore("context-1", _platform_context()) + timestamp = "2026-01-02T00:00:00+00:00" + checkpoints = { + "older": _checkpoint("older", timestamp="2026-01-01T00:00:00+00:00"), + "parent": _checkpoint("parent", timestamp=timestamp, previous_checkpoint_id="older"), + "child": _checkpoint("child", timestamp=timestamp, previous_checkpoint_id="parent"), + "grandchild": _checkpoint("grandchild", timestamp=timestamp, previous_checkpoint_id="child"), + } + listed = [checkpoints["older"], *(checkpoints[checkpoint_id] for checkpoint_id in list_order)] + monkeypatch.setattr(storage, "list_checkpoints", AsyncMock(return_value=listed)) + + latest = await storage.get_latest(workflow_name="workflow") + + assert latest is not None + assert latest.checkpoint_id == "grandchild" + + async def test_get_latest_returns_none_when_no_checkpoints_exist(monkeypatch: pytest.MonkeyPatch) -> None: storage = FoundryCheckpointStore("context-1", _platform_context()) monkeypatch.setattr(storage, "list_checkpoints", AsyncMock(return_value=[])) From ff52e601ef193defe3c6c37f793312e7736bda31 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:25:07 +0530 Subject: [PATCH 3/5] Python: Break same-timestamp ties between unsuperseded checkpoints by checkpoint_id Co-Authored-By: Claude Opus 5.5 --- .../_checkpoint_storage.py | 5 +++- .../tests/test_cosmos_checkpoint_storage.py | 22 +++++++++++++++++ .../agent_framework/_workflows/_checkpoint.py | 5 +++- .../core/tests/workflow/test_checkpoint.py | 24 +++++++++++++++++++ .../_state_store.py | 5 +++- .../foundry_hosting/tests/test_state_store.py | 22 +++++++++++++++++ 6 files changed, 80 insertions(+), 3 deletions(-) diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 64c2526e2a5..08ba98690fc 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py @@ -33,11 +33,14 @@ def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> Workflow 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} - return next((cp for cp in candidates if cp.checkpoint_id not in superseded_ids), candidates[0]) + 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): diff --git a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py index f81767b9f82..94dc945dee0 100644 --- a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py +++ b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py @@ -427,6 +427,28 @@ async def test_get_latest_resolves_timestamp_ties_through_lineage( 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 --- diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 3ada216377d..cbca7ef52e9 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -211,11 +211,14 @@ def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> Workflow 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} - return next((cp for cp in candidates if cp.checkpoint_id not in superseded_ids), candidates[0]) + 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: diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index d278494865d..e1c6a77f6ea 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -325,6 +325,30 @@ async def test_get_latest_resolves_timestamp_ties_through_lineage( 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 diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py index 2588fabc790..539287adc74 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py @@ -41,11 +41,14 @@ def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> Workflow 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} - return next((cp for cp in candidates if cp.checkpoint_id not in superseded_ids), candidates[0]) + 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: diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 1f5a7cdbd1c..5c622ca1e75 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -327,6 +327,28 @@ async def test_get_latest_resolves_timestamp_ties_through_lineage( assert latest.checkpoint_id == "grandchild" +@pytest.mark.parametrize("list_order", [("leaf-a", "leaf-b"), ("leaf-b", "leaf-a")]) +async def test_get_latest_breaks_ties_between_unsuperseded_checkpoints_by_id( + monkeypatch: pytest.MonkeyPatch, list_order: tuple[str, ...] +) -> None: + storage = FoundryCheckpointStore("context-1", _platform_context()) + timestamp = "2026-01-02T00:00:00+00:00" + checkpoints = { + "leaf-a": _checkpoint("leaf-a", timestamp=timestamp, previous_checkpoint_id="older"), + "leaf-b": _checkpoint("leaf-b", timestamp=timestamp, previous_checkpoint_id="older"), + } + listed = [ + _checkpoint("older", timestamp="2026-01-01T00:00:00+00:00"), + *(checkpoints[checkpoint_id] for checkpoint_id in list_order), + ] + monkeypatch.setattr(storage, "list_checkpoints", AsyncMock(return_value=listed)) + + latest = await storage.get_latest(workflow_name="workflow") + + assert latest is not None + assert latest.checkpoint_id == "leaf-b" + + async def test_get_latest_returns_none_when_no_checkpoints_exist(monkeypatch: pytest.MonkeyPatch) -> None: storage = FoundryCheckpointStore("context-1", _platform_context()) monkeypatch.setattr(storage, "list_checkpoints", AsyncMock(return_value=[])) From ef271e6a379e91af8f38608f50039f2485a6892c Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:42:04 +0530 Subject: [PATCH 4/5] Python: Read the newest Cosmos checkpoints with one ordered query get_latest found the newest timestamp with a MAX query and then read the rows at that timestamp with a second query. The two reads were not atomic: a checkpoint saved in between left the result pinned to the older timestamp, and deleting the newest rows in between returned None while older checkpoints remained. Read the rows newest first with a single ORDER BY c.timestamp DESC query and stop at the first row with an older timestamp, so the tied group comes from the same result set as the newest row. A small page size keeps the query from reading many full checkpoint documents it then discards. Co-Authored-By: Claude Opus 5.5 --- .../_checkpoint_storage.py | 47 +++---- .../tests/test_cosmos_checkpoint_storage.py | 123 +++++++++++++++--- 2 files changed, 123 insertions(+), 47 deletions(-) diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 08ba98690fc..b4158621d7a 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py @@ -25,6 +25,11 @@ logger = logging.getLogger(__name__) +# Page size for get_latest's newest-first query. It is normally stopped after the first row +# or a small group of rows sharing the newest timestamp, so a small page avoids reading many +# full checkpoint documents that are then discarded. +_LATEST_CHECKPOINT_PAGE_SIZE = 10 + def _select_latest_checkpoint(checkpoints: list[WorkflowCheckpoint]) -> WorkflowCheckpoint: """Return the latest checkpoint, resolving timestamp ties through the lineage chain. @@ -371,45 +376,35 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: """ await self._ensure_container_proxy() - # Checkpoints taken at the same superstep boundary can share a timestamp, so fetch + # Checkpoints taken at the same superstep boundary can share a timestamp, so collect # 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" + # instead of letting the backend pick one row. The newest timestamp and its rows come + # from one query ordered newest first rather than a MAX query followed by a second + # query, so a checkpoint saved or deleted in between cannot leave the result pinned to + # a stale timestamp, or empty while older checkpoints remain. Iteration stops at the + # first older row, so only the pages holding the tied group are read. + query = "SELECT * FROM c WHERE c.workflow_name = @workflow_name ORDER BY c.timestamp DESC" 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] query=query, parameters=parameters, partition_key=workflow_name, + max_item_count=_LATEST_CHECKPOINT_PAGE_SIZE, ) - candidates = [self._document_to_checkpoint(item) async for item in items] - if not candidates: + newest_documents: list[dict[str, Any]] = [] + async for item in items: + if newest_documents and item.get("timestamp") != newest_documents[0].get("timestamp"): + break + newest_documents.append(item) + + if not newest_documents: return None + candidates = [self._document_to_checkpoint(document) for document in newest_documents] checkpoint = _select_latest_checkpoint(candidates) logger.debug( "Latest checkpoint for workflow %s is %s", diff --git a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py index 94dc945dee0..fae83fe6c9c 100644 --- a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py +++ b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py @@ -81,10 +81,35 @@ 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)] +def _mock_get_latest_queries(container: MagicMock, docs: list[dict[str, Any]]) -> None: + """Answer get_latest's newest-first query the way Cosmos orders it, keeping tied rows in the given order.""" + newest_first = sorted(docs, key=lambda doc: doc["timestamp"], reverse=True) + container.query_items.return_value = _to_async_iter(newest_first) + + +class _ConcurrentlyWrittenContainer: + """A Cosmos container whose partition another writer changes right after each query is issued. + + Each query reads a snapshot of the partition, newest first, as Cosmos serves one result page. + """ + + def __init__(self, documents: list[dict[str, Any]], concurrent_write: Any) -> None: + self.documents = list(documents) + self._concurrent_write = concurrent_write + self.queries: list[str] = [] + + def query_items(self, *, query: str, partition_key: str, **kwargs: Any) -> AsyncIterator[Any]: + assert "ORDER BY c.timestamp DESC" in query + self.queries.append(query) + snapshot = sorted( + (doc for doc in self.documents if doc["workflow_name"] == partition_key), + key=lambda doc: doc["timestamp"], + reverse=True, + ) + write, self._concurrent_write = self._concurrent_write, None + if write is not None: + write(self.documents) + return _to_async_iter(snapshot) @pytest.fixture @@ -375,35 +400,91 @@ async def test_get_latest_returns_none_when_empty(mock_container: MagicMock) -> assert result is None -async def test_get_latest_fetches_every_checkpoint_at_the_newest_timestamp(mock_container: MagicMock) -> None: +async def test_get_latest_reads_the_newest_rows_with_one_ordered_query(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") - 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" + (query_call,) = mock_container.query_items.call_args_list + assert "ORDER BY c.timestamp DESC" in query_call.kwargs["query"] + assert "MAX(" not in query_call.kwargs["query"] + assert "LIMIT" not in query_call.kwargs["query"] + assert query_call.kwargs["parameters"] == [{"name": "@workflow_name", "value": "test-workflow"}] + assert query_call.kwargs["partition_key"] == "test-workflow" + assert query_call.kwargs["max_item_count"] > 0 -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([]), +async def test_get_latest_stops_reading_at_the_first_older_checkpoint(mock_container: MagicMock) -> None: + timestamp = "2025-06-01T00:00:00+00:00" + tied = [ + _checkpoint_to_cosmos_document(_make_checkpoint(checkpoint_id="parent", timestamp=timestamp)), + _checkpoint_to_cosmos_document( + _make_checkpoint(checkpoint_id="child", previous_checkpoint_id="parent", timestamp=timestamp) + ), ] + older = _checkpoint_to_cosmos_document( + _make_checkpoint(checkpoint_id="older", timestamp="2025-05-01T00:00:00+00:00") + ) + + async def _rows() -> AsyncIterator[Any]: + for doc in [*tied, older]: + yield doc + raise AssertionError("get_latest read past the first older checkpoint") + + mock_container.query_items.return_value = _rows() storage = CosmosCheckpointStorage(container_client=mock_container) + result = await storage.get_latest(workflow_name="test-workflow") + + assert result is not None + assert result.checkpoint_id == "child" + + +async def test_get_latest_is_not_pinned_to_a_stale_timestamp_by_a_concurrent_save() -> None: + older = _make_checkpoint(checkpoint_id="cp-1", timestamp="2025-05-01T00:00:00+00:00") + newest = _make_checkpoint( + checkpoint_id="cp-2", previous_checkpoint_id="cp-1", timestamp="2025-06-01T00:00:00+00:00" + ) + saved = _make_checkpoint(checkpoint_id="cp-3", previous_checkpoint_id="cp-2", timestamp="2025-07-01T00:00:00+00:00") + container = _ConcurrentlyWrittenContainer( + [_checkpoint_to_cosmos_document(older), _checkpoint_to_cosmos_document(newest)], + lambda documents: documents.append(_checkpoint_to_cosmos_document(saved)), + ) + storage = CosmosCheckpointStorage(container_client=container) # type: ignore[arg-type] + + during_save = await storage.get_latest(workflow_name="test-workflow") + after_save = await storage.get_latest(workflow_name="test-workflow") + + # The first call answers from the one read it made, before the save landed. + assert during_save is not None + assert during_save.checkpoint_id == "cp-2" + assert len(container.queries) == 2 + assert after_save is not None + assert after_save.checkpoint_id == "cp-3" + + +async def test_get_latest_falls_back_to_older_checkpoints_when_the_newest_is_deleted_concurrently() -> None: + older = _make_checkpoint(checkpoint_id="cp-1", timestamp="2025-05-01T00:00:00+00:00") + newest = _make_checkpoint( + checkpoint_id="cp-2", previous_checkpoint_id="cp-1", timestamp="2025-06-01T00:00:00+00:00" + ) + container = _ConcurrentlyWrittenContainer( + [_checkpoint_to_cosmos_document(older), _checkpoint_to_cosmos_document(newest)], + lambda documents: documents.remove(next(doc for doc in documents if doc["checkpoint_id"] == "cp-2")), + ) + storage = CosmosCheckpointStorage(container_client=container) # type: ignore[arg-type] + + during_delete = await storage.get_latest(workflow_name="test-workflow") + after_delete = await storage.get_latest(workflow_name="test-workflow") - assert await storage.get_latest(workflow_name="test-workflow") is None + # The first call returns the checkpoint its one read saw, never None while checkpoints remain. + assert during_delete is not None + assert during_delete.checkpoint_id == "cp-2" + assert len(container.queries) == 2 + assert after_delete is not None + assert after_delete.checkpoint_id == "cp-1" @pytest.mark.parametrize("row_order", [("parent", "child", "grandchild"), ("grandchild", "child", "parent")]) From 374f247f3603cf7dab78896c699580f367d4086a Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:54:49 +0530 Subject: [PATCH 5/5] Python: Cast the fake Cosmos container in the race tests for ty ty does not honour the mypy-specific type: ignore[arg-type] code. Co-Authored-By: Claude Opus 5.5 --- .../azure-cosmos/tests/test_cosmos_checkpoint_storage.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py index fae83fe6c9c..87a20fe65ff 100644 --- a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py +++ b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py @@ -7,7 +7,7 @@ from collections.abc import AsyncIterator from contextlib import suppress from dataclasses import dataclass -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -452,7 +452,7 @@ async def test_get_latest_is_not_pinned_to_a_stale_timestamp_by_a_concurrent_sav [_checkpoint_to_cosmos_document(older), _checkpoint_to_cosmos_document(newest)], lambda documents: documents.append(_checkpoint_to_cosmos_document(saved)), ) - storage = CosmosCheckpointStorage(container_client=container) # type: ignore[arg-type] + storage = CosmosCheckpointStorage(container_client=cast(Any, container)) during_save = await storage.get_latest(workflow_name="test-workflow") after_save = await storage.get_latest(workflow_name="test-workflow") @@ -474,7 +474,7 @@ async def test_get_latest_falls_back_to_older_checkpoints_when_the_newest_is_del [_checkpoint_to_cosmos_document(older), _checkpoint_to_cosmos_document(newest)], lambda documents: documents.remove(next(doc for doc in documents if doc["checkpoint_id"] == "cp-2")), ) - storage = CosmosCheckpointStorage(container_client=container) # type: ignore[arg-type] + storage = CosmosCheckpointStorage(container_client=cast(Any, container)) during_delete = await storage.get_latest(workflow_name="test-workflow") after_delete = await storage.get_latest(workflow_name="test-workflow")