Skip to content

Commit 1d0daeb

Browse files
GWealecopybara-github
authored andcommitted
fix(flows): hold a tool call that requires confirmation in the framework
Co-authored-by: George Weale <[email protected]> PiperOrigin-RevId: 983922758
1 parent 5bc9e8c commit 1d0daeb

2 files changed

Lines changed: 141 additions & 2 deletions

File tree

‎src/google/adk/flows/llm_flows/tools/_caller.py‎

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -711,6 +711,46 @@ async def _prepare_single(
711711
)
712712

713713

714+
async def _apply_confirmation_gate(
715+
tool: BaseTool,
716+
function_args: dict[str, Any],
717+
tool_context: ToolContext,
718+
) -> dict[str, str] | None:
719+
"""Answers a call whose tool is waiting on a human, instead of running it.
720+
721+
Args:
722+
tool: The tool the call names.
723+
function_args: The arguments the call carries.
724+
tool_context: The context the call will run in.
725+
726+
Returns:
727+
The response to answer the call with, or None if the call may proceed.
728+
"""
729+
if not await tool.check_require_confirmation(function_args, tool_context):
730+
return None
731+
732+
confirmation = tool_context.tool_confirmation
733+
if confirmation is None:
734+
tool_context.request_confirmation(
735+
hint=(
736+
f'Please approve or reject the tool call {tool.name}() by'
737+
' responding with a FunctionResponse with an expected'
738+
' ToolConfirmation payload.'
739+
),
740+
)
741+
# The pause is not a tool result for the model to summarize; without this
742+
# the flow re-invokes the model, which calls the tool again.
743+
tool_context.actions.skip_summarization = True
744+
return {
745+
'error': (
746+
'This tool call requires confirmation, please approve or reject.'
747+
)
748+
}
749+
if not confirmation.confirmed:
750+
return {'error': 'This tool call is rejected.'}
751+
return None
752+
753+
714754
async def _execute_single_prepared_call(
715755
invocation_context: InvocationContext,
716756
prepared_call: _PreparedFunctionCall,
@@ -793,10 +833,17 @@ async def _run_with_trace() -> Event | None:
793833
)
794834

795835
# Step 3: No before-tool callback answered the call, so proceed calling
796-
# the tool normally.
836+
# the tool normally. A tool that requires confirmation is answered by the
837+
# gate instead, so the gate holds for every tool rather than only the ones
838+
# that check it themselves, and a gate that raises is handled like a tool
839+
# that raises.
797840
if function_response is None:
798841
try:
799-
function_response = await tool_runner()
842+
function_response = await _apply_confirmation_gate(
843+
tool, function_args, tool_context
844+
)
845+
if function_response is None:
846+
function_response = await tool_runner()
800847
except Exception as tool_error:
801848
error_response = await _tool_error_handler.run_on_tool_error_callbacks(
802849
invocation_context=invocation_context,

‎tests/unittests/flows/llm_flows/tools/test_confirmation.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from google.adk.flows.llm_flows.tools._confirmation import _resolve_confirmation_targets
2222
from google.adk.flows.llm_flows.tools._confirmation import request_processor
2323
from google.adk.models.llm_request import LlmRequest
24+
from google.adk.tools.base_tool import BaseTool
2425
from google.adk.tools.function_tool import FunctionTool
2526
from google.adk.tools.tool_confirmation import ToolConfirmation
2627
from google.genai import types
@@ -1319,3 +1320,94 @@ async def test_resolve_confirmation_targets_requires_adk_name():
13191320

13201321
assert set(tool_confirmation_dict) == {"requested_fc_id"}
13211322
assert set(original_fcs_dict) == {"requested_fc_id"}
1323+
1324+
1325+
class _GatedTool(BaseTool):
1326+
"""A tool outside the FunctionTool hierarchy that asks for a human."""
1327+
1328+
def __init__(self):
1329+
super().__init__(name="gated_tool", description="Needs approval.")
1330+
self.runs = 0
1331+
1332+
async def check_require_confirmation(self, args, tool_context) -> bool:
1333+
return True
1334+
1335+
async def run_async(self, *, args, tool_context):
1336+
self.runs += 1
1337+
return {"status": "ran"}
1338+
1339+
1340+
async def _call_gated_tool(
1341+
tool: _GatedTool, tool_confirmation: ToolConfirmation | None = None
1342+
) -> Event:
1343+
"""Calls `tool` through the flow and returns its function response event."""
1344+
agent = LlmAgent(name="test_agent", tools=[tool])
1345+
invocation_context = await testing_utils.create_invocation_context(
1346+
agent=agent
1347+
)
1348+
function_call_event = Event(
1349+
invocation_id=invocation_context.invocation_id,
1350+
author=agent.name,
1351+
content=types.Content(
1352+
parts=[
1353+
types.Part(
1354+
function_call=types.FunctionCall(
1355+
id=MOCK_FUNCTION_CALL_ID, name=tool.name, args={}
1356+
)
1357+
)
1358+
]
1359+
),
1360+
)
1361+
event = await functions.handle_function_calls_async(
1362+
invocation_context,
1363+
function_call_event,
1364+
{tool.name: tool},
1365+
tool_confirmation_dict=(
1366+
{MOCK_FUNCTION_CALL_ID: tool_confirmation}
1367+
if tool_confirmation is not None
1368+
else None
1369+
),
1370+
)
1371+
assert event is not None
1372+
return event
1373+
1374+
1375+
@pytest.mark.asyncio
1376+
async def test_unconfirmed_tool_is_not_run():
1377+
"""A tool requiring confirmation must not run before a human answers."""
1378+
tool = _GatedTool()
1379+
1380+
event = await _call_gated_tool(tool)
1381+
1382+
assert tool.runs == 0
1383+
assert event.get_function_responses()[0].response == {
1384+
"error": "This tool call requires confirmation, please approve or reject."
1385+
}
1386+
assert set(event.actions.requested_tool_confirmations) == {
1387+
MOCK_FUNCTION_CALL_ID
1388+
}
1389+
assert event.actions.skip_summarization
1390+
1391+
1392+
@pytest.mark.asyncio
1393+
async def test_rejected_tool_is_not_run():
1394+
"""A tool the human refused must not run."""
1395+
tool = _GatedTool()
1396+
1397+
event = await _call_gated_tool(tool, ToolConfirmation(confirmed=False))
1398+
1399+
assert tool.runs == 0
1400+
assert event.get_function_responses()[0].response == {
1401+
"error": "This tool call is rejected."
1402+
}
1403+
1404+
1405+
@pytest.mark.asyncio
1406+
async def test_confirmed_tool_is_run():
1407+
"""A tool the human approved runs and returns its own result."""
1408+
tool = _GatedTool()
1409+
1410+
event = await _call_gated_tool(tool, ToolConfirmation(confirmed=True))
1411+
1412+
assert tool.runs == 1
1413+
assert event.get_function_responses()[0].response == {"status": "ran"}

0 commit comments

Comments
 (0)