Python: Allow hosted checkpoints to restore MessageRole#6049
Python: Allow hosted checkpoints to restore MessageRole#6049karimbaidar wants to merge 2 commits into
Conversation
Allow Responses hosting checkpoint storage to deserialize the Azure Responses MessageRole enum that hosted workflows can persist inside Agent Framework Message objects. Add regression coverage for both direct load() and the hosted get_latest() restore path, including the plain-storage failure mode where list_checkpoints logs the blocked type and get_latest() returns None. Ruff also normalizes a duplicate contextlib import in the touched hosting module.
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR updates checkpoint storage behavior to support persisting/restoring Azure’s MessageRole enum inside workflow checkpoints in hosted contexts, with tests to validate allowlisting and fallback behavior.
Changes:
- Allowlist
MessageRolefor hostedFileCheckpointStorageso Azure role enums can be deserialized during checkpoint restore. - Add tests covering restore behavior with/without the allowlist and
get_latest()skip behavior when restore fails.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py | Adds a narrow checkpoint-type allowlist entry for MessageRole in hosted checkpoint storage. |
| python/packages/foundry_hosting/tests/test_responses.py | Adds regression tests for restoring checkpoints containing Azure MessageRole, including get_latest() behavior. |
There was a problem hiding this comment.
The four new tests exercise _checkpoint_storage_for_context directly. Should we also drive one of these through ResponsesHostServer so the bug-motivating restore path (multi-turn previous_response_id) is actually covered? Every existing HTTP-level multi-turn test uses a MagicMock(spec=RawAgent) that ignores checkpoint_storage, so a regression on the host-side wiring (someone passing the wrong factory at _responses.py:553-556 or :651-658) would not be caught.
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| _AZURE_RESPONSES_MESSAGE_ROLE_TYPE = f"{MessageRole.__module__}:{MessageRole.__qualname__}" |
There was a problem hiding this comment.
Worth pinning the resolved string in a test? At runtime this is azure.ai.agentserver.responses.models._generated.sdk.models.models._enums:MessageRole. If the azure SDK ever re-homes the enum, __module__ would silently change and the allowlist key would still match the new pickle GLOBAL, so this is robust today, but a one-line assert _AZURE_RESPONSES_MESSAGE_ROLE_TYPE == "..." would catch the inverse case (re-export path stays the same but the canonical module moves and pickled module:qualname shifts).
Motivation and Context
Fixes #5810.
ResponsesHostServerowns workflow checkpoint storage for hosted Responses agents. Those checkpoints can persist the Azure ResponsesMessageRoleenum inside Agent FrameworkMessageobjects, but restrictedFileCheckpointStoragedeserialization did not allow that SDK enum. In the hosted restore path,get_latest()goes throughlist_checkpoints(), which logs the blocked type and returns no checkpoint, so resumable hosted workflows can fail to restore persisted state.Description
Configures the Responses hosting checkpoint storage to allow only the Azure Responses
MessageRoleenum during restricted checkpoint deserialization. This keeps the allowlist provider-specific and narrow instead of changing core checkpoint defaults or broadly allowing Azure SDK modules.Adds focused regression coverage for the direct
load()path and the hostedget_latest()restore path, including the plain-storage failure mode wherelist_checkpoints()logs the blockedMessageRoletype andget_latest()returnsNone.Contribution Checklist