From 16ee9a0a9cbfe70254cc464b7106b79aa1f3b9bc Mon Sep 17 00:00:00 2001
From: cyphercodes <cyphercodes@users.noreply.github.com>
Date: Tue, 2 Jun 2026 18:52:45 +0300
Subject: [PATCH] fix(ag-ui): preserve streamed text message id

---
 .../ag-ui/agent_framework_ag_ui/_agent_run.py | 10 +++++-----
 python/packages/ag-ui/tests/ag_ui/test_run.py | 20 +++++++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
index 75e8e242dc6..0d496088dbd 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
@@ -677,8 +677,11 @@ def _build_messages_snapshot(
 
     # Add assistant message with tool calls only (no content)
     if flow.pending_tool_calls:
+        tool_call_message_id = (
+            generate_event_id() if flow.accumulated_text else (flow.message_id or generate_event_id())
+        )
         tool_call_message = {
-            "id": flow.message_id or generate_event_id(),
+            "id": tool_call_message_id,
             "role": "assistant",
             "tool_calls": flow.pending_tool_calls.copy(),
         }
@@ -691,10 +694,7 @@ def _build_messages_snapshot(
     # This is a separate message from the tool calls message to maintain
     # the expected AG-UI protocol format (see issue #3619)
     if flow.accumulated_text:
-        # Use a new ID for the content message if we had tool calls (separate message)
-        content_message_id = (
-            generate_event_id() if flow.pending_tool_calls else (flow.message_id or generate_event_id())
-        )
+        content_message_id = flow.message_id or generate_event_id()
         all_messages.append(
             {
                 "id": content_message_id,
diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py
index a901b61e6e9..8dead7fe21a 100644
--- a/python/packages/ag-ui/tests/ag_ui/test_run.py
+++ b/python/packages/ag-ui/tests/ag_ui/test_run.py
@@ -685,6 +685,26 @@ def test_tool_calls_and_text_are_separate_messages(self):
         # The text message should have a different ID than the tool call message
         assert assistant_text_msg.id != assistant_tool_msg.id
 
+    def test_tool_calls_and_text_preserve_streamed_text_message_id(self):
+        """Mixed tool-call/text snapshots preserve the streamed text message ID."""
+        from agent_framework_ag_ui._agent_run import FlowState, _build_messages_snapshot
+
+        flow = FlowState()
+        flow.message_id = "streamed-text-msg"
+        flow.pending_tool_calls = [
+            {"id": "call_1", "function": {"name": "get_weather", "arguments": '{"city": "NYC"}'}},
+        ]
+        flow.accumulated_text = "Here is the weather information."
+        flow.tool_results = [{"id": "result-1", "role": "tool", "content": '{"temp": 72}', "toolCallId": "call_1"}]
+
+        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"
+
     def test_only_tool_calls_no_text(self):
         """Test snapshot with only tool calls and no accumulated text."""
         from agent_framework_ag_ui._agent_run import FlowState, _build_messages_snapshot