Python: fix: preserve tool call/result order in AG-UI message split

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

@@ -956,6 +956,92 @@ def _filter_modified_args(
     return result
 
 
+def _split_mixed_message_to_agui(msg: Message, role: str) -> list[dict[str, Any]]:
+    """Convert a Message that carries function_result content into ordered AG-UI messages.
+
+    A single Agent Framework message can interleave assistant content (text,
+    function_call) with one or more function_result (tool) contents -- for example a
+    parallel tool-call batch or a finalized turn. AG-UI needs each function_result as
+    its own ``tool`` message. Two ordering rules keep the transcript provider-valid:
+
+    * A ``function_call`` must precede its matching result, so a pending assistant
+      segment that carries tool calls is flushed (together with any buffered text) right
+      before the result.
+    * A text-only segment is NOT flushed before a result. Emitting a text-only assistant
+      message between an outstanding call and its result breaks the call -> result
+      adjacency providers require: ``_sanitize_tool_history`` then treats the earlier
+      call as abandoned, clears it, and drops the real result. Such text is deferred and
+      emitted after the results instead.
+
+    The source message id is kept on the first emitted message; every additional
+    message gets an independent generated id. Deriving suffixes from the source id
+    (e.g. ``f"{base_id}-1"``) risks colliding with a legitimate id elsewhere in the
+    history, which would let id-keyed clients re-collapse the split messages.
+    """
+    from ._utils import generate_event_id
+
+    messages: list[dict[str, Any]] = []
+    seg_text = ""
+    seg_tool_calls: list[dict[str, Any]] = []
+    source_id_available = bool(msg.message_id)
+
+    def next_id() -> str:
+        nonlocal source_id_available
+        if source_id_available and msg.message_id:
+            source_id_available = False
+            return msg.message_id
+        source_id_available = False
+        return generate_event_id()
+
+    def flush_segment(*, before_result: bool = False) -> None:
+        nonlocal seg_text, seg_tool_calls
+        # Before a result, only emit the segment if it carries tool calls (which must
+        # precede their results). A text-only assistant message here would separate an
+        # outstanding call from its result and cause the result to be dropped, so keep
+        # buffering that text until after the results.
+        if before_result and not seg_tool_calls:
+            return
+        if not seg_text and not seg_tool_calls:
+            return
+        assistant_msg: dict[str, Any] = {"id": next_id(), "role": role, "content": seg_text}
+        if seg_tool_calls:
+            assistant_msg["tool_calls"] = seg_tool_calls
+        messages.append(assistant_msg)
+        seg_text = ""
+        seg_tool_calls = []
+
+    for content in msg.contents:
+        if content.type == "text":
+            seg_text += content.text or ""
+        elif content.type == "function_call":
+            seg_tool_calls.append(
+                {
+                    "id": content.call_id,
+                    "type": "function",
+                    "function": {
+                        "name": content.name,
+                        "arguments": content.arguments,
+                    },
+                }
+            )
+        elif content.type == "function_result":
+            # Flush a pending assistant call (with any buffered text) so the matching
+            # call precedes its result; text-only content is deferred (see flush_segment).
+            flush_segment(before_result=True)
+            messages.append(
+                {
+                    "id": next_id(),
+                    "role": "tool",
+                    "content": content.result if content.result is not None else "",
+                    "toolCallId": content.call_id,
+                }
+            )
+
+    # Emit any deferred / trailing assistant text (e.g. a summary after the results).
+    flush_segment()
+    return messages
+
+
 def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]:
     """Convert Agent Framework messages to AG-UI format.
 
@@ -993,13 +1079,21 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An
         role_value: str = msg.role if hasattr(msg.role, "value") else msg.role
         role = FRAMEWORK_TO_AGUI_ROLE.get(role_value, "user")
 
+        # A message carrying function_result content may interleave assistant
+        # (text/function_call) and tool (function_result) segments -- e.g. parallel
+        # tool calls or a finalized turn. Split it into ordered AG-UI messages so no
+        # result is dropped and each result stays after its matching call. Messages
+        # with no result use the simple single-message form below.
+        if any(content.type == "function_result" for content in msg.contents):
+            result.extend(_split_mixed_message_to_agui(msg, role))
+            continue
+
         content_text = ""
         tool_calls: list[dict[str, Any]] = []
-        function_results: list[Any] = []
 
         for content in msg.contents:
             if content.type == "text":
-                content_text += content.text  # type: ignore[operator]
+                content_text += content.text or ""
             elif content.type == "function_call":
                 tool_calls.append(
                     {
@@ -1011,39 +1105,6 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An
                         },
                     }
                 )
-            elif content.type == "function_result":
-                function_results.append(content)
-
-        # A single Agent Framework message can carry several function_result
-        # contents (parallel tool calls). Emit one AG-UI tool message per result so
-        # none are dropped and each keeps its own toolCallId.
-        if function_results:
-            # Preserve the source id for the first result; give every additional
-            # message an independent generated id. Deriving suffixes from the source
-            # id (e.g. f"{base_id}-1") risks colliding with a legitimate id elsewhere
-            # in the history, which would let id-keyed clients re-collapse results.
-            for idx, fr in enumerate(function_results):
-                result.append(
-                    {
-                        "id": msg.message_id if (idx == 0 and msg.message_id) else generate_event_id(),
-                        "role": "tool",
-                        "content": fr.result if fr.result is not None else "",
-                        "toolCallId": fr.call_id,
-                    }
-                )
-            # A mixed message may also carry text / function_call contents alongside
-            # the tool results (e.g. a finalized assistant turn). Emit those as a
-            # separate, distinctly-identified message so they are not lost.
-            if content_text or tool_calls:
-                extra_msg: dict[str, Any] = {
-                    "id": generate_event_id(),
-                    "role": role,
-                    "content": content_text,
-                }
-                if tool_calls:
-                    extra_msg["tool_calls"] = tool_calls
-                result.append(extra_msg)
-            continue
 
         agui_msg: dict[str, Any] = {
             "id": msg.message_id if msg.message_id else generate_event_id(),  # Always include id

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

@@ -15,6 +15,7 @@
     agui_messages_to_agent_framework,
     agui_messages_to_snapshot_format,
     extract_text_from_contents,
+    normalize_agui_input_messages,
 )
 
 
@@ -963,6 +964,121 @@ def test_agent_framework_to_agui_function_result_with_text_preserves_both():
     assert text_msg["id"] != tool_msg["id"]
 
 
+def test_agent_framework_to_agui_function_call_precedes_its_result():
+    """A [function_call, function_result] message keeps the call before the result (no orphan)."""
+    msg = Message(
+        role="assistant",
+        contents=[
+            Content.from_function_call(call_id="call_a", name="get_weather", arguments={"city": "Seattle"}),
+            Content.from_function_result(call_id="call_a", result="Sunny"),
+        ],
+        message_id="mixed-order-1",
+    )
+
+    messages = agent_framework_messages_to_agui([msg])
+
+    assert len(messages) == 2
+    assistant_msg, tool_msg = messages
+    # The assistant call must be emitted before its result; a result ahead of its
+    # matching call would be an orphan that providers reject.
+    assert assistant_msg["role"] == "assistant"
+    assert [tc["id"] for tc in assistant_msg["tool_calls"]] == ["call_a"]
+    assert tool_msg["role"] == "tool"
+    assert tool_msg["toolCallId"] == "call_a"
+    assert tool_msg["content"] == "Sunny"
+    # First emitted message keeps the source id; the split-off message gets its own.
+    assert assistant_msg["id"] == "mixed-order-1"
+    assert tool_msg["id"] != assistant_msg["id"]
+
+
+def test_agent_framework_to_agui_call_result_text_order_preserved():
+    """[function_call, function_result, text] round-trips in order: call, result, then summary text."""
+    msg = Message(
+        role="assistant",
+        contents=[
+            Content.from_function_call(call_id="call_a", name="get_weather", arguments="{}"),
+            Content.from_function_result(call_id="call_a", result="Sunny"),
+            Content.from_text("It is sunny."),
+        ],
+        message_id="mixed-order-2",
+    )
+
+    messages = agent_framework_messages_to_agui([msg])
+
+    assert [m["role"] for m in messages] == ["assistant", "tool", "assistant"]
+    assert messages[0]["tool_calls"][0]["id"] == "call_a"
+    assert messages[1]["toolCallId"] == "call_a"
+    assert messages[2]["content"] == "It is sunny."
+    # Only the first emitted message reuses the source id, and all ids are distinct.
+    assert messages[0]["id"] == "mixed-order-2"
+    assert len({m["id"] for m in messages}) == 3
+
+
+def test_agent_framework_to_agui_text_before_result_deferred_after_result():
+    """Text preceding a result is emitted AFTER the result, never as an assistant-only message before it."""
+    msg = Message(
+        role="assistant",
+        contents=[
+            Content.from_text("Here is the weather."),
+            Content.from_function_result(call_id="weather-call", result="Sunny"),
+        ],
+        message_id="mixed-text-first",
+    )
+
+    messages = agent_framework_messages_to_agui([msg])
+
+    assert len(messages) == 2
+    tool_msg, text_msg = messages
+    # The tool result comes first; the text-only assistant message follows it, so it can
+    # never separate a prior outstanding call from this result.
+    assert tool_msg["role"] == "tool"
+    assert tool_msg["toolCallId"] == "weather-call"
+    assert text_msg["role"] == "assistant"
+    assert "tool_calls" not in text_msg
+    assert text_msg["content"] == "Here is the weather."
+
+
+def test_agent_framework_to_agui_text_before_result_round_trips_without_dropping():
+    """A prior call + a [text, result] message must not drop the result through sanitize_tool_history.
+
+    Regression for the MAF review finding: emitting a text-only assistant message between an
+    outstanding call and its result made ``_sanitize_tool_history`` clear the pending call and
+    drop the real result, leaving the provider with an unanswered tool call.
+    """
+    framework_messages = [
+        Message(
+            role="assistant",
+            contents=[Content.from_function_call(call_id="call_a", name="get_weather", arguments="{}")],
+            message_id="m1",
+        ),
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_text("Let me check the weather."),
+                Content.from_function_result(call_id="call_a", result="Sunny"),
+            ],
+            message_id="m2",
+        ),
+    ]
+
+    agui_messages = agent_framework_messages_to_agui(framework_messages)
+
+    # The call is immediately followed by its result (no assistant message in between).
+    assert [m["role"] for m in agui_messages] == ["assistant", "tool", "assistant"]
+    assert agui_messages[0]["tool_calls"][0]["id"] == "call_a"
+    assert agui_messages[1]["toolCallId"] == "call_a"
+
+    # Round-trip through provider normalization: the real result must survive.
+    provider_messages, _ = normalize_agui_input_messages(agui_messages, sanitize_tool_history=True)
+    surviving_result_ids = {
+        content.call_id
+        for message in provider_messages
+        for content in (message.contents or [])
+        if content.type == "function_result"
+    }
+    assert "call_a" in surviving_result_ids
+
+
 # Additional tests for better coverage