Python: add public TypedDict for AgentExecutor checkpoint state (#8201) - #8284
Conversation
…osoft#8201) Expose AgentExecutorCheckpointState / AgentSessionCheckpointState, validate restore field types, and cover partial/malformed/forward-compatible payloads.
GitHub Actions
Active
GitHub Actions
Active
GitHub Actions
Active
GitHub Actions
Active
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
🟡 Changes recommended
The root stub is unsynchronized, and validation does not fully enforce the published checkpoint schema.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds public checkpoint-state schemas for Python’s AgentExecutor, including restore validation and compatibility guidance.
Changes:
- Adds and exports two checkpoint
TypedDictcontracts. - Validates checkpoint fields during restoration.
- Adds schema, compatibility, and malformed-state tests.
File summaries
| File | Description |
|---|---|
test_agent_executor.py |
Tests checkpoint schemas and restoration behavior. |
_agent_executor.py |
Defines schemas and restore validation. |
__init__.py |
Exports the new public types. |
Review details
Suppressed comments (3)
python/packages/core/agent_framework/_workflows/_agent_executor.py:87
- Checking only
listdoes not validate the declared field types. For example,{"cache": ["bad"]}passes restore and is later sent torun_agentas though it containedMessageobjects; malformed pending responses similarly fail later when.typeis accessed instead of producing the promised checkpoint error. Validate each list element asMessageorContent, as appropriate.
list_keys = ("cache", "full_conversation", "pending_responses_to_agent")
for key in list_keys:
if key in state and state[key] is not None and not isinstance(state[key], list):
python/packages/core/agent_framework/_workflows/_agent_executor.py:93
- The mapping's declared
dict[str, Content]shape is not validated. A checkpoint with non-string keys or non-Contentvalues is accepted into_pending_agent_requests, which can leave pending requests that normal string request IDs can never remove. Validate both keys and values before restoring this field.
if "pending_agent_requests" in state and state["pending_agent_requests"] is not None:
if not isinstance(state["pending_agent_requests"], dict):
python/packages/core/agent_framework/_workflows/_agent_executor.py:100
- A dictionary is not sufficient for a valid
AgentSessionCheckpointState:session_idis required and must be a string. Payloads such as{}or{"session_id": 1}pass this check; the former is then caught and silently replaced with a new session, while the latter creates a session with an invalid ID. Validate the required nested field and raiseWorkflowCheckpointException.
if "agent_session" in state and state["agent_session"] is not None:
if not isinstance(state["agent_session"], dict):
- Files reviewed: 3/3 changed files
- Comments generated: 2
- Review effort level: Balanced
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
| "AgentExecutorCheckpointState", | ||
| "AgentExecutorRequest", | ||
| "AgentExecutorResponse", | ||
| "AgentSessionCheckpointState", |
There was a problem hiding this comment.
Follow-up: synced agent_framework/__init__.pyi with the new root exports (AgentExecutorCheckpointState, AgentSessionCheckpointState, AgentSessionDict). See FOWEPJF255#1 (f016657).
|
|
||
| type: NotRequired[str] | ||
| session_id: str | ||
| service_session_id: NotRequired[str | None] |
There was a problem hiding this comment.
Follow-up: AgentSessionDict.service_session_id is now str | ServiceSessionId | None, matching AgentSession.service_session_id / to_dict(). See FOWEPJF255#1.
| class AgentSessionCheckpointState(TypedDict): | ||
| """Serialized :class:`~agent_framework.AgentSession` payload (``AgentSession.to_dict()``). | ||
|
|
||
| ``state`` holds session-local data. When the session uses service-side storage, | ||
| local ``state`` may be incomplete relative to the remote conversation. | ||
| """ | ||
|
|
||
| type: NotRequired[str] | ||
| session_id: str | ||
| service_session_id: NotRequired[str | None] | ||
| state: NotRequired[dict[str, Any]] |
There was a problem hiding this comment.
can't we get this shape from the session itself, seems like a bad idea to maintain that twice
There was a problem hiding this comment.
Agreed — moved the shape onto AgentSession as AgentSessionDict (to_dict() return type) so AgentExecutor does not maintain a second copy. AgentSessionCheckpointState is an alias for the existing public name. Proposed in FOWEPJF255#1.
|
|
||
| @override | ||
| async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: | ||
| async def on_checkpoint_restore(self, state: AgentExecutorCheckpointState | dict[str, Any]) -> None: |
There was a problem hiding this comment.
why not just type this as the TypedDict, that's kind of the point, no?
There was a problem hiding this comment.
Updated: on_checkpoint_restore(self, state: AgentExecutorCheckpointState) — dropped the | dict[str, Any] union. Proposed in FOWEPJF255#1.
…wnership - Define AgentSessionDict on AgentSession (to_dict return) so the shape is not duplicated in AgentExecutor; keep AgentSessionCheckpointState as alias. - Include ServiceSessionId mapping in service_session_id. - Type on_checkpoint_restore as AgentExecutorCheckpointState only. - Sync root stub exports (__init__.pyi) with runtime __all__.
|
Left a follow-up PR into your branch addressing the open review threads:
Changes:
Feel free to merge #1 into your branch (or cherry-pick |
…wnership (#1) - Define AgentSessionDict on AgentSession (to_dict return) so the shape is not duplicated in AgentExecutor; keep AgentSessionCheckpointState as alias. - Include ServiceSessionId mapping in service_session_id. - Type on_checkpoint_restore as AgentExecutorCheckpointState only. - Sync root stub exports (__init__.pyi) with runtime __all__.
GitHub Actions
Active
|
Merged the TypedDict ownership follow-up on the PR branch (FOWEPJF255#1). Ready for another look when you have a moment. |
Motivation & Context
Applications that inspect, migrate, or associate
AgentExecutorcheckpoint data currently rely on an undocumenteddict[str, Any]shape. A public TypedDict contract improves static analysis and makes checkpoint persistence safer to evolve.Fixes #8201.
Description & Review Guide
AgentExecutorCheckpointStateandAgentSessionCheckpointStateTypedDicts.on_checkpoint_save/on_checkpoint_restoreagainst that schema.WorkflowCheckpointException); accept missing keys and ignore unknown keys for compatibility.agent_frameworkand add unit tests.Related Issue
Fixes #8201
No other open PR targets this issue.
Contribution Checklist
breaking changelabel (or add [BREAKING] to the title prefix, before or after any language prefix) 鈥?a workflow keeps the label and title prefix in sync automatically.