From 9a0974ab918b982a3aad0bf46e6c67e2fe1b20d0 Mon Sep 17 00:00:00 2001 From: atty57 <99388680+atty57@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:48:16 -0400 Subject: [PATCH 1/2] Python: Register built-in orchestration types for checkpoint restore Group chat, handoff, and Magentic orchestrations send framework-owned dataclasses as executor messages and request_info payloads. Those types live in agent_framework_orchestrations, which is outside the "agent_framework." module prefix that the restricted checkpoint unpickler auto-allows, so restoring a checkpoint failed unless the user listed the internal module paths in allowed_checkpoint_types. Register them with the existing register_checkpoint_type API at package import so built-in orchestrations restore without user configuration. Fixes #7789 --- .../__init__.py | 26 +++++++++ .../tests/test_checkpoint_types.py | 54 +++++++++++++++++++ 2 files changed, 80 insertions(+) create mode 100644 python/packages/orchestrations/tests/test_checkpoint_types.py diff --git a/python/packages/orchestrations/agent_framework_orchestrations/__init__.py b/python/packages/orchestrations/agent_framework_orchestrations/__init__.py index d1acb7af53e..fb3e456b99a 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/__init__.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/__init__.py @@ -17,10 +17,14 @@ except importlib.metadata.PackageNotFoundError: __version__ = "0.0.0" # Fallback for development mode +from agent_framework import register_checkpoint_type + from ._base_group_chat_orchestrator import ( BaseGroupChatOrchestrator, + GroupChatParticipantMessage, GroupChatRequestMessage, GroupChatRequestSentEvent, + GroupChatResponseMessage, GroupChatResponseReceivedEvent, TerminationCondition, ) @@ -59,6 +63,7 @@ MagenticProgressLedgerItem, MagenticResetSignal, StandardMagenticManager, + _MagenticTaskLedger, ) from ._orchestration_request_info import AgentRequestInfoResponse from ._orchestration_state import OrchestrationState @@ -108,3 +113,24 @@ "clean_conversation_for_handoff", "create_completion_message", ] + +# Framework-owned types that cross a checkpoint boundary: executor-to-executor message +# envelopes and request_info payloads/responses. Registering them here means built-in +# orchestrations restore without users maintaining their own `allowed_checkpoint_types` +# list of framework module paths. +for _checkpoint_type in ( + GroupChatRequestMessage, + GroupChatParticipantMessage, + GroupChatResponseMessage, + HandoffAgentUserRequest, + AgentRequestInfoResponse, + MagenticResetSignal, + MagenticPlanReviewRequest, + MagenticPlanReviewResponse, + MagenticProgressLedger, + MagenticProgressLedgerItem, + _MagenticTaskLedger, +): + register_checkpoint_type(_checkpoint_type) + +del _checkpoint_type diff --git a/python/packages/orchestrations/tests/test_checkpoint_types.py b/python/packages/orchestrations/tests/test_checkpoint_types.py new file mode 100644 index 00000000000..f21feea483a --- /dev/null +++ b/python/packages/orchestrations/tests/test_checkpoint_types.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft. All rights reserved. + +from typing import Any + +import pytest +from agent_framework import Message +from agent_framework._workflows._checkpoint_encoding import ( + _REGISTERED_CHECKPOINT_TYPE_KEYS, + decode_checkpoint_value, + encode_checkpoint_value, +) + +from agent_framework_orchestrations._base_group_chat_orchestrator import ( + GroupChatParticipantMessage, + GroupChatRequestMessage, + GroupChatResponseMessage, +) +from agent_framework_orchestrations._handoff import HandoffAgentUserRequest +from agent_framework_orchestrations._magentic import ( + MagenticPlanReviewRequest, + MagenticPlanReviewResponse, + MagenticResetSignal, +) +from agent_framework_orchestrations._orchestration_request_info import AgentRequestInfoResponse + + +@pytest.mark.parametrize( + "value", + [ + GroupChatRequestMessage(additional_instruction="go"), + GroupChatParticipantMessage(messages=[Message(role="user", contents=[])]), + GroupChatResponseMessage(message=Message(role="assistant", contents=[])), + MagenticResetSignal(), + ], +) +def test_builtin_envelopes_restore_without_extra_allowed_types(value: Any) -> None: + """Framework-owned envelopes decode under a restricted allowlist (issue #7789).""" + restored = decode_checkpoint_value(encode_checkpoint_value(value), allowed_types=frozenset()) + + assert type(restored) is type(value) + + +@pytest.mark.parametrize( + "cls", + [ + HandoffAgentUserRequest, + AgentRequestInfoResponse, + MagenticPlanReviewRequest, + MagenticPlanReviewResponse, + ], +) +def test_builtin_request_info_types_are_registered(cls: type) -> None: + """Request/response payloads that get persisted are trusted by default.""" + assert f"{cls.__module__}:{cls.__qualname__}" in _REGISTERED_CHECKPOINT_TYPE_KEYS From edb4b5c8109deba419d7f31ac753b83a0438563d Mon Sep 17 00:00:00 2001 From: Atharva Vichare <99388680+atty57@users.noreply.github.com> Date: Thu, 20 Aug 2026 12:55:39 -0400 Subject: [PATCH 2/2] Remove parameter from function as its already taken care of Remove parameter '_MagenticTaskLedger' from function as its already taken care of with an override for converting it to_dict and from_dict --- .../orchestrations/agent_framework_orchestrations/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/packages/orchestrations/agent_framework_orchestrations/__init__.py b/python/packages/orchestrations/agent_framework_orchestrations/__init__.py index fb3e456b99a..97e2c2485db 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/__init__.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/__init__.py @@ -129,7 +129,6 @@ MagenticPlanReviewResponse, MagenticProgressLedger, MagenticProgressLedgerItem, - _MagenticTaskLedger, ): register_checkpoint_type(_checkpoint_type)