fix(checkpoint): address #8284 moonbox3 session validation

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,21 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[
         )
 
 
+class AgentSessionDict(TypedDict):
+    """Serialized :class:`AgentSession` payload produced by :meth:`AgentSession.to_dict`.
+
+    ``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.
 
@@ -1744,7 +1760,7 @@ def session_id(self) -> str:
         """The unique identifier for this session."""
         return self._session_id
 
-    def to_dict(self) -> dict[str, Any]:
+    def to_dict(self) -> AgentSessionDict:
         """Serialize session to a plain dict for storage/transfer.
 
         Registered custom values use their configured codecs. Unregistered

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 resolve_agent_id
 from ._const import GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY
 from ._executor import Executor, handler
@@ -30,6 +31,127 @@
 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__}."
+            )
+        for index, item in enumerate(value):
+            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 = state["pending_responses_to_agent"]
+        if not isinstance(responses, list):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'pending_responses_to_agent' must be a list, "
+                f"got {type(responses).__name__}."
+            )
+        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 = state["pending_agent_requests"]
+        if not isinstance(pending, dict):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'pending_agent_requests' must be a dict, "
+                f"got {type(pending).__name__}."
+            )
+        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 = state["agent_session"]
+        if not isinstance(session, dict):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'agent_session' must be a dict, "
+                f"got {type(session).__name__}."
+            )
+        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:
@@ -347,32 +469,40 @@ async def _cancel_pending_request(
             await self._resume_with_pending_responses(ctx)
 
     @override
-    async def on_checkpoint_save(self) -> dict[str, Any]:
+    async def on_checkpoint_save(self) -> AgentExecutorCheckpointState:
         """Capture current executor state for checkpointing.
 
         NOTE: if the session uses service-side storage, the full session state
         may not be serialized locally.
 
         Returns:
-            Dict containing serialized cache and session state
+            :class:`AgentExecutorCheckpointState` with cache, conversation, session,
+            and pending request/response fields.
         """
         serialized_session = self._session.to_dict()
 
-        return {
-            "cache": self._cache,
-            "full_conversation": self._full_conversation,
-            "agent_session": serialized_session,
-            "pending_agent_requests": self._pending_agent_requests,
-            "pending_responses_to_agent": self._pending_responses_to_agent,
-        }
+        return AgentExecutorCheckpointState(
+            cache=self._cache,
+            full_conversation=self._full_conversation,
+            agent_session=serialized_session,
+            pending_agent_requests=self._pending_agent_requests,
+            pending_responses_to_agent=self._pending_responses_to_agent,
+        )
 
     @override
-    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
+    async def on_checkpoint_restore(self, state: AgentExecutorCheckpointState) -> 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 +514,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()