From 455bcd5e5ed4d7e91cdfb87df97d5f192a9de839 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 3 Aug 2026 14:01:35 +0900 Subject: [PATCH 1/4] Python: Fix Ollama approval resume message handling --- .../agent_framework_ollama/_chat_client.py | 3 ++ .../ollama/tests/test_ollama_chat_client.py | 41 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/python/packages/ollama/agent_framework_ollama/_chat_client.py b/python/packages/ollama/agent_framework_ollama/_chat_client.py index f8c7a2232ec..d27f440e314 100644 --- a/python/packages/ollama/agent_framework_ollama/_chat_client.py +++ b/python/packages/ollama/agent_framework_ollama/_chat_client.py @@ -474,6 +474,9 @@ def _format_system_message(self, message: Message) -> list[OllamaMessage]: def _format_user_message(self, message: Message) -> list[OllamaMessage]: if not any(c.type in {"text", "data"} for c in message.contents) and not message.text: + if message.contents and all(content.type == "function_approval_response" for content in message.contents): + # AG-UI resolves approvals in-process; this control-only message has no Ollama representation. + return [] raise ChatClientInvalidRequestException( "Ollama connector currently only supports user messages with TextContent or DataContent." ) diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index 70f578cc37a..c0e78885769 100644 --- a/python/packages/ollama/tests/test_ollama_chat_client.py +++ b/python/packages/ollama/tests/test_ollama_chat_client.py @@ -263,6 +263,47 @@ async def test_cmc( assert result.text == "test" +@patch.object(AsyncClient, "chat", new_callable=AsyncMock) +async def test_cmc_ignores_control_only_approval_resume_message( + mock_chat: AsyncMock, + ollama_unit_test_env: dict[str, str], +) -> None: + """A resolved approval control is not sent to an Ollama model as a user message.""" + mock_chat.return_value = OllamaChatResponse( + message=OllamaMessage(content="done", role="assistant"), + model="test", + ) + function_call = Content.from_function_call( + call_id="call_123", + name="dangerous_action", + arguments={"target": "production"}, + ) + approval_response = Content.from_function_approval_response( + approved=True, + id="approval_123", + function_call=function_call, + ) + messages = [ + Message(role="user", contents=[Content.from_text(text="Continue the task.")]), + Message(role="assistant", contents=[function_call]), + Message(role="user", contents=[approval_response]), + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_123", result="The task is complete.")], + ), + ] + + ollama_client = OllamaChatClient() + result = await ollama_client.get_response(messages=messages) + + assert result.text == "done" + assert mock_chat.await_args is not None + outgoing_messages = mock_chat.await_args.kwargs["messages"] + assert [message["role"] for message in outgoing_messages] == ["user", "assistant", "tool"] + assert outgoing_messages[0]["content"] == "Continue the task." + assert outgoing_messages[2]["content"] == "The task is complete." + + @patch.object(AsyncClient, "chat", new_callable=AsyncMock) async def test_cmc_maps_done_reason_to_finish_reason( mock_chat: AsyncMock, From 3e1ec5771200beb99cdf0f1e3966c1204c9968ba Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 3 Aug 2026 14:25:02 +0900 Subject: [PATCH 2/4] Python: Reject empty Ollama approval resume payload --- .../agent_framework_ollama/_chat_client.py | 2 +- .../ollama/tests/test_ollama_chat_client.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/python/packages/ollama/agent_framework_ollama/_chat_client.py b/python/packages/ollama/agent_framework_ollama/_chat_client.py index d27f440e314..b2815b7a96f 100644 --- a/python/packages/ollama/agent_framework_ollama/_chat_client.py +++ b/python/packages/ollama/agent_framework_ollama/_chat_client.py @@ -439,7 +439,7 @@ def _prepare_options(self, messages: Sequence[Message], options: Mapping[str, An # messages if messages and "messages" not in run_options: run_options["messages"] = self._prepare_messages_for_ollama(messages) - if "messages" not in run_options: + if not run_options.get("messages"): raise ChatClientInvalidRequestException("Messages are required for chat completions") # model diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index c0e78885769..a4b09d500de 100644 --- a/python/packages/ollama/tests/test_ollama_chat_client.py +++ b/python/packages/ollama/tests/test_ollama_chat_client.py @@ -304,6 +304,33 @@ async def test_cmc_ignores_control_only_approval_resume_message( assert outgoing_messages[2]["content"] == "The task is complete." +@patch.object(AsyncClient, "chat", new_callable=AsyncMock) +async def test_cmc_rejects_only_control_approval_resume_message( + mock_chat: AsyncMock, + ollama_unit_test_env: dict[str, str], +) -> None: + """A transcript with no Ollama-representable messages is rejected before the SDK call.""" + mock_chat.return_value = OllamaChatResponse( + message=OllamaMessage(content="unexpected", role="assistant"), + model="test", + ) + function_call = Content.from_function_call( + call_id="call_123", + name="dangerous_action", + arguments={"target": "production"}, + ) + approval_response = Content.from_function_approval_response( + approved=True, + id="approval_123", + function_call=function_call, + ) + + with pytest.raises(ChatClientInvalidRequestException, match="Messages are required for chat completions"): + await OllamaChatClient().get_response(messages=[Message(role="user", contents=[approval_response])]) + + mock_chat.assert_not_awaited() + + @patch.object(AsyncClient, "chat", new_callable=AsyncMock) async def test_cmc_maps_done_reason_to_finish_reason( mock_chat: AsyncMock, From b2d4435b0a8d258cc79690b2c46379b4b980bdf5 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Tue, 4 Aug 2026 09:47:15 +0900 Subject: [PATCH 3/4] Python: Keep AG-UI approval controls out of provider input --- .../specs/004-python-function-calling-loop.md | 6 + .../ag-ui/agent_framework_ag_ui/_agent_run.py | 427 +++++++++-- .../agent_framework_ag_ui/_run_common.py | 2 + .../tests/ag_ui/test_approval_result_event.py | 685 +++++++++++++++++- .../ag-ui/tests/ag_ui/test_endpoint.py | 230 ++++++ python/packages/ag-ui/tests/ag_ui/test_run.py | 208 ++++++ .../agent_framework_ollama/_chat_client.py | 5 +- .../ollama/tests/test_ollama_chat_client.py | 68 -- 8 files changed, 1502 insertions(+), 129 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 74c3ad7c155..8a1eb6595fe 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -353,6 +353,11 @@ that manually replay messages own the equivalent rule: do not resend an approval - `function_approval_request` and `function_approval_response` are control-plane contents, not durable model transcript items. - A current hosted approval response must be sent once on the immediate resume request. +- AG-UI removes a local approval response from its request and snapshot replay when a terminal result proves it was + already consumed, including result-before-response client replay; hosted approval responses remain provider + protocol data and pass through. +- Hosted AG-UI approval interrupts expose an accept/reject decision only; argument edits are rejected because the + hosted provider executes the server-owned request rather than client-edited arguments. - A server-issued approval request must not be replayed inline during service-side continuation. - History providers may retain approval control contents in their backing store for audit, but base history replay filters them before later model calls. @@ -439,6 +444,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Hosted server boundary | Standing approval does not cross `server_label`. | `test_tool_approval_middleware_standing_rules_include_hosted_server_boundary` | | Argument-scoped rule | Exact arguments are required; empty arguments are not tool-wide. | `test_tool_approval_middleware_always_approve_tool_with_arguments_rule`, `test_tool_approval_middleware_empty_arguments_rule_is_not_tool_wide` | | Provider-injected approval tool | A tool added during `before_run` defers to in-run resolution, executes once, and emits one result. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes` | +| AG-UI provider boundary | Completed local approval controls from AG-UI request and snapshot replay are absent from raw chat-client input while deferred and hosted approvals keep their respective in-run/provider paths. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_canonical_resume_preserves_hosted_approval_for_provider`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence`, `packages/ag-ui/tests/ag_ui/test_run.py::test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending` | ### Errors, control flow, and limits diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 479dcd67902..c8505da6c28 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -9,7 +9,8 @@ import logging import uuid from collections import OrderedDict -from collections.abc import AsyncIterable, Awaitable, Mapping +from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, TypedDict, cast from ag_ui.core import ( @@ -41,6 +42,7 @@ _ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY, # type: ignore _collect_approval_responses, # type: ignore _get_tool_map, # type: ignore + _is_hosted_tool_approval, # type: ignore _replace_approval_contents_with_results, # type: ignore _TOOL_APPROVAL_STATE_KEY, # type: ignore _try_execute_function_call_groups, # type: ignore @@ -99,6 +101,192 @@ _COLLECTED_APPROVAL_RESPONSES_KEY = "collected_approval_responses" +@dataclass +class _LocalApprovalOccurrence: + """One local function-call occurrence tracked across AG-UI replay.""" + + call_id: str + name: str | None + arguments: str | None + approval_ids: set[str] = field(default_factory=set) + closed: bool = False + + +def _local_approval_response_content_ids_to_remove(messages: Sequence[Message]) -> set[int]: + """Find completed and duplicate local approval responses by call occurrence.""" + occurrences_by_call_id: dict[str, list[_LocalApprovalOccurrence]] = {} + occurrences_by_approval_id: dict[str, list[_LocalApprovalOccurrence]] = {} + response_occurrence_by_approval_id: dict[str, _LocalApprovalOccurrence] = {} + response_content_id_by_approval_id: dict[str, int] = {} + response_occurrences: dict[int, _LocalApprovalOccurrence] = {} + duplicate_response_content_ids: set[int] = set() + + def add_occurrence(function_call: Content, *, closed: bool = False) -> _LocalApprovalOccurrence | None: + if function_call.call_id is None: + return None + occurrence = _LocalApprovalOccurrence( + call_id=function_call.call_id, + name=function_call.name, + arguments=canonical_function_arguments(function_call), + closed=closed, + ) + occurrences_by_call_id.setdefault(function_call.call_id, []).append(occurrence) + return occurrence + + def matching_occurrence( + candidates: Sequence[_LocalApprovalOccurrence], + function_call: Content, + *, + require_open: bool = False, + require_unbound: bool = False, + newest_first: bool = True, + prefer_open: bool = False, + ) -> _LocalApprovalOccurrence | None: + ordered = reversed(candidates) if newest_first else iter(candidates) + eligible = [ + occurrence + for occurrence in ordered + if (not require_open or not occurrence.closed) and (not require_unbound or not occurrence.approval_ids) + ] + arguments = canonical_function_arguments(function_call) + if prefer_open: + exact_open = next( + ( + occurrence + for occurrence in eligible + if not occurrence.closed + and occurrence.name == function_call.name + and occurrence.arguments == arguments + ), + None, + ) + if exact_open is not None: + return exact_open + if open_occurrence := next((occurrence for occurrence in eligible if not occurrence.closed), None): + return open_occurrence + exact = next( + ( + occurrence + for occurrence in eligible + if occurrence.name == function_call.name and occurrence.arguments == arguments + ), + None, + ) + return exact or next(iter(eligible), None) + + for message in messages: + for content in message.contents: + if content.type == "function_call": + add_occurrence(content) + continue + + if content.type == "function_approval_request": + function_call = content.function_call + if function_call is None or function_call.call_id is None or content.id is None: + continue + candidates = occurrences_by_call_id.get(function_call.call_id, []) + occurrence = matching_occurrence( + candidates, + function_call, + require_open=True, + require_unbound=True, + newest_first=False, + ) or add_occurrence(function_call) + if occurrence is not None: + occurrence.approval_ids.add(content.id) + occurrences_by_approval_id.setdefault(content.id, []).append(occurrence) + continue + + if content.type == "function_approval_response": + if _is_hosted_tool_approval(content): + continue + function_call = content.function_call + if function_call is None or function_call.call_id is None: + continue + previous_occurrence = response_occurrence_by_approval_id.get(content.id or "") + request_occurrences = occurrences_by_approval_id.get(content.id or "", []) + request_occurrence = matching_occurrence(request_occurrences, function_call, prefer_open=True) + call_occurrence = matching_occurrence( + occurrences_by_call_id.get(function_call.call_id, []), + function_call, + prefer_open=True, + ) + occurrence = ( + call_occurrence + if call_occurrence is not None + and not call_occurrence.closed + and (request_occurrence is None or request_occurrence.closed) + else request_occurrence or call_occurrence + ) + if occurrence is None: + occurrence = add_occurrence(function_call) + if previous_occurrence is not None and occurrence is previous_occurrence: + previous_content_id = response_content_id_by_approval_id.get(content.id or "") + if previous_content_id is not None: + duplicate_response_content_ids.add(previous_content_id) + if content.id is not None and occurrence is not None: + occurrence.approval_ids.add(content.id) + response_occurrence_by_approval_id[content.id] = occurrence + response_content_id_by_approval_id[content.id] = id(content) + if occurrence is not None: + response_occurrences[id(content)] = occurrence + continue + + if content.call_id is None: + continue + is_terminal_result = content.type == "function_result" and not ( + isinstance(content.result, str) and "[APPROVAL_PENDING]" in content.result + ) + is_follow_up_request = content.user_input_request and content.type not in { + "function_approval_request", + "function_approval_response", + } + if not (is_terminal_result or is_follow_up_request): + continue + occurrence = next( + (candidate for candidate in occurrences_by_call_id.get(content.call_id, []) if not candidate.closed), + None, + ) + if occurrence is None: + occurrence = _LocalApprovalOccurrence( + call_id=content.call_id, + name=None, + arguments=None, + ) + occurrences_by_call_id.setdefault(content.call_id, []).append(occurrence) + occurrence.closed = True + + return duplicate_response_content_ids | { + content_id for content_id, occurrence in response_occurrences.items() if occurrence.closed + } + + +def _filter_local_approval_responses_for_provider(messages: Sequence[Message]) -> list[Message]: + """Remove completed local approval controls from AG-UI provider input. + + A matching terminal result proves the local response has already been consumed, + even when client replay placed the result before the synthesized response. Local + responses without a result remain available to the in-run approval middleware, + while hosted-service responses remain provider protocol data. + """ + response_content_ids_to_remove = _local_approval_response_content_ids_to_remove(messages) + + filtered_messages: list[Message] = [] + for message in messages: + filtered_contents = [ + content for content in message.contents if id(content) not in response_content_ids_to_remove + ] + if len(filtered_contents) == len(message.contents): + filtered_messages.append(message) + continue + if not filtered_contents: + continue + filtered_message = copy.copy(message) + filtered_message.contents = filtered_contents + filtered_messages.append(filtered_message) + return filtered_messages + + def _build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any]: """Build metadata dict with string values for Azure compatibility. @@ -449,9 +637,10 @@ class _PendingApproval(TypedDict): class _PendingApprovalWithSiblings(_PendingApproval, total=False): - """Pending approval details including hidden already-approved sibling calls.""" + """Pending approval details including sibling calls and trusted hosted metadata.""" already_approved_requests: list[dict[str, Any]] + server_label: str PendingApprovalEntry = _PendingApprovalWithSiblings | str @@ -470,6 +659,7 @@ def _make_pending_approval_entry( request_id: str | None = None, interrupt_id: str | None = None, already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, ) -> _PendingApprovalWithSiblings: entry: _PendingApprovalWithSiblings = { "name": name, @@ -479,6 +669,8 @@ def _make_pending_approval_entry( } if already_approved_requests: entry["already_approved_requests"] = already_approved_requests + if server_label: + entry["server_label"] = server_label return entry @@ -511,6 +703,19 @@ def _pending_approval_already_approved_requests(entry: PendingApprovalEntry) -> return list(entry.get("already_approved_requests", [])) +def _pending_approval_server_label(entry: PendingApprovalEntry) -> str | None: + if isinstance(entry, str): + return None + return entry.get("server_label") + + +def _function_call_server_label(function_call: Content | None) -> str | None: + if function_call is None: + return None + server_label = function_call.additional_properties.get("server_label") + return server_label if isinstance(server_label, str) and server_label else None + + def _stored_already_approved_requests_for_visible_approval( session: AgentSession, *approval_ids: str | None, @@ -643,6 +848,7 @@ def _register_server_generated_approval_response( canonical_function_arguments(response.function_call), request_id=str(response.id) if response.id else None, interrupt_id=str(response.function_call.call_id) if response.function_call.call_id else None, + server_label=_function_call_server_label(response.function_call), ) _register_pending_approval_entry( pending_approvals, @@ -869,6 +1075,7 @@ def _register_pending_approval( request_id: str, interrupt_id: str | None, already_approved_requests: list[dict[str, Any]] | None = None, + server_label: str | None = None, ) -> None: """Register one pending approval under each distinct thread identity.""" keys = list( @@ -888,6 +1095,7 @@ def _register_pending_approval( request_id=request_id, interrupt_id=interrupt_id, already_approved_requests=already_approved_requests, + server_label=server_label, ) for key in keys: registry[key] = entry @@ -1094,6 +1302,18 @@ def _canonical_approval_resume_messages( pending_arguments = _pending_approval_arguments(pending_entry) original_arguments = _parse_json_object(pending_arguments) or {} edited_arguments = {key: value for key, value in payload.items() if key not in {"accepted", "approved"}} + if edited_arguments and _pending_approval_server_label(pending_entry): + return ( + [], + handled_ids, + cancelled_ids, + RunErrorEvent( + message=( + f"Hosted approval resume for interruptId '{interrupt_id}' does not support edited arguments." + ), + code="APPROVAL_RESUME_INVALID_RESPONSE", + ), + ) if not set(edited_arguments).issubset(set(original_arguments)): return ( [], @@ -1153,6 +1373,7 @@ def _canonical_approval_resume_messages( canonical_function_arguments(function_call), request_id=str(response.id) if response.id else None, interrupt_id=str(function_call.call_id) if function_call.call_id else None, + server_label=_function_call_server_label(function_call), ) _register_pending_approval_entry( pending_approvals, @@ -1198,6 +1419,7 @@ async def _resolve_approval_responses( run_kwargs: dict[str, Any], pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] | None = None, thread_id: str = "", + validated_approved_responses: list[Content] | None = None, ) -> list[Content]: """Execute approved function calls and replace approval content with results. @@ -1214,6 +1436,9 @@ async def _resolve_approval_responses( When provided, every approval response is validated against this registry to prevent bypass, function name spoofing, and replay. thread_id: The conversation thread ID used to scope registry keys. + validated_approved_responses: Optional collector for validated local + approval responses, including controls removed because the matching + call occurrence already has a terminal result. Returns: List of approved function_result Content objects only (empty if no @@ -1221,81 +1446,160 @@ async def _resolve_approval_responses( but are *not* included in the return value because they should not be emitted as TOOL_CALL_RESULT events. """ - fcc_todo = _collect_approval_responses(messages) - if not fcc_todo: - return [] - - approved_responses = [resp for resp in fcc_todo.values() if resp.approved] - rejected_responses = [resp for resp in fcc_todo.values() if not resp.approved] - - # Validate every approval response (approved AND rejected) against the - # pending approvals registry. Invalid responses are stripped from messages - # entirely — not converted to rejection results, which would inject - # attacker-controlled content into the LLM conversation. - if pending_approvals is not None and (approved_responses or rejected_responses): - validated: list[Any] = [] - validated_rejected: list[Any] = [] - invalid_ids: set[str] = set() - for resp in approved_responses + rejected_responses: - resp_id = resp.id or "" - resp_name = resp.function_call.name if resp.function_call else None + approval_responses: list[Content] = [] + responses_by_id: dict[str, list[Content]] = {} + for message in messages: + for content in message.contents: + if content.type == "function_approval_response" and content.id is not None: + approval_responses.append(content) + responses_by_id.setdefault(content.id, []).append(content) + + valid_response_content_ids: set[int] | None = None + response_content_ids_to_strip: set[int] = set() + if pending_approvals is not None: + valid_response_content_ids = set() + + def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: PendingApprovalEntry) -> bool: + if isinstance(candidate, str) and isinstance(expected, str): + return candidate == expected + return candidate is expected + + pending_response_groups: dict[tuple[str, object], tuple[PendingApprovalEntry, list[Content]]] = {} + for response in approval_responses: + resp_id = response.id + if resp_id is None: + continue registry_key = _pending_approval_key(thread_id, resp_id) + id_entry = pending_approvals.get(registry_key) + function_call_id = response.function_call.call_id if response.function_call else None + call_registry_key = ( + _pending_approval_key(thread_id, function_call_id) if function_call_id is not None else None + ) + call_entry = pending_approvals.get(call_registry_key) if call_registry_key is not None else None + pending_entry = id_entry or call_entry + if pending_entry is None: + if not _is_hosted_tool_approval(response): + logger.warning( + "Rejected approval response id=%s: no matching pending approval request", + resp_id, + ) + response_content_ids_to_strip.add(id(response)) + continue + + group_key: tuple[str, object] + if isinstance(pending_entry, str): + group_key = ("legacy", call_registry_key if call_entry is not None else registry_key) + else: + group_key = ("entry", id(pending_entry)) + group = pending_response_groups.get(group_key) + if group is None: + pending_response_groups[group_key] = (pending_entry, [response]) + else: + group[1].append(response) - if registry_key not in pending_approvals: + for pending_entry, responses in pending_response_groups.values(): + pending_name = _pending_approval_name(pending_entry) + # The canonical AG-UI approval id may be the provider call id, which can + # be reused by a later call occurrence, while provider request ids may + # alias that same pending entry. Only the latest response across every + # trusted alias can answer the current entry; earlier responses are + # stale replay controls and must not authorize a malformed fresh one. + primary_response = responses[-1] + response_content_ids_to_strip.update(id(response) for response in responses[:-1]) + resp_id = primary_response.id + registry_key = _pending_approval_key(thread_id, resp_id) if resp_id is not None else None + id_entry = pending_approvals.get(registry_key) if registry_key is not None else None + if not matches_pending_entry(id_entry, pending_entry): logger.warning( "Rejected approval response id=%s: no matching pending approval request", resp_id, ) - invalid_ids.add(resp_id) + response_content_ids_to_strip.add(id(primary_response)) continue - - pending_entry = pending_approvals[registry_key] - pending_name = _pending_approval_name(pending_entry) - if resp_name != pending_name: + function_call_id = primary_response.function_call.call_id if primary_response.function_call else None + call_registry_key = ( + _pending_approval_key(thread_id, function_call_id) if function_call_id is not None else None + ) + call_entry = pending_approvals.get(call_registry_key) if call_registry_key is not None else None + if not isinstance(pending_entry, str) and not matches_pending_entry(call_entry, pending_entry): + logger.warning( + "Rejected approval response id=%s: function call id mismatch (response=%s)", + resp_id, + function_call_id, + ) + response_content_ids_to_strip.add(id(primary_response)) + continue + response_name = primary_response.function_call.name if primary_response.function_call else None + if response_name != pending_name: logger.warning( "Rejected approval response id=%s: function name mismatch (response=%s, pending=%s)", resp_id, - resp_name, + response_name, pending_name, ) - invalid_ids.add(resp_id) + response_content_ids_to_strip.add(id(primary_response)) continue - pending_arguments = _pending_approval_arguments(pending_entry) - response_arguments = canonical_function_arguments(resp.function_call) + response_arguments = canonical_function_arguments(primary_response.function_call) if not _approval_arguments_match_pending(pending_arguments, response_arguments): - logger.warning( - "Rejected approval response id=%s: function arguments mismatch", - resp_id, - ) - invalid_ids.add(resp_id) + logger.warning("Rejected approval response id=%s: function arguments mismatch", resp_id) + response_content_ids_to_strip.add(id(primary_response)) continue - # Valid — consume entry to prevent replay + server_label = _pending_approval_server_label(pending_entry) + if primary_response.function_call is not None: + if server_label: + primary_response.function_call.additional_properties["server_label"] = server_label + else: + primary_response.function_call.additional_properties.pop("server_label", None) + valid_response_content_ids.add(id(primary_response)) _consume_pending_approval_entry( pending_approvals, thread_id, pending_entry, resp_id, - resp.function_call.call_id if resp.function_call else None, + primary_response.function_call.call_id if primary_response.function_call else None, ) - if resp.approved: - validated.append(resp) - else: - validated_rejected.append(resp) - - # Strip invalid approval responses from messages and fcc_todo so - # _replace_approval_contents_with_results never sees them. - if invalid_ids: - for inv_id in invalid_ids: - fcc_todo.pop(inv_id, None) - for msg in messages: - msg.contents = [ - c for c in msg.contents if not (c.type == "function_approval_response" and c.id in invalid_ids) - ] + if validated_approved_responses is not None and primary_response.approved and not server_label: + validated_approved_responses.append(primary_response) + elif validated_approved_responses is not None: + validated_approved_responses.extend( + responses[-1] + for responses in responses_by_id.values() + if responses[-1].approved and not _is_hosted_tool_approval(responses[-1]) + ) + + if response_content_ids_to_strip: + filtered_messages: list[Message] = [] + for message in messages: + filtered_contents = [ + content for content in message.contents if id(content) not in response_content_ids_to_strip + ] + if len(filtered_contents) == len(message.contents): + filtered_messages.append(message) + continue + if not filtered_contents: + continue + message.contents = filtered_contents + filtered_messages.append(message) + messages[:] = filtered_messages + + # A replayed terminal result can precede the synthesized approval response. + # Remove that completed control before static execution, and collapse duplicate + # responses for the same approval occurrence to one pending response. + messages[:] = _filter_local_approval_responses_for_provider(messages) + + fcc_todo = _collect_approval_responses(messages) + if valid_response_content_ids is not None: + fcc_todo = { + response_id: response + for response_id, response in fcc_todo.items() + if id(response) in valid_response_content_ids + } + if not fcc_todo: + return [] - approved_responses = validated - rejected_responses = validated_rejected + approved_responses = [resp for resp in fcc_todo.values() if resp.approved] approved_function_result_groups: list[list[Content]] = [] @@ -1305,7 +1609,7 @@ async def _resolve_approval_responses( for approval in approved_responses: tool_name = approval.function_call.name if approval.function_call else None - if tool_name in tool_map: + if tool_name in tool_map and not _is_hosted_tool_approval(approval): static_approved.append(approval) # Execute only statically-available approved tool calls @@ -1967,8 +2271,15 @@ async def run_agent_stream( # This must happen before running the agent so it sees the tool results tools_for_execution = tools if tools is not None else server_tools messages.extend(_pop_collected_tool_approval_response_messages(session, pending_approvals, approval_thread_id)) + validated_approved_responses: list[Content] = [] resolved_approval_results = await _resolve_approval_responses( - messages, tools_for_execution, agent, run_kwargs, pending_approvals, approval_thread_id + messages, + tools_for_execution, + agent, + run_kwargs, + pending_approvals, + approval_thread_id, + validated_approved_responses, ) # Defense-in-depth: replace approval payloads in snapshot with actual tool results @@ -1978,7 +2289,10 @@ async def run_agent_stream( _merge_resolved_approval_results_into_snapshot(snapshot_messages, messages) # Feature #3: Emit StateSnapshotEvent for approved state-changing tools before agent runs - approved_state_updates = _extract_approved_state_updates(messages, predictive_handler) + approved_state_updates = _extract_approved_state_updates( + [Message(role="user", contents=validated_approved_responses)], + predictive_handler, + ) approved_state_snapshot_emitted = False if approved_state_updates: flow.current_state.update(approved_state_updates) @@ -2090,6 +2404,7 @@ async def run_agent_stream( str(content.id), str(canonical_interrupt_id) if canonical_interrupt_id else None, ), + server_label=_function_call_server_label(content.function_call), ) # Evict oldest entries if the registry exceeds a safe bound (LRU) _evict_oldest_approvals(pending_approvals, max_size=10_000) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index cddd65e334e..fc887ee9198 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -836,10 +836,12 @@ def _emit_approval_request( ) interrupt_id = func_call_id or content.id if interrupt_id: + response_schema = _approval_response_schema() if func_call.additional_properties.get("server_label") else None flow.interrupts.append( _approval_interrupt_for_function_call( interrupt_id=str(interrupt_id), function_call=func_call, + response_schema=response_schema, ) ) 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 361cc7ad5cf..3e60017112c 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 @@ -11,7 +11,7 @@ from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from agent_framework_ag_ui._agent import AgentConfig -from agent_framework_ag_ui._agent_run import run_agent_stream +from agent_framework_ag_ui._agent_run import PendingApprovalEntry, PendingApprovalKey, run_agent_stream def _make_weather_tool() -> FunctionTool: @@ -608,6 +608,689 @@ async def fail_grouped_execution(**kwargs: Any) -> tuple[list[list[Content]], bo ] == ["Error: Tool call invocation failed."] +async def test_resolve_approval_responses_keeps_fresh_occurrence_when_canonical_id_is_reused() -> None: + """A completed occurrence cannot consume a later approval that reuses its canonical call id.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import ( + _make_pending_approval_entry, + _pending_approval_key, + _resolve_approval_responses, + ) + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool( + name="guarded_write", + description="Write a value", + func=guarded_write, + approval_mode="always_require", + ) + call_id = "call_reused" + first_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) + second_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) + messages = [ + Message(role="assistant", contents=[first_call]), + Message(role="tool", contents=[Content.from_function_result(call_id=call_id, result="already wrote")]), + Message( + role="user", + contents=[Content.from_function_approval_response(approved=True, id=call_id, function_call=first_call)], + ), + Message(role="assistant", contents=[second_call]), + Message( + role="user", + contents=[Content.from_function_approval_response(approved=True, id=call_id, function_call=second_call)], + ), + ] + thread_id = "thread-reused" + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { + _pending_approval_key(thread_id, call_id): _make_pending_approval_entry( + "guarded_write", + '{"value":"same"}', + request_id=call_id, + interrupt_id=call_id, + ) + } + agent = StubAgent(updates=[], default_options={"tools": [tool]}) + + results = await _resolve_approval_responses( + messages, + [tool], + agent, + {}, + pending_approvals, + thread_id, + ) + + assert executions == ["same"] + assert [result.result for result in results] == ["wrote:same"] + assert pending_approvals == {} + assert not [ + content for message in messages for content in message.contents if content.type == "function_approval_response" + ] + + +async def test_resolve_approval_responses_uses_fresh_decision_when_canonical_id_is_reused() -> None: + """A historical approval does not conflict with a fresh rejection for a reused call id.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import ( + _make_pending_approval_entry, + _pending_approval_key, + _resolve_approval_responses, + ) + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool( + name="guarded_write", + description="Write a value", + func=guarded_write, + approval_mode="always_require", + ) + call_id = "call_reused_decision" + first_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) + second_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) + messages = [ + Message(role="assistant", contents=[first_call]), + Message(role="tool", contents=[Content.from_function_result(call_id=call_id, result="already wrote")]), + Message( + role="user", + contents=[Content.from_function_approval_response(approved=True, id=call_id, function_call=first_call)], + ), + Message(role="assistant", contents=[second_call]), + Message( + role="user", + contents=[Content.from_function_approval_response(approved=False, id=call_id, function_call=second_call)], + ), + ] + thread_id = "thread-reused-decision" + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { + _pending_approval_key(thread_id, call_id): _make_pending_approval_entry( + "guarded_write", + '{"value":"same"}', + request_id=call_id, + interrupt_id=call_id, + ) + } + agent = StubAgent(updates=[], default_options={"tools": [tool]}) + + results = await _resolve_approval_responses( + messages, + [tool], + agent, + {}, + pending_approvals, + thread_id, + ) + + assert executions == [] + assert results == [] + assert pending_approvals == {} + assert [ + content.result for message in messages for content in message.contents if content.type == "function_result" + ] == ["already wrote", "Error: Tool call invocation was rejected by user."] + assert not [ + content for message in messages for content in message.contents if content.type == "function_approval_response" + ] + + +async def test_resolve_approval_responses_does_not_fall_back_when_fresh_reused_response_is_invalid() -> None: + """An invalid fresh response cannot consume pending state through a valid historical response.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import ( + _make_pending_approval_entry, + _pending_approval_key, + _resolve_approval_responses, + ) + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool( + name="guarded_write", + description="Write a value", + func=guarded_write, + approval_mode="always_require", + ) + call_id = "call_reused_invalid" + first_call = Content.from_function_call(call_id=call_id, name="guarded_write", arguments={"value": "same"}) + edited_second_call = Content.from_function_call( + call_id=call_id, + name="guarded_write", + arguments={"value": "tampered"}, + ) + messages = [ + Message(role="assistant", contents=[first_call]), + Message(role="tool", contents=[Content.from_function_result(call_id=call_id, result="already wrote")]), + Message( + role="user", + contents=[Content.from_function_approval_response(approved=True, id=call_id, function_call=first_call)], + ), + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id=call_id, + name="guarded_write", + arguments={"value": "same"}, + ) + ], + ), + Message( + role="user", + contents=[ + Content.from_function_approval_response( + approved=True, + id=call_id, + function_call=edited_second_call, + ) + ], + ), + ] + thread_id = "thread-reused-invalid" + pending_entry = _make_pending_approval_entry( + "guarded_write", + '{"value":"same"}', + request_id=call_id, + interrupt_id=call_id, + ) + pending_key = _pending_approval_key(thread_id, call_id) + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = {pending_key: pending_entry} + agent = StubAgent(updates=[], default_options={"tools": [tool]}) + + results = await _resolve_approval_responses( + messages, + [tool], + agent, + {}, + pending_approvals, + thread_id, + ) + + assert executions == [] + assert results == [] + assert pending_approvals == {pending_key: pending_entry} + assert not [ + content for message in messages for content in message.contents if content.type == "function_approval_response" + ] + + +async def test_resolve_approval_responses_consumes_trusted_hosted_pending_entry() -> None: + """A server-collected hosted response remains provider-bound but cannot be replayed.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import ( + _make_pending_approval_entry, + _pending_approval_key, + _resolve_approval_responses, + ) + + call_id = "mcpr_hosted" + server_label = "hosted-mcp" + hosted_call = Content.from_function_call( + call_id=call_id, + name="hosted_write", + arguments={"value": "same"}, + additional_properties={"server_label": server_label}, + ) + hosted_response = Content.from_function_approval_response( + approved=True, + id=call_id, + function_call=hosted_call, + ) + messages = [Message(role="user", contents=[hosted_response])] + thread_id = "thread-hosted-collected" + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { + _pending_approval_key(thread_id, call_id): _make_pending_approval_entry( + "hosted_write", + '{"value":"same"}', + request_id=call_id, + interrupt_id=call_id, + server_label=server_label, + ) + } + + results = await _resolve_approval_responses( + messages, + [], + StubAgent(updates=[]), + {}, + pending_approvals, + thread_id, + ) + + assert results == [] + assert pending_approvals == {} + assert len(messages) == 1 + assert messages[0].role == "user" + assert messages[0].contents == [hosted_response] + + +async def test_resolve_approval_responses_uses_fresh_response_across_pending_aliases() -> None: + """The interrupt-id response wins over historical replay under the request-id alias.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import ( + _make_pending_approval_entry, + _pending_approval_key, + _resolve_approval_responses, + ) + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) + request_id = "approval_alias" + interrupt_id = "call_alias" + function_call = Content.from_function_call( + call_id=interrupt_id, + name="guarded_write", + arguments={"value": "same"}, + ) + messages = [ + Message(role="assistant", contents=[function_call]), + Message( + role="user", + contents=[ + Content.from_function_approval_response( + approved=True, + id=request_id, + function_call=function_call, + ) + ], + ), + Message( + role="user", + contents=[ + Content.from_function_approval_response( + approved=False, + id=interrupt_id, + function_call=function_call, + ) + ], + ), + ] + thread_id = "thread-alias" + pending_entry = _make_pending_approval_entry( + "guarded_write", + '{"value":"same"}', + request_id=request_id, + interrupt_id=interrupt_id, + ) + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { + _pending_approval_key(thread_id, request_id): pending_entry, + _pending_approval_key(thread_id, interrupt_id): pending_entry, + } + + results = await _resolve_approval_responses( + messages, + [tool], + StubAgent(updates=[], default_options={"tools": [tool]}), + {}, + pending_approvals, + thread_id, + ) + + assert executions == [] + assert results == [] + assert pending_approvals == {} + assert [ + content.result for message in messages for content in message.contents if content.type == "function_result" + ] == ["Error: Tool call invocation was rejected by user."] + + +async def test_resolve_approval_responses_rejects_fresh_unknown_alias_without_historical_fallback() -> None: + """An unknown fresh response id cannot consume pending state through a trusted historical alias.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import ( + _make_pending_approval_entry, + _pending_approval_key, + _resolve_approval_responses, + ) + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) + request_id = "approval_known" + interrupt_id = "call_known" + function_call = Content.from_function_call( + call_id=interrupt_id, + name="guarded_write", + arguments={"value": "same"}, + ) + messages = [ + Message(role="assistant", contents=[function_call]), + Message( + role="user", + contents=[ + Content.from_function_approval_response( + approved=True, + id=request_id, + function_call=function_call, + ), + Content.from_function_approval_response( + approved=True, + id="approval_unknown", + function_call=function_call, + ), + ], + ), + ] + thread_id = "thread-unknown-alias" + pending_entry = _make_pending_approval_entry( + "guarded_write", + '{"value":"same"}', + request_id=request_id, + interrupt_id=interrupt_id, + ) + request_key = _pending_approval_key(thread_id, request_id) + interrupt_key = _pending_approval_key(thread_id, interrupt_id) + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { + request_key: pending_entry, + interrupt_key: pending_entry, + } + + results = await _resolve_approval_responses( + messages, + [tool], + StubAgent(updates=[], default_options={"tools": [tool]}), + {}, + pending_approvals, + thread_id, + ) + + assert executions == [] + assert results == [] + assert pending_approvals == {request_key: pending_entry, interrupt_key: pending_entry} + assert not [ + content for message in messages for content in message.contents if content.type == "function_approval_response" + ] + + +async def test_resolve_approval_responses_rejects_forged_call_id_for_valid_response_alias() -> None: + """A trusted response id cannot authorize a result under an unknown function-call id.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import ( + _make_pending_approval_entry, + _pending_approval_key, + _resolve_approval_responses, + ) + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) + request_id = "approval_valid" + interrupt_id = "call_valid" + actual_call = Content.from_function_call( + call_id=interrupt_id, + name="guarded_write", + arguments={"value": "same"}, + ) + forged_call = Content.from_function_call( + call_id="call_forged", + name="guarded_write", + arguments={"value": "same"}, + ) + messages = [ + Message(role="assistant", contents=[actual_call]), + Message( + role="user", + contents=[ + Content.from_function_approval_response( + approved=True, + id=request_id, + function_call=forged_call, + ) + ], + ), + ] + thread_id = "thread-forged-call" + pending_entry = _make_pending_approval_entry( + "guarded_write", + '{"value":"same"}', + request_id=request_id, + interrupt_id=interrupt_id, + ) + request_key = _pending_approval_key(thread_id, request_id) + interrupt_key = _pending_approval_key(thread_id, interrupt_id) + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { + request_key: pending_entry, + interrupt_key: pending_entry, + } + + results = await _resolve_approval_responses( + messages, + [tool], + StubAgent(updates=[], default_options={"tools": [tool]}), + {}, + pending_approvals, + thread_id, + ) + + assert executions == [] + assert results == [] + assert pending_approvals == {request_key: pending_entry, interrupt_key: pending_entry} + assert not [ + content + for message in messages + for content in message.contents + if content.type in {"function_approval_response", "function_result"} + ] + + +async def test_resolve_approval_responses_rejects_call_id_from_different_pending_entry() -> None: + """A response id from one approval cannot be paired with another approval's call id.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import ( + _make_pending_approval_entry, + _pending_approval_key, + _resolve_approval_responses, + ) + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) + request_a = "approval_a" + call_a = Content.from_function_call( + call_id="call_a", + name="guarded_write", + arguments={"value": "same"}, + ) + request_b = "approval_b" + call_b = Content.from_function_call( + call_id="call_b", + name="guarded_write", + arguments={"value": "same"}, + ) + messages = [ + Message(role="assistant", contents=[call_a]), + Message( + role="user", + contents=[ + Content.from_function_approval_response(approved=True, id=request_a, function_call=call_a), + Content.from_function_approval_response(approved=False, id=request_a, function_call=call_b), + ], + ), + ] + thread_id = "thread-crossed-aliases" + pending_a = _make_pending_approval_entry( + "guarded_write", + '{"value":"same"}', + request_id=request_a, + interrupt_id="call_a", + ) + pending_b = _make_pending_approval_entry( + "guarded_write", + '{"value":"same"}', + request_id=request_b, + interrupt_id="call_b", + ) + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { + _pending_approval_key(thread_id, request_a): pending_a, + _pending_approval_key(thread_id, "call_a"): pending_a, + _pending_approval_key(thread_id, request_b): pending_b, + _pending_approval_key(thread_id, "call_b"): pending_b, + } + expected_pending = dict(pending_approvals) + + results = await _resolve_approval_responses( + messages, + [tool], + StubAgent(updates=[], default_options={"tools": [tool]}), + {}, + pending_approvals, + thread_id, + ) + + assert executions == [] + assert results == [] + assert pending_approvals == expected_pending + assert not [ + content + for message in messages + for content in message.contents + if content.type in {"function_approval_response", "function_result"} + ] + + +async def test_resolve_approval_responses_without_registry_uses_latest_duplicate_decision() -> None: + """The optional no-registry path preserves the established last-response-wins behavior.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import _resolve_approval_responses + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) + call_id = "call_no_registry" + function_call = Content.from_function_call( + call_id=call_id, + name="guarded_write", + arguments={"value": "same"}, + ) + messages = [ + Message(role="assistant", contents=[function_call]), + Message( + role="user", + contents=[ + Content.from_function_approval_response(approved=True, id=call_id, function_call=function_call), + Content.from_function_approval_response(approved=False, id=call_id, function_call=function_call), + ], + ), + ] + + results = await _resolve_approval_responses( + messages, + [tool], + StubAgent(updates=[], default_options={"tools": [tool]}), + {}, + ) + + assert executions == [] + assert results == [] + assert [ + content.result for message in messages for content in message.contents if content.type == "function_result" + ] == ["Error: Tool call invocation was rejected by user."] + + +async def test_resolve_approval_responses_legacy_registry_uses_latest_duplicate_decision() -> None: + """Legacy string entries group duplicate decisions by their matched response id.""" + from agent_framework import Message + + from agent_framework_ag_ui._agent_run import _pending_approval_key, _resolve_approval_responses + + executions: list[str] = [] + + def guarded_write(value: str) -> str: + executions.append(value) + return f"wrote:{value}" + + tool = FunctionTool(name="guarded_write", description="Write a value", func=guarded_write) + approval_id = "approval_legacy" + first_call = Content.from_function_call( + call_id="call_legacy_old", + name="guarded_write", + arguments={"value": "same"}, + ) + latest_call = Content.from_function_call( + call_id="call_legacy_new", + name="guarded_write", + arguments={"value": "same"}, + ) + messages = [ + Message( + role="user", + contents=[ + Content.from_function_approval_response( + approved=True, + id=approval_id, + function_call=first_call, + ), + Content.from_function_approval_response( + approved=False, + id=approval_id, + function_call=latest_call, + ), + ], + ) + ] + thread_id = "thread-legacy" + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { + _pending_approval_key(thread_id, approval_id): "guarded_write" + } + + results = await _resolve_approval_responses( + messages, + [tool], + StubAgent(updates=[], default_options={"tools": [tool]}), + {}, + pending_approvals, + thread_id, + ) + + assert executions == [] + assert results == [] + assert pending_approvals == {} + assert [ + content.result for message in messages for content in message.contents if content.type == "function_result" + ] == ["Error: Tool call invocation was rejected by user."] + + class TestApprovalToolResultDisplayChannel: """Approved tools using ``state_update(..., tool_result=...)`` must route the display payload to the UI event while ``flow.tool_results`` still receives diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 0408de841d8..17b626fe168 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -5247,6 +5247,7 @@ async def test_endpoint_agent_approval_deferred_provider_tool_executes(streaming The deferred tool result must still be returned to AG-UI exactly once. """ side_effects: list[str] = [] + provider_messages: list[Message] = [] state = {"phase": "pause"} def provider_write() -> str: @@ -5279,6 +5280,7 @@ async def stream_fn( role="assistant", ) return + provider_messages[:] = list(messages) yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") # provider_write is intentionally NOT in the static tools list -- it is only injected via before_run. @@ -5331,7 +5333,235 @@ async def stream_fn( assert side_effects == ["wrote"] tool_results = [event for event in resume_events if event.get("type") == "TOOL_CALL_RESULT"] assert [(event["toolCallId"], event["content"]) for event in tool_results] == [("call_provider", "wrote to disk")] + assert not any( + content.type == "function_approval_response" for message in provider_messages for content in message.contents + ) # And it was neither rejected nor reported as a transport failure (the #7043 bug). assert "Tool call invocation was rejected" not in resume_text assert "Tool call invocation failed" not in resume_text assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + + +async def test_endpoint_canonical_resume_preserves_hosted_approval_for_provider( + streaming_chat_client_stub, +) -> None: + """Canonical AG-UI resume keeps trusted hosted metadata and never executes a local name collision.""" + call_id = "mcpr_docs" + server_label = "Microsoft_Learn_MCP" + state = {"phase": "pause"} + local_executions: list[str] = [] + provider_messages: list[Message] = [] + hosted_call = Content.from_function_call( + call_id=call_id, + name="docs_search", + arguments={"query": "azure"}, + additional_properties={"server_label": server_label}, + ) + + def docs_search(query: str) -> str: + local_executions.append(query) + return f"local:{query}" + + local_tool = FunctionTool( + name="docs_search", + description="A local tool whose name collides with the hosted tool.", + func=docs_search, + ) + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + if state["phase"] == "pause": + yield ChatResponseUpdate( + contents=[Content.from_function_approval_request(id=call_id, function_call=hosted_call)], + role="assistant", + ) + return + provider_messages[:] = list(messages) + yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") + + agent = Agent( + name="test_agent", + instructions="Test", + client=streaming_chat_client_stub(stream_fn), + tools=[local_tool], + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + client = TestClient(app) + + pause_response = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": "thread-hosted-approval", + "messages": [{"role": "user", "content": "Search the hosted docs"}], + }, + ) + assert pause_response.status_code == 200 + pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + pause_interrupts = _run_finished_interrupts(pause_finished[-1]) + assert [interrupt["id"] for interrupt in pause_interrupts] == [call_id] + assert set(pause_interrupts[0]["responseSchema"]["properties"]) == {"accepted"} + + state["phase"] = "resume" + resume_response = client.post( + "/approval", + json={ + "runId": "run-resume", + "threadId": "thread-hosted-approval", + "messages": [], + "resume": [{"interruptId": call_id, "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert resume_response.status_code == 200 + assert not [event for event in _decode_sse_events(resume_response) if event.get("type") == "RUN_ERROR"] + assert local_executions == [] + assert not wrapped_agent._pending_approvals # pyright: ignore[reportPrivateUsage] + approval_responses = [ + content + for message in provider_messages + for content in message.contents + if content.type == "function_approval_response" + ] + assert len(approval_responses) == 1 + assert approval_responses[0].id == call_id + assert approval_responses[0].approved is True + assert approval_responses[0].function_call is not None + assert approval_responses[0].function_call.additional_properties["server_label"] == server_label + + +async def test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client( + streaming_chat_client_stub, +) -> None: + """AG-UI consumes local approval controls before the raw provider boundary.""" + call_id = "call_local_approval" + state = {"phase": "pause"} + provider_messages: list[Message] = [] + local_executions: list[str] = [] + function_call = Content.from_function_call( + call_id=call_id, + name="local_action", + arguments={"document": "Approved draft"}, + ) + + def local_action(document: str) -> str: + local_executions.append(document) + return "Action executed again" + + local_tool = FunctionTool( + name="local_action", + description="A local action whose replayed result proves it already completed.", + func=local_action, + approval_mode="always_require", + ) + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + if state["phase"] == "pause": + yield ChatResponseUpdate( + contents=[Content.from_function_approval_request(id=call_id, function_call=function_call)], + role="assistant", + ) + return + provider_messages[:] = list(messages) + yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") + + chat_client = cast(Any, streaming_chat_client_stub(stream_fn)) + chat_client.function_invocation_configuration["enabled"] = False + agent = Agent( + name="test_agent", + instructions="Test", + client=chat_client, + tools=[local_tool], + ) + wrapped_agent = AgentFrameworkAgent( + agent=agent, + state_schema={"document": {"type": "string"}}, + predict_state_config={"document": {"tool": "local_action", "tool_argument": "document"}}, + require_confirmation=False, + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + wrapped_agent, + path="/approval", + ) + client = TestClient(app) + + pause_response = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": "thread-local-approval", + "messages": [{"role": "user", "content": "Run the local action"}], + }, + ) + assert pause_response.status_code == 200 + pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + assert [interrupt["id"] for interrupt in _run_finished_interrupts(pause_finished[-1])] == [call_id] + + state["phase"] = "resume" + resume_response = client.post( + "/approval", + json={ + "runId": "run-resume", + "threadId": "thread-local-approval", + "messages": [ + {"role": "user", "content": "Run the local action"}, + { + "role": "assistant", + "toolCalls": [ + { + "id": call_id, + "type": "function", + "function": { + "name": "local_action", + "arguments": '{"document":"Approved draft"}', + }, + } + ], + }, + {"role": "tool", "toolCallId": call_id, "content": "Action already completed"}, + { + "role": "user", + "function_approvals": [ + { + "id": call_id, + "call_id": call_id, + "name": "local_action", + "approved": True, + "arguments": {"document": "Approved draft"}, + } + ], + }, + ], + "state": {"document": "Old draft"}, + "resume": [{"interruptId": call_id, "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert resume_response.status_code == 200 + assert local_executions == [] + assert not wrapped_agent._pending_approvals # pyright: ignore[reportPrivateUsage] + state_snapshots = [ + event["snapshot"] for event in _decode_sse_events(resume_response) if event.get("type") == "STATE_SNAPSHOT" + ] + assert {"document": "Approved draft"} in state_snapshots + assert not any( + content.type == "function_approval_response" for message in provider_messages for content in message.contents + ) + assert any( + content.type == "function_result" and content.call_id == call_id + for message in provider_messages + for content in message.contents + ) diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 261f400e143..92446b623ec 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -30,6 +30,7 @@ _build_safe_metadata, _canonical_approval_resume_messages, _create_state_context_message, + _filter_local_approval_responses_for_provider, _inject_state_context, _make_pending_approval_entry, _normalize_response_stream, @@ -61,6 +62,180 @@ def _message_role(message: object) -> object: return getattr(message, "role", None) +def test_filter_local_approval_responses_for_provider_removes_only_completed_local_controls() -> None: + """Provider-bound filtering removes completed local controls without mutating caller messages.""" + local_call = Content.from_function_call(call_id="call_local_mixed", name="local_tool", arguments={}) + local_response = Content.from_function_approval_response( + approved=True, + id="approval_local_mixed", + function_call=local_call, + ) + control_call = Content.from_function_call(call_id="call_local_control", name="local_tool", arguments={}) + control_response = Content.from_function_approval_response( + approved=True, + id="approval_local_control", + function_call=control_call, + ) + hosted_call = Content.from_function_call( + call_id="call_hosted", + name="hosted_tool", + arguments={}, + additional_properties={"server_label": "hosted-server"}, + ) + hosted_response = Content.from_function_approval_response( + approved=True, + id="approval_hosted", + function_call=hosted_call, + ) + pending_call = Content.from_function_call(call_id="call_pending", name="pending_tool", arguments={}) + pending_response = Content.from_function_approval_response( + approved=True, + id="approval_pending", + function_call=pending_call, + ) + completed_message = Message( + role="tool", + contents=[ + Content.from_function_result(call_id="call_local_mixed", result="completed"), + Content.from_function_result(call_id="call_local_control", result="completed"), + ], + ) + mixed_message = Message( + role="user", + contents=[Content.from_text(text="Keep this text"), local_response], + ) + control_only_message = Message(role="user", contents=[control_response]) + hosted_message = Message(role="user", contents=[hosted_response]) + pending_message = Message(role="user", contents=[pending_response]) + empty_message = Message(role="user", contents=[]) + + filtered = _filter_local_approval_responses_for_provider( + [completed_message, mixed_message, control_only_message, hosted_message, pending_message, empty_message] + ) + + assert len(filtered) == 5 + assert filtered[0] is completed_message + assert filtered[1] is not mixed_message + assert [content.type for content in filtered[1].contents] == ["text"] + assert filtered[2] is hosted_message + assert filtered[3] is pending_message + assert filtered[4] is empty_message + assert [content.type for content in mixed_message.contents] == ["text", "function_approval_response"] + assert control_only_message.contents == [control_response] + + +def test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence() -> None: + """One completed occurrence does not erase a later approval that reused its call id.""" + call_id = "call_reused" + first_call = Content.from_function_call(call_id=call_id, name="local_tool", arguments={"turn": 1}) + first_request = Content.from_function_approval_request(id=call_id, function_call=first_call) + first_response = Content.from_function_approval_response( + approved=True, + id=call_id, + function_call=first_call, + ) + second_call = Content.from_function_call(call_id=call_id, name="local_tool", arguments={"turn": 2}) + second_response = Content.from_function_approval_response( + approved=True, + id=call_id, + function_call=second_call, + ) + first_call_message = Message(role="assistant", contents=[first_call, first_request]) + completed_message = Message( + role="tool", + contents=[Content.from_function_result(call_id=call_id, result="first completed")], + ) + first_response_message = Message(role="user", contents=[first_response]) + second_call_message = Message(role="assistant", contents=[second_call]) + second_response_message = Message(role="user", contents=[second_response]) + + filtered = _filter_local_approval_responses_for_provider( + [ + first_call_message, + completed_message, + first_response_message, + second_call_message, + second_response_message, + ] + ) + + assert filtered == [first_call_message, completed_message, second_call_message, second_response_message] + + +def test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls() -> None: + """All replayed responses for one completed approval occurrence are removed.""" + call = Content.from_function_call(call_id="call_duplicate", name="local_tool", arguments={}) + first_response = Content.from_function_approval_response( + approved=True, + id="approval_duplicate", + function_call=call, + ) + replayed_response = Content.from_dict(first_response.to_dict()) + call_message = Message(role="assistant", contents=[call]) + result_message = Message( + role="tool", + contents=[Content.from_function_result(call_id="call_duplicate", result="completed")], + ) + + filtered = _filter_local_approval_responses_for_provider( + [ + call_message, + result_message, + Message(role="user", contents=[first_response, replayed_response]), + ] + ) + + assert filtered == [call_message, result_message] + + +def test_filter_local_approval_responses_for_provider_prefers_fresh_reused_call_for_edited_response() -> None: + """An edited fresh response is not paired to an older exact-argument occurrence.""" + call_id = "call_reused_edited" + first_call = Content.from_function_call( + call_id=call_id, + name="local_tool", + arguments={"value": "approved"}, + ) + first_response = Content.from_function_approval_response( + approved=True, + id="approval_old", + function_call=first_call, + ) + second_call = Content.from_function_call( + call_id=call_id, + name="local_tool", + arguments={"value": "original"}, + ) + edited_second_response = Content.from_function_approval_response( + approved=True, + id="approval_fresh", + function_call=Content.from_function_call( + call_id=call_id, + name="local_tool", + arguments={"value": "approved"}, + ), + ) + first_call_message = Message(role="assistant", contents=[first_call]) + first_result_message = Message( + role="tool", + contents=[Content.from_function_result(call_id=call_id, result="completed")], + ) + second_call_message = Message(role="assistant", contents=[second_call]) + fresh_response_message = Message(role="user", contents=[edited_second_response]) + + filtered = _filter_local_approval_responses_for_provider( + [ + first_call_message, + first_result_message, + Message(role="user", contents=[first_response]), + second_call_message, + fresh_response_message, + ] + ) + + assert filtered == [first_call_message, first_result_message, second_call_message, fresh_response_message] + + class TestBuildSafeMetadata: """Tests for _build_safe_metadata function.""" @@ -870,6 +1045,39 @@ def test_canonical_approval_resume_does_not_mutate_arguments_until_batch_validat assert pending_entry["arguments"] == '{"city":"Seattle"}' +def test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending() -> None: + """Hosted approvals accept a decision only because providers ignore edited arguments.""" + pending_entry = _make_pending_approval_entry( + "docs_search", + '{"query":"azure"}', + request_id="mcpr_docs", + interrupt_id="mcpr_docs", + server_label="Microsoft_Learn_MCP", + ) + key = _pending_approval_key("thread-hosted", "mcpr_docs") + pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = {key: pending_entry} + + messages, handled_ids, cancelled_ids, error = _canonical_approval_resume_messages( + [ + { + "interruptId": "mcpr_docs", + "status": "resolved", + "payload": {"accepted": True, "query": "untrusted edit"}, + } + ], + pending_approvals, + "thread-hosted", + ) + + assert messages == [] + assert handled_ids == {"mcpr_docs"} + assert cancelled_ids == set() + assert error is not None + assert error.code == "APPROVAL_RESUME_INVALID_RESPONSE" + assert pending_entry["arguments"] == '{"query":"azure"}' + assert pending_approvals[key] is pending_entry + + def test_pending_approval_registry_scans_exact_thread_keys_with_colons(): """A thread id that prefixes another thread id must not inherit its pending approval contract.""" pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry] = { diff --git a/python/packages/ollama/agent_framework_ollama/_chat_client.py b/python/packages/ollama/agent_framework_ollama/_chat_client.py index b2815b7a96f..f8c7a2232ec 100644 --- a/python/packages/ollama/agent_framework_ollama/_chat_client.py +++ b/python/packages/ollama/agent_framework_ollama/_chat_client.py @@ -439,7 +439,7 @@ def _prepare_options(self, messages: Sequence[Message], options: Mapping[str, An # messages if messages and "messages" not in run_options: run_options["messages"] = self._prepare_messages_for_ollama(messages) - if not run_options.get("messages"): + if "messages" not in run_options: raise ChatClientInvalidRequestException("Messages are required for chat completions") # model @@ -474,9 +474,6 @@ def _format_system_message(self, message: Message) -> list[OllamaMessage]: def _format_user_message(self, message: Message) -> list[OllamaMessage]: if not any(c.type in {"text", "data"} for c in message.contents) and not message.text: - if message.contents and all(content.type == "function_approval_response" for content in message.contents): - # AG-UI resolves approvals in-process; this control-only message has no Ollama representation. - return [] raise ChatClientInvalidRequestException( "Ollama connector currently only supports user messages with TextContent or DataContent." ) diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index a4b09d500de..70f578cc37a 100644 --- a/python/packages/ollama/tests/test_ollama_chat_client.py +++ b/python/packages/ollama/tests/test_ollama_chat_client.py @@ -263,74 +263,6 @@ async def test_cmc( assert result.text == "test" -@patch.object(AsyncClient, "chat", new_callable=AsyncMock) -async def test_cmc_ignores_control_only_approval_resume_message( - mock_chat: AsyncMock, - ollama_unit_test_env: dict[str, str], -) -> None: - """A resolved approval control is not sent to an Ollama model as a user message.""" - mock_chat.return_value = OllamaChatResponse( - message=OllamaMessage(content="done", role="assistant"), - model="test", - ) - function_call = Content.from_function_call( - call_id="call_123", - name="dangerous_action", - arguments={"target": "production"}, - ) - approval_response = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=function_call, - ) - messages = [ - Message(role="user", contents=[Content.from_text(text="Continue the task.")]), - Message(role="assistant", contents=[function_call]), - Message(role="user", contents=[approval_response]), - Message( - role="tool", - contents=[Content.from_function_result(call_id="call_123", result="The task is complete.")], - ), - ] - - ollama_client = OllamaChatClient() - result = await ollama_client.get_response(messages=messages) - - assert result.text == "done" - assert mock_chat.await_args is not None - outgoing_messages = mock_chat.await_args.kwargs["messages"] - assert [message["role"] for message in outgoing_messages] == ["user", "assistant", "tool"] - assert outgoing_messages[0]["content"] == "Continue the task." - assert outgoing_messages[2]["content"] == "The task is complete." - - -@patch.object(AsyncClient, "chat", new_callable=AsyncMock) -async def test_cmc_rejects_only_control_approval_resume_message( - mock_chat: AsyncMock, - ollama_unit_test_env: dict[str, str], -) -> None: - """A transcript with no Ollama-representable messages is rejected before the SDK call.""" - mock_chat.return_value = OllamaChatResponse( - message=OllamaMessage(content="unexpected", role="assistant"), - model="test", - ) - function_call = Content.from_function_call( - call_id="call_123", - name="dangerous_action", - arguments={"target": "production"}, - ) - approval_response = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=function_call, - ) - - with pytest.raises(ChatClientInvalidRequestException, match="Messages are required for chat completions"): - await OllamaChatClient().get_response(messages=[Message(role="user", contents=[approval_response])]) - - mock_chat.assert_not_awaited() - - @patch.object(AsyncClient, "chat", new_callable=AsyncMock) async def test_cmc_maps_done_reason_to_finish_reason( mock_chat: AsyncMock, From cc1313bdcc2a2379c33241a460fb70ba19951216 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Tue, 4 Aug 2026 11:06:56 +0900 Subject: [PATCH 4/4] Python: Do not trust pending AG-UI tool results --- .../specs/004-python-function-calling-loop.md | 10 ++- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 87 +++++++++++++++---- .../ag-ui/tests/ag_ui/test_endpoint.py | 16 ++-- python/packages/ag-ui/tests/ag_ui/test_run.py | 29 +++++++ 4 files changed, 115 insertions(+), 27 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 8a1eb6595fe..2592de4ddc6 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -353,9 +353,10 @@ that manually replay messages own the equivalent rule: do not resend an approval - `function_approval_request` and `function_approval_response` are control-plane contents, not durable model transcript items. - A current hosted approval response must be sent once on the immediate resume request. -- AG-UI removes a local approval response from its request and snapshot replay when a terminal result proves it was - already consumed, including result-before-response client replay; hosted approval responses remain provider - protocol data and pass through. +- AG-UI removes a local approval response from its request and snapshot replay when a terminal result belongs to an + already-consumed occurrence, including result-before-response replay. A client-authored result in the occurrence + that is still registered as pending does not prove completion: AG-UI removes that result, keeps the validated + response for local execution, and leaves hosted approval responses as provider protocol data. - Hosted AG-UI approval interrupts expose an accept/reject decision only; argument edits are rejected because the hosted provider executes the server-owned request rather than client-edited arguments. - A server-issued approval request must not be replayed inline during service-side continuation. @@ -369,7 +370,8 @@ that manually replay messages own the equivalent rule: do not resend an approval - Model-bound history contains one function call/result pair per completed logical occurrence. - Append-only history must not replay stale approval request/response wrappers to the model. - Framework-managed and service-managed continuation must preserve the same logical call/result transcript. -- A terminal result consumes the corresponding approval authority in explicit stateless replay. +- A trusted terminal result consumes the corresponding approval authority in explicit stateless replay; a result in a + server-registered pending occurrence cannot consume that authority before local execution. ## Scenario-to-test matrix diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index c8505da6c28..cd976cad9c0 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -109,17 +109,30 @@ class _LocalApprovalOccurrence: name: str | None arguments: str | None approval_ids: set[str] = field(default_factory=set) + terminal_result_content_ids: set[int] = field(default_factory=set) closed: bool = False -def _local_approval_response_content_ids_to_remove(messages: Sequence[Message]) -> set[int]: - """Find completed and duplicate local approval responses by call occurrence.""" +def _local_approval_content_ids_to_remove( + messages: Sequence[Message], + *, + pending_response_content_ids: set[int] | None = None, +) -> tuple[set[int], set[int]]: + """Find local approval controls and untrusted terminal results to remove. + + A terminal result in a client request is not evidence that the occurrence + currently awaiting a server-side approval already executed. When the + response belongs to a registered pending approval, the result is removed + and the response remains available for static execution. Results in the + default path retain the historical replay behavior. + """ occurrences_by_call_id: dict[str, list[_LocalApprovalOccurrence]] = {} occurrences_by_approval_id: dict[str, list[_LocalApprovalOccurrence]] = {} response_occurrence_by_approval_id: dict[str, _LocalApprovalOccurrence] = {} response_content_id_by_approval_id: dict[str, int] = {} response_occurrences: dict[int, _LocalApprovalOccurrence] = {} duplicate_response_content_ids: set[int] = set() + untrusted_result_content_ids: set[int] = set() def add_occurrence(function_call: Content, *, closed: bool = False) -> _LocalApprovalOccurrence | None: if function_call.call_id is None: @@ -205,11 +218,18 @@ def matching_occurrence( continue previous_occurrence = response_occurrence_by_approval_id.get(content.id or "") request_occurrences = occurrences_by_approval_id.get(content.id or "", []) - request_occurrence = matching_occurrence(request_occurrences, function_call, prefer_open=True) + is_pending_response = ( + pending_response_content_ids is not None and id(content) in pending_response_content_ids + ) + request_occurrence = matching_occurrence( + request_occurrences, + function_call, + prefer_open=not is_pending_response, + ) call_occurrence = matching_occurrence( occurrences_by_call_id.get(function_call.call_id, []), function_call, - prefer_open=True, + prefer_open=not is_pending_response, ) occurrence = ( call_occurrence @@ -254,22 +274,49 @@ def matching_occurrence( arguments=None, ) occurrences_by_call_id.setdefault(content.call_id, []).append(occurrence) + occurrence.terminal_result_content_ids.add(id(content)) occurrence.closed = True - return duplicate_response_content_ids | { - content_id for content_id, occurrence in response_occurrences.items() if occurrence.closed - } + if pending_response_content_ids: + for response_content_id, occurrence in response_occurrences.items(): + if response_content_id not in pending_response_content_ids: + continue + # A client-supplied result cannot close the server-owned pending + # occurrence. Remove it before static execution so it cannot be + # mistaken for the result produced by the approved local tool. + untrusted_result_content_ids.update(occurrence.terminal_result_content_ids) + occurrence.closed = False + + return ( + duplicate_response_content_ids + | {content_id for content_id, occurrence in response_occurrences.items() if occurrence.closed}, + untrusted_result_content_ids, + ) -def _filter_local_approval_responses_for_provider(messages: Sequence[Message]) -> list[Message]: +def _local_approval_response_content_ids_to_remove(messages: Sequence[Message]) -> set[int]: + """Find completed and duplicate local approval responses by call occurrence.""" + response_content_ids, _ = _local_approval_content_ids_to_remove(messages) + return response_content_ids + + +def _filter_local_approval_responses_for_provider( + messages: Sequence[Message], + *, + pending_response_content_ids: set[int] | None = None, +) -> list[Message]: """Remove completed local approval controls from AG-UI provider input. - A matching terminal result proves the local response has already been consumed, - even when client replay placed the result before the synthesized response. Local - responses without a result remain available to the in-run approval middleware, - while hosted-service responses remain provider protocol data. + A matching terminal result only closes a local response when it is not part + of a server-registered pending approval occurrence. Client-authored results + for a still-pending occurrence are removed before execution, while hosted- + service responses remain provider protocol data. """ - response_content_ids_to_remove = _local_approval_response_content_ids_to_remove(messages) + response_content_ids_to_remove, untrusted_result_content_ids = _local_approval_content_ids_to_remove( + messages, + pending_response_content_ids=pending_response_content_ids, + ) + response_content_ids_to_remove.update(untrusted_result_content_ids) filtered_messages: list[Message] = [] for message in messages: @@ -1455,9 +1502,11 @@ async def _resolve_approval_responses( responses_by_id.setdefault(content.id, []).append(content) valid_response_content_ids: set[int] | None = None + pending_local_response_content_ids: set[int] | None = None response_content_ids_to_strip: set[int] = set() if pending_approvals is not None: valid_response_content_ids = set() + pending_local_response_content_ids = set() def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: PendingApprovalEntry) -> bool: if isinstance(candidate, str) and isinstance(expected, str): @@ -1553,6 +1602,8 @@ def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: Pend else: primary_response.function_call.additional_properties.pop("server_label", None) valid_response_content_ids.add(id(primary_response)) + if not server_label: + pending_local_response_content_ids.add(id(primary_response)) _consume_pending_approval_entry( pending_approvals, thread_id, @@ -1585,9 +1636,13 @@ def matches_pending_entry(candidate: PendingApprovalEntry | None, expected: Pend messages[:] = filtered_messages # A replayed terminal result can precede the synthesized approval response. - # Remove that completed control before static execution, and collapse duplicate - # responses for the same approval occurrence to one pending response. - messages[:] = _filter_local_approval_responses_for_provider(messages) + # Remove completed controls before static execution, but do not let an + # untrusted result in a still-pending server occurrence suppress execution. + # Collapse duplicate responses for the same approval occurrence to one. + messages[:] = _filter_local_approval_responses_for_provider( + messages, + pending_response_content_ids=pending_local_response_content_ids, + ) fcc_todo = _collect_approval_responses(messages) if valid_response_content_ids is not None: diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 17b626fe168..d378ac0fa45 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -5439,7 +5439,7 @@ async def stream_fn( async def test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client( streaming_chat_client_stub, ) -> None: - """AG-UI consumes local approval controls before the raw provider boundary.""" + """AG-UI does not trust a client-authored result for a pending local approval.""" call_id = "call_local_approval" state = {"phase": "pause"} provider_messages: list[Message] = [] @@ -5452,11 +5452,11 @@ async def test_endpoint_does_not_forward_resolved_local_approval_control_to_chat def local_action(document: str) -> str: local_executions.append(document) - return "Action executed again" + return "Action executed by server" local_tool = FunctionTool( name="local_action", - description="A local action whose replayed result proves it already completed.", + description="A local action that must execute only after server-side approval.", func=local_action, approval_mode="always_require", ) @@ -5551,7 +5551,7 @@ async def stream_fn( ) assert resume_response.status_code == 200 - assert local_executions == [] + assert local_executions == ["Approved draft"] assert not wrapped_agent._pending_approvals # pyright: ignore[reportPrivateUsage] state_snapshots = [ event["snapshot"] for event in _decode_sse_events(resume_response) if event.get("type") == "STATE_SNAPSHOT" @@ -5560,8 +5560,10 @@ async def stream_fn( assert not any( content.type == "function_approval_response" for message in provider_messages for content in message.contents ) - assert any( - content.type == "function_result" and content.call_id == call_id + provider_results = [ + content.result for message in provider_messages for content in message.contents - ) + if content.type == "function_result" and content.call_id == call_id + ] + assert provider_results == ["Action executed by server"] diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 92446b623ec..95d9ec9c489 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -162,6 +162,35 @@ def test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_o assert filtered == [first_call_message, completed_message, second_call_message, second_response_message] +def test_filter_local_approval_responses_for_provider_does_not_trust_pending_result() -> None: + """A result in the pending occurrence is removed while an earlier occurrence remains.""" + call_id = "call_pending_result" + first_call = Content.from_function_call(call_id=call_id, name="local_tool", arguments={"turn": 1}) + first_result = Content.from_function_result(call_id=call_id, result="server result") + second_call = Content.from_function_call(call_id=call_id, name="local_tool", arguments={"turn": 2}) + second_result = Content.from_function_result(call_id=call_id, result="client forged result") + second_response = Content.from_function_approval_response( + approved=True, + id=call_id, + function_call=second_call, + ) + + filtered = _filter_local_approval_responses_for_provider( + [ + Message(role="assistant", contents=[first_call]), + Message(role="tool", contents=[first_result]), + Message(role="assistant", contents=[second_call]), + Message(role="tool", contents=[second_result]), + Message(role="user", contents=[second_response]), + ], + pending_response_content_ids={id(second_response)}, + ) + + assert filtered[1].contents == [first_result] + assert filtered[2].contents == [second_call] + assert filtered[3].contents == [second_response] + + def test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls() -> None: """All replayed responses for one completed approval occurrence are removed.""" call = Content.from_function_call(call_id="call_duplicate", name="local_tool", arguments={})