Python: fix(ag-ui): preserve streamed text message id#6269
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Updates message snapshot generation so that, when both tool calls and streamed text are present, the streamed text message keeps the original flow.message_id while the tool-call message gets a different ID.
Changes:
- Adjust
_build_messages_snapshotmessage ID assignment to preserve streamed text message IDs in mixed tool-call/text snapshots. - Add a regression test covering the mixed tool-call + text case to ensure IDs are preserved/isolated correctly.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py | Changes assistant/tool-call vs assistant/text message ID selection logic in snapshots. |
| python/packages/ag-ui/tests/ag_ui/test_run.py | Adds a test asserting streamed text message ID is preserved when tool calls are also present. |
| result = _build_messages_snapshot(flow, []) | ||
|
|
||
| assistant_tool_msg = result.messages[0] | ||
| assistant_text_msg = result.messages[2] | ||
|
|
||
| assert assistant_text_msg.id == "streamed-text-msg" | ||
| assert assistant_tool_msg.id != "streamed-text-msg" |
| tool_call_message_id = ( | ||
| generate_event_id() if flow.accumulated_text else (flow.message_id or generate_event_id()) | ||
| ) |
|
@cyphercodes please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
Motivation and Context
Fixes #6266.
Mixed AG-UI snapshots with both pending tool calls and streamed trailing text were assigning the streamed text
message_idto the tool-call-only assistant message. The text message then received a generated ID, which broke reconciliation for the already-streamed text.Description
When a snapshot contains both tool calls and accumulated text, this changes the tool-call-only assistant message to use a generated ID and preserves
flow.message_idfor the text assistant message.A regression test covers the mixed tool-call/text snapshot and verifies that the text message keeps the streamed message ID while the tool-call message remains distinct.
Contribution Checklist
Validation:
uv run --group dev pytest packages/ag-ui/tests/ag_ui/test_run.py::TestBuildMessagesSnapshot -quv run --group dev poe syntax -P ag-uiuv run --group dev poe pyright -P ag-uiuv run --group dev poe test -P ag-ui