Python: Return caller-owned checkpoints from InMemoryCheckpointStorage - #7712
Python: Return caller-owned checkpoints from InMemoryCheckpointStorage#7712Shivani . (Shivani767) wants to merge 2 commits into
Conversation
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds an explicit “ownership” contract for CheckpointStorage backends and enforces it for the in-memory backend via defensive copying, with shared conformance tests to prevent regressions.
Changes:
- Documented the ownership/snapshot semantics in the
CheckpointStorageprotocol docstring. - Updated
InMemoryCheckpointStorageto return deep-copied checkpoints from read APIs. - Added conformance tests to validate backend behavior for both memory and file storages.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| python/packages/core/tests/workflow/test_checkpoint.py | Adds shared conformance tests asserting caller-owned copies and save-time snapshotting across backends. |
| python/packages/core/agent_framework/_workflows/_checkpoint.py | Documents ownership semantics and updates in-memory backend reads to copy.deepcopy to satisfy the contract. |
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
| 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] |
There was a problem hiding this comment.
Leaving this as a one-liner — ruff format (line-length 120) wants it on a single line. Wrapping it is what fails the formatter; that's also the likely cause of the code-quality check on #7697, which wrapped the same comprehension.
| # region checkpoint storage conformance | ||
|
|
||
| # These tests define the ownership contract that every CheckpointStorage backend must satisfy: | ||
| # 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. |
There was a problem hiding this comment.
Done — the conformance suite now lives in python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py. Per Tao, the same file is suggested on #7697 so the bug can land in one PR.
| def _conformance_checkpoint(workflow_name: str = "conformance-workflow") -> WorkflowCheckpoint: | ||
| """Build a checkpoint whose state 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"]}}, | ||
| }, | ||
| metadata={"tags": ["initial"]}, | ||
| ) |
There was a problem hiding this comment.
Done in the dedicated module. _conformance_checkpoint now includes nested mutable messages (WorkflowMessage.data["tags"]) and pending_request_info_events (WorkflowEvent.data["payload"]), and every ownership/snapshot test mutates those fields as well as state and metadata.
|
Thanks — I see #7685 was closed as a duplicate of #7683 after I'd started, and that #7697 already Happy to close this if it's just noise. Before I do, one thing here might be worth keeping The related point is that Would you like me to:
Happy to go whichever way is least disruptive. |
|
Thank you for your contributions and suggestions! I'd say let's do #2 so that we can address the bug in one PR. |
GitHub Actions
Active
|
Tao Chen (@TaoChenOSU) thanks — moved the conformance suite over to #7697 as a review suggestion: #7697 (review) Summary of what I suggested there:
I also applied the dedicated-module split on this PR so the file is easy to copy: Happy to close this once those tests are in #7697. |
Keep the ownership contract tests out of the already-large test_checkpoint.py and cover nested messages and pending request events as well as state/metadata.
efa9b8f to
c0585ff
Compare
GitHub Actions
Active
Motivation & Context
InMemoryCheckpointStoragehands back the checkpoint objects it stores. The other backendsreconstruct a checkpoint from its serialized form on every read, so their callers get an object
they own. In-memory callers instead get a shared reference into stored state.
This is easy to miss because in-memory is the backend most workflows are developed and tested
against. Mutating a loaded checkpoint silently rewrites what is stored, and two workflow instances
restored from one checkpoint share mutable state. The same code then behaves differently once a
real backend is configured.
save()already deep-copies, so the asymmetry is on the read side only.Description & Review Guide
What are the major changes?
load,list_checkpointsandget_latestonInMemoryCheckpointStoragereturn deep copies,matching the copy
savealready performs.CheckpointStorageprotocol now documents its ownership contract. It previously saidnothing about whether a returned checkpoint is owned by the caller, so there was no stated
contract for a backend to diverge from.
test_checkpoint.py— five tests coveringboth read and write ownership, run against each in-tree backend.
What is the impact of these changes?
FileCheckpointStorageor the Cosmos backend. They already satisfiedthe contract, and the new tests pass against them unmodified.
mainthe four read-path tests fail for in-memory and pass for file. The write-pathtest passes unfixed, which is what confirms
savewas already correct.benchmarked it and am not claiming a figure — happy to measure if you want it quantified.
What do you want reviewers to focus on?
test_checkpoint.py, so itonly covers the two in-tree backends.
CosmosCheckpointStoragelives inpackages/azure-cosmosand is tested separately. Making the suite importable across packages is a larger structural
change and I did not want to make that call unasked.
Related Issue
Fixes #7685
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.