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 9b1812a156d..539ddd21855 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,12 @@ from ag_ui.core import ( BaseEvent, MessagesSnapshotEvent, + ReasoningEncryptedValueEvent, + ReasoningEndEvent, + ReasoningMessageContentEvent, + ReasoningMessageEndEvent, + ReasoningMessageStartEvent, + ReasoningStartEvent, RunErrorEvent, RunFinishedEvent, RunStartedEvent, @@ -128,10 +134,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 @@ -173,11 +184,25 @@ 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._begin_reasoning_block() + 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 + messages = ( + self._emitted_messages if self._emitted_messages is not None else self._replayable_synthesized_messages() + ) persisted = _persistable_host_payload_history(messages) return AGUIThreadSnapshot( messages=_bound_host_payload_history(persisted), @@ -186,12 +211,22 @@ def build(self) -> AGUIThreadSnapshot: ) 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": ""} 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}" @@ -201,6 +236,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 @@ -218,7 +254,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, @@ -236,6 +272,9 @@ 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: + # `_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() message: dict[str, Any] = { "id": event.message_id, "role": "tool", @@ -252,7 +291,9 @@ def _observe_tool_call_result(self, event: ToolCallResultEvent) -> None: _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY, [{"type": "text", "text": "Tool result unavailable."}], ) - self._synthesized_messages.append(message) + # Through the helper, not `_synthesized_messages.append`, so the id index this + # branch relies on for collision detection stays in step. + self._append_synthesized_message(message) # A result closes the current tool-call group; later tool calls start a new # assistant message so replayed transcripts keep results adjacent to their # tool_calls message, which provider APIs require. @@ -262,11 +303,133 @@ 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: + # 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._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}" + + 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: + # 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 + # 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 + return + + 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. + 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: """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 0cde9c60020..4aa6a83ea27 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py @@ -221,6 +221,290 @@ 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_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 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 == [] + + def test_workflow_snapshot_builder_preserves_safe_bounded_mcp_replay() -> None: """Workflow event synthesis persists model-safe content and private Host replay data.""" from ag_ui.core import ToolCallResultEvent @@ -273,3 +557,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"}, + ] 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)