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,124 @@ def _filter_modified_args(
     return result
 
 
+def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, Any]]]:
+    """Encode assistant contents into an AG-UI ``(content, tool_calls)`` pair.
+
+    Shared by both the single-message path (``agent_framework_messages_to_agui``) and the
+    split path (``_split_mixed_message_to_agui``) so the text / function_call
+    serialization lives in one place. A future argument-format or supported-content
+    change then updates both paths at once instead of drifting between them.
+    """
+    text = ""
+    tool_calls: list[dict[str, Any]] = []
+    for content in contents:
+        if content.type == "text":
+            text += content.text or ""
+        elif content.type == "function_call":
+            tool_calls.append(
+                {
+                    "id": content.call_id,
+                    "type": "function",
+                    "function": {
+                        "name": content.name,
+                        "arguments": content.arguments,
+                    },
+                }
+            )
+    return text, tool_calls
+
+
+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. Three 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.
+    * A new assistant segment is NOT flushed while earlier emitted calls are still
+      awaiting their results (``unresolved_call_ids``). For an interleaved batch such as
+      ``[call A, call B, result A, call C, result B, result C]``, flushing ``assistant(C)``
+      before ``result B`` would separate the still-open call B from its result, and
+      ``_sanitize_tool_history`` would drop result B as orphaned. Deferring the new
+      segment yields ``[assistant(A,B), tool(A), tool(B), assistant(C), tool(C)]`` --
+      every call stays adjacent to its results.
+
+    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_contents: list[Content] = []
+    seg_has_call = False
+    # Call ids emitted in an assistant segment whose results have not been emitted yet.
+    unresolved_call_ids: set[str] = set()
+    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() -> None:
+        nonlocal seg_contents, seg_has_call
+        if not seg_contents:
+            return
+        seg_text, seg_tool_calls = _encode_agui_segment(seg_contents)
+        seg_contents = []
+        seg_has_call = False
+        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
+            unresolved_call_ids.update(str(tc["id"]) for tc in seg_tool_calls if tc["id"] is not None)
+        messages.append(assistant_msg)
+
+    for content in msg.contents:
+        if content.type in ("text", "function_call"):
+            seg_contents.append(content)
+            seg_has_call = seg_has_call or content.type == "function_call"
+        elif content.type == "function_result":
+            # Flush the buffered call-bearing segment before its result so the call
+            # precedes it -- but only when no earlier batch is still open. While
+            # unresolved_call_ids is non-empty, flushing a new assistant segment here
+            # would split those earlier calls from their results (see docstring), so the
+            # new segment stays buffered until the open batch's results are emitted.
+            # Text-only segments are likewise deferred (they carry no function_call).
+            if seg_has_call and not unresolved_call_ids:
+                flush_segment()
+            messages.append(
+                {
+                    "id": next_id(),
+                    "role": "tool",
+                    "content": content.result if content.result is not None else "",
+                    "toolCallId": content.call_id,
+                }
+            )
+            if content.call_id is not None:
+                unresolved_call_ids.discard(str(content.call_id))
+
+    # Emit any deferred / trailing segment: buffered text (e.g. a summary after the
+    # results) and/or a new-call segment whose results arrive in a later message.
+    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,58 +1111,17 @@ 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")
 
-        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]
-            elif content.type == "function_call":
-                tool_calls.append(
-                    {
-                        "id": content.call_id,
-                        "type": "function",
-                        "function": {
-                            "name": content.name,
-                            "arguments": content.arguments,
-                        },
-                    }
-                )
-            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)
+        # 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 = _encode_agui_segment(msg.contents)
+
         agui_msg: dict[str, Any] = {
             "id": msg.message_id if msg.message_id else generate_event_id(),  # Always include id
             "role": role,

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,191 @@ 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
+
+
+def test_agent_framework_to_agui_interleaved_parallel_batch_order_preserved():
+    """An interleaved parallel batch keeps every call adjacent to its results.
+
+    A new call (C) appearing before the preceding batch's results are all emitted must not
+    start a new assistant segment ahead of the still-open results (A, B). The split defers
+    ``assistant(C)`` until B's result has been emitted.
+    """
+    msg = Message(
+        role="assistant",
+        contents=[
+            Content.from_function_call(call_id="call_a", name="fa", arguments="{}"),
+            Content.from_function_call(call_id="call_b", name="fb", arguments="{}"),
+            Content.from_function_result(call_id="call_a", result="ra"),
+            Content.from_function_call(call_id="call_c", name="fc", arguments="{}"),
+            Content.from_function_result(call_id="call_b", result="rb"),
+            Content.from_function_result(call_id="call_c", result="rc"),
+        ],
+        message_id="interleaved-1",
+    )
+
+    messages = agent_framework_messages_to_agui([msg])
+
+    # assistant(A,B) -> tool(A) -> tool(B) -> assistant(C) -> tool(C): the new call C is
+    # deferred until the open batch {A, B} is fully resolved, so no assistant message ever
+    # separates B's call from B's result.
+    assert [m["role"] for m in messages] == ["assistant", "tool", "tool", "assistant", "tool"]
+    assert [tc["id"] for tc in messages[0]["tool_calls"]] == ["call_a", "call_b"]
+    assert messages[1]["toolCallId"] == "call_a"
+    assert messages[2]["toolCallId"] == "call_b"
+    assert [tc["id"] for tc in messages[3]["tool_calls"]] == ["call_c"]
+    assert messages[4]["toolCallId"] == "call_c"
+    # First emitted message keeps the source id; every other id is independent.
+    assert messages[0]["id"] == "interleaved-1"
+    assert len({m["id"] for m in messages}) == len(messages)
+
+
+def test_agent_framework_to_agui_interleaved_batch_round_trips_without_dropping():
+    """An interleaved parallel batch must not drop any result through sanitize_tool_history.
+
+    Regression for the review finding: with the naive split, ``[call A, call B, result A,
+    call C, result B, result C]`` became ``[assistant(A,B), tool(A), assistant(C), tool(B),
+    tool(C)]``; the intervening ``assistant(C)`` cleared the pending call B, so
+    ``_sanitize_tool_history`` dropped B's real result.
+    """
+    msg = Message(
+        role="assistant",
+        contents=[
+            Content.from_function_call(call_id="call_a", name="fa", arguments="{}"),
+            Content.from_function_call(call_id="call_b", name="fb", arguments="{}"),
+            Content.from_function_result(call_id="call_a", result="ra"),
+            Content.from_function_call(call_id="call_c", name="fc", arguments="{}"),
+            Content.from_function_result(call_id="call_b", result="rb"),
+            Content.from_function_result(call_id="call_c", result="rc"),
+        ],
+        message_id="interleaved-2",
+    )
+
+    agui_messages = agent_framework_messages_to_agui([msg])
+
+    # Round-trip through provider normalization: every 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 surviving_result_ids == {"call_a", "call_b", "call_c"}
+
+
 # Additional tests for better coverage