Python: fix: buffer out-of-order tool results in _sanitize_tool_history

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

@@ -30,6 +30,11 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
     """Normalize tool ordering and inject synthetic results for AG-UI edge cases."""
     sanitized: list[Message] = []
     pending_tool_call_ids: set[str] | None = None
+    # Buffer individual function_result Contents keyed by call_id for tool messages that
+    # arrive before their assistant message (out-of-order history). Buffering at the Content
+    # level (not the Message level) prevents a multi-result tool message from being
+    # re-injected multiple times or leaking unrelated results into the wrong assistant turn.
+    orphaned_tool_results: dict[str, Any] = {}
     pending_confirm_changes_id: str | None = None
 
     for msg in messages:
@@ -77,6 +82,20 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
             pending_confirm_changes_id = (
                 str(confirm_changes_call.call_id) if confirm_changes_call and confirm_changes_call.call_id else None
             )
+
+            # Re-inject any buffered tool results that belong to this assistant message.
+            # Build a single synthetic Message containing only the matched Contents so that
+            # unrelated results from the same original message are not re-emitted, and the
+            # same result is never appended more than once.
+            matched_contents = []
+            for call_id in list(tool_ids):
+                if call_id in orphaned_tool_results:
+                    matched_contents.append(orphaned_tool_results.pop(call_id))
+                    if pending_tool_call_ids:
+                        pending_tool_call_ids.discard(call_id)
+            if matched_contents:
+                sanitized.append(Message(role="tool", contents=matched_contents))
+
             continue
 
         if role_value == "user":
@@ -172,6 +191,12 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
 
         if role_value == "tool":
             if not pending_tool_call_ids:
+                # Tool result arrived before its assistant message — buffer each Content
+                # individually so re-injection can reconstruct a filtered Message containing
+                # only the results that belong to a given assistant turn.
+                for content in msg.contents or []:
+                    if content.type == "function_result" and content.call_id:
+                        orphaned_tool_results[str(content.call_id)] = content
                 continue
             keep = False
             for content in msg.contents or []:

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

@@ -984,6 +984,63 @@ def test_sanitize_json_confirm_changes_response():
     assert len(result) >= 1
 
 
+def test_sanitize_out_of_order_tool_result_is_preserved():
+    """Tool result arriving before its assistant message is buffered and re-injected after it."""
+    from agent_framework_ag_ui._message_adapters import _sanitize_tool_history
+
+    tool_msg = Message(
+        role="tool",
+        contents=[Content.from_function_result(call_id="c1", result="chart data")],
+    )
+    assistant_msg = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="c1", name="pieChart", arguments="{}")],
+    )
+
+    result = _sanitize_tool_history([tool_msg, assistant_msg])
+
+    roles = [m.role for m in result]
+    assert roles == ["assistant", "tool"], f"Expected [assistant, tool], got {roles}"
+    tool_results = [m for m in result if m.role == "tool"]
+    assert len(tool_results) == 1
+    assert tool_results[0].contents[0].call_id == "c1"
+
+
+def test_sanitize_out_of_order_multi_result_message_no_duplicates():
+    """A batched tool message with multiple results is split — each result goes to the right turn, no duplicates."""
+    from agent_framework_ag_ui._message_adapters import _sanitize_tool_history
+
+    # Single tool message carrying results for two separate tool calls
+    batched_tool_msg = Message(
+        role="tool",
+        contents=[
+            Content.from_function_result(call_id="c1", result="result_1"),
+            Content.from_function_result(call_id="c2", result="result_2"),
+        ],
+    )
+    assistant_1 = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="c1", name="tool_a", arguments="{}")],
+    )
+    assistant_2 = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="c2", name="tool_b", arguments="{}")],
+    )
+
+    result = _sanitize_tool_history([batched_tool_msg, assistant_1, assistant_2])
+
+    tool_messages = [m for m in result if m.role == "tool"]
+    # Each result should appear exactly once — no duplicates
+    all_call_ids = [
+        c.call_id
+        for m in tool_messages
+        for c in (m.contents or [])
+        if c.type == "function_result"
+    ]
+    assert sorted(all_call_ids) == ["c1", "c2"], f"Unexpected call_ids: {all_call_ids}"
+    assert len(all_call_ids) == 2, "Duplicate tool results detected"
+
+
 # ── Deduplication edge cases ──