Python: Add allowed_checkpoint_types support to CosmosCheckpointStorage for parity with FileCheckpointStorage#5202
Open
moonbox3 wants to merge 3 commits intomicrosoft:mainfrom
Open
Python: Add allowed_checkpoint_types support to CosmosCheckpointStorage for parity with FileCheckpointStorage#5202moonbox3 wants to merge 3 commits intomicrosoft:mainfrom
allowed_checkpoint_types support to CosmosCheckpointStorage for parity with FileCheckpointStorage#5202moonbox3 wants to merge 3 commits intomicrosoft:mainfrom
Conversation
…ge (microsoft#5200) Add allowed_checkpoint_types parameter to CosmosCheckpointStorage for parity with FileCheckpointStorage. This ensures both providers use the same restricted pickle deserialization by default. Changes: - Accept allowed_checkpoint_types kwarg in __init__, stored as frozenset - Convert _document_to_checkpoint from @staticmethod to instance method - Forward allowed_types to decode_checkpoint_value on all load paths - Update class docstring to describe the new parameter - Add tests covering built-in safe types, app type opt-in/blocking, and all load paths (load, list_checkpoints, get_latest) - Add changelog entry noting the breaking behavior change BREAKING CHANGE: CosmosCheckpointStorage now uses restricted pickle deserialization by default. Checkpoints containing application-defined types will require passing those types via allowed_checkpoint_types. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…orage` for parity with `FileCheckpointStorage` Fixes microsoft#5200
moonbox3
commented
Apr 10, 2026
Contributor
Author
moonbox3
left a comment
moonbox3
left a comment
There was a problem hiding this comment.
Automated Code Review
Reviewers: 4 | Confidence: 98% | Result: All clear
Reviewed: Correctness, Security Reliability, Test Coverage, Design Approach
Automated review by moonbox3's agents
Contributor
Author
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||
Contributor
There was a problem hiding this comment.
Pull request overview
This PR brings CosmosCheckpointStorage deserialization behavior into parity with FileCheckpointStorage by always using restricted pickle loading (with optional opt-in for application-defined types), and adds tests plus a changelog entry to document the breaking behavior change.
Changes:
- Add
allowed_checkpoint_typestoCosmosCheckpointStorage.__init__and flow it intodecode_checkpoint_value(...)for restricted deserialization on all load paths. - Add azure-cosmos unit tests covering built-in safe types, blocked custom types, and opt-in custom types across
load,list_checkpoints, andget_latest. - Add a changelog entry noting the breaking deserialization behavior change for the azure-cosmos package.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py | Adds allowed_checkpoint_types plumbing and updates class docstring to describe restricted deserialization behavior. |
| python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py | Adds test coverage for restricted deserialization and opt-in application types across Cosmos load paths. |
| python/CHANGELOG.md | Documents the breaking behavior change for CosmosCheckpointStorage. |
python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py
Show resolved
Hide resolved
python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py
Outdated
Show resolved
Hide resolved
…ples - Reintroduce explicit security warning about pickle deserialization risks - Convert Example:: block to .. code-block:: python with imports for consistency with other docstring examples - Note: PR title should be updated to include [BREAKING] prefix per changelog convention (comment #3, requires GitHub UI change) Fixes microsoft#5200 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
moonbox3
commented
Apr 10, 2026
Contributor
Author
moonbox3
left a comment
moonbox3
left a comment
There was a problem hiding this comment.
Automated Code Review
Reviewers: 4 | Confidence: 96% | Result: All clear
Reviewed: Correctness, Security Reliability, Test Coverage, Design Approach
Automated review by moonbox3's agents
eavanvalkenburg
approved these changes
Apr 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation and Context
CosmosCheckpointStoragecalleddecode_checkpoint_valuewithout anallowed_typesargument, causing it to use unrestrictedpickle.loadsinstead of theRestrictedUnpicklerthatFileCheckpointStoragealways uses. This meant the two providers—which implement the same protocol and share the same encoding—silently took different deserialization paths, making them non-interchangeable and leaving the Cosmos path without the restricted-type safety floor.Fixes #5200
Description
The fix adds an
allowed_checkpoint_types: list[str] | None = Nonekeyword argument toCosmosCheckpointStorage.__init__, stored as afrozensetidentical to theFileCheckpointStoragepattern._document_to_checkpointis changed from a@staticmethodto an instance method so it can forwardallowed_types=self._allowed_typesintodecode_checkpoint_value, ensuring all load paths (load,list_checkpoints,get_latest) use restricted deserialization. New tests cover built-in safe types loading without opt-in, application-defined types being blocked unless explicitly listed, and all three load paths. A changelog entry notes this as a breaking change for users who store custom types.Contribution Checklist