Python: add public TypedDict for AgentExecutor checkpoint state (#8201)

python/packages/core/agent_framework/__init__.py

@@ -195,6 +195,7 @@
     ),
     "._sessions": (
         "AgentSession",
+        "AgentSessionDict",
         "ContextProvider",
         "FileHistoryProvider",
         "FileSessionStore",
@@ -316,7 +317,12 @@
         "vectorstoremodel",
     ),
     "._workflows._agent": ("WorkflowAgent",),
-    "._workflows._agent_executor": ("AgentExecutor", "AgentExecutorRequest", "AgentExecutorResponse"),
+    "._workflows._agent_executor": (
+        "AgentExecutor",
+        "AgentExecutorCheckpointState",
+        "AgentExecutorRequest",
+        "AgentExecutorResponse",
+    ),
     "._workflows._agent_utils": ("resolve_agent_id",),
     "._workflows._checkpoint": (
         "CheckpointID",
@@ -423,6 +429,7 @@
     "AgentContext",
     "AgentEvalConverter",
     "AgentExecutor",
+    "AgentExecutorCheckpointState",
     "AgentExecutorRequest",
     "AgentExecutorResponse",
     "AgentFileStore",
@@ -436,6 +443,7 @@
     "AgentResponseUpdate",
     "AgentRunInputs",
     "AgentSession",
+    "AgentSessionDict",
     "AggregatingSkillsSource",
     "Annotation",
     "BackgroundAgentsProvider",

python/packages/core/agent_framework/__init__.pyi

@@ -156,6 +156,7 @@ from ._middleware import (
 from ._sessions import (
     MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY,
     AgentSession,
+    AgentSessionDict,
     ContextProvider,
     FileHistoryProvider,
     FileSessionStore,
@@ -273,7 +274,12 @@ from ._vectors import (
     vectorstoremodel,
 )
 from ._workflows._agent import WorkflowAgent
-from ._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
+from ._workflows._agent_executor import (
+    AgentExecutor,
+    AgentExecutorCheckpointState,
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+)
 from ._workflows._agent_utils import resolve_agent_id
 from ._workflows._checkpoint import (
     CheckpointID,
@@ -378,6 +384,7 @@ __all__ = [
     "AgentContext",
     "AgentEvalConverter",
     "AgentExecutor",
+    "AgentExecutorCheckpointState",
     "AgentExecutorRequest",
     "AgentExecutorResponse",
     "AgentFileStore",
@@ -391,6 +398,7 @@ __all__ = [
     "AgentResponseUpdate",
     "AgentRunInputs",
     "AgentSession",
+    "AgentSessionDict",
     "AggregatingSkillsSource",
     "Annotation",
     "BackgroundAgentsProvider",

python/packages/core/agent_framework/_sessions.py

@@ -36,6 +36,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeVar, cast
 
 import msgspec
+from typing_extensions import NotRequired, TypedDict
 
 from ._feature_stage import ExperimentalFeature, experimental
 from ._filesystem import (
@@ -1704,6 +1705,24 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[
         )
 
 
+class AgentSessionDict(TypedDict):
+    """Serialized :class:`AgentSession` payload shape produced by :meth:`AgentSession.to_dict`.
+
+    ``AgentSession.to_dict`` returns a plain ``dict[str, Any]`` that conforms to this
+    schema. Callers that need a TypedDict view can ``cast`` the result.
+
+    ``service_session_id`` may be a plain string or a structured
+    :data:`ServiceSessionId` mapping, matching :attr:`AgentSession.service_session_id`.
+    ``state`` holds session-local data and may be incomplete when the session uses
+    service-side storage.
+    """
+
+    type: NotRequired[str]
+    session_id: str
+    service_session_id: NotRequired[str | ServiceSessionId | None]
+    state: NotRequired[dict[str, Any]]
+
+
 class AgentSession:
     """A conversation session with an agent.
 
@@ -1747,6 +1766,11 @@ def session_id(self) -> str:
     def to_dict(self) -> dict[str, Any]:
         """Serialize session to a plain dict for storage/transfer.
 
+        The returned mapping matches :class:`AgentSessionDict`. The annotated
+        return type stays ``dict[str, Any]`` so subclasses and callers that
+        extend or pass the payload as a mutable ``dict`` remain type-correct
+        (TypedDict is not assignable to ``dict`` under pyright).
+
         Registered custom values use their configured codecs. Unregistered
         values defining ``to_dict`` retain the established dictionary behavior.
         Unregistered Pydantic models are still auto-registered temporarily but

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

@@ -3,17 +3,18 @@
 import inspect
 import logging
 import sys
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass
 from typing import Any, Literal, cast
 
-from typing_extensions import Never
+from typing_extensions import Never, TypedDict
 
 from agent_framework import Content
 
 from .._agents import SupportsAgentRun
-from .._sessions import AgentSession
+from .._sessions import AgentSession, AgentSessionDict
 from .._types import AgentResponse, AgentResponseUpdate, Message, ResponseStream
+from ..exceptions import WorkflowCheckpointException
 from ._agent_utils import prepare_agent_run_args, resolve_agent_id, resolve_executor_kwargs
 from ._const import INTERNAL_SOURCE_ID, RESOLVED_WORKFLOW_RUN_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY
 from ._executor import Executor, handler
@@ -30,6 +31,131 @@
 logger = logging.getLogger(__name__)
 
 
+class AgentExecutorCheckpointState(TypedDict, total=False):
+    """Public schema for state saved and restored by :class:`AgentExecutor`.
+
+    Stored under ``WorkflowCheckpoint.state["_executor_state"][executor_id]``.
+
+    ``on_checkpoint_save`` always writes every key below. Restore accepts partial
+    mappings for backward compatibility: missing keys reset to empty defaults.
+    Unknown keys are ignored so newer writers remain readable by older runtimes.
+
+    Compatibility:
+        - Add fields as ``total=False`` / ``NotRequired`` and treat absence as default.
+        - Do not rename or change the meaning of existing keys without a migration path.
+        - Custom executors should define their own TypedDict (or equivalent) and validate
+          types in ``on_checkpoint_restore``; prefer
+          :class:`~agent_framework.exceptions.WorkflowCheckpointException` for malformed data.
+
+    Keys:
+        cache: Messages buffered between runs before the next agent invocation.
+        full_conversation: Prior inputs plus assistant/tool outputs after the last run.
+        agent_session: Serialized session payload (:class:`~agent_framework.AgentSessionDict`).
+        pending_agent_requests: In-flight agent-owned user-input requests by request id.
+        pending_responses_to_agent: Queued content responses waiting to be sent to the agent.
+    """
+
+    cache: list[Message]
+    full_conversation: list[Message]
+    agent_session: AgentSessionDict
+    pending_agent_requests: dict[str, Content]
+    pending_responses_to_agent: list[Content]
+
+
+def _validate_agent_executor_checkpoint_state(state: Mapping[str, Any]) -> None:
+    """Raise :class:`WorkflowCheckpointException` when checkpoint payload types are wrong."""
+    if not isinstance(state, Mapping):
+        raise WorkflowCheckpointException(
+            f"AgentExecutor checkpoint state must be a mapping, got {type(state).__name__}."
+        )
+
+    message_list_keys = ("cache", "full_conversation")
+    for key in message_list_keys:
+        if key not in state or state[key] is None:
+            continue
+        value = state[key]
+        if not isinstance(value, list):
+            raise WorkflowCheckpointException(
+                f"AgentExecutor checkpoint field '{key}' must be a list, got {type(value).__name__}."
+            )
+        messages = cast(list[Any], value)
+        for index, item in enumerate(messages):
+            if not isinstance(item, Message):
+                raise WorkflowCheckpointException(
+                    f"AgentExecutor checkpoint field '{key}'[{index}] must be Message, "
+                    f"got {type(item).__name__}."
+                )
+
+    if "pending_responses_to_agent" in state and state["pending_responses_to_agent"] is not None:
+        responses_raw = state["pending_responses_to_agent"]
+        if not isinstance(responses_raw, list):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'pending_responses_to_agent' must be a list, "
+                f"got {type(responses_raw).__name__}."
+            )
+        responses = cast(list[Any], responses_raw)
+        for index, item in enumerate(responses):
+            if not isinstance(item, Content):
+                raise WorkflowCheckpointException(
+                    "AgentExecutor checkpoint field "
+                    f"'pending_responses_to_agent'[{index}] must be Content, "
+                    f"got {type(item).__name__}."
+                )
+
+    if "pending_agent_requests" in state and state["pending_agent_requests"] is not None:
+        pending_raw = state["pending_agent_requests"]
+        if not isinstance(pending_raw, dict):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'pending_agent_requests' must be a dict, "
+                f"got {type(pending_raw).__name__}."
+            )
+        pending = cast(dict[Any, Any], pending_raw)
+        for request_id, content in pending.items():
+            if not isinstance(request_id, str):
+                raise WorkflowCheckpointException(
+                    "AgentExecutor checkpoint field 'pending_agent_requests' keys must be str, "
+                    f"got {type(request_id).__name__}."
+                )
+            if not isinstance(content, Content):
+                raise WorkflowCheckpointException(
+                    "AgentExecutor checkpoint field "
+                    f"'pending_agent_requests[{request_id!r}]' must be Content, "
+                    f"got {type(content).__name__}."
+                )
+
+    if "agent_session" in state and state["agent_session"] is not None:
+        session_raw = state["agent_session"]
+        if not isinstance(session_raw, dict):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'agent_session' must be a dict, "
+                f"got {type(session_raw).__name__}."
+            )
+        session = cast(dict[str, Any], session_raw)
+        session_id = session.get("session_id")
+        if not isinstance(session_id, str):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'agent_session.session_id' must be a str, "
+                f"got {type(session_id).__name__}."
+            )
+        if "state" in session and session["state"] is not None and not isinstance(session["state"], dict):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'agent_session.state' must be a dict, "
+                f"got {type(session['state']).__name__}."
+            )
+        if "service_session_id" in session and session["service_session_id"] is not None:
+            service_session_id = session["service_session_id"]
+            if not isinstance(service_session_id, (str, Mapping)):
+                raise WorkflowCheckpointException(
+                    "AgentExecutor checkpoint field 'agent_session.service_session_id' must be "
+                    f"str, mapping, or None, got {type(service_session_id).__name__}."
+                )
+        if "type" in session and session["type"] is not None and not isinstance(session["type"], str):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'agent_session.type' must be a str, "
+                f"got {type(session['type']).__name__}."
+            )
+
+
 def _accepts_runtime_tools(agent: SupportsAgentRun) -> bool:
     """Return whether the agent run surface accepts a tools keyword."""
     try:
@@ -354,14 +480,15 @@ async def on_checkpoint_save(self) -> dict[str, Any]:
         may not be serialized locally.
 
         Returns:
-            Dict containing serialized cache and session state
+            A JSON-serializable ``dict`` matching :class:`AgentExecutorCheckpointState`
+            (cache, conversation, session, and pending request/response fields).
+            The return type remains ``dict[str, Any]`` so subclasses may extend the
+            payload and so the override stays compatible with :class:`Executor`.
         """
-        serialized_session = self._session.to_dict()
-
         return {
             "cache": self._cache,
             "full_conversation": self._full_conversation,
-            "agent_session": serialized_session,
+            "agent_session": self._session.to_dict(),
             "pending_agent_requests": self._pending_agent_requests,
             "pending_responses_to_agent": self._pending_responses_to_agent,
         }
@@ -371,8 +498,15 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
         """Restore executor state from checkpoint.
 
         Args:
-            state: Checkpoint data dict
+            state: Checkpoint payload matching :class:`AgentExecutorCheckpointState`.
+                Missing known keys use empty defaults; unknown keys are ignored.
+
+        Raises:
+            WorkflowCheckpointException: If ``state`` is not a mapping or a known
+                field has an incompatible type.
         """
+        _validate_agent_executor_checkpoint_state(state)
+
         cache_payload = state.get("cache")
         self._cache = cache_payload or []
 
@@ -384,8 +518,10 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
             try:
                 self._session = AgentSession.from_dict(session_payload)
             except Exception as exc:
-                logger.warning("Failed to restore agent session: %s", exc)
-                self._session = self._agent.create_session()
+                raise WorkflowCheckpointException(
+                    "AgentExecutor checkpoint field 'agent_session' could not be restored: "
+                    f"{exc}"
+                ) from exc
         else:
             self._session = self._agent.create_session()