diff --git a/python/packages/orchestrations/agent_framework_orchestrations/__init__.py b/python/packages/orchestrations/agent_framework_orchestrations/__init__.py index d1acb7af53..97e2c2485d 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,23 @@ "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, +): + 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 0000000000..f21feea483 --- /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