From 80d2a5e20df7f46debfdc88707aa4c7f9bd8b3a4 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 2 Sep 2026 13:27:08 +0530 Subject: [PATCH 1/4] Python: fix: preserve tool call/result order in AG-UI message split When a single Agent Framework message interleaved assistant content (function_call / text) with function_result content, the AG-UI converter emitted every tool result first and the assistant call/text segment afterwards. A [function_call, function_result] message therefore round-tripped as [function_result, function_call], reversing the transcript. Providers require a tool result to follow its matching assistant call, so the reversed order produced an orphan result (rejected or dropped) while the call stayed pending. Walk msg.contents in order and flush any accumulated assistant segment before each function_result, so calls always precede their results. Parallel results and trailing summary text are preserved, and the source message id is still kept on the first emitted message with independent generated ids for the rest. Follow-up to #7980 (addresses moonbox3's post-merge review feedback). --- .../_message_adapters.py | 120 +++++++++++++----- .../tests/ag_ui/test_message_adapters.py | 50 ++++++++ 2 files changed, 135 insertions(+), 35 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 b9db16161f..cb8b995623 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 @@ -956,6 +956,81 @@ def _filter_modified_args( return result +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, but the assistant call that produced a result must still + precede it; a result emitted ahead of its matching call is an orphan that providers + reject or drop. This walks ``msg.contents`` in order and flushes any accumulated + assistant segment (text/tool_calls) before each result, preserving the original + call -> result ordering. + + 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_text = "" + seg_tool_calls: list[dict[str, Any]] = [] + 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_text, seg_tool_calls + 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 + messages.append(assistant_msg) + seg_text = "" + seg_tool_calls = [] + + for content in msg.contents: + if content.type == "text": + seg_text += content.text or "" + elif content.type == "function_call": + seg_tool_calls.append( + { + "id": content.call_id, + "type": "function", + "function": { + "name": content.name, + "arguments": content.arguments, + }, + } + ) + elif content.type == "function_result": + # Flush any assistant call/text accumulated before this result so the + # matching call precedes it, then emit the result as its own tool message. + flush_segment() + messages.append( + { + "id": next_id(), + "role": "tool", + "content": content.result if content.result is not None else "", + "toolCallId": content.call_id, + } + ) + + # Emit any trailing assistant segment (e.g. a summary text after the results). + 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,13 +1068,21 @@ 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") + # 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: list[dict[str, Any]] = [] - function_results: list[Any] = [] for content in msg.contents: if content.type == "text": - content_text += content.text # type: ignore[operator] + content_text += content.text or "" elif content.type == "function_call": tool_calls.append( { @@ -1011,39 +1094,6 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An }, } ) - 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) - continue agui_msg: dict[str, Any] = { "id": msg.message_id if msg.message_id else generate_event_id(), # Always include id 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 1749cafe72..befd5ba9a8 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 @@ -963,6 +963,56 @@ 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 + + # Additional tests for better coverage From 60afad1cdb7e3aad4e2bbe1ae2ff7c117f7d5b13 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 2 Sep 2026 15:23:04 +0530 Subject: [PATCH 2/4] Python: defer text-only content before AG-UI tool results The mixed-message split flushed every accumulated segment before each function_result, including text-only segments. A [text, function_result] message whose result answers a call in a prior assistant message therefore emitted an assistant-only text message between the outstanding call and its result. _sanitize_tool_history then treats the call as abandoned, clears it, and drops the real result, leaving the provider with an unanswered tool call. Defer text-only content and emit it after the results; only a segment that carries tool calls is flushed before a result (with any buffered text coalesced), so a call still precedes its matching result and no assistant-only message is ever inserted between an outstanding call and its result. Addresses the MAF automated-review finding on this PR. --- .../_message_adapters.py | 31 ++++++--- .../tests/ag_ui/test_message_adapters.py | 66 +++++++++++++++++++ 2 files changed, 87 insertions(+), 10 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 cb8b995623..7319413619 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 @@ -962,11 +962,16 @@ def _split_mixed_message_to_agui(msg: Message, role: str) -> list[dict[str, Any] 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, but the assistant call that produced a result must still - precede it; a result emitted ahead of its matching call is an orphan that providers - reject or drop. This walks ``msg.contents`` in order and flushes any accumulated - assistant segment (text/tool_calls) before each result, preserving the original - call -> result ordering. + its own ``tool`` message. Two 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. 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 @@ -988,8 +993,14 @@ def next_id() -> str: source_id_available = False return generate_event_id() - def flush_segment() -> None: + def flush_segment(*, before_result: bool = False) -> None: nonlocal seg_text, seg_tool_calls + # Before a result, only emit the segment if it carries tool calls (which must + # precede their results). A text-only assistant message here would separate an + # outstanding call from its result and cause the result to be dropped, so keep + # buffering that text until after the results. + if before_result and not seg_tool_calls: + return if not seg_text and not seg_tool_calls: return assistant_msg: dict[str, Any] = {"id": next_id(), "role": role, "content": seg_text} @@ -1014,9 +1025,9 @@ def flush_segment() -> None: } ) elif content.type == "function_result": - # Flush any assistant call/text accumulated before this result so the - # matching call precedes it, then emit the result as its own tool message. - flush_segment() + # Flush a pending assistant call (with any buffered text) so the matching + # call precedes its result; text-only content is deferred (see flush_segment). + flush_segment(before_result=True) messages.append( { "id": next_id(), @@ -1026,7 +1037,7 @@ def flush_segment() -> None: } ) - # Emit any trailing assistant segment (e.g. a summary text after the results). + # Emit any deferred / trailing assistant text (e.g. a summary after the results). flush_segment() return messages 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 befd5ba9a8..713908a4a1 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 @@ -15,6 +15,7 @@ agui_messages_to_agent_framework, agui_messages_to_snapshot_format, extract_text_from_contents, + normalize_agui_input_messages, ) @@ -1013,6 +1014,71 @@ def test_agent_framework_to_agui_call_result_text_order_preserved(): 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 + + # Additional tests for better coverage From cbfdced00e7b4a7ef14b46427a1e2a4e8940aae4 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Thu, 3 Sep 2026 13:44:56 +0530 Subject: [PATCH 3/4] Python: fix AG-UI split for interleaved parallel tool-call batches Defer a new assistant segment while earlier emitted calls still await their results, so an interleaved batch like [call A, call B, result A, call C, result B, result C] no longer places assistant(C) between call B and result B (which _sanitize_tool_history would drop as orphaned). Extract a shared _encode_agui_segment() used by both the split and single-message paths so the text/function_call serialization can't drift. --- .../_message_adapters.py | 108 ++++++++++-------- .../tests/ag_ui/test_message_adapters.py | 70 ++++++++++++ 2 files changed, 132 insertions(+), 46 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 7319413619..e3e3cc34ca 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 @@ -956,13 +956,40 @@ 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. Two ordering rules keep the transcript provider-valid: + 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 @@ -972,6 +999,13 @@ def _split_mixed_message_to_agui(msg: Message, role: str) -> list[dict[str, Any] 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 @@ -981,8 +1015,10 @@ def _split_mixed_message_to_agui(msg: Message, role: str) -> list[dict[str, Any] from ._utils import generate_event_id messages: list[dict[str, Any]] = [] - seg_text = "" - seg_tool_calls: 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: @@ -993,41 +1029,34 @@ def next_id() -> str: source_id_available = False return generate_event_id() - def flush_segment(*, before_result: bool = False) -> None: - nonlocal seg_text, seg_tool_calls - # Before a result, only emit the segment if it carries tool calls (which must - # precede their results). A text-only assistant message here would separate an - # outstanding call from its result and cause the result to be dropped, so keep - # buffering that text until after the results. - if before_result and not seg_tool_calls: + 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) - seg_text = "" - seg_tool_calls = [] for content in msg.contents: - if content.type == "text": - seg_text += content.text or "" - elif content.type == "function_call": - seg_tool_calls.append( - { - "id": content.call_id, - "type": "function", - "function": { - "name": content.name, - "arguments": content.arguments, - }, - } - ) + 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 a pending assistant call (with any buffered text) so the matching - # call precedes its result; text-only content is deferred (see flush_segment). - flush_segment(before_result=True) + # 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(), @@ -1036,8 +1065,11 @@ def flush_segment(*, before_result: bool = False) -> None: "toolCallId": content.call_id, } ) + if content.call_id is not None: + unresolved_call_ids.discard(str(content.call_id)) - # Emit any deferred / trailing assistant text (e.g. a summary after the results). + # 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 @@ -1088,23 +1120,7 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An result.extend(_split_mixed_message_to_agui(msg, role)) continue - content_text = "" - tool_calls: list[dict[str, Any]] = [] - - for content in msg.contents: - if content.type == "text": - content_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, - }, - } - ) + 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 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 713908a4a1..525f838264 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 @@ -1079,6 +1079,76 @@ def test_agent_framework_to_agui_text_before_result_round_trips_without_dropping 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 From 01d4e3b4b630baa9393c4dfeea934f3b0d098d06 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Fri, 4 Sep 2026 07:28:35 +0530 Subject: [PATCH 4/4] Python: keep AG-UI tool results behind their calls across messages Carry the unresolved-call set across the whole conversion instead of rebuilding it per mixed message, so a call opened by an earlier message stays open and a later mixed message cannot place a new assistant segment between that call and its result. Queue a result whose own call is still buffered and release it only after that call is emitted, so a result never precedes its call. Add an exhaustive regression test over every valid call/result interleaving, in both single-message and split-message shapes, asserting that no result is dropped by _sanitize_tool_history and none is emitted ahead of its call. --- .../_message_adapters.py | 133 +++++++++++---- .../tests/ag_ui/test_message_adapters.py | 154 ++++++++++++++++++ 2 files changed, 252 insertions(+), 35 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 e3e3cc34ca..de3892b139 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 @@ -983,29 +983,38 @@ def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, A return text, tool_calls -def _split_mixed_message_to_agui(msg: Message, role: str) -> list[dict[str, Any]]: +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. Three ordering rules keep the transcript provider-valid: + 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. 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 + * 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 (``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. + 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 @@ -1017,8 +1026,9 @@ def _split_mixed_message_to_agui(msg: Message, role: str) -> list[dict[str, Any] 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() + 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: @@ -1036,6 +1046,7 @@ def flush_segment() -> None: 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} @@ -1044,33 +1055,60 @@ def flush_segment() -> None: 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) - seg_has_call = seg_has_call or content.type == "function_call" + 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": - # 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: + 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() - 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_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 (e.g. a summary after the - # results) and/or a new-call segment whose results arrive in a later message. + # 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 @@ -1086,6 +1124,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): @@ -1105,6 +1162,11 @@ 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 @@ -1117,7 +1179,7 @@ def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, An # 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)) + result.extend(_split_mixed_message_to_agui(msg, role, unresolved_call_ids)) continue content_text, tool_calls = _encode_agui_segment(msg.contents) @@ -1132,6 +1194,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 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 525f838264..a756009df6 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 @@ -5,6 +5,7 @@ import base64 import json import logging +from itertools import permutations from typing import Any import pytest @@ -1149,6 +1150,159 @@ def test_agent_framework_to_agui_interleaved_batch_round_trips_without_dropping( assert surviving_result_ids == {"call_a", "call_b", "call_c"} +def test_agent_framework_to_agui_pending_call_carries_across_messages(): + """A call opened by an earlier message stays open across the conversion. + + Regression: the unresolved-call set used to be rebuilt per mixed message, so a prior + ``assistant(call A)`` followed by ``[call C, result A, result C]`` emitted + ``assistant(A), assistant(C), tool(A), tool(C)``; the intervening ``assistant(C)`` + cleared pending call A and ``_sanitize_tool_history`` dropped A's real result. + """ + framework_messages = [ + Message( + role="assistant", + contents=[Content.from_function_call(call_id="call_a", name="fa", arguments="{}")], + message_id="m1", + ), + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="call_c", name="fc", arguments="{}"), + Content.from_function_result(call_id="call_a", result="ra"), + Content.from_function_result(call_id="call_c", result="rc"), + ], + message_id="m2", + ), + ] + + agui_messages = agent_framework_messages_to_agui(framework_messages) + + # assistant(A) -> tool(A) -> assistant(C) -> tool(C): the new call C is deferred until + # the earlier call A has received its result. + assert [m["role"] for m in agui_messages] == ["assistant", "tool", "assistant", "tool"] + assert [tc["id"] for tc in agui_messages[0]["tool_calls"]] == ["call_a"] + assert agui_messages[1]["toolCallId"] == "call_a" + assert [tc["id"] for tc in agui_messages[2]["tool_calls"]] == ["call_c"] + assert agui_messages[3]["toolCallId"] == "call_c" + + 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_c"} + + +def test_agent_framework_to_agui_result_for_buffered_call_waits_for_its_call(): + """A result whose own call is still buffered is held until that call is emitted. + + Regression: for ``[call A, call B, result A, call C, result C, result B]`` the split + correctly deferred ``assistant(C)`` but still emitted ``tool(C)`` immediately, leaving + C's result ahead of its call so ``_sanitize_tool_history`` dropped it. + """ + 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_c", result="rc"), + Content.from_function_result(call_id="call_b", result="rb"), + ], + message_id="out-of-order-1", + ) + + agui_messages = agent_framework_messages_to_agui([msg]) + + # Every tool message must come after an assistant message that declared its call. + declared: set[str] = set() + for agui_msg in agui_messages: + if agui_msg["role"] == "assistant": + declared.update(tc["id"] for tc in agui_msg.get("tool_calls") or []) + else: + assert agui_msg["toolCallId"] in declared, f"{agui_msg['toolCallId']} emitted before its call" + + 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"} + + +def _call_result_interleavings(call_ids: list[str]) -> list[tuple[tuple[str, str], ...]]: + """Every ordering of the given calls and their results where each call precedes its result.""" + events = [("call", call_id) for call_id in call_ids] + [("result", call_id) for call_id in call_ids] + orderings: list[tuple[tuple[str, str], ...]] = [] + for ordering in dict.fromkeys(permutations(events)): + emitted: set[str] = set() + for kind, call_id in ordering: + if kind == "call": + emitted.add(call_id) + elif call_id not in emitted: + break + else: + orderings.append(ordering) + return orderings + + +def _contents_for(ordering: tuple[tuple[str, str], ...]) -> list[Content]: + return [ + Content.from_function_call(call_id=call_id, name="f", arguments="{}") + if kind == "call" + else Content.from_function_result(call_id=call_id, result=f"r-{call_id}") + for kind, call_id in ordering + ] + + +def _assert_no_result_is_orphaned(agui_messages: list[dict[str, Any]], call_ids: list[str]) -> None: + """Every tool message follows an assistant message declaring its call, and no result is lost.""" + declared: set[str] = set() + for agui_msg in agui_messages: + if agui_msg["role"] == "assistant": + declared.update(tool_call["id"] for tool_call in agui_msg.get("tool_calls") or []) + elif agui_msg["role"] == "tool": + assert agui_msg["toolCallId"] in declared, f"{agui_msg['toolCallId']} emitted before its call" + + provider_messages, _ = normalize_agui_input_messages(agui_messages, sanitize_tool_history=True) + surviving = { + content.call_id + for message in provider_messages + for content in (message.contents or []) + if content.type == "function_result" + } + assert surviving == set(call_ids) + + +@pytest.mark.parametrize("call_ids", [["call_a", "call_b"], ["call_a", "call_b", "call_c"]]) +def test_agent_framework_to_agui_no_interleaving_drops_a_result(call_ids: list[str]): + """Exhaustive: no ordering of calls/results may lose a result or emit one ahead of its call. + + ``_sanitize_tool_history`` resets its pending-call set on every non-tool message, so any + assistant message emitted between a call and that call's result causes the result to be + dropped. This sweeps every valid interleaving -- in a single mixed message and split across + two messages -- because that whole class of ordering bug has recurred here repeatedly. + """ + for ordering in _call_result_interleavings(call_ids): + contents = _contents_for(ordering) + + single = [Message(role="assistant", contents=contents, message_id="m1")] + _assert_no_result_is_orphaned(agent_framework_messages_to_agui(single), call_ids) + + # Splitting the leading content into its own message exercises the cross-message + # pending-call state that a per-message set would lose. + split = [ + Message(role="assistant", contents=contents[:1], message_id="m1"), + Message(role="assistant", contents=contents[1:], message_id="m2"), + ] + _assert_no_result_is_orphaned(agent_framework_messages_to_agui(split), call_ids) + + # Additional tests for better coverage