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 7fe4132cf47..20aec002e1b 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 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 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: