diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 24dc9343659..e441975600d 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py @@ -54,8 +54,9 @@ class CosmosCheckpointStorage: By default, checkpoint deserialization is restricted to a built-in set of safe Python types (primitives, datetime, uuid, ...) and all ``agent_framework`` - internal types. To allow additional application-specific types, pass them via - the ``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format. + internal types. To allow additional application-specific types, register them + with ``register_checkpoint_type`` or pass them via the + ``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format. Example: diff --git a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py index 016220e6938..0fe47d5a5da 100644 --- a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py +++ b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py @@ -609,6 +609,13 @@ class _AppState: count: int +@dataclass +class _GloballyRegisteredAppState: + """Application-defined state type registered for all checkpoint backends.""" + + label: str + + _APP_STATE_TYPE_KEY = f"{_AppState.__module__}:{_AppState.__qualname__}" @@ -679,6 +686,21 @@ async def test_load_allows_listed_app_type(mock_container: MagicMock) -> None: assert loaded.state["data"].count == 7 +async def test_load_allows_globally_registered_app_type(mock_container: MagicMock) -> None: + """Registered application types load without configuring the Cosmos storage instance.""" + from agent_framework import register_checkpoint_type + + checkpoint = _make_checkpoint_with_state({"data": _GloballyRegisteredAppState(label="registered")}) + doc = _checkpoint_to_cosmos_document(checkpoint) + mock_container.query_items.return_value = _to_async_iter([doc]) + + register_checkpoint_type(_GloballyRegisteredAppState) + storage = CosmosCheckpointStorage(container_client=mock_container) + loaded = await storage.load(checkpoint.checkpoint_id) + + assert loaded.state["data"] == _GloballyRegisteredAppState(label="registered") + + async def test_list_checkpoints_blocks_unlisted_app_type(mock_container: MagicMock) -> None: """list_checkpoints skips documents with unlisted application types.""" checkpoint = _make_checkpoint_with_state({"data": _AppState(label="x", count=1)}) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 46e0d928572..caf1d87df91 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -293,6 +293,7 @@ "InMemoryCheckpointStorage", "WorkflowCheckpoint", ), + "._workflows._checkpoint_encoding": ("register_checkpoint_type",), "._workflows._const": ( "DEFAULT_MAX_ITERATIONS", "INTERNAL_SOURCE_ID", @@ -631,6 +632,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/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 4f4f10fb0ca..a382a6e5c87 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -259,6 +259,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, @@ -595,6 +596,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", diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 2b267979e99..3de9460c86d 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -256,8 +256,9 @@ class FileCheckpointStorage: 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. + (``openai.types``). To allow additional application-specific types, register them with + ``agent_framework.register_checkpoint_type`` or 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 0bec7a970d1..3a7d361b612 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -57,6 +57,28 @@ logger = logging.getLogger("agent_framework") +# Application-defined types registered for all restricted checkpoint decoders. +_REGISTERED_CHECKPOINT_TYPE_KEYS: set[str] = set() + + +def register_checkpoint_type(cls: type[Any]) -> None: + """Register an application type for restricted checkpoint deserialization. + + Registration applies process-wide to all checkpoint storage backends that + use :func:`decode_checkpoint_value` with a restricted allowlist, including + instances created before this function is called. + + Args: + cls: The application type to permit during checkpoint deserialization. + + Raises: + TypeError: If ``cls`` is not a class. + """ + if not isinstance(cls, type): + raise TypeError("Checkpoint types must be classes.") + _REGISTERED_CHECKPOINT_TYPE_KEYS.add(_type_to_key(cls)) + + # Marker to identify pickled values in serialized JSON _PICKLE_MARKER = "__pickled__" _TYPE_MARKER = "__type__" @@ -277,6 +299,8 @@ class MyState: ... data is malformed, or if a disallowed type is encountered during restricted deserialization. """ + if allowed_types is not None: + allowed_types = allowed_types | _REGISTERED_CHECKPOINT_TYPE_KEYS return _decode(value, allowed_types=allowed_types) 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 15e379bd7af..4e89dd1b2cc 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -22,7 +22,7 @@ import pytest -from agent_framework import WorkflowCheckpointException +from agent_framework import WorkflowCheckpointException, register_checkpoint_type from agent_framework._workflows._checkpoint import FileCheckpointStorage from agent_framework._workflows._checkpoint_encoding import ( _PICKLE_MARKER, @@ -210,6 +210,13 @@ class _AllowedTestState: value: int +@dataclass +class _GloballyRegisteredTestState: + """Test dataclass registered for process-wide checkpoint deserialization.""" + + name: str + + def test_restricted_decode_blocks_unlisted_user_type(): """User-defined types are blocked when not in allowed_checkpoint_types.""" original = _AllowedTestState(name="test", value=42) @@ -301,6 +308,25 @@ async def test_file_storage_allows_listed_user_type(): assert loaded.state["data"].value == 99 +async def test_file_storage_allows_globally_registered_user_type() -> None: + """A registered type can be restored without configuring the storage instance.""" + from agent_framework import WorkflowCheckpoint + + register_checkpoint_type(_GloballyRegisteredTestState) + + with tempfile.TemporaryDirectory() as tmpdir: + storage = FileCheckpointStorage(tmpdir) + checkpoint = WorkflowCheckpoint( + workflow_name="test", + graph_signature_hash="hash", + state={"data": _GloballyRegisteredTestState(name="registered")}, + ) + await storage.save(checkpoint) + loaded = await storage.load(checkpoint.checkpoint_id) + + assert loaded.state["data"] == _GloballyRegisteredTestState(name="registered") + + async def test_file_storage_round_trips_marker_shaped_dict_state() -> None: """FileCheckpointStorage preserves marker-shaped dictionaries as user data.""" from agent_framework import WorkflowCheckpoint diff --git a/python/samples/03-workflows/checkpoint/checkpoint_with_human_in_the_loop.py b/python/samples/03-workflows/checkpoint/checkpoint_with_human_in_the_loop.py index adb43a0921a..e5cbf1ddf38 100644 --- a/python/samples/03-workflows/checkpoint/checkpoint_with_human_in_the_loop.py +++ b/python/samples/03-workflows/checkpoint/checkpoint_with_human_in_the_loop.py @@ -20,6 +20,7 @@ WorkflowBuilder, WorkflowContext, handler, + register_checkpoint_type, response_handler, ) from agent_framework.foundry import FoundryChatClient @@ -27,9 +28,9 @@ from dotenv import load_dotenv if sys.version_info >= (3, 12): - from typing import override # type: ignore # pragma: no cover + from typing import override # pragma: no cover else: - from typing_extensions import override # type: ignore[import] # pragma: no cover + from typing_extensions import override # pragma: no cover # Load environment variables from .env file load_dotenv() @@ -42,8 +43,9 @@ 1. A brief is turned into a consistent prompt for an AI copywriter. 2. The copywriter (an `AgentExecutor`) drafts release notes. 3. A reviewer gateway sends a request for approval for every draft. -4. The workflow records checkpoints between each superstep so you can stop the - program, restart later, and optionally pre-supply human answers on resume. +4. An output executor emits the approved draft as the terminal workflow output. +5. The workflow records checkpoints between each superstep so you can stop the + program and restart later. Key concepts demonstrated ------------------------- @@ -55,10 +57,8 @@ 1. Run the workflow until a human approval request is emitted. 2. If the human is offline, exit the program. A checkpoint with ``status=awaiting human response`` now exists. -3. Later, restart the script, select that checkpoint, and provide the stored - human decision when prompted to pre-supply responses. - Doing so applies the answer immediately on resume, so the system does **not** - re-emit the same ``. +3. Later, restart the script and select that checkpoint. The workflow restores + and re-emits the pending request so the human can answer it. """ # Directory used for the sample's temporary checkpoint files. We isolate the @@ -107,8 +107,7 @@ class HumanApprovalRequest: """Request sent to the human reviewer.""" # These fields are intentionally simple because they are serialised into - # checkpoints. Keeping them primitive types guarantees the new - # `pending_requests_from_checkpoint` helper can reconstruct them on resume. + # checkpoints and reconstructed when the workflow resumes. prompt: str = "" draft: str = "" iteration: int = 0 @@ -193,7 +192,12 @@ def create_workflow(checkpoint_storage: FileCheckpointStorage) -> Workflow: prepare_brief = BriefPreparer(id="prepare_brief", agent_id="writer") workflow_builder = ( - WorkflowBuilder(max_iterations=6, start_executor=prepare_brief, checkpoint_storage=checkpoint_storage) + WorkflowBuilder( + max_iterations=6, + start_executor=prepare_brief, + checkpoint_storage=checkpoint_storage, + output_from=[review_gateway], + ) .add_edge(prepare_brief, writer) .add_edge(writer, review_gateway) .add_edge(review_gateway, writer) # revisions loop @@ -277,6 +281,11 @@ async def main() -> None: # deterministic even if the directory had stale checkpoints. file.unlink() + # Register the application-defined request type so file storage can reconstruct it when loading checkpoints. + # Alternatively, scope permission to this storage instance: + # allowed_types = [f"{HumanApprovalRequest.__module__}:{HumanApprovalRequest.__qualname__}"] + # storage = FileCheckpointStorage(storage_path=TEMP_DIR, allowed_checkpoint_types=allowed_types) + register_checkpoint_type(HumanApprovalRequest) storage = FileCheckpointStorage(storage_path=TEMP_DIR) workflow = create_workflow(checkpoint_storage=storage)