Python: Preserve AG-UI tool message IDs across snapshots

[jstar0]

Motivation & Context

When an AG-UI run streams assistant text followed by a tool call, the stream uses the open text message ID as the tool call's parent. The final MESSAGES_SNAPSHOT represents text and tool calls as separate messages, but currently assigns the tool-call message a new ID. The reference AG-UI client merges snapshots by ID, so the tool call and its result are appended after later assistant text instead of remaining in their original position.

This affects any AG-UI frontend that follows the reference ID-based merge behavior and is reproducible with interleaved assistant text and tool calls.

Description & Review Guide

• What are the major changes?
  • Allocate a message ID when a streamed tool-call segment starts.
  • Reuse that ID for ToolCallStartEvent.parent_message_id and the corresponding snapshot assistant message.
  • Add a regression test covering text, tool call, tool result, and trailing text.
• What is the impact of these changes?
  • Streamed tool calls and final snapshots now describe the same message identity, preserving client-side ordering.
  • Tool-only turns and existing legacy snapshot fallbacks retain their current behavior.
• What do you want reviewers to focus on?
  • Confirm that the segment ID allocation matches the AG-UI event/snapshot identity contract and does not merge tool calls into text messages.

Related Issue

Fixes #7491

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix) — a workflow keeps the label and title prefix in sync automatically.

[Copilot]

🟡 Not ready to approve

Some ToolCallStartEvent emitters in _run_common.py still parent tool calls to flow.message_id while snapshots use the newly allocated segment id, which can reintroduce the stream/snapshot ID mismatch for those tool-call paths.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Pull request overview

This PR updates the Python AG-UI streaming/snapshot implementation so tool-call “assistant messages” can keep a stable identity between streamed events and the final MESSAGES_SNAPSHOT, preserving client-side ordering for interleaved text/tool-call turns.

Changes:

• Allocate and track a dedicated message ID for streamed tool-call segments and use it as ToolCallStartEvent.parent_message_id.
• Reuse the tracked tool-call segment ID when building the corresponding snapshot assistant tool_calls message.
• Add a regression test asserting that the snapshot tool-call message ID matches the streamed tool-call parent message ID after leading text.

File summaries

File	Description
python/packages/ag-ui/tests/ag_ui/test_run.py	Adds regression coverage ensuring snapshot tool-call message IDs reuse the streamed tool-call parent ID.
python/packages/ag-ui/agent_framework_ag_ui/_run_common.py	Tracks per-tool-call-segment message IDs during streaming and uses them for ToolCallStartEvent.parent_message_id.
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py	Prefers the tracked tool-call segment ID when emitting snapshot tool_calls messages.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.