From af831b6434dbebcf6ea3ccf6e79b2b5757ccfd9f Mon Sep 17 00:00:00 2001 From: Shivani767 <101629653+Shivani767@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:57:24 +0530 Subject: [PATCH 1/2] Python: Return caller-owned checkpoints from InMemoryCheckpointStorage --- .../agent_framework/_workflows/_checkpoint.py | 17 ++- .../core/tests/workflow/test_checkpoint.py | 110 ++++++++++++++++++ 2 files changed, 123 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 3de9460c86d..4a4618d6aaa 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -127,7 +127,16 @@ def from_dict(cls, data: Mapping[str, Any]) -> WorkflowCheckpoint: class CheckpointStorage(Protocol): - """Protocol for checkpoint storage backends.""" + """Protocol for checkpoint storage backends. + + Ownership: + Checkpoints returned by ``load``, ``list_checkpoints`` and ``get_latest`` are owned by + the caller. Mutating a returned checkpoint must not change stored state, and repeated + reads must return independent objects. Symmetrically, ``save`` snapshots the checkpoint + at call time, so mutating the caller's object afterwards must not change what was stored. + Backends that serialize satisfy this implicitly, because decoding allocates a fresh object + graph; backends that retain live objects must copy explicitly. + """ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: """Save a checkpoint and return its ID. @@ -217,12 +226,12 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: checkpoint = self._checkpoints.get(checkpoint_id) if checkpoint: logger.debug(f"Loaded checkpoint {checkpoint_id} from memory") - return checkpoint + return copy.deepcopy(checkpoint) raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}") async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: """List checkpoint objects for a given workflow name.""" - return [cp for cp in self._checkpoints.values() if cp.workflow_name == workflow_name] + return [copy.deepcopy(cp) for cp in self._checkpoints.values() if cp.workflow_name == workflow_name] async def delete(self, checkpoint_id: CheckpointID) -> bool: """Delete a checkpoint by ID.""" @@ -239,7 +248,7 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: return None latest_checkpoint = max(checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp)) logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}") - return latest_checkpoint + return copy.deepcopy(latest_checkpoint) async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]: """List checkpoint IDs. If workflow_id is provided, filter by that workflow.""" diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 5f3da78cd1d..a8273520329 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -10,6 +10,7 @@ import pytest from agent_framework import ( + CheckpointStorage, FileCheckpointStorage, InMemoryCheckpointStorage, WorkflowCheckpoint, @@ -1766,3 +1767,112 @@ async def test_file_checkpoint_storage_roundtrip_empty_collections(): # endregion + + +# region checkpoint storage conformance + +# These tests define the ownership contract that every CheckpointStorage backend must satisfy: +# a checkpoint handed to the caller is owned by the caller, and a checkpoint handed to save() is +# snapshotted at call time. Backends that serialize (file, Cosmos) get this for free because +# decoding allocates a fresh object graph; backends that hold live objects must copy explicitly. + + +@pytest.fixture(params=["memory", "file"]) +def conformance_storage(request: pytest.FixtureRequest, tmp_path: Path) -> CheckpointStorage: + """Yield each in-tree checkpoint storage backend for the shared contract tests.""" + if request.param == "memory": + return InMemoryCheckpointStorage() + return FileCheckpointStorage(tmp_path / "conformance") + + +def _conformance_checkpoint(workflow_name: str = "conformance-workflow") -> WorkflowCheckpoint: + """Build a checkpoint whose state holds nested mutable containers.""" + return WorkflowCheckpoint( + workflow_name=workflow_name, + graph_signature_hash="conformance-hash", + state={ + "shared": {"counter": 0, "history": ["initial"]}, + "_executor_state": {"executor1": {"visits": ["first"]}}, + }, + metadata={"tags": ["initial"]}, + ) + + +async def test_conformance_load_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None: + """Mutating a loaded checkpoint must not alter what the backend has stored.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + loaded = await conformance_storage.load(checkpoint.checkpoint_id) + loaded.state["shared"]["counter"] = 999 + loaded.state["shared"]["history"].append("mutated") + loaded.state["_executor_state"]["executor1"]["visits"].append("mutated") + loaded.metadata["tags"].append("mutated") + + reloaded = await conformance_storage.load(checkpoint.checkpoint_id) + assert reloaded.state["shared"]["counter"] == 0 + assert reloaded.state["shared"]["history"] == ["initial"] + assert reloaded.state["_executor_state"]["executor1"]["visits"] == ["first"] + assert reloaded.metadata["tags"] == ["initial"] + + +async def test_conformance_repeated_loads_are_independent(conformance_storage: CheckpointStorage) -> None: + """Two loads of one checkpoint must not share mutable state.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + first = await conformance_storage.load(checkpoint.checkpoint_id) + second = await conformance_storage.load(checkpoint.checkpoint_id) + assert first is not second + + first.state["shared"]["counter"] = 42 + first.state["shared"]["history"].append("from-first") + + assert second.state["shared"]["counter"] == 0 + assert second.state["shared"]["history"] == ["initial"] + + +async def test_conformance_get_latest_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None: + """Mutating the result of get_latest must not alter stored state.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + latest = await conformance_storage.get_latest(workflow_name=checkpoint.workflow_name) + assert latest is not None + latest.state["shared"]["counter"] = 123 + + reloaded = await conformance_storage.get_latest(workflow_name=checkpoint.workflow_name) + assert reloaded is not None + assert reloaded.state["shared"]["counter"] == 0 + + +async def test_conformance_list_checkpoints_returns_caller_owned_copies( + conformance_storage: CheckpointStorage, +) -> None: + """Mutating a checkpoint from list_checkpoints must not alter stored state.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + listed = await conformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name) + assert len(listed) == 1 + listed[0].state["shared"]["history"].append("mutated") + + relisted = await conformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name) + assert len(relisted) == 1 + assert relisted[0].state["shared"]["history"] == ["initial"] + + +async def test_conformance_save_snapshots_state_at_call_time(conformance_storage: CheckpointStorage) -> None: + """Mutating the caller's object after save must not alter the stored checkpoint.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + checkpoint.state["shared"]["counter"] = 7 + checkpoint.state["shared"]["history"].append("after-save") + + loaded = await conformance_storage.load(checkpoint.checkpoint_id) + assert loaded.state["shared"]["counter"] == 0 + assert loaded.state["shared"]["history"] == ["initial"] + + +# endregion From c0585ffc2274163191fe8fed4d1bbe111ddda794 Mon Sep 17 00:00:00 2001 From: Shivani767 <101629653+Shivani767@users.noreply.github.com> Date: Thu, 20 Aug 2026 22:38:32 +0530 Subject: [PATCH 2/2] Python: Move checkpoint storage conformance tests to a dedicated module Keep the ownership contract tests out of the already-large test_checkpoint.py and cover nested messages and pending request events as well as state/metadata. --- .../core/tests/workflow/test_checkpoint.py | 110 ----------- .../test_checkpoint_storage_conformance.py | 185 ++++++++++++++++++ 2 files changed, 185 insertions(+), 110 deletions(-) create mode 100644 python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index a8273520329..5f3da78cd1d 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -10,7 +10,6 @@ import pytest from agent_framework import ( - CheckpointStorage, FileCheckpointStorage, InMemoryCheckpointStorage, WorkflowCheckpoint, @@ -1767,112 +1766,3 @@ async def test_file_checkpoint_storage_roundtrip_empty_collections(): # endregion - - -# region checkpoint storage conformance - -# These tests define the ownership contract that every CheckpointStorage backend must satisfy: -# a checkpoint handed to the caller is owned by the caller, and a checkpoint handed to save() is -# snapshotted at call time. Backends that serialize (file, Cosmos) get this for free because -# decoding allocates a fresh object graph; backends that hold live objects must copy explicitly. - - -@pytest.fixture(params=["memory", "file"]) -def conformance_storage(request: pytest.FixtureRequest, tmp_path: Path) -> CheckpointStorage: - """Yield each in-tree checkpoint storage backend for the shared contract tests.""" - if request.param == "memory": - return InMemoryCheckpointStorage() - return FileCheckpointStorage(tmp_path / "conformance") - - -def _conformance_checkpoint(workflow_name: str = "conformance-workflow") -> WorkflowCheckpoint: - """Build a checkpoint whose state holds nested mutable containers.""" - return WorkflowCheckpoint( - workflow_name=workflow_name, - graph_signature_hash="conformance-hash", - state={ - "shared": {"counter": 0, "history": ["initial"]}, - "_executor_state": {"executor1": {"visits": ["first"]}}, - }, - metadata={"tags": ["initial"]}, - ) - - -async def test_conformance_load_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None: - """Mutating a loaded checkpoint must not alter what the backend has stored.""" - checkpoint = _conformance_checkpoint() - await conformance_storage.save(checkpoint) - - loaded = await conformance_storage.load(checkpoint.checkpoint_id) - loaded.state["shared"]["counter"] = 999 - loaded.state["shared"]["history"].append("mutated") - loaded.state["_executor_state"]["executor1"]["visits"].append("mutated") - loaded.metadata["tags"].append("mutated") - - reloaded = await conformance_storage.load(checkpoint.checkpoint_id) - assert reloaded.state["shared"]["counter"] == 0 - assert reloaded.state["shared"]["history"] == ["initial"] - assert reloaded.state["_executor_state"]["executor1"]["visits"] == ["first"] - assert reloaded.metadata["tags"] == ["initial"] - - -async def test_conformance_repeated_loads_are_independent(conformance_storage: CheckpointStorage) -> None: - """Two loads of one checkpoint must not share mutable state.""" - checkpoint = _conformance_checkpoint() - await conformance_storage.save(checkpoint) - - first = await conformance_storage.load(checkpoint.checkpoint_id) - second = await conformance_storage.load(checkpoint.checkpoint_id) - assert first is not second - - first.state["shared"]["counter"] = 42 - first.state["shared"]["history"].append("from-first") - - assert second.state["shared"]["counter"] == 0 - assert second.state["shared"]["history"] == ["initial"] - - -async def test_conformance_get_latest_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None: - """Mutating the result of get_latest must not alter stored state.""" - checkpoint = _conformance_checkpoint() - await conformance_storage.save(checkpoint) - - latest = await conformance_storage.get_latest(workflow_name=checkpoint.workflow_name) - assert latest is not None - latest.state["shared"]["counter"] = 123 - - reloaded = await conformance_storage.get_latest(workflow_name=checkpoint.workflow_name) - assert reloaded is not None - assert reloaded.state["shared"]["counter"] == 0 - - -async def test_conformance_list_checkpoints_returns_caller_owned_copies( - conformance_storage: CheckpointStorage, -) -> None: - """Mutating a checkpoint from list_checkpoints must not alter stored state.""" - checkpoint = _conformance_checkpoint() - await conformance_storage.save(checkpoint) - - listed = await conformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name) - assert len(listed) == 1 - listed[0].state["shared"]["history"].append("mutated") - - relisted = await conformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name) - assert len(relisted) == 1 - assert relisted[0].state["shared"]["history"] == ["initial"] - - -async def test_conformance_save_snapshots_state_at_call_time(conformance_storage: CheckpointStorage) -> None: - """Mutating the caller's object after save must not alter the stored checkpoint.""" - checkpoint = _conformance_checkpoint() - await conformance_storage.save(checkpoint) - - checkpoint.state["shared"]["counter"] = 7 - checkpoint.state["shared"]["history"].append("after-save") - - loaded = await conformance_storage.load(checkpoint.checkpoint_id) - assert loaded.state["shared"]["counter"] == 0 - assert loaded.state["shared"]["history"] == ["initial"] - - -# endregion diff --git a/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py b/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py new file mode 100644 index 00000000000..383e9642c75 --- /dev/null +++ b/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py @@ -0,0 +1,185 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Conformance tests for the CheckpointStorage ownership contract. + +A checkpoint handed to the caller is owned by the caller, and a checkpoint +handed to ``save()`` is snapshotted at call time. Backends that serialize +(file, Cosmos) get this for free because decoding allocates a fresh object +graph; backends that hold live objects must copy explicitly. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, cast + +import pytest + +from agent_framework import ( + CheckpointStorage, + FileCheckpointStorage, + InMemoryCheckpointStorage, + WorkflowCheckpoint, + WorkflowCheckpointException, + WorkflowEvent, +) +from agent_framework._workflows._runner_context import WorkflowMessage + + +@pytest.fixture(params=["memory", "file"]) +def conformance_storage(request: pytest.FixtureRequest, tmp_path: Path) -> CheckpointStorage: + """Yield each in-tree checkpoint storage backend for the shared contract tests.""" + if request.param == "memory": + return InMemoryCheckpointStorage() + return FileCheckpointStorage(tmp_path / "conformance") + + +def _conformance_checkpoint(workflow_name: str = "conformance-workflow") -> WorkflowCheckpoint: + """Build a checkpoint whose object graph holds nested mutable containers.""" + return WorkflowCheckpoint( + workflow_name=workflow_name, + graph_signature_hash="conformance-hash", + state={ + "shared": {"counter": 0, "history": ["initial"]}, + "_executor_state": {"executor1": {"visits": ["first"]}}, + }, + messages={ + "executor1": [ + WorkflowMessage( + data={"text": "hello", "tags": ["initial"]}, + source_id="src", + target_id="tgt", + ) + ] + }, + pending_request_info_events={ + "req1": WorkflowEvent.request_info( + request_id="req1", + source_executor_id="executor1", + request_data={"payload": ["initial"]}, + response_type=str, + ), + }, + metadata={"tags": ["initial"]}, + ) + + +def _mutate(checkpoint: WorkflowCheckpoint) -> None: + """Mutate every nested container the ownership contract covers.""" + checkpoint.state["shared"]["counter"] = 999 + checkpoint.state["shared"]["history"].append("mutated") + checkpoint.state["_executor_state"]["executor1"]["visits"].append("mutated") + cast(dict[str, Any], checkpoint.messages["executor1"][0].data)["tags"].append("mutated") + cast(dict[str, Any], checkpoint.pending_request_info_events["req1"].data)["payload"].append("mutated") + checkpoint.metadata["tags"].append("mutated") + + +def _assert_pristine(checkpoint: WorkflowCheckpoint) -> None: + """Assert nested containers still hold their original values.""" + assert checkpoint.state["shared"]["counter"] == 0 + assert checkpoint.state["shared"]["history"] == ["initial"] + assert checkpoint.state["_executor_state"]["executor1"]["visits"] == ["first"] + assert cast(dict[str, Any], checkpoint.messages["executor1"][0].data)["tags"] == ["initial"] + assert cast(dict[str, Any], checkpoint.pending_request_info_events["req1"].data)["payload"] == ["initial"] + assert checkpoint.metadata["tags"] == ["initial"] + + +async def test_conformance_load_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None: + """Mutating a loaded checkpoint must not alter what the backend has stored.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + loaded = await conformance_storage.load(checkpoint.checkpoint_id) + _mutate(loaded) + + reloaded = await conformance_storage.load(checkpoint.checkpoint_id) + _assert_pristine(reloaded) + + +async def test_conformance_repeated_loads_are_independent(conformance_storage: CheckpointStorage) -> None: + """Two loads of one checkpoint must not share mutable state.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + first = await conformance_storage.load(checkpoint.checkpoint_id) + second = await conformance_storage.load(checkpoint.checkpoint_id) + assert first is not second + + _mutate(first) + _assert_pristine(second) + + +async def test_conformance_get_latest_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None: + """Mutating the result of get_latest must not alter stored state.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + latest = await conformance_storage.get_latest(workflow_name=checkpoint.workflow_name) + assert latest is not None + _mutate(latest) + + reloaded = await conformance_storage.get_latest(workflow_name=checkpoint.workflow_name) + assert reloaded is not None + _assert_pristine(reloaded) + + +async def test_conformance_list_checkpoints_returns_caller_owned_copies( + conformance_storage: CheckpointStorage, +) -> None: + """Mutating a checkpoint from list_checkpoints must not alter stored state.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + listed = await conformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name) + assert len(listed) == 1 + _mutate(listed[0]) + + relisted = await conformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name) + assert len(relisted) == 1 + _assert_pristine(relisted[0]) + + +async def test_conformance_save_snapshots_state_at_call_time(conformance_storage: CheckpointStorage) -> None: + """Mutating the caller's object after save must not alter the stored checkpoint.""" + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + _mutate(checkpoint) + + loaded = await conformance_storage.load(checkpoint.checkpoint_id) + _assert_pristine(loaded) + + +async def test_conformance_list_checkpoints_filters_by_workflow_name( + conformance_storage: CheckpointStorage, +) -> None: + """list_checkpoints must continue to filter by workflow_name after isolation.""" + first = _conformance_checkpoint(workflow_name="workflow-a") + second = _conformance_checkpoint(workflow_name="workflow-b") + await conformance_storage.save(first) + await conformance_storage.save(second) + + listed_a = await conformance_storage.list_checkpoints(workflow_name="workflow-a") + listed_b = await conformance_storage.list_checkpoints(workflow_name="workflow-b") + + assert {cp.checkpoint_id for cp in listed_a} == {first.checkpoint_id} + assert {cp.checkpoint_id for cp in listed_b} == {second.checkpoint_id} + + +async def test_conformance_get_latest_returns_none_when_empty(conformance_storage: CheckpointStorage) -> None: + """get_latest must return None for an unknown workflow, not raise.""" + result = await conformance_storage.get_latest(workflow_name="missing-workflow") + assert result is None + + +async def test_conformance_load_missing_id_raises(conformance_storage: CheckpointStorage) -> None: + """load must raise WorkflowCheckpointException for an unknown checkpoint id.""" + with pytest.raises(WorkflowCheckpointException): + await conformance_storage.load("does-not-exist") + + +async def test_conformance_save_returns_id(conformance_storage: CheckpointStorage) -> None: + """save must return the checkpoint_id.""" + checkpoint = _conformance_checkpoint() + result = await conformance_storage.save(checkpoint) + assert result == checkpoint.checkpoint_id