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

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

@@ -958,6 +958,162 @@ 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, unresolved_call_ids: set[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. ``_sanitize_tool_history`` resets its pending-call set on
+    every non-tool message, so a result is dropped as orphaned whenever an assistant
+    message lands between a call and that call's result. Four 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. 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. For ``[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; deferring it yields ``[assistant(A,B), tool(A),
+      tool(B), assistant(C), tool(C)]``.
+    * A result whose own call is still buffered (not yet emitted) is QUEUED rather than
+      emitted, because emitting it would leave it ahead of its call. For
+      ``[call A, call B, result A, call C, result C, result B]`` the queued ``result C``
+      is held until ``assistant(C)`` is flushed, giving ``[assistant(A,B), tool(A),
+      tool(B), assistant(C), tool(C)]``.
+
+    ``unresolved_call_ids`` is owned by ``agent_framework_messages_to_agui`` and carried
+    across the whole conversion, not rebuilt per message: a call emitted by an earlier
+    message stays open until its result is emitted, so a later mixed message cannot slip
+    a new assistant segment between that call and its result (e.g. a prior
+    ``assistant(call A)`` followed by ``[call C, result A, result C]``).
+
+    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
+    seg_call_ids: set[str] = set()
+    # Results whose own call is still buffered; emitted once that segment is flushed.
+    queued_results: list[Content] = []
+    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
+        seg_call_ids.clear()
+        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)
+
+    def emit_result(content: Content) -> None:
+        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))
+
+    def drain_queued() -> None:
+        """Flush the buffered segment, then release the results waiting on its calls."""
+        if not queued_results:
+            return
+        flush_segment()
+        for queued in queued_results:
+            emit_result(queued)
+        queued_results.clear()
+
+    for content in msg.contents:
+        if content.type in ("text", "function_call"):
+            seg_contents.append(content)
+            if content.type == "function_call":
+                seg_has_call = True
+                if content.call_id is not None:
+                    seg_call_ids.add(str(content.call_id))
+        elif content.type == "function_result":
+            call_id = str(content.call_id) if content.call_id is not None else None
+            if call_id is not None and call_id in unresolved_call_ids:
+                # Its call is already emitted and still open, so the result can go now.
+                emit_result(content)
+                if not unresolved_call_ids:
+                    drain_queued()
+            elif seg_has_call and not unresolved_call_ids:
+                # No older batch is open: flush the buffered segment so its calls precede
+                # this result, then emit it.
+                flush_segment()
+                emit_result(content)
+                drain_queued()
+            elif call_id is not None and call_id in seg_call_ids:
+                # This result's call is still buffered behind an open older batch.
+                # Emitting now would put the result ahead of its call, so hold it.
+                queued_results.append(content)
+            else:
+                # Its call came from an already-emitted message: emit in place.
+                emit_result(content)
+
+    # Emit any deferred / trailing segment (buffered text, or a new-call segment whose
+    # results arrive in a later message), then release anything still queued behind it.
+    flush_segment()
+    for queued in queued_results:
+        emit_result(queued)
+    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.
 
@@ -970,6 +1126,25 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An
     from ._utils import generate_event_id
 
     result: list[dict[str, Any]] = []
+    # Calls emitted so far whose results have not been emitted yet. Carried across every
+    # message (mirroring _sanitize_tool_history's pending set) so a mixed message never
+    # slips a new assistant segment between an earlier call and its result.
+    unresolved_call_ids: set[str] = set()
+
+    def track_emitted(
+        role_value: str | None, tool_calls: list[dict[str, Any]] | None, tool_call_id: Any = None
+    ) -> None:
+        """Mirror _sanitize_tool_history's pending-call bookkeeping for an emitted message."""
+        if role_value == "tool":
+            if tool_call_id:
+                unresolved_call_ids.discard(str(tool_call_id))
+            return
+        # Any non-tool message resets the pending set to the calls it introduces.
+        unresolved_call_ids.clear()
+        for tool_call in tool_calls or []:
+            if isinstance(tool_call, dict) and tool_call.get("id") is not None:
+                unresolved_call_ids.add(str(tool_call["id"]))
+
     for msg in messages:
         # If already a dict (AG-UI format), ensure it has an ID and normalize keys for Pydantic
         if isinstance(msg, dict):
@@ -989,64 +1164,28 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An
                     normalized_msg["toolCallId"] = ""
             # Always append the normalized copy, not the original
             result.append(normalized_msg)
+            track_emitted(
+                normalized_msg.get("role"),
+                normalized_msg.get("tool_calls"),
+                normalized_msg.get("toolCallId"),
+            )
             continue
 
         # Convert Message to AG-UI format
         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, unresolved_call_ids))
             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,
@@ -1057,6 +1196,7 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An
             agui_msg["tool_calls"] = tool_calls
 
         result.append(agui_msg)
+        track_emitted(role, tool_calls)
 
     return result