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 fb871cdb95b..78e4b2a026b 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 @@ -86,6 +86,7 @@ _has_only_tool_calls, # type: ignore _iterate_with_context, # 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 @@ -2002,11 +2003,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: @@ -2020,8 +2016,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 @@ -2037,7 +2033,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], } @@ -2784,10 +2780,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, @@ -2809,7 +2806,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 3606428f600..bb947873d6c 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 @@ -550,12 +550,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: @@ -585,6 +601,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. @@ -614,11 +643,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, ) ) @@ -629,7 +659,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 @@ -886,11 +915,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] = { @@ -911,7 +941,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 @@ -948,12 +977,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, ) ) @@ -973,7 +1003,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 b8075b2e482..68db77e4288 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, @@ -17,6 +18,7 @@ TextMessageEndEvent, TextMessageStartEvent, ToolCallArgsEvent, + ToolCallStartEvent, ) from agent_framework import AgentResponseUpdate, Content, Message, ResponseStream from agent_framework.exceptions import AgentInvalidResponseException @@ -664,6 +666,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() @@ -677,6 +697,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() @@ -706,20 +747,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(): @@ -943,6 +990,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_keeps_protocol_fields_when_tool_arguments_use_reserved_names() -> None: """Reserved protocol fields remain controls while editedArgs carries colliding tool arguments.""" flow = FlowState(message_id="msg-1") @@ -1043,6 +1110,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 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 + def test_resume_to_tool_messages_from_interrupts_payload(): """Resume payload interrupt responses map to tool messages.""" @@ -1782,6 +1857,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", @@ -1799,6 +1876,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()