From d4ec021521eae7b6a5791742eccfe5b3b178f391 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Fri, 31 Jul 2026 14:16:37 +0530 Subject: [PATCH 01/10] 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 02/10] 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 03/10] 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(":") From dbb4eb824d7a1a24118d7641190ae398443e5c9b Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Mon, 3 Aug 2026 08:20:03 +0530 Subject: [PATCH 04/10] types: expose register_checkpoint_type in __init__.pyi typed stubs --- python/packages/core/agent_framework/__init__.pyi | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 31fa53a56f..8aee69ab5d 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -257,6 +257,7 @@ from ._workflows._checkpoint import ( InMemoryCheckpointStorage, WorkflowCheckpoint, ) +from ._workflows._checkpoint_encoding import register_checkpoint_type from ._workflows._const import DEFAULT_MAX_ITERATIONS, INTERNAL_SOURCE_ID from ._workflows._edge import ( Case, @@ -588,6 +589,7 @@ __all__ = [ "normalize_tools", "prepend_agent_framework_to_user_agent", "prepend_instructions_to_messages", + "register_checkpoint_type", "register_state_type", "resolve_agent_id", "response_handler", From ca102f73c969275d80358e522e0be5add3ca63ee Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Mon, 3 Aug 2026 15:59:03 +0530 Subject: [PATCH 05/10] docs: update checkpoint deserialization blocked error message to mention register_checkpoint_type --- .../core/agent_framework/_workflows/_checkpoint_encoding.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 705f99523a..fba6cc599a 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -212,9 +212,9 @@ def find_class(self, module: str, name: str) -> Any: raise pickle.UnpicklingError( f"Checkpoint deserialization blocked for type '{type_key}'. " - f"To allow this type, either include its 'module:qualname' key in the " - f"'allowed_types' set passed to 'decode_checkpoint_value', or add it to " - f"'allowed_checkpoint_types' on your checkpoint storage " + 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')." ) From 378b4a0df9a9501493149fd69375caeca9d291dd Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Mon, 3 Aug 2026 19:39:39 +0530 Subject: [PATCH 06/10] validation: reject custom checkpoint type keys containing multiple colons --- .../core/agent_framework/_workflows/_checkpoint_encoding.py | 2 ++ .../core/tests/workflow/test_checkpoint_unrestricted_pickle.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index fba6cc599a..7bfce19cc6 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -410,6 +410,8 @@ 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 cls_or_key.count(":") != 1: + raise ValueError("Type key must be in the format 'module:qualname'.") module, sep, qualname = cls_or_key.partition(":") module = module.strip() qualname = qualname.strip() 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 8d32e02a71..fb87745ffb 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -593,3 +593,6 @@ def test_register_checkpoint_type_validation(): with pytest.raises(ValueError, match="Type key must be in the format"): register_checkpoint_type(":") + + with pytest.raises(ValueError, match="Type key must be in the format"): + register_checkpoint_type("module:qualname:extra") From 6d8d531fcacce622c5cc2bce3aa9d7610f1b13e2 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Mon, 3 Aug 2026 19:50:34 +0530 Subject: [PATCH 07/10] docs: clarify global registration scope/security tradeoffs and improve error guidance --- .../agent_framework/_workflows/_checkpoint_encoding.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 7bfce19cc6..1d475ef03d 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -214,8 +214,8 @@ def find_class(self, module: str, name: str) -> Any: f"Checkpoint deserialization blocked for type '{type_key}'. " 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')." + f"or pass 'allowed_checkpoint_types' to your checkpoint storage " + f"(for example, 'FileCheckpointStorage(..., allowed_checkpoint_types=[...])')." ) @@ -401,7 +401,11 @@ def _value_type_to_key(value: object) -> str: def register_checkpoint_type(cls_or_key: type[Any] | str) -> None: - """Register a custom type to be allowed during checkpoint deserialization. + """Register a custom type globally to be allowed during checkpoint deserialization. + + .. warning:: + Registration modifies a process-wide allowlist used during restricted unpickling. + Only register trusted application-defined types. Each registered type should be either a type class (e.g., custom models or dataclasses) or a ``"module:qualname"`` string. From 33eff8d2e30070f583ff3e309dd50472d1e70431 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Mon, 3 Aug 2026 20:05:25 +0530 Subject: [PATCH 08/10] types: ignore pyright reportPrivateUsage warning for _CUSTOM_ALLOWED_TYPES import in tests --- .../core/tests/workflow/test_checkpoint_unrestricted_pickle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 fb87745ffb..22c751cf6b 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -35,7 +35,7 @@ @pytest.fixture(autouse=True) def cleanup_custom_checkpoint_types(): - from agent_framework._workflows._checkpoint_encoding import _CUSTOM_ALLOWED_TYPES + from agent_framework._workflows._checkpoint_encoding import _CUSTOM_ALLOWED_TYPES # pyright: ignore[reportPrivateUsage] backup = set(_CUSTOM_ALLOWED_TYPES) yield From c1998bf408f2d7e99d7b6563dada4b71c673ccc3 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Mon, 3 Aug 2026 20:14:01 +0530 Subject: [PATCH 09/10] docs: update _RestrictedUnpickler docstring to mention globally registered types --- .../core/agent_framework/_workflows/_checkpoint_encoding.py | 4 ++-- 1 file changed, 2 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 1d475ef03d..db9b23cb2a 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -139,8 +139,8 @@ class _RestrictedUnpickler(pickle.Unpickler): # ruff:ignore[suspicious-pickle-u Only classes whose ``module:qualname`` key appears in the combined allow set (built-in safe types + framework types + OpenAI SDK types + - caller-specified extras) are permitted. All other classes raise - :class:`pickle.UnpicklingError`. + caller-specified extras + globally registered custom types) are permitted. + All other classes raise :class:`pickle.UnpicklingError`. """ def __init__(self, data: bytes, allowed_types: frozenset[str]) -> None: From ab477c11cbc47dcc00a7e886515c84e4bdad00b3 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Mon, 3 Aug 2026 20:30:42 +0530 Subject: [PATCH 10/10] tests: isolate custom checkpoint types test fixture and format imports --- .../tests/workflow/test_checkpoint_unrestricted_pickle.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 22c751cf6b..9a7604e7ce 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -35,9 +35,12 @@ @pytest.fixture(autouse=True) def cleanup_custom_checkpoint_types(): - from agent_framework._workflows._checkpoint_encoding import _CUSTOM_ALLOWED_TYPES # pyright: ignore[reportPrivateUsage] + from agent_framework._workflows._checkpoint_encoding import ( + _CUSTOM_ALLOWED_TYPES, # pyright: ignore[reportPrivateUsage] + ) backup = set(_CUSTOM_ALLOWED_TYPES) + _CUSTOM_ALLOWED_TYPES.clear() yield _CUSTOM_ALLOWED_TYPES.clear() _CUSTOM_ALLOWED_TYPES.update(backup)