diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 691c35e9110..3aaab93f164 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", @@ -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", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index b6ad94bdf11..20f1ef18545 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,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", 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 1c3a7203f49..326c04fe932 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, 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() diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py index 9e124db1d9c..2dfb2385964 100644 --- a/python/packages/core/tests/workflow/test_agent_executor.py +++ b/python/packages/core/tests/workflow/test_agent_executor.py @@ -578,6 +578,100 @@ 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, AgentSessionDict + + 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[AgentSessionDict] = AgentSessionDict + assert isinstance(state["agent_session"], dict) + + +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 Content, 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=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="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] + + 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.""" + 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 # ---------------------------------------------------------------------------