diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 2b267979e99..f3eb723160f 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -127,7 +127,13 @@ def from_dict(cls, data: Mapping[str, Any]) -> WorkflowCheckpoint: class CheckpointStorage(Protocol): - """Protocol for checkpoint storage backends.""" + """Protocol for checkpoint storage backends. + + Reads return objects owned by the caller: mutating a checkpoint returned by ``load``, + ``list_checkpoints`` or ``get_latest`` must not change what a later read returns. + Backends that serialize their storage satisfy this by construction; backends that hold + checkpoints in memory must copy on read. + """ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: """Save a checkpoint and return its ID. @@ -217,12 +223,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 +245,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..e7727547c2e 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1765,4 +1765,44 @@ async def test_file_checkpoint_storage_roundtrip_empty_collections(): assert loaded.pending_request_info_events == {} +async def test_memory_checkpoint_storage_load_returns_caller_owned_copy(): + """Mutating a loaded checkpoint must not change what a later read returns.""" + storage = InMemoryCheckpointStorage() + checkpoint = WorkflowCheckpoint( + workflow_name="test-workflow", + graph_signature_hash="test-hash", + state={"answer": "original"}, + ) + await storage.save(checkpoint) + + loaded = await storage.load(checkpoint.checkpoint_id) + loaded.state["answer"] = "mutated" + + reloaded = await storage.load(checkpoint.checkpoint_id) + assert reloaded.state["answer"] == "original" + + +async def test_memory_checkpoint_storage_list_and_get_latest_return_caller_owned_copies(): + """list_checkpoints and get_latest follow the same read-isolation contract as load.""" + storage = InMemoryCheckpointStorage() + checkpoint = WorkflowCheckpoint( + workflow_name="test-workflow", + graph_signature_hash="test-hash", + state={"answer": "original"}, + ) + await storage.save(checkpoint) + + listed = await storage.list_checkpoints(workflow_name="test-workflow") + listed[0].state["answer"] = "mutated-via-list" + + latest = await storage.get_latest(workflow_name="test-workflow") + assert latest is not None + assert latest.state["answer"] == "original" + + latest.state["answer"] = "mutated-via-get-latest" + + reloaded = await storage.load(checkpoint.checkpoint_id) + assert reloaded.state["answer"] == "original" + + # endregion