From 30cc3f5264f498b75affec3300405f17fee201e3 Mon Sep 17 00:00:00 2001 From: Liang Wu <18244712+wuliang229@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:03:05 +0000 Subject: [PATCH 1/4] feat(live): Support asynchronous non-blocking tools in v1 Port live API non-blocking function response scheduling and background task execution from main to v1 branch. - Add response_scheduling attribute to BaseTool - Mark response_scheduling & live streaming tools as NON_BLOCKING declarations - Execute non-blocking tools in background asyncio.Task without blocking live event stream - Add active_non_blocking_tool_tasks to InvocationContext for task tracking - Guard against emitting None function-response events in live stream --- src/google/adk/agents/invocation_context.py | 3 + .../adk/flows/llm_flows/base_llm_flow.py | 54 +++++-- src/google/adk/flows/llm_flows/functions.py | 132 +++++++++++++----- src/google/adk/tools/base_tool.py | 14 ++ .../flows/llm_flows/test_functions_simple.py | 79 +++++++++++ 5 files changed, 232 insertions(+), 50 deletions(-) diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 8dc1a01af28..75db1e808a9 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -185,6 +185,9 @@ class InvocationContext(BaseModel): active_streaming_tools: Optional[dict[str, ActiveStreamingTool]] = None """The running streaming tools of this invocation.""" + active_non_blocking_tool_tasks: Optional[dict[str, asyncio.Task[Any]]] = None + """The running non-blocking tool tasks of this invocation (Live only).""" + transcription_cache: Optional[list[TranscriptionEntry]] = None """Caches necessary data, audio or contents, that are needed by transcription.""" diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 9c0731e6a2b..ba33eb3a8a6 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -466,6 +466,27 @@ async def _process_agent_tools( tool_context=tool_context, llm_request=llm_request ) + if invocation_context.live_request_queue is not None: + _mark_live_async_tools_non_blocking(llm_request) + + +def _mark_live_async_tools_non_blocking(llm_request: LlmRequest) -> None: + """Marks live streaming and response-scheduling tools as NON_BLOCKING. + + These tools emit asynchronous FunctionResponses, which the Live API only + accepts for NON_BLOCKING declarations. + """ + if not llm_request.config.tools: + return + for gemini_tool in llm_request.config.tools: + for declaration in gemini_tool.function_declarations or []: + tool = (llm_request.tools_dict or {}).get(declaration.name) + if not tool: + continue + is_streaming = hasattr(tool, 'func') and inspect.isasyncgenfunction(tool.func) + if is_streaming or tool.response_scheduling is not None: + declaration.behavior = types.Behavior.NON_BLOCKING + class BaseLlmFlow(ABC): """A basic flow that calls the LLM in a loop until a final response is generated. @@ -1141,23 +1162,26 @@ async def _postprocess_live( # Handles function calls. if model_response_event.get_function_calls(): - function_response_event = await functions.handle_function_calls_live( + # handle_function_calls_live returns None when every call is deferred + # (e.g. all long-running or non-blocking), so guard before yielding to + # avoid emitting a None event into the live stream. + if function_response_event := await functions.handle_function_calls_live( invocation_context, model_response_event, llm_request.tools_dict - ) - # Always yield the function response event first - yield function_response_event - - # Check if this is a set_model_response function response - if json_response := _output_schema_processor.get_structured_model_response( - function_response_event ): - # Create and yield a final model response event - final_event = ( - _output_schema_processor.create_final_model_response_event( - invocation_context, json_response - ) - ) - yield final_event + # Always yield the function response event first + yield function_response_event + + # Check if this is a set_model_response function response + if json_response := _output_schema_processor.get_structured_model_response( + function_response_event + ): + # Create and yield a final model response event + final_event = ( + _output_schema_processor.create_final_model_response_event( + invocation_context, json_response + ) + ) + yield final_event async def _postprocess_run_processors_async( self, invocation_context: InvocationContext, llm_response: LlmResponse diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 588126698c7..85f69b9384b 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -613,8 +613,8 @@ async def handle_function_calls_live( if not function_calls: return None - # Create async lock for active_streaming_tools modifications - streaming_lock = asyncio.Lock() + # Create async lock for active_streaming_tools and active_non_blocking_tool_tasks modifications + active_tools_lock = asyncio.Lock() # Create tasks for parallel execution tasks = [ @@ -624,7 +624,7 @@ async def handle_function_calls_live( function_call, tools_dict, agent, - streaming_lock, + active_tools_lock, ) ) for function_call in function_calls @@ -671,7 +671,7 @@ async def _execute_single_function_call_live( function_call: types.FunctionCall, tools_dict: dict[str, BaseTool], agent: LlmAgent, - streaming_lock: asyncio.Lock, + active_tools_lock: asyncio.Lock, ) -> Optional[Event]: """Execute a single function call for live mode with thread safety.""" @@ -770,7 +770,7 @@ async def _run_with_trace(): function_call, function_args, invocation_context, - streaming_lock, + active_tools_lock, ) except Exception as tool_error: error_response = await _run_on_tool_error_callbacks( @@ -830,11 +830,62 @@ async def _run_with_trace(): ) return function_response_event - async with _instrumentation.record_tool_execution( - tool, agent, function_args - ) as tel_ctx: - tel_ctx.function_response_event = await _run_with_trace() - return tel_ctx.function_response_event + is_streaming = hasattr(tool, 'func') and inspect.isasyncgenfunction(tool.func) + is_non_blocking = not is_streaming and tool.response_scheduling is not None + if is_non_blocking: + task_key = f'{tool.name}_{function_call.id}' + + async def _background_task() -> None: + try: + async with _instrumentation.record_tool_execution( + tool, agent, function_args + ) as tel_ctx: + function_response_event = await _run_with_trace() + tel_ctx.function_response_event = function_response_event + + if function_response_event: + if ( + invocation_context.session_service + and invocation_context.session + ): + await invocation_context.session_service.append_event( + session=invocation_context.session, + event=function_response_event, + ) + if ( + invocation_context.live_request_queue + and function_response_event.content + ): + invocation_context.live_request_queue.send_content( + function_response_event.content + ) + except Exception as e: + logger.error( + 'Error executing non-blocking tool %s in background: %s', + tool.name, + e, + exc_info=True, + ) + finally: + async with active_tools_lock: + if ( + invocation_context.active_non_blocking_tool_tasks + and task_key in invocation_context.active_non_blocking_tool_tasks + ): + del invocation_context.active_non_blocking_tool_tasks[task_key] + + task = asyncio.create_task(_background_task()) + async with active_tools_lock: + if invocation_context.active_non_blocking_tool_tasks is None: + invocation_context.active_non_blocking_tool_tasks = {} + invocation_context.active_non_blocking_tool_tasks[task_key] = task + return None + else: + async with _instrumentation.record_tool_execution( + tool, agent, function_args + ) as tel_ctx: + tel_ctx.function_response_event = await _run_with_trace() + return tel_ctx.function_response_event async def _process_function_live_helper( @@ -843,7 +894,7 @@ async def _process_function_live_helper( function_call, function_args, invocation_context, - streaming_lock: asyncio.Lock, + active_tools_lock: asyncio.Lock, ): function_response = None # Check if this is a stop_streaming function call @@ -853,7 +904,7 @@ async def _process_function_live_helper( ): function_name = function_args['function_name'] # Thread-safe access to active_streaming_tools - async with streaming_lock: + async with active_tools_lock: active_tasks = invocation_context.active_streaming_tools if ( active_tasks @@ -886,7 +937,7 @@ async def _process_function_live_helper( } if not function_response: # Clean up the reference under lock - async with streaming_lock: + async with active_tools_lock: if ( invocation_context.active_streaming_tools and function_name in invocation_context.active_streaming_tools @@ -917,13 +968,8 @@ async def run_tool_and_update_queue(tool, function_args, tool_context): ) ) as agen: async for result in agen: - updated_content = types.Content( - role='user', - parts=[ - types.Part.from_text( - text=f'Function {tool.name} returned: {result}' - ) - ], + updated_content = _build_function_response_content( + tool, result, tool_context.function_call_id ) invocation_context.live_request_queue.send_content(updated_content) except asyncio.CancelledError: @@ -933,7 +979,7 @@ async def run_tool_and_update_queue(tool, function_args, tool_context): run_tool_and_update_queue(tool, function_args, tool_context) ) - async with streaming_lock: + async with active_tools_lock: if invocation_context.active_streaming_tools is None: invocation_context.active_streaming_tools = {} @@ -1112,26 +1158,17 @@ def __build_response_event( tool_context: ToolContext, invocation_context: InvocationContext, ) -> Event: - # Specs requires the result to be a dict. - if not isinstance(function_result, dict): - function_result = {'result': function_result} - function_response_parts = None if isinstance(tool, ComputerUseTool): function_response_parts = _try_decode_computer_use_image( tool, function_result ) - part_function_response = types.Part.from_function_response( - name=tool.name, - response=function_result, - parts=function_response_parts, - ) - part_function_response.function_response.id = tool_context.function_call_id - - content = types.Content( - role='user', - parts=[part_function_response], + content = _build_function_response_content( + tool, + function_result, + tool_context.function_call_id, + function_response_parts, ) function_response_event = Event( @@ -1145,6 +1182,31 @@ def __build_response_event( return function_response_event +def _build_function_response_content( + tool: BaseTool, + function_result: object, + function_call_id: Optional[str], + function_response_parts: Optional[list[types.FunctionResponsePart]] = None, +) -> types.Content: + """Builds the content carrying a tool result as a FunctionResponse.""" + # Specs requires the result to be a dict. + if not isinstance(function_result, dict): + function_result = {'result': function_result} + + part_function_response = types.Part.from_function_response( + name=tool.name, + response=function_result, + parts=function_response_parts, + ) + part_function_response.function_response.id = function_call_id + if tool.response_scheduling is not None: + part_function_response.function_response.scheduling = ( + tool.response_scheduling + ) + + return types.Content(role='user', parts=[part_function_response]) + + def deep_merge_dicts(d1: dict, d2: dict) -> dict: """Recursively merges d2 into d1.""" for key, value in d2.items(): diff --git a/src/google/adk/tools/base_tool.py b/src/google/adk/tools/base_tool.py index 8dd112a6c8a..e5c4bb73f98 100644 --- a/src/google/adk/tools/base_tool.py +++ b/src/google/adk/tools/base_tool.py @@ -65,6 +65,18 @@ class BaseTool(ABC): NOTE: the entire dict must be JSON serializable. """ + response_scheduling: Optional[types.FunctionResponseScheduling] = None + """Controls when the model reacts to the tool's response (Live API only). + + Applied to the emitted ``FunctionResponse`` for asynchronous function calling: + - ``SILENT``: feeds the response back without triggering a model turn. + - ``WHEN_IDLE``: defers the reaction until the model is idle. + - ``INTERRUPT``: reacts immediately. + + Ignored by models that don't support asynchronous function calling. ``None`` + preserves the default behavior. + """ + def __init__( self, *, @@ -72,11 +84,13 @@ def __init__( description, is_long_running: bool = False, custom_metadata: Optional[dict[str, Any]] = None, + response_scheduling: Optional[types.FunctionResponseScheduling] = None, ): self.name = name self.description = description self.is_long_running = is_long_running self.custom_metadata = custom_metadata + self.response_scheduling = response_scheduling def _get_declaration(self) -> Optional[types.FunctionDeclaration]: """Gets the OpenAPI specification of this tool in the form of a FunctionDeclaration. diff --git a/tests/unittests/flows/llm_flows/test_functions_simple.py b/tests/unittests/flows/llm_flows/test_functions_simple.py index be638de44cb..388b6d44c23 100644 --- a/tests/unittests/flows/llm_flows/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/test_functions_simple.py @@ -16,6 +16,7 @@ from typing import Any from typing import Callable +from google.adk.agents.live_request_queue import LiveRequestQueue from google.adk.agents.llm_agent import Agent from google.adk.events.event import Event from google.adk.events.event_actions import EventActions @@ -1319,3 +1320,81 @@ def simple_fn() -> dict[str, str]: assert result_parallel is not None assert result_parallel.live_session_id == 'test-live-session-id-parallel' + + +@pytest.mark.asyncio +async def test_response_scheduling_applied_to_function_response(): + """response_scheduling on a tool is stamped onto the FunctionResponse part.""" + + def simple_fn(**kwargs) -> dict: + return {'result': 'test'} + + tool = FunctionTool(simple_fn) + tool.response_scheduling = types.FunctionResponseScheduling.SILENT + 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='' + ) + + function_call = types.FunctionCall(name=tool.name, args={}, id='fc_test') + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + result_event = await handle_function_calls_async( + invocation_context, event, {tool.name: tool} + ) + + assert result_event is not None + function_response = result_event.content.parts[0].function_response + assert function_response.scheduling is types.FunctionResponseScheduling.SILENT + + +@pytest.mark.asyncio +async def test_non_blocking_tool_handled_asynchronously(): + """Tests that a NON_BLOCKING tool returns None inline and pushes to live request queue.""" + + async def slow_fn() -> dict[str, str]: + await asyncio.sleep(0.01) + return {'result': 'done'} + + tool = FunctionTool(slow_fn) + tool.response_scheduling = types.FunctionResponseScheduling.WHEN_IDLE + 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_non_blocking' + ) + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + result = await handle_function_calls_live( + invocation_context, event, {tool.name: tool} + ) + assert result is None + assert invocation_context.active_non_blocking_tool_tasks is not None + task_key = f'{tool.name}_fc_non_blocking' + assert task_key in invocation_context.active_non_blocking_tool_tasks + + request = await asyncio.wait_for( + invocation_context.live_request_queue._queue.get(), timeout=5 + ) + function_response = request.content.parts[0].function_response + assert function_response.response == {'result': 'done'} + assert function_response.id == 'fc_non_blocking' + assert ( + function_response.scheduling == types.FunctionResponseScheduling.WHEN_IDLE + ) + await asyncio.sleep(0) + assert task_key not in invocation_context.active_non_blocking_tool_tasks From 178334f563ba3e07c3109bd83100a1c7bbab4c90 Mon Sep 17 00:00:00 2001 From: Liang Wu <18244712+wuliang229@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:16:29 +0000 Subject: [PATCH 2/4] test(live): Port all unit tests for non-blocking tools to v1 Port complete unit test suite for response scheduling, non-blocking tool declarations, deferred function call handling, and parallel/exception background task execution from main to v1 branch. --- .../flows/llm_flows/test_base_llm_flow.py | 121 ++++++++++ .../flows/llm_flows/test_functions_simple.py | 219 ++++++++++++++++++ 2 files changed, 340 insertions(+) diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index 2b8eb92d3a6..35e2d8b8193 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -1421,3 +1421,124 @@ async def mock_receive(): call_req.live_connect_config.history_config.initial_history_in_client_content is False ) + + +async def _preprocess(agent, *, is_live: bool) -> LlmRequest: + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='test message' + ) + if is_live: + invocation_context.live_request_queue = LiveRequestQueue() + flow = BaseLlmFlowForTesting() + llm_request = LlmRequest() + async for _ in flow._preprocess_async(invocation_context, llm_request): + pass + return llm_request + + +def _declarations(llm_request: LlmRequest) -> dict: + return { + decl.name: decl + for decl in llm_request.config.tools[0].function_declarations + } + + +async def _streaming_tool(query: str): + """A streaming tool.""" + yield f'streaming: {query}' + + +def _scheduled_tool(query: str) -> str: + """A scheduled tool.""" + return f'scheduled: {query}' + + +@pytest.mark.asyncio +async def test_process_agent_tools_marks_streaming_tool_non_blocking_for_live(): + """Live streaming async-generator tools are marked NON_BLOCKING.""" + agent = Agent(name='test_agent', tools=[_streaming_tool]) + + llm_request = await _preprocess(agent, is_live=True) + + declaration = llm_request.config.tools[0].function_declarations[0] + assert declaration.behavior is types.Behavior.NON_BLOCKING + + +@pytest.mark.asyncio +async def test_process_agent_tools_marks_scheduled_tool_non_blocking_for_live(): + """Live response-scheduling tools are marked NON_BLOCKING.""" + from google.adk.tools.function_tool import FunctionTool + + tool = FunctionTool(func=_scheduled_tool) + tool.response_scheduling = types.FunctionResponseScheduling.SILENT + agent = Agent(name='test_agent', tools=[tool]) + + llm_request = await _preprocess(agent, is_live=True) + + declaration = llm_request.config.tools[0].function_declarations[0] + assert declaration.behavior is types.Behavior.NON_BLOCKING + + +@pytest.mark.asyncio +async def test_process_agent_tools_does_not_mark_non_blocking_for_non_live(): + """Non-live requests never set behavior, even for streaming tools.""" + from google.adk.tools.function_tool import FunctionTool + + scheduled = FunctionTool(func=_scheduled_tool) + scheduled.response_scheduling = types.FunctionResponseScheduling.SILENT + agent = Agent(name='test_agent', tools=[_streaming_tool, scheduled]) + + llm_request = await _preprocess(agent, is_live=False) + + declarations = _declarations(llm_request) + assert declarations['_streaming_tool'].behavior is None + assert declarations['_scheduled_tool'].behavior is None + + +@pytest.mark.asyncio +async def test_process_agent_tools_leaves_regular_tool_behavior_unset_for_live(): + """Regular (non-streaming, non-scheduled) live tools are left untouched.""" + agent = Agent(name='test_agent', tools=[_scheduled_tool]) + + llm_request = await _preprocess(agent, is_live=True) + + declaration = llm_request.config.tools[0].function_declarations[0] + assert declaration.behavior is None + + +@pytest.mark.asyncio +async def test_postprocess_live_skips_none_function_response_event(): + """When every live function call defers, no None event must be yielded.""" + from google.adk.flows.llm_flows import base_llm_flow as blf + + agent = Agent(name='test_agent', model='gemini-2.0-flash') + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + flow = BaseLlmFlowForTesting() + + fc_part = types.Part( + function_call=types.FunctionCall(name='lro', id='1', args={}) + ) + content = types.Content(role='model', parts=[fc_part]) + model_response_event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=content, + ) + llm_request = LlmRequest(model='gemini-2.0-flash') + llm_response = LlmResponse(content=content) + + with mock.patch.object( + blf.functions, + 'handle_function_calls_live', + new=AsyncMock(return_value=None), + ): + events = [ + event + async for event in flow._postprocess_live( + invocation_context, llm_request, llm_response, model_response_event + ) + ] + + assert all(event is not None for event in events) diff --git a/tests/unittests/flows/llm_flows/test_functions_simple.py b/tests/unittests/flows/llm_flows/test_functions_simple.py index 388b6d44c23..d8b1ad0c3c6 100644 --- a/tests/unittests/flows/llm_flows/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/test_functions_simple.py @@ -1398,3 +1398,222 @@ async def slow_fn() -> dict[str, str]: ) await asyncio.sleep(0) assert task_key not in invocation_context.active_non_blocking_tool_tasks + + +@pytest.mark.asyncio +async def test_response_scheduling_unset_by_default(): + """Without response_scheduling, the FunctionResponse part leaves it unset.""" + + def simple_fn(**kwargs) -> dict: + return {'result': 'test'} + + tool = FunctionTool(simple_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='' + ) + + function_call = types.FunctionCall(name=tool.name, args={}, id='fc_test') + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + result_event = await handle_function_calls_async( + invocation_context, event, {tool.name: tool} + ) + + assert result_event is not None + function_response = result_event.content.parts[0].function_response + assert function_response.scheduling is None + + +async def _drain_live_function_responses( + live_request_queue: LiveRequestQueue, + count: int, +) -> list[types.Content]: + """Drains ``count`` contents from a live request queue, ignoring acks.""" + contents = [] + while len(contents) < count: + request = await asyncio.wait_for(live_request_queue._queue.get(), timeout=5) + if request.content is not None: + contents.append(request.content) + return contents + + +@pytest.mark.asyncio +async def test_streaming_tool_with_scheduling_emits_function_response(): + """A streaming tool with response_scheduling relays yields as FunctionResponses.""" + + async def streaming_fn(**kwargs): + yield 'first' + yield 'second' + + tool = FunctionTool(streaming_fn) + tool.response_scheduling = types.FunctionResponseScheduling.SILENT + 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_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=2 + ) + + responses = [content.parts[0].function_response for content in contents] + assert [r.response for r in responses] == [ + {'result': 'first'}, + {'result': 'second'}, + ] + assert all(r.id == 'fc_stream' for r in responses) + assert all( + r.scheduling is types.FunctionResponseScheduling.SILENT for r in responses + ) + + +@pytest.mark.asyncio +async def test_streaming_tool_without_scheduling_emits_function_response(): + """A streaming tool without response_scheduling still relays yields as + FunctionResponses, leaving scheduling unset.""" + + async def streaming_fn(**kwargs): + yield 'hello' + + 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_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': 'hello'} + assert function_response.id == 'fc_stream' + assert function_response.scheduling is None + + +@pytest.mark.asyncio +async def test_non_blocking_tool_exception_handling_and_cleanup(): + """Tests that an exception in a NON_BLOCKING tool is logged and cleaned up.""" + + async def failing_fn() -> dict[str, str]: + await asyncio.sleep(0.05) + raise RuntimeError('Tool execution failed purposely') + + tool = FunctionTool(failing_fn) + tool.response_scheduling = types.FunctionResponseScheduling.WHEN_IDLE + 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_failing_non_blocking' + ) + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content(parts=[types.Part(function_call=function_call)]), + ) + + result = await handle_function_calls_live( + invocation_context, event, {tool.name: tool} + ) + assert result is None + task_key = f'{tool.name}_fc_failing_non_blocking' + assert ( + invocation_context.active_non_blocking_tool_tasks + and task_key in invocation_context.active_non_blocking_tool_tasks + ) + + await asyncio.sleep(0.1) + await asyncio.sleep(0) + assert task_key not in invocation_context.active_non_blocking_tool_tasks + + +@pytest.mark.asyncio +async def test_parallel_non_blocking_tools(): + """Tests that multiple NON_BLOCKING tools execute in parallel and clean up independently.""" + + async def slow_fn_1() -> dict[str, str]: + await asyncio.sleep(0.05) + return {'result': 'done_1'} + + async def slow_fn_2() -> dict[str, str]: + await asyncio.sleep(0.05) + return {'result': 'done_2'} + + tool1 = FunctionTool(slow_fn_1) + tool1.response_scheduling = types.FunctionResponseScheduling.WHEN_IDLE + tool2 = FunctionTool(slow_fn_2) + tool2.response_scheduling = types.FunctionResponseScheduling.SILENT + + model = testing_utils.MockModel.create(responses=[]) + agent = Agent(name='test_agent', model=model, tools=[tool1, tool2]) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, user_content='' + ) + invocation_context.live_request_queue = LiveRequestQueue() + + fc1 = types.FunctionCall(name=tool1.name, args={}, id='fc_1') + fc2 = types.FunctionCall(name=tool2.name, args={}, id='fc_2') + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content( + parts=[ + types.Part(function_call=fc1), + types.Part(function_call=fc2), + ] + ), + ) + + result = await handle_function_calls_live( + invocation_context, event, {tool1.name: tool1, tool2.name: tool2} + ) + assert result is None + assert len(invocation_context.active_non_blocking_tool_tasks) == 2 + key1 = f'{tool1.name}_fc_1' + key2 = f'{tool2.name}_fc_2' + assert key1 in invocation_context.active_non_blocking_tool_tasks + assert key2 in invocation_context.active_non_blocking_tool_tasks + + contents = await _drain_live_function_responses( + invocation_context.live_request_queue, count=2 + ) + responses = { + c.parts[0].function_response.id: c.parts[0].function_response + for c in contents + } + assert responses['fc_1'].response == {'result': 'done_1'} + assert responses['fc_2'].response == {'result': 'done_2'} + + await asyncio.sleep(0) + assert len(invocation_context.active_non_blocking_tool_tasks) == 0 From 89fef779f2c277249b1a296a69a3ec63ffdc8624 Mon Sep 17 00:00:00 2001 From: Liang Wu <18244712+wuliang229@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:33:57 +0000 Subject: [PATCH 3/4] fix(live): Add missing asyncio import and return type annotations for mypy - Import asyncio in invocation_context.py - Add -> Optional[Event] return type annotation to _run_with_trace in functions.py --- src/google/adk/agents/invocation_context.py | 1 + src/google/adk/flows/llm_flows/functions.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 75db1e808a9..614770e8cb6 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio from typing import Any from typing import cast from typing import Optional diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 85f69b9384b..7949d2c997f 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -504,7 +504,7 @@ async def _run_on_tool_error_callbacks( else: raise tool_error - async def _run_with_trace(): + async def _run_with_trace() -> Optional[Event]: nonlocal function_args # Step 1: Check if plugin before_tool_callback overrides the function @@ -733,7 +733,7 @@ async def _run_on_tool_error_callbacks( ) raise tool_error - async def _run_with_trace(): + async def _run_with_trace() -> Optional[Event]: nonlocal function_args # Do not use "args" as the variable name, because it is a reserved keyword From 72f60943fdee6eedff6abe499a45634883029a35 Mon Sep 17 00:00:00 2001 From: Liang Wu <18244712+wuliang229@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:41:51 +0000 Subject: [PATCH 4/4] style(live): Fix pyink line wrapping formatting in base_llm_flow.py --- src/google/adk/flows/llm_flows/base_llm_flow.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index ba33eb3a8a6..e2259cf3669 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -483,7 +483,9 @@ def _mark_live_async_tools_non_blocking(llm_request: LlmRequest) -> None: tool = (llm_request.tools_dict or {}).get(declaration.name) if not tool: continue - is_streaming = hasattr(tool, 'func') and inspect.isasyncgenfunction(tool.func) + is_streaming = hasattr(tool, 'func') and inspect.isasyncgenfunction( + tool.func + ) if is_streaming or tool.response_scheduling is not None: declaration.behavior = types.Behavior.NON_BLOCKING