fix(checkpoint): address review on session TypedDict ownership (#8201)

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",

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",

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

@@ -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 resolve_agent_id
@@ -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

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: