From 88951089cc0a5652fd9f592f17116b69ed67f5f6 Mon Sep 17 00:00:00 2001 From: ran Date: Fri, 27 Mar 2026 11:03:22 +0100 Subject: [PATCH 1/2] fix: buffer out-of-order tool results in _sanitize_tool_history --- .../agent_framework_ag_ui/_message_adapters.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) 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..5b2608523a 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,8 @@ 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 tool results that arrive before their assistant message (out-of-order history) + orphaned_tool_results: dict[str, Message] = {} pending_confirm_changes_id: str | None = None for msg in messages: @@ -77,6 +79,16 @@ 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. + # This handles out-of-order histories where tool results arrive before their + # assistant message (e.g. from CopilotKit's message snapshot merging). + for call_id in list(tool_ids): + if call_id in orphaned_tool_results: + sanitized.append(orphaned_tool_results.pop(call_id)) + if pending_tool_call_ids: + pending_tool_call_ids.discard(call_id) + continue if role_value == "user": @@ -172,6 +184,10 @@ 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 it for later. + for content in msg.contents or []: + if content.type == "function_result" and content.call_id: + orphaned_tool_results[str(content.call_id)] = msg continue keep = False for content in msg.contents or []: From 988515ee9ce32b76c553359bd84b055eaf3bb5a3 Mon Sep 17 00:00:00 2001 From: ran Date: Fri, 27 Mar 2026 12:02:34 +0100 Subject: [PATCH 2/2] fix: buffer at Content level to prevent duplicate/leaked results from batched tool messages --- .../_message_adapters.py | 23 +++++--- .../tests/ag_ui/test_message_adapters.py | 57 +++++++++++++++++++ 2 files changed, 73 insertions(+), 7 deletions(-) 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 5b2608523a..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,8 +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 tool results that arrive before their assistant message (out-of-order history) - orphaned_tool_results: dict[str, Message] = {} + # 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: @@ -81,13 +84,17 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]: ) # Re-inject any buffered tool results that belong to this assistant message. - # This handles out-of-order histories where tool results arrive before their - # assistant message (e.g. from CopilotKit's message snapshot merging). + # 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: - sanitized.append(orphaned_tool_results.pop(call_id)) + 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 @@ -184,10 +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 it for later. + # 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)] = msg + 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 ──