From f1a5a5115cf4ba154d8b5939bf52f8e50232ca2d Mon Sep 17 00:00:00 2001
From: Henry Su <henrysu4707@gmail.com>
Date: Fri, 24 Jul 2026 23:22:50 -0500
Subject: [PATCH] fix(ag-ui): clean confirm_changes approval payloads in
 snapshot

---
 .../ag-ui/agent_framework_ag_ui/_agent_run.py | 25 ++++--
 .../ag_ui/test_confirm_changes_snapshot.py    | 83 +++++++++++++++++++
 2 files changed, 99 insertions(+), 9 deletions(-)
 create mode 100644 python/packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py

diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
index 84820bba408..3ccd70d0321 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
@@ -1434,9 +1434,6 @@ def _clean_resolved_approvals_from_snapshot(
                 )
                 result_by_call_id[str(content.call_id)] = result_text
 
-    if not result_by_call_id:
-        return
-
     for snap_msg in snapshot_messages:
         if normalize_agui_role(snap_msg.get("role", "")) != "tool":
             continue
@@ -1455,12 +1452,22 @@ def _clean_resolved_approvals_from_snapshot(
         # Find matching tool result by toolCallId
         tool_call_id = snap_msg.get("toolCallId") or snap_msg.get("tool_call_id") or ""
         replacement = result_by_call_id.get(str(tool_call_id))
-        if replacement is not None:
-            snap_msg["content"] = replacement
-            logger.info(
-                "Replaced approval payload in snapshot for tool_call_id=%s with actual result",
-                tool_call_id,
-            )
+        if replacement is None:
+            # For confirm_changes synthetic tool calls, tool_call_id is the confirm_id
+            # while result_by_call_id is keyed by the original tool call id.
+            if parsed.get("accepted"):
+                replacement = (
+                    "\n\n".join(result_by_call_id.values())
+                    if result_by_call_id
+                    else "Changes confirmed and applied successfully."
+                )
+            else:
+                replacement = "Changes declined."
+        snap_msg["content"] = replacement
+        logger.info(
+            "Replaced approval payload in snapshot for tool_call_id=%s with actual result",
+            tool_call_id,
+        )
 
 
 def _snapshot_tool_call_ids(message: Mapping[str, Any]) -> list[str]:
diff --git a/python/packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py b/python/packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py
new file mode 100644
index 00000000000..02b9835e8ba
--- /dev/null
+++ b/python/packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py
@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AG-UI confirm_changes snapshot payload cleaning and tool result preservation."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from agent_framework import Content, Message
+
+from agent_framework_ag_ui._agent_run import _clean_resolved_approvals_from_snapshot
+
+
+def test_clean_resolved_approvals_from_snapshot_confirm_changes_flow() -> None:
+    """Verify that _clean_resolved_approvals_from_snapshot replaces confirm_changes
+    approval payloads (whose toolCallId is confirm_id) with the executed tool results
+    from resolved_messages (whose call_id is the original tool call ID).
+    """
+    original_call_id = "call_orig_123"
+    confirm_call_id = "call_confirm_456"
+
+    # Snapshot message sent by client containing confirm_changes response payload
+    snapshot_messages: list[dict[str, Any]] = [
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": original_call_id,
+                    "type": "function",
+                    "function": {"name": "apply_changes", "arguments": "{}"},
+                },
+                {
+                    "id": confirm_call_id,
+                    "type": "function",
+                    "function": {"name": "confirm_changes", "arguments": "{}"},
+                },
+            ],
+        },
+        {
+            "role": "tool",
+            "toolCallId": confirm_call_id,
+            "content": json.dumps({"accepted": True, "steps": []}),
+        },
+    ]
+
+    # Executed tool results returned after approval resolution
+    resolved_messages: list[Message] = [
+        Message(
+            role="tool",
+            contents=[Content.from_function_result(call_id=original_call_id, result="Changes applied successfully.")],
+        )
+    ]
+
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, resolved_messages)
+
+    # The confirm_changes tool message payload must no longer contain the raw {"accepted": true} JSON
+    tool_msg = snapshot_messages[1]
+    assert tool_msg["toolCallId"] == confirm_call_id
+    assert tool_msg["content"] == "Changes applied successfully."
+    assert "accepted" not in tool_msg["content"]
+
+
+def test_clean_resolved_approvals_from_snapshot_confirm_changes_rejection() -> None:
+    """Verify that rejected confirm_changes responses clean out the approval payload."""
+    confirm_call_id = "call_confirm_789"
+
+    snapshot_messages: list[dict[str, Any]] = [
+        {
+            "role": "tool",
+            "toolCallId": confirm_call_id,
+            "content": json.dumps({"accepted": False}),
+        },
+    ]
+
+    resolved_messages: list[Message] = []
+
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, resolved_messages)
+
+    tool_msg = snapshot_messages[0]
+    assert tool_msg["content"] == "Changes declined."
+    assert "accepted" not in tool_msg["content"]