chore: Register all streaming tool at runner

previously we only register streaming tool that accept stream input at runner, now uniformly register all streaming tool at runner.

Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com>
PiperOrigin-RevId: 869447996

Author: seanzhougoogle

src/google/adk/flows/llm_flows/functions.py

@@ -825,17 +825,11 @@ async def run_tool_and_update_queue(tool, function_args, tool_context):
         run_tool_and_update_queue(tool, function_args, tool_context)
     )
 
-    # Register streaming tool using original logic
+    # The tool is already registered in active_streaming_tools by
+    # runners.py at startup (all async-generator tools are registered
+    # there). Just attach the background task.
     async with streaming_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:
-        invocation_context.active_streaming_tools[tool.name] = (
-            ActiveStreamingTool(task=task)
-        )
+      invocation_context.active_streaming_tools[tool.name].task = task
 
     # Immediately return a pending response.
     # This is required by current live model.

src/google/adk/runners.py

@@ -1023,38 +1023,30 @@ async def run_live(
       canonical_tools = await invocation_context.agent.canonical_tools(
           invocation_context
       )
+      # Register all async-generator tools as streaming tools.
+      # A streaming tool is any tool whose underlying function is an
+      # async generator (i.e. uses `yield`). There are two sub-types:
+      #   1. Input-streaming tools: accept a `input_stream:
+      #      LiveRequestQueue` parameter to consume the live audio/video
+      #      stream. The stream is created lazily in `_call_live` when
+      #      the model actually calls the tool.
+      #   2. Output-streaming tools: async generators that yield results
+      #      over time but don't consume the live stream. They are run
+      #      as background tasks when called.
+      # Both types are registered here with `stream=None`. The
+      # distinction between them is made at call time.
       for tool in canonical_tools:
-        # We use `inspect.signature()` to examine the tool's underlying function (`tool.func`).
-        # This approach is deliberately chosen over `typing.get_type_hints()` for robustness.
-        #
-        # The Problem with `get_type_hints()`:
-        # `get_type_hints()` attempts to resolve forward-referenced (string-based) type
-        # annotations. This resolution can easily fail with a `NameError` (e.g., "Union not found")
-        # if the type isn't available in the scope where `get_type_hints()` is called.
-        # This is a common and brittle issue in framework code that inspects functions
-        # defined in separate user modules.
-        #
-        # Why `inspect.signature()` is Better Here:
-        # `inspect.signature()` does NOT resolve the annotations; it retrieves the raw
-        # annotation object as it was defined on the function. This allows us to
-        # perform a direct and reliable identity check (`param.annotation is LiveRequestQueue`)
-        # without risking a `NameError`.
         callable_to_inspect = tool.func if hasattr(tool, 'func') else tool
-        # Ensure the target is actually callable before inspecting to avoid errors.
         if not callable(callable_to_inspect):
           continue
-        for param in inspect.signature(callable_to_inspect).parameters.values():
-          if param.annotation is LiveRequestQueue:
-            if not invocation_context.active_streaming_tools:
-              invocation_context.active_streaming_tools = {}
-
-            logger.debug(
-                'Register streaming tool with input stream: %s', tool.name
-            )
-            active_streaming_tool = ActiveStreamingTool()
-            invocation_context.active_streaming_tools[tool.name] = (
-                active_streaming_tool
-            )
+        if inspect.isasyncgenfunction(callable_to_inspect):
+          if not invocation_context.active_streaming_tools:
+            invocation_context.active_streaming_tools = {}
+          logger.debug('Register streaming tool: %s', tool.name)
+          active_streaming_tool = ActiveStreamingTool()
+          invocation_context.active_streaming_tools[tool.name] = (
+              active_streaming_tool
+          )
 
     async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
       async with Aclosing(ctx.agent.run_live(ctx)) as agen:

tests/unittests/streaming/test_streaming.py

@@ -1497,3 +1497,51 @@ def capturing_method(*args, **kwargs):
   assert (
       active_tools['monitor_stock_price'].stream is None
   ), 'Expected stream to be reset to None after stop_streaming'
+
+
+def test_output_streaming_tool_registered_at_startup():
+  """Test that output-streaming tools (async generators without LiveRequestQueue) are registered at startup."""
+  response1 = LlmResponse(turn_complete=True)
+
+  mock_model = testing_utils.MockModel.create([response1])
+
+  async def monitor_stock_price(stock_symbol: str):
+    """Yield periodic price updates."""
+    yield f'price for {stock_symbol}'
+
+  root_agent = Agent(
+      name='root_agent',
+      model=mock_model,
+      tools=[monitor_stock_price],
+  )
+
+  runner = _LiveTestRunner(root_agent=root_agent)
+
+  # Capture invocation context to verify registration.
+  captured_context = None
+  original_method = runner.runner._new_invocation_context_for_live
+
+  def capturing_method(*args, **kwargs):
+    nonlocal captured_context
+    ctx = original_method(*args, **kwargs)
+    captured_context = ctx
+    return ctx
+
+  runner.runner._new_invocation_context_for_live = capturing_method
+
+  live_request_queue = LiveRequestQueue()
+  live_request_queue.send_realtime(
+      blob=types.Blob(data=b'test', mime_type='audio/pcm')
+  )
+
+  runner.run_live(live_request_queue, max_responses=1)
+
+  # Output-streaming tool should be registered with stream=None.
+  assert captured_context is not None
+  active_tools = captured_context.active_streaming_tools or {}
+  assert (
+      'monitor_stock_price' in active_tools
+  ), 'Expected output-streaming tool to be registered at startup'
+  assert (
+      active_tools['monitor_stock_price'].stream is None
+  ), 'Expected stream to be None for output-streaming tool'