From d4ec021521eae7b6a5791742eccfe5b3b178f391 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Fri, 31 Jul 2026 14:16:37 +0530 Subject: [PATCH 1/3] Python: allow registering custom types for checkpoint serialization --- .../packages/core/agent_framework/__init__.py | 2 + .../_workflows/_checkpoint_encoding.py | 24 ++++++- .../test_checkpoint_unrestricted_pickle.py | 65 +++++++++++++++++++ 3 files changed, 90 insertions(+), 1 deletion(-) 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..a6fabc1394 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,22 @@ 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): + if not cls_or_key or ":" not in cls_or_key: + raise ValueError("Type key must be in the format 'module:qualname'.") + _CUSTOM_ALLOWED_TYPES.add(cls_or_key) + 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..f736b785f7 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -509,3 +509,68 @@ 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") From f1b545dd6d6b796ad86d1483bf25992a5196b930 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan <152908416+Mahajan-Sachin@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:26:35 +0530 Subject: [PATCH 2/3] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../agent_framework/_workflows/_checkpoint_encoding.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index a6fabc1394..705f99523a 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -410,9 +410,12 @@ def register_checkpoint_type(cls_or_key: type[Any] | str) -> None: cls_or_key: The type class or module-qualified string to register. """ if isinstance(cls_or_key, str): - if not cls_or_key or ":" not in cls_or_key: + 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(cls_or_key) + _CUSTOM_ALLOWED_TYPES.add(f"{module}:{qualname}") elif isinstance(cls_or_key, type): _CUSTOM_ALLOWED_TYPES.add(_type_to_key(cls_or_key)) else: From 2e2a73965c66f77351c365b36474787f3211d30d Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Fri, 31 Jul 2026 14:30:37 +0530 Subject: [PATCH 3/3] tests: reset custom checkpoint types registry and extend validation tests --- .../test_checkpoint_unrestricted_pickle.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) 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 f736b785f7..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.""" @@ -574,3 +584,12 @@ def test_register_checkpoint_type_validation(): 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(":")