Python: [Bug]: `MessagesSnapshotEvent` reassigns the streamed text message ID to the tool-call message on mixed turns

[tschokokuki]

Description

Package: agent-framework-ag-ui==1.0.0rc1
File: agent_framework_ag_ui/_agent_run.py
Function: _build_messages_snapshot

Summary

When a single assistant turn produces both tool calls and trailing text,
_build_messages_snapshot assigns flow.message_id — the ID emitted on the
wire via TEXT_MESSAGE_START / TEXT_MESSAGE_CONTENT / TEXT_MESSAGE_END
for the streamed text — to the tool-call assistant message in the
snapshot, and mints a new ID for the text message.

The AG-UI protocol contract requires that a message streamed with ID X
keeps ID X in the subsequent MESSAGES_SNAPSHOT. Reassigning the
streaming ID to a different message breaks client reconciliation logic
(e.g. CopilotKit-style frontends that key off message IDs).

Reproduction (real captured SSE trace)

User turn: "saerch topbar" → the agent makes one tool call, receives the
result, then streams a short text reply.

Relevant events on the wire, in order:

TEXT_MESSAGE_START   messageId=c10b103b…                     (opened to carry parent_message_id for the tool call)
TOOL_CALL_START      toolCallId=call_apUV… parentMessageId=c10b103b…
TOOL_CALL_ARGS       … (streamed JSON args)
TOOL_CALL_END        toolCallId=call_apUV…
TOOL_CALL_RESULT     messageId=89c93e04…  toolCallId=call_apUV…
TEXT_MESSAGE_END     messageId=c10b103b…
TEXT_MESSAGE_START   messageId=dddc3821…
TEXT_MESSAGE_CONTENT messageId=dddc3821…  delta="I"
TEXT_MESSAGE_CONTENT messageId=dddc3821…  delta=" found"
...
TEXT_MESSAGE_END     messageId=dddc3821…
MESSAGES_SNAPSHOT
RUN_FINISHED

The emitted MESSAGES_SNAPSHOT:

{
  "type": "MESSAGES_SNAPSHOT",
  "messages": [
    { "id": "37b8beef…", "role": "user", "content": "saerch topbar" },
    { "id": "dddc3821…", "role": "assistant", "toolCalls": [ /* search_work_items */ ] },     // ← WRONG id reuse
    { "id": "89c93e04…", "role": "tool", "content": "Found 9 work item(s) …",
      "toolCallId": "call_apUV…" },
    { "id": "a4472f0d…", "role": "assistant",
      "content": "I found 9 work items matching \"topbar\". …" }                              // ← was streamed as dddc3821
  ]
}

Expected: the text message keeps dddc3821…; the tool-call assistant
message (which was never opened with a TEXT_MESSAGE_START of its own —
see "Secondary observation" below) gets its own ID.

Root cause

In _build_messages_snapshot (agent_framework_ag_ui/_agent_run.py, ~line 671):

# Add assistant message with tool calls only (no content)
if flow.pending_tool_calls:
    tool_call_message = {
        "id": flow.message_id or generate_event_id(),   # ← steals the streamed text id
        "role": "assistant",
        "tool_calls": flow.pending_tool_calls.copy(),
    }
    all_messages.append(tool_call_message)

# Add tool results
all_messages.extend(flow.tool_results)

# Add text-only assistant message if there is accumulated text
if flow.accumulated_text:
    content_message_id = (
        generate_event_id() if flow.pending_tool_calls else (flow.message_id or generate_event_id())
    )                                                   # ← gives the streamed text a fresh id
    all_messages.append({
        "id": content_message_id,
        "role": "assistant",
        "content": flow.accumulated_text,
    })

By the time _build_messages_snapshot runs, flow.message_id has been
reassigned by the most recent TextMessageStartEvent — in the mixed turn
shown above that is the trailing text's id (dddc3821…). Attaching it to
the tool-call message is wrong, and the text message then ends up with a
freshly minted id (a4472f0d…), so the snapshot id no longer matches the
id the client already saw on TEXT_MESSAGE_START / TEXT_MESSAGE_END.

Suggested fix

Minimal change: keep the streamed text id on the text message and give the
tool-call message a fresh id whenever both coexist.

if flow.pending_tool_calls:
    # Tool-call assistant messages are not streamed via TEXT_MESSAGE_START,
    # so they must not reuse flow.message_id (which belongs to the streamed
    # text message, if any).
    tool_call_message_id = (
        generate_event_id()
        if flow.accumulated_text
        else (flow.message_id or generate_event_id())
    )
    all_messages.append({
        "id": tool_call_message_id,
        "role": "assistant",
        "tool_calls": flow.pending_tool_calls.copy(),
    })

all_messages.extend(flow.tool_results)

if flow.accumulated_text:
    # Always preserve the streamed text id so the client can reconcile.
    content_message_id = flow.message_id or generate_event_id()
    all_messages.append({
        "id": content_message_id,
        "role": "assistant",
        "content": flow.accumulated_text,
    })

Behavior matrix after the fix:

Turn shape	Tool-call msg id	Text msg id
Text only	n/a	flow.message_id
Tool calls only	flow.message_id or new	n/a
Tool calls and text	fresh generate_event_id()	flow.message_id

Impact

Any AG-UI client that reconciles streamed text deltas with the post-stream
MESSAGES_SNAPSHOT by message ID will:

• duplicate the text message (once from the stream, once as a "new" message
  in the snapshot), and/or
• attach the streamed text to the wrong message (the tool-call one), losing
  it from the UI.

Secondary observation (not part of the proposed fix)

In the trace above, the tool-call sequence was wrapped in a
TEXT_MESSAGE_START / TEXT_MESSAGE_END pair with id c10b103b…, which
was also used as parentMessageId on TOOL_CALL_START. From the wire
contract perspective, that id arguably is the streamed identity of the
tool-call assistant message, and the snapshot's tool-call message id could
be expected to equal it. However, FlowState does not preserve that id
once the text wrapper is closed (the next TEXT_MESSAGE_START overwrites
flow.message_id), so honoring it would require an additional field and
is out of scope for the minimal fix proposed above. Flagging it here so
maintainers can decide whether to address it in the same change.

Workaround in downstream projects

Until a release is available, downstream projects can monkey-patch
agent_framework_ag_ui._agent_run._build_messages_snapshot at startup
with the replacement shown above.

Code Sample

Error Messages / Stack Traces

Package Versions

agent-framework-ag-ui==1.0.0rc1

Python Version

No response

Additional Context

No response

[github-actions]

🤖 Automated triage reproduction notes (agent-authored — trust but verify)

Agent analysis

Repro: _build_messages_snapshot() in python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py (line 679-696) assigns flow.message_id to the tool-call assistant message and mints a fresh ID for the text message when both flow.pending_tool_calls and flow.accumulated_text are non-empty. Trigger: any mixed assistant turn with tool calls followed by streamed text. Minimal repro: construct FlowState with message_id set, non-empty pending_tool_calls, and non-empty accumulated_text, then call _build_messages_snapshot and assert text_msg.id == flow.message_id.

• Failing test: python/packages/ag-ui/tests/ag_ui/test_repro_6266.py
• Files examined: python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py, python/packages/ag-ui/agent_framework_ag_ui/_run_common.py, python/packages/ag-ui/tests/ag_ui/test_run.py, python/packages/ag-ui/pyproject.toml
• Tests run: tests/ag_ui/test_repro_6266.py::TestMessageIdPreservationOnMixedTurns::test_streamed_text_id_preserved_when_both_tool_calls_and_text, tests/ag_ui/test_repro_6266.py::TestMessageIdPreservationOnMixedTurns::test_tool_call_msg_gets_fresh_id_when_text_also_present
• Reported version: 1.0.0rc1
• Current version: 1.0.0rc3