diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 7949d2c997f..e8803ee498a 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -26,7 +26,7 @@ import logging import threading from typing import Any -from typing import AsyncGenerator +from typing import Callable from typing import cast from typing import Dict from typing import Optional @@ -121,10 +121,10 @@ async def _call_tool_in_thread_pool( ) -> Any: """Runs a tool in a thread pool to avoid blocking the event loop. - For sync tools, this runs the tool's function directly in a background thread. - For async tools, this creates a new event loop in the background thread and - runs the async function there. This helps catch blocking I/O (like time.sleep, - network calls, file I/O) that was mistakenly used inside async functions. + The complete ``BaseTool.run_async`` contract is preserved. For synchronous + ``FunctionTool`` callables, tool-owned validation, authentication, and + confirmation stay on the caller loop while only synchronous callables enter + the pool. Other tools run their complete async contract in a worker loop. Note: Due to Python's GIL, this does NOT help with pure Python CPU-bound code. Thread pool only helps when the GIL is released (blocking I/O, C extensions). @@ -138,43 +138,39 @@ async def _call_tool_in_thread_pool( Returns: The result of running the tool. """ + from ...tools.function_tool import _use_sync_callable_runner from ...tools.function_tool import FunctionTool - ctx = contextvars.copy_context() loop = asyncio.get_running_loop() executor = _get_tool_thread_pool(max_workers) - if _is_sync_tool(tool): - if isinstance(tool, FunctionTool): - # For sync FunctionTool, call the underlying function directly. - def run_sync_tool(): - args_to_call = tool._preprocess_args(args) - signature = inspect.signature(tool.func) - valid_params = {param for param in signature.parameters} - if tool._context_param_name in valid_params: - args_to_call[tool._context_param_name] = tool_context - args_to_call = { - k: v for k, v in args_to_call.items() if k in valid_params - } - return tool.func(**args_to_call) + if _is_sync_tool(tool) and isinstance(tool, FunctionTool): + + async def run_sync_callable( + target: Callable[..., Any], call_args: dict[str, Any] + ) -> Any: + call_context = contextvars.copy_context() + + def invoke() -> Any: + with _use_sync_callable_runner(None): + return target(**call_args) return await loop.run_in_executor( - executor, lambda: ctx.run(run_sync_tool) + executor, + lambda: call_context.run(invoke), ) - else: - # For async tools, run them in a new event loop in a background thread. - # This helps when async functions contain blocking I/O (common user mistake) - # that would otherwise block the main event loop. - def run_async_tool_in_new_loop(): - # Create a new event loop for this thread - return asyncio.run(tool.run_async(args=args, tool_context=tool_context)) - - return await loop.run_in_executor( - executor, lambda: ctx.run(run_async_tool_in_new_loop) - ) - # Fall back to normal async execution for non-FunctionTool sync tools. - return await tool.run_async(args=args, tool_context=tool_context) + with _use_sync_callable_runner(run_sync_callable): + return await tool.run_async(args=args, tool_context=tool_context) + + ctx = contextvars.copy_context() + + def run_tool_in_new_loop() -> Any: + return asyncio.run(tool.run_async(args=args, tool_context=tool_context)) + + return await loop.run_in_executor( + executor, lambda: ctx.run(run_tool_in_new_loop) + ) def generate_client_function_call_id() -> str: @@ -715,6 +711,9 @@ async def _run_on_tool_error_callbacks( copy.deepcopy(function_call.args) if function_call.args else {} ) + # TODO: thread a ToolConfirmation through here so an approved tool can be + # re-executed in live mode. `tool_confirmation` is always None on this path, + # so a confirmation-gated tool can only ever be refused, never resumed. tool_context = _create_tool_context(invocation_context, function_call) try: @@ -958,23 +957,58 @@ async def _process_function_live_helper( # for streaming tool use case # we require the function to be an async generator function async def run_tool_and_update_queue(tool, function_args, tool_context): + live_request_queue = invocation_context.live_request_queue + if live_request_queue is None: + raise RuntimeError('Streaming tools require a live request queue.') try: - async with Aclosing( - __call_tool_live( - tool=tool, - args=function_args, - tool_context=tool_context, - invocation_context=invocation_context, - ) - ) as agen: - async for result in agen: - updated_content = _build_function_response_content( - tool, result, tool_context.function_call_id - ) - invocation_context.live_request_queue.send_content(updated_content) + res = await __call_tool_async( + tool=tool, + args=function_args, + tool_context=tool_context, + ) + if inspect.isasyncgen(res): + async with Aclosing(res) as agen: + async for result in agen: + updated_content = _build_function_response_content( + tool, result, tool_context.function_call_id + ) + live_request_queue.send_content(updated_content) + else: + # `res` is a single terminal payload (e.g. the error dict returned + # when confirmation is required/rejected or a mandatory argument is + # missing), not a chunk of a stream. + # TODO: for the confirmation-required case, hold the call pending + # (as long-running tools do) instead of relaying the error. Relaying + # it closes the call id with the model, so a later approval would + # have to send a second response reusing that same id. + updated_content = _build_function_response_content( + tool, res, tool_context.function_call_id + ) + live_request_queue.send_content(updated_content) except asyncio.CancelledError: raise # Re-raise to properly propagate the cancellation + except Exception: # pylint: disable=broad-except + # The model already got a `pending` response for this call, so it waits + # for a follow-up FunctionResponse. Swallowing the exception here would + # leave the live session hanging, so report the failure to the model. + # The exception text is deliberately not forwarded to the model: it can + # carry internal detail that is irrelevant to it. It is logged instead. + logger.exception('Error executing streaming tool %s.', tool.name) + error_content = _build_function_response_content( + tool, + { + 'error': ( + f'Invoking `{tool.name}()` failed with an internal error.' + ) + }, + tool_context.function_call_id, + ) + live_request_queue.send_content(error_content) + # TODO: resolve `require_confirmation` before spawning the task. The + # confirmation request is recorded on `tool_context.actions` by the + # background task while the caller builds the response event, and nothing + # orders the two, so the request can be missing from the emitted event. task = asyncio.create_task( run_tool_and_update_queue(tool, function_args, tool_context) ) @@ -1125,24 +1159,6 @@ def _try_decode_computer_use_image( return None -async def __call_tool_live( - tool: BaseTool, - args: dict[str, object], - tool_context: ToolContext, - invocation_context: InvocationContext, -) -> AsyncGenerator[Event, None]: - """Calls the tool asynchronously (awaiting the coroutine).""" - async with Aclosing( - tool._call_live( - args=args, - tool_context=tool_context, - invocation_context=invocation_context, - ) - ) as agen: - async for item in agen: - yield item - - async def __call_tool_async( tool: BaseTool, args: dict[str, Any], diff --git a/src/google/adk/tools/function_tool.py b/src/google/adk/tools/function_tool.py index cd77f7948f6..ee63fc3873c 100644 --- a/src/google/adk/tools/function_tool.py +++ b/src/google/adk/tools/function_tool.py @@ -14,25 +14,24 @@ from __future__ import annotations +from contextlib import contextmanager +import contextvars import inspect import logging from typing import Any +from typing import Awaitable from typing import Callable from typing import cast from typing import get_args from typing import get_origin +from typing import Iterator 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 -from ..utils.context_utils import Aclosing from ..utils.context_utils import find_context_parameter from ._automatic_function_calling_util import build_function_declaration from .base_tool import BaseTool @@ -40,6 +39,29 @@ logger = logging.getLogger('google_adk.' + __name__) +_SyncCallableRunner = Callable[ + [Callable[..., Any], dict[str, Any]], Awaitable[Any] +] +_SYNC_CALLABLE_RUNNER: contextvars.ContextVar[_SyncCallableRunner | None] = ( + contextvars.ContextVar('adk_sync_callable_runner', default=None) +) + + +@contextmanager +def _use_sync_callable_runner( + runner: _SyncCallableRunner | None = None, +) -> Iterator[None]: + """Binds the runner used for synchronous callables. + + Passing ``None`` clears the binding, which stops a worker-owned nested call + from reusing the caller's runner. + """ + token = _SYNC_CALLABLE_RUNNER.set(runner) + try: + yield + finally: + _SYNC_CALLABLE_RUNNER.reset(token) + class FunctionTool(BaseTool): """A tool that wraps a user-defined Python function. @@ -170,6 +192,19 @@ def _prepare_invocation_args( valid_params = set(signature.parameters.keys()) if self._context_param_name in valid_params: args_to_call[self._context_param_name] = tool_context + # In live mode (bidirectional streaming), tools may accept an 'input_stream' + # parameter (e.g., LiveRequestQueue) to receive real-time streaming data. + # When registered in _process_function_live_helper, the framework attaches + # the dedicated stream to invocation_context.active_streaming_tools[name]. + # If the tool signature expects 'input_stream', we inject that active stream. + if 'input_stream' in valid_params: + active_tools = tool_context._invocation_context.active_streaming_tools + if ( + active_tools is not None + and self.name in active_tools + and active_tools[self.name].stream is not None + ): + args_to_call['input_stream'] = active_tools[self.name].stream return {k: v for k, v in args_to_call.items() if k in valid_params} @override @@ -251,37 +286,10 @@ async def _invoke_callable( ) if is_async: return await target(**args_to_call) - else: - return target(**args_to_call) - - # TODO(hangfei): fix call live for function stream. - async def _call_live( - self, - *, - args: dict[str, Any], - tool_context: ToolContext, - 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 ( - 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 - ): - args_to_call['input_stream'] = invocation_context.active_streaming_tools[ - self.name - ].stream - if self._context_param_name in signature.parameters: - args_to_call[self._context_param_name] = tool_context - - # TODO: support tool confirmation for live mode. - async with Aclosing(self.func(**args_to_call)) as agen: - async for item in agen: - yield item + runner = _SYNC_CALLABLE_RUNNER.get() + if runner is not None: + return await runner(target, args_to_call) + return target(**args_to_call) def _get_mandatory_args( self, @@ -298,11 +306,17 @@ def _get_mandatory_args( # A parameter is mandatory if: # 1. It has no default value (param.default is inspect.Parameter.empty) # 2. It's not a variable positional (*args) or variable keyword (**kwargs) parameter + # 3. It's not an internal parameter to ignore (e.g. tool_context, input_stream) # # For more refer to: https://docs.python.org/3/library/inspect.html#inspect.Parameter.kind - if param.default == inspect.Parameter.empty and param.kind not in ( - inspect.Parameter.VAR_POSITIONAL, - inspect.Parameter.VAR_KEYWORD, + if ( + param.default == inspect.Parameter.empty + and name not in self._ignore_params + and param.kind + not in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ) ): mandatory_params.append(name) diff --git a/tests/unittests/flows/llm_flows/test_functions_simple.py b/tests/unittests/flows/llm_flows/test_functions_simple.py index d8b1ad0c3c6..9dbd2d3d4ec 100644 --- a/tests/unittests/flows/llm_flows/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/test_functions_simple.py @@ -1617,3 +1617,218 @@ async def slow_fn_2() -> dict[str, str]: await asyncio.sleep(0) assert len(invocation_context.active_non_blocking_tool_tasks) == 0 + + +@pytest.mark.asyncio +async def test_streaming_tool_with_input_stream_e2e_live(): + """A streaming tool that accepts input_stream receives the queue stream in live mode.""" + + async def streaming_fn(val: str, input_stream: LiveRequestQueue): + assert input_stream is not None + yield f'streamed_{val}' + + tool = FunctionTool(streaming_fn) + model = testing_utils.MockModel.create(responses=[]) + agent = Agent(name='test_agent', model=model, tools=[tool]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + invocation_context.live_request_queue = LiveRequestQueue() + + function_call = types.FunctionCall( + name=tool.name, args={'val': 'hello'}, id='fc_input_stream' + ) + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + await handle_function_calls_live(invocation_context, event, {tool.name: tool}) + contents = await _drain_live_function_responses( + invocation_context.live_request_queue, count=1 + ) + + function_response = contents[0].parts[0].function_response + assert function_response.response == {'result': 'streamed_hello'} + assert function_response.id == 'fc_input_stream' + + +@pytest.mark.asyncio +async def test_streaming_tool_missing_mandatory_arg_e2e_live(): + """A streaming tool with missing mandatory arguments emits an error FunctionResponse in live mode.""" + + async def streaming_fn(mandatory_arg: str): + yield f'res_{mandatory_arg}' + + tool = FunctionTool(streaming_fn) + model = testing_utils.MockModel.create(responses=[]) + agent = Agent(name='test_agent', model=model, tools=[tool]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + invocation_context.live_request_queue = LiveRequestQueue() + + function_call = types.FunctionCall( + name=tool.name, args={}, id='fc_missing_arg' + ) + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + await handle_function_calls_live(invocation_context, event, {tool.name: tool}) + contents = await _drain_live_function_responses( + invocation_context.live_request_queue, count=1 + ) + + function_response = contents[0].parts[0].function_response + assert ( + 'mandatory input parameters are not present' + in function_response.response['error'] + ) + assert function_response.id == 'fc_missing_arg' + + +@pytest.mark.asyncio +async def test_streaming_tool_sequential_branches_e2e_live(): + """Test sequential execution of both branches (error dict -> streaming generator) in live mode.""" + + async def streaming_fn(mandatory_arg: str): + yield f'res_{mandatory_arg}' + + tool = FunctionTool(streaming_fn) + model = testing_utils.MockModel.create(responses=[]) + agent = Agent(name='test_agent', model=model, tools=[tool]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + invocation_context.live_request_queue = LiveRequestQueue() + + # Turn 1: Call with missing argument -> exercises `else:` (dict error response) + function_call_err = types.FunctionCall( + name=tool.name, args={}, id='fc_turn1_err' + ) + event1 = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content( + parts=[types.Part(function_call=function_call_err)] + ), + ) + + await handle_function_calls_live( + invocation_context, event1, {tool.name: tool} + ) + contents1 = await _drain_live_function_responses( + invocation_context.live_request_queue, count=1 + ) + assert ( + 'mandatory input parameters are not present' + in contents1[0].parts[0].function_response.response['error'] + ) + assert contents1[0].parts[0].function_response.id == 'fc_turn1_err' + + # Turn 2: Call with mandatory argument provided -> exercises `if inspect.isasyncgen(res):` + function_call_success = types.FunctionCall( + name=tool.name, args={'mandatory_arg': 'world'}, id='fc_turn2_success' + ) + event2 = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content( + parts=[types.Part(function_call=function_call_success)] + ), + ) + + await handle_function_calls_live( + invocation_context, event2, {tool.name: tool} + ) + contents2 = await _drain_live_function_responses( + invocation_context.live_request_queue, count=1 + ) + assert contents2[0].parts[0].function_response.response == { + 'result': 'res_world' + } + assert contents2[0].parts[0].function_response.id == 'fc_turn2_success' + + +@pytest.mark.asyncio +async def test_streaming_tool_raising_reports_error_e2e_live(): + """A streaming tool that raises emits an error FunctionResponse in live mode.""" + + async def streaming_fn(val: str): + raise ValueError(f'sensitive_detail_{val}') + yield # pylint: disable=unreachable + + tool = FunctionTool(streaming_fn) + model = testing_utils.MockModel.create(responses=[]) + agent = Agent(name='test_agent', model=model, tools=[tool]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + invocation_context.live_request_queue = LiveRequestQueue() + + function_call = types.FunctionCall( + name=tool.name, args={'val': 'hello'}, id='fc_raises' + ) + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + await handle_function_calls_live(invocation_context, event, {tool.name: tool}) + contents = await _drain_live_function_responses( + invocation_context.live_request_queue, count=1 + ) + + function_response = contents[0].parts[0].function_response + assert function_response.response == { + 'error': f'Invoking `{tool.name}()` failed with an internal error.' + } + # The raw exception text is not leaked to the model. + assert 'sensitive_detail' not in function_response.response['error'] + assert function_response.id == 'fc_raises' + # The task completes instead of dying with an unretrieved exception. + task = invocation_context.active_streaming_tools[tool.name].task + assert task.done() + assert task.exception() is None + + +@pytest.mark.asyncio +async def test_streaming_tool_require_confirmation_refused_e2e_live(): + """A confirmation-gated streaming tool is refused, not run, in live mode.""" + calls = [] + + async def streaming_fn(val: str): + calls.append(val) + yield f'streamed_{val}' + + tool = FunctionTool(streaming_fn, require_confirmation=True) + model = testing_utils.MockModel.create(responses=[]) + agent = Agent(name='test_agent', model=model, tools=[tool]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + invocation_context.live_request_queue = LiveRequestQueue() + + function_call = types.FunctionCall( + name=tool.name, args={'val': 'hello'}, id='fc_confirmation' + ) + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + await handle_function_calls_live(invocation_context, event, {tool.name: tool}) + contents = await _drain_live_function_responses( + invocation_context.live_request_queue, count=1 + ) + + function_response = contents[0].parts[0].function_response + assert 'requires confirmation' in function_response.response['error'] + assert function_response.id == 'fc_confirmation' + assert calls == [] diff --git a/tests/unittests/flows/llm_flows/test_functions_thread_pool.py b/tests/unittests/flows/llm_flows/test_functions_thread_pool.py index 5ffd0f26d66..a2417c116e1 100644 --- a/tests/unittests/flows/llm_flows/test_functions_thread_pool.py +++ b/tests/unittests/flows/llm_flows/test_functions_thread_pool.py @@ -28,8 +28,8 @@ from google.adk.tools.base_tool import BaseTool from google.adk.tools.function_tool import FunctionTool from google.adk.tools.set_model_response_tool import SetModelResponseTool +from google.adk.tools.tool_confirmation import ToolConfirmation from google.adk.tools.tool_context import ToolContext -from google.genai import types from pydantic import BaseModel import pytest @@ -143,6 +143,99 @@ def sync_func() -> dict: assert tool_thread_id is not None assert tool_thread_id != main_thread_id + @pytest.mark.asyncio + async def test_sync_tool_preserves_function_tool_confirmation(self): + calls: list[str] = [] + + def write(value: str) -> dict[str, str]: + calls.append(value) + return {'value': value} + + tool = FunctionTool(write, require_confirmation=True) + agent = Agent( + name='test_agent', + model=testing_utils.MockModel.create(responses=[]), + tools=[tool], + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + tool_context = ToolContext( + invocation_context=invocation_context, + function_call_id='test_id', + ) + + result = await _call_tool_in_thread_pool( + tool, {'value': 'write'}, tool_context + ) + + assert result['error'].startswith('This tool call requires confirmation') + assert calls == [] + assert 'test_id' in tool_context.actions.requested_tool_confirmations + + @pytest.mark.asyncio + async def test_sync_tool_runs_in_thread_pool_once_confirmed(self): + main_thread_id = threading.current_thread().ident + tool_thread_id = None + calls: list[str] = [] + + def write(value: str) -> dict[str, str]: + nonlocal tool_thread_id + tool_thread_id = threading.current_thread().ident + calls.append(value) + return {'value': value} + + tool = FunctionTool(write, require_confirmation=True) + agent = Agent( + name='test_agent', + model=testing_utils.MockModel.create(responses=[]), + tools=[tool], + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + tool_context = ToolContext( + invocation_context=invocation_context, + function_call_id='test_id', + tool_confirmation=ToolConfirmation(confirmed=True), + ) + + result = await _call_tool_in_thread_pool( + tool, {'value': 'write'}, tool_context + ) + + assert result == {'value': 'write'} + assert calls == ['write'] + assert tool_thread_id is not None + assert tool_thread_id != main_thread_id + + @pytest.mark.asyncio + async def test_sync_tool_does_not_leak_runner_into_nested_function_tool(self): + inner_tool = FunctionTool(lambda: 'nested result') + agent = Agent( + name='test_agent', + model=testing_utils.MockModel.create(responses=[]), + tools=[inner_tool], + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + tool_context = ToolContext( + invocation_context=invocation_context, + function_call_id='test_id', + ) + + def outer() -> str: + return asyncio.run( + inner_tool.run_async(args={}, tool_context=tool_context) + ) + + result = await _call_tool_in_thread_pool( + FunctionTool(outer), {}, tool_context + ) + + assert result == 'nested result' + @pytest.mark.asyncio async def test_async_tool_runs_in_thread_pool(self): """Test that async tools run in a separate thread with new event loop.""" diff --git a/tests/unittests/tools/test_function_tool.py b/tests/unittests/tools/test_function_tool.py index 9c76529fb5d..fbb16821fc8 100644 --- a/tests/unittests/tools/test_function_tool.py +++ b/tests/unittests/tools/test_function_tool.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect +from typing import Any +from unittest import mock from unittest.mock import MagicMock from google.adk.agents.context import Context @@ -529,3 +532,115 @@ async def async_tool(query: str, context: Context) -> dict: assert result["query"] == "hello" assert result["context_type"] == "Context" + + +@pytest.mark.asyncio +async def test_run_async_with_async_generator_streaming_tool(mock_tool_context): + """Test that run_async returns an AsyncGenerator when wrapped function is an async generator.""" + + async def streaming_tool(val: int, tool_context: Context): + yield f"item_{val}" + yield f"item_{val + 1}" + + tool = FunctionTool(streaming_tool) + result = await tool.run_async( + args={"val": 10}, + tool_context=mock_tool_context, + ) + + items = [] + async for item in result: + items.append(item) + + assert items == ["item_10", "item_11"] + + +@pytest.mark.asyncio +async def test_run_async_with_streaming_tool_and_input_stream( + mock_tool_context, +): + """Test that run_async injects input_stream into args_to_call for a streaming tool.""" + mock_stream = mock.MagicMock() + mock_stream.read.return_value = "stream_data" + + mock_tool_context._invocation_context = mock.MagicMock() + mock_tool_context._invocation_context.active_streaming_tools = { + "streaming_tool_input": mock.MagicMock(stream=mock_stream) + } + + async def streaming_tool_input(val: int, input_stream: Any): + data = input_stream.read() + yield f"{data}_{val}" + + tool = FunctionTool(streaming_tool_input) + result = await tool.run_async( + args={"val": 42}, + tool_context=mock_tool_context, + ) + + items = [item async for item in result] + assert items == ["stream_data_42"] + + +@pytest.mark.asyncio +async def test_run_async_with_streaming_tool_require_confirmation( + mock_tool_context, +): + """Test e2e confirmation lifecycle for a streaming tool in run_async.""" + + async def streaming_tool_conf(val: int): + yield f"confirmed_{val}" + + tool = FunctionTool(streaming_tool_conf, require_confirmation=True) + mock_tool_context.function_call_id = "test_function_call_id" + + # Stage 1: Call without confirmation should request confirmation and return error dict + mock_tool_context.tool_confirmation = None + res_unconfirmed = await tool.run_async( + args={"val": 1}, + tool_context=mock_tool_context, + ) + assert isinstance(res_unconfirmed, dict) + assert "error" in res_unconfirmed + assert "requires confirmation" in res_unconfirmed["error"] + assert ( + "test_function_call_id" + in mock_tool_context.actions.requested_tool_confirmations + ) + + # Stage 2: Call with rejected confirmation + mock_tool_context.tool_confirmation = ToolConfirmation(confirmed=False) + res_rejected = await tool.run_async( + args={"val": 1}, + tool_context=mock_tool_context, + ) + assert res_rejected == {"error": "This tool call is rejected."} + + # Stage 3: Call with approved confirmation should return the AsyncGenerator + mock_tool_context.tool_confirmation = ToolConfirmation(confirmed=True) + res_confirmed = await tool.run_async( + args={"val": 1}, + tool_context=mock_tool_context, + ) + assert inspect.isasyncgen(res_confirmed) + items = [item async for item in res_confirmed] + assert items == ["confirmed_1"] + + +@pytest.mark.asyncio +async def test_run_async_with_streaming_tool_missing_mandatory_arg( + mock_tool_context, +): + """Test that missing mandatory parameters in a streaming tool return an error dict.""" + + async def streaming_tool_req(req_param: str): + yield req_param + + tool = FunctionTool(streaming_tool_req) + result = await tool.run_async( + args={}, + tool_context=mock_tool_context, + ) + assert isinstance(result, dict) + assert "error" in result + assert "mandatory input parameters are not present" in result["error"]