[BREAKING] Python: restrict checkpoint deserialization in FoundryCheckpointStore - #8045
Conversation
`FileCheckpointStorage` and the Cosmos checkpoint storage both take an `allowed_checkpoint_types` argument, hold it as `self._allowed_types`, and pass it to `decode_checkpoint_value`. `FoundryCheckpointStore` had neither, and called the decoder with the argument omitted in `load` and again in `list_checkpoints`. An omitted `allowed_types` means no restriction and falls through to plain `pickle.loads`, while the empty frozenset the other two stores pass by default selects the restricted unpickler. So the Foundry store was the only one of the three not following the module's own guidance that the argument be specified whenever possible. Give it the same argument and pass it on, so a workflow that restores on one store restores on the others. This is a behaviour change for applications on this store whose checkpoints hold their own types and which never registered them, since those are relying on the store being more permissive than the other two. They register the types with `register_checkpoint_type` or pass `allowed_checkpoint_types`, which an application on the file or Cosmos store already has to do.
GitHub Actions
Active
GitHub Actions
Active
GitHub Actions
Active
agent-framework-automation
Bot
added
python
GitHub Actions
Active
There was a problem hiding this comment.
Would it make sense for CheckpointStoreProvider to accept and forward allowed_checkpoint_types here? ResponsesHostServer creates stores through this provider, so hosted apps cannot use the new per-store option on the normal path and must either register types process-wide or replace the whole provider. Threading the list through CheckpointStoreProvider.__init__ would keep the new configuration reachable without changing its behavior.
There was a problem hiding this comment.
Good catch — you're right, and it made the change worse than it needed to be. With the option only on the store, a hosted app on the default path had no way to reach it, so its only recourse for a checkpoint carrying its own type was process-wide registration or replacing the provider.
Pushed 9fad234: CheckpointStoreProvider.__init__ takes allowed_checkpoint_types and forwards it to every store it creates. Keyword-only with a None default, so CheckpointStoreProvider() behaves exactly as before.
Two tests added: one asserting a type named through the provider survives a load, and one asserting a provider built with no arguments still restricts. The first fails without the change with TypeError: CheckpointStoreProvider() takes no arguments, which is the gap you described.
`ResponsesHostServer` builds a `CheckpointStoreProvider` itself on the default path, so an option settable only on `FoundryCheckpointStore` was out of reach for a hosted app: it would have had to register types process-wide or replace the whole provider. Take the list on the provider and pass it to each store it creates. The default is unchanged, so a provider built with no arguments still restricts exactly as before.
GitHub Actions
Active
Motivation & Context
FileCheckpointStorageand the Cosmos checkpoint storage both take anallowed_checkpoint_typesargument, hold it asself._allowed_types, and pass it todecode_checkpoint_value.FoundryCheckpointStorehad neither, and called the decoder with the argument omitted inloadand again inlist_checkpoints.An omitted
allowed_typesmeans no restriction and falls through to plainpickle.loads, while the empty frozenset the other two stores pass by default selects the restricted unpickler. The Foundry store was therefore the only one of the three not following the guidance in_checkpoint_encoding's own security notes, which asks that the argument be specified whenever possible.Description & Review Guide
What are the major changes?
FoundryCheckpointStore.__init__takesallowed_checkpoint_types, keeps it asself._allowed_types, and both decode sites pass it. The argument, its docstring and the frozenset conversion mirrorFileCheckpointStorageso the three stores read the same way.What is the impact of these changes? A workflow that restores on one checkpoint store restores on the others. It is a behaviour change for applications on this store whose checkpoints hold their own types and which never registered them, since those are relying on this store being more permissive than the other two. They register the types with
register_checkpoint_type, or passallowed_checkpoint_types, which an application on the file or Cosmos store already has to do. I have marked this as a breaking change on that basis; say the word if you would rather it were staged behind an argument that defaults to the old behaviour.What do you want reviewers to focus on? Whether aligning with the other two stores is the outcome you want here, or whether the Foundry store is permissive on purpose.
Testing
Three tests added to
packages/foundry_hosting/tests/test_state_store.py: a checkpoint carrying a type outside the allow set is refused onload, the same throughlist_checkpoints, and a caller naming that type throughallowed_checkpoint_typesgets it back. All three fail without the change.The file's 29 tests pass.
ruff check,ruff format --checkandmypyare clean on the touched files.Related Issue
Fixes #8044
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.