diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index ea9e4577a2..a341ef70eb 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -196,6 +196,7 @@ ), "._sessions": ( "AgentSession", + "AgentSessionDict", "ContextProvider", "FileHistoryProvider", "FileSessionStore", @@ -317,7 +318,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", @@ -424,6 +430,7 @@ "AgentContext", "AgentEvalConverter", "AgentExecutor", + "AgentExecutorCheckpointState", "AgentExecutorRequest", "AgentExecutorResponse", "AgentFileStore", @@ -437,6 +444,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 350a458a38..5a6912c791 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, @@ -379,6 +385,7 @@ __all__ = [ "AgentContext", "AgentEvalConverter", "AgentExecutor", + "AgentExecutorCheckpointState", "AgentExecutorRequest", "AgentExecutorResponse", "AgentFileStore", @@ -392,6 +399,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 9fcce18a07..b6c03b3524 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,24 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[ ) +class AgentSessionDict(TypedDict): + """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`. + ``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. @@ -1747,6 +1766,11 @@ def session_id(self) -> str: 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 6867bd73a0..d372e6e156 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 prepare_agent_run_args, resolve_agent_id, resolve_executor_kwargs from ._const import INTERNAL_SOURCE_ID, RESOLVED_WORKFLOW_RUN_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY from ._executor import Executor, handler @@ -30,6 +31,131 @@ 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__}." + ) + 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, " + f"got {type(item).__name__}." + ) + + if "pending_responses_to_agent" in state and state["pending_responses_to_agent"] is not None: + 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_raw).__name__}." + ) + responses = cast(list[Any], responses_raw) + 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_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_raw).__name__}." + ) + pending = cast(dict[Any, Any], pending_raw) + 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_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_raw).__name__}." + ) + session = cast(dict[str, Any], session_raw) + 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: @@ -354,14 +480,15 @@ async def on_checkpoint_save(self) -> dict[str, Any]: may not be serialized locally. Returns: - Dict containing serialized cache and session state + 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 { "cache": self._cache, "full_conversation": self._full_conversation, - "agent_session": serialized_session, + "agent_session": self._session.to_dict(), "pending_agent_requests": self._pending_agent_requests, "pending_responses_to_agent": self._pending_responses_to_agent, } @@ -371,8 +498,15 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> 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 +518,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 18d93cf37d..cda6bdad93 100644 --- a/python/packages/core/tests/workflow/test_agent_executor.py +++ b/python/packages/core/tests/workflow/test_agent_executor.py @@ -611,6 +611,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 # ---------------------------------------------------------------------------