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 ab074ae15c..10ce7adcb8 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 da7cbd7d36..706800cedb 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 0000000000..da4360832a --- /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 20f44361c9..abdfc96a38 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 0000000000..413da380d7 --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py @@ -0,0 +1,284 @@ +# 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, + MessagesSnapshotEvent, + RunStartedEvent, + StateSnapshotEvent, +) + +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, + ] + 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" + + 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"})