Python: allow registering custom types for checkpoint serialization and deserialization

python/packages/core/agent_framework/__init__.py

@@ -291,6 +291,7 @@
         "InMemoryCheckpointStorage",
         "WorkflowCheckpoint",
     ),
+    "._workflows._checkpoint_encoding": ("register_checkpoint_type",),
     "._workflows._const": (
         "DEFAULT_MAX_ITERATIONS",
         "INTERNAL_SOURCE_ID",
@@ -624,6 +625,7 @@
     "normalize_tools",
     "prepend_agent_framework_to_user_agent",
     "prepend_instructions_to_messages",
+    "register_checkpoint_type",
     "register_state_type",
     "resolve_agent_id",
     "response_handler",

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

@@ -131,6 +131,8 @@
     "__builtin__:getattr",
 })
 
+_CUSTOM_ALLOWED_TYPES: set[str] = set()
+
 
 class _RestrictedUnpickler(pickle.Unpickler):  # ruff:ignore[suspicious-pickle-usage]
     """Unpickler that restricts which classes may be instantiated.
@@ -150,6 +152,7 @@ def _is_allowed_type(self, resolved: type) -> bool:
         return (
             type_key in _BUILTIN_ALLOWED_TYPE_KEYS
             or type_key in self._allowed_types
+            or type_key in _CUSTOM_ALLOWED_TYPES
             or resolved.__module__.startswith(_FRAMEWORK_MODULE_PREFIX)
             or resolved.__module__.startswith(_OPENAI_MODULE_PREFIX)
         )
@@ -191,7 +194,7 @@ def find_class(self, module: str, name: str) -> Any:
         if type_key in _BUILTIN_ALLOWED_TYPE_KEYS:
             return super().find_class(module, name)  # nosec
 
-        if type_key in self._allowed_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
@@ -395,3 +398,25 @@ def _type_to_key(t: type[Any]) -> str:
 def _value_type_to_key(value: object) -> str:
     """Convert a value's type to a module:qualname string."""
     return _type_to_key(type(value))
+
+
+def register_checkpoint_type(cls_or_key: type[Any] | str) -> None:
+    """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.
+    """
+    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:
+            raise ValueError("Type key must be in the format 'module:qualname'.")
+        _CUSTOM_ALLOWED_TYPES.add(f"{module}:{qualname}")
+    elif isinstance(cls_or_key, type):
+        _CUSTOM_ALLOWED_TYPES.add(_type_to_key(cls_or_key))
+    else:
+        raise TypeError("Expected a type class or a 'module:qualname' string.")

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

@@ -33,6 +33,16 @@
 )
 
 
+@pytest.fixture(autouse=True)
+def cleanup_custom_checkpoint_types():
+    from agent_framework._workflows._checkpoint_encoding import _CUSTOM_ALLOWED_TYPES
+
+    backup = set(_CUSTOM_ALLOWED_TYPES)
+    yield
+    _CUSTOM_ALLOWED_TYPES.clear()
+    _CUSTOM_ALLOWED_TYPES.update(backup)
+
+
 class MaliciousPayload:
     """A class whose __reduce__ executes code during unpickling."""
 
@@ -509,3 +519,77 @@ def test_restricted_decode_allows_openai_response_types():
     assert isinstance(decoded, ResponseUsage)
     assert decoded.input_tokens == 10
     assert decoded.output_tokens == 20
+
+
+@dataclass
+class _GloballyAllowedTestState:
+    """Test state to verify global registration works."""
+
+    value: int
+
+
+@dataclass
+class _GloballyAllowedStringState:
+    """Test state to verify global registration by string key works."""
+
+    value: str
+
+
+def test_register_checkpoint_type_allows_type():
+    """Custom types registered globally are allowed during restricted deserialization."""
+    from agent_framework import register_checkpoint_type
+
+    # Verify that before registration, the type is blocked
+    original = _GloballyAllowedTestState(value=100)
+    encoded = encode_checkpoint_value(original)
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        decode_checkpoint_value(encoded, allowed_types=frozenset())
+
+    # Register the type
+    register_checkpoint_type(_GloballyAllowedTestState)
+
+    # Now it should be allowed
+    decoded = decode_checkpoint_value(encoded, allowed_types=frozenset())
+    assert isinstance(decoded, _GloballyAllowedTestState)
+    assert decoded.value == 100
+
+
+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")
+    encoded = encode_checkpoint_value(original)
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        decode_checkpoint_value(encoded, allowed_types=frozenset())
+
+    # Register by string key
+    type_key = f"{_GloballyAllowedStringState.__module__}:{_GloballyAllowedStringState.__qualname__}"
+    register_checkpoint_type(type_key)
+
+    # Now it should be allowed
+    decoded = decode_checkpoint_value(encoded, allowed_types=frozenset())
+    assert isinstance(decoded, _GloballyAllowedStringState)
+    assert decoded.value == "globally_allowed"
+
+
+def test_register_checkpoint_type_validation():
+    """register_checkpoint_type validates inputs properly."""
+    from agent_framework import register_checkpoint_type
+
+    with pytest.raises(TypeError, match="Expected a type class or a 'module:qualname' string"):
+        register_checkpoint_type(123)  # type: ignore
+
+    with pytest.raises(ValueError, match="Type key must be in the format"):
+        register_checkpoint_type("invalid_key")
+
+    with pytest.raises(ValueError, match="Type key must be in the format"):
+        register_checkpoint_type("module:")
+
+    with pytest.raises(ValueError, match="Type key must be in the format"):
+        register_checkpoint_type(":qualname")
+
+    with pytest.raises(ValueError, match="Type key must be in the format"):
+        register_checkpoint_type(":")