Python: preserve model emission order in AG-UI MESSAGES_SNAPSHOT

[he-yufeng]

Fixes #7223.

Problem

_build_messages_snapshot always emitted a turn's assistant text message after its tool-call message and tool results, regardless of the order the model produced them. Lead-in narration ("Let me research...") therefore rendered below the tool chips it introduces, and post-tool summaries only looked right by accident. The rc8 reasoning_messages accumulator had the same fixed-position issue.

Approach

Track segment boundaries as the run streams. FlowState gains a snapshot_segments list that the emit helpers append to: a text segment per opened text message, a tool-call segment per run of consecutive calls (shared by regular, MCP, and confirm_changes entries), and a reasoning segment per reasoning block. _build_messages_snapshot now emits those segments in recorded order, with tool results grouped after the tool-call message they answer.

The separate text/tool-call message split from #3619 is unchanged, and a tool-only turn keeps reusing the streamed message id. Flows that never went through the emit helpers fall back to the previous fixed layout, so nothing depending on it breaks.

Verification

• pytest packages/ag-ui/tests — 829 passed. (test_endpoint.py / test_handoff_replay.py fail collection on missing optional agent_framework_orchestrations / fastapi deps; identical on the base commit.)
• New tests cover lead-in before calls, text -> calls -> results -> text, tool-only id reuse, reasoning position, and the legacy fallback.
• The existing test_snapshot_reasoning_ordering encoded the fixed-position layout this issue reports; updated it to the stream-faithful order.
• pyright (package config): no new errors vs base. ruff check / format clean.

[Copilot]

Pull request overview

Preserves the model’s emission order when constructing AG-UI MESSAGES_SNAPSHOT in the Python ag-ui package, so lead-in text, tool calls/results, and reasoning appear in the same sequence the model streamed them (fixing #7223) while keeping the separate assistant text vs tool-call message split from #3619.

Changes:

• Added FlowState.snapshot_segments and streaming-time helpers to track ordered text/tool-calls/reasoning segments.
• Updated _build_messages_snapshot to emit messages based on recorded segments (with legacy fallback when no segments exist).
• Added/updated tests to assert correct snapshot ordering across mixed text/tool/reasoning scenarios.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
python/packages/ag-ui/tests/ag_ui/test_run.py	Adds tests validating snapshot ordering, tool-only ID reuse, reasoning placement, and legacy fallback behavior.
python/packages/ag-ui/agent_framework_ag_ui/_run_common.py	Introduces snapshot_segments and records segment boundaries during streaming emit helpers.
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py	Emits snapshot messages in recorded segment order via _append_segmented_snapshot_messages, with fallback to legacy behavior.

[github-actions]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/ag-ui/agent_framework_ag_ui
_agent_run.py	1140	124	89%	184–191, 238–239, 246, 355, 359, 361, 378, 405–406, 519, 533, 537, 541, 544, 549, 554, 563, 566, 573–579, 612, 624, 635, 638, 673, 727–731, 796, 811, 814, 816, 842, 868–870, 928, 930, 932, 935–939, 953, 961–966, 975–976, 1025–1028, 1039, 1047, 1079, 1094, 1108, 1120, 1150, 1154, 1157, 1159, 1199–1201, 1265, 1271–1272, 1277, 1281–1282, 1332–1334, 1346–1348, 1446, 1454, 1471, 1475, 1520, 1638, 1668–1669, 1791, 1937, 1997, 2014, 2034–2035, 2042, 2150, 2178, 2186, 2188, 2191, 2197, 2252, 2255, 2265–2266, 2273, 2320
_run_common.py	591	48	91%	76, 139–140, 142, 144, 147, 153, 155, 169, 176, 200–201, 214, 233, 235, 269, 313, 325, 327, 329, 332–336, 514, 794–795, 801–806, 1117–1118, 1123, 1125–1127, 1136, 1144, 1156, 1158–1161, 1224
TOTAL	45323	5071	88%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
9264	33 💤	0 ❌	0 🔥	2m 37s ⏱️

[moonbox3]

@he-yufeng please fix:

FAILED: pyright in packages/ag-ui
Poe => pyright
/home/runner/work/agent-framework/agent-framework/python/packages/ag-ui/agent_fr
amework_ag_ui/_agent_run.py
  /home/runner/work/agent-framework/agent-framework/python/packages/ag-ui/agent_
framework_ag_ui/_agent_run.py:1584:33 - error: Argument of type "set[Any | 
None]" cannot be assigned to parameter "s" of type "Iterable" in function 
"update"
    "set[Any | None]" is not assignable to "Iterable"
      Type parameter "_T_co@Iterable" is covariant, but "Any | None" is not a 
subtype of "str"
        Type "Any | None" is not assignable to type "str"
          "None" is not assignable to "str" (reportArgumentType)
1 error, 0 warnings, 0 informations

[he-yufeng]

Fixed in 6e5528d. The set was inferred as set[Any | None] because dict.get() keeps None in the value type even with the filter; the walrus binding narrows it properly, so update() now type-checks. Runtime behavior is unchanged.