diff --git a/src/google/adk/agents/active_streaming_tool.py b/src/google/adk/agents/active_streaming_tool.py index dce3a2321b9..f1000f3cb15 100644 --- a/src/google/adk/agents/active_streaming_tool.py +++ b/src/google/adk/agents/active_streaming_tool.py @@ -20,6 +20,7 @@ from pydantic import BaseModel from pydantic import ConfigDict +from pydantic import PrivateAttr from ..live.live_request_queue import LiveRequestQueue @@ -34,7 +35,68 @@ class ActiveStreamingTool(BaseModel): """The pydantic model config.""" task: Optional[asyncio.Task[Any]] = None - """The active task of this streaming tool.""" + """The most recently started task of this streaming tool.""" stream: Optional[LiveRequestQueue] = None - """The active (input) streams of this streaming tool.""" + """The input stream associated with the most recent task.""" + + _task_streams: dict[asyncio.Task[Any], LiveRequestQueue | None] = PrivateAttr( + default_factory=dict + ) + + def _track_task( + self, + task: asyncio.Task[Any], + stream: LiveRequestQueue | None = None, + ) -> None: + """Tracks one call and releases its resources when it completes.""" + self.task = task + self.stream = stream + self._task_streams[task] = stream + task.add_done_callback(self._discard_task) + + def _active_tasks(self) -> set[asyncio.Task[Any]]: + """Returns a snapshot of all running calls.""" + tasks = {task for task in self._task_streams if not task.done()} + if self.task is not None and not self.task.done(): + tasks.add(self.task) + return tasks + + def _active_streams(self) -> list[LiveRequestQueue]: + """Returns a snapshot of input streams for all running calls.""" + streams = [ + stream + for task, stream in self._task_streams.items() + if not task.done() and stream is not None + ] + if ( + not self._task_streams + and self.task is not None + and self.stream is not None + ): + streams.append(self.stream) + return streams + + def _discard_tasks(self, tasks: set[asyncio.Task[Any]]) -> None: + """Discards tracked calls without affecting calls started later.""" + for task in tasks: + self._task_streams.pop(task, None) + if not self._task_streams: + self.task = None + self.stream = None + elif self.task in tasks: + self._set_latest_task() + + def _discard_task(self, task: asyncio.Task[Any]) -> None: + self._task_streams.pop(task, None) + if self.task is task: + self._set_latest_task() + + def _set_latest_task(self) -> None: + if self._task_streams: + task = next(reversed(self._task_streams)) + self.task = task + self.stream = self._task_streams[task] + else: + self.task = None + self.stream = None diff --git a/src/google/adk/flows/llm_flows/_live_llm_flow.py b/src/google/adk/flows/llm_flows/_live_llm_flow.py index 9bbfafe7011..980b6835cf5 100644 --- a/src/google/adk/flows/llm_flows/_live_llm_flow.py +++ b/src/google/adk/flows/llm_flows/_live_llm_flow.py @@ -19,6 +19,7 @@ import asyncio import enum import logging +from typing import Any from typing import AsyncGenerator from typing import cast from typing import Optional @@ -99,11 +100,9 @@ async def stop_background_tool_tasks( ``_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: list[asyncio.Task[Any]] = [] + for active in (invocation_context.active_streaming_tools or {}).values(): + tasks.extend(active._active_tasks()) tasks.extend( (invocation_context.active_non_blocking_tool_tasks or {}).values() ) @@ -111,27 +110,30 @@ async def stop_background_tool_tasks( if not pending: return - from . import base_llm_flow + if pending: + from . import base_llm_flow - logger.debug('Stopping %d background tool task(s).', len(pending)) - for task in pending: - task.cancel() - stopped, still_running = await asyncio.wait( - pending, timeout=base_llm_flow._TOOL_SHUTDOWN_TIMEOUT_SECONDS - ) - for task in still_running: - logger.warning( - 'Tool task %s ignored cancellation and outlives its agent.', - task.get_name(), + logger.debug('Stopping %d background tool task(s).', len(pending)) + for task in pending: + task.cancel() + stopped, still_running = await asyncio.wait( + pending, timeout=base_llm_flow._TOOL_SHUTDOWN_TIMEOUT_SECONDS ) - 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( - 'Tool task %s failed.', task.get_name(), exc_info=task.exception() + 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( + 'Tool task %s failed.', + task.get_name(), + exc_info=task.exception(), + ) # Retire the registry entries: the run is over, so nothing it started is # current any more, whether or not the task honored the cancellation. @@ -202,8 +204,8 @@ async def send_to_model( for active_streaming_tool in ( invocation_context.active_streaming_tools ).values(): - if active_streaming_tool.stream: - active_streaming_tool.stream.send(live_request) + for input_stream in active_streaming_tool._active_streams(): + input_stream.send(live_request) # Yield to event loop for cooperative multitasking await asyncio.sleep(0) diff --git a/src/google/adk/flows/llm_flows/_tool_caller.py b/src/google/adk/flows/llm_flows/_tool_caller.py index 70ee2faa429..f63df5324ab 100644 --- a/src/google/adk/flows/llm_flows/_tool_caller.py +++ b/src/google/adk/flows/llm_flows/_tool_caller.py @@ -921,44 +921,43 @@ async def _process_function_live_helper( raise ValueError('stop_streaming requires a string function_name.') # Thread-safe access to active_streaming_tools async with active_tools_lock: - active_tasks = invocation_context.active_streaming_tools - active_task = ( - active_tasks[function_name].task - if active_tasks and function_name in active_tasks - else None + active_tools = invocation_context.active_streaming_tools + active_tool = ( + active_tools.get(function_name) if active_tools is not None else None ) - task = active_task if active_task and not active_task.done() else None - - if task: - task.cancel() - try: - # Wait for the task to be cancelled - await asyncio.wait_for(task, timeout=1.0) - except (asyncio.CancelledError, asyncio.TimeoutError): - # Log the specific condition - if task.cancelled(): - logging.info('Task %s was cancelled successfully', function_name) - elif task.done(): - logging.info('Task %s completed during cancellation', function_name) - else: - logging.warning( - 'Task %s might still be running after cancellation timeout', - function_name, - ) - function_response = { - 'status': f'The task is not cancelled yet for {function_name}.' - } - if not function_response: - # Clean up the reference under lock + tasks = active_tool._active_tasks() if active_tool is not None else set() + + if tasks: + for task in tasks: + task.cancel() + _, pending = await asyncio.wait(tasks, timeout=1.0) + if pending: + logging.warning( + '%d task(s) for %s might still be running after cancellation' + ' timeout', + len(pending), + function_name, + ) + function_response = { + 'status': f'The task is not cancelled yet for {function_name}.' + } + else: + logging.info( + '%d task(s) for %s stopped successfully', + len(tasks), + function_name, + ) + # Clean up references without discarding calls registered after this + # stop request took its snapshot. async with active_tools_lock: - if ( - invocation_context.active_streaming_tools - and function_name in invocation_context.active_streaming_tools - ): - invocation_context.active_streaming_tools[function_name].task = None - invocation_context.active_streaming_tools[function_name].stream = ( - None - ) + active_tools = invocation_context.active_streaming_tools + current = ( + active_tools.get(function_name) + if active_tools is not None + else None + ) + if current is not None: + current._discard_tasks(tasks) function_response = { 'status': f'Successfully stopped streaming function {function_name}' @@ -1040,6 +1039,15 @@ async def run_tool_and_update_queue( # 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. + sig = inspect.signature(streaming_tool.func) + input_stream = None + if 'input_stream' in sig.parameters and _is_live_request_queue_annotation( + sig.parameters['input_stream'] + ): + input_stream = LiveRequestQueue() + function_args = dict(function_args) + function_args['input_stream'] = input_stream + task = asyncio.create_task( run_tool_and_update_queue(streaming_tool, function_args, tool_context) ) @@ -1047,28 +1055,17 @@ async def run_tool_and_update_queue( async with active_tools_lock: if invocation_context.active_streaming_tools is None: invocation_context.active_streaming_tools = {} - if tool.name in invocation_context.active_streaming_tools: - invocation_context.active_streaming_tools[tool.name].task = task - else: + active_streaming_tool = invocation_context.active_streaming_tools.get( + tool.name + ) + if active_streaming_tool is None: # Register the streaming tool lazily when the model calls it. + active_streaming_tool = ActiveStreamingTool() invocation_context.active_streaming_tools[tool.name] = ( - ActiveStreamingTool(task=task) + active_streaming_tool ) logger.debug('Lazily registered streaming tool: %s', tool.name) - - # For input-streaming tools (those with `input_stream: - # LiveRequestQueue`), create a dedicated LiveRequestQueue so - # _send_to_model starts duplicating data to it. This also - # handles re-invocation after stop_streaming reset .stream - # to None. - sig = inspect.signature(streaming_tool.func) - if ( - 'input_stream' in sig.parameters - and _is_live_request_queue_annotation(sig.parameters['input_stream']) - ): - invocation_context.active_streaming_tools[tool.name].stream = ( - LiveRequestQueue() - ) + active_streaming_tool._track_task(task, input_stream) # Immediately return a pending response. # This is required by current live model. diff --git a/src/google/adk/tools/function_tool.py b/src/google/adk/tools/function_tool.py index 266a97c2d62..42d2a2fe98f 100644 --- a/src/google/adk/tools/function_tool.py +++ b/src/google/adk/tools/function_tool.py @@ -183,7 +183,7 @@ def _prepare_invocation_args( # 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: + if "input_stream" in valid_params and "input_stream" not in args_to_call: active_tools = tool_context._invocation_context.active_streaming_tools if ( active_tools is not None diff --git a/tests/unittests/flows/llm_flows/test_functions_simple.py b/tests/unittests/flows/llm_flows/test_functions_simple.py index 2ecd4048a37..b44783b186f 100644 --- a/tests/unittests/flows/llm_flows/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/test_functions_simple.py @@ -2472,10 +2472,10 @@ async def streaming_fn(val: str): # The raw exception text is not leaked to the model. assert 'sensitive_detail' not in function_response.response['error'] assert function_response.id == 'fc_raises' - # The task completes instead of dying with an unretrieved exception. - task = invocation_context.active_streaming_tools[tool.name].task - assert task.done() - assert task.exception() is None + # The task completes and releases its registry references. + active_tool = invocation_context.active_streaming_tools[tool.name] + await asyncio.sleep(0) + assert active_tool.task is None def _model_call_event(invocation_id: str, call_id: str) -> Event: diff --git a/tests/unittests/streaming/test_live_tool_shutdown.py b/tests/unittests/streaming/test_live_tool_shutdown.py index 66102ebea63..472c60c5cd6 100644 --- a/tests/unittests/streaming/test_live_tool_shutdown.py +++ b/tests/unittests/streaming/test_live_tool_shutdown.py @@ -35,6 +35,7 @@ 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.functions import handle_function_calls_live from google.adk.flows.llm_flows.single_flow import SingleFlow from google.adk.live import LiveRequestQueue from google.adk.models.llm_response import LlmResponse @@ -183,6 +184,90 @@ async def run() -> None: assert streaming_task.done() and non_blocking_task.done() +@pytest.mark.asyncio +@pytest.mark.parametrize('stop_via_tool', [False, True]) +async def test_all_parallel_calls_to_same_streaming_tool_stop( + stop_via_tool: bool, +): + """Both stop paths stop every call even when tool names are identical.""" + tasks: list[asyncio.Task[Any]] = [] + streams: list[LiveRequestQueue] = [] + both_started = asyncio.Event() + + async def monitor( + value: str, input_stream: LiveRequestQueue + ) -> AsyncGenerator[dict[str, str], None]: + tasks.append(asyncio.current_task()) + streams.append(input_stream) + if len(tasks) == 2: + both_started.set() + while True: + yield {'value': value} + await asyncio.sleep(60) + + tool = FunctionTool(monitor) + agent = Agent(name='agent', model=testing_utils.MockModel.create([])) + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + invocation_context.live_request_queue = LiveRequestQueue() + event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content( + parts=[ + types.Part.from_function_call( + name=tool.name, args={'value': 'first'} + ), + types.Part.from_function_call( + name=tool.name, args={'value': 'second'} + ), + ] + ), + ) + + try: + await handle_function_calls_live( + invocation_context, event, {tool.name: tool} + ) + await asyncio.wait_for(both_started.wait(), timeout=1) + assert streams[0] is not streams[1] + active_tool = invocation_context.active_streaming_tools[tool.name] + + if stop_via_tool: + + def stop_streaming(function_name: str) -> None: + pass + + stop_tool = FunctionTool(stop_streaming) + stop_event = Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + content=types.Content( + parts=[ + types.Part.from_function_call( + name='stop_streaming', + args={'function_name': tool.name}, + ) + ] + ), + ) + await handle_function_calls_live( + invocation_context, stop_event, {stop_tool.name: stop_tool} + ) + else: + await SingleFlow()._stop_background_tool_tasks(invocation_context) + + assert all(task.done() for task in tasks) + if stop_via_tool: + assert active_tool.task is None + assert active_tool.stream is None + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @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. diff --git a/tests/unittests/streaming/test_streaming.py b/tests/unittests/streaming/test_streaming.py index 4b8ce5dd159..756b1c8257a 100644 --- a/tests/unittests/streaming/test_streaming.py +++ b/tests/unittests/streaming/test_streaming.py @@ -1246,8 +1246,8 @@ def capturing_create(*args, **kwargs) -> Any: return captured_child_context.active_streaming_tools or {} -def test_input_streaming_tool_has_stream_set_at_registration(): - """Test that input-streaming tools get .stream set to a LiveRequestQueue during registration.""" +def test_completed_input_streaming_tool_releases_resources(): + """A completed input-streaming tool releases its task and stream.""" async def monitor_video_stream( input_stream: LiveRequestQueue, @@ -1259,16 +1259,9 @@ async def monitor_video_stream( monitor_video_stream, "monitor_video_stream" ) - assert ( - "monitor_video_stream" in active_tools - ), "Expected input-streaming tool to be registered when called" - # Stream should be a LiveRequestQueue, not None. - assert ( - active_tools["monitor_video_stream"].stream is not None - ), "Expected .stream to be set for input-streaming tool" - assert isinstance( - active_tools["monitor_video_stream"].stream, LiveRequestQueue - ), "Expected .stream to be a LiveRequestQueue instance" + assert "monitor_video_stream" in active_tools + assert active_tools["monitor_video_stream"].task is None + assert active_tools["monitor_video_stream"].stream is None def test_input_streaming_tool_stream_recreated_after_stop():