diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index ff0aa20e1a..9380a81bb8 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/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", diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 2c601a48d0..705f99523a 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/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.") diff --git a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py index 15e379bd7a..8d32e02a71 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/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(":")