Python: preserve model emission order in AG-UI MESSAGES_SNAPSHOT#7239
Python: preserve model emission order in AG-UI MESSAGES_SNAPSHOT#7239he-yufeng wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
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_segmentsand streaming-time helpers to track ordered text/tool-calls/reasoning segments. - Updated
_build_messages_snapshotto 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. |
| 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]}) | ||
| segment_call_ids = set(segment["call_ids"]) | ||
| emitted_call_ids.update(segment_call_ids) | ||
| all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in segment_call_ids) |
There was a problem hiding this comment.
Right, fixed in 24eb54f. Only actually-emitted calls get marked emitted now, so a stale segment id that never made it into tool_calls_by_id stays eligible for the leftover fallback instead of vanishing.
| leftover_calls = [tc for tc in flow.pending_tool_calls if tc.get("id") not in emitted_call_ids] | ||
| if leftover_calls: | ||
| all_messages.append( | ||
| { | ||
| "id": tool_open_id or generate_event_id(), | ||
| "role": "assistant", | ||
| "tool_calls": [call.copy() for call in leftover_calls], | ||
| } | ||
| ) | ||
| emitted_call_ids.update(call_id for call in leftover_calls if (call_id := call.get("id")) is not None) | ||
| all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") not in emitted_call_ids) |
There was a problem hiding this comment.
Fixed in 24eb54f. The leftover path now appends the results for those calls right after the message, so they are not excluded from the final append and dropped.
Python Test Coverage Report •
Python Unit Test Overview
|
|||||||||||||||||||||||||||||||||||
| """Build MessagesSnapshotEvent from current flow state.""" | ||
| all_messages = list(snapshot_messages) | ||
|
|
||
| if flow.snapshot_segments: |
There was a problem hiding this comment.
Could this preserve accumulated_text when no text segment exists? A valid async stream can emit a function-call-only update, which preopens flow.message_id without _open_text_segment, then a text update before a tool result. _emit_text streams and accumulates that text but records no segment, so this return omits it from MESSAGES_SNAPSHOT and persisted thread history.
There was a problem hiding this comment.
Good catch, that path was real. The tool-only detection preopens message_id without going through _emit_text, so the first text had no segment and got dropped. Fixed in 24eb54f: _emit_text now opens a text segment whenever the current message id has none, and there's a regression test for exactly this shape (tool call first, text after).
|
|
||
| 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": |
There was a problem hiding this comment.
Could a tool result close the current tool-call segment? For call A -> result A -> call B, this appends B to A because emitting A's result does not change snapshot_segments; the snapshot becomes calls(A,B) -> results(A,B). That rewrites a dependent call as preceding A's result, corrupting history replay on later turns.
There was a problem hiding this comment.
Done in 24eb54f. A tool result now closes the current tool-call segment, so call A -> result A -> call B snapshots as [calls(A), result(A), calls(B), result(B)] instead of grouping B with A. There's a test asserting that exact ordering.
- Preopened message ids (tool-only path) now open a text segment when the first text arrives, so their content can't drop out of the snapshot. - A tool result closes the current tool-call segment, so call A -> result A -> call B snapshots as two pairs in stream order. - emitted_call_ids only marks calls actually emitted, keeping stale segment ids eligible for the leftover fallback. - The leftover path carries its tool results too instead of dropping them.
Fixes #7223.
Problem
_build_messages_snapshotalways 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 rc8reasoning_messagesaccumulator had the same fixed-position issue.Approach
Track segment boundaries as the run streams.
FlowStategains asnapshot_segmentslist 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, andconfirm_changesentries), and a reasoning segment per reasoning block._build_messages_snapshotnow 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.pyfail collection on missing optionalagent_framework_orchestrations/fastapideps; identical on the base commit.)test_snapshot_reasoning_orderingencoded the fixed-position layout this issue reports; updated it to the stream-faithful order.