diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 7ec28187998..3037b7f2dfa 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -203,6 +203,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 1debe440044..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,6 +349,26 @@ async def run_async( if not tools_to_resume_with_confirmation: return + 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 the confirmed tools. from .. import functions diff --git a/tests/unittests/flows/llm_flows/tools/test_confirmation.py b/tests/unittests/flows/llm_flows/tools/test_confirmation.py index 16cab223363..b4c15c09844 100644 --- a/tests/unittests/flows/llm_flows/tools/test_confirmation.py +++ b/tests/unittests/flows/llm_flows/tools/test_confirmation.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio from unittest.mock import create_autospec from unittest.mock import patch @@ -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."""