1 Commits

• 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>

  Evan Mattson · 2026-04-10 00:04:17 +00:00
  4dbe696e0e