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 33d3bede30e..3b07d9770b8 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -6,9 +6,11 @@ import logging import uuid +from collections import Counter from collections.abc import AsyncGenerator, Callable from typing import Any, cast +import msgspec from ag_ui.core import ( BaseEvent, MessagesSnapshotEvent, @@ -29,16 +31,18 @@ ToolCallResultEvent, ToolCallStartEvent, ) -from agent_framework import CheckpointID, CheckpointStorage, Workflow, WorkflowCheckpoint +from agent_framework import CheckpointID, CheckpointStorage, Content, Message, Workflow, WorkflowCheckpoint from agent_framework._telemetry import mark_feature_used from ._feature_usage import FeatureIndex -from ._message_adapters import agui_messages_to_snapshot_format +from ._message_adapters import agent_framework_messages_to_agui, agui_messages_to_snapshot_format from ._run_common import ( _cancelled_resume_interrupt_ids, _extract_resume_payload, _normalize_resume_interrupts, _reconstruct_messages_from_thread_snapshot, + _snapshot_messages_match, + _snapshot_overlap_length, ) from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts from ._snapshots import ( @@ -67,6 +71,155 @@ _CHECKPOINT_REQUEST_OWNER_KEY = "ag_ui_workflow_request_owner" +def _hashable_message_content(content: Any) -> Any: + """Return a hashable, order-stable form of snapshot message content.""" + if isinstance(content, (str, int, float, bool)) or content is None: + return content + return msgspec.json.encode(make_json_safe(content), order="deterministic") + + +def _snapshot_messages_from_resume_value( + value: Any, +) -> list[dict[str, Any]]: + """Convert a resolved workflow resume value into snapshot chat messages when user-visible. + + Values have already been coerced and validated by the execution path. Only + ``user`` turns are projected so resume cannot + forge assistant/system/tool history into the backend-owned snapshot. + """ + if isinstance(value, Content): + value = [value] + if isinstance(value, list) and value and all(isinstance(item, Content) for item in value): + # Reuse the user-content encoder for ordered text/media, excluding tool and + # approval controls before they can be reinterpreted as snapshot history. + contents = [item for item in value if item.type in {"text", "uri", "data"}] + if not contents: + return [] + return agui_messages_to_snapshot_format( + agent_framework_messages_to_agui([Message(role="user", contents=contents)]) + ) + value = make_json_safe(value) + if isinstance(value, bool) or value is None: + return [] + if isinstance(value, str): + text = value.strip() + if not text: + return [] + return [{"role": "user", "content": text}] + if isinstance(value, dict): + # Function-approval style payloads are not chat turns. + if any(key in value for key in ("approved", "accepted", "functionCall", "function_call")): + return [] + if value.get("role") == "user": + return agui_messages_to_snapshot_format([_resume_message_to_agui_dict(value)]) + return [] + if isinstance(value, list): + message_like = [ + _resume_message_to_agui_dict(item) + for item in value + if isinstance(item, dict) and item.get("role") == "user" + ] + if message_like: + return agui_messages_to_snapshot_format(message_like) + return [] + + +_RESUME_USER_ALLOWED_KEYS = frozenset({"id", "role", "content", "contents", "name"}) + + +def _resume_message_to_agui_dict(message: dict[str, Any]) -> dict[str, Any]: + """Normalize resume user turns for snapshot encoding. + + Only chat-safe user fields are retained. Control fields such as ``tool_calls`` / + ``toolCalls``, ``actionExecutionId``, or ``function_approvals`` must not enter the + backend-owned thread snapshot — later hydrate/replay can otherwise reinterpret a + user turn as assistant/tool-control history. + """ + normalized = {key: value for key, value in message.items() if key in _RESUME_USER_ALLOWED_KEYS} + if not normalized.get("id") and message.get("message_id"): + normalized["id"] = message["message_id"] + normalized["role"] = "user" + if normalized.get("content") not in (None, ""): + normalized.pop("contents", None) + return normalized + contents = normalized.get("contents") + if isinstance(contents, list): + # Let the snapshot adapter normalize text and media without losing their order. + normalized["content"] = contents + normalized.pop("contents", None) + return normalized + + +def _message_identity(message: dict[str, Any]) -> tuple[Any, ...]: + """Stable identity for deduping resume-synthesized turns against request messages.""" + return (message.get("role"), message.get("id"), _hashable_message_content(message.get("content"))) + + +def _message_content_identity(message: dict[str, Any]) -> tuple[Any, ...]: + """Role+content identity used when message IDs differ across messages vs resume.""" + return (message.get("role"), _hashable_message_content(message.get("content"))) + + +def _current_turn_client_messages( + stored_messages: list[dict[str, Any]], client_messages: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """Exclude confirmed prior transcript overlap before resume content matching.""" + if not stored_messages or not client_messages: + return client_messages + matched_count = max( + ( + overlap + for start, stored_message in enumerate(stored_messages) + if _snapshot_messages_match(stored_message, client_messages[0]) + and ( + (overlap := _snapshot_overlap_length(stored_messages, client_messages, start=start)) >= 2 + or (stored_message.get("id") and stored_message.get("id") == client_messages[0].get("id")) + ) + ), + default=0, + ) + stored_ids = {message.get("id") for message in stored_messages if message.get("id")} + return [message for message in client_messages[matched_count:] if message.get("id") not in stored_ids] + + +def _append_unique_snapshot_messages( + existing: list[dict[str, Any]], + incoming: list[dict[str, Any]], + *, + content_dedupe_against: list[dict[str, Any]] | None = None, +) -> list[dict[str, Any]]: + """Append resume-derived turns that are not already present in the seed. + + Prefer id equality against ``existing``. Role/content fallback is limited to + ``content_dedupe_against`` (current-turn client overlap / id remaps). Callers + must not pass replayed prior transcript rows here — those keep their ids and + would otherwise let an earlier user ``"yes"`` consume a later resume interrupt + with the same text. When the list is empty or omitted, identical replies across + separate HITL turns on ``messages: []`` resumes are kept. + """ + seen_ids = {message.get("id") for message in existing if message.get("id")} + content_source = content_dedupe_against if content_dedupe_against is not None else [] + # Reserve exact-id overlap before content matching, regardless of incoming order. + # A client turn matched by id must not also consume a distinct identical reply. + matched_ids = seen_ids.intersection(message.get("id") for message in incoming if message.get("id")) + remaining_content = Counter( + _message_content_identity(message) for message in content_source if message.get("id") not in matched_ids + ) + merged = list(existing) + for message in incoming: + message_id = message.get("id") + if message_id and message_id in seen_ids: + continue + content_key = _message_content_identity(message) + if remaining_content[content_key] > 0: + remaining_content[content_key] -= 1 + continue + if message_id: + seen_ids.add(message_id) + merged.append(message) + return merged + + def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: """Read an optional checkpoint id to resume from out of the AG-UI forwarded props.""" forwarded_props = input_data.get("forwarded_props") or input_data.get("forwardedProps") @@ -146,6 +299,19 @@ def __init__(self, raw_messages: list[dict[str, Any]]) -> None: self.state: dict[str, Any] | None = None self.interrupt: list[dict[str, Any]] | None = None + def append_resume_messages( + self, messages: list[dict[str, Any]], *, content_dedupe_against: list[dict[str, Any]] + ) -> None: + """Seed validated user replies before any workflow output is observed.""" + self._synthesized_messages = _append_unique_snapshot_messages( + self._synthesized_messages, + agui_messages_to_snapshot_format(messages), + content_dedupe_against=content_dedupe_against, + ) + self._synthesized_message_ids = { + message_id for message in self._synthesized_messages if (message_id := message.get("id")) + } + def observe(self, event: BaseEvent) -> None: """Fold one replayable AG-UI event into the latest snapshot state.""" if isinstance(event, StateSnapshotEvent): @@ -543,6 +709,9 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: run_id = str(input_data.get("run_id") or input_data.get("runId") or uuid.uuid4()) 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 [])) + # Preserve the client-supplied transcript for content-only dedupe of HITL + # resume turns. Stored history is not a confirmed client-replay overlap. + client_request_messages = list(raw_messages) resume_payload = _extract_resume_payload(input_data) snapshot_session = await ThreadSnapshotSession.open( store=self.snapshot_store, @@ -571,7 +740,7 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: "AgentFrameworkWorkflow (or the AG-UI endpoint), or WorkflowBuilder " "checkpoint storage on the workflow instance." ) - live_pending_events = await _pending_request_events(self.workflow) if self.workflow is not None else {} + live_pending_events = await _pending_request_events(workflow) if self.workflow is not None and checkpoint_id is None: for request_event in live_pending_events.values(): owner = getattr(request_event, _REQUEST_OWNER_ATTRIBUTE, None) @@ -605,6 +774,11 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: code="WORKFLOW_RESUME_NOT_FOUND", ) return + effective_pending_events = live_pending_events + # Resolve persisted requests for response validation, but let a pure cold + # restore reach the core so it can re-emit its pending interrupt cards. + if checkpoint_id is not None and resume_payload is not None: + effective_pending_events = dict(checkpoint.pending_request_info_events or {}) if self.workflow is not None and checkpoint_id is None: for interrupt_id in resume_interrupt_ids: request_event = live_pending_events.get(interrupt_id) @@ -657,6 +831,21 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: else: builder_seed_messages = snapshot_session.resume_seeded_messages(builder_seed_messages) snapshot_builder = _WorkflowSnapshotBuilder(builder_seed_messages) if snapshot_session.enabled else None + + def append_conversational_responses(values: list[Any]) -> None: + # Execution owns response validation and approval classification. Project only + # validated conversational values from its effective live/checkpoint requests. + if snapshot_builder is None or resume_payload is None: + return + hitl_messages = [message for value in values for message in _snapshot_messages_from_resume_value(value)] + current_turn_client_messages = _current_turn_client_messages( + stored_snapshot.messages if stored_snapshot is not None else [], client_request_messages + ) + snapshot_builder.append_resume_messages( + hitl_messages, + content_dedupe_against=agui_messages_to_snapshot_format(current_turn_client_messages), + ) + 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. @@ -665,7 +854,12 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: snapshot_builder.state = cast(dict[str, Any], state_snapshot) run_error_emitted = False async for event in run_workflow_stream( - input_data, workflow, checkpoint_storage=run_checkpoint_storage, checkpoint_id=checkpoint_id + input_data, + workflow, + checkpoint_storage=run_checkpoint_storage, + checkpoint_id=checkpoint_id, + pending_request_events=effective_pending_events, + on_conversational_responses=append_conversational_responses, ): if snapshot_builder is not None: snapshot_builder.observe(event) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 1f01c63c9f1..7017931580b 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -874,12 +874,15 @@ def _coerce_responses_for_pending_requests( def _coerce_responses_for_pending_requests_strict( responses: dict[str, Any], pending_events: dict[str, Any], + *, + on_conversational_responses: Callable[[list[Any]], None] | None = None, ) -> tuple[dict[str, Any], RunErrorEvent | None]: """Coerce resume responses or return RUN_ERROR for invalid pending response payloads.""" if not responses or not pending_events: return responses, None normalized: dict[str, Any] = {} + conversational_values: list[Any] = [] pending_by_id: dict[str, Any] = {} for request_id, event in pending_events.items(): pending_by_id[str(request_id)] = event @@ -920,6 +923,14 @@ def _coerce_responses_for_pending_requests_strict( ), ) normalized[request_key] = coerced_value + # Typed execution controls must never become conversational snapshot history. + # Use the response execution actually accepts, including JSON-string Content. + if not isinstance(coerced_value, bool) and not ( + isinstance(coerced_value, Content) and coerced_value.type == "function_approval_response" + ): + conversational_values.append(coerced_value) + if on_conversational_responses is not None: + on_conversational_responses(conversational_values) return normalized, None @@ -1136,6 +1147,8 @@ async def run_workflow_stream( *, checkpoint_storage: CheckpointStorage | None = None, checkpoint_id: str | None = None, + pending_request_events: dict[str, Any] | None = None, + on_conversational_responses: Callable[[list[Any]], None] | None = None, ) -> AsyncGenerator[BaseEvent]: """Run a Workflow and emit AG-UI protocol events. @@ -1153,6 +1166,8 @@ async def run_workflow_stream( mirroring ``Workflow.run(checkpoint_id=...)``. Any incoming messages are treated as request-info responses (or ignored) rather than a new start-executor message, so resume stays consistent with the core API. + pending_request_events: Optional effective pending requests already resolved by the wrapper. + on_conversational_responses: Optional callback for validated non-control responses. """ supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") thread_id = supplied_thread_id or str(uuid.uuid4()) @@ -1179,7 +1194,9 @@ async def run_workflow_stream( # persisted pending set instead. Only do so when a resume payload is present, so a # pure checkpoint restore still surfaces its pending interrupts instead of tripping # the "resume required" contract. - if checkpoint_id is not None and resume_payload is not None: + if pending_request_events is not None: + pending_before_run = pending_request_events + elif checkpoint_id is not None and resume_payload is not None: # Prefer the explicit AG-UI storage argument; otherwise allow the workflow's # builder/runtime storage so builder-emitted pause IDs remain round-trippable. if checkpoint_storage is None and not workflow._runner.context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] @@ -1227,7 +1244,9 @@ async def run_workflow_stream( if request_id in pending_interrupt_ids } responses = _merge_workflow_response_sources(resume_responses, message_responses) - responses, response_error = _coerce_responses_for_pending_requests_strict(responses, pending_before_run) + responses, response_error = _coerce_responses_for_pending_requests_strict( + responses, pending_before_run, on_conversational_responses=on_conversational_responses + ) if response_error is not None: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield response_error diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py index 615d3ce4f1a..6f5df70d923 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py @@ -4,22 +4,48 @@ from __future__ import annotations +from copy import deepcopy +from datetime import datetime, timezone from typing import Any, cast +from unittest.mock import Mock import pytest from agent_framework import ( + Content, Executor, InMemoryCheckpointStorage, + Message, Workflow, WorkflowBuilder, WorkflowContext, executor, handler, + response_handler, ) from agent_framework_ag_ui import AgentFrameworkWorkflow +class _HitlMessageRequestExecutor(Executor): + """Minimal HITL executor shared by snapshot resume regression tests.""" + + def __init__(self) -> None: + super().__init__(id="message_request_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info({"prompt": "Need user follow-up"}, list[Message], request_id="handoff-user-input") + + @response_handler + async def handle_user_input( + self, original_request: dict[str, Any], response: list[Message], ctx: WorkflowContext[Any, str] + ) -> None: + del original_request + user_text = response[0].text if response else "" + await ctx.yield_output(f"Captured response: {user_text}") + + async def _run(agent: AgentFrameworkWorkflow, payload: dict[str, Any]) -> list[Any]: return [event async for event in agent.run(payload)] @@ -453,3 +479,991 @@ async def finalizer(message: str, ctx: WorkflowContext[None, str]) -> None: # The final assistant text is still preserved alongside it. assert any(message.get("content") == "Final answer." for message in snapshot.messages) + + +async def test_workflow_hitl_resume_persists_user_text_in_thread_snapshot() -> None: + """HITL resume with messages:[] must still record the user reply in the snapshot (#8160).""" + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshot + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_HitlMessageRequestExecutor()).build() + store = InMemoryAGUIThreadSnapshotStore() + agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store, checkpoint_storage=storage) + + first_events = await _run( + agent, + { + "thread_id": "thread-hitl", + "run_id": "run-1", + "messages": [{"id": "user-1", "role": "user", "content": "start"}], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + await store.save( + scope="tenant-a", + thread_id="thread-hitl", + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "user-1", "role": "user", "content": "start"}, + {"id": "assistant-1", "role": "assistant", "content": "Need more detail"}, + ], + state=None, + interrupt=None, + ), + ) + + resumed_events = await _run( + agent, + { + "thread_id": "thread-hitl", + "run_id": "run-2", + "messages": [], + "resume": { + "interrupts": [ + { + "id": "handoff-user-input", + "value": [ + { + "role": "user", + "contents": [{"type": "text", "text": "Please ship a replacement instead."}], + } + ], + } + ] + }, + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in resumed_events] + + snapshot = await store.get(scope="tenant-a", thread_id="thread-hitl") + assert snapshot is not None + contents = [message.get("content") for message in snapshot.messages] + assert "start" in contents + assert any(isinstance(content, str) and "replacement" in content for content in contents) + assert any( + message.get("role") == "user" + and isinstance(message.get("content"), str) + and "replacement" in message["content"] + for message in snapshot.messages + ) + + +async def test_workflow_hitl_resume_keeps_repeated_yes_on_empty_messages() -> None: + """A second HITL 'yes' with messages:[] must not be dropped as a content duplicate.""" + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshot + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_HitlMessageRequestExecutor()).build() + store = InMemoryAGUIThreadSnapshotStore() + agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store, checkpoint_storage=storage) + + # Establish a live pending request_info interrupt on this workflow instance. + first_events = await _run( + agent, + { + "thread_id": "thread-hitl-yes", + "run_id": "run-yes-1", + "messages": [{"id": "user-1", "role": "user", "content": "start"}], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + # Simulate a thread that already persisted an earlier identical "yes". + await store.save( + scope="tenant-a", + thread_id="thread-hitl-yes", + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "user-1", "role": "user", "content": "start"}, + {"id": "assistant-1", "role": "assistant", "content": "confirm?"}, + {"id": "user-yes-1", "role": "user", "content": "yes"}, + {"id": "assistant-2", "role": "assistant", "content": "confirm again?"}, + ], + state=None, + interrupt=None, + ), + ) + + resumed_events = await _run( + agent, + { + "thread_id": "thread-hitl-yes", + "run_id": "run-yes-2", + "messages": [], + "resume": { + "interrupts": [ + { + "id": "handoff-user-input", + "value": [ + { + "role": "user", + "contents": [{"type": "text", "text": "yes"}], + } + ], + } + ] + }, + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in resumed_events] + + snapshot = await store.get(scope="tenant-a", thread_id="thread-hitl-yes") + assert snapshot is not None + yes_turns = [ + message for message in snapshot.messages if message.get("role") == "user" and message.get("content") == "yes" + ] + assert len(yes_turns) >= 2 + + +@pytest.mark.parametrize("replay_ids", [True, False], ids=["with-ids", "without-ids"]) +async def test_workflow_hitl_resume_keeps_yes_when_messages_replay_prior_yes(replay_ids: bool) -> None: + """Client-replayed prior 'yes' in messages must not drop a new resume interrupt yes.""" + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshot + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_HitlMessageRequestExecutor()).build() + store = InMemoryAGUIThreadSnapshotStore() + agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store, checkpoint_storage=storage) + + prior = [ + {"id": "user-1", "role": "user", "content": "start"}, + {"id": "assistant-1", "role": "assistant", "content": "confirm?"}, + {"id": "user-yes-1", "role": "user", "content": "yes"}, + {"id": "assistant-2", "role": "assistant", "content": "confirm again?"}, + ] + + first_events = await _run( + agent, + { + "thread_id": "thread-hitl-replay-yes", + "run_id": "run-yes-setup", + "messages": [{"id": "user-1", "role": "user", "content": "start"}], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + await store.save( + scope="tenant-a", + thread_id="thread-hitl-replay-yes", + snapshot=AGUIThreadSnapshot(messages=prior, state=None, interrupt=None), + ) + + resumed_events = await _run( + agent, + { + "thread_id": "thread-hitl-replay-yes", + "run_id": "run-yes-replay", + "messages": [ + {key: value for key, value in message.items() if replay_ids or key != "id"} for message in prior + ], + "resume": { + "interrupts": [ + { + "id": "handoff-user-input", + "value": [ + { + "role": "user", + "contents": [{"type": "text", "text": "yes"}], + } + ], + } + ] + }, + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in resumed_events] + + snapshot = await store.get(scope="tenant-a", thread_id="thread-hitl-replay-yes") + assert snapshot is not None + yes_turns = [ + message for message in snapshot.messages if message.get("role") == "user" and message.get("content") == "yes" + ] + assert len(yes_turns) == 2 + + +@pytest.mark.parametrize("workflow_mode", ["live", "factory", "cold-checkpoint"]) +@pytest.mark.parametrize("reply_kind", ["text", "text-list", "media-list", "audio", "control", "mixed-control"]) +async def test_workflow_hitl_content_reply_survives_snapshot_hydration(workflow_mode: str, reply_kind: str) -> None: + """Accepted Content replies survive hydration as one ordered user turn, without control parts.""" + import msgspec + + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY + + handled: list[list[Content]] = [] + + class ContentRequestExecutor(Executor): + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info("Send a reply", Content, request_id="content-input") + + @response_handler + async def respond(self, original_request: str, response: Content, ctx: WorkflowContext[Any, str]) -> None: + del original_request + handled.append([response]) + await ctx.yield_output("Reply received") + + class ContentListRequestExecutor(Executor): + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info("Send a reply", list[Content], request_id="content-input") + + @response_handler + async def respond(self, original_request: str, response: list[Content], ctx: WorkflowContext[Any, str]) -> None: + del original_request + handled.append(response) + await ctx.yield_output("Reply received") + + text = Content.from_text("approved", raw_representation={"private": "not-chat-history"}) + audio = Content.from_data(b"audio", "audio/wav") + image = Content.from_uri("https://example.com/reply.png", media_type="image/png") + control = Content.from_function_call(call_id="control-call", name="refund", arguments={"amount": 5}) + replies = { + "text": [text], + "text-list": [text, Content.from_text(" again")], + "media-list": [text, image, Content.from_text(" after")], + "audio": [audio], + "control": [control], + "mixed-control": [control, text], + }[reply_kind] + is_list = reply_kind in {"text-list", "media-list", "mixed-control"} + executor_type = ContentListRequestExecutor if is_list else ContentRequestExecutor + + def build_workflow() -> Workflow: + return WorkflowBuilder(start_executor=executor_type(id="content-requester")).build() + + workflow = build_workflow() + store = InMemoryAGUIThreadSnapshotStore() + storage = InMemoryCheckpointStorage() + agent = AgentFrameworkWorkflow( + workflow=workflow if workflow_mode != "factory" else None, + workflow_factory=(lambda _thread_id: build_workflow()) if workflow_mode == "factory" else None, + snapshot_store=store, + checkpoint_storage=storage, + ) + start = {"id": "start", "role": "user", "content": "start"} + setup = await _run(agent, {"thread_id": "content-thread", "messages": [start], _SNAPSHOT_SCOPE_INPUT_KEY: "scope"}) + assert not [event.model_dump() for event in setup if event.type == "RUN_ERROR"] + finished = next(event for event in setup if event.type == "RUN_FINISHED") + assert [interrupt["id"] for interrupt in _interrupts_from_finished(finished)] == ["content-input"] + # Exercise the wire format the public resume path coerces into Content objects. + wire_parts = [part.to_dict() for part in replies] + wire_with_raw = deepcopy(wire_parts) + raw_representation = {"private": "not-chat-history"} + wire_with_raw[0]["raw_representation"] = raw_representation + wire_value = msgspec.json.encode(wire_with_raw if is_list else wire_with_raw[0]).decode() + payload: dict[str, Any] = { + "thread_id": "content-thread", + "messages": [], + _SNAPSHOT_SCOPE_INPUT_KEY: "scope", + "resume": {"interrupts": [{"id": "content-input", "value": wire_value}]}, + } + original_payload = deepcopy(payload) + if workflow_mode == "cold-checkpoint": + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + pending = next( + checkpoint for checkpoint in checkpoints if "content-input" in checkpoint.pending_request_info_events + ) + payload["forwarded_props"] = {"checkpoint_id": pending.checkpoint_id} + original_payload = deepcopy(payload) + agent = AgentFrameworkWorkflow(workflow=build_workflow(), snapshot_store=store, checkpoint_storage=storage) + resumed = await _run(agent, payload) + assert not [event.model_dump() for event in resumed if event.type == "RUN_ERROR"] + assert [part.to_dict() for part in handled[0]] == wire_parts + assert handled[0][0].raw_representation == raw_representation + assert payload == original_payload + + expected: Any = { + "text": "approved", + "text-list": "approved again", + "media-list": [ + {"type": "text", "text": "approved"}, + {"type": "binary", "mimeType": "image/png", "url": "https://example.com/reply.png"}, + {"type": "text", "text": " after"}, + ], + "audio": [{"type": "binary", "mimeType": "audio/wav", "data": "YXVkaW8="}], + "control": None, + "mixed-control": "approved", + }[reply_kind] + snapshot = await store.get(scope="scope", thread_id="content-thread") + assert snapshot is not None + user_replies = [ + message for message in snapshot.messages if message.get("role") == "user" and message["id"] != "start" + ] + assert [message["content"] for message in user_replies] == ([] if expected is None else [expected]) + assert all(set(message) == {"id", "role", "content"} for message in user_replies) + hydrated = await _run(agent, {"thread_id": "content-thread", "messages": [], _SNAPSHOT_SCOPE_INPUT_KEY: "scope"}) + replay = next(event for event in hydrated if event.type == "MESSAGES_SNAPSHOT") + replay_replies = [message for message in replay.messages if message.role == "user" and message.id != "start"] + assert [message.model_dump(by_alias=True, exclude_none=True) for message in replay_replies] == user_replies + assert len(handled) == 1 + + +@pytest.mark.parametrize( + "client_id", + [None, "client-reply", "resume-1", "resume-2"], + ids=["without-id", "remapped-id", "matching-first-id", "matching-last-id"], +) +@pytest.mark.parametrize( + ("parts", "expected_content"), + [ + pytest.param([{"type": "input_text", "text": "yes"}], "yes", id="text-parts"), + pytest.param( + [ + {"type": "input_text", "text": "Before"}, + { + "type": "image", + "source": {"type": "url", "value": "https://example.com/image.png", "mimeType": "image/png"}, + }, + {"type": "text", "text": "After"}, + ], + [ + {"type": "text", "text": "Before"}, + {"type": "binary", "url": "https://example.com/image.png", "mimeType": "image/png"}, + {"type": "text", "text": "After"}, + ], + id="mixed-image", + ), + pytest.param( + [{"type": "audio", "source": {"type": "base64", "value": "YWJj", "mimeType": "audio/wav"}}], + [{"type": "binary", "data": "YWJj", "mimeType": "audio/wav"}], + id="media-only-audio", + ), + ], +) +async def test_workflow_hitl_resume_dedupes_normalized_current_turn_content( + client_id: str | None, parts: list[dict[str, Any]], expected_content: Any +) -> None: + """Only current-turn overlap consumes one matching reply after content normalization.""" + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshot + + handled: list[list[dict[str, Any]]] = [] + + class MessageRequestExecutor(Executor): + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info("Send replies", list[dict[str, Any]], request_id="user-input") + + @response_handler + async def capture( + self, original_request: str, response: list[dict[str, Any]], ctx: WorkflowContext[Any, str] + ) -> None: + del original_request + handled.append(response) + await ctx.yield_output("Captured replies") + + store = InMemoryAGUIThreadSnapshotStore() + workflow = WorkflowBuilder(start_executor=MessageRequestExecutor(id="message-request")).build() + agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store) + start_message = {"id": "start", "role": "user", "content": "start"} + first_events = await _run( + agent, + { + "thread_id": "overlap-thread", + "messages": [start_message], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert not [event.model_dump() for event in first_events if event.type == "RUN_ERROR"] + first_finished = next(event for event in first_events if event.type == "RUN_FINISHED") + assert _interrupts_from_finished(first_finished)[0]["id"] == "user-input" + + # A prior identical reply must not consume either of this interrupt's replies. + await store.save( + scope="tenant-a", + thread_id="overlap-thread", + snapshot=AGUIThreadSnapshot( + messages=[start_message, {"id": "prior-reply", "role": "user", "content": expected_content}], + state=None, + interrupt=None, + ), + ) + client_reply: dict[str, Any] = {"role": "user", "content": deepcopy(parts)} + if client_id is not None: + client_reply["id"] = client_id + resume_replies = [ + {"id": "resume-1", "role": "user", "contents": deepcopy(parts)}, + {"id": "resume-2", "role": "user", "contents": deepcopy(parts)}, + ] + payload = { + "thread_id": "overlap-thread", + "messages": [ + start_message, + {"id": "prior-reply", "role": "user", "content": deepcopy(parts)}, + client_reply, + ], + "resume": {"interrupts": [{"id": "user-input", "value": resume_replies}]}, + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + } + original_payload = deepcopy(payload) + resumed_events = await _run(agent, payload) + assert not [event.model_dump() for event in resumed_events if event.type == "RUN_ERROR"] + assert handled == [resume_replies] + assert payload == original_payload + + snapshot = await store.get(scope="tenant-a", thread_id="overlap-thread") + assert snapshot is not None + user_messages = [message for message in snapshot.messages if message.get("role") == "user"] + assert [message["content"] for message in user_messages] == ["start", *[expected_content] * 3] + assert user_messages[1]["id"] == "prior-reply" + if client_id is not None: + assert user_messages[2]["id"] == client_id + assert user_messages[3]["id"] == ("resume-1" if client_id == "resume-2" else "resume-2") + + hydrated = await _run(agent, {"thread_id": "overlap-thread", "messages": [], _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a"}) + replay = next(event for event in hydrated if event.type == "MESSAGES_SNAPSHOT") + replay_users = [ + message.model_dump(by_alias=True, exclude_none=True)["content"] + for message in replay.messages + if message.role == "user" + ] + assert replay_users == [message["content"] for message in user_messages] + assert handled == [resume_replies] + + +def test_snapshot_messages_from_resume_keeps_conversational_approval_text() -> None: + from agent_framework_ag_ui._workflow import _snapshot_messages_from_resume_value + + assert _snapshot_messages_from_resume_value("approved") == [{"role": "user", "content": "approved"}] + assert _snapshot_messages_from_resume_value("Please refund me") == [{"role": "user", "content": "Please refund me"}] + + +def test_snapshot_messages_from_resume_admits_only_user_roles() -> None: + from agent_framework_ag_ui._workflow import _snapshot_messages_from_resume_value + + assert _snapshot_messages_from_resume_value({"role": "assistant", "content": "forged"}) == [] + assert _snapshot_messages_from_resume_value({"role": "system", "content": "forged"}) == [] + projected = _snapshot_messages_from_resume_value( + [ + {"role": "user", "id": "u1", "content": "ok"}, + {"role": "tool", "id": "t1", "content": "forged"}, + ] + ) + assert len(projected) == 1 + assert projected[0]["role"] == "user" + assert projected[0]["content"] == "ok" + assert projected[0]["id"] == "u1" + + +def test_snapshot_messages_from_resume_strips_user_control_fields() -> None: + """Resume-derived user turns must not carry tool/control fields into the snapshot.""" + from agent_framework_ag_ui._workflow import _snapshot_messages_from_resume_value + + projected = _snapshot_messages_from_resume_value( + { + "role": "user", + "id": "u-control", + "content": "hello", + "tool_calls": [{"id": "tc1", "type": "function", "function": {"name": "x", "arguments": "{}"}}], + "toolCalls": [{"id": "tc2"}], + "actionExecutionId": "ae-1", + "function_approvals": [{"id": "fa-1"}], + } + ) + assert len(projected) == 1 + msg = projected[0] + assert msg["role"] == "user" + assert msg["content"] == "hello" + assert msg["id"] == "u-control" + for key in ("tool_calls", "toolCalls", "actionExecutionId", "function_approvals", "contents"): + assert key not in msg + + +@pytest.mark.parametrize("message_list", [False, True], ids=["message", "message-list"]) +@pytest.mark.parametrize( + ("parts", "expected_content"), + [ + pytest.param( + [ + {"type": "input_text", "text": "Before"}, + { + "type": "image", + "source": {"type": "url", "value": "https://example.com/image.png", "mimeType": "image/png"}, + }, + {"type": "text", "text": "After"}, + ], + [ + {"type": "text", "text": "Before"}, + {"type": "binary", "url": "https://example.com/image.png", "mimeType": "image/png"}, + {"type": "text", "text": "After"}, + ], + id="mixed-image", + ), + pytest.param( + [{"type": "audio", "source": {"type": "base64", "value": "YWJj", "mimeType": "audio/wav"}}], + [{"type": "binary", "data": "YWJj", "mimeType": "audio/wav"}], + id="media-only-audio", + ), + pytest.param( + [{"type": "binary", "url": "https://example.com/video.mp4", "mimeType": "video/mp4"}], + [{"type": "binary", "url": "https://example.com/video.mp4", "mimeType": "video/mp4"}], + id="legacy-binary", + ), + pytest.param( + [{"type": "input_text", "text": "First"}, {"type": "text", "text": "Second"}], + "FirstSecond", + id="text-only", + ), + ], +) +def test_snapshot_messages_from_resume_preserves_contents( + message_list: bool, parts: list[dict[str, Any]], expected_content: Any +) -> None: + """Resume contents retain media and ordering without admitting user control fields.""" + from agent_framework_ag_ui._workflow import _snapshot_messages_from_resume_value + + message = { + "id": "u-media", + "role": "user", + "name": "user", + "contents": parts, + "tool_calls": [{"id": "tc1", "type": "function", "function": {"name": "x", "arguments": "{}"}}], + "toolCalls": [{"id": "tc2"}], + "actionExecutionId": "ae-1", + "function_approvals": [{"id": "fa-1"}], + } + + projected = _snapshot_messages_from_resume_value([message] if message_list else message) + + assert projected == [{"id": "u-media", "role": "user", "name": "user", "content": expected_content}] + + +@pytest.mark.parametrize("content", [None, "", "Preferred content"]) +def test_snapshot_messages_from_resume_contents_respects_content_precedence(content: str | None) -> None: + """Contents supply missing/empty content while explicit content retains precedence.""" + from agent_framework_ag_ui._workflow import _snapshot_messages_from_resume_value + + projected = _snapshot_messages_from_resume_value( + {"id": "u1", "role": "user", "content": content, "contents": [{"type": "text", "text": "Fallback"}]} + ) + + assert projected == [{"id": "u1", "role": "user", "content": content or "Fallback"}] + + +async def test_workflow_hitl_resume_persists_media_only_contents_for_replay() -> None: + """A media-only HITL reply survives persistence and conversion from stored history.""" + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._message_adapters import agui_messages_to_agent_framework + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY + + class MediaRequestExecutor(Executor): + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info("Send media", dict[str, Any], request_id="media-input") + + @response_handler + async def capture( + self, original_request: str, response: dict[str, Any], ctx: WorkflowContext[Any, str] + ) -> None: + del original_request, response + await ctx.yield_output("Captured media") + + store = InMemoryAGUIThreadSnapshotStore() + workflow = WorkflowBuilder(start_executor=MediaRequestExecutor(id="media-request")).build() + agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store) + first_events = await _run( + agent, + { + "thread_id": "media-thread", + "messages": [{"role": "user", "content": "start"}], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + resumed_events = await _run( + agent, + { + "thread_id": "media-thread", + "messages": [], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + "resume": { + "interrupts": [ + { + "id": "media-input", + "value": { + "id": "media-reply", + "role": "user", + "contents": [ + { + "type": "audio", + "source": {"type": "base64", "value": "YWJj", "mimeType": "audio/wav"}, + } + ], + }, + } + ] + }, + }, + ) + assert "RUN_ERROR" not in [event.type for event in resumed_events] + + snapshot = await store.get(scope="tenant-a", thread_id="media-thread") + assert snapshot is not None + reply = next(message for message in snapshot.messages if message.get("id") == "media-reply") + assert reply["content"] == [{"type": "binary", "data": "YWJj", "mimeType": "audio/wav"}] + replayed = agui_messages_to_agent_framework([reply]) + assert replayed[0].contents[0].type == "data" + assert replayed[0].contents[0].uri == "data:audio/wav;base64,YWJj" + assert replayed[0].contents[0].media_type == "audio/wav" + + +def test_message_identity_supports_multimodal_content() -> None: + from agent_framework_ag_ui._workflow import _append_unique_snapshot_messages, _message_identity + + multimodal = { + "id": "m1", + "role": "user", + "content": [{"type": "text", "text": "hi"}, {"type": "image", "url": "x"}], + } + # Must be hashable for set membership during resume dedupe. + assert _message_identity(multimodal) in {_message_identity(multimodal)} + assert _append_unique_snapshot_messages([multimodal], [multimodal]) == [multimodal] + + +def test_message_identity_reuses_model_serialization() -> None: + from pydantic import BaseModel + + from agent_framework_ag_ui._workflow import _message_content_identity + + class TextPart(BaseModel): + type: str = "text" + text: str + + modeled = {"role": "user", "content": [TextPart(text="yes")]} + wire = {"role": "user", "content": [{"text": "yes", "type": "text"}]} + assert _message_content_identity(modeled) == _message_content_identity(wire) + + +@pytest.mark.parametrize("with_client_overlap", [False, True]) +async def test_workflow_hitl_resume_preserves_typed_message_id(with_client_overlap: bool) -> None: + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY + + class TypedMessageRequestExecutor(Executor): + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info("first", list[Message], request_id="first-input") + + @response_handler + async def reply(self, original_request: str, response: list[Message], ctx: WorkflowContext[Any, str]) -> None: + assert response + if original_request == "first": + await ctx.request_info("second", list[Message], request_id="second-input") + else: + await ctx.yield_output("Done") + + store = InMemoryAGUIThreadSnapshotStore() + agent = AgentFrameworkWorkflow( + workflow=WorkflowBuilder(start_executor=TypedMessageRequestExecutor(id="typed-request")).build(), + snapshot_store=store, + ) + start = {"id": "start", "role": "user", "content": "start"} + reply = Message(role="user", contents=["yes"], message_id="typed-reply") + setup = await _run(agent, {"thread_id": "typed-thread", "messages": [start], _SNAPSHOT_SCOPE_INPUT_KEY: "scope"}) + assert not [event for event in setup if event.type == "RUN_ERROR"] + current = {"id": "typed-reply", "role": "user", "content": "yes"} + payload = { + "thread_id": "typed-thread", + "messages": [start, current] if with_client_overlap else [], + "resume": {"interrupts": [{"id": "first-input", "value": [reply.to_dict()]}]}, + _SNAPSHOT_SCOPE_INPUT_KEY: "scope", + } + original_payload = deepcopy(payload) + resumed = await _run(agent, payload) + assert not [event for event in resumed if event.type == "RUN_ERROR"] + assert payload == original_payload + snapshot = await store.get(scope="scope", thread_id="typed-thread") + assert snapshot is not None + user_replies = [message for message in snapshot.messages if message.get("content") == "yes"] + assert user_replies == [current] + prior_transcript = deepcopy(snapshot.messages) + + replay = await _run( + agent, + { + "thread_id": "typed-thread", + "messages": prior_transcript, + "resume": { + "interrupts": [ + {"id": "second-input", "value": [Message("user", ["next"], message_id="next").to_dict()]} + ] + }, + _SNAPSHOT_SCOPE_INPUT_KEY: "scope", + }, + ) + assert not [event for event in replay if event.type == "RUN_ERROR"] + snapshot = await store.get(scope="scope", thread_id="typed-thread") + assert snapshot is not None + assert [message for message in snapshot.messages if message.get("content") == "yes"] == [current] + + +def test_append_unique_snapshot_messages_dedupes_client_replay() -> None: + from agent_framework_ag_ui._workflow import _append_unique_snapshot_messages + + existing = [{"id": "u1", "role": "user", "content": "already present"}] + incoming = [ + {"id": "u1", "role": "user", "content": "already present"}, + {"id": "u2", "role": "user", "content": "new reply"}, + ] + assert _append_unique_snapshot_messages(existing, incoming) == [ + {"id": "u1", "role": "user", "content": "already present"}, + {"id": "u2", "role": "user", "content": "new reply"}, + ] + + +def test_append_unique_snapshot_messages_dedupes_different_ids_same_content() -> None: + from agent_framework_ag_ui._workflow import _append_unique_snapshot_messages + + existing = [{"id": "client-id", "role": "user", "content": "same turn"}] + incoming = [{"id": "generated-id", "role": "user", "content": "same turn"}] + # Content fallback only applies against confirmed client-replay overlap. + assert ( + _append_unique_snapshot_messages( + existing, + incoming, + content_dedupe_against=existing, + ) + == existing + ) + + +def test_append_unique_snapshot_messages_keeps_intentional_repeated_replies() -> None: + from agent_framework_ag_ui._workflow import _append_unique_snapshot_messages + + existing = [{"id": "u0", "role": "user", "content": "hello"}] + incoming = [ + {"id": "r1", "role": "user", "content": "repeat"}, + {"id": "r2", "role": "user", "content": "repeat"}, + ] + merged = _append_unique_snapshot_messages(existing, incoming) + assert [m["id"] for m in merged] == ["u0", "r1", "r2"] + + +def test_append_unique_snapshot_messages_keeps_current_overlap_after_historical_id_match() -> None: + """Matching historical ids must not consume a current client turn's content allowance.""" + from agent_framework_ag_ui._workflow import _append_unique_snapshot_messages + + history = [{"id": "prior-reply", "role": "user", "content": "yes"}] + current = [{"id": "client-reply", "role": "user", "content": "yes"}] + incoming = [ + {"id": "prior-reply", "role": "user", "content": "yes"}, + {"id": "resume-reply", "role": "user", "content": "yes"}, + ] + + assert _append_unique_snapshot_messages(history + current, incoming, content_dedupe_against=current) == ( + history + current + ) + + +def test_append_unique_snapshot_messages_keeps_second_hitl_yes_without_client_replay() -> None: + """messages:[] HITL resumes must not collapse a later identical reply against history.""" + from agent_framework_ag_ui._workflow import _append_unique_snapshot_messages + + history = [ + {"id": "u0", "role": "user", "content": "start"}, + {"id": "a0", "role": "assistant", "content": "confirm?"}, + {"id": "u1", "role": "user", "content": "yes"}, + {"id": "a1", "role": "assistant", "content": "confirm again?"}, + ] + second_yes = [{"id": "generated-yes-2", "role": "user", "content": "yes"}] + merged = _append_unique_snapshot_messages( + history, + second_yes, + content_dedupe_against=[], + ) + assert [m["id"] for m in merged] == ["u0", "a0", "u1", "a1", "generated-yes-2"] + + +def test_append_unique_snapshot_messages_keeps_resume_yes_when_client_replays_prior_yes() -> None: + """Replayed prior-yes rows (same ids as stored) must not be passed as content fallback.""" + from agent_framework_ag_ui._workflow import _append_unique_snapshot_messages + + history = [ + {"id": "u0", "role": "user", "content": "start"}, + {"id": "a0", "role": "assistant", "content": "confirm?"}, + {"id": "u1", "role": "user", "content": "yes"}, + {"id": "a1", "role": "assistant", "content": "confirm again?"}, + ] + stored_ids = {message.get("id") for message in history if message.get("id")} + client_replay = list(history) + # Mirrors the run() call site: only count client turns not already stored by id. + current_turn_client_messages = [ + message for message in client_replay if not (message.get("id") and message.get("id") in stored_ids) + ] + second_yes = [{"id": "generated-yes-2", "role": "user", "content": "yes"}] + merged = _append_unique_snapshot_messages( + history, + second_yes, + content_dedupe_against=current_turn_client_messages, + ) + assert current_turn_client_messages == [] + assert [m["id"] for m in merged] == ["u0", "a0", "u1", "a1", "generated-yes-2"] + + +@pytest.mark.parametrize( + "workflow_mode", ["live", "factory", "cold-checkpoint", "cold-restore", "cold-factory-restore"] +) +@pytest.mark.parametrize("response_kind", ["content-approval", "bool-approval", "text"]) +async def test_workflow_snapshot_uses_validated_resume_responses( + workflow_mode: str, response_kind: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Snapshots preserve conversational replies and pending cards across resume and cold restore.""" + import json + + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY + + # A coarse clock can give the initial and interrupted checkpoints the same timestamp. + checkpoint_clock = Mock(wraps=datetime) + checkpoint_clock.now.return_value = datetime(2026, 1, 1, tzinfo=timezone.utc) + monkeypatch.setattr("agent_framework._workflows._checkpoint.datetime", checkpoint_clock) + + handled: list[Any] = [] + + class ContentApprovalExecutor(Executor): + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + call = Content.from_function_call(call_id="call-1", name="refund", arguments={"amount": 5}) + request = Content.from_function_approval_request(id="reply-1", function_call=call) + await ctx.request_info(request, Content, request_id="reply-1") + + @response_handler + async def respond(self, original_request: Content, response: Content, ctx: WorkflowContext[Any, str]) -> None: + del original_request + handled.append(response) + assert response.type == "function_approval_response" + assert response.approved is True + await ctx.yield_output("Decision processed") + + class BoolApprovalExecutor(Executor): + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info("Approve?", bool, request_id="reply-1") + + @response_handler + async def respond(self, original_request: str, response: bool, ctx: WorkflowContext[Any, str]) -> None: + del original_request + handled.append(response) + assert response is True + await ctx.yield_output("Decision processed") + + class TextExecutor(Executor): + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info("Say something", str, request_id="reply-1") + + @response_handler + async def respond(self, original_request: str, response: str, ctx: WorkflowContext[Any, str]) -> None: + del original_request + handled.append(response) + assert response == "approved" + await ctx.yield_output("Reply processed") + + executor_type: type[Executor] = { + "content-approval": ContentApprovalExecutor, + "bool-approval": BoolApprovalExecutor, + "text": TextExecutor, + }[response_kind] + + def build_workflow() -> Workflow: + return WorkflowBuilder(start_executor=executor_type(id="requester")).build() + + storage = InMemoryCheckpointStorage() + store = InMemoryAGUIThreadSnapshotStore() + workflow = build_workflow() + agent = AgentFrameworkWorkflow( + workflow=workflow if workflow_mode != "factory" else None, + workflow_factory=(lambda _thread_id: build_workflow()) if workflow_mode == "factory" else None, + snapshot_store=store, + checkpoint_storage=storage, + ) + first_events = await _run( + agent, + { + "thread_id": "approval-thread", + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + "messages": [{"id": "start", "role": "user", "content": "start"}], + }, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + first_finished = next(event for event in first_events if event.type == "RUN_FINISHED") + interrupt = _interrupts_from_finished(first_finished)[0] + request = interrupt["metadata"]["agent_framework"]["value"] + if isinstance(request, str): + request = json.loads(request) + if response_kind == "content-approval": + # The execution path accepts a JSON-string Content response; snapshotting + # must consume its validated Content instead of persisting that raw string. + value = json.dumps( + { + "type": "function_approval_response", + "id": "reply-1", + "approved": True, + "function_call": request["function_call"], + } + ) + elif response_kind == "bool-approval": + value = "true" + else: + value = "approved" + payload: dict[str, Any] = { + "thread_id": "approval-thread", + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + "messages": [], + "resume": {"interrupts": [{"id": "reply-1", "value": value}]}, + } + restore_only = workflow_mode in {"cold-restore", "cold-factory-restore"} + if workflow_mode == "cold-checkpoint" or restore_only: + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + assert len(checkpoints) > 1 + assert len({checkpoint.timestamp for checkpoint in checkpoints}) == 1 + pending_checkpoint = next( + checkpoint for checkpoint in checkpoints if "reply-1" in checkpoint.pending_request_info_events + ) + payload["forwarded_props"] = {"checkpoint_id": pending_checkpoint.checkpoint_id} + agent = AgentFrameworkWorkflow( + workflow=build_workflow() if workflow_mode != "cold-factory-restore" else None, + workflow_factory=(lambda _thread_id: build_workflow()) if workflow_mode == "cold-factory-restore" else None, + snapshot_store=store, + checkpoint_storage=storage, + ) + if restore_only: + del payload["resume"] + resumed = await _run(agent, payload) + assert not [event.model_dump() for event in resumed if event.type == "RUN_ERROR"] + if restore_only: + restored_finished = next(event for event in resumed if event.type == "RUN_FINISHED") + assert [pending["id"] for pending in _interrupts_from_finished(restored_finished)] == ["reply-1"] + assert len(handled) == (0 if restore_only else 1) + snapshot = await store.get(scope="tenant-a", thread_id="approval-thread") + assert snapshot is not None + user_contents = [message["content"] for message in snapshot.messages if message.get("role") == "user"] + assert user_contents == (["start", "approved"] if response_kind == "text" and not restore_only else ["start"]) + # Hydration must replay the same safe history without executing a pending response. + hydrated = await _run( + agent, {"thread_id": "approval-thread", _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", "messages": []} + ) + replay = next(event for event in hydrated if event.type == "MESSAGES_SNAPSHOT") + replay_users = [message.content for message in replay.messages if message.role == "user"] + assert replay_users == user_contents + assert len(handled) == (0 if restore_only else 1)