Python: Add a global workflow checkpoint type registry (#7636)

* Add a glocal checkpoint type registry

* Update samples

* Revert uv.lock

* Address comments

* Revert uv.lock

* Revert uv.lock

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:
 

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)})

python/packages/core/agent_framework/__init__.py

@@ -293,6 +293,7 @@ _LAZY_MODULE_EXPORTS: Final[Mapping[str, tuple[str, ...]]] = {
         "InMemoryCheckpointStorage",
         "WorkflowCheckpoint",
     ),
+    "._workflows._checkpoint_encoding": ("register_checkpoint_type",),
     "._workflows._const": (
         "DEFAULT_MAX_ITERATIONS",
         "INTERNAL_SOURCE_ID",
@@ -631,6 +632,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",

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",

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::
 

python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py

@@ -57,6 +57,28 @@ from ..exceptions import WorkflowCheckpointException
 
 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 @@ def decode_checkpoint_value(value: Any, *, allowed_types: frozenset[str] | None
             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)
 
 

python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py

@@ -22,7 +22,7 @@ from typing import Any
 
 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

python/samples/03-workflows/checkpoint/checkpoint_with_human_in_the_loop.py

@@ -20,6 +20,7 @@ from agent_framework import (
     WorkflowBuilder,
     WorkflowContext,
     handler,
+    register_checkpoint_type,
     response_handler,
 )
 from agent_framework.foundry import FoundryChatClient
@@ -27,9 +28,9 @@ from azure.identity import AzureCliCredential
 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 @@ This getting-started sample keeps the moving pieces to a minimum:
 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 @@ Typical pause/resume flow
 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)