From 51a6201b5a33981c1ecf8027af99220291fc5378 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 3 Aug 2026 10:39:38 +0900 Subject: [PATCH 1/2] Python: Give the AG-UI Thread Snapshot lifecycle a single owner module Both the agent and workflow runners independently implemented the thread snapshot lifecycle: hydration replay, the load-once stored read, resume message seeding, the stored/request/deferred-default state overlay, and the save whose storage failures must never surface on an already-streamed run. The two copies had already drifted in small ways (one hydrate helper re-checked a store the caller had verified; the two cancelled-resume-id helpers differed on missing-id handling). Introduce ThreadSnapshotSession in _snapshot_session.py as the one owner of that lifecycle, opened once per run and inert when no store or scope is configured so callers stop branching on configuration. Rewire both runners onto it, consolidate _cancelled_resume_interrupt_ids in _run_common (defensive variant) and _event_messages_to_snapshot_dicts in the new module, and delete the superseded per-runner copies. The session interface is covered by dedicated tests; existing suites pin runner behavior. Public exports are unchanged. --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 160 ++-------- .../agent_framework_ag_ui/_run_common.py | 12 + .../_snapshot_session.py | 183 ++++++++++++ .../ag-ui/agent_framework_ag_ui/_workflow.py | 159 +++------- .../tests/ag_ui/test_snapshot_session.py | 275 ++++++++++++++++++ 5 files changed, 539 insertions(+), 250 deletions(-) create mode 100644 python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py create mode 100644 python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index ab074ae15cd..10ce7adcb87 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -58,6 +58,7 @@ _approval_interrupt_for_function_call, # type: ignore _approval_steps_response_schema, # type: ignore _build_run_finished_event, # type: ignore + _cancelled_resume_interrupt_ids, # type: ignore _close_reasoning_block, # type: ignore _emit_content, # type: ignore _extract_resume_payload, # type: ignore @@ -74,8 +75,8 @@ _DEFAULT_STATE_INPUT_KEY, _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshot, - _clear_thread_snapshot_interrupt, ) +from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts from ._utils import ( canonical_function_arguments, convert_agui_tools_to_agent_framework, @@ -775,18 +776,6 @@ def _approval_state_tool_call_ids( return call_ids -def _cancelled_resume_interrupt_ids(resume_payload: Any) -> set[str]: - """Return cancelled canonical resume interrupt ids.""" - interrupt_ids: set[str] = set() - for interrupt in _normalize_resume_interrupts(resume_payload): - if interrupt.get("status") != "cancelled": - continue - interrupt_id = interrupt.get("id") - if interrupt_id: - interrupt_ids.add(str(interrupt_id)) - return interrupt_ids - - def _tool_approval_state_exists_for_cancelled_resume( resume_payload: Any, approval_state_store: InMemoryAGUIApprovalStateStore | None, @@ -1679,14 +1668,6 @@ def _build_messages_snapshot( return MessagesSnapshotEvent(messages=all_messages) # type: ignore[arg-type] -def _event_messages_to_snapshot_dicts(messages: list[Any]) -> list[dict[str, Any]]: - """Convert AG-UI message event models back to plain snapshot dictionaries.""" - safe_messages = make_json_safe(messages) - if not isinstance(safe_messages, list): - return [] - return [cast(dict[str, Any], message) for message in safe_messages if isinstance(message, dict)] - - def _text_events_to_snapshot_messages(events: list[BaseEvent]) -> list[dict[str, Any]]: """Convert streamed text-message events into snapshot message dictionaries.""" messages: list[dict[str, Any]] = [] @@ -1703,67 +1684,6 @@ def _text_events_to_snapshot_messages(events: list[BaseEvent]) -> list[dict[str, return [message for message in messages if message.get("content")] -async def _hydrate_thread_snapshot( - *, - config: AgentConfig, - scope: str, - thread_id: str, - run_id: str, -) -> AsyncGenerator[BaseEvent]: - """Replay the latest stored AG-UI Thread Snapshot without invoking the agent.""" - yield RunStartedEvent(run_id=run_id, thread_id=thread_id) - if config.snapshot_store is None: - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) - return - - snapshot = await config.snapshot_store.get(scope=scope, thread_id=thread_id) - if snapshot is None: - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) - return - - if snapshot.state is not None: - yield StateSnapshotEvent(snapshot=snapshot.state) - if snapshot.messages: - yield MessagesSnapshotEvent(messages=snapshot.messages) # type: ignore[arg-type] - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=snapshot.interrupt) - - -async def _save_thread_snapshot( - *, - config: AgentConfig, - scope: str | None, - thread_id: str, - messages: list[dict[str, Any]], - state: dict[str, Any] | None, - interrupt: list[dict[str, Any]] | None, - session_state: dict[str, Any] | None, -) -> None: - """Save the latest AG-UI Thread Snapshot when persistence is configured.""" - if config.snapshot_store is None or scope is None: - return - - try: - await config.snapshot_store.save( - scope=scope, - thread_id=thread_id, - snapshot=AGUIThreadSnapshot( - messages=messages, - state=state, - interrupt=interrupt, - session_state=session_state, - ), - ) - except Exception: - # The run itself already streamed successfully; a transient store failure - # must not surface as RUN_ERROR for a completed run. The previous snapshot - # stays available for hydration. - logger.exception( - "Failed to save AG-UI Thread Snapshot for scope=%s thread_id=%s; keeping previous snapshot.", - scope, - thread_id, - ) - - def _restore_session_continuation_state(session: AgentSession, snapshot: AGUIThreadSnapshot | None) -> None: """Restore typed private state from trusted snapshot storage.""" if snapshot is None or snapshot.session_state is None: @@ -1888,50 +1808,38 @@ async def run_agent_stream( available_interrupts = input_data.get("available_interrupts") or input_data.get("availableInterrupts") raw_messages: list[dict[str, Any]] = input_data.get("messages", []) or [] resume_payload = _extract_resume_payload(input_data) - if config.snapshot_store is not None and snapshot_scope is not None and not raw_messages and resume_payload is None: - async for event in _hydrate_thread_snapshot( - config=config, - scope=snapshot_scope, - thread_id=thread_id, - run_id=run_id, - ): + snapshot_session = await ThreadSnapshotSession.open( + store=config.snapshot_store, + scope=snapshot_scope, + thread_id=thread_id, + ) + if snapshot_session.enabled and not raw_messages and resume_payload is None: + async for event in snapshot_session.hydrate_events(run_id=run_id): yield event return - stored_snapshot: AGUIThreadSnapshot | None = None + stored_snapshot = snapshot_session.stored stored_pending_approval_interrupt_ids: set[str] = set() seeded_resume_from_snapshot = False - if config.snapshot_store is not None and snapshot_scope is not None: - stored_snapshot = await config.snapshot_store.get(scope=snapshot_scope, thread_id=thread_id) - if stored_snapshot is not None: - stored_pending_approval_interrupt_ids = _stored_pending_approval_interrupt_ids(stored_snapshot.interrupt) - if stored_snapshot is not None and resume_payload is not None and stored_pending_approval_interrupt_ids: - raw_messages = [copy.deepcopy(message) for message in stored_snapshot.messages] + raw_messages + if stored_snapshot is not None: + stored_pending_approval_interrupt_ids = _stored_pending_approval_interrupt_ids(stored_snapshot.interrupt) + if resume_payload is not None and stored_pending_approval_interrupt_ids: + raw_messages = snapshot_session.resume_seeded_messages(raw_messages) seeded_resume_from_snapshot = True - elif stored_snapshot is not None: + else: raw_messages = _reconstruct_messages_from_thread_snapshot( stored_messages=stored_snapshot.messages, incoming_messages=raw_messages, stored_interrupt=stored_snapshot.interrupt, ) - # Initialize flow state with stored state plus request-provided overrides. + # Initialize flow state with stored state plus request-provided overrides; + # endpoint-deferred defaults apply only to keys missing from both. flow = FlowState() - request_state = input_data.get("state") - if stored_snapshot is not None and stored_snapshot.state is not None: - flow.current_state = dict(stored_snapshot.state) - if isinstance(request_state, dict): - flow.current_state.update(request_state) - elif isinstance(request_state, dict): - flow.current_state = dict(request_state) - - # Apply endpoint-deferred defaults only for keys missing from both the stored - # snapshot state and the request state, so defaults never reset persisted state. - deferred_default_state = cast(dict[str, Any] | None, input_data.get(_DEFAULT_STATE_INPUT_KEY)) - if deferred_default_state: - for key, value in deferred_default_state.items(): - if key not in flow.current_state: - flow.current_state[key] = copy.deepcopy(value) + flow.current_state = snapshot_session.effective_state( + request_state=input_data.get("state"), + deferred_defaults=cast(dict[str, Any] | None, input_data.get(_DEFAULT_STATE_INPUT_KEY)), + ) # Apply schema defaults for missing state keys if state_schema: @@ -1971,13 +1879,7 @@ async def run_agent_stream( if should_clear_tool_approval_state: _clear_tool_approval_state(approval_state_store, approval_thread_id) if resume_error_code == "APPROVAL_RESUME_CANCELLED": - if config.snapshot_store is not None and snapshot_scope is not None: - await _clear_thread_snapshot_interrupt( - snapshot_store=config.snapshot_store, - scope=snapshot_scope, - thread_id=thread_id, - interrupt_ids=cancelled_resume_ids or None, - ) + await snapshot_session.clear_interrupts(interrupt_ids=cancelled_resume_ids or None) yield resume_error return resume_messages = _resume_to_tool_messages(resume_payload, exclude_interrupt_ids=handled_resume_ids) @@ -2094,14 +1996,11 @@ async def run_agent_stream( # Persist the completed confirmation turn with interrupt=None so hydration # does not replay the stale pending interrupt after the user responded. persisted_messages = snapshot_messages + _text_events_to_snapshot_messages(confirmation_events) - if resume_payload is not None and stored_snapshot is not None and not seeded_resume_from_snapshot: + if resume_payload is not None and not seeded_resume_from_snapshot: # Generic resume requests carry only the synthesized response, so prepend # stored history unless this run already seeded raw messages from it. - persisted_messages = [copy.deepcopy(message) for message in stored_snapshot.messages] + persisted_messages - await _save_thread_snapshot( - config=config, - scope=snapshot_scope, - thread_id=thread_id, + persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages) + await snapshot_session.save( messages=persisted_messages, state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None, interrupt=None, @@ -2397,14 +2296,11 @@ async def run_agent_stream( # Always emit RunFinished - confirm_changes tool call is complete (Start -> Args -> End) # The UI will show confirmation dialog and send a new request when user responds persisted_messages = latest_messages_snapshot - if resume_payload is not None and stored_snapshot is not None and not seeded_resume_from_snapshot: + if resume_payload is not None and not seeded_resume_from_snapshot: # Generic resume requests carry only the synthesized response, so prepend # stored history unless this run already seeded raw messages from it. - persisted_messages = [copy.deepcopy(message) for message in stored_snapshot.messages] + persisted_messages - await _save_thread_snapshot( - config=config, - scope=snapshot_scope, - thread_id=thread_id, + persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages) + await snapshot_session.save( messages=persisted_messages, state=latest_state_snapshot, interrupt=flow.interrupts or None, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index da7cbd7d36f..706800cedb4 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -108,6 +108,18 @@ def _normalize_resume_interrupts(resume_payload: Any) -> list[dict[str, Any]]: return normalized +def _cancelled_resume_interrupt_ids(resume_payload: Any) -> set[str]: + """Return cancelled canonical resume interrupt ids.""" + interrupt_ids: set[str] = set() + for interrupt in _normalize_resume_interrupts(resume_payload): + if interrupt.get("status") != "cancelled": + continue + interrupt_id = interrupt.get("id") + if interrupt_id: + interrupt_ids.add(str(interrupt_id)) + return interrupt_ids + + def _extract_resume_payload(input_data: dict[str, Any]) -> Any: """Extract resume payload from standard and forwarded-props request locations.""" resume_payload = input_data.get("resume") diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py new file mode 100644 index 00000000000..da4360832ac --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py @@ -0,0 +1,183 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Single owner of the AG-UI Thread Snapshot lifecycle for a run. + +A ThreadSnapshotSession is opened once per run and owns every interaction +with the AG-UI Thread Snapshot store: the load-once read, hydration replay, +the effective-state overlay, resume message seeding, and the save whose +storage failures must never surface on an already-streamed run. +""" + +from __future__ import annotations + +import copy +import logging +from collections.abc import AsyncGenerator +from typing import Any, cast + +from ag_ui.core import ( + BaseEvent, + MessagesSnapshotEvent, + RunStartedEvent, + StateSnapshotEvent, +) + +from ._run_common import _build_run_finished_event +from ._snapshots import ( + AGUIThreadSnapshot, + AGUIThreadSnapshotStore, + _clear_thread_snapshot_interrupt, +) +from ._utils import make_json_safe + +logger = logging.getLogger(__name__) + + +def _event_messages_to_snapshot_dicts(messages: list[Any]) -> list[dict[str, Any]]: + """Convert AG-UI message event models back to plain snapshot dictionaries.""" + safe_messages = make_json_safe(messages) + if not isinstance(safe_messages, list): + return [] + return [cast(dict[str, Any], message) for message in safe_messages if isinstance(message, dict)] + + +class ThreadSnapshotSession: + """Per-run owner of one scoped AG-UI Thread Snapshot. + + Open with :meth:`open`. When the store or scope is not configured the + session is disabled: reads return nothing and writes are no-ops, so + callers never branch on configuration themselves. + """ + + def __init__( + self, + *, + store: AGUIThreadSnapshotStore | None, + scope: str | None, + thread_id: str, + stored: AGUIThreadSnapshot | None, + ) -> None: + self._store = store + self._scope = scope + self._thread_id = thread_id + self._stored = stored + + @classmethod + async def open( + cls, + *, + store: AGUIThreadSnapshotStore | None, + scope: str | None, + thread_id: str, + ) -> ThreadSnapshotSession: + """Open the session, loading the stored snapshot once when scoped.""" + stored: AGUIThreadSnapshot | None = None + if store is not None and scope is not None: + stored = await store.get(scope=scope, thread_id=thread_id) + return cls(store=store, scope=scope, thread_id=thread_id, stored=stored) + + @property + def enabled(self) -> bool: + """Whether a store and scope are both configured.""" + return self._store is not None and self._scope is not None + + @property + def stored(self) -> AGUIThreadSnapshot | None: + """The snapshot loaded at open, or ``None``.""" + return self._stored + + async def hydrate_events(self, *, run_id: str) -> AsyncGenerator[BaseEvent]: + """Replay the stored snapshot as a complete run without invoking the agent.""" + yield RunStartedEvent(run_id=run_id, thread_id=self._thread_id) + snapshot = self._stored + if snapshot is None: + yield _build_run_finished_event(run_id=run_id, thread_id=self._thread_id) + return + + if snapshot.state is not None: + yield StateSnapshotEvent(snapshot=snapshot.state) + if snapshot.messages: + yield MessagesSnapshotEvent(messages=snapshot.messages) # type: ignore[arg-type] + yield _build_run_finished_event(run_id=run_id, thread_id=self._thread_id, interrupts=snapshot.interrupt) + + def effective_state( + self, + *, + request_state: Any, + deferred_defaults: dict[str, Any] | None, + ) -> dict[str, Any]: + """Overlay request state onto stored state, then fill missing keys with defaults. + + Request values overlay stored values with the same keys; endpoint-deferred + defaults apply only to keys missing from both, so defaults never reset + persisted state. Default values are copied, never aliased. + """ + state: dict[str, Any] = {} + if self._stored is not None and self._stored.state is not None: + state.update(self._stored.state) + if isinstance(request_state, dict): + state.update(request_state) + if deferred_defaults: + for key, value in deferred_defaults.items(): + if key not in state: + state[key] = copy.deepcopy(value) + return state + + def resume_seeded_messages(self, incoming: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Prepend copies of stored thread history to a resume request's messages. + + Resume requests carry only the synthesized interrupt response; seeding + with stored history keeps the persisted thread from being truncated. + """ + if self._stored is None: + return incoming + return [copy.deepcopy(message) for message in self._stored.messages] + incoming + + async def save( + self, + *, + messages: list[dict[str, Any]], + state: dict[str, Any] | None, + interrupt: list[dict[str, Any]] | None, + session_state: dict[str, Any] | None, + ) -> None: + """Commit the latest thread snapshot in one write when persistence is configured. + + The run has already streamed by the time this is called, so a store + failure is logged and swallowed; the previous snapshot stays + authoritative for hydration. + """ + if self._store is None or self._scope is None: + return + try: + await self._store.save( + scope=self._scope, + thread_id=self._thread_id, + snapshot=AGUIThreadSnapshot( + messages=messages, + state=state, + interrupt=interrupt, + session_state=session_state, + ), + ) + except Exception: + logger.exception( + "Failed to save AG-UI Thread Snapshot for scope=%s thread_id=%s; keeping previous snapshot.", + self._scope, + self._thread_id, + ) + + async def clear_interrupts(self, *, interrupt_ids: set[str] | None = None) -> None: + """Remove completed interrupts from the latest stored snapshot. + + Clears all interrupts when ``interrupt_ids`` is omitted. Failures are + logged and swallowed for the same reason as :meth:`save`. + """ + if self._store is None or self._scope is None: + return + await _clear_thread_snapshot_interrupt( + snapshot_store=self._store, + scope=self._scope, + thread_id=self._thread_id, + interrupt_ids=interrupt_ids, + ) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index 20f44361c91..abdfc96a386 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -4,7 +4,6 @@ from __future__ import annotations -import copy import logging import uuid from collections.abc import AsyncGenerator, Callable @@ -15,7 +14,6 @@ MessagesSnapshotEvent, RunErrorEvent, RunFinishedEvent, - RunStartedEvent, StateSnapshotEvent, TextMessageContentEvent, TextMessageEndEvent, @@ -30,17 +28,16 @@ from ._feature_usage import FeatureIndex from ._message_adapters import agui_messages_to_snapshot_format from ._run_common import ( - _build_run_finished_event, + _cancelled_resume_interrupt_ids, _extract_resume_payload, - _normalize_resume_interrupts, _reconstruct_messages_from_thread_snapshot, ) +from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshot, AGUIThreadSnapshotStore, - _clear_thread_snapshot_interrupt, ) from ._utils import generate_event_id, make_json_safe from ._workflow_run import run_workflow_stream @@ -50,23 +47,6 @@ WorkflowFactory = Callable[[str], Workflow] -def _cancelled_resume_interrupt_ids(resume_payload: Any) -> set[str]: - """Return cancelled interrupt ids from a resume payload.""" - return { - str(interrupt["id"]) - for interrupt in _normalize_resume_interrupts(resume_payload) - if interrupt.get("status") == "cancelled" - } - - -def _event_messages_to_snapshot_dicts(messages: list[Any]) -> list[dict[str, Any]]: - """Convert AG-UI message event models to plain snapshot dictionaries.""" - safe_messages = make_json_safe(messages) - if not isinstance(safe_messages, list): - return [] - return [cast(dict[str, Any], message) for message in safe_messages if isinstance(message, dict)] - - class _WorkflowSnapshotBuilder: """Capture replayable workflow protocol output without retaining raw events.""" @@ -198,27 +178,6 @@ def _flush_open_text_message(self) -> None: self._open_text_message = None -async def _hydrate_workflow_thread_snapshot( - *, - snapshot_store: AGUIThreadSnapshotStore, - scope: str, - thread_id: str, - run_id: str, -) -> AsyncGenerator[BaseEvent]: - """Replay the latest stored workflow AG-UI Thread Snapshot without invoking the workflow.""" - yield RunStartedEvent(run_id=run_id, thread_id=thread_id) - snapshot = await snapshot_store.get(scope=scope, thread_id=thread_id) - if snapshot is None: - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) - return - - if snapshot.state is not None: - yield StateSnapshotEvent(snapshot=snapshot.state) - if snapshot.messages: - yield MessagesSnapshotEvent(messages=snapshot.messages) # type: ignore[arg-type] - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=snapshot.interrupt) - - class AgentFrameworkWorkflow: """Base AG-UI workflow wrapper. @@ -306,60 +265,42 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY)) raw_messages = list(cast(list[dict[str, Any]], input_data.get("messages", []) or [])) resume_payload = _extract_resume_payload(input_data) - snapshot_store = self.snapshot_store - - if snapshot_store is not None and snapshot_scope is not None and not raw_messages and resume_payload is None: - async for event in _hydrate_workflow_thread_snapshot( - snapshot_store=snapshot_store, - scope=snapshot_scope, - thread_id=thread_id, - run_id=run_id, - ): + snapshot_session = await ThreadSnapshotSession.open( + store=self.snapshot_store, + scope=snapshot_scope, + thread_id=thread_id, + ) + + if snapshot_session.enabled and not raw_messages and resume_payload is None: + async for event in snapshot_session.hydrate_events(run_id=run_id): yield event return - # Load the stored snapshot for follow-up turns so the workflow runs with the - # full persisted thread history instead of just the latest request messages. - stored_snapshot: AGUIThreadSnapshot | None = None - if snapshot_store is not None and snapshot_scope is not None: - stored_snapshot = await snapshot_store.get(scope=snapshot_scope, thread_id=thread_id) - if stored_snapshot is not None and resume_payload is None: - raw_messages = _reconstruct_messages_from_thread_snapshot( - stored_messages=stored_snapshot.messages, - incoming_messages=raw_messages, - stored_interrupt=stored_snapshot.interrupt, - ) - input_data["messages"] = raw_messages - - # Merge stored state with request overrides, then fill endpoint-deferred - # defaults only for keys missing from both. - request_state = input_data.get("state") - deferred_default_state = cast(dict[str, Any] | None, input_data.get(_DEFAULT_STATE_INPUT_KEY)) - effective_state: dict[str, Any] = {} - if stored_snapshot is not None and stored_snapshot.state is not None: - effective_state.update(stored_snapshot.state) - if isinstance(request_state, dict): - effective_state.update(cast(dict[str, Any], request_state)) - if deferred_default_state: - for key, value in deferred_default_state.items(): - if key not in effective_state: - effective_state[key] = copy.deepcopy(value) + # Seed follow-up turns so the workflow runs with the full persisted thread + # history instead of just the latest request messages. + stored_snapshot = snapshot_session.stored + if stored_snapshot is not None and resume_payload is None: + raw_messages = _reconstruct_messages_from_thread_snapshot( + stored_messages=stored_snapshot.messages, + incoming_messages=raw_messages, + stored_interrupt=stored_snapshot.interrupt, + ) + input_data["messages"] = raw_messages + + effective_state = snapshot_session.effective_state( + request_state=input_data.get("state"), + deferred_defaults=cast(dict[str, Any] | None, input_data.get(_DEFAULT_STATE_INPUT_KEY)), + ) if effective_state: input_data["state"] = effective_state workflow = self._resolve_workflow(thread_id, snapshot_scope) builder_seed_messages = raw_messages - if resume_payload is not None and stored_snapshot is not None: + if resume_payload is not None: # Resume requests carry only the synthesized interrupt response, so seed # the builder with stored history to avoid persisting a truncated thread. - builder_seed_messages = [ - copy.deepcopy(message) for message in stored_snapshot.messages - ] + builder_seed_messages - snapshot_builder = ( - _WorkflowSnapshotBuilder(builder_seed_messages) - if snapshot_store is not None and snapshot_scope is not None - else None - ) + builder_seed_messages = snapshot_session.resume_seeded_messages(builder_seed_messages) + snapshot_builder = _WorkflowSnapshotBuilder(builder_seed_messages) if snapshot_session.enabled else None if snapshot_builder is not None and effective_state: # Seed builder state so a run that emits no StateSnapshotEvent still # persists the latest known Shared State instead of dropping it. @@ -372,37 +313,19 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: snapshot_builder.observe(event) if isinstance(event, RunErrorEvent): run_error_emitted = True - if ( - getattr(event, "code", None) == "WORKFLOW_RESUME_CANCELLED" - and snapshot_store is not None - and snapshot_scope is not None - ): - await _clear_thread_snapshot_interrupt( - snapshot_store=snapshot_store, - scope=snapshot_scope, - thread_id=thread_id, - interrupt_ids=_cancelled_resume_interrupt_ids(resume_payload), + if getattr(event, "code", None) == "WORKFLOW_RESUME_CANCELLED": + await snapshot_session.clear_interrupts( + interrupt_ids=_cancelled_resume_interrupt_ids(resume_payload) ) yield event - if ( - snapshot_builder is not None - and not run_error_emitted - and snapshot_store is not None - and snapshot_scope is not None - ): - try: - await snapshot_store.save( - scope=snapshot_scope, - thread_id=thread_id, - snapshot=snapshot_builder.build(), - ) - except Exception: - # RUN_FINISHED has already been yielded; a store failure must not - # surface as a second terminal RUN_ERROR event. The previous - # snapshot stays available for hydration. - logger.exception( - "Failed to save AG-UI Thread Snapshot for scope=%s thread_id=%s; keeping previous snapshot.", - snapshot_scope, - thread_id, - ) + if snapshot_builder is not None and not run_error_emitted: + # RUN_FINISHED has already been yielded; the session swallows store + # failures so they never surface as a second terminal RUN_ERROR event. + built = snapshot_builder.build() + await snapshot_session.save( + messages=built.messages, + state=built.state, + interrupt=built.interrupt, + session_state=built.session_state, + ) diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py new file mode 100644 index 00000000000..8f8e518672f --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py @@ -0,0 +1,275 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for the ThreadSnapshotSession interface. + +ThreadSnapshotSession is the single owner of the AG-UI Thread Snapshot +lifecycle: load-once at open, hydration replay, effective-state overlay, +resume message seeding, and save-with-swallow semantics. These tests drive +that interface only; runner integration is covered by the existing suite. +""" + +import pytest +from ag_ui.core import EventType + +from agent_framework_ag_ui import AGUIThreadSnapshot, InMemoryAGUIThreadSnapshotStore +from agent_framework_ag_ui._snapshot_session import ThreadSnapshotSession + + +async def make_store_with( + scope: str, + thread_id: str, + snapshot: AGUIThreadSnapshot, +) -> InMemoryAGUIThreadSnapshotStore: + store = InMemoryAGUIThreadSnapshotStore() + await store.save(scope=scope, thread_id=thread_id, snapshot=snapshot) + return store + + +class TestOpenUnscoped: + """A session without a store or scope is inert but safe to call.""" + + async def test_open_without_store_is_disabled(self) -> None: + session = await ThreadSnapshotSession.open(store=None, scope="user-1", thread_id="t1") + assert session.enabled is False + assert session.stored is None + + async def test_open_without_scope_is_disabled(self) -> None: + store = InMemoryAGUIThreadSnapshotStore() + session = await ThreadSnapshotSession.open(store=store, scope=None, thread_id="t1") + assert session.enabled is False + assert session.stored is None + + +class TestOpenScoped: + """A scoped session loads the stored snapshot exactly once at open.""" + + async def test_open_loads_stored_snapshot(self) -> None: + snapshot = AGUIThreadSnapshot( + messages=[{"id": "m1", "role": "user", "content": "hi"}], + state={"counter": 1}, + ) + store = await make_store_with("user-1", "t1", snapshot) + + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + assert session.enabled is True + assert session.stored is not None + assert session.stored.messages == [{"id": "m1", "role": "user", "content": "hi"}] + assert session.stored.state == {"counter": 1} + + async def test_open_with_no_prior_snapshot_is_enabled_but_empty(self) -> None: + store = InMemoryAGUIThreadSnapshotStore() + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + assert session.enabled is True + assert session.stored is None + + +class TestHydrateEvents: + """Hydration replays the stored snapshot as a complete run, no agent invoked.""" + + async def test_full_snapshot_replays_state_messages_and_interrupts(self) -> None: + interrupts = [{"id": "int-1", "type": "approval"}] + snapshot = AGUIThreadSnapshot( + messages=[{"id": "m1", "role": "user", "content": "hi"}], + state={"counter": 1}, + interrupt=interrupts, + ) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + events = [event async for event in session.hydrate_events(run_id="r1")] + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.STATE_SNAPSHOT, + EventType.MESSAGES_SNAPSHOT, + EventType.RUN_FINISHED, + ] + assert events[0].run_id == "r1" + assert events[0].thread_id == "t1" + assert events[1].snapshot == {"counter": 1} + assert [message.id for message in events[2].messages] == ["m1"] + outcome = getattr(events[3], "outcome", None) + assert getattr(outcome, "type", None) == "interrupt" + + async def test_no_stored_snapshot_replays_empty_run(self) -> None: + store = InMemoryAGUIThreadSnapshotStore() + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + events = [event async for event in session.hydrate_events(run_id="r1")] + + assert [event.type for event in events] == [EventType.RUN_STARTED, EventType.RUN_FINISHED] + + async def test_snapshot_without_state_or_interrupts_replays_messages_only(self) -> None: + snapshot = AGUIThreadSnapshot(messages=[{"id": "m1", "role": "user", "content": "hi"}]) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + events = [event async for event in session.hydrate_events(run_id="r1")] + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.MESSAGES_SNAPSHOT, + EventType.RUN_FINISHED, + ] + + +class TestEffectiveState: + """Request values overlay stored values; defaults never reset either.""" + + async def test_request_overlays_stored_and_defaults_fill_missing(self) -> None: + snapshot = AGUIThreadSnapshot(messages=[], state={"a": "stored", "b": "stored"}) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + state = session.effective_state( + request_state={"b": "request", "c": "request"}, + deferred_defaults={"a": "default", "c": "default", "d": "default"}, + ) + + assert state == {"a": "stored", "b": "request", "c": "request", "d": "default"} + + async def test_defaults_are_copied_not_aliased(self) -> None: + store = InMemoryAGUIThreadSnapshotStore() + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + defaults = {"items": ["seed"]} + + state = session.effective_state(request_state=None, deferred_defaults=defaults) + state["items"].append("mutated") + + assert defaults == {"items": ["seed"]} + + async def test_non_dict_request_state_is_ignored(self) -> None: + snapshot = AGUIThreadSnapshot(messages=[], state={"a": 1}) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + assert session.effective_state(request_state="bogus", deferred_defaults=None) == {"a": 1} + + async def test_disabled_session_uses_request_and_defaults_only(self) -> None: + session = await ThreadSnapshotSession.open(store=None, scope=None, thread_id="t1") + + state = session.effective_state( + request_state={"a": "request"}, + deferred_defaults={"a": "default", "b": "default"}, + ) + + assert state == {"a": "request", "b": "default"} + + +class TestResumeSeededMessages: + """Resume requests carry only the interrupt response; stored history is prepended.""" + + async def test_prepends_stored_messages_to_incoming(self) -> None: + snapshot = AGUIThreadSnapshot(messages=[{"id": "m1", "role": "user", "content": "hi"}]) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + incoming = [{"id": "m2", "role": "tool", "content": "approved"}] + + seeded = session.resume_seeded_messages(incoming) + + assert [message["id"] for message in seeded] == ["m1", "m2"] + + async def test_seeded_copies_do_not_alias_stored_snapshot(self) -> None: + snapshot = AGUIThreadSnapshot(messages=[{"id": "m1", "role": "user", "content": "hi"}]) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + seeded = session.resume_seeded_messages([]) + seeded[0]["content"] = "mutated" + + assert session.stored is not None + assert session.stored.messages[0]["content"] == "hi" + + async def test_without_stored_snapshot_returns_incoming_unchanged(self) -> None: + session = await ThreadSnapshotSession.open(store=None, scope=None, thread_id="t1") + incoming = [{"id": "m2", "role": "user", "content": "hello"}] + + assert session.resume_seeded_messages(incoming) == incoming + + +class FailingStore: + """Store whose writes always fail, for exercising save-failure semantics.""" + + async def save(self, *, scope, thread_id, snapshot) -> None: + raise RuntimeError("storage down") + + async def get(self, *, scope, thread_id): + return None + + async def delete(self, *, scope, thread_id) -> bool: + return False + + async def clear(self, *, scope=None) -> None: + return None + + +class TestSave: + """One snapshot write commits messages, state, interrupts, and session state together.""" + + async def test_save_persists_full_snapshot(self) -> None: + store = InMemoryAGUIThreadSnapshotStore() + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + await session.save( + messages=[{"id": "m1", "role": "user", "content": "hi"}], + state={"counter": 2}, + interrupt=[{"id": "int-1"}], + session_state={"provider": {"k": "v"}}, + ) + + saved = await store.get(scope="user-1", thread_id="t1") + assert saved is not None + assert saved.messages == [{"id": "m1", "role": "user", "content": "hi"}] + assert saved.state == {"counter": 2} + assert saved.interrupt == [{"id": "int-1"}] + assert saved.session_state == {"provider": {"k": "v"}} + + async def test_save_on_disabled_session_is_a_noop(self) -> None: + store = InMemoryAGUIThreadSnapshotStore() + session = await ThreadSnapshotSession.open(store=store, scope=None, thread_id="t1") + + await session.save(messages=[{"id": "m1"}], state=None, interrupt=None, session_state=None) + + assert await store.get(scope="unused", thread_id="t1") is None + + async def test_store_failure_is_swallowed_and_logged(self, caplog: pytest.LogCaptureFixture) -> None: + session = await ThreadSnapshotSession.open(store=FailingStore(), scope="user-1", thread_id="t1") + + with caplog.at_level("ERROR"): + await session.save(messages=[], state=None, interrupt=None, session_state=None) + + assert any("keeping previous snapshot" in record.message for record in caplog.records) + + +class TestClearInterrupts: + """Completed interrupts are removed from the latest stored snapshot.""" + + async def test_clears_only_matching_interrupt_ids(self) -> None: + snapshot = AGUIThreadSnapshot( + messages=[{"id": "m1", "role": "user", "content": "hi"}], + interrupt=[{"id": "int-1"}, {"id": "int-2"}], + ) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + await session.clear_interrupts(interrupt_ids={"int-1"}) + + saved = await store.get(scope="user-1", thread_id="t1") + assert saved is not None + assert saved.interrupt == [{"id": "int-2"}] + + async def test_clears_all_interrupts_when_ids_omitted(self) -> None: + snapshot = AGUIThreadSnapshot(messages=[], interrupt=[{"id": "int-1"}]) + store = await make_store_with("user-1", "t1", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="user-1", thread_id="t1") + + await session.clear_interrupts() + + saved = await store.get(scope="user-1", thread_id="t1") + assert saved is not None + assert saved.interrupt is None + + async def test_disabled_session_clear_is_a_noop(self) -> None: + session = await ThreadSnapshotSession.open(store=None, scope=None, thread_id="t1") + await session.clear_interrupts(interrupt_ids={"int-1"}) From 07becc68cb136bc17aec4fc79f2d4836b9c11700 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 3 Aug 2026 11:24:07 +0900 Subject: [PATCH 2/2] Python: Narrow AG-UI event types in snapshot session tests The hydration test accessed run_id, snapshot, and messages on values typed as BaseEvent, which fails the tests/samples type checkers. Narrow each event with isinstance assertions before reading its fields. --- .../tests/ag_ui/test_snapshot_session.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py index 8f8e518672f..413da380d75 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py @@ -9,7 +9,12 @@ """ import pytest -from ag_ui.core import EventType +from ag_ui.core import ( + EventType, + MessagesSnapshotEvent, + RunStartedEvent, + StateSnapshotEvent, +) from agent_framework_ag_ui import AGUIThreadSnapshot, InMemoryAGUIThreadSnapshotStore from agent_framework_ag_ui._snapshot_session import ThreadSnapshotSession @@ -85,10 +90,14 @@ async def test_full_snapshot_replays_state_messages_and_interrupts(self) -> None EventType.MESSAGES_SNAPSHOT, EventType.RUN_FINISHED, ] - assert events[0].run_id == "r1" - assert events[0].thread_id == "t1" - assert events[1].snapshot == {"counter": 1} - assert [message.id for message in events[2].messages] == ["m1"] + run_started, state_snapshot, messages_snapshot = events[0], events[1], events[2] + assert isinstance(run_started, RunStartedEvent) + assert run_started.run_id == "r1" + assert run_started.thread_id == "t1" + assert isinstance(state_snapshot, StateSnapshotEvent) + assert state_snapshot.snapshot == {"counter": 1} + assert isinstance(messages_snapshot, MessagesSnapshotEvent) + assert [message.id for message in messages_snapshot.messages] == ["m1"] outcome = getattr(events[3], "outcome", None) assert getattr(outcome, "type", None) == "interrupt"