diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index c4d2e9b2cd3..e89f7046b71 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -26,21 +26,64 @@ logger = logging.getLogger(__name__) +def _append_synthetic_tool_results( + sanitized: list[Message], + pending_tool_call_ids: list[str], + result: str, +) -> None: + for pending_call_id in pending_tool_call_ids: + logger.info("Injecting synthetic tool result for pending call_id=%s", pending_call_id) + sanitized.append( + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id=pending_call_id, + result=result, + ) + ], + ) + ) + + +def _ordered_unique_tool_call_ids(contents: list[Content]) -> list[str]: + tool_ids: list[str] = [] + seen: set[str] = set() + for content in contents: + if content.type != "function_call" or not content.call_id: + continue + tool_id = str(content.call_id) + if tool_id in seen: + continue + tool_ids.append(tool_id) + seen.add(tool_id) + return tool_ids + + def _sanitize_tool_history(messages: list[Message]) -> list[Message]: """Normalize tool ordering and inject synthetic results for AG-UI edge cases.""" sanitized: list[Message] = [] - pending_tool_call_ids: set[str] | None = None + pending_tool_call_ids: list[str] | None = None pending_confirm_changes_id: str | None = None for msg in messages: role_value = get_role_value(msg) if role_value == "assistant": - tool_ids = { - str(content.call_id) - for content in msg.contents or [] - if content.type == "function_call" and content.call_id - } + if pending_tool_call_ids: + logger.info( + "Assistant message arrived with %d pending tool calls - injecting synthetic results", + len(pending_tool_call_ids), + ) + _append_synthetic_tool_results( + sanitized, + pending_tool_call_ids, + "Tool execution skipped - assistant continued before the tool result was available.", + ) + pending_tool_call_ids = None + pending_confirm_changes_id = None + + tool_ids = _ordered_unique_tool_call_ids(msg.contents or []) confirm_changes_call = None for content in msg.contents or []: if content.type == "function_call" and content.name == "confirm_changes": @@ -67,7 +110,8 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]: # Remove confirm_changes from tool_ids since we filtered it from the message if confirm_changes_call.call_id: - tool_ids.discard(str(confirm_changes_call.call_id)) + confirm_call_id = str(confirm_changes_call.call_id) + tool_ids = [tool_id for tool_id in tool_ids if tool_id != confirm_call_id] # Don't set pending_confirm_changes_id - we don't want a synthetic result confirm_changes_call = None else: @@ -92,7 +136,9 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]: approval_accepted = approval_accepted and bool(content.approved) if approval_call_ids and pending_tool_call_ids: - pending_tool_call_ids -= approval_call_ids + pending_tool_call_ids = [ + call_id for call_id in pending_tool_call_ids if call_id not in approval_call_ids + ] logger.info( f"function_approval_response content found for call_ids={sorted(approval_call_ids)} - " "framework will handle execution" @@ -111,7 +157,9 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]: ) sanitized.append(synthetic_result) if pending_tool_call_ids: - pending_tool_call_ids.discard(pending_confirm_changes_id) + pending_tool_call_ids = [ + call_id for call_id in pending_tool_call_ids if call_id != pending_confirm_changes_id + ] pending_confirm_changes_id = None if pending_confirm_changes_id: @@ -140,7 +188,9 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]: ) sanitized.append(synthetic_result) if pending_tool_call_ids: - pending_tool_call_ids.discard(pending_confirm_changes_id) + pending_tool_call_ids = [ + call_id for call_id in pending_tool_call_ids if call_id != pending_confirm_changes_id + ] pending_confirm_changes_id = None continue except (json.JSONDecodeError, KeyError) as exc: @@ -151,18 +201,11 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]: f"User message arrived with {len(pending_tool_call_ids)} pending tool calls - " "injecting synthetic results" ) - for pending_call_id in pending_tool_call_ids: - logger.info(f"Injecting synthetic tool result for pending call_id={pending_call_id}") - synthetic_result = Message( - role="tool", - contents=[ - Content.from_function_result( - call_id=pending_call_id, - result="Tool execution skipped - user provided follow-up message", - ) - ], - ) - sanitized.append(synthetic_result) + _append_synthetic_tool_results( + sanitized, + pending_tool_call_ids, + "Tool execution skipped - user provided follow-up message", + ) pending_tool_call_ids = None pending_confirm_changes_id = None @@ -182,7 +225,9 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]: # Remove the call_id from pending since we now have its result. # This prevents duplicate synthetic "skipped" results from being # injected when a user message arrives later. - pending_tool_call_ids.discard(call_id) + pending_tool_call_ids = [ + pending_id for pending_id in pending_tool_call_ids if pending_id != call_id + ] if call_id == pending_confirm_changes_id: pending_confirm_changes_id = None break @@ -190,10 +235,33 @@ def _sanitize_tool_history(messages: list[Message]) -> list[Message]: sanitized.append(msg) continue + if pending_tool_call_ids: + logger.info( + "%s message arrived with %d pending tool calls - injecting synthetic results", + role_value, + len(pending_tool_call_ids), + ) + _append_synthetic_tool_results( + sanitized, + pending_tool_call_ids, + "Tool execution skipped - conversation continued before the tool result was available.", + ) + sanitized.append(msg) pending_tool_call_ids = None pending_confirm_changes_id = None + if pending_tool_call_ids: + logger.info( + "History ended with %d pending tool calls - injecting synthetic results", + len(pending_tool_call_ids), + ) + _append_synthetic_tool_results( + sanitized, + pending_tool_call_ids, + "Tool execution skipped - conversation ended before the tool result was available.", + ) + return sanitized diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index b83eec10d24..6aa370368e9 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -155,6 +155,80 @@ async def test_approval_resume_result_has_content() -> None: assert result_event.content == "Sunny in Portland" +async def test_approval_resume_snapshot_replaces_approval_payload_with_tool_result() -> None: + """Approved HITL tools persist their executed result in MESSAGES_SNAPSHOT for replay.""" + from agent_framework_ag_ui._message_adapters import normalize_agui_input_messages + + call_id = "call_snapshot_replay" + weather_tool = _make_weather_tool() + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="The weather is sunny.")], role="assistant")], + default_options={"tools": [weather_tool]}, + ) + config = AgentConfig() + resume_messages: list[dict[str, Any]] = [ + {"role": "user", "content": "What's the weather in Seattle?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"city": "Seattle"}), + }, + } + ], + }, + { + "role": "tool", + "content": json.dumps({"accepted": True}), + "toolCallId": call_id, + }, + ] + + events: list[Any] = [] + async for event in run_agent_stream( + { + "thread_id": "thread-snapshot-replay", + "run_id": "run-snapshot-replay", + "messages": resume_messages, + }, + agent, + config, + ): + events.append(event) + + snapshots = [event.messages for event in events if getattr(event, "type", None) == "MESSAGES_SNAPSHOT"] + assert snapshots + snapshot_messages = [ + message.model_dump(by_alias=True, exclude_none=True) if hasattr(message, "model_dump") else message + for message in snapshots[-1] + ] + tool_messages = [message for message in snapshot_messages if message.get("role") == "tool"] + assert any( + message.get("toolCallId") == call_id and message.get("content") == "Sunny in Seattle" + for message in tool_messages + ) + assert not any(message.get("content") == json.dumps({"accepted": True}) for message in tool_messages) + + replay_messages = snapshot_messages + [{"role": "user", "content": "What is the weather now?"}] + provider_messages, _ = normalize_agui_input_messages(replay_messages) + + assert not any( + content.type == "function_approval_response" + for message in provider_messages + for content in message.contents or [] + ) + assert any( + content.type == "function_result" and content.call_id == call_id and content.result == "Sunny in Seattle" + for message in provider_messages + for content in message.contents or [] + ) + + async def test_no_approval_no_extra_tool_result() -> None: """When no approval response is present, no extra TOOL_CALL_RESULT events should be emitted.""" agent = StubAgent(updates=[AgentResponseUpdate(contents=[Content.from_text(text="Hello.")], role="assistant")]) diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index 69f7b7bdb35..19d2f3233dd 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py @@ -994,6 +994,50 @@ def test_sanitize_pending_tool_skip_on_user_followup(): assert "skipped" in str(tool_results[0].contents[0].result).lower() +def test_sanitize_consecutive_assistant_tool_calls_closes_previous_call(): + """Consecutive assistant tool-call messages are separated by a synthetic tool result.""" + from agent_framework_ag_ui._message_adapters import _sanitize_tool_history + + first_assistant = Message( + role="assistant", + contents=[Content.from_function_call(call_id="c1", name="get_weather", arguments="{}")], + ) + second_assistant = Message( + role="assistant", + contents=[Content.from_function_call(call_id="c2", name="get_time", arguments="{}")], + ) + user_msg = Message(role="user", contents=[Content.from_text(text="Continue")]) + + result = _sanitize_tool_history([first_assistant, second_assistant, user_msg]) + + roles = [msg.role for msg in result] + assert roles == ["assistant", "tool", "assistant", "tool", "user"] + assert result[1].contents[0].type == "function_result" + assert result[1].contents[0].call_id == "c1" + assert result[3].contents[0].type == "function_result" + assert result[3].contents[0].call_id == "c2" + + +def test_sanitize_history_ending_with_tool_call_closes_pending_call(): + """History ending with assistant tool calls gets synthetic results before provider replay.""" + from agent_framework_ag_ui._message_adapters import _sanitize_tool_history + + assistant_msg = Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c1", name="get_weather", arguments="{}"), + Content.from_function_call(call_id="c2", name="get_time", arguments="{}"), + Content.from_function_call(call_id="c1", name="get_weather", arguments="{}"), + ], + ) + + result = _sanitize_tool_history([assistant_msg]) + + assert [msg.role for msg in result] == ["assistant", "tool", "tool"] + assert [result[1].contents[0].call_id, result[2].contents[0].call_id] == ["c1", "c2"] + assert all("skipped" in str(msg.contents[0].result).lower() for msg in result[1:]) + + def test_sanitize_tool_result_clears_pending_confirm(): """Tool result for pending confirm_changes call_id clears pending state.""" from agent_framework_ag_ui._message_adapters import _sanitize_tool_history @@ -1015,7 +1059,7 @@ def test_sanitize_tool_result_clears_pending_confirm(): def test_sanitize_non_standard_role_resets_state(): - """System message between assistant+user resets pending tool state.""" + """System message between assistant+user closes pending tool state.""" from agent_framework_ag_ui._message_adapters import _sanitize_tool_history assistant_msg = Message( @@ -1026,9 +1070,9 @@ def test_sanitize_non_standard_role_resets_state(): user_msg = Message(role="user", contents=[Content.from_text(text="Continue")]) result = _sanitize_tool_history([assistant_msg, system_msg, user_msg]) - # System message should reset pending state, so no synthetic tool results tool_results = [m for m in result if m.role == "tool"] - assert len(tool_results) == 0 + assert len(tool_results) == 1 + assert tool_results[0].contents[0].call_id == "c1" def test_sanitize_json_confirm_changes_response():