Python: Fix AG-UI tool history replay sanitization

python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py

@@ -26,21 +26,64 @@
 logger = logging.getLogger(__name__)
 
 
+def _append_synthetic_tool_results(
+    sanitized: list[Message],
+    pending_tool_call_ids: list[str],
+    result: str,
+) -> None:
+    for pending_call_id in pending_tool_call_ids:
+        logger.info("Injecting synthetic tool result for pending call_id=%s", pending_call_id)
+        sanitized.append(
+            Message(
+                role="tool",
+                contents=[
+                    Content.from_function_result(
+                        call_id=pending_call_id,
+                        result=result,
+                    )
+                ],
+            )
+        )
+
+
+def _ordered_unique_tool_call_ids(contents: list[Content]) -> list[str]:
+    tool_ids: list[str] = []
+    seen: set[str] = set()
+    for content in contents:
+        if content.type != "function_call" or not content.call_id:
+            continue
+        tool_id = str(content.call_id)
+        if tool_id in seen:
+            continue
+        tool_ids.append(tool_id)
+        seen.add(tool_id)
+    return tool_ids
+
+
 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
+    pending_tool_call_ids: list[str] | None = None
     pending_confirm_changes_id: str | None = None
 
     for msg in messages:
         role_value = get_role_value(msg)
 
         if role_value == "assistant":
-            tool_ids = {
-                str(content.call_id)
-                for content in msg.contents or []
-                if content.type == "function_call" and content.call_id
-            }
+            if pending_tool_call_ids:
+                logger.info(
+                    "Assistant message arrived with %d pending tool calls - injecting synthetic results",
+                    len(pending_tool_call_ids),
+                )
+                _append_synthetic_tool_results(
+                    sanitized,
+                    pending_tool_call_ids,
+                    "Tool execution skipped - assistant continued before the tool result was available.",
+                )
+                pending_tool_call_ids = None
+                pending_confirm_changes_id = None
+
+            tool_ids = _ordered_unique_tool_call_ids(msg.contents or [])
             confirm_changes_call = None
             for content in msg.contents or []:
                 if content.type == "function_call" and content.name == "confirm_changes":
@@ -67,7 +110,8 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
 
                 # Remove confirm_changes from tool_ids since we filtered it from the message
                 if confirm_changes_call.call_id:
-                    tool_ids.discard(str(confirm_changes_call.call_id))
+                    confirm_call_id = str(confirm_changes_call.call_id)
+                    tool_ids = [tool_id for tool_id in tool_ids if tool_id != confirm_call_id]
                 # Don't set pending_confirm_changes_id - we don't want a synthetic result
                 confirm_changes_call = None
             else:
@@ -92,7 +136,9 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
                         approval_accepted = approval_accepted and bool(content.approved)
 
             if approval_call_ids and pending_tool_call_ids:
-                pending_tool_call_ids -= approval_call_ids
+                pending_tool_call_ids = [
+                    call_id for call_id in pending_tool_call_ids if call_id not in approval_call_ids
+                ]
                 logger.info(
                     f"function_approval_response content found for call_ids={sorted(approval_call_ids)} - "
                     "framework will handle execution"
@@ -111,7 +157,9 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
                 )
                 sanitized.append(synthetic_result)
                 if pending_tool_call_ids:
-                    pending_tool_call_ids.discard(pending_confirm_changes_id)
+                    pending_tool_call_ids = [
+                        call_id for call_id in pending_tool_call_ids if call_id != pending_confirm_changes_id
+                    ]
                 pending_confirm_changes_id = None
 
             if pending_confirm_changes_id:
@@ -140,7 +188,9 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
                         )
                         sanitized.append(synthetic_result)
                         if pending_tool_call_ids:
-                            pending_tool_call_ids.discard(pending_confirm_changes_id)
+                            pending_tool_call_ids = [
+                                call_id for call_id in pending_tool_call_ids if call_id != pending_confirm_changes_id
+                            ]
                         pending_confirm_changes_id = None
                         continue
                 except (json.JSONDecodeError, KeyError) as exc:
@@ -151,18 +201,11 @@ 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)
+                _append_synthetic_tool_results(
+                    sanitized,
+                    pending_tool_call_ids,
+                    "Tool execution skipped - user provided follow-up message",
+                )
                 pending_tool_call_ids = None
                 pending_confirm_changes_id = None
 
@@ -182,18 +225,43 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
                         # Remove the call_id from pending since we now have its result.
                         # This prevents duplicate synthetic "skipped" results from being
                         # injected when a user message arrives later.
-                        pending_tool_call_ids.discard(call_id)
+                        pending_tool_call_ids = [
+                            pending_id for pending_id in pending_tool_call_ids if pending_id != call_id
+                        ]
                         if call_id == pending_confirm_changes_id:
                             pending_confirm_changes_id = None
                         break
             if keep:
                 sanitized.append(msg)
             continue
 
+        if pending_tool_call_ids:
+            logger.info(
+                "%s message arrived with %d pending tool calls - injecting synthetic results",
+                role_value,
+                len(pending_tool_call_ids),
+            )
+            _append_synthetic_tool_results(
+                sanitized,
+                pending_tool_call_ids,
+                "Tool execution skipped - conversation continued before the tool result was available.",
+            )
+
         sanitized.append(msg)
         pending_tool_call_ids = None
         pending_confirm_changes_id = None
 
+    if pending_tool_call_ids:
+        logger.info(
+            "History ended with %d pending tool calls - injecting synthetic results",
+            len(pending_tool_call_ids),
+        )
+        _append_synthetic_tool_results(
+            sanitized,
+            pending_tool_call_ids,
+            "Tool execution skipped - conversation ended before the tool result was available.",
+        )
+
     return sanitized
 
 

python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py

@@ -155,6 +155,80 @@ async def test_approval_resume_result_has_content() -> None:
     assert result_event.content == "Sunny in Portland"
 
 
+async def test_approval_resume_snapshot_replaces_approval_payload_with_tool_result() -> None:
+    """Approved HITL tools persist their executed result in MESSAGES_SNAPSHOT for replay."""
+    from agent_framework_ag_ui._message_adapters import normalize_agui_input_messages
+
+    call_id = "call_snapshot_replay"
+    weather_tool = _make_weather_tool()
+    agent = StubAgent(
+        updates=[AgentResponseUpdate(contents=[Content.from_text(text="The weather is sunny.")], role="assistant")],
+        default_options={"tools": [weather_tool]},
+    )
+    config = AgentConfig()
+    resume_messages: list[dict[str, Any]] = [
+        {"role": "user", "content": "What's the weather in Seattle?"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": call_id,
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": json.dumps({"city": "Seattle"}),
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": True}),
+            "toolCallId": call_id,
+        },
+    ]
+
+    events: list[Any] = []
+    async for event in run_agent_stream(
+        {
+            "thread_id": "thread-snapshot-replay",
+            "run_id": "run-snapshot-replay",
+            "messages": resume_messages,
+        },
+        agent,
+        config,
+    ):
+        events.append(event)
+
+    snapshots = [event.messages for event in events if getattr(event, "type", None) == "MESSAGES_SNAPSHOT"]
+    assert snapshots
+    snapshot_messages = [
+        message.model_dump(by_alias=True, exclude_none=True) if hasattr(message, "model_dump") else message
+        for message in snapshots[-1]
+    ]
+    tool_messages = [message for message in snapshot_messages if message.get("role") == "tool"]
+    assert any(
+        message.get("toolCallId") == call_id and message.get("content") == "Sunny in Seattle"
+        for message in tool_messages
+    )
+    assert not any(message.get("content") == json.dumps({"accepted": True}) for message in tool_messages)
+
+    replay_messages = snapshot_messages + [{"role": "user", "content": "What is the weather now?"}]
+    provider_messages, _ = normalize_agui_input_messages(replay_messages)
+
+    assert not any(
+        content.type == "function_approval_response"
+        for message in provider_messages
+        for content in message.contents or []
+    )
+    assert any(
+        content.type == "function_result" and content.call_id == call_id and content.result == "Sunny in Seattle"
+        for message in provider_messages
+        for content in message.contents or []
+    )
+
+
 async def test_no_approval_no_extra_tool_result() -> None:
     """When no approval response is present, no extra TOOL_CALL_RESULT events should be emitted."""
     agent = StubAgent(updates=[AgentResponseUpdate(contents=[Content.from_text(text="Hello.")], role="assistant")])

python/packages/ag-ui/tests/ag_ui/test_message_adapters.py

@@ -994,6 +994,50 @@ 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_closes_previous_call():
+    """Consecutive assistant tool-call messages are separated by a synthetic tool result."""
+    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="get_weather", arguments="{}")],
+    )
+    second_assistant = Message(
+        role="assistant",
+        contents=[Content.from_function_call(call_id="c2", name="get_time", arguments="{}")],
+    )
+    user_msg = Message(role="user", contents=[Content.from_text(text="Continue")])
+
+    result = _sanitize_tool_history([first_assistant, second_assistant, user_msg])
+
+    roles = [msg.role for msg in result]
+    assert roles == ["assistant", "tool", "assistant", "tool", "user"]
+    assert result[1].contents[0].type == "function_result"
+    assert result[1].contents[0].call_id == "c1"
+    assert result[3].contents[0].type == "function_result"
+    assert result[3].contents[0].call_id == "c2"
+
+
+def test_sanitize_history_ending_with_tool_call_closes_pending_call():
+    """History ending with assistant tool calls gets synthetic results before provider replay."""
+    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="get_weather", arguments="{}"),
+            Content.from_function_call(call_id="c2", name="get_time", arguments="{}"),
+            Content.from_function_call(call_id="c1", name="get_weather", arguments="{}"),
+        ],
+    )
+
+    result = _sanitize_tool_history([assistant_msg])
+
+    assert [msg.role for msg in result] == ["assistant", "tool", "tool"]
+    assert [result[1].contents[0].call_id, result[2].contents[0].call_id] == ["c1", "c2"]
+    assert all("skipped" in str(msg.contents[0].result).lower() for msg in result[1:])
+
+
 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
@@ -1015,7 +1059,7 @@ def test_sanitize_tool_result_clears_pending_confirm():
 
 
 def test_sanitize_non_standard_role_resets_state():
-    """System message between assistant+user resets pending tool state."""
+    """System message between assistant+user closes pending tool state."""
     from agent_framework_ag_ui._message_adapters import _sanitize_tool_history
 
     assistant_msg = Message(
@@ -1026,9 +1070,9 @@ def test_sanitize_non_standard_role_resets_state():
     user_msg = Message(role="user", contents=[Content.from_text(text="Continue")])
 
     result = _sanitize_tool_history([assistant_msg, system_msg, user_msg])
-    # System message should reset pending state, so no synthetic tool results
     tool_results = [m for m in result if m.role == "tool"]
-    assert len(tool_results) == 0
+    assert len(tool_results) == 1
+    assert tool_results[0].contents[0].call_id == "c1"
 
 
 def test_sanitize_json_confirm_changes_response():