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 67d99ae625..ab074ae15c 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 @@ -68,6 +68,7 @@ _resume_contract_error, # type: ignore _resolve_ui_payload, # type: ignore _stringify_tool_result, # type: ignore + _track_tool_call_segment, # type: ignore ) from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, @@ -1577,6 +1578,61 @@ def _merge_resolved_approval_results_into_snapshot( snapshot_messages[:] = merged_messages +def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict[str, Any]]) -> None: + """Append this turn's messages in the order the model emitted them. + + Segments tracked during streaming record whether text came before or after + tool calls (issue #7223); tool results still follow the tool-call message + 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: + kind = segment["kind"] + if kind == "text": + if segment["text"]: + all_messages.append({"id": segment["id"], "role": "assistant", "content": segment["text"]}) + elif kind == "tool_calls": + calls = [ + flow.tool_calls_by_id[call_id] for call_id in segment["call_ids"] if call_id in flow.tool_calls_by_id + ] + if not calls: + continue + message_id = 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 + # never made it into tool_calls_by_id must stay eligible for the + # leftover path below rather than vanishing silently. + emitted_ids = {call["id"] for call in calls} + emitted_call_ids.update(emitted_ids) + all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in emitted_ids) + elif kind == "reasoning": + all_messages.extend(entry for entry in flow.reasoning_messages if entry.get("id") == segment["id"]) + + leftover_calls = [tc for tc in flow.pending_tool_calls if tc.get("id") not in emitted_call_ids] + if leftover_calls: + 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(), + "role": "assistant", + "tool_calls": [call.copy() for call in leftover_calls], + } + ) + # Their results ride along too; without this they would be marked + # emitted above and then excluded from the final append below. + all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in leftover_ids) + emitted_call_ids.update(leftover_ids) + all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") not in emitted_call_ids) + + def _build_messages_snapshot( flow: FlowState, snapshot_messages: list[dict[str, Any]], @@ -1584,6 +1640,10 @@ def _build_messages_snapshot( """Build MessagesSnapshotEvent from current flow state.""" all_messages = list(snapshot_messages) + if flow.snapshot_segments: + _append_segmented_snapshot_messages(flow, all_messages) + return MessagesSnapshotEvent(messages=all_messages) # type: ignore[arg-type] + # Add assistant message with tool calls only (no content) if flow.pending_tool_calls: tool_call_message_id = ( @@ -2287,6 +2347,7 @@ 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 e59ea0941f..da7cbd7d36 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 @@ -467,6 +467,10 @@ class FlowState: reasoning_messages: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType] accumulated_reasoning: dict[str, str] = field(default_factory=dict) # pyright: ignore[reportUnknownVariableType] reasoning_message_id: str | None = None + # Ordered text/tool_calls/reasoning segment boundaries in the order the + # model emitted them, so MESSAGES_SNAPSHOT can preserve that order + # (issue #7223) instead of the fixed tool-calls -> results -> text layout. + snapshot_segments: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType] def get_tool_name(self, call_id: str | None) -> str | None: """Get tool name by call ID.""" @@ -480,6 +484,37 @@ def get_pending_without_end(self) -> list[dict[str, Any]]: return [tc for tc in self.pending_tool_calls if tc.get("id") not in self.tool_calls_ended] +def _open_text_segment(flow: FlowState, message_id: str) -> dict[str, Any]: + """Start a new snapshot text segment for a freshly opened text message.""" + segment = {"kind": "text", "id": message_id, "text": ""} + flow.snapshot_segments.append(segment) + return segment + + +def _text_segment_for(flow: FlowState, message_id: str) -> dict[str, Any] | None: + """Find the text segment belonging to the given message id.""" + for segment in reversed(flow.snapshot_segments): + if segment["kind"] == "text" and segment["id"] == message_id: + return segment + 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.""" + if flow.snapshot_segments and flow.snapshot_segments[-1]["kind"] == "tool_calls": + flow.snapshot_segments[-1]["call_ids"].append(tool_call_id) + else: + flow.snapshot_segments.append({"kind": "tool_calls", "call_ids": [tool_call_id]}) + + +def _track_reasoning_segment(flow: FlowState, message_id: str) -> None: + """Record where a reasoning block sits in the snapshot emission order.""" + for segment in flow.snapshot_segments: + if segment["kind"] == "reasoning" and segment["id"] == message_id: + return + flow.snapshot_segments.append({"kind": "reasoning", "id": message_id}) + + def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]: """Emit TextMessage events for TextContent.""" if not content.text: @@ -492,14 +527,23 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li if not flow.message_id: flow.message_id = generate_event_id() flow.accumulated_text = "" + _open_text_segment(flow, flow.message_id) events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant")) elif flow.accumulated_text and content.text == flow.accumulated_text: # Guard against full-message replay chunks that can appear after streaming deltas. logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id) return [] + # 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. + segment = _text_segment_for(flow, flow.message_id) + if segment is None: + segment = _open_text_segment(flow, flow.message_id) + events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text)) flow.accumulated_text += content.text + segment["text"] = flow.accumulated_text return events @@ -534,6 +578,7 @@ 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 @@ -682,6 +727,10 @@ def _emit_tool_result_common( "content": result_content, } ) + # A result closes the current tool-call segment: a later call opens a new + # one, so `call A -> result A -> call B` snapshots as two call/result pairs + # in stream order instead of grouping B with A (moonbox3's replay concern). + flow.snapshot_segments.append({"kind": "tool_results"}) if predictive_handler: predictive_handler.apply_pending_updates() @@ -809,6 +858,7 @@ 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 @@ -870,6 +920,7 @@ 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 @@ -1020,6 +1071,7 @@ def _emit_text_reasoning(content: Content, flow: FlowState | None = None) -> lis if content.protected_data is not None: reasoning_entry["encryptedValue"] = content.protected_data flow.reasoning_messages.append(reasoning_entry) + _track_reasoning_segment(flow, message_id) else: existing_entry["content"] = full_text if content.protected_data is not None: 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 bd5853f439..217983c3d9 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -26,6 +26,7 @@ from agent_framework_ag_ui._agent_run import ( PendingApprovalEntry, PendingApprovalKey, + _build_messages_snapshot, _build_safe_metadata, _canonical_approval_resume_messages, _create_state_context_message, @@ -407,6 +408,162 @@ def test_emit_text_skips_when_waiting_for_approval(): assert len(events) == 0 +def _snapshot_kinds(event): + """Flatten a MessagesSnapshotEvent into (kind, message) pairs for order assertions.""" + kinds = [] + for message in event.messages: + dumped = message if isinstance(message, dict) else message.model_dump(exclude_none=True) + if dumped.get("role") == "assistant" and dumped.get("tool_calls"): + kinds.append(("tool_calls", dumped)) + elif dumped.get("role") == "assistant": + kinds.append(("text", dumped)) + elif dumped.get("role") == "tool": + kinds.append(("result", dumped)) + elif dumped.get("role") == "reasoning": + kinds.append(("reasoning", dumped)) + return kinds + + +def test_snapshot_places_lead_in_text_before_tool_calls(): + """Lead-in narration streamed before the tool calls stays in front of them.""" + flow = FlowState() + _emit_text(Content.from_text("Let me research the docs first."), flow) + _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow) + + event = _build_messages_snapshot(flow, []) + + kinds = _snapshot_kinds(event) + assert [kind for kind, _ in kinds] == ["text", "tool_calls"] + assert kinds[0][1]["content"] == "Let me research the docs first." + assert kinds[0][1]["id"] == flow.message_id + assert kinds[1][1]["tool_calls"][0]["id"] == "call_1" + + +def test_snapshot_preserves_stream_order_around_tool_results(): + """A turn shaped text -> calls -> results -> text snapshots in that order.""" + flow = FlowState() + _emit_text(Content.from_text("First, the plan."), flow) + _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow) + _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 kinds[0][1]["content"] == "First, the plan." + assert kinds[2][1].get("toolCallId", kinds[2][1].get("tool_call_id")) == "call_1" + assert kinds[3][1]["content"] == "And the summary." + assert kinds[3][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() + flow.message_id = "tool-only-msg" + _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow) + + event = _build_messages_snapshot(flow, []) + + kinds = _snapshot_kinds(event) + assert [kind for kind, _ in kinds] == ["tool_calls"] + assert kinds[0][1]["id"] == "tool-only-msg" + + +def test_snapshot_keeps_reasoning_in_emission_order(): + """Reasoning blocks keep their streamed position instead of always trailing.""" + flow = FlowState() + _emit_text_reasoning(Content.from_text("thinking out loud"), flow) + _emit_text(Content.from_text("Answer incoming."), flow) + _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow) + + event = _build_messages_snapshot(flow, []) + + kinds = _snapshot_kinds(event) + assert [kind for kind, _ in kinds] == ["reasoning", "text", "tool_calls"] + + +def test_snapshot_without_segment_tracking_keeps_legacy_layout(): + """Flows assembled without the emit helpers keep the old fixed grouping.""" + flow = FlowState() + flow.accumulated_text = "late text" + tool_entry = {"id": "call_1", "type": "function", "function": {"name": "docs_fetch", "arguments": ""}} + flow.pending_tool_calls.append(tool_entry) + flow.tool_calls_by_id["call_1"] = tool_entry + flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_1", "content": "done"}) + + event = _build_messages_snapshot(flow, []) + + kinds = _snapshot_kinds(event) + assert [kind for kind, _ in kinds] == ["tool_calls", "result", "text"] + + +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.""" + 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) + + 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" + + +def test_snapshot_separates_calls_across_results(): + """call A -> result A -> call B -> result B snapshots as two pairs, not grouped.""" + flow = FlowState() + _emit_tool_call(Content.from_function_call(call_id="call_a", name="tool_a", arguments="{}"), flow) + _emit_tool_result(Content.from_function_result(call_id="call_a", result="a done"), flow) + _emit_tool_call(Content.from_function_call(call_id="call_b", name="tool_b", arguments="{}"), flow) + _emit_tool_result(Content.from_function_result(call_id="call_b", result="b done"), flow) + + event = _build_messages_snapshot(flow, []) + + kinds = _snapshot_kinds(event) + assert [kind for kind, _ in kinds] == ["tool_calls", "result", "tool_calls", "result"] + assert kinds[0][1]["tool_calls"][0]["id"] == "call_a" + assert kinds[1][1].get("toolCallId", kinds[1][1].get("tool_call_id")) == "call_a" + assert kinds[2][1]["tool_calls"][0]["id"] == "call_b" + assert kinds[3][1].get("toolCallId", kinds[3][1].get("tool_call_id")) == "call_b" + + +def test_snapshot_leftover_call_keeps_its_result(): + """A pending call untracked by segments still appears with its result.""" + flow = FlowState() + tool_entry = {"id": "call_x", "type": "function", "function": {"name": "tool_x", "arguments": ""}} + flow.pending_tool_calls.append(tool_entry) + flow.tool_calls_by_id["call_x"] = tool_entry + flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_x", "content": "done"}) + + event = _build_messages_snapshot(flow, []) + + kinds = _snapshot_kinds(event) + assert [kind for kind, _ in kinds] == ["tool_calls", "result"] + assert kinds[1][1].get("toolCallId", kinds[1][1].get("tool_call_id")) == "call_x" + + +def test_snapshot_stale_segment_id_falls_back_to_leftover(): + """A segment id missing from tool_calls_by_id stays eligible for the fallback.""" + flow = FlowState() + flow.snapshot_segments.append({"kind": "tool_calls", "call_ids": ["call_x"]}) + tool_entry = {"id": "call_x", "type": "function", "function": {"name": "tool_x", "arguments": ""}} + flow.pending_tool_calls.append(tool_entry) + # tool_calls_by_id intentionally lacks call_x, so the segment emits nothing + flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_x", "content": "done"}) + + event = _build_messages_snapshot(flow, []) + + kinds = _snapshot_kinds(event) + assert [kind for kind, _ in kinds] == ["tool_calls", "result"] + assert kinds[0][1]["tool_calls"][0]["id"] == "call_x" + + def test_emit_text_skips_when_skip_text_flag(): """Test _emit_text skips with skip_text flag.""" flow = FlowState() @@ -1619,7 +1776,7 @@ def test_reasoning_without_flow_does_not_error(self): assert isinstance(events[0], ReasoningStartEvent) def test_snapshot_reasoning_ordering(self): - """Reasoning messages appear after assistant text in snapshot.""" + """Reasoning keeps its streamed position in the snapshot.""" from agent_framework_ag_ui._agent_run import _build_messages_snapshot flow = FlowState() @@ -1631,10 +1788,10 @@ def test_snapshot_reasoning_ordering(self): snapshot = _build_messages_snapshot(flow, [{"id": "u1", "role": "user", "content": "Hi"}]) - # user -> assistant text -> reasoning + # reasoning streamed before the answer, so it snapshots before it too assert len(snapshot.messages) == 3 roles = [_message_role(m) for m in snapshot.messages] - assert roles == ["user", "assistant", "reasoning"] + assert roles == ["user", "reasoning", "assistant"] def test_reasoning_accumulates_incremental_deltas(self): """Multiple reasoning deltas with the same id accumulate into one entry."""