Skip to content

Commit 4e25917

Browse files
authored
Python: Fix AG-UI message handling and MCP tool double-call bug (#3635)
* AG-UI bug fixes * Fixes * Fixes * Revert human_in_the_loop_agent.py changes * Address copilot feedback * PR feedback addressed
1 parent a971d24 commit 4e25917

File tree

5 files changed

+737
-52
lines changed

5 files changed

+737
-52
lines changed

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

Lines changed: 48 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,32 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
4444
confirm_changes_call = content
4545
break
4646

47-
sanitized.append(msg)
47+
# Filter out confirm_changes from assistant messages before sending to LLM.
48+
# confirm_changes is a synthetic tool for the approval UI flow - the LLM shouldn't
49+
# see it because it may contain stale function_arguments that confuse the model
50+
# (e.g., showing 5 steps when only 2 were approved).
51+
# When we filter out confirm_changes, we also remove it from tool_ids and don't
52+
# set pending_confirm_changes_id, so no synthetic result is injected for it.
53+
# This is required because OpenAI validates that every tool result has a matching
54+
# tool call in the previous assistant message.
55+
if confirm_changes_call:
56+
filtered_contents = [
57+
c for c in (msg.contents or []) if not (c.type == "function_call" and c.name == "confirm_changes")
58+
]
59+
if filtered_contents:
60+
# Create a new message without confirm_changes to avoid mutating the input
61+
filtered_msg = ChatMessage(role=msg.role, contents=filtered_contents)
62+
sanitized.append(filtered_msg)
63+
# If no contents left after filtering, don't append anything
64+
65+
# Remove confirm_changes from tool_ids since we filtered it from the message
66+
if confirm_changes_call.call_id:
67+
tool_ids.discard(str(confirm_changes_call.call_id))
68+
# Don't set pending_confirm_changes_id - we don't want a synthetic result
69+
confirm_changes_call = None
70+
else:
71+
sanitized.append(msg)
72+
4873
pending_tool_call_ids = tool_ids if tool_ids else None
4974
pending_confirm_changes_id = (
5075
str(confirm_changes_call.call_id) if confirm_changes_call and confirm_changes_call.call_id else None
@@ -66,7 +91,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
6691
if approval_call_ids and pending_tool_call_ids:
6792
pending_tool_call_ids -= approval_call_ids
6893
logger.info(
69-
f"FunctionApprovalResponseContent found for call_ids={sorted(approval_call_ids)} - "
94+
f"function_approval_response content found for call_ids={sorted(approval_call_ids)} - "
7095
"framework will handle execution"
7196
)
7297

@@ -93,6 +118,8 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
93118
user_text = content.text # type: ignore[assignment]
94119
break
95120

121+
if not user_text:
122+
continue
96123
try:
97124
parsed = json.loads(user_text) # type: ignore[arg-type]
98125
if "accepted" in parsed:
@@ -149,6 +176,10 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
149176
call_id = str(content.call_id)
150177
if call_id in pending_tool_call_ids:
151178
keep = True
179+
# Remove the call_id from pending since we now have its result.
180+
# This prevents duplicate synthetic "skipped" results from being
181+
# injected when a user message arrives later.
182+
pending_tool_call_ids.discard(call_id)
152183
if call_id == pending_confirm_changes_id:
153184
pending_confirm_changes_id = None
154185
break
@@ -337,7 +368,7 @@ def _filter_modified_args(
337368
result: list[ChatMessage] = []
338369
for msg in messages:
339370
# Handle standard tool result messages early (role="tool") to preserve provider invariants
340-
# This path maps AG‑UI tool messages to FunctionResultContent with the correct tool_call_id
371+
# This path maps AG‑UI tool messages to function_result content with the correct tool_call_id
341372
role_str = normalize_agui_role(msg.get("role", "user"))
342373
if role_str == "tool":
343374
# Prefer explicit tool_call_id fields; fall back to backend fields only if necessary
@@ -370,7 +401,7 @@ def _filter_modified_args(
370401

371402
if is_approval:
372403
# Look for the matching function call in previous messages to create
373-
# a proper FunctionApprovalResponseContent. This enables the agent framework
404+
# proper function_approval_response content. This enables the agent framework
374405
# to execute the approved tool (fix for GitHub issue #3034).
375406
accepted = parsed.get("accepted", False) if parsed is not None else False
376407
approval_payload_text = result_content if isinstance(result_content, str) else json.dumps(parsed)
@@ -447,11 +478,17 @@ def _filter_modified_args(
447478
merged_args["steps"] = merged_steps
448479
state_args = merged_args
449480

450-
# Keep the original tool call and AG-UI snapshot in sync with approved args.
451-
updated_args = (
452-
json.dumps(merged_args) if isinstance(matching_func_call.arguments, str) else merged_args
481+
# Update the ChatMessage tool call with only enabled steps (for LLM context).
482+
# The LLM should only see the steps that were actually approved/executed.
483+
updated_args_for_llm = (
484+
json.dumps(filtered_args)
485+
if isinstance(matching_func_call.arguments, str)
486+
else filtered_args
453487
)
454-
matching_func_call.arguments = updated_args
488+
matching_func_call.arguments = updated_args_for_llm
489+
490+
# Update raw messages with all steps + status (for MESSAGES_SNAPSHOT display).
491+
# This allows the UI to show which steps were enabled/disabled.
455492
_update_tool_call_arguments(messages, str(approval_call_id), merged_args)
456493
# Create a new FunctionCallContent with the modified arguments
457494
func_call_for_approval = Content.from_function_call(
@@ -464,7 +501,7 @@ def _filter_modified_args(
464501
# No modified arguments - use the original function call
465502
func_call_for_approval = matching_func_call
466503

467-
# Create FunctionApprovalResponseContent for the agent framework
504+
# Create function_approval_response content for the agent framework
468505
approval_response = Content.from_function_approval_response(
469506
approved=accepted,
470507
id=str(approval_call_id),
@@ -488,7 +525,7 @@ def _filter_modified_args(
488525
result.append(chat_msg)
489526
continue
490527

491-
# Cast result_content to acceptable type for FunctionResultContent
528+
# Cast result_content to acceptable type for function_result content
492529
func_result: str | dict[str, Any] | list[Any]
493530
if isinstance(result_content, str):
494531
func_result = result_content
@@ -565,7 +602,7 @@ def _filter_modified_args(
565602

566603
# Check if this message contains function approvals
567604
if "function_approvals" in msg and msg["function_approvals"]:
568-
# Convert function approvals to FunctionApprovalResponseContent
605+
# Convert function approvals to function_approval_response content
569606
approval_contents: list[Any] = []
570607
for approval in msg["function_approvals"]:
571608
# Create FunctionCallContent with the modified arguments

0 commit comments

Comments
 (0)