Python: Fix ag-ui regressions (#2114)

* Bump ag-ui package version. Update CHANGELOG

* Fix ag-ui bugs

* Revert port test change

* Cleanup

* Intro factory funcs for samples

* Revert package ver change
This commit is contained in:
Evan Mattson
2025-11-13 00:18:24 +00:00
committed by GitHub
parent 348ac764e6
commit 5537b1da79
21 changed files with 5174 additions and 3759 deletions
@@ -85,6 +85,7 @@ class AgentFrameworkEventBridge:
self.input_messages = input_messages or []
self.pending_tool_calls: list[dict[str, Any]] = [] # Track tool calls for assistant message
self.tool_results: list[dict[str, Any]] = [] # Track tool results
self.tool_calls_ended: set[str] = set() # Track which tool calls have had ToolCallEndEvent emitted
async def from_agent_run_update(self, update: AgentRunResponseUpdate) -> list[BaseEvent]:
"""
@@ -118,12 +119,14 @@ class AgentFrameworkEventBridge:
message_id=self.current_message_id,
role="assistant",
)
logger.debug(f"Emitting TextMessageStartEvent with message_id={self.current_message_id}")
events.append(start_event)
event = TextMessageContentEvent(
message_id=self.current_message_id,
delta=content.text,
)
logger.debug(f"Emitting TextMessageContentEvent with delta: {content.text}")
events.append(event)
elif isinstance(content, FunctionCallContent):
@@ -378,6 +381,7 @@ class AgentFrameworkEventBridge:
)
logger.info(f"Emitting ToolCallEndEvent for completed tool call '{content.call_id}'")
events.append(end_event)
self.tool_calls_ended.add(content.call_id) # Track that we emitted end event
# Log total StateDeltaEvent count for this tool call
if self.state_delta_count > 0:
@@ -617,6 +621,7 @@ class AgentFrameworkEventBridge:
f"Emitting ToolCallEndEvent for approval-required tool '{content.function_call.call_id}'"
)
events.append(end_event)
self.tool_calls_ended.add(content.function_call.call_id) # Track that we emitted end event
# Emit custom event for approval request
# Note: In AG-UI protocol, the frontend handles interrupts automatically
@@ -38,22 +38,69 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
"""
result: list[ChatMessage] = []
for msg in messages:
# Check for backend tool rendering results FIRST (may not have role field)
if "actionExecutionId" in msg or "actionName" in msg:
# Backend tool rendering - convert to FunctionResultContent
from agent_framework import FunctionResultContent
# Handle standard tool result messages early (role="tool") to preserve provider invariants
# This path maps AGUI tool messages to FunctionResultContent with the correct tool_call_id
role_str = msg.get("role", "user")
if role_str == "tool":
# Prefer explicit tool_call_id fields; fall back to backend fields only if necessary
tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
tool_call_id = msg.get("actionExecutionId", "")
# If no explicit tool_call_id, treat as backend tool rendering payloads where
# AGUI may send actionExecutionId/actionName. This must still map to the
# assistant's tool call id to satisfy provider requirements.
if not tool_call_id:
tool_call_id = msg.get("actionExecutionId") or ""
# Extract raw content text
result_content = msg.get("content")
if result_content is None:
result_content = msg.get("result", "")
# Distinguish approval payloads from actual tool results
is_approval = False
if isinstance(result_content, str) and result_content:
import json as _json
try:
parsed = _json.loads(result_content)
is_approval = isinstance(parsed, dict) and "accepted" in parsed
except Exception:
is_approval = False
if is_approval:
# Approval responses should be treated as user messages to trigger human-in-the-loop flow
chat_msg = ChatMessage(
role=Role.USER,
contents=[TextContent(text=str(result_content))],
additional_properties={"is_tool_result": True, "tool_call_id": str(tool_call_id or "")},
)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
continue
chat_msg = ChatMessage(
role=Role.TOOL,
contents=[FunctionResultContent(call_id=str(tool_call_id), result=result_content)],
)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
continue
# Backend tool rendering payloads without an explicit role
# Prefer standard tool mapping above; this block only covers legacy/minimal payloads
if "actionExecutionId" in msg or "actionName" in msg:
# Prefer toolCallId if present; otherwise fall back to actionExecutionId
tool_call_id = msg.get("toolCallId") or msg.get("tool_call_id") or msg.get("actionExecutionId", "")
result_content = msg.get("result", msg.get("content", ""))
chat_msg = ChatMessage(
role=Role.TOOL, # Tool results must be tool role
contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
role=Role.TOOL,
contents=[FunctionResultContent(call_id=str(tool_call_id), result=result_content)],
)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
continue
@@ -93,55 +140,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
result.append(chat_msg)
continue
role_str = msg.get("role", "user")
# Handle tool result messages (with role="tool")
if role_str == "tool":
# Check if this is a standard tool result (has tool_call_id or toolCallId)
tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
result_content = msg.get("content", "")
# Distinguish between backend tool results and approval responses
# Approval responses have {"accepted": ...} structure
is_approval = False
if result_content:
import json
try:
parsed_content = json.loads(result_content)
is_approval = "accepted" in parsed_content
except (json.JSONDecodeError, TypeError):
is_approval = False
# Backend tool results have non-empty content WITHOUT "accepted" field
if tool_call_id and result_content and not is_approval:
# Tool execution result - convert to FunctionResultContent with correct role
from agent_framework import FunctionResultContent
chat_msg = ChatMessage(
role=Role.TOOL,
contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
continue
else:
# Human-in-the-loop approval response - mark for special handling
content = msg.get("content", "")
chat_msg = ChatMessage(
role=Role.USER, # Approval responses are user messages
contents=[TextContent(text=content)],
additional_properties={"is_tool_result": True, "tool_call_id": msg.get("toolCallId", "")},
)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
continue
# No special handling required for assistant/plain messages here
role = _AGUI_TO_FRAMEWORK_ROLE.get(role_str, Role.USER)
@@ -16,7 +16,15 @@ from ag_ui.core import (
TextMessageEndEvent,
TextMessageStartEvent,
)
from agent_framework import AgentProtocol, AgentThread, ChatAgent, TextContent
from agent_framework import (
AgentProtocol,
AgentThread,
ChatAgent,
ChatMessage,
FunctionCallContent,
FunctionResultContent,
TextContent,
)
from ._utils import convert_agui_tools_to_agent_framework, generate_event_id
@@ -276,6 +284,98 @@ class DefaultOrchestrator(Orchestrator):
response_format = context.agent.chat_options.response_format
skip_text_content = response_format is not None
# Sanitizer: ensure tool results only follow assistant tool calls
# Also inject synthetic tool results for confirm_changes
def sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
sanitized: list[ChatMessage] = []
pending_tool_call_ids: set[str] | None = None
pending_confirm_changes_id: str | None = None
for msg in messages:
role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
if role_value == "assistant":
tool_ids = {
str(content.call_id)
for content in msg.contents or []
if isinstance(content, FunctionCallContent) and content.call_id
}
# Check for confirm_changes tool call
confirm_changes_call = None
for content in msg.contents or []:
if isinstance(content, FunctionCallContent) and content.name == "confirm_changes":
confirm_changes_call = content
break
sanitized.append(msg)
pending_tool_call_ids = tool_ids if tool_ids else None
pending_confirm_changes_id = (
str(confirm_changes_call.call_id)
if confirm_changes_call and confirm_changes_call.call_id
else None
)
continue
if role_value == "user" and pending_confirm_changes_id:
# Check if this is a confirm_changes response (JSON with "accepted" field)
user_text = ""
for content in msg.contents or []:
if isinstance(content, TextContent):
user_text = content.text
break
try:
parsed = json.loads(user_text)
if "accepted" in parsed:
# This is a confirm_changes response - inject synthetic tool result
logger.info(
f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}"
)
synthetic_result = ChatMessage(
role="tool",
contents=[
FunctionResultContent(
call_id=pending_confirm_changes_id,
result="Confirmed" if parsed.get("accepted") else "Rejected",
)
],
)
sanitized.append(synthetic_result)
if pending_tool_call_ids:
pending_tool_call_ids.discard(pending_confirm_changes_id)
pending_confirm_changes_id = None
# Don't add the user message to sanitized - it's been converted to tool result
continue
except (json.JSONDecodeError, KeyError) as e:
# Failed to parse user message as confirm_changes response; continue normal processing
logger.debug(f"Could not parse user message as confirm_changes response: {e}")
# Not a confirm_changes response, continue normal processing
sanitized.append(msg)
pending_tool_call_ids = None
pending_confirm_changes_id = None
continue
if role_value == "tool":
if not pending_tool_call_ids:
continue
keep = False
for content in msg.contents or []:
if isinstance(content, FunctionResultContent):
call_id = str(content.call_id)
if call_id in pending_tool_call_ids:
keep = True
break
if keep:
sanitized.append(msg)
continue
sanitized.append(msg)
pending_tool_call_ids = None
pending_confirm_changes_id = None
return sanitized
# Create event bridge
event_bridge = AgentFrameworkEventBridge(
run_id=context.run_id,
@@ -328,22 +428,151 @@ class DefaultOrchestrator(Orchestrator):
if current_state:
thread.metadata["current_state"] = current_state # type: ignore[attr-defined]
# Add incoming AG-UI messages to the thread history
if context.messages:
await thread.on_new_messages(context.messages)
# Use the full incoming message batch to preserve tool-call adjacency
if not context.messages:
raw_messages = context.messages or []
if not raw_messages:
logger.warning("No messages provided in AG-UI input")
yield event_bridge.create_run_finished_event()
return
logger.info(f"Received {len(raw_messages)} raw messages from client")
for i, msg in enumerate(raw_messages):
role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
msg_id = getattr(msg, "message_id", None)
logger.info(f" Raw message {i}: role={role}, id={msg_id}")
if hasattr(msg, "contents") and msg.contents:
for j, content in enumerate(msg.contents):
content_type = type(content).__name__
if isinstance(content, TextContent):
logger.debug(f" Content {j}: {content_type} - {content.text}")
elif isinstance(content, FunctionCallContent):
logger.debug(f" Content {j}: {content_type} - {content.name}({content.arguments})")
elif isinstance(content, FunctionResultContent):
logger.debug(
f" Content {j}: {content_type} - call_id={content.call_id}, result={content.result}"
)
else:
logger.debug(f" Content {j}: {content_type} - {content}")
# After getting sanitized_messages, deduplicate them
def deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
"""Remove duplicate messages while preserving order.
For tool results with the same call_id, prefer the one with actual data.
"""
seen_keys: dict[Any, int] = {} # key -> index in unique_messages (key can be various tuple types)
unique_messages: list[ChatMessage] = []
for idx, msg in enumerate(messages):
role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
# For tool messages, use call_id as unique key
if role_value == "tool" and msg.contents and isinstance(msg.contents[0], FunctionResultContent):
call_id = str(msg.contents[0].call_id)
key: Any = (role_value, call_id)
# Check if we already have this tool result
if key in seen_keys:
existing_idx = seen_keys[key]
existing_msg = unique_messages[existing_idx]
# Compare results - prefer non-empty over empty
existing_result = None
if existing_msg.contents and isinstance(existing_msg.contents[0], FunctionResultContent):
existing_result = existing_msg.contents[0].result
new_result = msg.contents[0].result
# Replace if existing is empty/None and new has data
if (not existing_result or existing_result == "") and new_result:
logger.info(
f"Replacing empty tool result at index {existing_idx} with data from index {idx}"
)
unique_messages[existing_idx] = msg
else:
logger.info(f"Skipping duplicate tool result at index {idx}: call_id={call_id}")
continue
seen_keys[key] = len(unique_messages)
unique_messages.append(msg)
elif (
role_value == "assistant"
and msg.contents
and any(isinstance(c, FunctionCallContent) for c in msg.contents)
):
# For assistant messages with tool_calls, use the tool call IDs
tool_call_ids = tuple(
sorted(str(c.call_id) for c in msg.contents if isinstance(c, FunctionCallContent) and c.call_id)
)
key = (role_value, tool_call_ids)
if key in seen_keys:
logger.info(f"Skipping duplicate assistant tool call at index {idx}")
continue
seen_keys[key] = len(unique_messages)
unique_messages.append(msg)
else:
# For other messages (system, user, assistant without tools), hash the content
content_str = str([str(c) for c in msg.contents]) if msg.contents else ""
key = (role_value, hash(content_str))
if key in seen_keys:
logger.info(f"Skipping duplicate message at index {idx}: role={role_value}")
continue
seen_keys[key] = len(unique_messages)
unique_messages.append(msg)
return unique_messages
# Then use it:
sanitized_messages = sanitize_tool_history(raw_messages)
provider_messages = deduplicate_messages(sanitized_messages)
if not provider_messages:
logger.info("No provider-eligible messages after filtering; finishing run without invoking agent.")
yield event_bridge.create_run_finished_event()
return
logger.info(f"Processing {len(provider_messages)} provider messages after sanitization/deduplication")
for i, msg in enumerate(provider_messages):
role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
logger.info(f" Message {i}: role={role}")
if hasattr(msg, "contents") and msg.contents:
for j, content in enumerate(msg.contents):
content_type = type(content).__name__
if isinstance(content, TextContent):
logger.info(f" Content {j}: {content_type} - {content.text}")
elif isinstance(content, FunctionCallContent):
logger.info(f" Content {j}: {content_type} - {content.name}({content.arguments})")
elif isinstance(content, FunctionResultContent):
logger.info(
f" Content {j}: {content_type} - call_id={content.call_id}, result={content.result}"
)
else:
logger.info(f" Content {j}: {content_type} - {content}")
# NOTE: For AG-UI, the client sends the full conversation history on each request.
# We should NOT add to thread.on_new_messages() as that would cause duplication.
# Instead, we pass messages directly to the agent via messages_to_run.
# Inject current state as system message context if we have state
messages_to_run: list[Any] = []
if current_state and context.config.state_schema:
state_json = json.dumps(current_state, indent=2)
from agent_framework import ChatMessage
conversation_has_tool_calls = False
logger.debug(f"Checking {len(provider_messages)} provider messages for tool calls")
for i, msg in enumerate(provider_messages):
logger.debug(
f" Message {i}: role={msg.role.value}, contents={len(msg.contents) if hasattr(msg, 'contents') and msg.contents else 0}"
)
for msg in provider_messages:
if msg.role.value == "assistant" and hasattr(msg, "contents") and msg.contents:
if any(isinstance(content, FunctionCallContent) for content in msg.contents):
conversation_has_tool_calls = True
break
if current_state and context.config.state_schema and not conversation_has_tool_calls:
state_json = json.dumps(current_state, indent=2)
state_context_msg = ChatMessage(
role="system",
contents=[
@@ -359,9 +588,9 @@ Never replace existing data - always append or merge."""
)
messages_to_run.append(state_context_msg)
# Preserve order from client to satisfy provider constraints (assistant tool_calls must
# immediately precede tool result messages). Using the full batch avoids reordering.
messages_to_run.extend(context.messages)
# Add all provider messages to messages_to_run
# AG-UI sends full conversation history on each request, so we pass it directly to the agent
messages_to_run.extend(provider_messages)
# Handle client tools for hybrid execution
# Client sends tool metadata, server merges with its own tools.
@@ -370,11 +599,23 @@ Never replace existing data - always append or merge."""
from agent_framework import BaseChatClient
client_tools = convert_agui_tools_to_agent_framework(context.input_data.get("tools"))
logger.info(f"[TOOLS] Client sent {len(client_tools) if client_tools else 0} tools")
if client_tools:
for tool in client_tools:
tool_name = getattr(tool, "name", "unknown")
declaration_only = getattr(tool, "declaration_only", None)
logger.info(f"[TOOLS] - Client tool: {tool_name}, declaration_only={declaration_only}")
# Extract server tools - use type narrowing when possible
server_tools: list[Any] = []
if isinstance(context.agent, ChatAgent):
server_tools = context.agent.chat_options.tools or []
tools_from_agent = context.agent.chat_options.tools
server_tools = list(tools_from_agent) if tools_from_agent else []
logger.info(f"[TOOLS] Agent has {len(server_tools)} configured tools")
for tool in server_tools:
tool_name = getattr(tool, "name", "unknown")
approval_mode = getattr(tool, "approval_mode", None)
logger.info(f"[TOOLS] - {tool_name}: approval_mode={approval_mode}")
else:
# AgentProtocol allows duck-typed implementations - fallback to attribute access
# This supports test mocks and custom agent implementations
@@ -412,15 +653,37 @@ Never replace existing data - always append or merge."""
except AttributeError:
pass
combined_tools: list[Any] = []
if server_tools:
combined_tools.extend(server_tools)
# For tools parameter: only pass if we have client tools to add
# If we pass tools=, it overrides the agent's configured tools and loses metadata like approval_mode
# So only pass tools when we need to add client tools on top of server tools
# IMPORTANT: Don't include client tools that duplicate server tools (same name)
tools_param = None
if client_tools:
combined_tools.extend(client_tools)
# Get server tool names
server_tool_names = {getattr(tool, "name", None) for tool in server_tools}
# Filter out client tools that duplicate server tools
unique_client_tools = [
tool for tool in client_tools if getattr(tool, "name", None) not in server_tool_names
]
if unique_client_tools:
combined_tools: list[Any] = []
if server_tools:
combined_tools.extend(server_tools)
combined_tools.extend(unique_client_tools)
tools_param = combined_tools
logger.info(
f"[TOOLS] Passing tools= parameter with {len(combined_tools)} tools ({len(server_tools)} server + {len(unique_client_tools)} unique client)"
)
else:
logger.info("[TOOLS] All client tools duplicate server tools - not passing tools= parameter")
else:
logger.info("[TOOLS] No client tools - not passing tools= parameter (using agent's configured tools)")
# Collect all updates to get the final structured output
all_updates: list[Any] = []
async for update in context.agent.run_stream(messages_to_run, thread=thread, tools=combined_tools or None):
async for update in context.agent.run_stream(messages_to_run, thread=thread, tools=tools_param):
all_updates.append(update)
events = await event_bridge.from_agent_run_update(update)
for event in events:
@@ -432,6 +695,27 @@ Never replace existing data - always append or merge."""
yield event_bridge.create_run_finished_event()
return
# Check if there are pending tool calls (declaration-only tools that weren't executed)
# These need ToolCallEndEvent to signal the client to execute them
# Only emit for tool calls that haven't already had ToolCallEndEvent emitted
# (approval-required tools already had their end event emitted)
if event_bridge.pending_tool_calls:
pending_without_end = [
tc for tc in event_bridge.pending_tool_calls if tc.get("id") not in event_bridge.tool_calls_ended
]
if pending_without_end:
logger.info(
f"Found {len(pending_without_end)} pending tool calls without end event - emitting ToolCallEndEvent"
)
for tool_call in pending_without_end:
tool_call_id = tool_call.get("id")
if tool_call_id:
from ag_ui.core import ToolCallEndEvent
end_event = ToolCallEndEvent(tool_call_id=tool_call_id)
logger.info(f"Emitting ToolCallEndEvent for declaration-only tool call '{tool_call_id}'")
yield end_event
# After streaming completes, check if agent has response_format and extract structured output
if all_updates and response_format:
from agent_framework import AgentRunResponse