From 9cdd95abe1fef36cdeab242b79acbbbb7e1ddeed Mon Sep 17 00:00:00 2001 From: Liang Wu <18244712+wuliang229@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:54:33 +0000 Subject: [PATCH 1/3] fix(live): stop background tool tasks when a live agent run ends Streaming tools and non-blocking tools run as background asyncio tasks that nothing owns: only an explicit `stop_streaming` call ever cancelled one. When a live agent run ended for any other reason -- a handoff to another agent, `task_completed`, the model closing the connection, the caller walking away -- its tools kept running, feeding function responses into a live request queue that by then belonged to another agent, which never made those calls, or to nobody at all. The tasks were never awaited either, so a failure in one surfaced only as a stray asyncio warning. `BaseLlmFlow.run_live` now cancels the background tools its run started before it returns, and at a handoff it cancels them before the sub agent takes over the live request queue rather than when the enclosing run eventually finishes. Cancellation is best effort: a task that ignores it is logged instead of stalling the handoff or the caller's teardown. Stopped tools also give up their registry entry. `_send_to_model` copies every live request into each registered stream, so a stream left behind by a tool that is no longer reading kept growing for the rest of the session, one entry per audio chunk the user spoke. Co-authored-by: Liang Wu PiperOrigin-RevId: 964915958 --- .../adk/flows/llm_flows/base_llm_flow.py | 455 ++++++++++------- .../streaming/test_live_tool_shutdown.py | 471 ++++++++++++++++++ tests/unittests/streaming/test_streaming.py | 17 +- 3 files changed, 754 insertions(+), 189 deletions(-) create mode 100644 tests/unittests/streaming/test_live_tool_shutdown.py 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 23ec1f19747..8291e646ee5 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -76,6 +76,11 @@ class _ReconnectSentinel(Event): DEFAULT_TRANSFER_AGENT_DELAY = 1.0 DEFAULT_TASK_COMPLETION_DELAY = 1.0 +# How long a live run waits for a background tool task to honor cancellation +# before giving up on it. Matches the budget `stop_streaming` already gives a +# streaming tool it cancels. +_TOOL_SHUTDOWN_TIMEOUT_SECONDS = 1.0 + DEFAULT_MAX_RECONNECT_ATTEMPTS = 5 # Statistics configuration @@ -530,211 +535,293 @@ async def run_live( ) attempt = 1 - while True: - try: - # On subsequent attempts, use the saved token to reconnect - if invocation_context.live_session_resumption_handle: - logger.info('Attempting to reconnect (Attempt %s)...', attempt) - attempt += 1 - if not llm_request.live_connect_config: - llm_request.live_connect_config = types.LiveConnectConfig() - if not llm_request.live_connect_config.session_resumption: - llm_request.live_connect_config.session_resumption = ( - types.SessionResumptionConfig() - ) - llm_request.live_connect_config.session_resumption.handle = ( - invocation_context.live_session_resumption_handle - ) - # Only set transparent=True for Vertex AI backend, as the Gemini API - # backend explicitly rejects it. - if ( - isinstance(llm, Gemini) - and llm._api_backend == GoogleLLMVariant.VERTEX_AI # pylint: disable=protected-access - ): - session_resumption = ( - llm_request.live_connect_config.session_resumption - ) - if session_resumption.transparent is None: - session_resumption.transparent = True - - # When seeding a fresh connection with prior conversation history, set - # initial_history_in_client_content to True. This tells the Live server - # that the provided history already includes the model's past responses, - # preventing the server from generating duplicate responses for those replayed turns. - if ( - llm_request.contents - and not invocation_context.live_session_resumption_handle - ): - if not llm_request.live_connect_config: - llm_request.live_connect_config = types.LiveConnectConfig() - if not llm_request.live_connect_config.history_config: - llm_request.live_connect_config.history_config = ( - types.HistoryConfig() - ) - if ( - llm_request.live_connect_config.history_config.initial_history_in_client_content - is None - ): - llm_request.live_connect_config.history_config.initial_history_in_client_content = ( - True + try: + while True: + try: + # On subsequent attempts, use the saved token to reconnect + if invocation_context.live_session_resumption_handle: + logger.info('Attempting to reconnect (Attempt %s)...', attempt) + attempt += 1 + if not llm_request.live_connect_config: + llm_request.live_connect_config = types.LiveConnectConfig() + if not llm_request.live_connect_config.session_resumption: + llm_request.live_connect_config.session_resumption = ( + types.SessionResumptionConfig() + ) + llm_request.live_connect_config.session_resumption.handle = ( + invocation_context.live_session_resumption_handle ) + # Only set transparent=True for Vertex AI backend, as the Gemini API + # backend explicitly rejects it. + if ( + isinstance(llm, Gemini) + and llm._api_backend == GoogleLLMVariant.VERTEX_AI # pylint: disable=protected-access + ): + session_resumption = ( + llm_request.live_connect_config.session_resumption + ) + if session_resumption.transparent is None: + session_resumption.transparent = True - logger.info( - 'Establishing live connection for agent: %s', - invocation_context.agent.name, - ) - async with llm.connect(llm_request) as llm_connection: - # Reset attempt counter on successful connection. - attempt = 1 - # Skip sending history if we are resuming a session. The server - # already has the state associated with the resumption handle. + # When seeding a fresh connection with prior conversation history, set + # initial_history_in_client_content to True. This tells the Live server + # that the provided history already includes the model's past responses, + # preventing the server from generating duplicate responses for those replayed turns. if ( llm_request.contents and not invocation_context.live_session_resumption_handle ): - # Sends the conversation history to the model. - with tracer.start_as_current_span('send_data'): - # Combine regular contents with audio/transcription from session - logger.debug('Sending history to model: %s', llm_request.contents) - await llm_connection.send_history(llm_request.contents) - trace_send_data( - invocation_context, event_id, llm_request.contents + if not llm_request.live_connect_config: + llm_request.live_connect_config = types.LiveConnectConfig() + if not llm_request.live_connect_config.history_config: + llm_request.live_connect_config.history_config = ( + types.HistoryConfig() + ) + if ( + llm_request.live_connect_config.history_config.initial_history_in_client_content + is None + ): + llm_request.live_connect_config.history_config.initial_history_in_client_content = ( + True ) - send_task = asyncio.create_task( - self._send_to_model(llm_connection, invocation_context) + logger.info( + 'Establishing live connection for agent: %s', + invocation_context.agent.name, ) - - should_reconnect = False - try: - async with Aclosing( - self._receive_from_model( - llm_connection, - event_id, - invocation_context, - llm_request, + async with llm.connect(llm_request) as llm_connection: + # Reset attempt counter on successful connection. + attempt = 1 + # Skip sending history if we are resuming a session. The server + # already has the state associated with the resumption handle. + if ( + llm_request.contents + and not invocation_context.live_session_resumption_handle + ): + # Sends the conversation history to the model. + with tracer.start_as_current_span('send_data'): + # Combine regular contents with audio/transcription from session + logger.debug( + 'Sending history to model: %s', llm_request.contents ) - ) as agen: - async for event in agen: - if isinstance(event, _ReconnectSentinel): - should_reconnect = True - break - # Empty event means the queue is closed. - if not event: - break - logger.debug('Receive new event: %s', event) - yield event - # send back the function response to models - if event.get_function_responses(): - logger.debug( - 'Sending back last function response event: %s', event - ) - invocation_context.live_request_queue.send_content( - event.content - ) - # We handle agent transfer here in `run_live` rather than - # in `_postprocess_live` to prevent duplication of function - # response processing. If agent transfer were handled in - # `_postprocess_live`, events yielded from child agent's - # `run_live` would bubble up to parent agent's `run_live`, - # causing `event.get_function_responses()` to be true in both - # child and parent, and `send_content()` to be called twice for - # the same function response. By handling agent transfer here, - # we ensure that only child agent processes its own function - # responses after the transfer. - # - # The transfer is gated on the `transfer_to_agent` action - # rather than on the position of the `transfer_to_agent` - # function response: the model may issue the transfer alongside - # other function calls, whose responses are merged into a - # single event in call order, so the transfer response is not - # necessarily `parts[0]`. Gating on the action matches - # `_postprocess_handle_function_calls_async`, and also covers - # tools that request a transfer by setting the action directly - # instead of calling `transfer_to_agent`. - transfer_to_agent = event.actions.transfer_to_agent - if transfer_to_agent: - await asyncio.sleep(DEFAULT_TRANSFER_AGENT_DELAY) - # cancel the tasks that belongs to the closed connection. - send_task.cancel() - logger.debug('Closing live connection') - await llm_connection.close() - logger.debug('Live connection closed.') - # transfer to the sub agent. - logger.debug('Transferring to agent: %s', transfer_to_agent) - agent_to_run = self._get_agent_to_run( - invocation_context, transfer_to_agent + await llm_connection.send_history(llm_request.contents) + trace_send_data( + invocation_context, event_id, llm_request.contents + ) + + send_task = asyncio.create_task( + self._send_to_model(llm_connection, invocation_context) + ) + + should_reconnect = False + try: + async with Aclosing( + self._receive_from_model( + llm_connection, + event_id, + invocation_context, + llm_request, ) - child_ctx = invocation_context.model_copy() - # Child Live agent should start a new Live session. - # Do not reuse the parent session's resumption handle. - child_ctx.live_session_resumption_handle = None - - if child_ctx.run_config: - child_ctx.run_config = child_ctx.run_config.model_copy( - deep=True + ) as agen: + async for event in agen: + if isinstance(event, _ReconnectSentinel): + should_reconnect = True + break + # Empty event means the queue is closed. + if not event: + break + logger.debug('Receive new event: %s', event) + yield event + # send back the function response to models + if event.get_function_responses(): + logger.debug( + 'Sending back last function response event: %s', event ) - if child_ctx.run_config.session_resumption: - child_ctx.run_config.session_resumption.handle = None - - async with Aclosing(agent_to_run.run_live(child_ctx)) as agen: - async for item in agen: - yield item - if ( - event.content - and event.content.parts - and event.content.parts[0].function_response - and event.content.parts[0].function_response.name - == 'task_completed' - ): - # this is used for sequential agent to signal the end of the agent. - await asyncio.sleep(DEFAULT_TASK_COMPLETION_DELAY) - # cancel the tasks that belongs to the closed connection. - send_task.cancel() - return - finally: - # Clean up - if not send_task.done(): - send_task.cancel() - try: - await send_task - except asyncio.CancelledError: - pass - if should_reconnect: - continue - break - except (ConnectionClosed, ConnectionClosedOK) as e: - # If we have a session resumption handle, we attempt to reconnect. - # This handle is updated dynamically during the session. - if invocation_context.live_session_resumption_handle: - if attempt > DEFAULT_MAX_RECONNECT_ATTEMPTS: - logger.error('Max reconnection attempts reached (%s).', e) - raise - logger.info( - 'Connection closed (%s), reconnecting with session handle.', e - ) - continue - logger.error('Connection closed: %s.', e) - raise - except errors.APIError as e: - # Error code 1000 and 1006 indicates a recoverable connection drop. - # In that case, we attempt to reconnect with session handle if available. - if e.code in [1000, 1006]: + invocation_context.live_request_queue.send_content( + event.content + ) + # We handle agent transfer here in `run_live` rather than + # in `_postprocess_live` to prevent duplication of function + # response processing. If agent transfer were handled in + # `_postprocess_live`, events yielded from child agent's + # `run_live` would bubble up to parent agent's `run_live`, + # causing `event.get_function_responses()` to be true in both + # child and parent, and `send_content()` to be called twice for + # the same function response. By handling agent transfer here, + # we ensure that only child agent processes its own function + # responses after the transfer. + # + # The transfer is gated on the `transfer_to_agent` action + # rather than on the position of the `transfer_to_agent` + # function response: the model may issue the transfer alongside + # other function calls, whose responses are merged into a + # single event in call order, so the transfer response is not + # necessarily `parts[0]`. Gating on the action matches + # `_postprocess_handle_function_calls_async`, and also covers + # tools that request a transfer by setting the action directly + # instead of calling `transfer_to_agent`. + transfer_to_agent = event.actions.transfer_to_agent + if transfer_to_agent: + await asyncio.sleep(DEFAULT_TRANSFER_AGENT_DELAY) + # cancel the tasks that belongs to the closed connection. + send_task.cancel() + logger.debug('Closing live connection') + await llm_connection.close() + logger.debug('Live connection closed.') + # Stop this agent's background tools before the child agent + # starts reading and writing the live request queue. + await self._stop_background_tool_tasks(invocation_context) + # transfer to the sub agent. + logger.debug('Transferring to agent: %s', transfer_to_agent) + agent_to_run = self._get_agent_to_run( + invocation_context, transfer_to_agent + ) + child_ctx = invocation_context.model_copy() + # Child Live agent should start a new Live session. + # Do not reuse the parent session's resumption handle. + child_ctx.live_session_resumption_handle = None + + if child_ctx.run_config: + child_ctx.run_config = child_ctx.run_config.model_copy( + deep=True + ) + if child_ctx.run_config.session_resumption: + child_ctx.run_config.session_resumption.handle = None + + async with Aclosing( + agent_to_run.run_live(child_ctx) + ) as agen: + async for item in agen: + yield item + if ( + event.content + and event.content.parts + and event.content.parts[0].function_response + and event.content.parts[0].function_response.name + == 'task_completed' + ): + # this is used for sequential agent to signal the end of the agent. + await asyncio.sleep(DEFAULT_TASK_COMPLETION_DELAY) + # cancel the tasks that belongs to the closed connection. + send_task.cancel() + return + finally: + # Clean up + if not send_task.done(): + send_task.cancel() + try: + await send_task + except asyncio.CancelledError: + pass + if should_reconnect: + continue + break + except (ConnectionClosed, ConnectionClosedOK) as e: + # If we have a session resumption handle, we attempt to reconnect. + # This handle is updated dynamically during the session. if invocation_context.live_session_resumption_handle: if attempt > DEFAULT_MAX_RECONNECT_ATTEMPTS: logger.error('Max reconnection attempts reached (%s).', e) raise logger.info( - 'Connection lost (%s), reconnecting with session handle.', e + 'Connection closed (%s), reconnecting with session handle.', e ) continue - logger.error('APIError in live flow: %s', e) - raise - except Exception as e: + logger.error('Connection closed: %s.', e) + raise + except errors.APIError as e: + # Error code 1000 and 1006 indicates a recoverable connection drop. + # In that case, we attempt to reconnect with session handle if available. + if e.code in [1000, 1006]: + if invocation_context.live_session_resumption_handle: + if attempt > DEFAULT_MAX_RECONNECT_ATTEMPTS: + logger.error('Max reconnection attempts reached (%s).', e) + raise + logger.info( + 'Connection lost (%s), reconnecting with session handle.', e + ) + continue + logger.error('APIError in live flow: %s', e) + raise + except Exception as e: + logger.error( + 'An unexpected error occurred in live flow: %s', e, exc_info=True + ) + raise + + finally: + await self._stop_background_tool_tasks(invocation_context) + + async def _stop_background_tool_tasks( + self, invocation_context: InvocationContext + ) -> None: + """Cancels the background tool tasks this live run started. + + A live run starts two kinds of tools as bare asyncio tasks: streaming + tools (``active_streaming_tools``) and non-blocking tools + (``active_non_blocking_tool_tasks``). Nothing tied either to the lifetime + of the run that started it — only an explicit ``stop_streaming`` call ever + cancelled one — so a tool kept running after its agent was done, feeding + function responses into a live request queue that by then belonged to + another agent, or to nobody at all. + + The tools stop when the run that started them ends, whether that is a + handoff to another agent, ``task_completed``, the connection closing, or + the caller walking away. Tying this to the agent run rather than to the + whole invocation is what keeps a tool from reaching the model of the + agent that comes after it. + + Cancellation is best effort: a task that does not stop within + ``_TOOL_SHUTDOWN_TIMEOUT_SECONDS`` is logged and left behind rather than + stalling the handoff or the caller's teardown on it. + """ + tasks = [ + active.task + for active in (invocation_context.active_streaming_tools or {}).values() + if active.task is not None + ] + tasks.extend( + (invocation_context.active_non_blocking_tool_tasks or {}).values() + ) + pending = [task for task in tasks if not task.done()] + if not pending: + return + + logger.debug('Stopping %d background tool task(s).', len(pending)) + for task in pending: + task.cancel() + stopped, still_running = await asyncio.wait( + pending, timeout=_TOOL_SHUTDOWN_TIMEOUT_SECONDS + ) + for task in still_running: + logger.warning( + 'Tool task %s ignored cancellation and outlives its agent.', + task.get_name(), + ) + for task in stopped: + # A tool reports its own failures to the model, so an exception here is + # unexpected. Retrieve it anyway: an unread one is reported by asyncio + # itself, out of context, when the task is garbage collected. + if not task.cancelled() and task.exception() is not None: logger.error( - 'An unexpected error occurred in live flow: %s', e, exc_info=True + 'Tool task %s failed.', task.get_name(), exc_info=task.exception() ) - raise + + # Retire the registry entries: the run is over, so nothing it started is + # current any more, whether or not the task honored the cancellation. + # (``stop_streaming`` blanks an entry's fields and keeps the key, because + # the model it answers to is still running and may ask again. Here nobody + # is coming back for it.) Letting go of the streams is what matters most: + # ``_send_to_model`` copies every live request into each registered + # stream, so one left behind by a tool that no longer reads it grows for + # as long as the session lasts, an entry per audio chunk the user speaks. + if invocation_context.active_streaming_tools: + invocation_context.active_streaming_tools.clear() + # A non-blocking tool drops its own entry in its `finally`, so that one is + # usually empty already; it has something to remove only when the task + # never got there, because it ignored the cancellation or died first. + if invocation_context.active_non_blocking_tool_tasks: + invocation_context.active_non_blocking_tool_tasks.clear() async def _send_to_model( self, diff --git a/tests/unittests/streaming/test_live_tool_shutdown.py b/tests/unittests/streaming/test_live_tool_shutdown.py new file mode 100644 index 00000000000..ce09327f5af --- /dev/null +++ b/tests/unittests/streaming/test_live_tool_shutdown.py @@ -0,0 +1,471 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests that a live agent stops its background tools when its run ends. + +Live mode runs two kinds of tools as background tasks that outlive the model +turn that started them: streaming tools and non-blocking tools. Both belong to +the agent run that started them, and both stop when it ends -- including when +it ends by handing off to another agent, which is when the next agent takes +over the live request queue they write to. +""" + +from __future__ import annotations + +import asyncio +from contextlib import aclosing +from typing import Any +from typing import AsyncGenerator +from typing import Callable + +from google.adk.agents.active_streaming_tool import ActiveStreamingTool +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.live_request_queue import LiveRequestQueue +from google.adk.agents.llm_agent import Agent +from google.adk.agents.run_config import RunConfig +from google.adk.events.event import Event +from google.adk.flows.llm_flows import base_llm_flow +from google.adk.flows.llm_flows.single_flow import SingleFlow +from google.adk.models.llm_response import LlmResponse +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.tools.function_tool import FunctionTool +from google.adk.tools.tool_context import ToolContext +from google.genai import types +import pytest + +from .. import testing_utils + +_MONITOR = 'monitor' +_MAX_EVENTS = 50 +# Slow enough that a monitor left running is unmistakable, yet quick enough +# that a turn produces only a handful of ticks before it ends. +_TICK_SECONDS = 0.1 + + +def _call(name: str) -> LlmResponse: + """A model turn that calls ``name`` with no arguments.""" + return LlmResponse( + content=types.Content( + role='model', + parts=[types.Part.from_function_call(name=name, args={})], + ), + turn_complete=False, + ) + + +async def _run_live_turn( + tools: list[Any], + *, + calls: list[str], + stop_when: Callable[[list[Event]], bool] | None = None, + timeout: float = 5.0, +) -> tuple[list[Event], bool]: + """Runs a live turn in which the model makes ``calls``, in order. + + Returns the events the caller saw and whether the stream ended on its own. + It does not end on its own if ``stop_when`` asked to stop early, or if the + turn is still producing events once ``timeout`` elapses. + """ + agent = Agent( + name='root_agent', + model=testing_utils.MockModel.create([_call(name) for name in calls]), + tools=tools, + ) + session_service = InMemorySessionService() + session = await session_service.create_session(app_name='app', user_id='u') + runner = Runner(app_name='app', agent=agent, session_service=session_service) + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + types.Blob(data=b'question', mime_type='audio/pcm') + ) + + events: list[Event] = [] + ended = False + + async def _consume() -> None: + nonlocal ended + async with aclosing( + runner.run_live( + user_id='u', + session_id=session.id, + live_request_queue=live_request_queue, + run_config=RunConfig(response_modalities=['TEXT']), + ) + ) as agen: + async for event in agen: + events.append(event) + if len(events) >= _MAX_EVENTS or (stop_when and stop_when(events)): + return + ended = True + + try: + await asyncio.wait_for(_consume(), timeout=timeout) + except asyncio.TimeoutError: + pass + # Let the teardown that the generator's closure kicked off finish. + for _ in range(10): + await asyncio.sleep(0) + return events, ended + + +def _task_completed() -> str: + """The signal a live agent uses to end its own turn.""" + return 'done' + + +@pytest.mark.asyncio +async def test_teardown_empties_both_registries( + monkeypatch: pytest.MonkeyPatch, +): + """Neither registry keeps a tool of a run that is over. + + Tools that stop on request retire themselves, so this uses two that refuse + to: what is left behind is exactly what teardown has to sweep up. + """ + monkeypatch.setattr(base_llm_flow, '_TOOL_SHUTDOWN_TIMEOUT_SECONDS', 0.05) + + def refuses_to_stop() -> Any: + """Ignores the first cancellation; honors the second, so this test ends.""" + swallowed = False + + async def run() -> None: + nonlocal swallowed + while True: + try: + await asyncio.sleep(0.01) + except asyncio.CancelledError: + if swallowed: + raise + swallowed = True + + return run() + + streaming_task = asyncio.create_task(refuses_to_stop()) + non_blocking_task = asyncio.create_task(refuses_to_stop()) + await asyncio.sleep(0) + + invocation_context = await testing_utils.create_invocation_context( + agent=Agent(name='agent', model=testing_utils.MockModel.create([])) + ) + invocation_context.active_streaming_tools = { + _MONITOR: ActiveStreamingTool( + task=streaming_task, stream=LiveRequestQueue() + ) + } + invocation_context.active_non_blocking_tool_tasks = { + 'lookup_1': non_blocking_task + } + + await SingleFlow()._stop_background_tool_tasks(invocation_context) + + assert not invocation_context.active_streaming_tools + assert not invocation_context.active_non_blocking_tool_tasks + + # The registry no longer holds them, so this test owns their disposal: a + # task left pending here would stall the event loop's shutdown. + for task in (streaming_task, non_blocking_task): + task.cancel() + await asyncio.gather( + streaming_task, non_blocking_task, return_exceptions=True + ) + assert streaming_task.done() and non_blocking_task.done() + + +@pytest.mark.asyncio +async def test_streaming_tool_stops_when_its_agent_hands_off(): + """A handoff ends the agent's run, so its background tools end with it. + + The sub agent takes over the live request queue: a tool still running for + the previous agent would push function responses at a model that never + called it. + """ + tasks: list[asyncio.Task[Any]] = [] + ticks = 0 + seen_by_sub_agent: dict[str, Any] = {} + + async def monitor() -> AsyncGenerator[Any, None]: + nonlocal ticks + tasks.append(asyncio.current_task()) + while True: + ticks += 1 + yield {'tick': ticks} + await asyncio.sleep(_TICK_SECONDS) + + def report() -> str: + """Records, from inside the sub agent, what the handoff left running.""" + seen_by_sub_agent['monitor_stopped'] = tasks[0].done() + seen_by_sub_agent['ticks'] = ticks + return 'reported' + + sub_agent = Agent( + name='sub_agent', + model=testing_utils.MockModel.create([_call('report')]), + tools=[report], + ) + root_agent = Agent( + name='root_agent', + model=testing_utils.MockModel.create([ + _call(_MONITOR), + LlmResponse( + content=types.Content( + role='model', + parts=[ + types.Part.from_function_call( + name='transfer_to_agent', + args={'agent_name': 'sub_agent'}, + ) + ], + ), + turn_complete=False, + ), + ]), + tools=[monitor], + sub_agents=[sub_agent], + ) + + session_service = InMemorySessionService() + session = await session_service.create_session(app_name='app', user_id='u') + runner = Runner( + app_name='app', agent=root_agent, session_service=session_service + ) + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + types.Blob(data=b'question', mime_type='audio/pcm') + ) + + async def _consume() -> None: + async with aclosing( + runner.run_live( + user_id='u', + session_id=session.id, + live_request_queue=live_request_queue, + run_config=RunConfig(response_modalities=['TEXT']), + ) + ) as agen: + seen = 0 + async for _ in agen: + seen += 1 + # Stop once the sub agent has run, or the replaying mock loops. + if 'monitor_stopped' in seen_by_sub_agent or seen >= _MAX_EVENTS: + return + + try: + await asyncio.wait_for(_consume(), timeout=10.0) + except asyncio.TimeoutError: + pass + + assert seen_by_sub_agent.get('monitor_stopped'), ( + 'the monitor was still running while the sub agent held the live' + ' request queue' + ) + # It stopped at the handoff, not merely by the end of the session. + await asyncio.sleep(_TICK_SECONDS * 3) + assert ticks == seen_by_sub_agent['ticks'] + + +@pytest.mark.asyncio +async def test_handoff_stops_feeding_the_stopped_tools_stream(): + """A stopped tool's stream is dropped, not left collecting live input. + + ``_send_to_model`` duplicates every live request into each registered + stream, so an entry left behind after the tool is gone grows for the rest of + the session -- one entry per audio chunk the user speaks. + """ + contexts: list[InvocationContext] = [] + handed_off = asyncio.Event() + + async def monitor( + tool_context: ToolContext, input_stream: LiveRequestQueue + ) -> AsyncGenerator[Any, None]: + # Declaring `input_stream` is what gets this tool a dedicated queue. + contexts.append(tool_context._invocation_context) + while True: + await input_stream.get() + yield {'saw': 'input'} + + def report() -> str: + handed_off.set() + return 'sub agent is live' + + sub_agent = Agent( + name='sub_agent', + model=testing_utils.MockModel.create([_call('report')]), + tools=[report], + ) + root_agent = Agent( + name='root_agent', + model=testing_utils.MockModel.create([ + _call(_MONITOR), + LlmResponse( + content=types.Content( + role='model', + parts=[ + types.Part.from_function_call( + name='transfer_to_agent', + args={'agent_name': 'sub_agent'}, + ) + ], + ), + turn_complete=False, + ), + ]), + tools=[monitor], + sub_agents=[sub_agent], + ) + + session_service = InMemorySessionService() + session = await session_service.create_session(app_name='app', user_id='u') + runner = Runner( + app_name='app', agent=root_agent, session_service=session_service + ) + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + types.Blob(data=b'question', mime_type='audio/pcm') + ) + + async def _consume() -> None: + async with aclosing( + runner.run_live( + user_id='u', + session_id=session.id, + live_request_queue=live_request_queue, + run_config=RunConfig(response_modalities=['TEXT']), + ) + ) as agen: + seen = 0 + async for _ in agen: + seen += 1 + if handed_off.is_set(): + # The user keeps talking while the sub agent is in charge. + for _ in range(25): + live_request_queue.send_realtime( + types.Blob(data=b'...', mime_type='audio/pcm') + ) + await asyncio.sleep(0.005) + return + if seen >= _MAX_EVENTS: + return + + try: + await asyncio.wait_for(_consume(), timeout=10.0) + except asyncio.TimeoutError: + pass + + assert _MONITOR not in (contexts[0].active_streaming_tools or {}), ( + 'the stopped tool is still registered, so every live request the user' + ' sends for the rest of the session is copied into its stream' + ) + + +@pytest.mark.asyncio +async def test_streaming_tool_stops_when_the_live_turn_ends(): + """A streaming tool that never stops on its own is stopped for it.""" + tasks: list[asyncio.Task[Any]] = [] + ticks = 0 + started = asyncio.Event() + + async def monitor() -> AsyncGenerator[Any, None]: + nonlocal ticks + tasks.append(asyncio.current_task()) + started.set() + while True: + ticks += 1 + yield {'tick': ticks} + await asyncio.sleep(_TICK_SECONDS) + + async def task_completed() -> str: + # Ends the turn only once the monitor is up, so the turn cannot end + # before there is anything to stop. + await started.wait() + return _task_completed() + + events, ended = await _run_live_turn( + [monitor, task_completed], calls=[_MONITOR, 'task_completed'] + ) + + assert ended, ( + 'the live stream never ended: the streaming tool kept producing after' + f' the agent turn was over. Saw: {len(events)} events.' + ) + assert tasks[0].done() + # And it really is stopped, not merely between ticks. + ticks_at_the_end = ticks + await asyncio.sleep(_TICK_SECONDS * 3) + assert ticks == ticks_at_the_end + + +@pytest.mark.asyncio +async def test_streaming_tool_stops_when_the_caller_stops_listening(): + """Abandoning the stream stops the tool too, rather than leaking it.""" + tasks: list[asyncio.Task[Any]] = [] + ticks = 0 + started = asyncio.Event() + + async def monitor() -> AsyncGenerator[Any, None]: + nonlocal ticks + tasks.append(asyncio.current_task()) + started.set() + while True: + ticks += 1 + yield {'tick': ticks} + await asyncio.sleep(_TICK_SECONDS) + + async def sync() -> str: + # Answers only once the monitor is up, so the event that makes the caller + # walk away cannot arrive before there is something to leak. + await started.wait() + return 'ok' + + _, ended = await _run_live_turn( + [monitor, sync], + calls=[_MONITOR, 'sync'], + stop_when=lambda _: started.is_set(), + ) + + assert not ended # The caller walked away mid-stream. + assert tasks[0].done() + ticks_at_the_end = ticks + await asyncio.sleep(_TICK_SECONDS * 3) + assert ticks == ticks_at_the_end + + +@pytest.mark.asyncio +async def test_non_blocking_tool_stops_when_the_live_turn_ends(): + """A non-blocking tool's task is cancelled with the invocation.""" + started = asyncio.Event() + cancelled = asyncio.Event() + + async def slow_lookup() -> str: + started.set() + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + cancelled.set() + raise + return 'never' + + async def task_completed() -> str: + await started.wait() + return _task_completed() + + scheduled = FunctionTool(func=slow_lookup) + scheduled.response_scheduling = types.FunctionResponseScheduling.SILENT + + _, ended = await _run_live_turn( + [scheduled, task_completed], calls=['slow_lookup', 'task_completed'] + ) + + assert ended + assert cancelled.is_set() diff --git a/tests/unittests/streaming/test_streaming.py b/tests/unittests/streaming/test_streaming.py index 409243a09e7..eef5c0974a3 100644 --- a/tests/unittests/streaming/test_streaming.py +++ b/tests/unittests/streaming/test_streaming.py @@ -1678,6 +1678,7 @@ def test_input_streaming_tool_stream_recreated_after_stop(): ) call_count = 0 + streams: list[LiveRequestQueue] = [] async def monitor_video( input_stream: LiveRequestQueue, @@ -1685,6 +1686,7 @@ async def monitor_video( """Simulate an input-streaming tool that tracks invocation count.""" nonlocal call_count call_count += 1 + streams.append(input_stream) yield f"started (call {call_count})" while True: await asyncio.sleep(0.1) @@ -1730,13 +1732,18 @@ def capturing_create(*args, **kwargs) -> Any: call_names.count("monitor_video") >= 2 ), f"Expected monitor_video called at least twice, got: {call_names}" - # After re-invocation, stream should be set again (not None). + # After re-invocation the tool is handed a stream again, and a fresh one: + # not the queue that stop_streaming tore down. This is asserted from what + # the tool was passed rather than from the registry, because the registry + # entry is retired when the live run ends. assert captured_child_context is not None - active_tools = captured_child_context.active_streaming_tools or {} - assert "monitor_video" in active_tools assert ( - active_tools["monitor_video"].stream is not None - ), "Expected .stream to be recreated after stop + re-invocation" + len(streams) >= 2 + ), f"Expected monitor_video to run twice, got {len(streams)} stream(s)" + assert all(stream is not None for stream in streams) + assert ( + streams[0] is not streams[1] + ), "Expected a new stream to be created after stop + re-invocation" def test_async_gen_with_input_stream_wrong_annotation_gets_no_stream(): From 0c35638ef614f53e50cf27299cfed1ef274f66c3 Mon Sep 17 00:00:00 2001 From: Liang Wu <18244712+wuliang229@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:54:37 +0000 Subject: [PATCH 2/3] feat(live): use RunConfig.session_resumption.handle when opening a live session A caller that resumes an earlier live session by passing a handle in `RunConfig.session_resumption.handle` only got that handle onto the wire. The basic request processor forwards it to `LiveConnectConfig`, but every other part of the run keys off `InvocationContext.live_session_resumption_handle`, which was populated exclusively from a server-issued `session_resumption_update`. The run therefore behaved as if the session were new: it replayed the whole conversation through `send_history()` even though the server already held that state, declared that history as `initial_history_in_client_content`, left `transparent` unset on the Vertex AI backend, and raised instead of reconnecting when the socket dropped before the server issued its first handle. Seed `InvocationContext.live_session_resumption_handle` in `BaseLlmFlow.run_live` from the handle request assembly has already put on `llm_request.live_connect_config.session_resumption`, before the connect loop, so the first connection is treated as a resumption in the same way as any mid-session reconnect. Reading the assembled request rather than the `RunConfig` keeps the seed on the same object the reconnect path goes on to write, and honors a handle set by any request processor rather than only one set through `RunConfig`. Agent transfer is unaffected because it already clears both the invocation handle and the deep-copied run config handle, so a child agent still starts a fresh live session. Co-authored-by: Liang Wu PiperOrigin-RevId: 966066295 --- .../adk/flows/llm_flows/base_llm_flow.py | 30 +- .../flows/llm_flows/test_base_llm_flow.py | 352 ++++++++++++++++++ 2 files changed, 379 insertions(+), 3 deletions(-) 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 8291e646ee5..40c47354f55 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -534,8 +534,32 @@ async def run_live( llm_request, ) - attempt = 1 try: + # A caller can resume an earlier live session by handing the flow a + # handle it obtained from a previous run, which request assembly has by + # now put on the connect config (from `RunConfig.session_resumption`). + # Seed the invocation with it so the very first connection is treated as + # a resumption like any mid-session reconnect: the history the server + # already holds is not replayed, and if this connection drops before the + # server has pushed its first `session_resumption_update`, the reconnect + # path still has the caller's handle to retry with instead of failing the + # run. Without this the handle only reached the connect config while the + # rest of the run still behaved as if the session were new. + session_resumption = ( + llm_request.live_connect_config.session_resumption + if llm_request.live_connect_config + else None + ) + if ( + not invocation_context.live_session_resumption_handle + and session_resumption is not None + and session_resumption.handle + ): + invocation_context.live_session_resumption_handle = ( + session_resumption.handle + ) + + attempt = 1 while True: try: # On subsequent attempts, use the saved token to reconnect @@ -730,9 +754,9 @@ async def run_live( logger.error('Connection closed: %s.', e) raise except errors.APIError as e: - # Error code 1000 and 1006 indicates a recoverable connection drop. + # Error code 1000, 1006 and 1011 indicates a recoverable connection drop. # In that case, we attempt to reconnect with session handle if available. - if e.code in [1000, 1006]: + if e.code in [1000, 1006, 1011]: if invocation_context.live_session_resumption_handle: if attempt > DEFAULT_MAX_RECONNECT_ATTEMPTS: logger.error('Max reconnection attempts reached (%s).', e) 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 6cc86066410..cb4de478d08 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -682,6 +682,358 @@ async def mock_receive(): mock_connection.send_history.assert_not_called() +async def _mock_preprocess_basic(ctx, req): + """Preprocess stub that runs only the real live connect config assembly. + + `BaseLlmFlow` carries no request processors of its own, so without this the + RunConfig never reaches `llm_request.live_connect_config` and a test cannot + exercise anything that reads from it. + """ + from google.adk.flows.llm_flows.basic import _build_basic_request + + _build_basic_request(ctx, req) + if False: # pylint: disable=using-constant-test + yield + + +async def _mock_preprocess_with_history(ctx, req): + """Preprocess stub that seeds history and builds the live connect config.""" + from google.adk.flows.llm_flows.basic import _build_basic_request + + req.contents = [types.Content(parts=[types.Part.from_text(text='history')])] + _build_basic_request(ctx, req) + if False: # pylint: disable=using-constant-test + yield + + +@pytest.mark.asyncio +async def test_run_live_resumes_from_run_config_handle(): + """A caller-supplied RunConfig handle starts the session as a resumption.""" + + real_model = Gemini() + mock_connection = mock.AsyncMock() + + agent = Agent(name='test_agent', model=real_model) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, + run_config=RunConfig( + session_resumption=types.SessionResumptionConfig( + handle='caller_handle' + ) + ), + ) + invocation_context.live_request_queue = LiveRequestQueue() + + flow = BaseLlmFlowForTesting() + + with mock.patch.object( + flow, '_preprocess_async', side_effect=_mock_preprocess_with_history + ): + with mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock): + + class StopError(Exception): + pass + + async def mock_receive(): + yield LlmResponse( + content=types.Content(parts=[types.Part.from_text(text='hi')]) + ) + raise StopError('stop') + + mock_connection.receive = mock.Mock(side_effect=mock_receive) + + with mock.patch( + 'google.adk.models.google_llm.Gemini.connect' + ) as mock_connect: + mock_connect.return_value.__aenter__.return_value = mock_connection + + try: + async for _ in flow.run_live(invocation_context): + pass + except StopError: + pass + + # The handle is adopted by the invocation, so the rest of the run treats + # this session as resumed. + assert invocation_context.live_session_resumption_handle == 'caller_handle' + assert mock_connect.call_count == 1 + connect_request = mock_connect.call_args[0][0] + assert ( + connect_request.live_connect_config.session_resumption.handle + == 'caller_handle' + ) + # The server already holds the conversation, so history is neither replayed + # nor declared as client-provided initial history. + mock_connection.send_history.assert_not_called() + assert connect_request.live_connect_config.history_config is None + + +@pytest.mark.asyncio +async def test_run_live_without_run_config_handle_still_sends_history(): + """A resumption config carrying no handle does not start a resumed session.""" + + real_model = Gemini() + mock_connection = mock.AsyncMock() + + agent = Agent(name='test_agent', model=real_model) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, + run_config=RunConfig( + session_resumption=types.SessionResumptionConfig(transparent=True) + ), + ) + invocation_context.live_request_queue = LiveRequestQueue() + + flow = BaseLlmFlowForTesting() + + with mock.patch.object( + flow, '_preprocess_async', side_effect=_mock_preprocess_with_history + ): + with mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock): + + class StopError(Exception): + pass + + async def mock_receive(): + yield LlmResponse( + content=types.Content(parts=[types.Part.from_text(text='hi')]) + ) + raise StopError('stop') + + mock_connection.receive = mock.Mock(side_effect=mock_receive) + + with mock.patch( + 'google.adk.models.google_llm.Gemini.connect' + ) as mock_connect: + mock_connect.return_value.__aenter__.return_value = mock_connection + + try: + async for _ in flow.run_live(invocation_context): + pass + except StopError: + pass + + assert invocation_context.live_session_resumption_handle is None + mock_connection.send_history.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'drop', + [ + 'connection_closed', + 'api_error_1000', + 'api_error_1006', + 'api_error_1011', + ], +) +async def test_run_live_reconnects_with_run_config_handle_before_first_update( + drop, +): + """A caller-supplied handle lets the first connection drop be recovered. + + Without it there is no handle until the server sends one, so a drop on the + very first connection has nothing to reconnect with. `ConnectionClosed` and + the 1006/1011 API errors then propagate, and a 1000 API error is read as a + clean end-of-session and silently ends the stream, which is the more + damaging outcome because the caller sees a truncated session rather than an + error. Both gates on the handle are separate branches, so cover each drop. + """ + from google.genai.errors import APIError + + real_model = Gemini() + mock_connection = mock.AsyncMock() + + def _raise_drop(): + if drop == 'connection_closed': + raise ConnectionClosed(None, None) + raise APIError(int(drop.removeprefix('api_error_')), {}) + + async def mock_receive(): + # The connection drops before the server ever issues its own handle. + _raise_drop() + yield # pylint: disable=unreachable + + mock_connection.receive = mock.Mock(side_effect=mock_receive) + + agent = Agent(name='test_agent', model=real_model) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, + run_config=RunConfig( + session_resumption=types.SessionResumptionConfig( + handle='caller_handle' + ) + ), + ) + invocation_context.live_request_queue = LiveRequestQueue() + + flow = BaseLlmFlowForTesting() + + with ( + mock.patch.object( + flow, '_preprocess_async', side_effect=_mock_preprocess_basic + ), + mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock), + ): + mock_connection_2 = mock.AsyncMock() + + class NonRetryableError(Exception): + pass + + async def mock_receive_2(): + yield LlmResponse( + content=types.Content(parts=[types.Part.from_text(text='hi')]) + ) + raise NonRetryableError('stop') + + mock_connection_2.receive = mock.Mock(side_effect=mock_receive_2) + + mock_aenter = mock.AsyncMock() + mock_aenter.side_effect = [mock_connection, mock_connection_2] + + with mock.patch( + 'google.adk.models.google_llm.Gemini.connect' + ) as mock_connect: + mock_connect.return_value.__aenter__ = mock_aenter + + try: + async for _ in flow.run_live(invocation_context): + pass + except NonRetryableError: + pass + + assert mock_connect.call_count == 2 + assert invocation_context.live_session_resumption_handle == ( + 'caller_handle' + ) + + +@pytest.mark.asyncio +async def test_run_live_run_config_handle_sets_transparent_on_vertex(): + """A caller-supplied handle gets transparent defaulted on the Vertex backend.""" + + real_model = Gemini() + mock_connection = mock.AsyncMock() + + agent = Agent(name='test_agent', model=real_model) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, + run_config=RunConfig( + session_resumption=types.SessionResumptionConfig( + handle='caller_handle' + ) + ), + ) + invocation_context.live_request_queue = LiveRequestQueue() + + flow = BaseLlmFlowForTesting() + + with mock.patch.object( + flow, '_preprocess_async', side_effect=_mock_preprocess_with_history + ): + with mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock): + + class StopError(Exception): + pass + + async def mock_receive(): + yield LlmResponse( + content=types.Content(parts=[types.Part.from_text(text='hi')]) + ) + raise StopError('stop') + + mock_connection.receive = mock.Mock(side_effect=mock_receive) + + with mock.patch( + 'google.adk.models.google_llm.Gemini.connect' + ) as mock_connect: + mock_connect.return_value.__aenter__.return_value = mock_connection + + with mock.patch.object( + Gemini, + '_api_backend', + new_callable=mock.PropertyMock, + return_value=GoogleLLMVariant.VERTEX_AI, + ): + try: + async for _ in flow.run_live(invocation_context): + pass + except StopError: + pass + + connect_request = mock_connect.call_args[0][0] + assert connect_request.live_connect_config.session_resumption.transparent + + +@pytest.mark.asyncio +async def test_run_live_server_handle_supersedes_run_config_handle(): + """Reconnects use the newest server handle, not the caller-supplied one.""" + + real_model = Gemini() + mock_connection = mock.AsyncMock() + + async def mock_receive(): + yield LlmResponse( + live_session_resumption_update=types.LiveServerSessionResumptionUpdate( + new_handle='server_handle' + ) + ) + raise ConnectionClosed(None, None) + + mock_connection.receive = mock.Mock(side_effect=mock_receive) + + agent = Agent(name='test_agent', model=real_model) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, + run_config=RunConfig( + session_resumption=types.SessionResumptionConfig( + handle='caller_handle' + ) + ), + ) + invocation_context.live_request_queue = LiveRequestQueue() + + flow = BaseLlmFlowForTesting() + + with mock.patch.object(flow, '_send_to_model', new_callable=AsyncMock): + mock_connection_2 = mock.AsyncMock() + + class NonRetryableError(Exception): + pass + + async def mock_receive_2(): + yield LlmResponse( + content=types.Content(parts=[types.Part.from_text(text='hi')]) + ) + raise NonRetryableError('stop') + + mock_connection_2.receive = mock.Mock(side_effect=mock_receive_2) + + mock_aenter = mock.AsyncMock() + mock_aenter.side_effect = [mock_connection, mock_connection_2] + + with mock.patch( + 'google.adk.models.google_llm.Gemini.connect' + ) as mock_connect: + mock_connect.return_value.__aenter__ = mock_aenter + + try: + async for _ in flow.run_live(invocation_context): + pass + except NonRetryableError: + pass + + assert mock_connect.call_count == 2 + assert invocation_context.live_session_resumption_handle == ( + 'server_handle' + ) + second_request = mock_connect.call_args_list[1][0][0] + assert ( + second_request.live_connect_config.session_resumption.handle + == 'server_handle' + ) + + @pytest.mark.asyncio async def test_live_session_resumption_go_away(): """Test that go_away triggers reconnection.""" From 84ad97627004c20da67979063d038af3a435c350 Mon Sep 17 00:00:00 2001 From: Liang Wu <18244712+wuliang229@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:15:55 +0000 Subject: [PATCH 3/3] fix(live): add 1011 to recoverable APIError codes for live connection --- src/google/adk/flows/llm_flows/base_llm_flow.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 3877ced4d65..40c47354f55 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -754,9 +754,9 @@ async def run_live( logger.error('Connection closed: %s.', e) raise except errors.APIError as e: - # Error code 1000 and 1006 indicates a recoverable connection drop. + # Error code 1000, 1006 and 1011 indicates a recoverable connection drop. # In that case, we attempt to reconnect with session handle if available. - if e.code in [1000, 1006]: + if e.code in [1000, 1006, 1011]: if invocation_context.live_session_resumption_handle: if attempt > DEFAULT_MAX_RECONNECT_ATTEMPTS: logger.error('Max reconnection attempts reached (%s).', e)