diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 2b267979e99..0e0c27ad00c 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -213,16 +213,28 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: return checkpoint.checkpoint_id async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: - """Load a checkpoint by ID.""" + """Load a checkpoint by ID. + + Returns a deep copy of the stored checkpoint so that callers can + freely mutate the returned object (and any nested containers) without + affecting the storage backend's internal representation. This makes + the returned checkpoint symmetric with what :meth:`save` accepted. + """ 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] + """List checkpoint objects for a given workflow name. + + Returns deep copies so callers cannot mutate the storage backend's + internal checkpoint objects through the returned list. + """ + 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.""" @@ -233,13 +245,17 @@ async def delete(self, checkpoint_id: CheckpointID) -> bool: return False async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: - """Get the latest checkpoint for a given workflow name.""" + """Get the latest checkpoint for a given workflow name. + + Returns a deep copy, matching the isolation guarantees of + :meth:`load` and :meth:`list_checkpoints`. + """ 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)) 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/agent_framework/_workflows/_state.py b/python/packages/core/agent_framework/_workflows/_state.py index 093cfea8b67..c8f3e2abc3e 100644 --- a/python/packages/core/agent_framework/_workflows/_state.py +++ b/python/packages/core/agent_framework/_workflows/_state.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import copy from typing import Any @@ -104,18 +105,31 @@ def discard(self) -> None: self._pending.clear() def export_state(self) -> dict[str, Any]: - """Export a serialized copy of the committed state. + """Export an isolated copy of the committed state. Note: Does not include pending changes. + + The returned dict (and any mutable containers reachable through its + values) is a deep copy, so later mutations to the live ``State`` — + including in-place mutation of values handed out by :meth:`get` — + will not be reflected in the exported snapshot. This matches the + isolation that :class:`RunnerContext` already applies to ``messages`` + in the checkpoint-construction path and is what makes a snapshot + safe to share with another workflow instance. """ - return dict(self._committed) + return copy.deepcopy(self._committed) def import_state(self, state: dict[str, Any]) -> None: """Import state from a serialized dictionary. Merges into committed state. Does not affect pending changes. + + The incoming ``state`` dict is deep-copied before merge, so later + in-place mutations of the caller's dict (or any mutable container + reachable through its values) will not leak into the committed + state. This is the import-side counterpart to :meth:`export_state`. """ - self._committed.update(state) + self._committed.update(copy.deepcopy(state)) class _DeleteSentinelType: diff --git a/python/packages/core/tests/workflow/test_checkpoint_isolation_7683.py b/python/packages/core/tests/workflow/test_checkpoint_isolation_7683.py new file mode 100644 index 00000000000..5db80010b11 --- /dev/null +++ b/python/packages/core/tests/workflow/test_checkpoint_isolation_7683.py @@ -0,0 +1,198 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Regression tests for https://github.com/microsoft/agent-framework/issues/7683. + +The InMemoryCheckpointStorage must return deep copies from +``load()``, ``list_checkpoints()``, and ``get_latest()`` so that callers +cannot mutate the storage backend's internal checkpoint objects. This +matches the defensive copy already applied in ``save()``. + +These tests import directly from ``_workflows._checkpoint`` and +``agent_framework.exceptions`` to avoid pulling the top-level +``agent_framework`` package, which has heavy runtime dependencies +(opentelemetry, etc.) that are not relevant to the isolation contract +under test. + +The tests are written as synchronous functions that drive an asyncio +event loop internally. This keeps the test file runnable both locally +(without ``pytest-asyncio``) and in CI (where the project's standard +``asyncio_mode = "auto"`` setting would also be honoured). +""" + +from __future__ import annotations + +import asyncio +from datetime import datetime, timezone + +from agent_framework._workflows._checkpoint import ( + InMemoryCheckpointStorage, + WorkflowCheckpoint, +) +from agent_framework.exceptions import WorkflowCheckpointException + + +def _make_checkpoint( + *, + checkpoint_id: str, + workflow_name: str = "demo", + state: dict | None = None, + messages: dict | None = None, + timestamp: str | None = None, +) -> WorkflowCheckpoint: + return WorkflowCheckpoint( + workflow_name=workflow_name, + graph_signature_hash="hash-1", + checkpoint_id=checkpoint_id, + previous_checkpoint_id=None, + timestamp=timestamp or datetime.now(timezone.utc).isoformat(), + messages=messages if messages is not None else {}, + state=state if state is not None else {"history": ["step-1"]}, + pending_request_info_events={}, + iteration_count=1, + ) + + +def test_load_returns_isolated_checkpoint() -> None: + """Mutating a loaded checkpoint must not affect the storage backend.""" + + async def _run() -> None: + storage = InMemoryCheckpointStorage() + await storage.save(_make_checkpoint(checkpoint_id="cp-1")) + + loaded = await storage.load("cp-1") + loaded.state["history"].append("step-2") # mutate the returned snapshot + + reloaded = await storage.load("cp-1") + assert reloaded.state == {"history": ["step-1"]} + + asyncio.run(_run()) + + +def test_load_returns_isolated_messages() -> None: + """The ``messages`` field is also nested mutable state and must be isolated.""" + + async def _run() -> None: + storage = InMemoryCheckpointStorage() + checkpoint = _make_checkpoint( + checkpoint_id="cp-1", + state={"history": ["step-1"]}, + messages={"src": ["msg-1"]}, + ) + await storage.save(checkpoint) + + loaded = await storage.load("cp-1") + loaded.messages["src"].append("msg-2") + + reloaded = await storage.load("cp-1") + assert reloaded.messages == {"src": ["msg-1"]} + + asyncio.run(_run()) + + +def test_list_checkpoints_returns_isolated_snapshots() -> None: + """Mutating checkpoints returned from list_checkpoints must not affect the backend.""" + + async def _run() -> None: + storage = InMemoryCheckpointStorage() + await storage.save(_make_checkpoint(checkpoint_id="cp-1")) + await storage.save(_make_checkpoint(checkpoint_id="cp-2")) + + listed = await storage.list_checkpoints(workflow_name="demo") + assert len(listed) == 2 + + for cp in listed: + cp.state["history"].append("step-2") # mutate each returned snapshot + + # Reload: storage must be untouched. + cp1 = await storage.load("cp-1") + cp2 = await storage.load("cp-2") + assert cp1.state == {"history": ["step-1"]} + assert cp2.state == {"history": ["step-1"]} + + asyncio.run(_run()) + + +def test_list_checkpoints_filters_by_workflow_name() -> None: + """list_checkpoints must continue to filter by workflow_name after isolation.""" + + async def _run() -> None: + storage = InMemoryCheckpointStorage() + await storage.save(_make_checkpoint(checkpoint_id="cp-1", workflow_name="a")) + await storage.save(_make_checkpoint(checkpoint_id="cp-2", workflow_name="b")) + + listed_a = await storage.list_checkpoints(workflow_name="a") + listed_b = await storage.list_checkpoints(workflow_name="b") + + assert {cp.checkpoint_id for cp in listed_a} == {"cp-1"} + assert {cp.checkpoint_id for cp in listed_b} == {"cp-2"} + + asyncio.run(_run()) + + +def test_get_latest_returns_isolated_snapshot() -> None: + """Mutating the result of get_latest must not affect any stored checkpoint.""" + + async def _run() -> None: + storage = InMemoryCheckpointStorage() + cp1 = _make_checkpoint( + checkpoint_id="cp-1", + timestamp="2026-01-01T00:00:00+00:00", + ) + cp2 = _make_checkpoint( + checkpoint_id="cp-2", + timestamp="2026-02-01T00:00:00+00:00", + ) + await storage.save(cp1) + await storage.save(cp2) + + latest = await storage.get_latest(workflow_name="demo") + assert latest is not None + assert latest.checkpoint_id == "cp-2" + + latest.state["history"].append("step-2") # mutate the latest snapshot + + # Both stored checkpoints must be unchanged. + reloaded_latest = await storage.get_latest(workflow_name="demo") + assert reloaded_latest is not None + assert reloaded_latest.state == {"history": ["step-1"]} + # And the older one wasn't touched either. + reloaded_cp1 = await storage.load("cp-1") + assert reloaded_cp1.state == {"history": ["step-1"]} + + asyncio.run(_run()) + + +def test_get_latest_returns_none_when_empty() -> None: + """get_latest must return None for an unknown workflow, not raise.""" + + async def _run() -> None: + storage = InMemoryCheckpointStorage() + result = await storage.get_latest(workflow_name="demo") + assert result is None + + asyncio.run(_run()) + + +def test_load_missing_id_raises() -> None: + """load must raise WorkflowCheckpointException for an unknown checkpoint id.""" + + async def _run() -> None: + storage = InMemoryCheckpointStorage() + try: + await storage.load("does-not-exist") + except WorkflowCheckpointException: + return + raise AssertionError("Expected WorkflowCheckpointException for missing checkpoint id") + + asyncio.run(_run()) + + +def test_save_returns_id() -> None: + """save must return the checkpoint_id (and the contract must hold post-isolation).""" + + async def _run() -> None: + storage = InMemoryCheckpointStorage() + result = await storage.save(_make_checkpoint(checkpoint_id="cp-1")) + assert result == "cp-1" + + asyncio.run(_run()) diff --git a/python/packages/core/tests/workflow/test_state.py b/python/packages/core/tests/workflow/test_state.py index 7781eb4141a..8275c575d71 100644 --- a/python/packages/core/tests/workflow/test_state.py +++ b/python/packages/core/tests/workflow/test_state.py @@ -301,3 +301,101 @@ def test_import_does_not_affect_pending(self) -> None: # Pending is still there assert state.get("pending_key") == "pending_value" assert "pending_key" in state._pending # pyright: ignore[reportPrivateUsage] + + +class TestStateIsolation: + """Tests for isolation between State instances across export/import. + + Regression tests for https://github.com/microsoft/agent-framework/issues/7683 + — checkpoint state must be isolated from live workflow state across + restoration and storage boundaries, so a resumed workflow mutating its + own state cannot corrupt the snapshot it was restored from. + """ + + def test_export_isolates_mutable_list_value(self) -> None: + """Mutating a list value on the source State must not mutate the exported snapshot.""" + source = State() + source.set("history", ["step-1"]) + source.commit() + + snapshot = source.export_state() + + # Read, mutate in place, write back on the live state. + history = source.get("history") + assert history is not None + history.append("step-2") + source.set("history", history) + source.commit() + + # The snapshot taken before the mutation must be unchanged. + assert snapshot == {"history": ["step-1"]} + + def test_export_isolates_mutable_dict_value(self) -> None: + """Mutating a dict value on the source State must not mutate the exported snapshot.""" + source = State() + source.set("counters", {"a": 1}) + source.commit() + + snapshot = source.export_state() + + counters = source.get("counters") + assert counters is not None + counters["b"] = 2 + source.set("counters", counters) + source.commit() + + assert snapshot == {"counters": {"a": 1}} + + def test_import_does_not_alias_caller_state(self) -> None: + """Mutating the caller's dict after import_state must not affect committed state.""" + state = State() + incoming = {"history": ["step-1"]} + state.import_state(incoming) + + # Caller mutates their dict in place after import. + incoming["history"].append("step-2") + incoming["extra"] = "leaked" + + # Committed state must be unaffected. + assert state.get("history") == ["step-1"] + assert state.has("extra") is False + + def test_roundtrip_restoration_isolates_checkpoint(self) -> None: + """The full export → import round-trip must keep a checkpoint stable under in-place mutation. + + Mirrors the checkpoint build / restore path used by RunnerContext.build_checkpoint + and Runner.restore_checkpoint, where the workflow does + history = restored.get("history") + history.append("...") + restored.set("history", history) + after being resumed. + """ + source = State() + source.set("history", ["step-1"]) + source.commit() + checkpoint_state = source.export_state() + + restored = State() + restored.import_state(checkpoint_state) + + history = restored.get("history") + assert history is not None + history.append("step-2") + restored.set("history", history) + restored.commit() + + # Checkpoint snapshot must remain at the original value. + assert checkpoint_state == {"history": ["step-1"]} + + def test_export_returns_independent_dict(self) -> None: + """Mutating the returned dict itself must not affect the source State.""" + source = State() + source.set("a", 1) + source.commit() + + snapshot = source.export_state() + snapshot["b"] = 2 + snapshot.pop("a") + + assert source.get("a") == 1 + assert source.has("b") is False