Python: preserve model emission order in AG-UI MESSAGES_SNAPSHOT

python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py

@@ -68,6 +68,7 @@
     _resume_contract_error,  # type: ignore
     _resolve_ui_payload,  # type: ignore
     _stringify_tool_result,  # type: ignore
+    _track_tool_call_segment,  # type: ignore
 )
 from ._snapshots import (
     _DEFAULT_STATE_INPUT_KEY,
@@ -1577,13 +1578,72 @@ def _merge_resolved_approval_results_into_snapshot(
     snapshot_messages[:] = merged_messages
 
 
+def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict[str, Any]]) -> None:
+    """Append this turn's messages in the order the model emitted them.
+
+    Segments tracked during streaming record whether text came before or after
+    tool calls (issue #7223); tool results still follow the tool-call message
+    they answer. Anything not covered by segment tracking falls back to the
+    legacy grouping so no content is dropped.
+    """
+    text_message_ids = {segment["id"] for segment in flow.snapshot_segments if segment["kind"] == "text"}
+    # A tool-only opening message (TextMessageStart with no text segment) lets
+    # the first tool-call message reuse the streamed message id, matching the
+    # legacy layout; every other tool message gets a fresh id.
+    tool_open_id = flow.message_id if flow.message_id and flow.message_id not in text_message_ids else None
+    emitted_call_ids: set[str] = set()
+
+    for segment in flow.snapshot_segments:
+        kind = segment["kind"]
+        if kind == "text":
+            if segment["text"]:
+                all_messages.append({"id": segment["id"], "role": "assistant", "content": segment["text"]})
+        elif kind == "tool_calls":
+            calls = [
+                flow.tool_calls_by_id[call_id] for call_id in segment["call_ids"] if call_id in flow.tool_calls_by_id
+            ]
+            if not calls:
+                continue
+            message_id = tool_open_id or generate_event_id()
+            tool_open_id = None
+            all_messages.append({"id": message_id, "role": "assistant", "tool_calls": [call.copy() for call in calls]})
+            # Only mark the calls we actually emitted; a stale segment id that
+            # never made it into tool_calls_by_id must stay eligible for the
+            # leftover path below rather than vanishing silently.
+            emitted_ids = {call["id"] for call in calls}
+            emitted_call_ids.update(emitted_ids)
+            all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in emitted_ids)
+        elif kind == "reasoning":
+            all_messages.extend(entry for entry in flow.reasoning_messages if entry.get("id") == segment["id"])
+
+    leftover_calls = [tc for tc in flow.pending_tool_calls if tc.get("id") not in emitted_call_ids]
+    if leftover_calls:
+        leftover_ids = {cid for call in leftover_calls if (cid := call.get("id")) is not None}
+        all_messages.append(
+            {
+                "id": tool_open_id or generate_event_id(),
+                "role": "assistant",
+                "tool_calls": [call.copy() for call in leftover_calls],
+            }
+        )
+        # Their results ride along too; without this they would be marked
+        # emitted above and then excluded from the final append below.
+        all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in leftover_ids)
+        emitted_call_ids.update(leftover_ids)
+    all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") not in emitted_call_ids)
+
+
 def _build_messages_snapshot(
     flow: FlowState,
     snapshot_messages: list[dict[str, Any]],
 ) -> MessagesSnapshotEvent:
     """Build MessagesSnapshotEvent from current flow state."""
     all_messages = list(snapshot_messages)
 
+    if flow.snapshot_segments:
+        _append_segmented_snapshot_messages(flow, all_messages)
+        return MessagesSnapshotEvent(messages=all_messages)  # type: ignore[arg-type]
+
     # Add assistant message with tool calls only (no content)
     if flow.pending_tool_calls:
         tool_call_message_id = (
@@ -2287,6 +2347,7 @@ async def run_agent_stream(
                         flow.pending_tool_calls.append(confirm_entry)
                         flow.tool_calls_by_id[confirm_id] = confirm_entry
                         flow.tool_calls_ended.add(confirm_id)  # Mark as ended since we emit End event
+                        _track_tool_call_segment(flow, confirm_id)
                         flow.waiting_for_approval = True
                         flow.interrupts.append(
                             _approval_interrupt_for_function_call(

python/packages/ag-ui/agent_framework_ag_ui/_run_common.py

@@ -467,6 +467,10 @@ class FlowState:
     reasoning_messages: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
     accumulated_reasoning: dict[str, str] = field(default_factory=dict)  # pyright: ignore[reportUnknownVariableType]
     reasoning_message_id: str | None = None
+    # Ordered text/tool_calls/reasoning segment boundaries in the order the
+    # model emitted them, so MESSAGES_SNAPSHOT can preserve that order
+    # (issue #7223) instead of the fixed tool-calls -> results -> text layout.
+    snapshot_segments: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
 
     def get_tool_name(self, call_id: str | None) -> str | None:
         """Get tool name by call ID."""
@@ -480,6 +484,37 @@ def get_pending_without_end(self) -> list[dict[str, Any]]:
         return [tc for tc in self.pending_tool_calls if tc.get("id") not in self.tool_calls_ended]
 
 
+def _open_text_segment(flow: FlowState, message_id: str) -> dict[str, Any]:
+    """Start a new snapshot text segment for a freshly opened text message."""
+    segment = {"kind": "text", "id": message_id, "text": ""}
+    flow.snapshot_segments.append(segment)
+    return segment
+
+
+def _text_segment_for(flow: FlowState, message_id: str) -> dict[str, Any] | None:
+    """Find the text segment belonging to the given message id."""
+    for segment in reversed(flow.snapshot_segments):
+        if segment["kind"] == "text" and segment["id"] == message_id:
+            return segment
+    return None
+
+
+def _track_tool_call_segment(flow: FlowState, tool_call_id: str) -> None:
+    """Record a tool call in the current tool segment, opening one if needed."""
+    if flow.snapshot_segments and flow.snapshot_segments[-1]["kind"] == "tool_calls":
+        flow.snapshot_segments[-1]["call_ids"].append(tool_call_id)
+    else:
+        flow.snapshot_segments.append({"kind": "tool_calls", "call_ids": [tool_call_id]})
+
+
+def _track_reasoning_segment(flow: FlowState, message_id: str) -> None:
+    """Record where a reasoning block sits in the snapshot emission order."""
+    for segment in flow.snapshot_segments:
+        if segment["kind"] == "reasoning" and segment["id"] == message_id:
+            return
+    flow.snapshot_segments.append({"kind": "reasoning", "id": message_id})
+
+
 def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]:
     """Emit TextMessage events for TextContent."""
     if not content.text:
@@ -492,14 +527,23 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li
     if not flow.message_id:
         flow.message_id = generate_event_id()
         flow.accumulated_text = ""
+        _open_text_segment(flow, flow.message_id)
         events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant"))
     elif flow.accumulated_text and content.text == flow.accumulated_text:
         # Guard against full-message replay chunks that can appear after streaming deltas.
         logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id)
         return []
 
+    # The message may have been pre-opened by the tool-only path, which never
+    # goes through this function, so the first text arriving later has no
+    # segment yet; without one the snapshot would drop it.
+    segment = _text_segment_for(flow, flow.message_id)
+    if segment is None:
+        segment = _open_text_segment(flow, flow.message_id)
+
     events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text))
     flow.accumulated_text += content.text
+    segment["text"] = flow.accumulated_text
     return events
 
 
@@ -534,6 +578,7 @@ def _emit_tool_call(
         }
         flow.pending_tool_calls.append(tool_entry)
         flow.tool_calls_by_id[tool_call_id] = tool_entry
+        _track_tool_call_segment(flow, tool_call_id)
 
     elif tool_call_id:
         flow.tool_call_id = tool_call_id
@@ -682,6 +727,10 @@ def _emit_tool_result_common(
             "content": result_content,
         }
     )
+    # A result closes the current tool-call segment: a later call opens a new
+    # one, so `call A -> result A -> call B` snapshots as two call/result pairs
+    # in stream order instead of grouping B with A (moonbox3's replay concern).
+    flow.snapshot_segments.append({"kind": "tool_results"})
 
     if predictive_handler:
         predictive_handler.apply_pending_updates()
@@ -809,6 +858,7 @@ def _emit_approval_request(
         flow.pending_tool_calls.append(confirm_entry)
         flow.tool_calls_by_id[confirm_id] = confirm_entry
         flow.tool_calls_ended.add(confirm_id)
+        _track_tool_call_segment(flow, confirm_id)
 
     flow.waiting_for_approval = True
     return events
@@ -870,6 +920,7 @@ def _emit_mcp_tool_call(content: Content, flow: FlowState) -> list[BaseEvent]:
     }
     flow.pending_tool_calls.append(tool_entry)
     flow.tool_calls_by_id[tool_call_id] = tool_entry
+    _track_tool_call_segment(flow, tool_call_id)
 
     return events
 
@@ -1020,6 +1071,7 @@ def _emit_text_reasoning(content: Content, flow: FlowState | None = None) -> lis
             if content.protected_data is not None:
                 reasoning_entry["encryptedValue"] = content.protected_data
             flow.reasoning_messages.append(reasoning_entry)
+            _track_reasoning_segment(flow, message_id)
         else:
             existing_entry["content"] = full_text
             if content.protected_data is not None:

python/packages/ag-ui/tests/ag_ui/test_run.py

@@ -26,6 +26,7 @@
 from agent_framework_ag_ui._agent_run import (
     PendingApprovalEntry,
     PendingApprovalKey,
+    _build_messages_snapshot,
     _build_safe_metadata,
     _canonical_approval_resume_messages,
     _create_state_context_message,
@@ -407,6 +408,162 @@ def test_emit_text_skips_when_waiting_for_approval():
     assert len(events) == 0
 
 
+def _snapshot_kinds(event):
+    """Flatten a MessagesSnapshotEvent into (kind, message) pairs for order assertions."""
+    kinds = []
+    for message in event.messages:
+        dumped = message if isinstance(message, dict) else message.model_dump(exclude_none=True)
+        if dumped.get("role") == "assistant" and dumped.get("tool_calls"):
+            kinds.append(("tool_calls", dumped))
+        elif dumped.get("role") == "assistant":
+            kinds.append(("text", dumped))
+        elif dumped.get("role") == "tool":
+            kinds.append(("result", dumped))
+        elif dumped.get("role") == "reasoning":
+            kinds.append(("reasoning", dumped))
+    return kinds
+
+
+def test_snapshot_places_lead_in_text_before_tool_calls():
+    """Lead-in narration streamed before the tool calls stays in front of them."""
+    flow = FlowState()
+    _emit_text(Content.from_text("Let me research the docs first."), flow)
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["text", "tool_calls"]
+    assert kinds[0][1]["content"] == "Let me research the docs first."
+    assert kinds[0][1]["id"] == flow.message_id
+    assert kinds[1][1]["tool_calls"][0]["id"] == "call_1"
+
+
+def test_snapshot_preserves_stream_order_around_tool_results():
+    """A turn shaped text -> calls -> results -> text snapshots in that order."""
+    flow = FlowState()
+    _emit_text(Content.from_text("First, the plan."), flow)
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+    _emit_tool_result(Content.from_function_result(call_id="call_1", result="done"), flow)
+    _emit_text(Content.from_text("And the summary."), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["text", "tool_calls", "result", "text"]
+    assert kinds[0][1]["content"] == "First, the plan."
+    assert kinds[2][1].get("toolCallId", kinds[2][1].get("tool_call_id")) == "call_1"
+    assert kinds[3][1]["content"] == "And the summary."
+    assert kinds[3][1]["id"] != kinds[0][1]["id"]
+
+
+def test_snapshot_tool_only_message_reuses_stream_message_id():
+    """Tool-only turns keep the message id the stream opened with."""
+    flow = FlowState()
+    flow.message_id = "tool-only-msg"
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls"]
+    assert kinds[0][1]["id"] == "tool-only-msg"
+
+
+def test_snapshot_keeps_reasoning_in_emission_order():
+    """Reasoning blocks keep their streamed position instead of always trailing."""
+    flow = FlowState()
+    _emit_text_reasoning(Content.from_text("thinking out loud"), flow)
+    _emit_text(Content.from_text("Answer incoming."), flow)
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["reasoning", "text", "tool_calls"]
+
+
+def test_snapshot_without_segment_tracking_keeps_legacy_layout():
+    """Flows assembled without the emit helpers keep the old fixed grouping."""
+    flow = FlowState()
+    flow.accumulated_text = "late text"
+    tool_entry = {"id": "call_1", "type": "function", "function": {"name": "docs_fetch", "arguments": ""}}
+    flow.pending_tool_calls.append(tool_entry)
+    flow.tool_calls_by_id["call_1"] = tool_entry
+    flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_1", "content": "done"})
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "result", "text"]
+
+
+def test_snapshot_includes_text_when_message_preopened_by_tool_only_path():
+    """Text that arrives after a tool-only preopen still lands in the snapshot."""
+    flow = FlowState()
+    # The tool-only detection in agent_run.py preopens message_id without
+    # going through _emit_text, so the first text has no segment yet.
+    flow.message_id = "preopened"
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+    _emit_text(Content.from_text("Let me check the docs."), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "text"]
+    assert kinds[1][1]["content"] == "Let me check the docs."
+    assert kinds[1][1]["id"] == "preopened"
+
+
+def test_snapshot_separates_calls_across_results():
+    """call A -> result A -> call B -> result B snapshots as two pairs, not grouped."""
+    flow = FlowState()
+    _emit_tool_call(Content.from_function_call(call_id="call_a", name="tool_a", arguments="{}"), flow)
+    _emit_tool_result(Content.from_function_result(call_id="call_a", result="a done"), flow)
+    _emit_tool_call(Content.from_function_call(call_id="call_b", name="tool_b", arguments="{}"), flow)
+    _emit_tool_result(Content.from_function_result(call_id="call_b", result="b done"), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "result", "tool_calls", "result"]
+    assert kinds[0][1]["tool_calls"][0]["id"] == "call_a"
+    assert kinds[1][1].get("toolCallId", kinds[1][1].get("tool_call_id")) == "call_a"
+    assert kinds[2][1]["tool_calls"][0]["id"] == "call_b"
+    assert kinds[3][1].get("toolCallId", kinds[3][1].get("tool_call_id")) == "call_b"
+
+
+def test_snapshot_leftover_call_keeps_its_result():
+    """A pending call untracked by segments still appears with its result."""
+    flow = FlowState()
+    tool_entry = {"id": "call_x", "type": "function", "function": {"name": "tool_x", "arguments": ""}}
+    flow.pending_tool_calls.append(tool_entry)
+    flow.tool_calls_by_id["call_x"] = tool_entry
+    flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_x", "content": "done"})
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "result"]
+    assert kinds[1][1].get("toolCallId", kinds[1][1].get("tool_call_id")) == "call_x"
+
+
+def test_snapshot_stale_segment_id_falls_back_to_leftover():
+    """A segment id missing from tool_calls_by_id stays eligible for the fallback."""
+    flow = FlowState()
+    flow.snapshot_segments.append({"kind": "tool_calls", "call_ids": ["call_x"]})
+    tool_entry = {"id": "call_x", "type": "function", "function": {"name": "tool_x", "arguments": ""}}
+    flow.pending_tool_calls.append(tool_entry)
+    # tool_calls_by_id intentionally lacks call_x, so the segment emits nothing
+    flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_x", "content": "done"})
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "result"]
+    assert kinds[0][1]["tool_calls"][0]["id"] == "call_x"
+
+
 def test_emit_text_skips_when_skip_text_flag():
     """Test _emit_text skips with skip_text flag."""
     flow = FlowState()
@@ -1619,7 +1776,7 @@ def test_reasoning_without_flow_does_not_error(self):
         assert isinstance(events[0], ReasoningStartEvent)
 
     def test_snapshot_reasoning_ordering(self):
-        """Reasoning messages appear after assistant text in snapshot."""
+        """Reasoning keeps its streamed position in the snapshot."""
         from agent_framework_ag_ui._agent_run import _build_messages_snapshot
 
         flow = FlowState()
@@ -1631,10 +1788,10 @@ def test_snapshot_reasoning_ordering(self):
 
         snapshot = _build_messages_snapshot(flow, [{"id": "u1", "role": "user", "content": "Hi"}])
 
-        # user -> assistant text -> reasoning
+        # reasoning streamed before the answer, so it snapshots before it too
         assert len(snapshot.messages) == 3
         roles = [_message_role(m) for m in snapshot.messages]
-        assert roles == ["user", "assistant", "reasoning"]
+        assert roles == ["user", "reasoning", "assistant"]
 
     def test_reasoning_accumulates_incremental_deltas(self):
         """Multiple reasoning deltas with the same id accumulate into one entry."""