Python: Correlate AG-UI confirm_changes snapshots by call id

docs/specs/004-python-function-calling-loop.md

@@ -474,6 +474,8 @@ that manually replay messages own the equivalent rule: do not resend an approval
 | AG-UI approval-time follow-up | The full grouped user-input pause remains in message history and emits no synthetic `TOOL_CALL_RESULT`. | `test_resolve_approval_responses_preserves_follow_up_user_input_group` |
 | AG-UI approval execution failure | A grouped executor failure becomes one deterministic terminal error result for the approved call. | `test_resolve_approval_responses_returns_failure_when_grouped_execution_raises` |
 | AG-UI no-approval path | Ordinary tool results do not gain an extra approval result event. | `test_no_approval_no_extra_tool_result` |
+| AG-UI `confirm_changes` snapshot | An accepted synthetic confirmation is replaced only when its original function call has a real result; rejection is cleaned explicitly, and missing accepted results remain inert. | `packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py` |
+| AG-UI malformed `confirm_changes` metadata | Non-list tool-call metadata and malformed argument JSON are ignored without guessing a target call. | `test_confirm_changes_target_ignores_non_list_tool_calls`, `test_confirm_changes_target_rejects_malformed_arguments_json` |
 | Compaction pair integrity | Function call/result groups remain atomic. | `packages/core/tests/core/test_compaction.py::test_group_annotations_keep_tool_call_and_tool_result_atomic`, `test_group_annotations_include_reasoning_in_tool_call_group` |
 
 ## Required coverage gaps
@@ -487,7 +489,6 @@ These scenarios are required but are not fully covered by merged tests on `main`
 | Service-owned `previous_response_id` continuation cannot execute a terminal approval again on a later turn. | #6851 |
 | A provider that ignores `tool_choice="none"` after an invocation limit cannot expose an unanswered call. | #7045 |
 | Declaration-only streaming preserves request metadata without duplicating arguments. | #6973 |
-| AG-UI `confirm_changes` cleanup correlates one result by original function call id when several results exist. | #6828 |
 
 Do not mark these rows covered by nearby tests; each needs a dedicated regression at the owning layer.
 
@@ -537,6 +538,7 @@ Before accepting an update, reviewers must confirm:
 - #7241 — approval-resolution result streaming
 - #7267 / #7271 and #7304 — replayed calls and reused ids
 - #7043 — provider-injected approval execution
+- #6828 — AG-UI `confirm_changes` snapshot correlation
 - #6851 — duplicate side effects after approval continuation
 - #7383 — bind approval responses to framework-issued requests after this foundation merges
 - #6963 / #7095 — opaque reasoning-signature replay

python/packages/ag-ui/AGENTS.md

@@ -33,6 +33,8 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard.
   resumed messages, while `TOOL_CALL_RESULT` events are emitted only for terminal `function_result` contents.
 - Approval responses for tools injected during `before_run` are deferred to the in-run approval middleware rather
   than executed or rejected by the transport before those tools exist.
+- `confirm_changes` snapshot cleanup resolves the synthetic confirmation back to its original `function_call_id`;
+  it must never concatenate unrelated tool results or record accepted changes without a matching real result.
 - SSE keepalive is endpoint-owned transport behavior configured through
   `add_agent_framework_fastapi_endpoint(keepalive_seconds=...)`. It emits SSE comments only; do not add `PING`,
   `HEARTBEAT`, or `KEEPALIVE` AG-UI events, and do not add runner-level keepalive settings.

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

@@ -1406,6 +1406,39 @@ def _convert_approval_results_to_tool_messages(messages: list[Message]) -> None:
     messages[:] = result
 
 
+def _confirm_changes_target_call_id(
+    snapshot_messages: list[dict[str, Any]],
+    confirm_call_id: str,
+    approval_payload: Mapping[str, Any],
+) -> str | None:
+    explicit_call_id = approval_payload.get("function_call_id")
+    if explicit_call_id:
+        return str(explicit_call_id)
+
+    for snapshot_message in snapshot_messages:
+        if normalize_agui_role(snapshot_message.get("role", "")) != "assistant":
+            continue
+        tool_calls = snapshot_message.get("tool_calls") or snapshot_message.get("toolCalls")
+        if not isinstance(tool_calls, list):
+            continue
+        for tool_call in tool_calls:
+            if not isinstance(tool_call, Mapping) or str(tool_call.get("id") or "") != confirm_call_id:
+                continue
+            function = tool_call.get("function")
+            if not isinstance(function, Mapping) or function.get("name") != "confirm_changes":
+                return None
+            arguments = function.get("arguments")
+            if isinstance(arguments, str):
+                try:
+                    arguments = json.loads(arguments)
+                except json.JSONDecodeError:
+                    return None
+            if isinstance(arguments, Mapping) and arguments.get("function_call_id"):
+                return str(arguments["function_call_id"])
+            return None
+    return None
+
+
 def _clean_resolved_approvals_from_snapshot(
     snapshot_messages: list[dict[str, Any]],
     resolved_messages: list[Message],
@@ -1439,9 +1472,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
@@ -1460,12 +1490,25 @@ 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:
+            target_call_id = _confirm_changes_target_call_id(
+                snapshot_messages,
+                str(tool_call_id),
+                parsed,
             )
+            if target_call_id is None:
+                continue
+            if parsed.get("accepted"):
+                replacement = result_by_call_id.get(target_call_id)
+                if replacement is None:
+                    continue
+            else:
+                replacement = "Changes declined."
+        snap_msg["content"] = replacement
+        logger.info(
+            "Replaced approval payload in snapshot for tool_call_id=%s with resolved content",
+            tool_call_id,
+        )
 
 
 def _snapshot_tool_call_ids(message: Mapping[str, Any]) -> list[str]:

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

@@ -0,0 +1,157 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+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,
+    _confirm_changes_target_call_id,
+)
+
+
+def _confirm_snapshot(
+    *,
+    original_call_id: str,
+    confirm_call_id: str,
+    accepted: bool,
+) -> list[dict[str, Any]]:
+    return [
+        {
+            "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": json.dumps({"function_call_id": original_call_id}),
+                    },
+                },
+            ],
+        },
+        {
+            "role": "tool",
+            "toolCallId": confirm_call_id,
+            "content": json.dumps({"accepted": accepted, "steps": []}),
+        },
+    ]
+
+
+def test_confirm_changes_target_ignores_non_list_tool_calls() -> None:
+    snapshot_messages: list[dict[str, Any]] = [
+        {
+            "role": "assistant",
+            "tool_calls": {"id": "call_confirm"},
+        }
+    ]
+
+    target_call_id = _confirm_changes_target_call_id(snapshot_messages, "call_confirm", {})
+
+    assert target_call_id is None
+
+
+def test_confirm_changes_target_rejects_malformed_arguments_json() -> None:
+    snapshot_messages: list[dict[str, Any]] = [
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "id": "call_confirm",
+                    "type": "function",
+                    "function": {
+                        "name": "confirm_changes",
+                        "arguments": "{not-json",
+                    },
+                }
+            ],
+        }
+    ]
+
+    target_call_id = _confirm_changes_target_call_id(snapshot_messages, "call_confirm", {})
+
+    assert target_call_id is None
+
+
+def test_confirm_changes_snapshot_uses_original_call_id_with_multiple_results() -> None:
+    snapshot_messages = _confirm_snapshot(
+        original_call_id="call_target",
+        confirm_call_id="call_confirm",
+        accepted=True,
+    )
+    resolved_messages = [
+        Message(
+            role="tool",
+            contents=[
+                Content.from_function_result(call_id="call_old", result="old result"),
+                Content.from_function_result(call_id="call_target", result="target result"),
+            ],
+        )
+    ]
+
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, resolved_messages)
+
+    assert snapshot_messages[1]["content"] == "target result"
+
+
+def test_confirm_changes_snapshot_accepts_explicit_payload_call_id() -> None:
+    snapshot_messages: list[dict[str, Any]] = [
+        {
+            "role": "tool",
+            "toolCallId": "call_confirm",
+            "content": json.dumps({"accepted": True, "function_call_id": "call_target"}),
+        }
+    ]
+    resolved_messages = [
+        Message(
+            role="tool",
+            contents=[Content.from_function_result(call_id="call_target", result="target result")],
+        )
+    ]
+
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, resolved_messages)
+
+    assert snapshot_messages[0]["content"] == "target result"
+
+
+def test_confirm_changes_snapshot_cleans_rejection_without_results() -> None:
+    snapshot_messages = _confirm_snapshot(
+        original_call_id="call_target",
+        confirm_call_id="call_confirm",
+        accepted=False,
+    )
+
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, [])
+
+    assert snapshot_messages[1]["content"] == "Changes declined."
+
+
+def test_confirm_changes_snapshot_keeps_accepted_payload_when_target_result_is_missing() -> None:
+    snapshot_messages = _confirm_snapshot(
+        original_call_id="call_target",
+        confirm_call_id="call_confirm",
+        accepted=True,
+    )
+    original_content = snapshot_messages[1]["content"]
+    resolved_messages = [
+        Message(
+            role="tool",
+            contents=[
+                Content.from_function_result(call_id="call_old_1", result="old result 1"),
+                Content.from_function_result(call_id="call_old_2", result="old result 2"),
+            ],
+        )
+    ]
+
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, resolved_messages)
+
+    assert snapshot_messages[1]["content"] == original_content