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 775215785fc..58df96a38b7 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() ) @@ -202,8 +201,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/tools/_caller.py b/src/google/adk/flows/llm_flows/tools/_caller.py index 985faf71f0d..f949c87b50c 100644 --- a/src/google/adk/flows/llm_flows/tools/_caller.py +++ b/src/google/adk/flows/llm_flows/tools/_caller.py @@ -995,44 +995,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}' @@ -1114,6 +1113,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) ) @@ -1121,28 +1129,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/live/_active_streaming_tool.py b/src/google/adk/live/_active_streaming_tool.py index 85661ec1aa9..bcc11a237e4 100644 --- a/src/google/adk/live/_active_streaming_tool.py +++ b/src/google/adk/live/_active_streaming_tool.py @@ -22,6 +22,7 @@ from pydantic import BaseModel from pydantic import ConfigDict +from pydantic import PrivateAttr from .live_request_queue import LiveRequestQueue @@ -31,12 +32,80 @@ class ActiveStreamingTool(BaseModel): model_config = ConfigDict( arbitrary_types_allowed=True, - extra="forbid", + extra='forbid', ) """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 model_post_init(self, __context: Any) -> None: + """Adds a task supplied through the compatibility fields.""" + del __context + if self.task is not None: + self._track_task(self.task, self.stream) + + 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 + if task not in self._task_streams: + task.add_done_callback(self._discard_task) + self._task_streams[task] = stream + + 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/tools/function_tool.py b/src/google/adk/tools/function_tool.py index 4314c4d7062..3bf681de63f 100644 --- a/src/google/adk/tools/function_tool.py +++ b/src/google/adk/tools/function_tool.py @@ -290,7 +290,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/tools/test_functions_simple.py b/tests/unittests/flows/llm_flows/tools/test_functions_simple.py index 9f359016788..171b1a09563 100644 --- a/tests/unittests/flows/llm_flows/tools/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/tools/test_functions_simple.py @@ -2536,10 +2536,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/live/test_active_streaming_tool.py b/tests/unittests/live/test_active_streaming_tool.py index f5c68c0870d..0e80ef27f14 100644 --- a/tests/unittests/live/test_active_streaming_tool.py +++ b/tests/unittests/live/test_active_streaming_tool.py @@ -55,6 +55,8 @@ async def _dummy(): assert tool.task is task assert tool.stream is queue + assert tool._active_tasks() == {task} + assert tool._active_streams() == [queue] await task @@ -62,3 +64,58 @@ def test_active_streaming_tool_extra_fields_forbidden(): """Verifies that extra attributes are rejected by pydantic configuration.""" with pytest.raises(ValidationError): ActiveStreamingTool(unexpected_arg="not_allowed") + + +@pytest.mark.asyncio +async def test_active_streaming_tool_tracks_concurrent_calls(): + """Tracks independent streams and releases each completed call.""" + release = asyncio.Event() + + async def _wait(): + await release.wait() + + first_task = asyncio.create_task(_wait()) + second_task = asyncio.create_task(_wait()) + first_stream = LiveRequestQueue() + second_stream = LiveRequestQueue() + tool = ActiveStreamingTool() + tool._track_task(first_task, first_stream) + tool._track_task(second_task, second_stream) + + assert tool._active_tasks() == {first_task, second_task} + assert tool._active_streams() == [first_stream, second_stream] + + second_task.cancel() + await asyncio.gather(second_task, return_exceptions=True) + await asyncio.sleep(0) + assert tool._active_tasks() == {first_task} + assert tool.task is first_task + assert tool.stream is first_stream + + release.set() + await first_task + await asyncio.sleep(0) + assert tool._active_tasks() == set() + assert tool._active_streams() == [] + assert tool.task is None + assert tool.stream is None + + +@pytest.mark.asyncio +async def test_discard_snapshot_preserves_later_call(): + """Discarding a stop snapshot does not remove a later registration.""" + first_task = asyncio.create_task(asyncio.sleep(60)) + second_task = asyncio.create_task(asyncio.sleep(60)) + tool = ActiveStreamingTool() + tool._track_task(first_task) + snapshot = tool._active_tasks() + tool._track_task(second_task) + + try: + tool._discard_tasks(snapshot) + assert tool._active_tasks() == {second_task} + assert tool.task is second_task + finally: + first_task.cancel() + second_task.cancel() + await asyncio.gather(first_task, second_task, return_exceptions=True) diff --git a/tests/unittests/streaming/test_live_tool_shutdown.py b/tests/unittests/streaming/test_live_tool_shutdown.py index 11b4e0fb197..ba677608679 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.live._active_streaming_tool import ActiveStreamingTool @@ -184,6 +185,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():