Python: register built-in orchestration types for checkpoint restore

[Manjunath Janardhan (manjunathshiva)]

Motivation & Context

Checkpoint restore runs pickle through a restricted unpickler. Its default allowlist auto-trusts the
agent_framework. prefix -- dotted, so it covers core but not sibling distributions like
agent_framework_orchestrations. The built-in group-chat, handoff and Magentic envelopes are
therefore rejected on restore.

Who is actually affected, and how. I built the end-to-end scenario rather than trusting the
issue text, and it narrowed the trigger in one direction and widened the impact in another.

The orchestrator only wraps traffic in the group-chat envelopes for participants that are not
agents (_base_group_chat_orchestrator.py:434,467); agent participants get a core
AgentExecutorRequest, which is already trusted. So a group chat built purely from agents never
hits this, which is why the existing suite did not catch it -- and the suite's other checkpoint
tests use InMemoryCheckpointStorage, which deepcopys and never reaches the unpickler at all.

The symptom is worse than a failed load(). Running a two-round group chat with one custom
executor participant and a FileCheckpointStorage writes six checkpoints; on main, four of them
cannot be decoded, and list_checkpoints logs and skips them rather than raising:

Failed to read checkpoint file ...json: Failed to decode pickled checkpoint data:
Checkpoint deserialization blocked for type
'agent_framework_orchestrations._base_group_chat_orchestrator:GroupChatResponseMessage'.

So the listing silently returns 2 of 6 and get_latest hands back a stale checkpoint instead of
failing. That is the same silent-invisibility behaviour reported in #7831, which cross-references
this issue. With the registration in place all six list and load.

The only workaround today is for users to hand-maintain allowed_checkpoint_types with
framework-internal module paths that can move between releases.

Worth flagging for sequencing: #8214 adds save-time encode/decode validation, which turns this into
a save-time failure. After it lands a affected workflow fails on its first checkpoint rather
than losing them quietly -- louder, and better, but it makes this more urgent, not less.

Description & Review Guide

• What are the major changes? Ten framework-owned types are registered through the existing
  public register_checkpoint_type when agent_framework_orchestrations is imported. No core
  change, no new API, no behaviour change to the orchestrations themselves.

• Why these ten. Each was verified against the tree rather than inferred from the issue text:

  Type	Crosses the boundary via
  GroupChatRequestMessage	send_message, _base_group_chat_orchestrator.py:486
  GroupChatParticipantMessage	send_message, same file
  GroupChatResponseMessage	handle_participant_response parameter, :238
  MagenticResetSignal	send_message and its handler
  HandoffAgentUserRequest	ctx.request_info(...), _handoff.py:439
  AgentRequestInfoResponse	request_info response type, _orchestration_request_info.py:101
  MagenticPlanReviewRequest / Response	ctx.request_info(...), _magentic.py:1048
  MagenticProgressLedger / Item	nested inside MagenticPlanReviewRequest.current_progress

  _MagenticTaskLedger is excluded on purpose — it is persisted through to_dict()/from_dict()
  and never reaches the unpickler.

• The alternative I did not take. Dropping the dot from _FRAMEWORK_MODULE_PREFIX fixes the
  whole class in one line and would cover declarative, ag-ui and every future sibling too. I think
  it should not be done: this is a pickle allowlist, and the undotted prefix auto-trusts any
  installed distribution merely named agent_framework_*, which is a name that can be taken on
  PyPI. That is a security-boundary change and a maintainer decision, so it is not in this PR. Happy
  to open it as a separate discussion if you want the structural fix.

• What is the impact of these changes? Built-in orchestrations restore from checkpoints without
  users supplying framework-internal module paths. Two properties worth stating plainly: the
  allowlist is process-wide, so importing this package widens what any checkpoint storage in the
  process will unpickle; and registration happens at import, so a process restoring one of these
  checkpoints must import agent_framework_orchestrations. In practice it already does — the
  checkpoints cannot exist otherwise — but it is the cost of opting in by name instead of widening
  the prefix, and there is a test pinning it. All ten registered types are plain dataclasses with no
  __reduce__, __setstate__ or __new__ overrides, so none of them adds a construction hook.

• Credit. Python: Register built-in orchestration types for checkpoint restore #7791 by Atharva Vichare (@atty57) reached the same shape first and was closed on 2026-09-04 as a stale
  draft rather than on approach. This follows that approach and addresses the Copilot review comment
  left open there — that a test asserting only the outer request class would still pass with the
  nested ledger types unregistered. The tests here round-trip fully populated payloads through a
  storage built with no allowed_checkpoint_types; removing only MagenticProgressLedger and
  MagenticProgressLedgerItem from the registration leaves every other case green and fails exactly
  three tests.

• What do you want reviewers to focus on? Whether the ten are the right ten — I would rather
  register too few and be corrected than widen the unpickling surface speculatively. Also whether
  GroupChatParticipantMessage and GroupChatResponseMessage should join __all__; they are
  imported here for registration but deliberately left unexported, since expanding public API is not
  a bug fix's business.

  Validation: orchestrations 206 passing with poe check -P orchestrations clean across all five
  type checkers, and core 5128 passing unchanged. Every new test fails against the unregistered
  tree.

  Related, not claimed: Python: Python : checkpoint restore +pending request deserialization failure when hosting workflow as agent with ResponsesHostserver in foundry #7618 reports the same root cause reaching foundry hosting. I have not
  verified its downstream unknown request Id cascade, so this links only Python: Built-in group chat message types are rejected during checkpoint restore #7789.

Related Issue

Fixes #7789

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change.

[Copilot]

🟢 Approval recommended

The registrations match checkpoint-boundary usage and are comprehensively covered by focused restoration tests.

Pull request overview

Registers built-in orchestration payload types so restricted checkpoint deserialization can restore group-chat, handoff, and Magentic workflows without user-managed allowlists.

Changes:

• Registers ten checkpoint-crossing orchestration types at package import.
• Adds round-trip and nested-ledger restoration tests.
• Verifies package import populates the global registry.

File summaries

File	Description
python/packages/orchestrations/agent_framework_orchestrations/__init__.py	Registers trusted orchestration checkpoint types.
python/packages/orchestrations/tests/test_checkpoint_types.py	Tests default restore behavior and registration completeness.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Balanced

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.