Python: return caller-owned checkpoints from InMemoryCheckpointStorage reads

[Oleg Solozobov (dev404ai)]

Motivation & Context

InMemoryCheckpointStorage handed out the checkpoint objects it stores, so a caller that mutated a loaded checkpoint also mutated the storage. The other two implementations of the same CheckpointStorage protocol, FileCheckpointStorage and the Azure Cosmos backend, rebuild every returned checkpoint from its serialized form, so their reads are isolated by construction.

That made the in-memory implementation the odd one out on a protocol three backends share, and it is the one tests and samples reach for first, so a workflow that mutates restored state could behave differently depending on which backend was configured.

The protocol itself did not say which behavior was correct, which is how the two were able to drift apart.

Description & Review Guide

• What are the major changes?

  InMemoryCheckpointStorage now returns copy.deepcopy from its three read paths, load, list_checkpoints and get_latest, matching the deep copy that save already stored. Isolation on write was deliberate, and only the read side was missing.

  The CheckpointStorage protocol docstring now states the contract the fix restores: reads return objects owned by the caller, backends that serialize satisfy this by construction, and backends holding checkpoints in memory must copy on read.

• What is the impact of these changes?

  All three implementations now agree. Callers may treat a returned checkpoint as their own, which is what the file and Cosmos backends already allowed. No public signature changes, and no behavior changes for callers that do not mutate what they read.

  get_latest copies only the single checkpoint it returns rather than the whole candidate list, so listing cost is unchanged.

• What do you want reviewers to focus on?

  Whether the added protocol sentence states the intended contract. If the intended contract is instead that callers must copy defensively, then the file and Cosmos backends are the ones that diverge and this change should be reversed in favor of documenting that. I read the existing copy.deepcopy in save as evidence that caller isolation was already intended.

Related Issue

Fixes #7685

Contribution Checklist

• The code builds clean without any errors or warnings
• I have added tests that prove my fix is effective
• I have added the relevant documentation where applicable

Two tests were added to the workflow checkpoint suite:

• test_memory_checkpoint_storage_load_returns_caller_owned_copy — mutating a loaded checkpoint does not change what a later read returns.
• test_memory_checkpoint_storage_list_and_get_latest_return_caller_owned_copies — the same contract for the other two read paths.

Both fail without the fix. With it, uv run poe test over the workflow suite reports 926 passed, 2 skipped and 2 xfailed, while ruff check and ruff format --check are clean on both changed files.

[Tao Chen (TaoChenOSU)]

Closing in flavor of #7697 as it's a more comprehensive solution.