Python: allow registering custom types for checkpoint serialization and deserialization

[Mahajan-Sachin]

Motivation & Context

In hosting environments like Foundry, developers don't manage the checkpoint store directly. If custom types are used in their workflows, checkpoint serialization/deserialization will fail because the restricted unpickler rejects any unknown type for security reasons. Currently, there is no way to extend allowed_checkpoint_types without direct access to the checkpoint storage configuration.

This PR introduces a process-wide global registry for checkpoint serialization, similar to the existing register_state_type utility used in agent sessions, allowing developers to register custom classes or string type keys.

Description & Review Guide

What are the major changes?

• Created a global _CUSTOM_ALLOWED_TYPES registry set in packages/core/agent_framework/_workflows/_checkpoint_encoding.py.
• Added the public register_checkpoint_type(cls_or_key: type[Any] | str) function to allow developers to register classes or string keys.
• Updated _RestrictedUnpickler.find_class and _is_allowed_type to allow types registered in _CUSTOM_ALLOWED_TYPES.
• Exposed register_checkpoint_type under the core package root __init__.py.
• Added unit tests in packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py verifying registration behavior, string key registration, and input validation.

What is the impact of these changes?
Developers can now call register_checkpoint_type(MyCustomClass) at module import time. This allows their custom types to pass through checkpoint serialization in hosting environments without needing direct access to storage initializers.

What do you want reviewers to focus on?

• The integration of _CUSTOM_ALLOWED_TYPES into both find_class and _is_allowed_type in the restricted unpickler.
• The lazy-loading export structure in __init__.py.

Related Issue

Fixes #7413

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue.

[Copilot]

Pull request overview

Adds a process-wide registry to extend the restricted checkpoint unpickler’s allowlist so applications in hosted environments can permit custom workflow state types without modifying checkpoint storage configuration.

Changes:

• Introduces a global _CUSTOM_ALLOWED_TYPES registry and a public register_checkpoint_type(...) API in checkpoint encoding.
• Extends restricted unpickling checks to permit globally registered types.
• Exposes register_checkpoint_type from the package root and adds unit tests covering registration and validation.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py	Adds global allowlist + registration API and integrates it into restricted unpickling.
python/packages/core/agent_framework/init.py	Re-exports register_checkpoint_type via the lazy export mechanism and __all__.
python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py	Adds tests validating that global registration allows restricted deserialization.

Suppressed comments (1)

python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py:552

• This test also mutates the process-wide checkpoint type registry and should restore it to avoid leaking global state to other tests.

def test_register_checkpoint_type_allows_string_key():
    """Custom types registered globally by string key are allowed during restricted deserialization."""
    from agent_framework import register_checkpoint_type

    original = _GloballyAllowedStringState(value="globally_allowed")

[Mahajan-Sachin]

@copilot please review

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py:201

• The fallback UnpicklingError guidance is now incomplete: with global registration supported, users should be told they can also allow a type via agent_framework.register_checkpoint_type(...) (in addition to allowed_types / allowed_checkpoint_types).

        if type_key in self._allowed_types or type_key in _CUSTOM_ALLOWED_TYPES:
            resolved = super().find_class(module, name)  # nosec
            if isinstance(resolved, type):
                return resolved
            raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for non-type global '{type_key}'.")

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py:418

• register_checkpoint_type accepts strings with multiple ':' characters (e.g. "pkg.mod:Type:Extra") because it uses partition(":") and doesn't validate the remainder. This will silently register a key that can never match the module:name keys produced by pickle, making registration ineffective and hard to debug. Consider rejecting keys that don't contain exactly one ':' before stripping/normalizing.

    if isinstance(cls_or_key, str):
        module, sep, qualname = cls_or_key.partition(":")
        module = module.strip()
        qualname = qualname.strip()
        if not sep or not module or not qualname:

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (2)

python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py:218

• The unblock guidance mentions FileCheckpointStorage.allowed_checkpoint_types, but FileCheckpointStorage doesn’t expose an allowed_checkpoint_types attribute (it’s a constructor parameter stored internally as _allowed_types). This example is misleading; suggest pointing to passing allowed_checkpoint_types to the constructor instead.

            f"To allow this type, register it globally via 'agent_framework.register_checkpoint_type', "
            f"include its 'module:qualname' key in the 'allowed_types' set passed to 'decode_checkpoint_value', "
            f"or add it to 'allowed_checkpoint_types' on your checkpoint storage "
            f"(for example, 'FileCheckpointStorage.allowed_checkpoint_types')."

python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py:411

• register_checkpoint_type mutates a process-wide global allowlist, but the docstring doesn’t call out the global/process-wide scope or the security implications (expands what the restricted unpickler will instantiate). Clarifying this helps callers understand when to call it and the risk tradeoff.

    """Register a custom type to be allowed during checkpoint deserialization.

    Each registered type should be either a type class (e.g., custom models or
    dataclasses) or a ``"module:qualname"`` string.

    Args:
        cls_or_key: The type class or module-qualified string to register.
    """

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py:38

• Importing the private symbol _CUSTOM_ALLOWED_TYPES from agent_framework._workflows._checkpoint_encoding will likely trigger pyright's reportPrivateUsage (this file already suppresses that check for _base64_to_unpickle). Add an inline pyright ignore to avoid type-check failures in CI.

    from agent_framework._workflows._checkpoint_encoding import _CUSTOM_ALLOWED_TYPES

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py:138

• The _RestrictedUnpickler summary docstring no longer reflects behavior now that a process-wide custom-type registry is part of the allowlist. Consider mentioning the global registration mechanism in the docstring summary so readers don’t miss that types can be permitted outside allowed_types.

    """Unpickler that restricts which classes may be instantiated.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py:43

• This fixture restores the global registry after the test, but it never clears it before the test runs. That makes the new tests depend on _CUSTOM_ALLOWED_TYPES already being empty (or at least not containing these types) when the test starts, which can become order-dependent if any other test registers checkpoint types globally.

    backup = set(_CUSTOM_ALLOWED_TYPES)
    yield
    _CUSTOM_ALLOWED_TYPES.clear()
    _CUSTOM_ALLOWED_TYPES.update(backup)

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

[Mahajan-Sachin]

Hi @moonbox3 and @TaoChenOSU,

I've updated the PR to address all review feedback and edge cases:

1. Added Typed Stubs: Imported and exposed register_checkpoint_type in agent_framework/__init__.pyi to ensure IDEs and static type checkers (Pyright/Mypy) can resolve the symbol correctly.
2. Improved Key Validation: Added strict validation in register_checkpoint_type to ensure string keys contain exactly one colon separator (rejecting partial formats or values with multiple colons like module:qualname:extra). Added test coverage for these cases.
3. Isolated Test State: Updated the test fixture to clear the global _CUSTOM_ALLOWED_TYPES registry both before and after each test runs to guarantee strict state isolation and prevent any potential order-dependency.
4. Enhanced Error Messages: Corrected the fallback UnpicklingError message format to point to the constructor parameters and added the global registration mechanism to the guide text.
5. Addressed Docstring Warnings: Clarified the global process-wide scope and security tradeoffs in both the registry function and _RestrictedUnpickler docstrings.

All unit tests and type-checks are passing cleanly. Ready for your review!