diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 2b267979e9..47293fe2aa 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -255,9 +255,9 @@ class FileCheckpointStorage: for human-readable checkpoint files while preserving the ability to store complex Python objects. By default, checkpoint deserialization is restricted to a built-in set of safe Python types - (primitives, datetime, uuid, ...), all ``agent_framework`` internal types, and OpenAI SDK types - (``openai.types``). To allow additional application-specific types, pass them via the - ``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format. + (primitives, datetime, uuid, ...), all ``agent_framework`` and ``agent_framework_orchestrations`` + internal types, and OpenAI SDK types (``openai.types``). To allow additional application-specific + types, pass them via the ``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format. Example:: diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 0bec7a970d..5c668596cf 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -10,9 +10,10 @@ When ``allowed_types`` is supplied to :func:`decode_checkpoint_value`, a ``RestrictedUnpickler`` is used that limits which classes may be instantiated during deserialization. The default built-in safe set covers common Python -value types (primitives, datetime, uuid, ...), all ``agent_framework`` internal -types, and all ``openai.types`` types. Callers can extend the set by passing -additional ``"module:qualname"`` strings. +value types (primitives, datetime, uuid, ...), all ``agent_framework`` and +``agent_framework_orchestrations`` internal types, and all ``openai.types`` +types. Callers can extend the set by passing additional ``"module:qualname"`` +strings. Security Model -------------- @@ -68,8 +69,11 @@ # Types that are natively JSON-serializable and don't need pickling _JSON_NATIVE_TYPES = (str, int, float, bool, type(None)) -# Module prefix for framework-internal types that are always allowed -_FRAMEWORK_MODULE_PREFIX = "agent_framework." +# Module prefixes for framework-internal types that are always allowed +_FRAMEWORK_MODULE_PREFIXES = ( + "agent_framework.", + "agent_framework_orchestrations.", +) # Module prefix for OpenAI SDK types that are always allowed _OPENAI_MODULE_PREFIX = "openai.types." @@ -150,7 +154,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 resolved.__module__.startswith(_FRAMEWORK_MODULE_PREFIX) + or resolved.__module__.startswith(_FRAMEWORK_MODULE_PREFIXES) or resolved.__module__.startswith(_OPENAI_MODULE_PREFIX) ) @@ -197,7 +201,7 @@ def find_class(self, module: str, name: str) -> Any: return resolved raise pickle.UnpicklingError(f"Checkpoint deserialization blocked for non-type global '{type_key}'.") - if module.startswith(_FRAMEWORK_MODULE_PREFIX) or module.startswith(_OPENAI_MODULE_PREFIX): + if module.startswith(_FRAMEWORK_MODULE_PREFIXES) or module.startswith(_OPENAI_MODULE_PREFIX): # Pickle dotted names traverse attributes on an allowed module; keep the prefix allowlist to concrete # top-level classes rather than helper callables reachable through module attributes. if "." in name: diff --git a/python/packages/core/tests/workflow/test_checkpoint_orchestrations_deserialization.py b/python/packages/core/tests/workflow/test_checkpoint_orchestrations_deserialization.py new file mode 100644 index 0000000000..490cd08aee --- /dev/null +++ b/python/packages/core/tests/workflow/test_checkpoint_orchestrations_deserialization.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft. All rights reserved. + +import pytest + +pytest.importorskip("agent_framework_orchestrations") + +from agent_framework import AgentResponse, Message +from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value +from agent_framework.orchestrations import HandoffAgentUserRequest, MagenticPlanReviewRequest + + +@pytest.mark.parametrize( + ("request_value", "request_type"), + [ + ( + HandoffAgentUserRequest( + agent_response=AgentResponse(messages=[Message("assistant", ["handoff response"])]) + ), + HandoffAgentUserRequest, + ), + ( + MagenticPlanReviewRequest( + plan=Message("assistant", ["review this plan"]), + current_progress=None, + is_stalled=False, + ), + MagenticPlanReviewRequest, + ), + ], +) +def test_restricted_decode_roundtrips_orchestration_requests(request_value: object, request_type: type[object]) -> None: + """Pending orchestration requests can be restored from a restricted checkpoint.""" + encoded = encode_checkpoint_value(request_value) + + decoded = decode_checkpoint_value(encoded, allowed_types=frozenset()) + + assert isinstance(decoded, request_type)