fix: track concurrent streaming tool calls

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

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,39 +100,40 @@ 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()
   )
   pending = [task for task in tasks if not task.done()]
   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)
 

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,35 +1039,33 @@ 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)
     )
 
     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.

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

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: