Files
agent-framework/python/packages
T
Evan Mattson 4dbe696e0e Python: Restrict persisted checkpoint deserialization by default (#4941)
* Harden Python checkpoint persistence defaults

Add RestrictedUnpickler to _checkpoint_encoding.py that limits which
types may be instantiated during pickle deserialization.  By default
FileCheckpointStorage now uses the restricted unpickler, allowing only:

- Built-in Python value types (primitives, datetime, uuid, decimal,
  collections, etc.)
- All agent_framework.* internal types
- Additional types specified via the new allowed_checkpoint_types
  parameter on FileCheckpointStorage

This narrows the default type surface area for persisted checkpoints
while keeping framework-owned scenarios working without extra
configuration.  Developers can extend the allowed set by passing
"module:qualname" strings to allowed_checkpoint_types.

The decode_checkpoint_value function retains backward-compatible
unrestricted behavior when called without the new allowed_types kwarg.

Fixes #4894

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: resolve mypy no-any-return error in checkpoint encoding

Add explicit type annotation for super().find_class() return value
to satisfy mypy's no-any-return check.

Fixes #4894

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Simplify find_class return in _RestrictedUnpickler (#4894)

Remove unnecessary intermediate variable and apply # noqa: S301 # nosec
directly on the super().find_class() call, matching the established
pattern used on the pickle.loads() call in the same file.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback for #4894: Python: Harden Python checkpoint persistence defaults

* Restore # noqa: S301 on line 102 of _checkpoint_encoding.py (#4894)

The review feedback correctly identified that removing the # noqa: S301
suppression from the find_class return statement would cause a ruff S301
lint failure, since the project enables bandit ("S") rules. This
restores consistency with lines 82 and 246 in the same file.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback for #4894: Python: Harden Python checkpoint persistence defaults

* Address PR review comments on checkpoint encoding (#4894)

- Move module docstring to proper position after __future__ import
- Fix find_class return type annotation to type[Any]
- Add missing # noqa: S301 pragma on find_class return
- Improve error message to reference both allowed_types param and
  FileCheckpointStorage.allowed_checkpoint_types
- Add -> None return annotation to FileCheckpointStorage.__init__
- Replace tempfile.mktemp with TemporaryDirectory in test
- Replace contextlib.suppress with pytest.raises for precise assertion
- Remove unused contextlib import

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR #4941 review comments: fix docstring position and return type

- Move module docstring before 'from __future__' import so it populates
  __doc__ (comment #4)
- Change find_class return annotation from type[Any] to type to avoid
  misleading callers about non-type returns like copyreg._reconstructor
  (comment #2)

Comments #1, #3, #5, #6, #7, #8 were already addressed in the current code.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback for #4894: review comment fixes

* fix: use pickle.UnpicklingError in RestrictedUnpickler and improve docstring (#4894)

- Change _RestrictedUnpickler.find_class to raise pickle.UnpicklingError
  instead of WorkflowCheckpointException, since it is pickle-level concern
  that gets wrapped by the caller in _base64_to_unpickle.
- Remove now-unnecessary WorkflowCheckpointException re-raise in
  _base64_to_unpickle (pickle.UnpicklingError is caught by the generic
  except Exception handler and wrapped).
- Expand decode_checkpoint_value docstring to show a concrete example of
  the module:qualname format with a user-defined class.
- Add regression test verifying find_class raises pickle.UnpicklingError.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: address PR #4941 review comments for checkpoint encoding

- Comment 1 (line 103): Already resolved in prior commit — _RestrictedUnpickler
  now raises pickle.UnpicklingError instead of WorkflowCheckpointException.

- Comment 2 (line 140): Add concrete usage examples to decode_checkpoint_value
  docstring showing both direct allowed_types usage and FileCheckpointStorage
  allowed_checkpoint_types usage. Rename 'SafeState' to 'MyState' across all
  docstrings for consistency, making it clear this is a user-defined class name.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: replace deprecated 'builtin' repo with pre-commit-hooks in pre-commit config

pre-commit 4.x no longer supports 'repo: builtin'. Merge those hooks into
the existing pre-commit-hooks repo entry.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* style: apply pyupgrade formatting to docstring example

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: resolve pre-commit hook paths for monorepo git root

The poe-check and bandit hooks referenced paths relative to python/
but pre-commit runs hooks from the git root (monorepo root). Fix
poe-check entry to cd into python/ first, and update bandit config
path to python/pyproject.toml.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix pre-commit config paths for prek --cd python execution

Revert bandit config path from 'python/pyproject.toml' to 'pyproject.toml'
and poe-check entry from explicit 'cd python' wrapper to direct invocation,
since prek --cd python already sets the working directory to python/.

Also apply ruff formatting fixes to cosmos checkpoint storage files.

Fixes #4894

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: add builtins:getattr to checkpoint deserialization allowlist

Pickle uses builtins:getattr to reconstruct enum members (e.g.,
WorkflowMessage.type which is a MessageType enum). Without it in the
allowlist, checkpoint roundtrip tests fail with
WorkflowCheckpointException.

Fixes #4894

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback for #4894: review comment fixes

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
4dbe696e0e · 2026-04-10 00:04:17 +00:00
History
..