Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 16 additions & 9 deletions python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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."
Comment on lines +1455 to +1465
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]:
Expand Down
83 changes: 83 additions & 0 deletions python/packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py
Original file line number Diff line number Diff line change
@@ -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": "{}"},
},
Comment on lines +34 to +38
],
},
{
"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"]
Loading