Python: Prevent workflow checkpoints from being mutated outside of storage implementations

python/packages/core/agent_framework/_harness/_loop.py

@@ -34,7 +34,9 @@
 from pydantic import BaseModel, Field
 from typing_extensions import Self
 
-from .._agents import _LOOP_ITERATION_TOKEN_KEY  # pyright: ignore[reportPrivateUsage] -- shared loop-turn marker, see _agents.py
+from .._agents import (
+    _LOOP_ITERATION_TOKEN_KEY,  # pyright: ignore[reportPrivateUsage] -- shared loop-turn marker, see _agents.py
+)
 from .._feature_stage import ExperimentalFeature, experimental
 from .._middleware import AgentContext, AgentMiddleware, MiddlewareTermination
 from .._sessions import SessionContext
@@ -490,8 +492,7 @@ async def _fire_turn_scoped_after_providers(
         if response is None or run_after is None:
             return
         if not any(
-            getattr(provider, "after_run_once_per_turn", False)
-            for provider in getattr(agent, "context_providers", [])
+            getattr(provider, "after_run_once_per_turn", False) for provider in getattr(agent, "context_providers", [])
         ):
             return
         session_context = SessionContext(

python/packages/core/agent_framework/_workflows/_agent.py

@@ -713,11 +713,7 @@ def _process_request_info_event(
         Note:
             Text requests use the function-call envelope so callers can reply with a matching function result.
         """
-        if (
-            isinstance(event.data, Content)
-            and event.data.user_input_request
-            and event.data.type != "text"
-        ):
+        if isinstance(event.data, Content) and event.data.user_input_request and event.data.type != "text":
             # Preserve specialized requests that callers already understand how to present.
             return event.data
 

python/packages/core/agent_framework/_workflows/_checkpoint.py

@@ -130,7 +130,7 @@ class CheckpointStorage(Protocol):
     """Protocol for checkpoint storage backends."""
 
     async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
-        """Save a checkpoint and return its ID.
+        """Create a copy of the given checkpoint and store it, returning its ID.
 
         Args:
             checkpoint: The WorkflowCheckpoint object to save.
@@ -147,7 +147,7 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
             checkpoint_id: The unique ID of the checkpoint to load.
 
         Returns:
-            The WorkflowCheckpoint object corresponding to the given ID.
+            A copy of the WorkflowCheckpoint object corresponding to the given ID.
 
         Raises:
             WorkflowCheckpointException: If no checkpoint with the given ID exists.
@@ -161,7 +161,7 @@ async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoi
             workflow_name: The name of the workflow to list checkpoints for.
 
         Returns:
-            A list of WorkflowCheckpoint objects for the specified workflow name.
+            A list of copies of WorkflowCheckpoint objects for the specified workflow name.
         """
         ...
 
@@ -183,7 +183,8 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None:
             workflow_name: The name of the workflow to get the latest checkpoint for.
 
         Returns:
-            The latest WorkflowCheckpoint object for the specified workflow name, or None if no checkpoints exist.
+            A copy of the latest WorkflowCheckpoint object for the specified workflow name,
+            or None if no checkpoints exist.
         """
         ...
 
@@ -207,7 +208,7 @@ def __init__(self) -> None:
         self._checkpoints: dict[CheckpointID, WorkflowCheckpoint] = {}
 
     async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
-        """Save a checkpoint and return its ID."""
+        """Create a copy of the given checkpoint and store it, returning its ID."""
         self._checkpoints[checkpoint.checkpoint_id] = copy.deepcopy(checkpoint)
         logger.debug(f"Saved checkpoint {checkpoint.checkpoint_id} to memory")
         return checkpoint.checkpoint_id
@@ -217,12 +218,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 +240,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."""

python/packages/core/agent_framework/_workflows/_state.py

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import copy
 from typing import Any
 
 
@@ -31,7 +32,8 @@ def set(self, key: str, value: Any) -> None:
         """Set a value in the pending state buffer.
 
         The value will be visible to subsequent `get()` calls but won't be
-        committed to the actual state until `commit()` is called.
+        committed to the actual state until `commit()` is called. A deep copy is
+        stored so later caller mutations do not change pending state.
 
         Note:
             When multiple executors run concurrently within the same superstep,
@@ -40,7 +42,7 @@ def set(self, key: str, value: Any) -> None:
             the .NET behavior and the superstep execution model where all executors
             in a superstep see the same committed state at the start.
         """
-        self._pending[key] = value
+        self._pending[key] = copy.deepcopy(value)
 
     def get(self, key: str, default: Any = None) -> Any:
         """Get a value from state, checking pending first then committed.
@@ -50,14 +52,17 @@ def get(self, key: str, default: Any = None) -> Any:
             default: Value to return if key is not found. Defaults to None.
 
         Returns:
-            The value if found, otherwise the default value.
+            A deep copy of the value if found, otherwise the default value. Mutate
+            the returned value and pass it to :meth:`set` to update workflow state.
         """
         if key in self._pending:
             value = self._pending[key]
             if value is _DeleteSentinel:
                 return default
-            return value
-        return self._committed.get(key, default)
+            return copy.deepcopy(value)
+        if key in self._committed:
+            return copy.deepcopy(self._committed[key])
+        return default
 
     def has(self, key: str) -> bool:
         """Check if a key exists in pending or committed state."""
@@ -104,18 +109,20 @@ def discard(self) -> None:
         self._pending.clear()
 
     def export_state(self) -> dict[str, Any]:
-        """Export a serialized copy of the committed state.
+        """Export a deepcopy of the committed state.
 
-        Note: Does not include pending changes.
+        Note:
+            Does not include pending changes. Values must support :func:`copy.deepcopy`.
         """
-        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.
+        Merges a deepcopy into committed state. Does not affect pending changes.
+        Values must support :func:`copy.deepcopy`.
         """
-        self._committed.update(state)
+        self._committed.update(copy.deepcopy(state))
 
 
 class _DeleteSentinelType:

python/packages/core/tests/core/test_harness_loop.py

@@ -32,9 +32,9 @@
     background_tasks_running,
     background_tasks_running_message,
     set_agent_mode,
-    tool,
     todos_remaining,
     todos_remaining_message,
+    tool,
 )
 from agent_framework._harness._loop import (
     DEFAULT_JUDGE_MAX_ITERATIONS,

python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py

@@ -0,0 +1,154 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Conformance tests for the CheckpointStorage ownership contract."""
+
+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."""
+    if request.param == "memory":
+        return InMemoryCheckpointStorage()
+    return FileCheckpointStorage(tmp_path / "conformance")
+
+
+def _conformance_checkpoint(workflow_name: str = "conformance-workflow") -> WorkflowCheckpoint:
+    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:
+    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 checkpoint.state["shared"] == {"counter": 0, "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_load_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None:
+    checkpoint = _conformance_checkpoint()
+    await conformance_storage.save(checkpoint)
+
+    loaded = await conformance_storage.load(checkpoint.checkpoint_id)
+    _mutate(loaded)
+
+    _assert_pristine(await conformance_storage.load(checkpoint.checkpoint_id))
+
+
+async def test_repeated_loads_are_independent(conformance_storage: CheckpointStorage) -> None:
+    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_get_latest_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None:
+    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_list_checkpoints_returns_caller_owned_copies(conformance_storage: CheckpointStorage) -> None:
+    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_save_snapshots_at_call_time(conformance_storage: CheckpointStorage) -> None:
+    checkpoint = _conformance_checkpoint()
+    await conformance_storage.save(checkpoint)
+
+    _mutate(checkpoint)
+
+    _assert_pristine(await conformance_storage.load(checkpoint.checkpoint_id))
+
+
+async def test_list_checkpoints_filters_by_workflow_name(conformance_storage: CheckpointStorage) -> None:
+    first = _conformance_checkpoint("workflow-a")
+    second = _conformance_checkpoint("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 {checkpoint.checkpoint_id for checkpoint in listed_a} == {first.checkpoint_id}
+    assert {checkpoint.checkpoint_id for checkpoint in listed_b} == {second.checkpoint_id}
+
+
+async def test_get_latest_returns_none_when_empty(conformance_storage: CheckpointStorage) -> None:
+    assert await conformance_storage.get_latest(workflow_name="missing-workflow") is None
+
+
+async def test_load_missing_id_raises(conformance_storage: CheckpointStorage) -> None:
+    with pytest.raises(WorkflowCheckpointException):
+        await conformance_storage.load("does-not-exist")
+
+
+async def test_save_returns_id(conformance_storage: CheckpointStorage) -> None:
+    checkpoint = _conformance_checkpoint()
+    assert await conformance_storage.save(checkpoint) == checkpoint.checkpoint_id

python/packages/core/tests/workflow/test_runner.py

@@ -542,21 +542,27 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
 
     # Establish some state to capture.
     executor.count = 7
-    state.set("shared_key", "shared_value")
+    state.set("shared_key", {"history": ["shared_value"]})
     state.commit()
 
     checkpoint = await runner.build_checkpoint()
     assert checkpoint.graph_signature_hash == "test_hash"
 
     # Mutate after capture; restoring must roll back to the captured snapshot.
     executor.count = 999
-    state.set("shared_key", "mutated")
+    shared_state = state.get("shared_key")
+    shared_state["history"].append("mutated")
+    state.set("shared_key", shared_state)
     state.commit()
 
     await runner.restore_checkpoint(checkpoint)
 
     assert executor.count == 7
-    assert state.get("shared_key") == "shared_value"
+    assert state.get("shared_key") == {"history": ["shared_value"]}
+    restored_state = state.get("shared_key")
+    restored_state["history"].append("restored-mutation")
+    state.set("shared_key", restored_state)
+    assert checkpoint.state["shared_key"] == {"history": ["shared_value"]}
     assert runner._previous_checkpoint_id == checkpoint.checkpoint_id  # pyright: ignore[reportPrivateUsage]