|
21 | 21 | from google.adk.flows.llm_flows.tools._confirmation import _resolve_confirmation_targets |
22 | 22 | from google.adk.flows.llm_flows.tools._confirmation import request_processor |
23 | 23 | from google.adk.models.llm_request import LlmRequest |
| 24 | +from google.adk.tools.base_tool import BaseTool |
24 | 25 | from google.adk.tools.function_tool import FunctionTool |
25 | 26 | from google.adk.tools.tool_confirmation import ToolConfirmation |
26 | 27 | from google.genai import types |
@@ -1319,3 +1320,94 @@ async def test_resolve_confirmation_targets_requires_adk_name(): |
1319 | 1320 |
|
1320 | 1321 | assert set(tool_confirmation_dict) == {"requested_fc_id"} |
1321 | 1322 | 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