Skip to content

Commit 435012e

Browse files
FOWEPJF255lhl-lhiktz03moonbox3eavanvalkenburg
authored
Python: add public TypedDict for AgentExecutor checkpoint state (#8201) (#8284)
* Python: add public TypedDict for AgentExecutor checkpoint state (#8201) Expose AgentExecutorCheckpointState / AgentSessionCheckpointState, validate restore field types, and cover partial/malformed/forward-compatible payloads. * fix(checkpoint): address #8284 review on session TypedDict ownership (#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__. * fix(checkpoint): validate AgentExecutor restore element types (#8284) Enforce Message/Content element types, string pending-request keys, and required agent_session.session_id so malformed checkpoints fail at restore. * fix(checkpoint): validate AgentSessionDict fields on restore - Validate agent_session.state / service_session_id types - Raise WorkflowCheckpointException when session restore fails - Drop unshipped AgentSessionCheckpointState alias; use AgentSessionDict * fix(checkpoint): keep session/executor checkpoint payloads as dict TypedDict schemas remain public for documentation and validation, but to_dict and on_checkpoint_save/restore stay dict[str, Any] so subclasses, samples, and pyright stay compatible. * fix(checkpoint): keep AgentSessionDict optional keys under postponed annotations Use a required base plus total=False subclass, stop exporting AgentSessionDict from the public package root, and tighten AgentExecutor restore validation with walrus checks. --------- Co-authored-by: minelhi <[email protected]> Co-authored-by: LI <[email protected]> Co-authored-by: Evan Mattson <[email protected]> Co-authored-by: Eduard van Valkenburg <[email protected]>
1 parent f80498b commit 435012e

5 files changed

Lines changed: 287 additions & 12 deletions

File tree

‎python/packages/core/agent_framework/__init__.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -322,7 +322,12 @@
322322
"vectorstoremodel",
323323
),
324324
"._workflows._agent": ("WorkflowAgent",),
325-
"._workflows._agent_executor": ("AgentExecutor", "AgentExecutorRequest", "AgentExecutorResponse"),
325+
"._workflows._agent_executor": (
326+
"AgentExecutor",
327+
"AgentExecutorCheckpointState",
328+
"AgentExecutorRequest",
329+
"AgentExecutorResponse",
330+
),
326331
"._workflows._agent_utils": ("resolve_agent_id",),
327332
"._workflows._checkpoint": (
328333
"CheckpointID",
@@ -429,6 +434,7 @@
429434
"AgentContext",
430435
"AgentEvalConverter",
431436
"AgentExecutor",
437+
"AgentExecutorCheckpointState",
432438
"AgentExecutorRequest",
433439
"AgentExecutorResponse",
434440
"AgentFileStore",

‎python/packages/core/agent_framework/__init__.pyi‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,12 @@ from ._vectors import (
278278
vectorstoremodel,
279279
)
280280
from ._workflows._agent import WorkflowAgent
281-
from ._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
281+
from ._workflows._agent_executor import (
282+
AgentExecutor,
283+
AgentExecutorCheckpointState,
284+
AgentExecutorRequest,
285+
AgentExecutorResponse,
286+
)
282287
from ._workflows._agent_utils import resolve_agent_id
283288
from ._workflows._checkpoint import (
284289
CheckpointID,
@@ -384,6 +389,7 @@ __all__ = [
384389
"AgentContext",
385390
"AgentEvalConverter",
386391
"AgentExecutor",
392+
"AgentExecutorCheckpointState",
387393
"AgentExecutorRequest",
388394
"AgentExecutorResponse",
389395
"AgentFileStore",

‎python/packages/core/agent_framework/_sessions.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeVar, cast
3737

3838
import msgspec
39+
from typing_extensions import TypedDict
3940

4041
from ._feature_stage import ExperimentalFeature, experimental
4142
from ._filesystem import (
@@ -1754,6 +1755,32 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[
17541755
)
17551756

17561757

1758+
class _AgentSessionDictRequired(TypedDict):
1759+
"""Required fields for a serialized :class:`AgentSession`."""
1760+
1761+
session_id: str
1762+
1763+
1764+
class AgentSessionDict(_AgentSessionDictRequired, total=False):
1765+
"""Serialized :class:`AgentSession` payload shape produced by :meth:`AgentSession.to_dict`.
1766+
1767+
``AgentSession.to_dict`` returns a plain ``dict[str, Any]`` that conforms to this
1768+
schema. Callers that need a TypedDict view can ``cast`` the result.
1769+
1770+
Built as a required base plus ``total=False`` optional fields so postponed
1771+
annotations do not turn optional keys into required runtime metadata.
1772+
1773+
``service_session_id`` may be a plain string or a structured
1774+
:data:`ServiceSessionId` mapping, matching :attr:`AgentSession.service_session_id`.
1775+
``state`` holds session-local data and may be incomplete when the session uses
1776+
service-side storage.
1777+
"""
1778+
1779+
type: str
1780+
service_session_id: str | ServiceSessionId | None
1781+
state: dict[str, Any]
1782+
1783+
17571784
class AgentSession:
17581785
"""A conversation session with an agent.
17591786
@@ -1797,6 +1824,11 @@ def session_id(self) -> str:
17971824
def to_dict(self) -> dict[str, Any]:
17981825
"""Serialize session to a plain dict for storage/transfer.
17991826
1827+
The returned mapping matches :class:`AgentSessionDict`. The annotated
1828+
return type stays ``dict[str, Any]`` so subclasses and callers that
1829+
extend or pass the payload as a mutable ``dict`` remain type-correct
1830+
(TypedDict is not assignable to ``dict`` under pyright).
1831+
18001832
Registered custom values use their configured codecs. Unregistered
18011833
values defining ``to_dict`` retain the established dictionary behavior.
18021834
Unregistered Pydantic models are still auto-registered temporarily but

‎python/packages/core/agent_framework/_workflows/_agent_executor.py‎

Lines changed: 143 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,18 @@
33
import inspect
44
import logging
55
import sys
6-
from collections.abc import Awaitable, Callable
6+
from collections.abc import Awaitable, Callable, Mapping
77
from dataclasses import dataclass
88
from typing import Any, Literal, cast
99

10-
from typing_extensions import Never
10+
from typing_extensions import Never, TypedDict
1111

1212
from agent_framework import Content
1313

1414
from .._agents import SupportsAgentRun
15-
from .._sessions import AgentSession
15+
from .._sessions import AgentSession, AgentSessionDict
1616
from .._types import AgentResponse, AgentResponseUpdate, Message, ResponseStream
17+
from ..exceptions import WorkflowCheckpointException
1718
from ._agent_utils import prepare_agent_run_args, resolve_agent_id, resolve_executor_kwargs
1819
from ._const import INTERNAL_SOURCE_ID, RESOLVED_WORKFLOW_RUN_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
1920
from ._executor import Executor, handler
@@ -30,6 +31,128 @@
3031
logger = logging.getLogger(__name__)
3132

3233

34+
class AgentExecutorCheckpointState(TypedDict, total=False):
35+
"""Public schema for state saved and restored by :class:`AgentExecutor`.
36+
37+
Stored under ``WorkflowCheckpoint.state["_executor_state"][executor_id]``.
38+
39+
``on_checkpoint_save`` always writes every key below. Restore accepts partial
40+
mappings for backward compatibility: missing keys reset to empty defaults.
41+
Unknown keys are ignored so newer writers remain readable by older runtimes.
42+
43+
Compatibility:
44+
- Add fields as ``total=False`` / ``NotRequired`` and treat absence as default.
45+
- Do not rename or change the meaning of existing keys without a migration path.
46+
- Custom executors should define their own TypedDict (or equivalent) and validate
47+
types in ``on_checkpoint_restore``; prefer
48+
:class:`~agent_framework.exceptions.WorkflowCheckpointException` for malformed data.
49+
50+
Keys:
51+
cache: Messages buffered between runs before the next agent invocation.
52+
full_conversation: Prior inputs plus assistant/tool outputs after the last run.
53+
agent_session: Serialized session payload (:class:`~agent_framework._sessions.AgentSessionDict`).
54+
pending_agent_requests: In-flight agent-owned user-input requests by request id.
55+
pending_responses_to_agent: Queued content responses waiting to be sent to the agent.
56+
"""
57+
58+
cache: list[Message]
59+
full_conversation: list[Message]
60+
agent_session: AgentSessionDict
61+
pending_agent_requests: dict[str, Content]
62+
pending_responses_to_agent: list[Content]
63+
64+
65+
def _validate_agent_executor_checkpoint_state(state: Mapping[str, Any]) -> None:
66+
"""Raise :class:`WorkflowCheckpointException` when checkpoint payload types are wrong."""
67+
if not isinstance(state, Mapping):
68+
raise WorkflowCheckpointException(
69+
f"AgentExecutor checkpoint state must be a mapping, got {type(state).__name__}."
70+
)
71+
72+
message_list_keys = ("cache", "full_conversation")
73+
for key in message_list_keys:
74+
if key not in state or state[key] is None:
75+
continue
76+
value = state[key]
77+
if not isinstance(value, list):
78+
raise WorkflowCheckpointException(
79+
f"AgentExecutor checkpoint field '{key}' must be a list, got {type(value).__name__}."
80+
)
81+
messages = cast(list[Any], value)
82+
for index, item in enumerate(messages):
83+
if not isinstance(item, Message):
84+
raise WorkflowCheckpointException(
85+
f"AgentExecutor checkpoint field '{key}'[{index}] must be Message, "
86+
f"got {type(item).__name__}."
87+
)
88+
89+
if (responses_raw := state.get("pending_responses_to_agent")) is not None:
90+
if not isinstance(responses_raw, list):
91+
raise WorkflowCheckpointException(
92+
"AgentExecutor checkpoint field 'pending_responses_to_agent' must be a list, "
93+
f"got {type(responses_raw).__name__}."
94+
)
95+
responses = cast(list[Any], responses_raw)
96+
for index, item in enumerate(responses):
97+
if not isinstance(item, Content):
98+
raise WorkflowCheckpointException(
99+
"AgentExecutor checkpoint field "
100+
f"'pending_responses_to_agent'[{index}] must be Content, "
101+
f"got {type(item).__name__}."
102+
)
103+
104+
if (pending_raw := state.get("pending_agent_requests")) is not None:
105+
if not isinstance(pending_raw, dict):
106+
raise WorkflowCheckpointException(
107+
"AgentExecutor checkpoint field 'pending_agent_requests' must be a dict, "
108+
f"got {type(pending_raw).__name__}."
109+
)
110+
pending = cast(dict[Any, Any], pending_raw)
111+
for request_id, content in pending.items():
112+
if not isinstance(request_id, str):
113+
raise WorkflowCheckpointException(
114+
"AgentExecutor checkpoint field 'pending_agent_requests' keys must be str, "
115+
f"got {type(request_id).__name__}."
116+
)
117+
if not isinstance(content, Content):
118+
raise WorkflowCheckpointException(
119+
"AgentExecutor checkpoint field "
120+
f"'pending_agent_requests[{request_id!r}]' must be Content, "
121+
f"got {type(content).__name__}."
122+
)
123+
124+
if (session_raw := state.get("agent_session")) is not None:
125+
if not isinstance(session_raw, dict):
126+
raise WorkflowCheckpointException(
127+
"AgentExecutor checkpoint field 'agent_session' must be a dict, "
128+
f"got {type(session_raw).__name__}."
129+
)
130+
session = cast(dict[str, Any], session_raw)
131+
session_id = session.get("session_id")
132+
if not isinstance(session_id, str):
133+
raise WorkflowCheckpointException(
134+
"AgentExecutor checkpoint field 'agent_session.session_id' must be a str, "
135+
f"got {type(session_id).__name__}."
136+
)
137+
if "state" in session and session["state"] is not None and not isinstance(session["state"], dict):
138+
raise WorkflowCheckpointException(
139+
"AgentExecutor checkpoint field 'agent_session.state' must be a dict, "
140+
f"got {type(session['state']).__name__}."
141+
)
142+
if "service_session_id" in session and session["service_session_id"] is not None:
143+
service_session_id = session["service_session_id"]
144+
if not isinstance(service_session_id, (str, Mapping)):
145+
raise WorkflowCheckpointException(
146+
"AgentExecutor checkpoint field 'agent_session.service_session_id' must be "
147+
f"str, mapping, or None, got {type(service_session_id).__name__}."
148+
)
149+
if "type" in session and session["type"] is not None and not isinstance(session["type"], str):
150+
raise WorkflowCheckpointException(
151+
"AgentExecutor checkpoint field 'agent_session.type' must be a str, "
152+
f"got {type(session['type']).__name__}."
153+
)
154+
155+
33156
def _accepts_runtime_tools(agent: SupportsAgentRun) -> bool:
34157
"""Return whether the agent run surface accepts a tools keyword."""
35158
try:
@@ -368,14 +491,15 @@ async def on_checkpoint_save(self) -> dict[str, Any]:
368491
may not be serialized locally.
369492
370493
Returns:
371-
Dict containing serialized cache and session state
494+
A JSON-serializable ``dict`` matching :class:`AgentExecutorCheckpointState`
495+
(cache, conversation, session, and pending request/response fields).
496+
The return type remains ``dict[str, Any]`` so subclasses may extend the
497+
payload and so the override stays compatible with :class:`Executor`.
372498
"""
373-
serialized_session = self._session.to_dict()
374-
375499
return {
376500
"cache": self._cache,
377501
"full_conversation": self._full_conversation,
378-
"agent_session": serialized_session,
502+
"agent_session": self._session.to_dict(),
379503
"pending_agent_requests": self._pending_agent_requests,
380504
"pending_responses_to_agent": self._pending_responses_to_agent,
381505
}
@@ -385,8 +509,15 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
385509
"""Restore executor state from checkpoint.
386510
387511
Args:
388-
state: Checkpoint data dict
512+
state: Checkpoint payload matching :class:`AgentExecutorCheckpointState`.
513+
Missing known keys use empty defaults; unknown keys are ignored.
514+
515+
Raises:
516+
WorkflowCheckpointException: If ``state`` is not a mapping or a known
517+
field has an incompatible type.
389518
"""
519+
_validate_agent_executor_checkpoint_state(state)
520+
390521
cache_payload = state.get("cache")
391522
self._cache = cache_payload or []
392523

@@ -398,8 +529,10 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
398529
try:
399530
self._session = AgentSession.from_dict(session_payload)
400531
except Exception as exc:
401-
logger.warning("Failed to restore agent session: %s", exc)
402-
self._session = self._agent.create_session()
532+
raise WorkflowCheckpointException(
533+
"AgentExecutor checkpoint field 'agent_session' could not be restored: "
534+
f"{exc}"
535+
) from exc
403536
else:
404537
self._session = self._agent.create_session()
405538

0 commit comments

Comments
 (0)