From 0545c4fe55d5ac29c5ce1f107138c0bede4d7b16 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Fri, 24 Jul 2026 13:56:29 -0700 Subject: [PATCH 1/2] fix: Prevent continuation forgery in tool confirmation An attacker who could manipulate or inject events into the session history could execute unauthorized tools by forging a tool confirmation response. This fixes the vulnerability by: - When resolving confirmation targets, the processor verifies if the tool is registered in the executing agent's tools_dict - Validate that the tool actually requires confirmation, supporting both static definitions and dynamic confirmation requests - Verify that the original tool call event exists in the session history with the matching ID, and that its name and arguments match the confirmation request's originalFunctionCall exactly to prevent argument tampering. Co-authored-by: Xuan Yang PiperOrigin-RevId: 953540969 Change-Id: Iff6e8c861605fafafce4985ee9a269274d6d789c --- src/google/adk/agents/invocation_context.py | 7 +- .../flows/llm_flows/request_confirmation.py | 162 ++++- src/google/adk/tools/base_tool.py | 6 + src/google/adk/tools/function_tool.py | 52 +- src/google/adk/tools/mcp_tool/mcp_tool.py | 93 ++- src/google/adk/tools/tool_confirmation.py | 19 +- .../llm_flows/test_request_confirmation.py | 646 +++++++++++++++++- 7 files changed, 895 insertions(+), 90 deletions(-) diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 614770e8cb6..34e234ff0b3 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -365,7 +365,12 @@ def _get_events( if event.invocation_id == self.invocation_id ] if current_branch: - results = [event for event in results if event.branch == self.branch] + results = [ + event + for event in results + if event.branch == self.branch + or (event.branch is None and event.author == "user") + ] return results def should_pause_invocation(self, event: Event) -> bool: diff --git a/src/google/adk/flows/llm_flows/request_confirmation.py b/src/google/adk/flows/llm_flows/request_confirmation.py index d066db791df..2de3a4a90e1 100644 --- a/src/google/adk/flows/llm_flows/request_confirmation.py +++ b/src/google/adk/flows/llm_flows/request_confirmation.py @@ -13,7 +13,6 @@ # limitations under the License. from __future__ import annotations -import json import logging from typing import Any from typing import AsyncGenerator @@ -27,7 +26,9 @@ from ...agents.readonly_context import ReadonlyContext from ...events.event import Event from ...models.llm_request import LlmRequest +from ...tools.base_tool import BaseTool from ...tools.tool_confirmation import ToolConfirmation +from ...tools.tool_context import ToolContext from ._base_llm_processor import BaseLlmRequestProcessor from .functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME @@ -35,60 +36,147 @@ from ...agents.llm_agent import LlmAgent -logger = logging.getLogger('google_adk.' + __name__) +logger = logging.getLogger("google_adk." + __name__) def _parse_tool_confirmation(response: dict[str, Any]) -> ToolConfirmation: - """Parse ToolConfirmation from a function response dict. + """Parses ToolConfirmation from a function response dict.""" + return ToolConfirmation.from_response_dict(response) - Handles both the direct dict format and the ADK client's - ``{'response': json_string}`` wrapper format. - """ - if response and len(response.values()) == 1 and 'response' in response.keys(): - return ToolConfirmation.model_validate(json.loads(response['response'])) - return ToolConfirmation.model_validate(response) - - -def _resolve_confirmation_targets( +async def _resolve_confirmation_targets( + invocation_context: InvocationContext, events: list[Event], confirmation_fc_ids: set[str], confirmations_by_fc_id: dict[str, ToolConfirmation], + tools_dict: dict[str, BaseTool], ) -> tuple[dict[str, ToolConfirmation], dict[str, types.FunctionCall]]: - """Find original function calls for confirmed tools. + """Find original function calls for confirmed tools and validate them. Scans events for ``adk_request_confirmation`` function calls whose IDs are in *confirmation_fc_ids*, extracts the ``originalFunctionCall`` from - their args, and maps each confirmation to the original FC ID. + their args, validates that they are registered, actually require confirmation, + and match the original function calls in history, and maps each confirmation + to the original FC ID. Args: + invocation_context: Current invocation context. events: Session events to scan. confirmation_fc_ids: IDs of ``adk_request_confirmation`` function calls. confirmations_by_fc_id: Mapping of confirmation FC ID -> ``ToolConfirmation``. + tools_dict: Dictionary of registered tools. Returns: Tuple of ``(tool_confirmation_dict, original_fcs_dict)`` where both are keyed by the ORIGINAL function call IDs. + + Raises: + ValueError: If validation of any confirmation target fails. """ tool_confirmation_dict: dict[str, ToolConfirmation] = {} original_fcs_dict: dict[str, types.FunctionCall] = {} + history_fcs = { + fc.id: (fc, ev) + for ev in events + for fc in ev.get_function_calls() + if fc.id and fc.name != REQUEST_CONFIRMATION_FUNCTION_CALL_NAME + } + history_fr_events = { + fr.id: ev for ev in events for fr in ev.get_function_responses() if fr.id + } + for event in events: event_function_calls = event.get_function_calls() if not event_function_calls: continue for function_call in event_function_calls: - if function_call.id not in confirmation_fc_ids: + if not function_call.id or function_call.id not in confirmation_fc_ids: continue args = function_call.args - if 'originalFunctionCall' not in args: + if not args or "originalFunctionCall" not in args: continue original_function_call = types.FunctionCall( - **args['originalFunctionCall'] + **args["originalFunctionCall"] + ) + if not original_function_call.id: + raise ValueError("Original function call ID is missing.") + tool_name = original_function_call.name + if not tool_name: + raise ValueError("Original function call name is missing.") + + # Check 1: Is the tool registered? + original_fc_info = history_fcs.get(original_function_call.id) + if not original_fc_info: + raise ValueError( + f"Original function call for ID '{original_function_call.id}' not" + " found in session history." + ) + original_fc_in_history, original_fc_event = original_fc_info + + # If this tool call was authored by another agent, skip it to let that + # agent's processor handle it. + agent = invocation_context.agent + if agent and original_fc_event.author != agent.name: + continue + + tool = tools_dict.get(tool_name) + if not tool: + raise ValueError( + f"Tool '{original_function_call.name}' is not registered." + ) + + # Check 2: Does the tool require confirmation for these arguments? + # We check if it is either statically required, or if it was dynamically + # requested in the session history. + temp_tool_context = ToolContext( + invocation_context=invocation_context, + function_call_id=original_function_call.id, ) + requires_confirmation = await tool.check_require_confirmation( + original_function_call.args or {}, temp_tool_context + ) + + requested_in_history = False + if not requires_confirmation: + # Search the history for the response event of the original tool call + original_response_event = history_fr_events.get( + original_function_call.id + ) + if ( + original_response_event + and original_response_event.actions.requested_tool_confirmations + ): + requested_in_history = ( + original_function_call.id + in original_response_event.actions.requested_tool_confirmations + ) + + if not requires_confirmation and not requested_in_history: + raise ValueError( + f"Tool '{original_function_call.name}' does not require" + " confirmation." + ) + + # Check 3: Does the original function call match name and arguments? + if original_fc_in_history.name != original_function_call.name: + raise ValueError( + f"Function call name mismatch for ID '{original_function_call.id}':" + f" history has '{original_fc_in_history.name}', confirmation has" + f" '{original_function_call.name}'." + ) + + hist_args = original_fc_in_history.args or {} + conf_args = original_function_call.args or {} + if hist_args != conf_args: + raise ValueError( + "Function call arguments mismatch for ID" + f" '{original_function_call.id}'." + ) + tool_confirmation_dict[original_function_call.id] = ( confirmations_by_fc_id[function_call.id] ) @@ -116,10 +204,9 @@ async def run_async( # Step 1: Find the last user-authored event and parse confirmation # responses from it. confirmations_by_fc_id: dict[str, ToolConfirmation] = {} - confirmation_event_index = -1 for k in range(len(events) - 1, -1, -1): event = events[k] - if not event.author or event.author != 'user': + if not event.author or event.author != "user": continue responses = event.get_function_responses() if not responses: @@ -128,20 +215,35 @@ async def run_async( for function_response in responses: if function_response.name != REQUEST_CONFIRMATION_FUNCTION_CALL_NAME: continue + if not function_response.id or function_response.response is None: + continue confirmations_by_fc_id[function_response.id] = _parse_tool_confirmation( function_response.response ) - confirmation_event_index = k break if not confirmations_by_fc_id: return + # Resolve all canonical tools and build tools_dict + tools_dict = {} + if agent is not None and hasattr(agent, "canonical_tools"): + tools_dict = { + tool.name: tool + for tool in await agent.canonical_tools( + ReadonlyContext(invocation_context) + ) + } + # Step 2: Resolve confirmation targets using extracted helper. confirmation_fc_ids = set(confirmations_by_fc_id.keys()) tools_to_resume_with_confirmation, tools_to_resume_with_args = ( - _resolve_confirmation_targets( - events, confirmation_fc_ids, confirmations_by_fc_id + await _resolve_confirmation_targets( + invocation_context, + events, + confirmation_fc_ids, + confirmations_by_fc_id, + tools_dict, ) ) @@ -149,8 +251,9 @@ async def run_async( return # Step 3: Remove tools that have already been confirmed (dedup). - for i in range(len(events) - 1, confirmation_event_index, -1): - event = events[i] + for event in reversed(events): + if event.author == "user": + break fr_list = event.get_function_responses() if not fr_list: continue @@ -168,14 +271,9 @@ async def run_async( # Step 4: Re-execute the confirmed tools. if function_response_event := await functions.handle_function_call_list_async( invocation_context, - tools_to_resume_with_args.values(), - { - tool.name: tool - for tool in await agent.canonical_tools( - ReadonlyContext(invocation_context) - ) - }, - tools_to_resume_with_confirmation.keys(), + list(tools_to_resume_with_args.values()), + tools_dict, + set(tools_to_resume_with_confirmation.keys()), tools_to_resume_with_confirmation, ): yield function_response_event diff --git a/src/google/adk/tools/base_tool.py b/src/google/adk/tools/base_tool.py index e5c4bb73f98..7c14ca2573b 100644 --- a/src/google/adk/tools/base_tool.py +++ b/src/google/adk/tools/base_tool.py @@ -142,6 +142,12 @@ async def process_llm_request( # Use the consolidated logic in LlmRequest.append_tools llm_request.append_tools([self]) + async def check_require_confirmation( + self, args: dict[str, Any], tool_context: ToolContext + ) -> bool: + """Returns whether the tool requires confirmation for the given args.""" + return False + @property def _api_variant(self) -> GoogleLLMVariant: return get_google_llm_variant() diff --git a/src/google/adk/tools/function_tool.py b/src/google/adk/tools/function_tool.py index 10e32a5473d..cd77f7948f6 100644 --- a/src/google/adk/tools/function_tool.py +++ b/src/google/adk/tools/function_tool.py @@ -18,11 +18,16 @@ import logging from typing import Any from typing import Callable +from typing import cast from typing import get_args from typing import get_origin from typing import Optional +from typing import TYPE_CHECKING from typing import Union +if TYPE_CHECKING: + from ..agents.invocation_context import InvocationContext + from google.genai import types import pydantic from typing_extensions import override @@ -156,20 +161,35 @@ def _preprocess_args(self, args: dict[str, Any]) -> dict[str, Any]: return converted_args - @override - async def run_async( - self, *, args: dict[str, Any], tool_context: ToolContext - ) -> Any: - # Preprocess arguments (includes Pydantic model conversion) + def _prepare_invocation_args( + self, args: dict[str, Any], tool_context: ToolContext + ) -> dict[str, Any]: + """Prepare args for function invocation (preprocesses, injects context and filters).""" args_to_call = self._preprocess_args(args) - signature = inspect.signature(self.func) - valid_params = {param for param in signature.parameters} + valid_params = set(signature.parameters.keys()) if self._context_param_name in valid_params: args_to_call[self._context_param_name] = tool_context + return {k: v for k, v in args_to_call.items() if k in valid_params} + + @override + async def check_require_confirmation( + self, args: dict[str, Any], tool_context: ToolContext + ) -> bool: + if callable(self._require_confirmation): + args_to_call = self._prepare_invocation_args(args, tool_context) + return cast( + bool, + await self._invoke_callable(self._require_confirmation, args_to_call), + ) + return bool(self._require_confirmation) - # Filter args_to_call to only include valid parameters for the function - args_to_call = {k: v for k, v in args_to_call.items() if k in valid_params} + @override + async def run_async( + self, *, args: dict[str, Any], tool_context: ToolContext + ) -> Any: + # Preprocess arguments (includes Pydantic model conversion) + args_to_call = self._prepare_invocation_args(args, tool_context) # Before invoking the function, we check for if the list of args passed in # has all the mandatory arguments or not. @@ -188,12 +208,9 @@ async def run_async( You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters.""" return {'error': error_str} - if isinstance(self._require_confirmation, Callable): - require_confirmation = await self._invoke_callable( - self._require_confirmation, args_to_call - ) - else: - require_confirmation = bool(self._require_confirmation) + require_confirmation = await self.check_require_confirmation( + args, tool_context + ) if require_confirmation: if not tool_context.tool_confirmation: @@ -243,14 +260,15 @@ async def _call_live( *, args: dict[str, Any], tool_context: ToolContext, - invocation_context, + invocation_context: InvocationContext, ) -> Any: args_to_call = args.copy() signature = inspect.signature(self.func) # For input-streaming tools, the stream is created during # registration in _process_function_live_helper. Pass it here. if ( - self.name in invocation_context.active_streaming_tools + invocation_context.active_streaming_tools is not None + and self.name in invocation_context.active_streaming_tools and invocation_context.active_streaming_tools[self.name].stream is not None ): diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 6a24651f923..7c7a2bdd9f5 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -21,6 +21,7 @@ import os from typing import Any from typing import Callable +from typing import cast from typing import Dict from typing import List from typing import Optional @@ -292,43 +293,61 @@ async def _invoke_callable( else: return target(**args_to_call) + def _prepare_callable_args( + self, + target: Callable[..., Any], + args: dict[str, Any], + tool_context: ToolContext, + ) -> dict[str, Any]: + """Prepares arguments for invoking a user-provided callable.""" + args_to_call = args.copy() + try: + signature = inspect.signature(target) + except (ValueError, TypeError): + return args_to_call + + valid_params = set(signature.parameters.keys()) + has_kwargs = any( + param.kind == inspect.Parameter.VAR_KEYWORD + for param in signature.parameters.values() + ) + + # Detect context parameter by type or fallback to 'tool_context' name + context_param = find_context_parameter(target) or "tool_context" + if context_param in valid_params or has_kwargs: + args_to_call[context_param] = tool_context + + # Filter args_to_call only if there's no **kwargs + if not has_kwargs: + # Add context param to valid_params if it was added to args_to_call + if context_param in args_to_call: + valid_params.add(context_param) + args_to_call = { + k: v for k, v in args_to_call.items() if k in valid_params + } + return args_to_call + + @override + async def check_require_confirmation( + self, args: dict[str, Any], tool_context: ToolContext + ) -> bool: + if callable(self._require_confirmation): + args_to_call = self._prepare_callable_args( + self._require_confirmation, args, tool_context + ) + return cast( + bool, + await self._invoke_callable(self._require_confirmation, args_to_call), + ) + return bool(self._require_confirmation) + @override async def run_async( self, *, args: dict[str, Any], tool_context: ToolContext ) -> Any: - if isinstance(self._require_confirmation, Callable): - args_to_call = args.copy() - try: - signature = inspect.signature(self._require_confirmation) - valid_params = set(signature.parameters.keys()) - has_kwargs = any( - param.kind == inspect.Parameter.VAR_KEYWORD - for param in signature.parameters.values() - ) - - # Detect context parameter by type or fallback to 'tool_context' name - context_param = ( - find_context_parameter(self._require_confirmation) or "tool_context" - ) - if context_param in valid_params or has_kwargs: - args_to_call[context_param] = tool_context - - # Filter args_to_call only if there's no **kwargs - if not has_kwargs: - # Add context param to valid_params if it was added to args_to_call - if context_param in args_to_call: - valid_params.add(context_param) - args_to_call = { - k: v for k, v in args_to_call.items() if k in valid_params - } - except ValueError: - args_to_call = args - - require_confirmation = await self._invoke_callable( - self._require_confirmation, args_to_call - ) - else: - require_confirmation = bool(self._require_confirmation) + require_confirmation = await self.check_require_confirmation( + args, tool_context + ) if require_confirmation: if not tool_context.tool_confirmation: @@ -371,7 +390,11 @@ async def run_async( @retry_on_errors @override async def _run_async_impl( - self, *, args, tool_context: ToolContext, credential: AuthCredential + self, + *, + args: dict[str, Any], + tool_context: ToolContext, + credential: AuthCredential, ) -> Dict[str, Any]: """Runs the tool asynchronously. @@ -588,7 +611,7 @@ async def _get_headers( class MCPTool(McpTool): """Deprecated name, use `McpTool` instead.""" - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: warnings.warn( "MCPTool class is deprecated, use `McpTool` instead.", DeprecationWarning, diff --git a/src/google/adk/tools/tool_confirmation.py b/src/google/adk/tools/tool_confirmation.py index 683da17cebb..f1f2490b54c 100644 --- a/src/google/adk/tools/tool_confirmation.py +++ b/src/google/adk/tools/tool_confirmation.py @@ -11,9 +11,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - from __future__ import annotations +import json from typing import Any from typing import Optional @@ -43,3 +43,20 @@ class ToolConfirmation(BaseModel): payload: Optional[Any] = None """The custom data payload needed from the user to continue the flow. It should be JSON serializable.""" + + @classmethod + def from_response_dict(cls, response: dict[str, Any]) -> ToolConfirmation: + """Parse ToolConfirmation from a function response dict. + + Handles both the direct dict format and the ADK client's + ``{'response': json_string}`` wrapper format. + """ + if response and len(response) == 1 and "response" in response: + parsed = cls.model_validate(json.loads(response["response"])) + else: + parsed = cls.model_validate(response) + if isinstance(parsed, ToolConfirmation): + return parsed + raise TypeError( + f"Expected ToolConfirmation instance, got {type(parsed).__name__}" + ) diff --git a/tests/unittests/flows/llm_flows/test_request_confirmation.py b/tests/unittests/flows/llm_flows/test_request_confirmation.py index 39b35454b75..c8b55c47d31 100644 --- a/tests/unittests/flows/llm_flows/test_request_confirmation.py +++ b/tests/unittests/flows/llm_flows/test_request_confirmation.py @@ -17,9 +17,12 @@ from google.adk.agents.llm_agent import LlmAgent from google.adk.events.event import Event +from google.adk.events.event import EventActions from google.adk.flows.llm_flows import functions +from google.adk.flows.llm_flows.request_confirmation import _resolve_confirmation_targets from google.adk.flows.llm_flows.request_confirmation import request_processor from google.adk.models.llm_request import LlmRequest +from google.adk.tools.function_tool import FunctionTool from google.adk.tools.tool_confirmation import ToolConfirmation from google.genai import types import pytest @@ -112,7 +115,10 @@ async def test_request_confirmation_processor_no_confirmation_function_response( @pytest.mark.asyncio async def test_request_confirmation_processor_success(): """Test the successful processing of a tool confirmation.""" - agent = LlmAgent(name="test_agent", tools=[mock_tool]) + agent = LlmAgent( + name="test_agent", + tools=[FunctionTool(mock_tool, require_confirmation=True)], + ) invocation_context = await testing_utils.create_invocation_context( agent=agent ) @@ -122,6 +128,16 @@ async def test_request_confirmation_processor_success(): name=MOCK_TOOL_NAME, args={"param1": "test"}, id=MOCK_FUNCTION_CALL_ID ) + # Add original tool call to history + invocation_context.session.events.append( + Event( + author=agent.name, + content=types.Content( + parts=[types.Part(function_call=original_function_call)] + ), + ) + ) + tool_confirmation = ToolConfirmation(confirmed=False, hint="test hint") tool_confirmation_args = { "originalFunctionCall": original_function_call.model_dump( @@ -135,7 +151,7 @@ async def test_request_confirmation_processor_success(): # Event with the request for confirmation invocation_context.session.events.append( Event( - author="agent", + author=agent.name, content=types.Content( parts=[ types.Part( @@ -213,7 +229,10 @@ async def test_request_confirmation_processor_success(): @pytest.mark.asyncio async def test_request_confirmation_processor_tool_not_confirmed(): """Test when the tool execution is not confirmed by the user.""" - agent = LlmAgent(name="test_agent", tools=[mock_tool]) + agent = LlmAgent( + name="test_agent", + tools=[FunctionTool(mock_tool, require_confirmation=True)], + ) invocation_context = await testing_utils.create_invocation_context( agent=agent ) @@ -223,6 +242,16 @@ async def test_request_confirmation_processor_tool_not_confirmed(): name=MOCK_TOOL_NAME, args={"param1": "test"}, id=MOCK_FUNCTION_CALL_ID ) + # Add original tool call to history + invocation_context.session.events.append( + Event( + author=agent.name, + content=types.Content( + parts=[types.Part(function_call=original_function_call)] + ), + ) + ) + tool_confirmation = ToolConfirmation(confirmed=False, hint="test hint") tool_confirmation_args = { "originalFunctionCall": original_function_call.model_dump( @@ -235,7 +264,7 @@ async def test_request_confirmation_processor_tool_not_confirmed(): invocation_context.session.events.append( Event( - author="agent", + author=agent.name, content=types.Content( parts=[ types.Part( @@ -300,3 +329,612 @@ async def test_request_confirmation_processor_tool_not_confirmed(): assert ( args[4][MOCK_FUNCTION_CALL_ID] == user_confirmation ) # tool_confirmation_dict + + +@pytest.mark.asyncio +async def test_request_confirmation_processor_finds_user_confirmation_in_default_branch(): + """Processor finds user confirmation in default branch when agent is in child branch. + + Setup: + - Agent in 'child_branch'. + - RequestConfirmation event in 'child_branch'. + - User response event in default branch (None). + Act: Run request_processor. + Assert: Processor finds the response and triggers tool execution. + """ + # Arrange + agent = LlmAgent( + name="test_agent", + tools=[FunctionTool(mock_tool, require_confirmation=True)], + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + # Set branch for the agent context + invocation_context.branch = "child_branch" + llm_request = LlmRequest() + + original_function_call = types.FunctionCall( + name=MOCK_TOOL_NAME, args={"param1": "test"}, id=MOCK_FUNCTION_CALL_ID + ) + + # Add original tool call to history + invocation_context.session.events.append( + Event( + author=agent.name, + branch="child_branch", + content=types.Content( + parts=[types.Part(function_call=original_function_call)] + ), + ) + ) + + tool_confirmation = ToolConfirmation(confirmed=False, hint="test hint") + tool_confirmation_args = { + "originalFunctionCall": original_function_call.model_dump( + exclude_none=True, by_alias=True + ), + "toolConfirmation": tool_confirmation.model_dump( + by_alias=True, exclude_none=True + ), + } + + # Event with the request for confirmation (in child branch) + invocation_context.session.events.append( + Event( + author=agent.name, + branch="child_branch", + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall( + name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, + args=tool_confirmation_args, + id=MOCK_CONFIRMATION_FUNCTION_CALL_ID, + ) + ) + ] + ), + ) + ) + + # Event with the user's confirmation (in default branch, branch=None) + user_confirmation = ToolConfirmation(confirmed=True) + invocation_context.session.events.append( + Event( + author="user", + branch=None, + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, + id=MOCK_CONFIRMATION_FUNCTION_CALL_ID, + response={ + "response": user_confirmation.model_dump_json() + }, + ) + ) + ] + ), + ) + ) + + expected_event = Event( + author="agent", + branch="child_branch", + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=MOCK_TOOL_NAME, + id=MOCK_FUNCTION_CALL_ID, + response={"result": "Mock tool result with test"}, + ) + ) + ] + ), + ) + + # Act & Assert + with patch( + "google.adk.flows.llm_flows.functions.handle_function_call_list_async" + ) as mock_handle_function_call_list_async: + mock_handle_function_call_list_async.return_value = expected_event + + events = [] + async for event in request_processor.run_async( + invocation_context, llm_request + ): + events.append(event) + + assert len(events) == 1 + assert events[0] == expected_event + + +@pytest.mark.asyncio +async def test_request_confirmation_processor_dynamic_success(): + """Test successful processing of dynamic tool confirmation (require_confirmation=False).""" + agent = LlmAgent( + name="test_agent", + tools=[FunctionTool(mock_tool, require_confirmation=False)], + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + llm_request = LlmRequest() + + original_function_call = types.FunctionCall( + name=MOCK_TOOL_NAME, args={"param1": "test"}, id=MOCK_FUNCTION_CALL_ID + ) + + # 1. Event with the original tool call + invocation_context.session.events.append( + Event( + author=agent.name, + content=types.Content( + parts=[types.Part(function_call=original_function_call)] + ), + ) + ) + + # 2. Event with the tool's response requesting confirmation dynamically. + # This event needs to have actions.requested_tool_confirmations. + tool_confirmation_request = ToolConfirmation( + confirmed=False, hint="dynamic hint" + ) + original_response_event = Event( + author="user", + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=MOCK_TOOL_NAME, + id=MOCK_FUNCTION_CALL_ID, + response={"status": "waiting_for_confirm"}, + ) + ) + ] + ), + actions=EventActions( + requested_tool_confirmations={ + MOCK_FUNCTION_CALL_ID: tool_confirmation_request + } + ), + ) + invocation_context.session.events.append(original_response_event) + + # 3. Confirmation request event from the agent to the client. + tool_confirmation_args = { + "originalFunctionCall": original_function_call.model_dump( + exclude_none=True, by_alias=True + ), + "toolConfirmation": tool_confirmation_request.model_dump( + by_alias=True, exclude_none=True + ), + } + invocation_context.session.events.append( + Event( + author=agent.name, + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall( + name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, + args=tool_confirmation_args, + id=MOCK_CONFIRMATION_FUNCTION_CALL_ID, + ) + ) + ] + ), + ) + ) + + # 4. Event with the user's confirmation response. + user_confirmation = ToolConfirmation(confirmed=True) + invocation_context.session.events.append( + Event( + author="user", + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, + id=MOCK_CONFIRMATION_FUNCTION_CALL_ID, + response={ + "response": user_confirmation.model_dump_json() + }, + ) + ) + ] + ), + ) + ) + + expected_event = Event( + author="agent", + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=MOCK_TOOL_NAME, + id=MOCK_FUNCTION_CALL_ID, + response={"result": "Mock tool result with test"}, + ) + ) + ] + ), + ) + + with patch( + "google.adk.flows.llm_flows.functions.handle_function_call_list_async" + ) as mock_handle_function_call_list_async: + mock_handle_function_call_list_async.return_value = expected_event + + events = [] + async for event in request_processor.run_async( + invocation_context, llm_request + ): + events.append(event) + + assert len(events) == 1 + assert events[0] == expected_event + + mock_handle_function_call_list_async.assert_called_once() + args, _ = mock_handle_function_call_list_async.call_args + + assert list(args[1]) == [original_function_call] # function_calls + assert args[3] == {MOCK_FUNCTION_CALL_ID} # tools_to_confirm + assert ( + args[4][MOCK_FUNCTION_CALL_ID] == user_confirmation + ) # tool_confirmation_dict + + +@pytest.mark.parametrize( + "tools, original_args, confirmation_args, expected_exception_match", + [ + ( + [], + {"param1": "test"}, + {"param1": "test"}, + "is not registered", + ), + ( + [FunctionTool(mock_tool, require_confirmation=False)], + {"param1": "test"}, + {"param1": "test"}, + "does not require confirmation", + ), + ( + [FunctionTool(mock_tool, require_confirmation=True)], + {"param1": "test"}, + {"param1": "tampered"}, + "arguments mismatch", + ), + ], +) +@pytest.mark.asyncio +async def test_request_confirmation_processor_rejections( + tools, original_args, confirmation_args, expected_exception_match +): + """Test various validation rejections in request confirmation processor.""" + agent = LlmAgent(name="test_agent", tools=tools) + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + llm_request = LlmRequest() + + original_function_call = types.FunctionCall( + name=MOCK_TOOL_NAME, args=original_args, id=MOCK_FUNCTION_CALL_ID + ) + + # 1. Event with the original tool call + invocation_context.session.events.append( + Event( + author=agent.name, + content=types.Content( + parts=[types.Part(function_call=original_function_call)] + ), + ) + ) + + # 2. Confirmation request event from the agent to the client. + confirmation_function_call = types.FunctionCall( + name=MOCK_TOOL_NAME, args=confirmation_args, id=MOCK_FUNCTION_CALL_ID + ) + tool_confirmation = ToolConfirmation(confirmed=False, hint="test hint") + tool_confirmation_args = { + "originalFunctionCall": confirmation_function_call.model_dump( + exclude_none=True, by_alias=True + ), + "toolConfirmation": tool_confirmation.model_dump( + by_alias=True, exclude_none=True + ), + } + + invocation_context.session.events.append( + Event( + author=agent.name, + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall( + name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, + args=tool_confirmation_args, + id=MOCK_CONFIRMATION_FUNCTION_CALL_ID, + ) + ) + ] + ), + ) + ) + + # 3. Event with the user's confirmation response. + user_confirmation = ToolConfirmation(confirmed=True) + invocation_context.session.events.append( + Event( + author="user", + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, + id=MOCK_CONFIRMATION_FUNCTION_CALL_ID, + response={ + "response": user_confirmation.model_dump_json() + }, + ) + ) + ] + ), + ) + ) + + with pytest.raises(ValueError, match=expected_exception_match): + async for _ in request_processor.run_async(invocation_context, llm_request): + pass + + +def _build_consumed_dynamic_confirmation_events( + agent_name: str, +) -> list[Event]: + """Builds a session where a dynamic confirmation was already acted on. + + Reproduces the state the processor sees on the *second* LLM step of a turn: + a tool was gated at runtime by a policy plugin, the user approved, the + processor re-executed the tool, and the model then made one more tool call — + which sends the flow through preprocessing again while the approval is still + the last user event. + + Args: + agent_name: Author to use for the agent-authored events. + + Returns: + The session events, in order. + """ + original_function_call = types.FunctionCall( + name=MOCK_TOOL_NAME, args={"param1": "test"}, id=MOCK_FUNCTION_CALL_ID + ) + tool_confirmation_request = ToolConfirmation( + confirmed=False, hint="dynamic hint" + ) + return [ + # 1. The model calls the tool. + Event( + author=agent_name, + content=types.Content( + parts=[types.Part(function_call=original_function_call)] + ), + ), + # 2. The tool is gated at runtime and requests confirmation. + Event( + author=agent_name, + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=MOCK_TOOL_NAME, + id=MOCK_FUNCTION_CALL_ID, + response={"status": "waiting_for_confirm"}, + ) + ) + ] + ), + actions=EventActions( + requested_tool_confirmations={ + MOCK_FUNCTION_CALL_ID: tool_confirmation_request + } + ), + ), + # 3. ADK asks the client to confirm. + Event( + author=agent_name, + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall( + name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, + id=MOCK_CONFIRMATION_FUNCTION_CALL_ID, + args={ + "originalFunctionCall": ( + original_function_call.model_dump( + exclude_none=True, by_alias=True + ) + ), + "toolConfirmation": ( + tool_confirmation_request.model_dump( + by_alias=True, exclude_none=True + ) + ), + }, + ) + ) + ] + ), + ), + # 4. The user approves. + Event( + author="user", + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=functions.REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, + id=MOCK_CONFIRMATION_FUNCTION_CALL_ID, + response={ + "response": ( + ToolConfirmation( + confirmed=True + ).model_dump_json() + ) + }, + ) + ) + ] + ), + ), + # 5. The processor re-executed the tool. Note this response carries no + # `requested_tool_confirmations`. + Event( + author=agent_name, + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=MOCK_TOOL_NAME, + id=MOCK_FUNCTION_CALL_ID, + response={"result": "Mock tool result with test"}, + ) + ) + ] + ), + ), + # 6. The model makes one more tool call, forcing another LLM step. + Event( + author=agent_name, + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall( + name="another_tool", id="another_function_call_id" + ) + ) + ] + ), + ), + ] + + +@pytest.mark.asyncio +async def test_request_confirmation_processor_consumed_dynamic_confirmation_is_noop(): + """A dynamic confirmation already acted on must not be processed again.""" + agent = LlmAgent( + name="test_agent", + tools=[FunctionTool(mock_tool, require_confirmation=False)], + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + invocation_context.session.events.extend( + _build_consumed_dynamic_confirmation_events(agent.name) + ) + + events = [] + async for event in request_processor.run_async( + invocation_context, LlmRequest() + ): + events.append(event) + + assert not events + + +@pytest.mark.asyncio +async def test_request_confirmation_processor_consumed_confirmation_ignores_deregistered_tool(): + """A consumed confirmation must not fail when the toolset has moved on. + + Toolsets are resolved per step, so a tool present when the user approved can + be gone by the next step (e.g. a disconnected MCP toolset). That must not + abort the invocation. + """ + agent = LlmAgent(name="test_agent", tools=[]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + invocation_context.session.events.extend( + _build_consumed_dynamic_confirmation_events(agent.name) + ) + + events = [] + async for event in request_processor.run_async( + invocation_context, LlmRequest() + ): + events.append(event) + + assert not events + + +@pytest.mark.asyncio +async def test_request_confirmation_processor_consumed_confirmation_skips_revalidation(): + """A consumed confirmation must not re-invoke `check_require_confirmation`. + + It is a user-overridable hook that may be expensive or have side effects, so + it must not run once per LLM step for the rest of the turn. + """ + check_require_confirmation_calls = [] + + class _CountingFunctionTool(FunctionTool): + + async def check_require_confirmation(self, args, tool_context) -> bool: + check_require_confirmation_calls.append(args) + return False + + agent = LlmAgent( + name="test_agent", + tools=[_CountingFunctionTool(mock_tool, require_confirmation=False)], + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + invocation_context.session.events.extend( + _build_consumed_dynamic_confirmation_events(agent.name) + ) + + async for _ in request_processor.run_async(invocation_context, LlmRequest()): + pass + + assert not check_require_confirmation_calls + + +@pytest.mark.asyncio +async def test_resolve_confirmation_targets_after_reexecution(): + """The re-execution response must not shadow the original confirmation request. + + `_resolve_confirmation_targets` is also called directly by out-of-tree + callers that have no dedup of their own, so it has to stay correct once the + confirmed tool has produced a second response under the same call ID. + """ + tool = FunctionTool(mock_tool, require_confirmation=False) + agent = LlmAgent(name="test_agent", tools=[tool]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + invocation_context.session.events.extend( + _build_consumed_dynamic_confirmation_events(agent.name) + ) + + tool_confirmation_dict, original_fcs_dict = ( + await _resolve_confirmation_targets( + invocation_context, + invocation_context.session.events, + {MOCK_CONFIRMATION_FUNCTION_CALL_ID}, + { + MOCK_CONFIRMATION_FUNCTION_CALL_ID: ToolConfirmation( + confirmed=True + ) + }, + {MOCK_TOOL_NAME: tool}, + ) + ) + + assert set(tool_confirmation_dict) == {MOCK_FUNCTION_CALL_ID} + assert set(original_fcs_dict) == {MOCK_FUNCTION_CALL_ID} From 7801db0bd4577673e3a03f91e5cc0479ca8b6fa8 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Thu, 30 Jul 2026 10:41:19 -0700 Subject: [PATCH 2/2] fix: Stop re-validating already-consumed tool confirmations Co-authored-by: Xuan Yang PiperOrigin-RevId: 956611754 --- .../flows/llm_flows/request_confirmation.py | 161 +++++++++++++----- 1 file changed, 119 insertions(+), 42 deletions(-) diff --git a/src/google/adk/flows/llm_flows/request_confirmation.py b/src/google/adk/flows/llm_flows/request_confirmation.py index 2de3a4a90e1..895a609948e 100644 --- a/src/google/adk/flows/llm_flows/request_confirmation.py +++ b/src/google/adk/flows/llm_flows/request_confirmation.py @@ -16,6 +16,7 @@ import logging from typing import Any from typing import AsyncGenerator +from typing import Optional from typing import TYPE_CHECKING from google.genai import types @@ -44,6 +45,31 @@ def _parse_tool_confirmation(response: dict[str, Any]) -> ToolConfirmation: return ToolConfirmation.from_response_dict(response) +def _get_original_function_call_args( + function_call: types.FunctionCall, +) -> Optional[dict[str, Any]]: + """Returns the raw ``originalFunctionCall`` payload of a confirmation call. + + Both the dedup pre-pass and ``_resolve_confirmation_targets`` read the + original function call out of an ``adk_request_confirmation`` call's args. + They must agree on what counts as a well-formed payload, otherwise a + confirmation could be skipped by one and processed by the other. + + Args: + function_call: An ``adk_request_confirmation`` function call. + + Returns: + The ``originalFunctionCall`` dict, or ``None`` if it is absent or malformed. + """ + args = function_call.args + if not args: + return None + original_function_call = args.get("originalFunctionCall") + if not isinstance(original_function_call, dict): + return None + return original_function_call + + async def _resolve_confirmation_targets( invocation_context: InvocationContext, events: list[Event], @@ -83,9 +109,19 @@ async def _resolve_confirmation_targets( for fc in ev.get_function_calls() if fc.id and fc.name != REQUEST_CONFIRMATION_FUNCTION_CALL_NAME } - history_fr_events = { - fr.id: ev for ev in events for fr in ev.get_function_responses() if fr.id - } + # IDs of function calls for which a tool dynamically requested confirmation. + # This accumulates over ALL events rather than keeping one event per ID: once + # the confirmed tool is re-executed it emits a second function response with + # the same ID and no `requested_tool_confirmations`, which would otherwise + # shadow the original request. + dynamically_requested_fc_ids: set[str] = set() + for ev in events: + requested_tool_confirmations = ev.actions.requested_tool_confirmations or {} + if not requested_tool_confirmations: + continue + for fr in ev.get_function_responses(): + if fr.id and fr.id in requested_tool_confirmations: + dynamically_requested_fc_ids.add(fr.id) for event in events: event_function_calls = event.get_function_calls() @@ -96,12 +132,12 @@ async def _resolve_confirmation_targets( if not function_call.id or function_call.id not in confirmation_fc_ids: continue - args = function_call.args - if not args or "originalFunctionCall" not in args: - continue - original_function_call = types.FunctionCall( - **args["originalFunctionCall"] + original_function_call_args = _get_original_function_call_args( + function_call ) + if original_function_call_args is None: + continue + original_function_call = types.FunctionCall(**original_function_call_args) if not original_function_call.id: raise ValueError("Original function call ID is missing.") tool_name = original_function_call.name @@ -140,20 +176,9 @@ async def _resolve_confirmation_targets( original_function_call.args or {}, temp_tool_context ) - requested_in_history = False - if not requires_confirmation: - # Search the history for the response event of the original tool call - original_response_event = history_fr_events.get( - original_function_call.id - ) - if ( - original_response_event - and original_response_event.actions.requested_tool_confirmations - ): - requested_in_history = ( - original_function_call.id - in original_response_event.actions.requested_tool_confirmations - ) + requested_in_history = ( + original_function_call.id in dynamically_requested_fc_ids + ) if not requires_confirmation and not requested_in_history: raise ValueError( @@ -185,6 +210,44 @@ async def _resolve_confirmation_targets( return tool_confirmation_dict, original_fcs_dict +def _map_confirmation_to_original_fc_ids( + events: list[Event], + confirmation_fc_ids: set[str], +) -> dict[str, str]: + """Maps each confirmation function call ID to its original function call ID. + + This is a cheap, validation-free pre-pass so that already-consumed + confirmations can be dropped *before* the expensive and strict + ``_resolve_confirmation_targets``. + + Args: + events: Session events to scan. + confirmation_fc_ids: IDs of ``adk_request_confirmation`` function calls. + + Returns: + Mapping of confirmation FC ID -> original FC ID. Confirmations whose + original function call cannot be determined are omitted. + """ + mapping: dict[str, str] = {} + for event in events: + for function_call in event.get_function_calls(): + if not function_call.id or function_call.id not in confirmation_fc_ids: + continue + original_function_call_args = _get_original_function_call_args( + function_call + ) + # Mirror the `is None` check in `_resolve_confirmation_targets`: an empty + # payload must reach the strict validation there and be rejected, not be + # quietly dropped here (dropping it would skip the dedup and produce a + # confusing downstream error instead). + if original_function_call_args is None: + continue + original_fc_id = original_function_call_args.get("id") + if original_fc_id: + mapping[function_call.id] = original_fc_id + return mapping + + class _RequestConfirmationLlmRequestProcessor(BaseLlmRequestProcessor): """Handles tool confirmation information to build the LLM request.""" @@ -225,7 +288,39 @@ async def run_async( if not confirmations_by_fc_id: return - # Resolve all canonical tools and build tools_dict + # Step 2: Drop confirmations that have already been consumed. + # + # This must happen BEFORE resolving targets. The processor re-runs on every + # LLM step of the invocation, and the approval stays the last user event for + # the rest of the turn, so a confirmation the previous step already acted on + # is seen again here. Re-validating consumed state is not just wasted work: + # the session and the toolset have moved on since the approval, so the + # strict checks in `_resolve_confirmation_targets` can now legitimately fail + # and abort the invocation. + confirmation_to_original_fc_id = _map_confirmation_to_original_fc_ids( + events, set(confirmations_by_fc_id.keys()) + ) + responded_fc_ids: set[str] = set() + for event in reversed(events): + if event.author == "user": + break + for function_response in event.get_function_responses(): + if function_response.id: + responded_fc_ids.add(function_response.id) + + confirmations_by_fc_id = { + confirmation_fc_id: confirmation + for confirmation_fc_id, confirmation in confirmations_by_fc_id.items() + if confirmation_to_original_fc_id.get(confirmation_fc_id) + not in responded_fc_ids + } + + if not confirmations_by_fc_id: + return + + # Resolve all canonical tools and build tools_dict. Deliberately after the + # dedup above so a consumed confirmation does not force a toolset + # resolution, which can be a remote call for e.g. MCP toolsets. tools_dict = {} if agent is not None and hasattr(agent, "canonical_tools"): tools_dict = { @@ -235,7 +330,7 @@ async def run_async( ) } - # Step 2: Resolve confirmation targets using extracted helper. + # Step 3: Resolve confirmation targets using extracted helper. confirmation_fc_ids = set(confirmations_by_fc_id.keys()) tools_to_resume_with_confirmation, tools_to_resume_with_args = ( await _resolve_confirmation_targets( @@ -247,24 +342,6 @@ async def run_async( ) ) - if not tools_to_resume_with_confirmation: - return - - # Step 3: Remove tools that have already been confirmed (dedup). - for event in reversed(events): - if event.author == "user": - break - fr_list = event.get_function_responses() - if not fr_list: - continue - - for function_response in fr_list: - if function_response.id in tools_to_resume_with_confirmation: - tools_to_resume_with_confirmation.pop(function_response.id) - tools_to_resume_with_args.pop(function_response.id) - if not tools_to_resume_with_confirmation: - break - if not tools_to_resume_with_confirmation: return