From 0765b49c0fa56c693d2012307b0543794d8a2559 Mon Sep 17 00:00:00 2001 From: King Star Date: Tue, 4 Aug 2026 23:51:44 +0800 Subject: [PATCH 1/4] fix(ag-ui): preserve streamed tool message IDs --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 2 +- .../agent_framework_ag_ui/_run_common.py | 25 ++++++++++++++----- python/packages/ag-ui/tests/ag_ui/test_run.py | 19 ++++++++++++++ 3 files changed, 39 insertions(+), 7 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 cd976cad9c0..9622c4262d3 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 @@ -1952,7 +1952,7 @@ def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict ] if not calls: continue - message_id = tool_open_id or generate_event_id() + message_id = segment.get("id") or tool_open_id or generate_event_id() tool_open_id = None 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 diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index fc887ee9198..d797fda830c 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -511,12 +511,25 @@ 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 an ID that is distinct from any streamed text segment.""" + text_message_ids = {segment.get("id") for segment in flow.snapshot_segments if segment["kind"] == "text"} + if flow.message_id and flow.message_id not in text_message_ids: + return flow.message_id + return generate_event_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 +588,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 +604,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 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 95d9ec9c489..fe90352721a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/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() From 53282e28e3f2af35917cee2961b34a7f29c157fb Mon Sep 17 00:00:00 2001 From: King Star Date: Wed, 5 Aug 2026 00:17:06 +0800 Subject: [PATCH 2/4] fix(ag-ui): align approval and MCP tool message IDs --- .../agent_framework_ag_ui/_run_common.py | 8 +++--- python/packages/ag-ui/tests/ag_ui/test_run.py | 28 +++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index d797fda830c..8df4bf5f1d1 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -860,11 +860,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] = { @@ -885,7 +886,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 @@ -922,12 +922,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, ) ) @@ -947,7 +948,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 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 fe90352721a..22040fb26e0 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -956,6 +956,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") @@ -1602,6 +1622,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", @@ -1619,6 +1641,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() From c0e01c27d65c64c0f64e08a2dfaed26a39b554f0 Mon Sep 17 00:00:00 2001 From: King Star Date: Sun, 9 Aug 2026 22:03:46 +0800 Subject: [PATCH 3/4] fix(ag-ui): ensure unique tool segment IDs --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 16 ++++------ .../agent_framework_ag_ui/_run_common.py | 11 ++++--- python/packages/ag-ui/tests/ag_ui/test_run.py | 29 +++++++++++++++++++ 3 files changed, 42 insertions(+), 14 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 9622c4262d3..628c9db5508 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 @@ -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 = segment.get("id") or 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( diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index 8df4bf5f1d1..f993e7de266 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -512,11 +512,14 @@ def _text_segment_for(flow: FlowState, message_id: str) -> dict[str, Any] | None def _new_tool_call_segment_id(flow: FlowState) -> str: - """Allocate an ID that is distinct from any streamed text segment.""" - text_message_ids = {segment.get("id") for segment in flow.snapshot_segments if segment["kind"] == "text"} - if flow.message_id and flow.message_id not in text_message_ids: + """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 - return generate_event_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: 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 22040fb26e0..ee9c2070a19 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -693,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() @@ -1048,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.""" From a4e5a48841a8b8623ab726c127f899fb97dec876 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Wed, 19 Aug 2026 13:59:35 +0900 Subject: [PATCH 4/4] fix(ag-ui): keep tool and text snapshot IDs unique --- .../ag-ui/agent_framework_ag_ui/_run_common.py | 13 +++++++++++++ python/packages/ag-ui/tests/ag_ui/test_run.py | 15 +++++++++++---- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index f993e7de266..fb4e49dd57d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -562,6 +562,19 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id) return [] + # A tool-only response may pre-open a message before its tool-call segment + # is tracked. If that segment claims the pre-opened ID, rotate to a fresh + # text message before recording the text so snapshot IDs stay unique. + current_message_id = flow.message_id + if current_message_id and any( + segment.get("kind") == "tool_calls" and segment.get("id") == current_message_id + for segment in flow.snapshot_segments + ): + flow.message_id = generate_event_id() + flow.accumulated_text = "" + events.append(TextMessageEndEvent(message_id=current_message_id)) + events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant")) + # The message may have been pre-opened by the tool-only path, which never # goes through this function, so the first text arriving later has no # segment yet; without one the snapshot would drop it. 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 ee9c2070a19..b3eaf577116 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -7,6 +7,7 @@ import pytest from ag_ui.core import ( CustomEvent, + MessagesSnapshotEvent, ReasoningEncryptedValueEvent, ReasoningEndEvent, ReasoningMessageContentEvent, @@ -743,20 +744,26 @@ def test_snapshot_without_segment_tracking_keeps_legacy_layout(): def test_snapshot_includes_text_when_message_preopened_by_tool_only_path(): - """Text that arrives after a tool-only preopen still lands in the snapshot.""" + """Text after a tool-only preopen gets a unique snapshot message ID.""" flow = FlowState() # The tool-only detection in agent_run.py preopens message_id without # going through _emit_text, so the first text has no segment yet. flow.message_id = "preopened" _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow) - _emit_text(Content.from_text("Let me check the docs."), flow) + text_events = _emit_text(Content.from_text("Let me check the docs."), flow) event = _build_messages_snapshot(flow, []) kinds = _snapshot_kinds(event) assert [kind for kind, _ in kinds] == ["tool_calls", "text"] assert kinds[1][1]["content"] == "Let me check the docs." - assert kinds[1][1]["id"] == "preopened" + assert kinds[0][1]["id"] == "preopened" + assert kinds[1][1]["id"] != kinds[0][1]["id"] + assert len({message["id"] for _, message in kinds}) == len(kinds) + assert isinstance(text_events[0], TextMessageEndEvent) + assert isinstance(text_events[1], TextMessageStartEvent) + assert text_events[0].message_id == "preopened" + assert text_events[1].message_id == kinds[1][1]["id"] def test_snapshot_separates_calls_across_results(): @@ -1072,7 +1079,7 @@ async def test_predictive_confirmation_run_finished_interrupt_links_tool_call(): 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"] + snapshots = [event for event in events if isinstance(event, MessagesSnapshotEvent)] 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