fix: Route thread-pool and live-mode tool dispatch through run_async (v1) (#6799)

Co-authored-by: GWeale <GWeale@users.noreply.github.com>

src/google/adk/flows/llm_flows/functions.py

@@ -26,7 +26,7 @@ import inspect
 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
@@ -46,6 +46,8 @@ from ...telemetry import _instrumentation
 from ...telemetry.tracing import trace_merged_tool_calls
 from ...telemetry.tracing import tracer
 from ...tools.base_tool import BaseTool
+from ...tools.function_tool import _use_sync_callable_runner
+from ...tools.function_tool import FunctionTool
 from ...tools.tool_confirmation import ToolConfirmation
 from ...tools.tool_context import ToolContext
 from ...utils.context_utils import Aclosing
@@ -121,10 +123,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 +140,36 @@ async def _call_tool_in_thread_pool(
   Returns:
     The result of running the tool.
   """
-  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)
-    )
+    with _use_sync_callable_runner(run_sync_callable):
+      return await tool.run_async(args=args, tool_context=tool_context)
 
-  # Fall back to normal async execution for non-FunctionTool sync tools.
-  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 +710,9 @@ async def _execute_single_function_call_live(
       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 +956,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 +1158,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],

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 @@ from .tool_context import ToolContext
 
 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 @@ class FunctionTool(BaseTool):
     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 @@ You could retry calling this tool, but it is IMPORTANT for you to provide all th
     )
     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 @@ You could retry calling this tool, but it is IMPORTANT for you to provide all th
       # 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)
 

tests/unittests/flows/llm_flows/test_functions_simple.py

@@ -1617,3 +1617,218 @@ async def test_parallel_non_blocking_tools():
 
   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 == []

tests/unittests/flows/llm_flows/test_functions_thread_pool.py

@@ -28,8 +28,8 @@ from google.adk.flows.llm_flows.functions import _is_sync_tool
 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 @@ class TestCallToolInThreadPool:
     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."""

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 test_run_async_with_context_type_annotation(mock_tool_context):
 
   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"]