From 3557a4139707950a647849aa3c9bce44bcae5533 Mon Sep 17 00:00:00 2001 From: AnvitDevadiga Date: Thu, 17 Sep 2026 08:49:51 +0530 Subject: [PATCH 1/4] fix: enforce denied tool confirmations centrally --- .../flows/llm_flows/tools/_confirmation.py | 52 +++++++++++++++---- .../llm_flows/tools/test_confirmation.py | 32 +++--------- .../runners/test_run_tool_confirmation.py | 14 +++-- 3 files changed, 60 insertions(+), 38 deletions(-) diff --git a/src/google/adk/flows/llm_flows/tools/_confirmation.py b/src/google/adk/flows/llm_flows/tools/_confirmation.py index 80a98b87258..133d3b269b3 100644 --- a/src/google/adk/flows/llm_flows/tools/_confirmation.py +++ b/src/google/adk/flows/llm_flows/tools/_confirmation.py @@ -354,17 +354,51 @@ async def run_async( if not tools_to_resume_with_confirmation: return - # Step 4: Re-execute the confirmed tools. + # Step 4: Re-execute only confirmed tools. Denials are handled here at the + # framework boundary so custom BaseTool implementations cannot accidentally + # perform a side effect after the user declines confirmation. from .. import functions - if function_response_event := await functions.handle_function_call_list_async( - invocation_context, - 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 + denied_parts: list[types.Part] = [] + confirmed_args: list[types.FunctionCall] = [] + confirmed_ids: set[str] = set() + confirmed_tools: dict[str, ToolConfirmation] = {} + for function_call_id, function_call in tools_to_resume_with_args.items(): + confirmation = tools_to_resume_with_confirmation[function_call_id] + if confirmation.confirmed: + confirmed_args.append(function_call) + confirmed_ids.add(function_call_id) + confirmed_tools[function_call_id] = confirmation + else: + denied_parts.append( + types.Part( + function_response=types.FunctionResponse( + name=function_call.name, + id=function_call_id, + response={"error": "Tool execution not confirmed"}, + ) + ) + ) + + if confirmed_args: + if function_response_event := await functions.handle_function_call_list_async( + invocation_context, + confirmed_args, + tools_dict, + confirmed_ids, + confirmed_tools, + ): + denied_parts.extend(function_response_event.content.parts) + yield function_response_event.model_copy(update={ + "content": types.Content(parts=denied_parts) + }) + return + + if denied_parts: + yield Event( + author=invocation_context.agent.name if invocation_context.agent else "agent", + content=types.Content(parts=denied_parts), + ) return diff --git a/tests/unittests/flows/llm_flows/tools/test_confirmation.py b/tests/unittests/flows/llm_flows/tools/test_confirmation.py index 0acd73915bf..65ae514ca78 100644 --- a/tests/unittests/flows/llm_flows/tools/test_confirmation.py +++ b/tests/unittests/flows/llm_flows/tools/test_confirmation.py @@ -301,21 +301,6 @@ async def test_request_confirmation_processor_tool_not_confirmed(): 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 = Event( - author="agent", - content=types.Content( - parts=[ - types.Part( - function_response=types.FunctionResponse( - name=MOCK_TOOL_NAME, - id=MOCK_FUNCTION_CALL_ID, - response={"error": "Tool execution not confirmed"}, - ) - ) - ] - ), - ) - events = [] async for event in request_processor.run_async( invocation_context, llm_request @@ -323,11 +308,10 @@ async def test_request_confirmation_processor_tool_not_confirmed(): events.append(event) assert len(events) == 1 - mock_handle_function_call_list_async.assert_called_once() - args, _ = mock_handle_function_call_list_async.call_args - assert ( - args[4][MOCK_FUNCTION_CALL_ID] == user_confirmation - ) # tool_confirmation_dict + mock_handle_function_call_list_async.assert_not_called() + assert events[0].content.parts[0].function_response.response == { + "error": "Tool execution not confirmed" + } TRANSFER_TOOL_NAME = "transfer_to_agent" @@ -509,10 +493,10 @@ async def test_request_confirmation_transfer_to_agent_rejected(): events.append(event) assert len(events) == 1 - mock_handle.assert_called_once() - args, _ = mock_handle.call_args - tools_dict = args[2] - assert TRANSFER_TOOL_NAME in tools_dict + mock_handle.assert_not_called() + assert events[0].content.parts[0].function_response.response == { + "error": "Tool execution not confirmed" + } @pytest.mark.asyncio diff --git a/tests/unittests/runners/test_run_tool_confirmation.py b/tests/unittests/runners/test_run_tool_confirmation.py index 005fb98b5af..7dc383bce33 100644 --- a/tests/unittests/runners/test_run_tool_confirmation.py +++ b/tests/unittests/runners/test_run_tool_confirmation.py @@ -211,7 +211,7 @@ async def test_confirmation_flow( name=tools[0].name, response={"result": f"confirmed={tool_call_confirmed}"} if tool_call_confirmed - else {"error": "This tool call is rejected."}, + else {"error": "Tool execution not confirmed"}, ) ), ), @@ -352,10 +352,14 @@ async def test_confirmation_flow( ) events = await runner.run_async(user_confirmation) - expected_response = { - "result": f"confirmed={tool_call_confirmed}", - "custom_payload": custom_payload, - } + expected_response = ( + { + "result": "confirmed=True", + "custom_payload": custom_payload, + } + if tool_call_confirmed + else {"error": "Tool execution not confirmed"} + ) expected_parts_final = [ ( agent.name, From 4311da1cffa0062f060602f2b0a9f4d9ed8e6186 Mon Sep 17 00:00:00 2001 From: AnvitDevadiga Date: Wed, 23 Sep 2026 10:24:57 +0530 Subject: [PATCH 2/4] fix: prevent replayed tool confirmations --- src/google/adk/agents/invocation_context.py | 13 +++++++++++ .../flows/llm_flows/tools/_confirmation.py | 23 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index b332af8b2e6..61a26577730 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -177,6 +177,19 @@ class InvocationContext(BaseModel): None for non-workflow agents. """ + _consumed_tool_confirmation_ids: set[str] = PrivateAttr(default_factory=set) + _tool_confirmation_consume_lock: asyncio.Lock = PrivateAttr( + default_factory=asyncio.Lock + ) + + async def _consume_tool_confirmation(self, function_call_id: str) -> bool: + """Atomically claim a confirmation so it can only resume a tool once.""" + async with self._tool_confirmation_consume_lock: + if function_call_id in self._consumed_tool_confirmation_ids: + return False + self._consumed_tool_confirmation_ids.add(function_call_id) + return True + agent_states: dict[str, dict[str, Any]] = Field(default_factory=dict) """The state of the agent for this invocation.""" diff --git a/src/google/adk/flows/llm_flows/tools/_confirmation.py b/src/google/adk/flows/llm_flows/tools/_confirmation.py index 133d3b269b3..5c2e05e74f5 100644 --- a/src/google/adk/flows/llm_flows/tools/_confirmation.py +++ b/src/google/adk/flows/llm_flows/tools/_confirmation.py @@ -354,6 +354,29 @@ async def run_async( if not tools_to_resume_with_confirmation: return + # Claim confirmations atomically before dispatch. Event-history dedup is + # useful for normal reruns, but cannot prevent two concurrent resume calls + # from both observing the same still-unresponded confirmation. + claimed_ids = { + function_call_id + for function_call_id in tools_to_resume_with_confirmation + if await invocation_context._consume_tool_confirmation(function_call_id) + } + if not claimed_ids: + return + tools_to_resume_with_confirmation = { + function_call_id: confirmation + for function_call_id, confirmation in ( + tools_to_resume_with_confirmation.items() + ) + if function_call_id in claimed_ids + } + tools_to_resume_with_args = { + function_call_id: function_call + for function_call_id, function_call in tools_to_resume_with_args.items() + if function_call_id in claimed_ids + } + # Step 4: Re-execute only confirmed tools. Denials are handled here at the # framework boundary so custom BaseTool implementations cannot accidentally # perform a side effect after the user declines confirmation. From 90bec5e99eca40623fc2c692542d0322aaa44bc5 Mon Sep 17 00:00:00 2001 From: AnvitDevadiga Date: Wed, 23 Sep 2026 18:34:12 +0530 Subject: [PATCH 3/4] test: cover concurrent tool confirmation claims --- .../flows/llm_flows/tools/test_confirmation.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/unittests/flows/llm_flows/tools/test_confirmation.py b/tests/unittests/flows/llm_flows/tools/test_confirmation.py index d34541f35f7..ff82d4aab19 100644 --- a/tests/unittests/flows/llm_flows/tools/test_confirmation.py +++ b/tests/unittests/flows/llm_flows/tools/test_confirmation.py @@ -13,6 +13,7 @@ # limitations under the License. from unittest.mock import create_autospec +import asyncio from unittest.mock import patch from google.adk.agents.llm_agent import LlmAgent @@ -40,6 +41,22 @@ def mock_tool(param1: str): return f"Mock tool result with {param1}" +@pytest.mark.asyncio +async def test_tool_confirmation_claim_is_atomic(): + """Only one concurrent resume may claim a function call confirmation.""" + agent = LlmAgent(name="test_agent") + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + + claims = await asyncio.gather( + invocation_context._consume_tool_confirmation(MOCK_FUNCTION_CALL_ID), + invocation_context._consume_tool_confirmation(MOCK_FUNCTION_CALL_ID), + ) + + assert sorted(claims) == [False, True] + + @pytest.mark.asyncio async def test_request_confirmation_processor_no_events(): """Test that the processor returns None when there are no events.""" From 4beee3a49f21564e8153c1acc78fe541bd9352ad Mon Sep 17 00:00:00 2001 From: AnvitDevadiga Date: Fri, 25 Sep 2026 10:49:32 +0530 Subject: [PATCH 4/4] fix: make confirmation replay protection durable --- .../flows/llm_flows/tools/_confirmation.py | 68 ++++--------------- .../llm_flows/tools/test_confirmation.py | 34 +++++++--- .../runners/test_run_tool_confirmation.py | 14 ++-- 3 files changed, 42 insertions(+), 74 deletions(-) diff --git a/src/google/adk/flows/llm_flows/tools/_confirmation.py b/src/google/adk/flows/llm_flows/tools/_confirmation.py index 91ef026b870..4c56b604adb 100644 --- a/src/google/adk/flows/llm_flows/tools/_confirmation.py +++ b/src/google/adk/flows/llm_flows/tools/_confirmation.py @@ -293,20 +293,13 @@ async def run_async( # 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. + # This must happen BEFORE resolving targets. Persisted event history is the + # durable source of truth when a later run rebuilds InvocationContext. 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 event in events: for function_response in event.get_function_responses(): if function_response.id: responded_fc_ids.add(function_response.id) @@ -356,9 +349,6 @@ async def run_async( if not tools_to_resume_with_confirmation: return - # Claim confirmations atomically before dispatch. Event-history dedup is - # useful for normal reruns, but cannot prevent two concurrent resume calls - # from both observing the same still-unresponded confirmation. claimed_ids = { function_call_id for function_call_id in tools_to_resume_with_confirmation @@ -379,51 +369,17 @@ async def run_async( if function_call_id in claimed_ids } - # Step 4: Re-execute only confirmed tools. Denials are handled here at the - # framework boundary so custom BaseTool implementations cannot accidentally - # perform a side effect after the user declines confirmation. + # Step 4: Re-execute the confirmed tools. from .. import functions - denied_parts: list[types.Part] = [] - confirmed_args: list[types.FunctionCall] = [] - confirmed_ids: set[str] = set() - confirmed_tools: dict[str, ToolConfirmation] = {} - for function_call_id, function_call in tools_to_resume_with_args.items(): - confirmation = tools_to_resume_with_confirmation[function_call_id] - if confirmation.confirmed: - confirmed_args.append(function_call) - confirmed_ids.add(function_call_id) - confirmed_tools[function_call_id] = confirmation - else: - denied_parts.append( - types.Part( - function_response=types.FunctionResponse( - name=function_call.name, - id=function_call_id, - response={"error": "Tool execution not confirmed"}, - ) - ) - ) - - if confirmed_args: - if function_response_event := await functions.handle_function_call_list_async( - invocation_context, - confirmed_args, - tools_dict, - confirmed_ids, - confirmed_tools, - ): - denied_parts.extend(function_response_event.content.parts) - yield function_response_event.model_copy(update={ - "content": types.Content(parts=denied_parts) - }) - return - - if denied_parts: - yield Event( - author=invocation_context.agent.name if invocation_context.agent else "agent", - content=types.Content(parts=denied_parts), - ) + if function_response_event := await functions.handle_function_call_list_async( + invocation_context, + 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 return diff --git a/tests/unittests/flows/llm_flows/tools/test_confirmation.py b/tests/unittests/flows/llm_flows/tools/test_confirmation.py index ff82d4aab19..b4c15c09844 100644 --- a/tests/unittests/flows/llm_flows/tools/test_confirmation.py +++ b/tests/unittests/flows/llm_flows/tools/test_confirmation.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import create_autospec import asyncio +from unittest.mock import create_autospec from unittest.mock import patch from google.adk.agents.llm_agent import LlmAgent @@ -320,6 +320,21 @@ async def test_request_confirmation_processor_tool_not_confirmed(): 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 = Event( + author="agent", + content=types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=MOCK_TOOL_NAME, + id=MOCK_FUNCTION_CALL_ID, + response={"error": "Tool execution not confirmed"}, + ) + ) + ] + ), + ) + events = [] async for event in request_processor.run_async( invocation_context, llm_request @@ -327,10 +342,11 @@ async def test_request_confirmation_processor_tool_not_confirmed(): events.append(event) assert len(events) == 1 - mock_handle_function_call_list_async.assert_not_called() - assert events[0].content.parts[0].function_response.response == { - "error": "Tool execution not confirmed" - } + mock_handle_function_call_list_async.assert_called_once() + args, _ = mock_handle_function_call_list_async.call_args + assert ( + args[4][MOCK_FUNCTION_CALL_ID] == user_confirmation + ) # tool_confirmation_dict TRANSFER_TOOL_NAME = "transfer_to_agent" @@ -512,10 +528,10 @@ async def test_request_confirmation_transfer_to_agent_rejected(): events.append(event) assert len(events) == 1 - mock_handle.assert_not_called() - assert events[0].content.parts[0].function_response.response == { - "error": "Tool execution not confirmed" - } + mock_handle.assert_called_once() + args, _ = mock_handle.call_args + tools_dict = args[2] + assert TRANSFER_TOOL_NAME in tools_dict @pytest.mark.asyncio diff --git a/tests/unittests/runners/test_run_tool_confirmation.py b/tests/unittests/runners/test_run_tool_confirmation.py index 7dc383bce33..005fb98b5af 100644 --- a/tests/unittests/runners/test_run_tool_confirmation.py +++ b/tests/unittests/runners/test_run_tool_confirmation.py @@ -211,7 +211,7 @@ async def test_confirmation_flow( name=tools[0].name, response={"result": f"confirmed={tool_call_confirmed}"} if tool_call_confirmed - else {"error": "Tool execution not confirmed"}, + else {"error": "This tool call is rejected."}, ) ), ), @@ -352,14 +352,10 @@ async def test_confirmation_flow( ) events = await runner.run_async(user_confirmation) - expected_response = ( - { - "result": "confirmed=True", - "custom_payload": custom_payload, - } - if tool_call_confirmed - else {"error": "Tool execution not confirmed"} - ) + expected_response = { + "result": f"confirmed={tool_call_confirmed}", + "custom_payload": custom_payload, + } expected_parts_final = [ ( agent.name,