Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
GWeale wants to merge 2 commits into
base: v1
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
27f26ae
fix: preserve function-tool preflight in thread pool (v1)
GWeale Aug 17, 2026
7ea01ca
fix: stop live tool execution from bypassing run_async (v1)
GWeale Aug 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
144 changes: 80 additions & 64 deletions src/google/adk/flows/llm_flows/functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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).
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
)
Expand Down Expand Up @@ -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],
Expand Down
92 changes: 53 additions & 39 deletions src/google/adk/tools/function_tool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,32 +14,54 @@

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
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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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)

Expand Down
Loading
Loading