From 47086ab9e31cd26dc1853bf07ad54cb7919338d7 Mon Sep 17 00:00:00 2001
From: minelhi <3417378192@qq.com>
Date: Fri, 11 Sep 2026 07:19:52 +0000
Subject: [PATCH 1/5] Python: add public TypedDict for AgentExecutor checkpoint
 state (#8201)

Expose AgentExecutorCheckpointState / AgentSessionCheckpointState, validate restore field types, and cover partial/malformed/forward-compatible payloads.
---
 .../packages/core/agent_framework/__init__.py |  10 +-
 .../_workflows/_agent_executor.py             | 109 +++++++++++++++---
 .../tests/workflow/test_agent_executor.py     |  72 ++++++++++++
 3 files changed, 177 insertions(+), 14 deletions(-)

diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py
index 691c35e9110..a6cfd0b72ae 100644
--- a/python/packages/core/agent_framework/__init__.py
+++ b/python/packages/core/agent_framework/__init__.py
@@ -316,7 +316,13 @@
         "vectorstoremodel",
     ),
     "._workflows._agent": ("WorkflowAgent",),
-    "._workflows._agent_executor": ("AgentExecutor", "AgentExecutorRequest", "AgentExecutorResponse"),
+    "._workflows._agent_executor": (
+        "AgentExecutor",
+        "AgentExecutorCheckpointState",
+        "AgentExecutorRequest",
+        "AgentExecutorResponse",
+        "AgentSessionCheckpointState",
+    ),
     "._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",
+    "AgentSessionCheckpointState",
     "AggregatingSkillsSource",
     "Annotation",
     "BackgroundAgentsProvider",
diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py
index cd69504c71d..24c2459d931 100644
--- a/python/packages/core/agent_framework/_workflows/_agent_executor.py
+++ b/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, NotRequired, TypedDict
 
 from agent_framework import Content
 
 from .._agents import SupportsAgentRun
 from .._sessions import AgentSession
 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, WORKFLOW_RUN_KWARGS_KEY
 from ._executor import Executor, handler
@@ -30,6 +31,79 @@
 logger = logging.getLogger(__name__)
 
 
+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]]
+
+
+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:`AgentSessionCheckpointState`).
+        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: AgentSessionCheckpointState
+    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__}."
+        )
+
+    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):
+            raise WorkflowCheckpointException(
+                f"AgentExecutor checkpoint field '{key}' must be a list, got {type(state[key]).__name__}."
+            )
+
+    if "pending_agent_requests" in state and state["pending_agent_requests"] is not None:
+        if not isinstance(state["pending_agent_requests"], dict):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'pending_agent_requests' must be a dict, "
+                f"got {type(state['pending_agent_requests']).__name__}."
+            )
+
+    if "agent_session" in state and state["agent_session"] is not None:
+        if not isinstance(state["agent_session"], dict):
+            raise WorkflowCheckpointException(
+                "AgentExecutor checkpoint field 'agent_session' must be a dict, "
+                f"got {type(state['agent_session']).__name__}."
+            )
+
+
 def _accepts_runtime_tools(agent: SupportsAgentRun) -> bool:
     """Return whether the agent run surface accepts a tools keyword."""
     try:
@@ -347,32 +421,41 @@ 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 | dict[str, Any]) -> None:
         """Restore executor state from checkpoint.
 
         Args:
-            state: Checkpoint data dict
+            state: Checkpoint payload matching :class:`AgentExecutorCheckpointState`
+                (or a compatible mapping). 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 []
 
diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py
index 9e124db1d9c..b0ba8e62bd3 100644
--- a/python/packages/core/tests/workflow/test_agent_executor.py
+++ b/python/packages/core/tests/workflow/test_agent_executor.py
@@ -578,6 +578,78 @@ async def test_checkpoint_restore_works_without_context_mode_in_state() -> None:
     assert executor._context_mode == "last_agent"  # pyright: ignore[reportPrivateUsage]
 
 
+async def test_agent_executor_checkpoint_state_public_schema_keys() -> None:
+    """Saved AgentExecutor checkpoint state exposes the public TypedDict keys."""
+    from agent_framework import AgentExecutorCheckpointState, AgentSessionCheckpointState
+
+    agent = _CountingAgent(id="schema_agent", name="SchemaAgent")
+    executor = AgentExecutor(agent)
+    executor._cache = [Message(role="user", contents=["hello"])]  # pyright: ignore[reportPrivateUsage]
+
+    state = await executor.on_checkpoint_save()
+
+    assert set(state) == {
+        "cache",
+        "full_conversation",
+        "agent_session",
+        "pending_agent_requests",
+        "pending_responses_to_agent",
+    }
+    assert isinstance(state, dict)
+    assert len(state["cache"]) == 1
+    assert "session_id" in state["agent_session"]
+    # Public types remain importable for static analysis / migrations.
+    _: type[AgentExecutorCheckpointState] = AgentExecutorCheckpointState
+    __: type[AgentSessionCheckpointState] = AgentSessionCheckpointState
+
+
+async def test_agent_executor_checkpoint_restore_missing_optional_fields() -> None:
+    """Restore accepts older partial payloads (missing optional TypedDict fields)."""
+    agent = _CountingAgent(id="partial_agent", name="PartialAgent")
+    executor = AgentExecutor(agent)
+    executor._cache = [Message(role="user", contents=["stale"])]  # pyright: ignore[reportPrivateUsage]
+
+    await executor.on_checkpoint_restore({})
+
+    assert executor._cache == []  # pyright: ignore[reportPrivateUsage]
+    assert executor._full_conversation == []  # pyright: ignore[reportPrivateUsage]
+    assert executor._pending_agent_requests == {}  # pyright: ignore[reportPrivateUsage]
+    assert executor._pending_responses_to_agent == []  # pyright: ignore[reportPrivateUsage]
+
+
+async def test_agent_executor_checkpoint_restore_rejects_malformed_fields() -> None:
+    """Restore raises WorkflowCheckpointException for wrong field types."""
+    from agent_framework import WorkflowCheckpointException
+
+    agent = _CountingAgent(id="bad_agent", name="BadAgent")
+    executor = AgentExecutor(agent)
+
+    with pytest.raises(WorkflowCheckpointException, match="cache"):
+        await executor.on_checkpoint_restore({"cache": "not-a-list"})  # type: ignore[typeddict-item]
+
+    with pytest.raises(WorkflowCheckpointException, match="agent_session"):
+        await executor.on_checkpoint_restore({"agent_session": "not-a-dict"})  # type: ignore[typeddict-item]
+
+    with pytest.raises(WorkflowCheckpointException, match="pending_agent_requests"):
+        await executor.on_checkpoint_restore({"pending_agent_requests": []})  # type: ignore[typeddict-item]
+
+
+async def test_agent_executor_checkpoint_restore_ignores_unknown_keys() -> None:
+    """Forward-compatible restore ignores unknown checkpoint keys."""
+    agent = _CountingAgent(id="fwd_agent", name="FwdAgent")
+    executor = AgentExecutor(agent)
+
+    await executor.on_checkpoint_restore(
+        {
+            "cache": [Message(role="user", contents=["ok"])],
+            "future_field": {"ignored": True},
+        }
+    )
+
+    assert len(executor._cache) == 1  # pyright: ignore[reportPrivateUsage]
+    assert executor._cache[0].text == "ok"  # pyright: ignore[reportPrivateUsage]
+
+
 # ---------------------------------------------------------------------------
 # Per-executor kwargs resolution tests
 # ---------------------------------------------------------------------------

From 07c0177a5ed1a3680fbbfe60e668649c1a40da62 Mon Sep 17 00:00:00 2001
From: LI <2484593937@qq.com>
Date: Fri, 11 Sep 2026 21:36:04 +0800
Subject: [PATCH 2/5] 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__.
---
 .../packages/core/agent_framework/__init__.py |  2 ++
 .../core/agent_framework/__init__.pyi         | 12 +++++++-
 .../core/agent_framework/_sessions.py         | 18 +++++++++++-
 .../_workflows/_agent_executor.py             | 29 ++++++-------------
 .../tests/workflow/test_agent_executor.py     |  3 +-
 5 files changed, 41 insertions(+), 23 deletions(-)

diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py
index a6cfd0b72ae..c4027f0f7f6 100644
--- a/python/packages/core/agent_framework/__init__.py
+++ b/python/packages/core/agent_framework/__init__.py
@@ -195,6 +195,7 @@
     ),
     "._sessions": (
         "AgentSession",
+        "AgentSessionDict",
         "ContextProvider",
         "FileHistoryProvider",
         "FileSessionStore",
@@ -444,6 +445,7 @@
     "AgentRunInputs",
     "AgentSession",
     "AgentSessionCheckpointState",
+    "AgentSessionDict",
     "AggregatingSkillsSource",
     "Annotation",
     "BackgroundAgentsProvider",
diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi
index b6ad94bdf11..ac2ebf332a9 100644
--- a/python/packages/core/agent_framework/__init__.pyi
+++ b/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,13 @@ 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,
+    AgentSessionCheckpointState,
+)
 from ._workflows._agent_utils import resolve_agent_id
 from ._workflows._checkpoint import (
     CheckpointID,
@@ -378,6 +385,7 @@ __all__ = [
     "AgentContext",
     "AgentEvalConverter",
     "AgentExecutor",
+    "AgentExecutorCheckpointState",
     "AgentExecutorRequest",
     "AgentExecutorResponse",
     "AgentFileStore",
@@ -391,6 +399,8 @@ __all__ = [
     "AgentResponseUpdate",
     "AgentRunInputs",
     "AgentSession",
+    "AgentSessionCheckpointState",
+    "AgentSessionDict",
     "AggregatingSkillsSource",
     "Annotation",
     "BackgroundAgentsProvider",
diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py
index 9fcce18a071..160db5228b8 100644
--- a/python/packages/core/agent_framework/_sessions.py
+++ b/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
diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py
index 24c2459d931..9e4e26a4dd1 100644
--- a/python/packages/core/agent_framework/_workflows/_agent_executor.py
+++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py
@@ -7,12 +7,12 @@
 from dataclasses import dataclass
 from typing import Any, Literal, cast
 
-from typing_extensions import Never, NotRequired, TypedDict
+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
@@ -30,18 +30,8 @@
 
 logger = logging.getLogger(__name__)
 
-
-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]]
+# Alias kept for the public PR surface; the canonical shape lives on AgentSession.
+AgentSessionCheckpointState = AgentSessionDict
 
 
 class AgentExecutorCheckpointState(TypedDict, total=False):
@@ -63,14 +53,14 @@ class AgentExecutorCheckpointState(TypedDict, total=False):
     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:`AgentSessionCheckpointState`).
+        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: AgentSessionCheckpointState
+    agent_session: AgentSessionDict
     pending_agent_requests: dict[str, Content]
     pending_responses_to_agent: list[Content]
 
@@ -442,13 +432,12 @@ async def on_checkpoint_save(self) -> AgentExecutorCheckpointState:
         )
 
     @override
-    async def on_checkpoint_restore(self, state: AgentExecutorCheckpointState | dict[str, Any]) -> None:
+    async def on_checkpoint_restore(self, state: AgentExecutorCheckpointState) -> None:
         """Restore executor state from checkpoint.
 
         Args:
-            state: Checkpoint payload matching :class:`AgentExecutorCheckpointState`
-                (or a compatible mapping). Missing known keys use empty defaults;
-                unknown keys are ignored.
+            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
diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py
index b0ba8e62bd3..4aab5259a91 100644
--- a/python/packages/core/tests/workflow/test_agent_executor.py
+++ b/python/packages/core/tests/workflow/test_agent_executor.py
@@ -580,7 +580,7 @@ async def test_checkpoint_restore_works_without_context_mode_in_state() -> None:
 
 async def test_agent_executor_checkpoint_state_public_schema_keys() -> None:
     """Saved AgentExecutor checkpoint state exposes the public TypedDict keys."""
-    from agent_framework import AgentExecutorCheckpointState, AgentSessionCheckpointState
+    from agent_framework import AgentExecutorCheckpointState, AgentSessionCheckpointState, AgentSessionDict
 
     agent = _CountingAgent(id="schema_agent", name="SchemaAgent")
     executor = AgentExecutor(agent)
@@ -601,6 +601,7 @@ async def test_agent_executor_checkpoint_state_public_schema_keys() -> None:
     # Public types remain importable for static analysis / migrations.
     _: type[AgentExecutorCheckpointState] = AgentExecutorCheckpointState
     __: type[AgentSessionCheckpointState] = AgentSessionCheckpointState
+    assert AgentSessionCheckpointState is AgentSessionDict
 
 
 async def test_agent_executor_checkpoint_restore_missing_optional_fields() -> None:

From 01fcd2354ee23b9c080a44c9f256a706efdd251f Mon Sep 17 00:00:00 2001
From: minelhi <3417378192@qq.com>
Date: Sun, 13 Sep 2026 00:41:39 +0800
Subject: [PATCH 3/5] 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.
---
 .../_workflows/_agent_executor.py             | 60 ++++++++++++++++---
 .../tests/workflow/test_agent_executor.py     | 20 ++++++-
 2 files changed, 71 insertions(+), 9 deletions(-)

diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py
index 9e4e26a4dd1..5ab6a7bac15 100644
--- a/python/packages/core/agent_framework/_workflows/_agent_executor.py
+++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py
@@ -72,25 +72,69 @@ def _validate_agent_executor_checkpoint_state(state: Mapping[str, Any]) -> None:
             f"AgentExecutor checkpoint state must be a mapping, got {type(state).__name__}."
         )
 
-    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):
+    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(state[key]).__name__}."
+                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:
-        if not isinstance(state["pending_agent_requests"], dict):
+        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(state['pending_agent_requests']).__name__}."
+                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:
-        if not isinstance(state["agent_session"], dict):
+        session = state["agent_session"]
+        if not isinstance(session, dict):
             raise WorkflowCheckpointException(
                 "AgentExecutor checkpoint field 'agent_session' must be a dict, "
-                f"got {type(state['agent_session']).__name__}."
+                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__}."
             )
 
 
diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py
index 4aab5259a91..f06b0c2a13d 100644
--- a/python/packages/core/tests/workflow/test_agent_executor.py
+++ b/python/packages/core/tests/workflow/test_agent_executor.py
@@ -620,7 +620,7 @@ async def test_agent_executor_checkpoint_restore_missing_optional_fields() -> No
 
 async def test_agent_executor_checkpoint_restore_rejects_malformed_fields() -> None:
     """Restore raises WorkflowCheckpointException for wrong field types."""
-    from agent_framework import WorkflowCheckpointException
+    from agent_framework import Content, WorkflowCheckpointException
 
     agent = _CountingAgent(id="bad_agent", name="BadAgent")
     executor = AgentExecutor(agent)
@@ -628,12 +628,30 @@ async def test_agent_executor_checkpoint_restore_rejects_malformed_fields() -> N
     with pytest.raises(WorkflowCheckpointException, match="cache"):
         await executor.on_checkpoint_restore({"cache": "not-a-list"})  # type: ignore[typeddict-item]
 
+    with pytest.raises(WorkflowCheckpointException, match=r"'cache'\[0\]"):
+        await executor.on_checkpoint_restore({"cache": ["not-a-message"]})  # type: ignore[typeddict-item]
+
+    with pytest.raises(WorkflowCheckpointException, match="pending_responses_to_agent"):
+        await executor.on_checkpoint_restore({"pending_responses_to_agent": ["bad"]})  # type: ignore[typeddict-item]
+
     with pytest.raises(WorkflowCheckpointException, match="agent_session"):
         await executor.on_checkpoint_restore({"agent_session": "not-a-dict"})  # type: ignore[typeddict-item]
 
+    with pytest.raises(WorkflowCheckpointException, match="agent_session.session_id"):
+        await executor.on_checkpoint_restore({"agent_session": {}})  # type: ignore[typeddict-item]
+
+    with pytest.raises(WorkflowCheckpointException, match="agent_session.session_id"):
+        await executor.on_checkpoint_restore({"agent_session": {"session_id": 1}})  # type: ignore[typeddict-item]
+
     with pytest.raises(WorkflowCheckpointException, match="pending_agent_requests"):
         await executor.on_checkpoint_restore({"pending_agent_requests": []})  # type: ignore[typeddict-item]
 
+    with pytest.raises(WorkflowCheckpointException, match="pending_agent_requests"):
+        await executor.on_checkpoint_restore({"pending_agent_requests": {1: Content(type="text", text="x")}})  # type: ignore[typeddict-item]
+
+    with pytest.raises(WorkflowCheckpointException, match="pending_agent_requests"):
+        await executor.on_checkpoint_restore({"pending_agent_requests": {"req": "not-content"}})  # type: ignore[typeddict-item]
+
 
 async def test_agent_executor_checkpoint_restore_ignores_unknown_keys() -> None:
     """Forward-compatible restore ignores unknown checkpoint keys."""

From 893438264cdc7876251f61fcbf9ca5962bcab099 Mon Sep 17 00:00:00 2001
From: LI <2484593937@qq.com>
Date: Mon, 14 Sep 2026 20:36:00 +0800
Subject: [PATCH 4/5] 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
---
 .../packages/core/agent_framework/__init__.py |  2 --
 .../core/agent_framework/__init__.pyi         |  2 --
 .../_workflows/_agent_executor.py             | 26 +++++++++++++++----
 .../tests/workflow/test_agent_executor.py     |  9 ++++---
 4 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py
index c4027f0f7f6..3aaab93f164 100644
--- a/python/packages/core/agent_framework/__init__.py
+++ b/python/packages/core/agent_framework/__init__.py
@@ -322,7 +322,6 @@
         "AgentExecutorCheckpointState",
         "AgentExecutorRequest",
         "AgentExecutorResponse",
-        "AgentSessionCheckpointState",
     ),
     "._workflows._agent_utils": ("resolve_agent_id",),
     "._workflows._checkpoint": (
@@ -444,7 +443,6 @@
     "AgentResponseUpdate",
     "AgentRunInputs",
     "AgentSession",
-    "AgentSessionCheckpointState",
     "AgentSessionDict",
     "AggregatingSkillsSource",
     "Annotation",
diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi
index ac2ebf332a9..20f1ef18545 100644
--- a/python/packages/core/agent_framework/__init__.pyi
+++ b/python/packages/core/agent_framework/__init__.pyi
@@ -279,7 +279,6 @@ from ._workflows._agent_executor import (
     AgentExecutorCheckpointState,
     AgentExecutorRequest,
     AgentExecutorResponse,
-    AgentSessionCheckpointState,
 )
 from ._workflows._agent_utils import resolve_agent_id
 from ._workflows._checkpoint import (
@@ -399,7 +398,6 @@ __all__ = [
     "AgentResponseUpdate",
     "AgentRunInputs",
     "AgentSession",
-    "AgentSessionCheckpointState",
     "AgentSessionDict",
     "AggregatingSkillsSource",
     "Annotation",
diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py
index 5ab6a7bac15..6bed2fee914 100644
--- a/python/packages/core/agent_framework/_workflows/_agent_executor.py
+++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py
@@ -30,9 +30,6 @@
 
 logger = logging.getLogger(__name__)
 
-# Alias kept for the public PR surface; the canonical shape lives on AgentSession.
-AgentSessionCheckpointState = AgentSessionDict
-
 
 class AgentExecutorCheckpointState(TypedDict, total=False):
     """Public schema for state saved and restored by :class:`AgentExecutor`.
@@ -136,6 +133,23 @@ def _validate_agent_executor_checkpoint_state(state: Mapping[str, Any]) -> None:
                 "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:
@@ -500,8 +514,10 @@ async def on_checkpoint_restore(self, state: AgentExecutorCheckpointState) -> No
             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()
 
diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py
index f06b0c2a13d..2dfb2385964 100644
--- a/python/packages/core/tests/workflow/test_agent_executor.py
+++ b/python/packages/core/tests/workflow/test_agent_executor.py
@@ -580,7 +580,7 @@ async def test_checkpoint_restore_works_without_context_mode_in_state() -> None:
 
 async def test_agent_executor_checkpoint_state_public_schema_keys() -> None:
     """Saved AgentExecutor checkpoint state exposes the public TypedDict keys."""
-    from agent_framework import AgentExecutorCheckpointState, AgentSessionCheckpointState, AgentSessionDict
+    from agent_framework import AgentExecutorCheckpointState, AgentSessionDict
 
     agent = _CountingAgent(id="schema_agent", name="SchemaAgent")
     executor = AgentExecutor(agent)
@@ -600,8 +600,8 @@ async def test_agent_executor_checkpoint_state_public_schema_keys() -> None:
     assert "session_id" in state["agent_session"]
     # Public types remain importable for static analysis / migrations.
     _: type[AgentExecutorCheckpointState] = AgentExecutorCheckpointState
-    __: type[AgentSessionCheckpointState] = AgentSessionCheckpointState
-    assert AgentSessionCheckpointState is AgentSessionDict
+    __: type[AgentSessionDict] = AgentSessionDict
+    assert isinstance(state["agent_session"], dict)
 
 
 async def test_agent_executor_checkpoint_restore_missing_optional_fields() -> None:
@@ -643,6 +643,9 @@ async def test_agent_executor_checkpoint_restore_rejects_malformed_fields() -> N
     with pytest.raises(WorkflowCheckpointException, match="agent_session.session_id"):
         await executor.on_checkpoint_restore({"agent_session": {"session_id": 1}})  # type: ignore[typeddict-item]
 
+    with pytest.raises(WorkflowCheckpointException, match="agent_session.state"):
+        await executor.on_checkpoint_restore({"agent_session": {"session_id": "s", "state": []}})  # type: ignore[typeddict-item]
+
     with pytest.raises(WorkflowCheckpointException, match="pending_agent_requests"):
         await executor.on_checkpoint_restore({"pending_agent_requests": []})  # type: ignore[typeddict-item]
 

From 76985acdd883feda0263fdd421f63b5c76457463 Mon Sep 17 00:00:00 2001
From: minelhi <3417378192@qq.com>
Date: Tue, 15 Sep 2026 20:25:03 +0800
Subject: [PATCH 5/5] 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.
---
 .../core/agent_framework/_sessions.py         | 12 ++++-
 .../_workflows/_agent_executor.py             | 50 ++++++++++---------
 2 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py
index 160db5228b8..b6c03b35245 100644
--- a/python/packages/core/agent_framework/_sessions.py
+++ b/python/packages/core/agent_framework/_sessions.py
@@ -1706,7 +1706,10 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[
 
 
 class AgentSessionDict(TypedDict):
-    """Serialized :class:`AgentSession` payload produced by :meth:`AgentSession.to_dict`.
+    """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`.
@@ -1760,9 +1763,14 @@ def session_id(self) -> str:
         """The unique identifier for this session."""
         return self._session_id
 
-    def to_dict(self) -> AgentSessionDict:
+    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
diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py
index fb96872fd03..d372e6e1560 100644
--- a/python/packages/core/agent_framework/_workflows/_agent_executor.py
+++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py
@@ -78,7 +78,8 @@ def _validate_agent_executor_checkpoint_state(state: Mapping[str, Any]) -> None:
             raise WorkflowCheckpointException(
                 f"AgentExecutor checkpoint field '{key}' must be a list, got {type(value).__name__}."
             )
-        for index, item in enumerate(value):
+        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, "
@@ -86,12 +87,13 @@ def _validate_agent_executor_checkpoint_state(state: Mapping[str, Any]) -> None:
                 )
 
     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):
+        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).__name__}."
+                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(
@@ -101,12 +103,13 @@ def _validate_agent_executor_checkpoint_state(state: Mapping[str, Any]) -> None:
                 )
 
     if "pending_agent_requests" in state and state["pending_agent_requests"] is not None:
-        pending = state["pending_agent_requests"]
-        if not isinstance(pending, dict):
+        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).__name__}."
+                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(
@@ -121,12 +124,13 @@ def _validate_agent_executor_checkpoint_state(state: Mapping[str, Any]) -> None:
                 )
 
     if "agent_session" in state and state["agent_session"] is not None:
-        session = state["agent_session"]
-        if not isinstance(session, dict):
+        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).__name__}."
+                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(
@@ -469,28 +473,28 @@ async def _cancel_pending_request(
             await self._resume_with_pending_responses(ctx)
 
     @override
-    async def on_checkpoint_save(self) -> AgentExecutorCheckpointState:
+    async def on_checkpoint_save(self) -> dict[str, Any]:
         """Capture current executor state for checkpointing.
 
         NOTE: if the session uses service-side storage, the full session state
         may not be serialized locally.
 
         Returns:
-            :class:`AgentExecutorCheckpointState` with cache, conversation, session,
-            and pending request/response fields.
+            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 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,
-        )
+        return {
+            "cache": self._cache,
+            "full_conversation": self._full_conversation,
+            "agent_session": self._session.to_dict(),
+            "pending_agent_requests": self._pending_agent_requests,
+            "pending_responses_to_agent": self._pending_responses_to_agent,
+        }
 
     @override
-    async def on_checkpoint_restore(self, state: AgentExecutorCheckpointState) -> None:
+    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
         """Restore executor state from checkpoint.
 
         Args: