From 2de5ddaa9b67a0d99dc6dbfcaa912bf9db03bd75 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Fri, 15 May 2026 00:02:58 +0800
Subject: [PATCH] Python: fix AG-UI tool history replay

---
 .../_message_adapters.py                      | 45 +++++++++++--------
 .../tests/ag_ui/test_message_adapters.py      | 40 +++++++++++++++++
 2 files changed, 67 insertions(+), 18 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 c4d2e9b2cd..35b8d66ba3 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
@@ -32,10 +32,34 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
     pending_tool_call_ids: set[str] | None = None
     pending_confirm_changes_id: str | None = None
 
+    def flush_pending_tool_calls(reason: str) -> None:
+        nonlocal pending_tool_call_ids, pending_confirm_changes_id
+        if not pending_tool_call_ids:
+            return
+
+        logger.info(f"{reason} with {len(pending_tool_call_ids)} pending tool calls - injecting synthetic results")
+        for pending_call_id in pending_tool_call_ids:
+            logger.info(f"Injecting synthetic tool result for pending call_id={pending_call_id}")
+            sanitized.append(
+                Message(
+                    role="tool",
+                    contents=[
+                        Content.from_function_result(
+                            call_id=pending_call_id,
+                            result="Tool execution skipped - history continued without matching tool result",
+                        )
+                    ],
+                )
+            )
+        pending_tool_call_ids = None
+        pending_confirm_changes_id = None
+
     for msg in messages:
         role_value = get_role_value(msg)
 
         if role_value == "assistant":
+            flush_pending_tool_calls("Assistant message arrived")
+
             tool_ids = {
                 str(content.call_id)
                 for content in msg.contents or []
@@ -147,24 +171,7 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
                     logger.debug(f"Could not parse user message as confirm_changes response: {type(exc).__name__}")
 
             if pending_tool_call_ids:
-                logger.info(
-                    f"User message arrived with {len(pending_tool_call_ids)} pending tool calls - "
-                    "injecting synthetic results"
-                )
-                for pending_call_id in pending_tool_call_ids:
-                    logger.info(f"Injecting synthetic tool result for pending call_id={pending_call_id}")
-                    synthetic_result = Message(
-                        role="tool",
-                        contents=[
-                            Content.from_function_result(
-                                call_id=pending_call_id,
-                                result="Tool execution skipped - user provided follow-up message",
-                            )
-                        ],
-                    )
-                    sanitized.append(synthetic_result)
-                pending_tool_call_ids = None
-                pending_confirm_changes_id = None
+                flush_pending_tool_calls("User message arrived")
 
             sanitized.append(msg)
             pending_confirm_changes_id = None
@@ -194,6 +201,8 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
         pending_tool_call_ids = None
         pending_confirm_changes_id = None
 
+    flush_pending_tool_calls("History ended")
+
     return sanitized
 
 
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 69f7b7bdb3..7409c725b5 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
@@ -994,6 +994,46 @@ def test_sanitize_pending_tool_skip_on_user_followup():
     assert "skipped" in str(tool_results[0].contents[0].result).lower()
 
 
+def test_sanitize_consecutive_assistant_tool_calls_flushes_previous_pending():
+    """A second assistant tool-call message must not orphan the first call."""
+    from agent_framework_ag_ui._message_adapters import _sanitize_tool_history
+
+    first_assistant = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="c1", name="first_tool", arguments="{}")],
+    )
+    second_assistant = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="c2", name="second_tool", arguments="{}")],
+    )
+    user_msg = Message(role="user", contents=[Content.from_text(text="continue")])
+
+    result = _sanitize_tool_history([first_assistant, second_assistant, user_msg])
+
+    assert [m.role for m in result] == ["assistant", "tool", "assistant", "tool", "user"]
+    tool_results = [m.contents[0] for m in result if m.role == "tool"]
+    assert [content.call_id for content in tool_results] == ["c1", "c2"]
+
+
+def test_sanitize_history_end_flushes_pending_tool_calls():
+    """History ending after assistant tool calls still needs matching tool results."""
+    from agent_framework_ag_ui._message_adapters import _sanitize_tool_history
+
+    assistant_msg = Message(
+        role="assistant",
+        contents=[
+            Content.from_function_call(call_id="c1", name="first_tool", arguments="{}"),
+            Content.from_function_call(call_id="c2", name="second_tool", arguments="{}"),
+        ],
+    )
+
+    result = _sanitize_tool_history([assistant_msg])
+
+    assert [m.role for m in result] == ["assistant", "tool", "tool"]
+    tool_results = [m.contents[0] for m in result if m.role == "tool"]
+    assert {content.call_id for content in tool_results} == {"c1", "c2"}
+
+
 def test_sanitize_tool_result_clears_pending_confirm():
     """Tool result for pending confirm_changes call_id clears pending state."""
     from agent_framework_ag_ui._message_adapters import _sanitize_tool_history