diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index 5e4fced97c..2e0f975aae 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/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 []: diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index 9508b53085..e93ad5ad08 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/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 ──