Python: Preserve AG-UI tool message IDs across snapshots

python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py

@@ -67,6 +67,7 @@
     _extract_tool_result_display,  # type: ignore
     _has_only_tool_calls,  # type: ignore
     _normalize_resume_interrupts,  # type: ignore
+    _new_tool_call_segment_id,  # type: ignore
     _reconstruct_messages_from_thread_snapshot,  # type: ignore
     _resume_contract_error,  # type: ignore
     _resolve_ui_payload,  # type: ignore
@@ -1934,11 +1935,6 @@ def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict
     they answer. Anything not covered by segment tracking falls back to the
     legacy grouping so no content is dropped.
     """
-    text_message_ids = {segment["id"] for segment in flow.snapshot_segments if segment["kind"] == "text"}
-    # A tool-only opening message (TextMessageStart with no text segment) lets
-    # the first tool-call message reuse the streamed message id, matching the
-    # legacy layout; every other tool message gets a fresh id.
-    tool_open_id = flow.message_id if flow.message_id and flow.message_id not in text_message_ids else None
     emitted_call_ids: set[str] = set()
 
     for segment in flow.snapshot_segments:
@@ -1952,8 +1948,8 @@ def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict
             ]
             if not calls:
                 continue
-            message_id = tool_open_id or generate_event_id()
-            tool_open_id = None
+            message_id = str(segment.get("id") or _new_tool_call_segment_id(flow))
+            segment["id"] = message_id
             all_messages.append({"id": message_id, "role": "assistant", "tool_calls": [call.copy() for call in calls]})
             # Only mark the calls we actually emitted; a stale segment id that
             # never made it into tool_calls_by_id must stay eligible for the
@@ -1969,7 +1965,7 @@ def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict
         leftover_ids = {cid for call in leftover_calls if (cid := call.get("id")) is not None}
         all_messages.append(
             {
-                "id": tool_open_id or generate_event_id(),
+                "id": _new_tool_call_segment_id(flow),
                 "role": "assistant",
                 "tool_calls": [call.copy() for call in leftover_calls],
             }
@@ -2591,10 +2587,11 @@ async def run_agent_stream(
 
                         # Emit confirm_changes tool call
                         confirm_id = generate_event_id()
+                        confirm_message_id = _track_tool_call_segment(flow, confirm_id)
                         yield ToolCallStartEvent(
                             tool_call_id=confirm_id,
                             tool_call_name="confirm_changes",
-                            parent_message_id=flow.message_id,
+                            parent_message_id=confirm_message_id,
                         )
                         confirm_args = {
                             "function_name": tool_name,
@@ -2616,7 +2613,6 @@ async def run_agent_stream(
                         flow.pending_tool_calls.append(confirm_entry)
                         flow.tool_calls_by_id[confirm_id] = confirm_entry
                         flow.tool_calls_ended.add(confirm_id)  # Mark as ended since we emit End event
-                        _track_tool_call_segment(flow, confirm_id)
                         flow.waiting_for_approval = True
                         flow.interrupts.append(
                             _approval_interrupt_for_function_call(

python/packages/ag-ui/agent_framework_ag_ui/_run_common.py

@@ -511,12 +511,28 @@ def _text_segment_for(flow: FlowState, message_id: str) -> dict[str, Any] | None
     return None
 
 
-def _track_tool_call_segment(flow: FlowState, tool_call_id: str) -> None:
-    """Record a tool call in the current tool segment, opening one if needed."""
+def _new_tool_call_segment_id(flow: FlowState) -> str:
+    """Allocate a snapshot ID, reusing a tool-only opening ID at most once."""
+    segment_ids = {segment.get("id") for segment in flow.snapshot_segments if segment.get("id")}
+    if flow.message_id and flow.message_id not in segment_ids:
+        return flow.message_id
+    message_id = generate_event_id()
+    while message_id in segment_ids:
+        message_id = generate_event_id()
+    return message_id
+
+
+def _track_tool_call_segment(flow: FlowState, tool_call_id: str) -> str:
+    """Record a tool call and return the message ID used by its stream events."""
+    segment: dict[str, Any]
     if flow.snapshot_segments and flow.snapshot_segments[-1]["kind"] == "tool_calls":
-        flow.snapshot_segments[-1]["call_ids"].append(tool_call_id)
+        segment = flow.snapshot_segments[-1]
+        segment.setdefault("id", _new_tool_call_segment_id(flow))
     else:
-        flow.snapshot_segments.append({"kind": "tool_calls", "call_ids": [tool_call_id]})
+        segment = {"kind": "tool_calls", "id": _new_tool_call_segment_id(flow), "call_ids": []}
+        flow.snapshot_segments.append(segment)
+    segment["call_ids"].append(tool_call_id)
+    return str(segment["id"])
 
 
 def _track_reasoning_segment(flow: FlowState, message_id: str) -> None:
@@ -575,11 +591,12 @@ def _emit_tool_call(
         if predictive_handler:
             predictive_handler.reset_streaming()
 
+        tool_message_id = _track_tool_call_segment(flow, tool_call_id)
         events.append(
             ToolCallStartEvent(
                 tool_call_id=tool_call_id,
                 tool_call_name=content.name,
-                parent_message_id=flow.message_id,
+                parent_message_id=tool_message_id,
             )
         )
 
@@ -590,7 +607,6 @@ def _emit_tool_call(
         }
         flow.pending_tool_calls.append(tool_entry)
         flow.tool_calls_by_id[tool_call_id] = tool_entry
-        _track_tool_call_segment(flow, tool_call_id)
 
     elif tool_call_id:
         flow.tool_call_id = tool_call_id
@@ -847,11 +863,12 @@ def _emit_approval_request(
 
     if require_confirmation:
         confirm_id = generate_event_id()
+        confirm_message_id = _track_tool_call_segment(flow, confirm_id)
         events.append(
             ToolCallStartEvent(
                 tool_call_id=confirm_id,
                 tool_call_name="confirm_changes",
-                parent_message_id=flow.message_id,
+                parent_message_id=confirm_message_id,
             )
         )
         args: dict[str, Any] = {
@@ -872,7 +889,6 @@ def _emit_approval_request(
         flow.pending_tool_calls.append(confirm_entry)
         flow.tool_calls_by_id[confirm_id] = confirm_entry
         flow.tool_calls_ended.add(confirm_id)
-        _track_tool_call_segment(flow, confirm_id)
 
     flow.waiting_for_approval = True
     return events
@@ -909,12 +925,13 @@ def _emit_mcp_tool_call(content: Content, flow: FlowState) -> list[BaseEvent]:
     tool_name = content.tool_name or "mcp_tool"
 
     display_name = tool_name
+    tool_message_id = _track_tool_call_segment(flow, tool_call_id)
 
     events.append(
         ToolCallStartEvent(
             tool_call_id=tool_call_id,
             tool_call_name=display_name,
-            parent_message_id=flow.message_id,
+            parent_message_id=tool_message_id,
         )
     )
 
@@ -934,7 +951,6 @@ def _emit_mcp_tool_call(content: Content, flow: FlowState) -> list[BaseEvent]:
     }
     flow.pending_tool_calls.append(tool_entry)
     flow.tool_calls_by_id[tool_call_id] = tool_entry
-    _track_tool_call_segment(flow, tool_call_id)
 
     return events
 

python/packages/ag-ui/tests/ag_ui/test_run.py

@@ -17,6 +17,7 @@
     TextMessageEndEvent,
     TextMessageStartEvent,
     ToolCallArgsEvent,
+    ToolCallStartEvent,
 )
 from agent_framework import AgentResponseUpdate, Content, Message, ResponseStream
 from agent_framework.exceptions import AgentInvalidResponseException
@@ -661,6 +662,24 @@ def test_snapshot_preserves_stream_order_around_tool_results():
     assert kinds[3][1]["id"] != kinds[0][1]["id"]
 
 
+def test_snapshot_reuses_streamed_tool_message_id_after_text():
+    """Tool-call snapshots reuse the stream ID used by the reference client merge."""
+    flow = FlowState()
+    _emit_text(Content.from_text("First, the plan."), flow)
+    tool_events = _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+    tool_start = next(event for event in tool_events if isinstance(event, ToolCallStartEvent))
+    _emit_tool_result(Content.from_function_result(call_id="call_1", result="done"), flow)
+    _emit_text(Content.from_text("And the summary."), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["text", "tool_calls", "result", "text"]
+    assert tool_start.parent_message_id is not None
+    assert kinds[1][1]["id"] == tool_start.parent_message_id
+    assert kinds[1][1]["id"] != kinds[0][1]["id"]
+
+
 def test_snapshot_tool_only_message_reuses_stream_message_id():
     """Tool-only turns keep the message id the stream opened with."""
     flow = FlowState()
@@ -674,6 +693,27 @@ def test_snapshot_tool_only_message_reuses_stream_message_id():
     assert kinds[0][1]["id"] == "tool-only-msg"
 
 
+def test_snapshot_tool_only_segments_get_unique_ids_across_reasoning():
+    """A tool-only opening ID is consumed once across separated tool segments."""
+    flow = FlowState(message_id="tool-only-msg")
+    first_events = _emit_tool_call(
+        Content.from_function_call(call_id="call_1", name="first_tool", arguments="{}"), flow
+    )
+    _emit_text_reasoning(Content.from_text("Thinking between calls."), flow)
+    second_events = _emit_tool_call(
+        Content.from_function_call(call_id="call_2", name="second_tool", arguments="{}"), flow
+    )
+
+    snapshot = _build_messages_snapshot(flow, [])
+    tool_messages = [message for kind, message in _snapshot_kinds(snapshot) if kind == "tool_calls"]
+    first_start = next(event for event in first_events if isinstance(event, ToolCallStartEvent))
+    second_start = next(event for event in second_events if isinstance(event, ToolCallStartEvent))
+
+    assert [message["id"] for message in tool_messages] == ["tool-only-msg", second_start.parent_message_id]
+    assert first_start.parent_message_id == "tool-only-msg"
+    assert second_start.parent_message_id != first_start.parent_message_id
+
+
 def test_snapshot_keeps_reasoning_in_emission_order():
     """Reasoning blocks keep their streamed position instead of always trailing."""
     flow = FlowState()
@@ -937,6 +977,26 @@ def test_emit_approval_request_populates_interrupt_metadata():
     }
 
 
+def test_emit_approval_request_reuses_confirmation_message_id_in_snapshot():
+    """Confirmation tool events and snapshots share the same message ID."""
+    flow = FlowState()
+    _emit_text(Content.from_text("Before approval."), flow)
+    text_message_id = flow.message_id
+    function_call = Content.from_function_call(call_id="call_123", name="write_doc", arguments={"content": "x"})
+    approval_content = Content.from_function_approval_request(id="approval_1", function_call=function_call)
+
+    events = _emit_approval_request(approval_content, flow)
+    confirm_start = next(
+        event for event in events if isinstance(event, ToolCallStartEvent) and event.tool_call_name == "confirm_changes"
+    )
+    snapshot = _build_messages_snapshot(flow, [])
+    kinds = _snapshot_kinds(snapshot)
+
+    assert [kind for kind, _ in kinds] == ["text", "tool_calls"]
+    assert confirm_start.parent_message_id == kinds[1][1]["id"]
+    assert confirm_start.parent_message_id != text_message_id
+
+
 def test_emit_approval_request_accumulates_multiple_interrupts():
     """Multiple approval requests in the same turn should accumulate in flow.interrupts."""
     flow = FlowState(message_id="msg-1")
@@ -1009,6 +1069,14 @@ async def test_predictive_confirmation_run_finished_interrupt_links_tool_call():
         "arguments": {"content": "Draft"},
     }
 
+    confirm_start = next(
+        event for event in events if isinstance(event, ToolCallStartEvent) and event.tool_call_name == "confirm_changes"
+    )
+    snapshots = [event for event in events if getattr(event, "type", None) == "MESSAGES_SNAPSHOT"]
+    assert snapshots
+    snapshot_tool_message = next(message for message in snapshots[-1].messages if getattr(message, "tool_calls", None))
+    assert confirm_start.parent_message_id == snapshot_tool_message.id
+
 
 def test_resume_to_tool_messages_from_interrupts_payload():
     """Resume payload interrupt responses map to tool messages."""
@@ -1583,6 +1651,8 @@ class TestEmitMcpToolCall:
     def test_produces_start_and_args_events(self):
         """MCP tool call emits ToolCallStart + ToolCallArgs events."""
         flow = FlowState()
+        _emit_text(Content.from_text("Before MCP call."), flow)
+        text_message_id = flow.message_id
         content = Content.from_mcp_server_tool_call(
             call_id="mcp_call_1",
             tool_name="search",
@@ -1600,6 +1670,12 @@ def test_produces_start_and_args_events(self):
         assert events[1].tool_call_id == "mcp_call_1"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
         assert "weather" in events[1].delta  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
 
+        snapshot = _build_messages_snapshot(flow, [])
+        kinds = _snapshot_kinds(snapshot)
+        assert [kind for kind, _ in kinds] == ["text", "tool_calls"]
+        assert events[0].parent_message_id == kinds[1][1]["id"]  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert events[0].parent_message_id != text_message_id  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
     def test_tracks_in_flow_state(self):
         """MCP tool call is tracked in flow.pending_tool_calls and tool_calls_by_id."""
         flow = FlowState()