diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index dd1fb3d704..979389f49e 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -99,14 +99,24 @@ def __init__(self, data: bytes, allowed_types: frozenset[str]) -> None: def find_class(self, module: str, name: str) -> type: type_key = f"{module}:{name}" - if ( - type_key in _BUILTIN_ALLOWED_TYPE_KEYS - or type_key in self._allowed_types - or module.startswith(_FRAMEWORK_MODULE_PREFIX) - or module.startswith(_OPENAI_MODULE_PREFIX) - ): + if type_key in _BUILTIN_ALLOWED_TYPE_KEYS or type_key in self._allowed_types: return super().find_class(module, name) # type: ignore[no-any-return] # nosec + if module.startswith(_FRAMEWORK_MODULE_PREFIX) or module.startswith( + _OPENAI_MODULE_PREFIX + ): + resolved = super().find_class(module, name) # nosec + # Reject non-type globals from allowed package prefixes. A broad + # module-prefix allowlist returns arbitrary module attributes such + # as functions and sub-modules, which can be combined with + # builtins.getattr to call unrestricted pickle.loads. + if isinstance(resolved, type): + return resolved # type: ignore[return-value] + raise pickle.UnpicklingError( + f"Checkpoint deserialization blocked for non-type global " + f"'{type_key}'." + ) + raise pickle.UnpicklingError( f"Checkpoint deserialization blocked for type '{type_key}'. " f"To allow this type, either include its 'module:qualname' key in the "