Uh oh!
There was an error while loading. Please reload this page.
Python: Isolate checkpoint state from live workflow state across restoration and storage boundaries - #7697
Conversation
…and storage boundaries Closesmicrosoft#7683. State.export_state() and State.import_state() used to perform shallow copies, so any mutable container (list, dict) in a checkpoint snapshot remained aliased by reference between the snapshot and the live State it was built from or restored into. InMemoryCheckpointStorage.load(), list_checkpoints(), and get_latest() similarly returned the internally stored checkpoint object by reference, asymmetric with save() which already deep-copies. Both gaps let a resumed workflow silently corrupt a checkpoint snapshot through an ordinary read/mutate/write-back pattern, with no exception raised. Fix: - State.export_state() returns copy.deepcopy(self._committed) and State.import_state() deep-copies its argument before merging, so a snapshot is independent of any State it crosses. - InMemoryCheckpointStorage.load(), list_checkpoints(), and get_latest() now return deep copies of their stored checkpoints, matching the existing save() isolation and closing the storage leak. Tests: - Added 5 isolation tests to tests/workflow/test_state.py covering mutable list / dict values, caller-side aliasing, full round-trip restoration, and dict-level independence. - Added tests/workflow/test_checkpoint_isolation_7683.py with 8 tests covering load(), list_checkpoints(), get_latest(), workflow_name filtering, missing-id error path, messages field isolation, and the save() id return contract. All new tests pass; the 23 pre-existing state tests continue to pass.
Tao Chen (TaoChenOSU)
commented
Aug 17, 2026
/review |
There was a problem hiding this comment.
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.
| def test_load_returns_isolated_checkpoint() -> None: | ||
| """Mutating a loaded checkpoint must not affect the storage backend.""" | ||
| async def _run() -> None: |
There was a problem hiding this comment.
We can make the tests themselves async.
Python Test Coverage Report •
Python Unit Test Overview
| |||||||||||||||||||||||||||||||||||
There was a problem hiding this comment.
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
| safe to share with another workflow instance. | ||
| """ | ||
| return dict(self._committed) | ||
| return copy.deepcopy(self._committed) |
There was a problem hiding this comment.
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)
commented
Aug 17, 2026
Karan Dhaodiyal (@karandhaodiyal28-hash) Thank you for your contributions! Please address open comments and fix the failing code quality check. |
Shivani . (Shivani767)
left a comment
There was a problem hiding this comment.
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):
- Protocol ownership contract. Please document ownership on
CheckpointStorageitself, 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:
classCheckpointStorage(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. """- Replace
test_checkpoint_isolation_7683.pywithtest_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 nestedmessagesandpending_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.
| return [ | ||
| copy.deepcopy(cp) for cp in self._checkpoints.values() if cp.workflow_name == workflow_name | ||
| ] |
There was a problem hiding this comment.
ruff format (line-length 120) collapses this back to a single line, which is likely the failing code quality check.
| return [ | |
| copy.deepcopy(cp) forcpinself._checkpoints.values() ifcp.workflow_name==workflow_name | |
| ] | |
| return [copy.deepcopy(cp) forcpinself._checkpoints.values() ifcp.workflow_name==workflow_name] |
| @@ -0,0 +1,198 @@ | |||
| # Copyright (c) Microsoft. All rights reserved. | |||
There was a problem hiding this comment.
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 checkpointhanded to ``save()`` is snapshotted at call time. Backends that serialize(file, Cosmos) get this for free because decoding allocates a fresh objectgraph; backends that hold live objects must copy explicitly."""from __future__ importannotationsfrompathlibimportPathfromtypingimportAny, castimportpytestfromagent_frameworkimport (
CheckpointStorage,
FileCheckpointStorage,
InMemoryCheckpointStorage,
WorkflowCheckpoint,
WorkflowCheckpointException,
WorkflowEvent,
)
fromagent_framework._workflows._runner_contextimportWorkflowMessage@pytest.fixture(params=["memory", "file"])defconformance_storage(request: pytest.FixtureRequest, tmp_path: Path) ->CheckpointStorage:
"""Yield each in-tree checkpoint storage backend for the shared contract tests."""ifrequest.param=="memory":
returnInMemoryCheckpointStorage()
returnFileCheckpointStorage(tmp_path/"conformance")
def_conformance_checkpoint(workflow_name: str="conformance-workflow") ->WorkflowCheckpoint:
"""Build a checkpoint whose object graph holds nested mutable containers."""returnWorkflowCheckpoint(
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"] =999checkpoint.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."""assertcheckpoint.state["shared"]["counter"] ==0assertcheckpoint.state["shared"]["history"] == ["initial"]
assertcheckpoint.state["_executor_state"]["executor1"]["visits"] == ["first"]
assertcast(dict[str, Any], checkpoint.messages["executor1"][0].data)["tags"] == ["initial"]
assertcast(dict[str, Any], checkpoint.pending_request_info_events["req1"].data)["payload"] == ["initial"]
assertcheckpoint.metadata["tags"] == ["initial"]
asyncdeftest_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()
awaitconformance_storage.save(checkpoint)
loaded=awaitconformance_storage.load(checkpoint.checkpoint_id)
_mutate(loaded)
reloaded=awaitconformance_storage.load(checkpoint.checkpoint_id)
_assert_pristine(reloaded)
asyncdeftest_conformance_repeated_loads_are_independent(conformance_storage: CheckpointStorage) ->None:
"""Two loads of one checkpoint must not share mutable state."""checkpoint=_conformance_checkpoint()
awaitconformance_storage.save(checkpoint)
first=awaitconformance_storage.load(checkpoint.checkpoint_id)
second=awaitconformance_storage.load(checkpoint.checkpoint_id)
assertfirstisnotsecond_mutate(first)
_assert_pristine(second)
asyncdeftest_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()
awaitconformance_storage.save(checkpoint)
latest=awaitconformance_storage.get_latest(workflow_name=checkpoint.workflow_name)
assertlatestisnotNone_mutate(latest)
reloaded=awaitconformance_storage.get_latest(workflow_name=checkpoint.workflow_name)
assertreloadedisnotNone_assert_pristine(reloaded)
asyncdeftest_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()
awaitconformance_storage.save(checkpoint)
listed=awaitconformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name)
assertlen(listed) ==1_mutate(listed[0])
relisted=awaitconformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name)
assertlen(relisted) ==1_assert_pristine(relisted[0])
asyncdeftest_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()
awaitconformance_storage.save(checkpoint)
_mutate(checkpoint)
loaded=awaitconformance_storage.load(checkpoint.checkpoint_id)
_assert_pristine(loaded)
asyncdeftest_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")
awaitconformance_storage.save(first)
awaitconformance_storage.save(second)
listed_a=awaitconformance_storage.list_checkpoints(workflow_name="workflow-a")
listed_b=awaitconformance_storage.list_checkpoints(workflow_name="workflow-b")
assert {cp.checkpoint_idforcpinlisted_a} == {first.checkpoint_id}
assert {cp.checkpoint_idforcpinlisted_b} == {second.checkpoint_id}
asyncdeftest_conformance_get_latest_returns_none_when_empty(conformance_storage: CheckpointStorage) ->None:
"""get_latest must return None for an unknown workflow, not raise."""result=awaitconformance_storage.get_latest(workflow_name="missing-workflow")
assertresultisNoneasyncdeftest_conformance_load_missing_id_raises(conformance_storage: CheckpointStorage) ->None:
"""load must raise WorkflowCheckpointException for an unknown checkpoint id."""withpytest.raises(WorkflowCheckpointException):
awaitconformance_storage.load("does-not-exist")
asyncdeftest_conformance_save_returns_id(conformance_storage: CheckpointStorage) ->None:
"""save must return the checkpoint_id."""checkpoint=_conformance_checkpoint()
result=awaitconformance_storage.save(checkpoint)
assertresult==checkpoint.checkpoint_id
Fixes#7683.
Summary
WorkflowCheckpointis documented as a stable, shareable snapshot ofworkflow 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:
Restoration boundary.
State.export_state()andState.import_state()were shallow copies (dict(self._committed)and
self._committed.update(state)). Mutable values in a checkpointremained aliased by reference between the snapshot and the live
Stateit was built from or restored into.Storage boundary.
InMemoryCheckpointStorage.save()alreadydeep-copies before storing, but
load(),list_checkpoints(), andget_latest()returned the internally stored object by reference —asymmetric with
save().This is inconsistent with the isolation that
RunnerContext.build_checkpointalready explicitly applies to
messages(via a per-sourcelist(messages)copy), and with the docstring on
WorkflowCheckpointthat says checkpoints"can be shared and restored across different workflow instances of the same
workflow definition."
Fix
State.export_state()returnscopy.deepcopy(self._committed).State.import_state()deep-copies its argument before merging.InMemoryCheckpointStorage.load(),list_checkpoints(), andget_latest()each return a deep copy, matchingsave().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-existingtests continue to pass):
test_export_isolates_mutable_list_value— list values in an exportedsnapshot 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 notaliased 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 itselfcannot mutate the source.
tests/workflow/test_checkpoint_isolation_7683.py(new, 8 tests, allpass):
test_load_returns_isolated_checkpoint— Repro 2: mutating a loadedcheckpoint must not affect the backend.
test_load_returns_isolated_messages— themessagesfield is alsonested mutable state and must be isolated.
test_list_checkpoints_returns_isolated_snapshots— mutating anycheckpoint returned from
list_checkpointsmust not affect thebackend.
test_list_checkpoints_filters_by_workflow_name— workflow_namefiltering still works after isolation.
test_get_latest_returns_isolated_snapshot— mutating the result ofget_latestmust 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 fulltest_checkpoint.pysuiteneeds
pytest-asyncioand other deps, which are not in this local envbut 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 crossingthe boundary, or a
frozen-like wrapper — can be a follow-up ifprofiling shows the cost is real.