Python: [Bug]: AG-UI snapshot places assistant text after tool calls, discarding the model's emission order #7223

Description

Summary

_build_messages_snapshot always emits a turn's assistant text message after its assistant tool-call message (and after the tool results):

# agent_framework_ag_ui/_agent_run.py  (_build_messages_snapshot, ~L1532–1571 on 1.0.0rc8)
# Add assistant message with tool calls only (no content)
if flow.pending_tool_calls:
    all_messages.append(tool_call_message)          # tool-call message first
# Add tool results
all_messages.extend(flow.tool_results)
# Add text-only assistant message if there is accumulated text
# This is a separate message from the tool calls message to maintain
# the expected AG-UI protocol format (see issue #3619)
if flow.accumulated_text:
    all_messages.append({... "content": flow.accumulated_text})   # text message last
# (1.0.0rc8 adds a third accumulator with the same fixed-position problem — see
# "New in rc8" below)
all_messages.extend(flow.reasoning_messages)

A client that renders messages in array order and pins each tool's generative UI to its tool-call message (e.g. CopilotKit) therefore shows the model's lead-in narration below the tool chips it introduces:

[tool: microsoft_docs_fetch] [tool: microsoft_docs_fetch]
"Now let me research the key Microsoft Learn pages…"     ← reads as if it comes after the fetches

Keeping the text as a separate message (issue #3619) is fine. The problem is the order: the builder unconditionally puts tool calls first, even when the model produced the text before the tool calls.

Why a naive "text always first" fix is wrong

The text that accompanies a tool-using turn is not always a lead-in. Surveying real advisor threads, the assistant text that follows a tool-call message is a mix:

• Lead-in — "Now let me research…", "Let me load the landing-zone skill…" → should render before the tool chips.
• Post-tool summary — "I've loaded the skill and reviewed the CAF methodologies…", "Great — I've gathered the research…", "Here's the plan I built…" → should render after the tool chips.

In our sample the summaries outnumbered the lead-ins. So simply reordering every text message before its tool-call message would hoist summaries above the tools that produced them — worse than the status quo. A downstream consumer cannot fix this safely: lead-in vs summary is not recoverable from the normalized snapshot (both look like "text after tool-call").

The signal exists at stream time and is discarded

During streaming, the relative order is known: a lead-in arrives as TEXT_MESSAGE_* events before the TOOL_CALL_* events; a post-tool summary arrives after the tool results. _build_messages_snapshot collapses this by accumulating all text into flow.accumulated_text and all calls into flow.pending_tool_calls, then emitting tool-calls-then-text regardless of the order they streamed in.

The turn's true shape can be a sequence — e.g. text₁ → tool_calls → (results) → text₂ — but the snapshot flattens it to tool_calls → results → (text₁+text₂).

New in rc8: a third accumulator with the same problem

Between rc7 and rc8, FlowState gained reasoning_messages (agent_framework_ag_ui/_run_common.py:467) so reasoning content survives a MESSAGES_SNAPSHOT reconciliation. _build_messages_snapshot appends it last, unconditionally, after the text message — the identical fixed-position pattern this report is about, now with three message kinds (tool-calls, text, reasoning) instead of two. It doesn't change the fix: preserving stream order for reasoning_messages alongside pending_tool_calls/accumulated_text falls out of the same segment-tracking approach proposed below rather than needing separate handling.

Code Sample

Error Messages / Stack Traces

Package Versions

agent-framework-ag-ui: 1.0.0rc8, agent-framework-core: 1.11.0

Python Version

No response

Additional Context

Proposed fix

Preserve the order in which the model emitted text vs. tool calls, instead of always tool-calls-first. Concretely, track segment boundaries as the run streams (a text run, a tool-call run, another text run …) and, in _build_messages_snapshot, emit the separate assistant messages in that recorded order rather than in the fixed tool-calls-then-results-then-text order. Tool results still follow their tool-call message; only the position of text message(s) relative to the tool-call message follows the model.

This keeps the #3619 separation, adds no new protocol surface, and makes both cases correct with no downstream heuristics:

• lead-in text streamed before the calls → emitted before the tool-call message → renders above the chips;
• summary text streamed after the results → emitted after them → renders below the chips.

Alternatives, if preserving order is too invasive:

• B. Tag the text message with the stream position (e.g. additional_properties = {"emitted_before_tool_calls": bool}) so order-sensitive clients can place it correctly without guessing.
• C. A host option to choose text-before-tool-calls vs the current text-after, for apps that know their model narrates lead-ins. (Coarse; does not handle mixed turns.)

Option A is preferred because it needs no client cooperation and handles a turn that contains both a lead-in and a summary.

Evidence / references (re-verified against agent-framework-ag-ui 1.0.0rc8)