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 b029d3cb182..f9e8bac626f 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 @@ -1076,8 +1076,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(), } @@ -1090,10 +1093,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 be26196d03a..bd5853f439a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -833,6 +833,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