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	1179	118	89%	185–192, 239–240, 247, 356, 360, 362, 379, 406–407, 520, 534, 538, 542, 545, 550, 555, 564, 567, 574–580, 613, 625, 636, 639, 674, 728–732, 797, 812, 815, 817, 843, 869–871, 929, 931, 933, 936–940, 954, 962–967, 976–977, 1026–1029, 1040, 1048, 1080, 1095, 1109, 1121, 1151, 1155, 1158, 1160, 1200–1202, 1266, 1272–1273, 1278, 1282–1283, 1481, 1489, 1519, 1523, 1568, 1686, 1716–1717, 1839, 1985, 2045, 2062, 2082–2083, 2090, 2198, 2226, 2234, 2236, 2239, 2245, 2300, 2303, 2313–2314, 2321, 2368
_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	46163	4480	90%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
9554	34 💤	0 ❌	0 🔥	2m 26s ⏱️

[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.