From 7a88fb77e12a43f0588f90eddb8d974d4dc8d971 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Wed, 20 May 2026 23:12:37 +0800
Subject: [PATCH] Python: fix AG-UI pending tool history replay

---
 .../_message_adapters.py                      | 48 +++++++++++++------
 .../tests/ag_ui/test_message_adapters.py      | 37 ++++++++++++++
 2 files changed, 71 insertions(+), 14 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..ec39670f70 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,36 @@ 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_results(result: str) -> None:
+        nonlocal pending_confirm_changes_id, pending_tool_call_ids
+
+        if not pending_tool_call_ids:
+            pending_confirm_changes_id = None
+            return
+
+        for pending_call_id in sorted(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=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":
+            if pending_tool_call_ids:
+                logger.info(
+                    f"Assistant message arrived with {len(pending_tool_call_ids)} pending tool calls - "
+                    "injecting synthetic results"
+                )
+                flush_pending_tool_results("Tool execution skipped - assistant provided another message")
+
             tool_ids = {
                 str(content.call_id)
                 for content in msg.contents or []
@@ -151,20 +177,7 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
                     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_results("Tool execution skipped - user provided follow-up message")
 
             sanitized.append(msg)
             pending_confirm_changes_id = None
@@ -194,6 +207,13 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
         pending_tool_call_ids = None
         pending_confirm_changes_id = None
 
+    if pending_tool_call_ids:
+        logger.info(
+            f"Message history ended with {len(pending_tool_call_ids)} pending tool calls - "
+            "injecting synthetic results"
+        )
+        flush_pending_tool_results("Tool execution skipped - missing tool result")
+
     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..afc5ac6184 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,43 @@ def test_sanitize_pending_tool_skip_on_user_followup():
     assert "skipped" in str(tool_results[0].contents[0].result).lower()
 
 
+def test_sanitize_pending_tool_skip_before_next_assistant():
+    """A new assistant message first closes any previous pending tool calls."""
+    from agent_framework_ag_ui._message_adapters import _sanitize_tool_history
+
+    first_assistant = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="call_1", name="search", arguments="{}")],
+    )
+    second_assistant = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="call_2", name="write_file", arguments="{}")],
+    )
+    user_msg = Message(role="user", contents=[Content.from_text(text="continue")])
+
+    result = _sanitize_tool_history([first_assistant, second_assistant, user_msg])
+
+    assert [msg.role for msg in result] == ["assistant", "tool", "assistant", "tool", "user"]
+    assert result[1].contents[0].call_id == "call_1"
+    assert result[3].contents[0].call_id == "call_2"
+
+
+def test_sanitize_pending_tool_skip_at_end_of_history():
+    """History ending with an assistant tool call is closed before replay."""
+    from agent_framework_ag_ui._message_adapters import _sanitize_tool_history
+
+    assistant_msg = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="call_1", name="search", arguments="{}")],
+    )
+
+    result = _sanitize_tool_history([assistant_msg])
+
+    assert [msg.role for msg in result] == ["assistant", "tool"]
+    assert result[1].contents[0].call_id == "call_1"
+    assert "skipped" in str(result[1].contents[0].result).lower()
+
+
 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