Python: Isolate checkpoint state from live workflow state across restoration and storage boundaries

[Karan Dhaodiyal (karandhaodiyal28-hash)]

Fixes #7683.

Summary

WorkflowCheckpoint is documented as a stable, shareable snapshot of
workflow state — but the actual implementation leaked mutable state in
two places. A resumed workflow could silently corrupt a checkpoint
snapshot (or the storage backend's internal representation) through an
ordinary read / mutate / write-back pattern, with no exception raised.

Root cause

Two related but distinct isolation gaps, both surfaced in #7683:

1. Restoration boundary. State.export_state() and
   State.import_state() were shallow copies (dict(self._committed)
   and self._committed.update(state)). Mutable values in a checkpoint
   remained aliased by reference between the snapshot and the live
   State it was built from or restored into.

2. Storage boundary. InMemoryCheckpointStorage.save() already
   deep-copies before storing, but load(), list_checkpoints(), and
   get_latest() returned the internally stored object by reference —
   asymmetric with save().

This is inconsistent with the isolation that RunnerContext.build_checkpoint
already explicitly applies to messages (via a per-source list(messages)
copy), and with the docstring on WorkflowCheckpoint that says checkpoints
"can be shared and restored across different workflow instances of the same
workflow definition."

Fix

• State.export_state() returns copy.deepcopy(self._committed).
• State.import_state() deep-copies its argument before merging.
• InMemoryCheckpointStorage.load(), list_checkpoints(), and
  get_latest() each return a deep copy, matching save().

The change is localized to the existing checkpoint / state boundaries
and uses only copy.deepcopy (already imported in _checkpoint.py,
added to _state.py).

Tests

Two existing repros from the issue are turned into regression tests,
plus additional coverage for edge cases.

tests/workflow/test_state.py (5 new tests, all pass; 23 pre-existing
tests continue to pass):

• test_export_isolates_mutable_list_value — list values in an exported
  snapshot are independent of the source State.
• test_export_isolates_mutable_dict_value — same for dict values.
• test_import_does_not_alias_caller_state — caller-side dict is not
  aliased to committed state.
• test_roundtrip_restoration_isolates_checkpoint — full export →
  import round-trip keeps a snapshot stable under in-place mutation
  (the original Repro 1).
• test_export_returns_independent_dict — the exported dict itself
  cannot mutate the source.

tests/workflow/test_checkpoint_isolation_7683.py (new, 8 tests, all
pass):

• test_load_returns_isolated_checkpoint — Repro 2: mutating a loaded
  checkpoint must not affect the backend.
• test_load_returns_isolated_messages — the messages field is also
  nested mutable state and must be isolated.
• test_list_checkpoints_returns_isolated_snapshots — mutating any
  checkpoint returned from list_checkpoints must not affect the
  backend.
• test_list_checkpoints_filters_by_workflow_name — workflow_name
  filtering still works after isolation.
• test_get_latest_returns_isolated_snapshot — mutating the result of
  get_latest must not affect any stored checkpoint.
• test_get_latest_returns_none_when_empty — empty-storage contract.
• test_load_missing_id_raises — error path.
• test_save_returns_id — save() id-return contract.

Local verification: 36 passed, 0 failed in test_state.py +
test_checkpoint_isolation_7683.py (the full test_checkpoint.py suite
needs pytest-asyncio and other deps, which are not in this local env
but are present in CI).

Notes

The trade-off mentioned in the issue (extra copy cost on every
checkpoint build / restore / read) is acknowledged but should be
acceptable for the typical checkpoint cadence (superstep boundary, not
hot path), and matches the explicit isolation already applied to
messages. A targeted optimization — e.g. copy only the values crossing
the boundary, or a frozen-like wrapper — can be a follow-up if
profiling shows the cost is real.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Tao Chen (TaoChenOSU)]

/review

[Tao Chen (TaoChenOSU)]

Prefer adding the new tests to the existing test_checkpoint.py tests or a new test_immutable_checkpoint_test.py. Please do not include the issue number in the file name.

[Tao Chen (TaoChenOSU)]

We can make the tests themselves async.

[github-actions]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/core/agent_framework/_workflows
_checkpoint.py	158	1	99%	325
_state.py	44	0	100%
TOTAL	46780	4357	90%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
9465	36 💤	0 ❌	0 🔥	2m 38s ⏱️

[github-actions]

MAF Automated Review — Iteration 1

Result: Findings reported
Scope: full PR (1 commit(s)): 820acb5b06d6
Model: gpt-5.6-sol

Overview

The PR consistently applies deep-copy isolation at state export/import and at every object-returning in-memory storage boundary, with focused regression tests for nested mutable containers and existing storage contracts. Those guards address ordinary list and dictionary aliasing, but making deepcopy mandatory at the state boundary narrows the existing checkpoint value contract and can disable checkpoint creation or restoration for otherwise supported custom values.

Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
1 verified finding remained after source verification (1 medium) across 1 file. Details are attached to the affected lines below.

Affected areas: python/packages/core/agent_framework/_workflows/_state.py

[github-actions]

Checkpoint state supports custom pickle-serializable values through allowed_checkpoint_types, but those values are not necessarily deepcopyable. A value whose __deepcopy__ raises could previously be saved by file/Cosmos storage; this line now makes checkpoint creation fail before serialization, and the matching import copy makes existing checkpoints unrestorable. Please preserve isolation without narrowing the supported checkpoint-value contract, and cover a pickleable non-deepcopyable value in both build and restore paths.

[Tao Chen (TaoChenOSU)]

Karan Dhaodiyal (@karandhaodiyal28-hash) Thank you for your contributions! Please address open comments and fix the failing code quality check.

[Shivani . (Shivani767)]

Karan Dhaodiyal (@karandhaodiyal28-hash) Tao Chen (@TaoChenOSU) — following Tao's request on #7712 to fold the backend-parametrized ownership tests into this PR so the bug is addressed in one place.

What this PR already has that #7712 does not: the State.export_state / import_state isolation, which is the other half of #7683.

What #7712 uniquely adds (inline below):

1. Protocol ownership contract. Please document ownership on CheckpointStorage itself, not only on the in-memory methods. File and Cosmos already satisfy it by reconstructing from serialized form; without a stated contract a new backend can silently diverge again. Suggested protocol docstring:

class CheckpointStorage(Protocol):
    """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.
    """

2. Replace test_checkpoint_isolation_7683.py with test_checkpoint_storage_conformance.py (no issue number; async tests — matching your earlier comments here). Parametrized over in-memory and file storage so the same ownership / snapshot contract cannot regress in one backend only. Also covers nested messages and pending_request_info_events.

I ran the replacement suite locally: 18 passed (9 tests × 2 backends). Same file is in #7712 if that's easier to copy.

A related nit: ruff format wants the list_checkpoints comprehension on one line. That wrap is likely the failing code quality check.

[Shivani . (Shivani767)]

ruff format (line-length 120) collapses this back to a single line, which is likely the failing code quality check.

Suggested change

	return [
	copy.deepcopy(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]

[Shivani . (Shivani767)]

Please replace this file with python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py (no issue number, async tests). That also keeps these out of the already-large test_checkpoint.py.

The suite is parametrized over memory and file backends so the ownership contract is asserted for every in-tree storage, not only InMemoryCheckpointStorage. It covers nested state, metadata, messages, and pending_request_info_events, plus the filter / missing-id / empty-latest / save-returns-id cases from this PR.

Same file is in #7712

# 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