From ec0a6861951df0c1e751e4d120e5f9f65227108a Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Fri, 4 Sep 2026 16:16:49 +0530 Subject: [PATCH 1/3] Python: fix: keep AG-UI workflow reasoning in thread snapshots `_emit_text_reasoning` persists each workflow reasoning message into `flow.reasoning_messages`, and `_WorkflowSnapshotBuilder.observe` folded `TextMessage*` and `ToolCall*` events into the synthesized snapshot but had no `Reasoning*` branch. Reasoning events fell through silently, so intermediate output rendered live and then vanished when the thread was hydrated from a snapshot -- while the same run through the agent path kept it. Fold reasoning into the builder: accumulate deltas per message_id and emit entries in the shape the agent path already produces ({"id", "role": "reasoning", "content", ["encryptedValue"]}). Reasoning is flushed at build() and at text-message and tool-call starts so a block that streamed before other output replays in the position it streamed in. Reasoning deliberately does not close an open tool-call group: it is UI-only state that `agui_messages_to_agent_framework` drops, so it cannot break the tool_calls/result adjacency providers require. Fixes #8054 Co-Authored-By: Claude Opus 5 --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 57 ++++++ .../ag-ui/tests/ag_ui/test_snapshots.py | 163 ++++++++++++++++++ .../ag-ui/tests/ag_ui/test_workflow_agent.py | 61 +++++++ 3 files changed, 281 insertions(+) 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 73b2a1d7ec0..8644b2ed50a 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -12,6 +12,11 @@ from ag_ui.core import ( BaseEvent, MessagesSnapshotEvent, + ReasoningEncryptedValueEvent, + ReasoningEndEvent, + ReasoningMessageContentEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, RunErrorEvent, RunFinishedEvent, RunStartedEvent, @@ -122,6 +127,7 @@ def __init__(self, raw_messages: list[dict[str, Any]]) -> None: self._synthesized_messages = agui_messages_to_snapshot_format(raw_messages) self._emitted_messages: list[dict[str, Any]] | None = None self._open_text_message: dict[str, Any] | None = None + self._open_reasoning_message: dict[str, Any] | None = None self._tool_call_message: dict[str, Any] | None = None self._tool_calls_by_id: dict[str, dict[str, Any]] = {} self.state: dict[str, Any] | None = None @@ -165,14 +171,24 @@ def observe(self, event: BaseEvent) -> None: self._observe_tool_call_args(event) elif isinstance(event, ToolCallResultEvent): self._observe_tool_call_result(event) + elif isinstance(event, ReasoningMessageStartEvent): + self._observe_reasoning_start(event) + elif isinstance(event, ReasoningMessageContentEvent): + self._observe_reasoning_content(event) + elif isinstance(event, (ReasoningMessageEndEvent, ReasoningEndEvent)): + self._observe_reasoning_end(event) + elif isinstance(event, ReasoningEncryptedValueEvent): + self._observe_reasoning_encrypted_value(event) def build(self) -> AGUIThreadSnapshot: """Return the replayable thread snapshot.""" + self._flush_open_reasoning_message() self._flush_open_text_message() messages = self._emitted_messages if self._emitted_messages is not None else self._synthesized_messages return AGUIThreadSnapshot(messages=messages, state=self.state, interrupt=self.interrupt) def _observe_text_start(self, event: TextMessageStartEvent) -> None: + self._flush_open_reasoning_message() if self._open_text_message is not None and self._open_text_message.get("id") != event.message_id: self._flush_open_text_message() self._open_text_message = {"id": event.message_id, "role": event.role, "content": ""} @@ -188,6 +204,7 @@ def _observe_text_end(self, event: TextMessageEndEvent) -> None: self._flush_open_text_message() def _observe_tool_call_start(self, event: ToolCallStartEvent) -> None: + self._flush_open_reasoning_message() parent_message_id = event.parent_message_id if ( self._open_text_message is not None @@ -245,6 +262,46 @@ def _flush_open_text_message(self) -> None: self._tool_call_message = None self._open_text_message = None + def _observe_reasoning_start(self, event: ReasoningMessageStartEvent) -> None: + if self._open_reasoning_message is not None and self._open_reasoning_message.get("id") != event.message_id: + self._flush_open_reasoning_message() + self._open_reasoning_message = {"id": event.message_id, "role": "reasoning", "content": ""} + + def _observe_reasoning_content(self, event: ReasoningMessageContentEvent) -> None: + if self._open_reasoning_message is None or self._open_reasoning_message.get("id") != event.message_id: + self._flush_open_reasoning_message() + self._open_reasoning_message = {"id": event.message_id, "role": "reasoning", "content": ""} + self._open_reasoning_message["content"] = f"{self._open_reasoning_message.get('content', '')}{event.delta}" + + def _observe_reasoning_end(self, event: ReasoningMessageEndEvent | ReasoningEndEvent) -> None: + if self._open_reasoning_message is None or self._open_reasoning_message.get("id") != event.message_id: + return + self._flush_open_reasoning_message() + + def _observe_reasoning_encrypted_value(self, event: ReasoningEncryptedValueEvent) -> None: + # Only message-scoped encrypted values belong on a reasoning message; the + # protocol also uses this event for other subtypes. + if event.subtype != "message": + return + if self._open_reasoning_message is not None and self._open_reasoning_message.get("id") == event.entity_id: + self._open_reasoning_message["encryptedValue"] = event.encrypted_value + return + # The block may already be closed (encrypted value trailing its own + # REASONING_MESSAGE_END); attach it to the message we already flushed. + for message in reversed(self._synthesized_messages): + if message.get("role") == "reasoning" and message.get("id") == event.entity_id: + message["encryptedValue"] = event.encrypted_value + return + + def _flush_open_reasoning_message(self) -> None: + if self._open_reasoning_message is None: + return + # An encrypted-value-only block carries no display text but still has to + # survive hydration, so it counts as content worth keeping. + if self._open_reasoning_message.get("content") or self._open_reasoning_message.get("encryptedValue"): + self._synthesized_messages.append(self._open_reasoning_message) + self._open_reasoning_message = None + class AgentFrameworkWorkflow: """Base AG-UI workflow wrapper. diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py index e2ea85e37df..7b1bc8f0fed 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py @@ -185,6 +185,169 @@ def test_workflow_snapshot_builder_splits_tool_call_groups() -> None: ] +def test_workflow_snapshot_builder_folds_reasoning_into_snapshot() -> None: + """Streamed reasoning deltas accumulate into a replayable reasoning message.""" + from ag_ui.core import ( + ReasoningMessageContentEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="step one ")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="step two")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + + assert builder.build().messages == [{"id": "reason-1", "role": "reasoning", "content": "step one step two"}] + + +def test_workflow_snapshot_builder_keeps_reasoning_in_emission_order() -> None: + """Reasoning is replayed where it streamed, not appended after the visible output.""" + from ag_ui.core import ( + ReasoningMessageContentEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + TextMessageContentEvent, + TextMessageEndEvent, + TextMessageStartEvent, + ToolCallResultEvent, + ToolCallStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="planning")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + builder.observe(ToolCallStartEvent(tool_call_id="call-a", tool_call_name="toolA")) + builder.observe(ToolCallResultEvent(message_id="result-a", tool_call_id="call-a", content="resA")) + builder.observe(TextMessageStartEvent(message_id="text-1", role="assistant")) + builder.observe(TextMessageContentEvent(message_id="text-1", delta="done")) + builder.observe(TextMessageEndEvent(message_id="text-1")) + + messages = builder.build().messages + assert [ + ( + message.get("role"), + [tool_call["id"] for tool_call in message.get("tool_calls", [])] + or message.get("toolCallId") + or message.get("content"), + ) + for message in messages + ] == [ + ("reasoning", "planning"), + ("assistant", ["call-a"]), + ("tool", "call-a"), + ("assistant", "done"), + ] + + +def test_workflow_snapshot_builder_captures_reasoning_encrypted_value() -> None: + """Encrypted reasoning payloads survive hydration under the protocol's camelCase key.""" + from ag_ui.core import ( + ReasoningEncryptedValueEvent, + ReasoningMessageContentEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="hidden")) + builder.observe(ReasoningEncryptedValueEvent(subtype="message", entity_id="reason-1", encrypted_value="cipher")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + + assert builder.build().messages == [ + {"id": "reason-1", "role": "reasoning", "content": "hidden", "encryptedValue": "cipher"} + ] + + +def test_workflow_snapshot_builder_flushes_reasoning_left_open_at_build() -> None: + """A run that ends without REASONING_MESSAGE_END still snapshots what streamed.""" + from ag_ui.core import ReasoningMessageContentEvent, ReasoningMessageStartEvent + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="unterminated")) + + assert builder.build().messages == [{"id": "reason-1", "role": "reasoning", "content": "unterminated"}] + + +def test_workflow_snapshot_builder_separates_consecutive_reasoning_blocks() -> None: + """A new reasoning message id closes the previous block instead of merging into it.""" + from ag_ui.core import ReasoningMessageContentEvent, ReasoningMessageStartEvent + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="first")) + # Second block opens without the first ever being closed. + builder.observe(ReasoningMessageStartEvent(message_id="reason-2", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-2", delta="second")) + + assert builder.build().messages == [ + {"id": "reason-1", "role": "reasoning", "content": "first"}, + {"id": "reason-2", "role": "reasoning", "content": "second"}, + ] + + +def test_workflow_snapshot_builder_folds_reasoning_content_without_a_start_event() -> None: + """Reasoning deltas are kept even if the opening REASONING_MESSAGE_START was missed.""" + from ag_ui.core import ReasoningMessageContentEvent + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="orphaned")) + + assert builder.build().messages == [{"id": "reason-1", "role": "reasoning", "content": "orphaned"}] + + +def test_workflow_snapshot_builder_attaches_encrypted_value_after_block_closed() -> None: + """An encrypted value trailing REASONING_MESSAGE_END still lands on its message.""" + from ag_ui.core import ( + ReasoningEncryptedValueEvent, + ReasoningMessageContentEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="hidden")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + builder.observe( + ReasoningEncryptedValueEvent(subtype="message", entity_id="reason-1", encrypted_value="late-cipher") + ) + + assert builder.build().messages == [ + {"id": "reason-1", "role": "reasoning", "content": "hidden", "encryptedValue": "late-cipher"} + ] + + +def test_workflow_snapshot_builder_ignores_tool_call_encrypted_values() -> None: + """A tool-call-scoped encrypted value must not be folded in as reasoning content.""" + from ag_ui.core import ReasoningEncryptedValueEvent + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningEncryptedValueEvent(subtype="tool-call", entity_id="call-a", encrypted_value="cipher")) + + assert builder.build().messages == [] + + async def test_in_memory_snapshot_store_rejects_invalid_keys() -> None: """Key parts must be non-empty strings for every store operation.""" import pytest 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 4a4ec411605..8702adcd047 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 @@ -359,3 +359,64 @@ async def test_workflow_checkpoint_only_resume_preserves_thread_snapshot() -> No assert "Earlier reply" in contents # ...plus the newly produced output from the resumed run. assert any(isinstance(content, str) and "done" in content for content in contents) + + +async def test_workflow_snapshot_preserves_streamed_reasoning() -> None: + """Reasoning that streamed during a workflow run survives thread hydration. + + Regression test for workflow reasoning rendering live and then vanishing from the + stored AG-UI Thread Snapshot, so a hydrated thread lost the intermediate output + that the agent path retains. + """ + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY + + @executor(id="thinker") + async def thinker(message: Any, ctx: WorkflowContext[str, str]) -> None: + await ctx.yield_output("Weighing the options...") + await ctx.send_message("go") + + @executor(id="finalizer") + async def finalizer(message: str, ctx: WorkflowContext[None, str]) -> None: + await ctx.yield_output("Final answer.") + + workflow = ( + WorkflowBuilder( + start_executor=thinker, + output_from=[finalizer], + intermediate_output_from=[thinker], + ) + .add_edge(thinker, finalizer) + .build() + ) + + store = InMemoryAGUIThreadSnapshotStore() + agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store) + + events = await _run( + agent, + { + "thread_id": "thread-reasoning", + "run_id": "run-1", + "messages": [{"id": "user-1", "role": "user", "content": "Decide"}], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in events] + + # The reasoning really did stream ... + reasoning_deltas = [ + event.delta # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + for event in events + if event.type == "REASONING_MESSAGE_CONTENT" + ] + assert "Weighing the options..." in reasoning_deltas + + # ... and it is still there once the thread is hydrated from the snapshot. + snapshot = await store.get(scope="tenant-a", thread_id="thread-reasoning") + assert snapshot is not None + reasoning_messages = [message for message in snapshot.messages if message.get("role") == "reasoning"] + assert [message.get("content") for message in reasoning_messages] == ["Weighing the options..."] + + # The final assistant text is still preserved alongside it. + assert any(message.get("content") == "Final answer." for message in snapshot.messages) From 55be3e9e6cc8d2da517b1c2cf455276513bf2183 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Fri, 4 Sep 2026 16:43:59 +0530 Subject: [PATCH 2/3] Python: fix: keep encrypted-only AG-UI workflow reasoning on message end Address review: an encrypted-value-only reasoning message was dropped for the event order `_emit_text_reasoning` produces without a flow, where REASONING_MESSAGE_END precedes REASONING_ENCRYPTED_VALUE. The message carries no display text, so closing it at REASONING_MESSAGE_END discarded it and left the encrypted value with nothing to attach to. The underlying mistake was conflating two protocol levels. REASONING_MESSAGE_END closes the message; REASONING_END closes the block; and an encrypted value is block-scoped, so it legitimately trails the message end. Keep the message open past REASONING_MESSAGE_END and let REASONING_END, a new REASONING_START, intervening text/tool output, or build() finalize it. Handle REASONING_START so a new block also closes anything the previous one left open. Adds regression coverage for the reported order, for an empty block with neither text nor an encrypted value still being dropped, for a block closed by the next REASONING_START, for an encrypted value arriving after intervening text, and for end events naming a message that was never opened. Co-Authored-By: Claude Opus 5 --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 15 ++- .../ag-ui/tests/ag_ui/test_snapshots.py | 121 ++++++++++++++++++ 2 files changed, 134 insertions(+), 2 deletions(-) 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 8644b2ed50a..5b35f8d00ce 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -17,6 +17,7 @@ ReasoningMessageContentEvent, ReasoningMessageEndEvent, ReasoningMessageStartEvent, + ReasoningStartEvent, RunErrorEvent, RunFinishedEvent, RunStartedEvent, @@ -171,6 +172,9 @@ def observe(self, event: BaseEvent) -> None: self._observe_tool_call_args(event) elif isinstance(event, ToolCallResultEvent): self._observe_tool_call_result(event) + elif isinstance(event, ReasoningStartEvent): + # A new reasoning block supersedes anything still open from the last one. + self._flush_open_reasoning_message() elif isinstance(event, ReasoningMessageStartEvent): self._observe_reasoning_start(event) elif isinstance(event, ReasoningMessageContentEvent): @@ -276,6 +280,13 @@ def _observe_reasoning_content(self, event: ReasoningMessageContentEvent) -> Non def _observe_reasoning_end(self, event: ReasoningMessageEndEvent | ReasoningEndEvent) -> None: if self._open_reasoning_message is None or self._open_reasoning_message.get("id") != event.message_id: return + if isinstance(event, ReasoningMessageEndEvent): + # REASONING_MESSAGE_END closes the message, not the block, and an + # encrypted value is block-scoped so it legitimately trails it -- that is + # the order `_emit_text_reasoning` produces without a flow. Keep the + # message open so protected-data-only reasoning still has somewhere to + # land; REASONING_END, the next block, or build() finalizes it. + return self._flush_open_reasoning_message() def _observe_reasoning_encrypted_value(self, event: ReasoningEncryptedValueEvent) -> None: @@ -286,8 +297,8 @@ def _observe_reasoning_encrypted_value(self, event: ReasoningEncryptedValueEvent if self._open_reasoning_message is not None and self._open_reasoning_message.get("id") == event.entity_id: self._open_reasoning_message["encryptedValue"] = event.encrypted_value return - # The block may already be closed (encrypted value trailing its own - # REASONING_MESSAGE_END); attach it to the message we already flushed. + # Intervening text or tool output can flush the message before its encrypted + # value arrives; attach it to the message we already synthesized. for message in reversed(self._synthesized_messages): if message.get("role") == "reasoning" and message.get("id") == event.entity_id: message["encryptedValue"] = event.encrypted_value diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py index 7b1bc8f0fed..5b4e8e694f5 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py @@ -336,6 +336,127 @@ def test_workflow_snapshot_builder_attaches_encrypted_value_after_block_closed() ] +def test_workflow_snapshot_builder_keeps_encrypted_only_reasoning_ended_before_its_value() -> None: + """Protected-data-only reasoning survives the no-flow order, where END precedes the value. + + `_emit_text_reasoning` without a flow emits REASONING_MESSAGE_END before + REASONING_ENCRYPTED_VALUE. The message carries no display text, so closing it at + REASONING_MESSAGE_END would discard it and leave the encrypted value nothing to + attach to. + """ + from ag_ui.core import ( + ReasoningEncryptedValueEvent, + ReasoningEndEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + builder.observe(ReasoningEncryptedValueEvent(subtype="message", entity_id="reason-1", encrypted_value="cipher")) + builder.observe(ReasoningEndEvent(message_id="reason-1")) + + assert builder.build().messages == [ + {"id": "reason-1", "role": "reasoning", "content": "", "encryptedValue": "cipher"} + ] + + +def test_workflow_snapshot_builder_drops_reasoning_with_neither_text_nor_encrypted_value() -> None: + """An empty reasoning block with nothing to replay is not synthesized into a message.""" + from ag_ui.core import ( + ReasoningEndEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + builder.observe(ReasoningEndEvent(message_id="reason-1")) + + assert builder.build().messages == [] + + +def test_workflow_snapshot_builder_closes_previous_block_on_new_reasoning_start() -> None: + """A new REASONING_START finalizes a message the previous block left open.""" + from ag_ui.core import ( + ReasoningMessageContentEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="first")) + # No REASONING_END: the next block's start has to close this one. + builder.observe(ReasoningStartEvent(message_id="reason-2")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-2", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-2", delta="second")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-2")) + + assert builder.build().messages == [ + {"id": "reason-1", "role": "reasoning", "content": "first"}, + {"id": "reason-2", "role": "reasoning", "content": "second"}, + ] + + +def test_workflow_snapshot_builder_attaches_encrypted_value_after_intervening_output() -> None: + """Text arriving between a reasoning message and its encrypted value does not lose the value.""" + from ag_ui.core import ( + ReasoningEncryptedValueEvent, + ReasoningMessageContentEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + TextMessageContentEvent, + TextMessageStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="thinking")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + # Text output flushes the reasoning message before the encrypted value shows up. + builder.observe(TextMessageStartEvent(message_id="text-1", role="assistant")) + builder.observe(TextMessageContentEvent(message_id="text-1", delta="answer")) + builder.observe(ReasoningEncryptedValueEvent(subtype="message", entity_id="reason-1", encrypted_value="cipher")) + + assert builder.build().messages == [ + {"id": "reason-1", "role": "reasoning", "content": "thinking", "encryptedValue": "cipher"}, + {"id": "text-1", "role": "assistant", "content": "answer"}, + ] + + +def test_workflow_snapshot_builder_ignores_reasoning_end_for_unopened_message() -> None: + """A reasoning end event for a message that was never opened is a no-op.""" + from ag_ui.core import ReasoningEndEvent, ReasoningMessageContentEvent, ReasoningMessageStartEvent + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningEndEvent(message_id="never-opened")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="kept")) + # An end event for a different id must not close the open message. + builder.observe(ReasoningEndEvent(message_id="other")) + + assert builder.build().messages == [{"id": "reason-1", "role": "reasoning", "content": "kept"}] + + def test_workflow_snapshot_builder_ignores_tool_call_encrypted_values() -> None: """A tool-call-scoped encrypted value must not be folded in as reasoning content.""" from ag_ui.core import ReasoningEncryptedValueEvent From b7c83fb85639cd27f4d565b955c3df812a25612e Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Tue, 8 Sep 2026 19:37:26 +0530 Subject: [PATCH 3/3] Python: replay AG-UI workflow reasoning in the order it streamed Review on #8058 found two ordering defects in `_WorkflowSnapshotBuilder`. An `output` event followed by a later `intermediate` event left both the open text message and the open reasoning message unflushed, so `build()`'s fixed flush order decided the sequence and hydration reversed what streamed. Nothing flushed open text when reasoning opened, even though `_observe_text_start` and `_observe_tool_call_start` already flush open reasoning. Reasoning is now flushed wherever other output is appended and text is flushed wherever reasoning opens, so the two slots can never both be open and `build()`'s order cannot matter. That covers the `_observe_text_content` resume path and `_observe_tool_call_result`, which had the same gap. Reasoning is UI-only and dropped before provider conversion, so a block that now lands between a tool call and its result still converts back with the two adjacent, which provider APIs require; a test pins that across the two modules. A block carrying only protected data has no content event at all, so output arriving before its encrypted value flushed an empty message that the late `ReasoningEncryptedValueEvent` could no longer find, and the protected value was lost. The empty message now stays addressable in the position it streamed and `build()` filters it out only while nothing has claimed it, so genuinely empty reasoning is still absent from the snapshot. Filtering rather than deleting is not observable through the one call path today, which builds once per run; it keeps `build()` a projection of accumulated state so the choice does not have to be revisited if the builder is ever driven incrementally. A live fan-out workflow on gpt-5-mini shows both preconditions for that loss. Reasoning arrives with an encrypted value and no visible text, the value is emitted twice for one item, and with a flow it lands between REASONING_MESSAGE_START and REASONING_MESSAGE_END rather than after it. Concurrent executors interleave their events, so a text message can start in the gap before the value arrives; replaying the captured order with that interleaving loses the encrypted reasoning before this change and keeps it after. That same interleaving exposed a second loss on the no-start text path: `_observe_text_content` opened a message over one already open under a different id, discarding its content, where `_observe_text_start` flushes first. It now flushes as well. Splitting an open text message on reasoning meant a resumed message reusing its id replayed twice under that id; the later fragment is now re-identified the way `_observe_tool_call_start` already re-identifies a split message, and only when a collision is real. A repeated start for the block already open is now a no-op rather than silently discarding the deltas folded into it. Detecting that collision by scanning the accumulated messages on every flush made snapshot building quadratic -- flat at about 8us per message before, rising to 124us per message by 4000 messages. Appends now go through one method that maintains an id index, which restores linear scaling, and a test pins the index against a future append that bypasses it. --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 115 ++++++- .../ag-ui/tests/ag_ui/test_snapshots.py | 296 ++++++++++++++++++ 2 files changed, 400 insertions(+), 11 deletions(-) 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 5b35f8d00ce..88d3907c658 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -126,11 +126,15 @@ class _WorkflowSnapshotBuilder: def __init__(self, raw_messages: list[dict[str, Any]]) -> None: self._synthesized_messages = agui_messages_to_snapshot_format(raw_messages) + self._synthesized_message_ids: set[Any] = { + message_id for message in self._synthesized_messages if (message_id := message.get("id")) + } self._emitted_messages: list[dict[str, Any]] | None = None self._open_text_message: dict[str, Any] | None = None self._open_reasoning_message: dict[str, Any] | None = None self._tool_call_message: dict[str, Any] | None = None self._tool_calls_by_id: dict[str, dict[str, Any]] = {} + self._provisional_reasoning_messages: list[dict[str, Any]] = [] self.state: dict[str, Any] | None = None self.interrupt: list[dict[str, Any]] | None = None @@ -174,7 +178,7 @@ def observe(self, event: BaseEvent) -> None: self._observe_tool_call_result(event) elif isinstance(event, ReasoningStartEvent): # A new reasoning block supersedes anything still open from the last one. - self._flush_open_reasoning_message() + self._begin_reasoning_block() elif isinstance(event, ReasoningMessageStartEvent): self._observe_reasoning_start(event) elif isinstance(event, ReasoningMessageContentEvent): @@ -188,7 +192,9 @@ def build(self) -> AGUIThreadSnapshot: """Return the replayable thread snapshot.""" self._flush_open_reasoning_message() self._flush_open_text_message() - messages = self._emitted_messages if self._emitted_messages is not None else self._synthesized_messages + messages = ( + self._emitted_messages if self._emitted_messages is not None else self._replayable_synthesized_messages() + ) return AGUIThreadSnapshot(messages=messages, state=self.state, interrupt=self.interrupt) def _observe_text_start(self, event: TextMessageStartEvent) -> None: @@ -199,6 +205,15 @@ def _observe_text_start(self, event: TextMessageStartEvent) -> None: def _observe_text_content(self, event: TextMessageContentEvent) -> None: if self._open_text_message is None or self._open_text_message.get("id") != event.message_id: + # Text arriving without a start event still opens a message, so it has to + # close what is open the way `_observe_text_start` does. Without the + # reasoning flush the reasoning outlives it and `build()` replays it after + # text that streamed later; without the text flush a message already open + # under a different id was overwritten and its content lost, which + # concurrently scheduled executors can trigger by interleaving content + # events. + self._flush_open_reasoning_message() + self._flush_open_text_message() self._open_text_message = {"id": event.message_id, "role": "assistant", "content": ""} self._open_text_message["content"] = f"{self._open_text_message.get('content', '')}{event.delta}" @@ -226,7 +241,7 @@ def _observe_tool_call_start(self, event: ToolCallStartEvent) -> None: "role": "assistant", "tool_calls": [], } - self._synthesized_messages.append(self._tool_call_message) + self._append_synthesized_message(self._tool_call_message) tool_call = { "id": event.tool_call_id, @@ -244,7 +259,10 @@ def _observe_tool_call_args(self, event: ToolCallArgsEvent) -> None: function_payload["arguments"] = f"{function_payload.get('arguments', '')}{event.delta}" def _observe_tool_call_result(self, event: ToolCallResultEvent) -> None: - self._synthesized_messages.append( + # `_observe_tool_call_start` closes open reasoning but the result path did not, + # so reasoning that streamed before a result replayed after it. + self._flush_open_reasoning_message() + self._append_synthesized_message( { "id": event.message_id, "role": "tool", @@ -261,19 +279,63 @@ def _flush_open_text_message(self) -> None: if self._open_text_message is None: return if self._open_text_message.get("content"): - self._synthesized_messages.append(self._open_text_message) + # A text message resumed after an interleaved reasoning block reuses its + # `message_id`, so flushing both fragments would put two messages under one + # id. Re-id the later fragment the way `_observe_tool_call_start` re-ids a + # split message, and only when a collision is real, so the common case + # keeps the id it streamed under. + if self._has_synthesized_message_id(self._open_text_message.get("id")): + self._open_text_message["id"] = generate_event_id() + self._append_synthesized_message(self._open_text_message) # Text between tool calls closes the current tool-call group as well. self._tool_call_message = None self._open_text_message = None + def _append_synthesized_message(self, message: dict[str, Any]) -> None: + """Append a replayable message, keeping the id index in step. + + Every append goes through here so ``_has_synthesized_message_id`` can stay a set + lookup. Scanning the accumulated list on each flush instead made snapshot + building quadratic in message count -- about 8us per message at any length + before, rising to 124us per message by 4000 messages. + """ + self._synthesized_messages.append(message) + message_id = message.get("id") + if message_id: + self._synthesized_message_ids.add(message_id) + + def _has_synthesized_message_id(self, message_id: Any) -> bool: + """Whether a message already replays under *message_id*. + + A falsy id never collides: it would otherwise match every message that carries + no id of its own. + """ + return bool(message_id) and message_id in self._synthesized_message_ids + + def _begin_reasoning_block(self) -> None: + """Close whatever is open so a reasoning block replays where it streamed. + + ``_observe_text_start`` and ``_observe_tool_call_start`` already flush open + reasoning, but nothing flushed open text, so an ``output`` event followed by a + later ``intermediate`` event left both open and let ``build()``'s flush order + decide the sequence -- reversing what streamed. Flushing here keeps the two + slots peers, so they can never both be open and that order cannot matter. + """ + self._flush_open_reasoning_message() + self._flush_open_text_message() + def _observe_reasoning_start(self, event: ReasoningMessageStartEvent) -> None: - if self._open_reasoning_message is not None and self._open_reasoning_message.get("id") != event.message_id: - self._flush_open_reasoning_message() + if self._open_reasoning_message is not None and self._open_reasoning_message.get("id") == event.message_id: + # A repeated start for the block already open is a no-op: reopening it would + # discard the deltas already folded in, and flushing first would replay two + # messages under one id. + return + self._begin_reasoning_block() self._open_reasoning_message = {"id": event.message_id, "role": "reasoning", "content": ""} def _observe_reasoning_content(self, event: ReasoningMessageContentEvent) -> None: if self._open_reasoning_message is None or self._open_reasoning_message.get("id") != event.message_id: - self._flush_open_reasoning_message() + self._begin_reasoning_block() self._open_reasoning_message = {"id": event.message_id, "role": "reasoning", "content": ""} self._open_reasoning_message["content"] = f"{self._open_reasoning_message.get('content', '')}{event.delta}" @@ -307,11 +369,42 @@ def _observe_reasoning_encrypted_value(self, event: ReasoningEncryptedValueEvent def _flush_open_reasoning_message(self) -> None: if self._open_reasoning_message is None: return + message = self._open_reasoning_message + self._open_reasoning_message = None # An encrypted-value-only block carries no display text but still has to # survive hydration, so it counts as content worth keeping. - if self._open_reasoning_message.get("content") or self._open_reasoning_message.get("encryptedValue"): - self._synthesized_messages.append(self._open_reasoning_message) - self._open_reasoning_message = None + self._append_synthesized_message(message) + if not (message.get("content") or message.get("encryptedValue")): + # A block whose only payload is protected data emits no content event, so an + # empty message here may still be waiting for its encrypted value. Measured + # on a live fan-out workflow: with a flow the value arrives between + # REASONING_MESSAGE_START and REASONING_MESSAGE_END, so output interleaved by + # a concurrent executor can flush this message before it lands; without a + # flow `_emit_text_reasoning` emits REASONING_MESSAGE_END first, widening the + # same window. Keep it addressable in the position it streamed rather than + # dropping it -- attaching the value on arrival instead would replay the + # reasoning after the output that flushed it. `build()` leaves it out of the + # snapshot while nothing has claimed it. + self._provisional_reasoning_messages.append(message) + + def _replayable_synthesized_messages(self) -> list[dict[str, Any]]: + """Synthesized messages without reasoning shells that nothing ever claimed. + + Filters rather than deletes so ``build()`` stays a projection of accumulated + state. The single caller builds once per run, so that is not observable today; it + keeps a shell excluded from one snapshot addressable afterwards, which is the + behaviour to want if the builder is ever driven incrementally. + """ + # Compare by identity: a caller-supplied message from `raw_messages` may share + # an id with one we synthesized, and must not be filtered out with it. + unclaimed = { + id(message) + for message in self._provisional_reasoning_messages + if not (message.get("content") or message.get("encryptedValue")) + } + if not unclaimed: + return self._synthesized_messages + return [message for message in self._synthesized_messages if id(message) not in unclaimed] class AgentFrameworkWorkflow: diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py index 5b4e8e694f5..a0f979669b9 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py @@ -488,3 +488,299 @@ async def test_in_memory_snapshot_store_rejects_invalid_keys() -> None: await store.delete(scope=None, thread_id="thread-1") # type: ignore[arg-type] # ty: ignore[invalid-argument-type] with pytest.raises(ValueError): await store.clear(scope="") + + +def test_workflow_snapshot_builder_keeps_output_before_later_reasoning() -> None: + """An output event that streams before reasoning replays before it, not after.""" + from ag_ui.core import ( + ReasoningMessageContentEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + TextMessageContentEvent, + TextMessageStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + # An `output` event streams assistant text. The workflow keeps running, so no + # TextMessageEndEvent arrives before the later `intermediate` event opens reasoning. + builder.observe(TextMessageStartEvent(message_id="text-1", role="assistant")) + builder.observe(TextMessageContentEvent(message_id="text-1", delta="visible output")) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="later thought")) + + messages = builder.build().messages + assert [(message["role"], message.get("content")) for message in messages] == [ + ("assistant", "visible output"), + ("reasoning", "later thought"), + ] + + +def test_workflow_snapshot_builder_keeps_encrypted_only_reasoning_addressable_across_text() -> None: + """A protected-data-only block flushed by intervening text still receives its value.""" + from ag_ui.core import ( + ReasoningEncryptedValueEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + TextMessageContentEvent, + TextMessageEndEvent, + TextMessageStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + # `_emit_text_reasoning` without a flow emits REASONING_MESSAGE_END before the + # encrypted value, and emits no content event at all when the reasoning carries + # only protected data. Intervening text must not discard the message the value + # still has to land on. + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + builder.observe(TextMessageStartEvent(message_id="text-1", role="assistant")) + builder.observe(TextMessageContentEvent(message_id="text-1", delta="visible")) + builder.observe(TextMessageEndEvent(message_id="text-1")) + builder.observe(ReasoningEncryptedValueEvent(subtype="message", entity_id="reason-1", encrypted_value="enc-1")) + + messages = builder.build().messages + reasoning = [message for message in messages if message.get("role") == "reasoning"] + assert reasoning == [{"id": "reason-1", "role": "reasoning", "content": "", "encryptedValue": "enc-1"}] + assert [message.get("role") for message in messages] == ["reasoning", "assistant"] + + +def test_workflow_snapshot_builder_splits_text_resumed_after_reasoning() -> None: + """Text resumed after a reasoning block replays around it under distinct ids.""" + from ag_ui.core import ( + ReasoningMessageContentEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + TextMessageContentEvent, + TextMessageEndEvent, + TextMessageStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(TextMessageStartEvent(message_id="text-1", role="assistant")) + builder.observe(TextMessageContentEvent(message_id="text-1", delta="part one ")) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="thought")) + # The provider reuses the text message id when the visible message resumes. + builder.observe(TextMessageContentEvent(message_id="text-1", delta="part two")) + builder.observe(TextMessageEndEvent(message_id="text-1")) + + messages = builder.build().messages + assert [(message["role"], message.get("content")) for message in messages] == [ + ("assistant", "part one "), + ("reasoning", "thought"), + ("assistant", "part two"), + ] + ids = [message.get("id") for message in messages] + assert len(set(ids)) == len(ids), f"snapshot replays two messages under one id: {ids}" + assert ids[0] == "text-1" + + +def test_workflow_snapshot_builder_keeps_reasoning_before_a_tool_result() -> None: + """Reasoning that streams before a tool result replays before it.""" + from ag_ui.core import ( + ReasoningMessageContentEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + ToolCallResultEvent, + ToolCallStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ToolCallStartEvent(tool_call_id="call-1", tool_call_name="toolA", parent_message_id="msg-1")) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="deciding")) + builder.observe(ToolCallResultEvent(message_id="result-1", tool_call_id="call-1", content="done")) + + assert [message.get("role") for message in builder.build().messages] == ["assistant", "reasoning", "tool"] + + +def test_workflow_snapshot_builder_ignores_repeated_reasoning_start_for_open_block() -> None: + """A duplicate start for the open block keeps its deltas under a single id.""" + from ag_ui.core import ReasoningMessageContentEvent, ReasoningMessageStartEvent + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="first ")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="second")) + + assert builder.build().messages == [{"id": "reason-1", "role": "reasoning", "content": "first second"}] + + +def test_workflow_snapshot_reasoning_between_tool_call_and_result_stays_convertible() -> None: + """Reasoning between a tool call and its result must not break provider adjacency.""" + from ag_ui.core import ( + ReasoningMessageContentEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + ToolCallArgsEvent, + ToolCallResultEvent, + ToolCallStartEvent, + ) + + from agent_framework_ag_ui._message_adapters import agui_messages_to_agent_framework + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ToolCallStartEvent(tool_call_id="call-1", tool_call_name="toolA", parent_message_id="msg-1")) + builder.observe(ToolCallArgsEvent(tool_call_id="call-1", delta='{"a": 1}')) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="deciding")) + builder.observe(ToolCallResultEvent(message_id="result-1", tool_call_id="call-1", content="ok")) + + snapshot = builder.build().messages + assert [message.get("role") for message in snapshot] == ["assistant", "reasoning", "tool"] + + # Reasoning is UI-only and dropped before provider conversion, so the result has to + # come back adjacent to the call it belongs to. + converted = agui_messages_to_agent_framework(snapshot) + shapes = [(message.role, [content.type for content in message.contents]) for message in converted] + call_indexes = [index for index, (_, types) in enumerate(shapes) if "function_call" in types] + result_indexes = [index for index, (_, types) in enumerate(shapes) if "function_result" in types] + assert call_indexes and result_indexes, shapes + assert result_indexes[0] == call_indexes[0] + 1, shapes + + +def test_workflow_snapshot_builder_keeps_reasoning_shell_addressable_after_a_build() -> None: + """A value arriving after a snapshot is taken still lands on the next one.""" + from ag_ui.core import ( + ReasoningEncryptedValueEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageEndEvent(message_id="reason-1")) + + # An unclaimed shell is excluded from the snapshot rather than deleted, so a store + # that saves mid-run does not make the protected value unrecoverable. + assert builder.build().messages == [] + + builder.observe(ReasoningEncryptedValueEvent(subtype="message", entity_id="reason-1", encrypted_value="enc-late")) + assert builder.build().messages == [ + {"id": "reason-1", "role": "reasoning", "content": "", "encryptedValue": "enc-late"} + ] + + +def test_workflow_snapshot_builder_message_id_index_tracks_every_append() -> None: + """The id index must cover every appended message, or collisions go undetected. + + `_has_synthesized_message_id` is a set lookup because scanning the message list on + each flush made snapshot building quadratic. That only stays correct while every + append goes through `_append_synthesized_message`, so pin the invariant rather than + the timing. + """ + from ag_ui.core import ( + ReasoningMessageContentEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + TextMessageContentEvent, + TextMessageEndEvent, + TextMessageStartEvent, + ToolCallArgsEvent, + ToolCallResultEvent, + ToolCallStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([{"id": "seed-1", "role": "user", "content": "hi"}]) + builder.observe(ToolCallStartEvent(tool_call_id="call-1", tool_call_name="toolA", parent_message_id="msg-1")) + builder.observe(ToolCallArgsEvent(tool_call_id="call-1", delta="{}")) + builder.observe(ReasoningStartEvent(message_id="reason-1")) + builder.observe(ReasoningMessageStartEvent(message_id="reason-1", role="reasoning")) + builder.observe(ReasoningMessageContentEvent(message_id="reason-1", delta="why")) + builder.observe(ToolCallResultEvent(message_id="result-1", tool_call_id="call-1", content="ok")) + builder.observe(TextMessageStartEvent(message_id="text-1", role="assistant")) + builder.observe(TextMessageContentEvent(message_id="text-1", delta="answer")) + builder.observe(TextMessageEndEvent(message_id="text-1")) + + messages = builder.build().messages + appended_ids = {message_id for message in messages if (message_id := message.get("id"))} + assert appended_ids <= builder._synthesized_message_ids, ( # pyright: ignore[reportPrivateUsage] + "an append bypassed `_append_synthesized_message`, so the id index is stale: " + f"missing {appended_ids - builder._synthesized_message_ids}" # pyright: ignore[reportPrivateUsage] + ) + + +def test_workflow_snapshot_builder_keeps_encrypted_reasoning_across_concurrent_text() -> None: + """Concurrent output must not discard a reasoning block awaiting its encrypted value. + + Event order captured from a live fan-out workflow on gpt-5-mini: with a flow, the + encrypted value is emitted right after REASONING_MESSAGE_START and before + REASONING_MESSAGE_END, and the provider sends it twice. The reasoning carries no + visible text, so a concurrent executor's text message landing in the gap before the + value arrives used to flush an empty message and lose the value. + """ + from ag_ui.core import ( + ReasoningEncryptedValueEvent, + ReasoningEndEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, + TextMessageContentEvent, + TextMessageEndEvent, + TextMessageStartEvent, + ) + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + reasoning_id = "rs_0773abc2200cb36b006aa02287c0d0819" + builder = _WorkflowSnapshotBuilder([]) + builder.observe(ReasoningStartEvent(message_id=reasoning_id)) + builder.observe(ReasoningMessageStartEvent(message_id=reasoning_id, role="reasoning")) + # A concurrently scheduled executor emits its own message here. + builder.observe(TextMessageStartEvent(message_id="concurrent-1", role="assistant")) + builder.observe(TextMessageContentEvent(message_id="concurrent-1", delta="from the other branch")) + builder.observe(TextMessageEndEvent(message_id="concurrent-1")) + # The provider emits the value twice for one item; both must be idempotent. + builder.observe(ReasoningEncryptedValueEvent(subtype="message", entity_id=reasoning_id, encrypted_value="enc")) + builder.observe(ReasoningEncryptedValueEvent(subtype="message", entity_id=reasoning_id, encrypted_value="enc")) + builder.observe(ReasoningMessageEndEvent(message_id=reasoning_id)) + builder.observe(ReasoningEndEvent(message_id=reasoning_id)) + + messages = builder.build().messages + assert [message.get("role") for message in messages] == ["reasoning", "assistant"] + assert messages[0] == {"id": reasoning_id, "role": "reasoning", "content": "", "encryptedValue": "enc"} + + +def test_workflow_snapshot_builder_keeps_both_interleaved_text_messages_without_starts() -> None: + """Content events for a second message must not overwrite the first one's text. + + `_observe_text_start` flushes an open message under a different id; the no-start path + used to overwrite it and lose the content. Concurrently scheduled executors reach + this by interleaving content events, which a live fan-out workflow does. + """ + from ag_ui.core import TextMessageContentEvent + + from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder + + builder = _WorkflowSnapshotBuilder([]) + builder.observe(TextMessageContentEvent(message_id="from-agent-a", delta="branch A")) + builder.observe(TextMessageContentEvent(message_id="from-agent-b", delta="branch B")) + + assert builder.build().messages == [ + {"id": "from-agent-a", "role": "assistant", "content": "branch A"}, + {"id": "from-agent-b", "role": "assistant", "content": "branch B"}, + ]