Merge branch 'main' into feat/allow-thinking-config-4108

Author: Akshat8510

src/google/adk/runners.py

@@ -1015,12 +1015,15 @@ async def run_live(
     # Pre-processing for live streaming tools
     # Inspect the tool's parameters to find if it uses LiveRequestQueue
     invocation_context.active_streaming_tools = {}
-    # TODO(hangfei): switch to use canonical_tools.
-    # for shell agents, there is no tools associated with it so we should skip.
-    if hasattr(invocation_context.agent, 'tools'):
+    # For shell agents, there is no canonical_tools method so we should skip.
+    if hasattr(invocation_context.agent, 'canonical_tools'):
       import inspect
 
-      for tool in invocation_context.agent.tools:
+      # Use canonical_tools to get properly wrapped BaseTool instances
+      canonical_tools = await invocation_context.agent.canonical_tools(
+          invocation_context
+      )
+      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.
         #
@@ -1044,10 +1047,14 @@ async def run_live(
           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(
                 stream=LiveRequestQueue()
             )
-            invocation_context.active_streaming_tools[tool.__name__] = (
+            invocation_context.active_streaming_tools[tool.name] = (
                 active_streaming_tool
             )
 

tests/unittests/test_runners.py

@@ -23,6 +23,7 @@
 from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.context_cache_config import ContextCacheConfig
 from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.live_request_queue import LiveRequestQueue
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.agents.run_config import RunConfig
 from google.adk.apps.app import App
@@ -34,6 +35,7 @@
 from google.adk.runners import Runner
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.sessions.session import Session
+from google.adk.tools.function_tool import FunctionTool
 from google.genai import types
 import pytest
 
@@ -358,6 +360,88 @@ async def test_run_live_auto_create_session():
   assert session is not None
 
 
+@pytest.mark.asyncio
+async def test_run_live_detects_streaming_tools_with_canonical_tools():
+  """run_live should detect streaming tools using canonical_tools and tool.name."""
+
+  # Define streaming tools - one as raw function, one wrapped in FunctionTool
+  async def raw_streaming_tool(
+      input_stream: LiveRequestQueue,
+  ) -> AsyncGenerator[str, None]:
+    """A raw streaming tool function."""
+    yield "test"
+
+  async def wrapped_streaming_tool(
+      input_stream: LiveRequestQueue,
+  ) -> AsyncGenerator[str, None]:
+    """A streaming tool wrapped in FunctionTool."""
+    yield "test"
+
+  def non_streaming_tool(param: str) -> str:
+    """A regular non-streaming tool."""
+    return param
+
+  # Create a mock LlmAgent that yields an event and captures invocation context
+  captured_context = {}
+
+  class StreamingToolsAgent(LlmAgent):
+
+    async def _run_live_impl(
+        self, invocation_context: InvocationContext
+    ) -> AsyncGenerator[Event, None]:
+      # Capture the active_streaming_tools for verification
+      captured_context["active_streaming_tools"] = (
+          invocation_context.active_streaming_tools
+      )
+      yield Event(
+          invocation_id=invocation_context.invocation_id,
+          author=self.name,
+          content=types.Content(
+              role="model", parts=[types.Part(text="streaming test")]
+          ),
+      )
+
+  agent = StreamingToolsAgent(
+      name="streaming_agent",
+      model="gemini-2.0-flash",
+      tools=[
+          raw_streaming_tool,  # Raw function
+          FunctionTool(wrapped_streaming_tool),  # Wrapped in FunctionTool
+          non_streaming_tool,  # Non-streaming tool (should not be detected)
+      ],
+  )
+
+  session_service = InMemorySessionService()
+  artifact_service = InMemoryArtifactService()
+  runner = Runner(
+      app_name="streaming_test_app",
+      agent=agent,
+      session_service=session_service,
+      artifact_service=artifact_service,
+      auto_create_session=True,
+  )
+
+  live_queue = LiveRequestQueue()
+
+  agen = runner.run_live(
+      user_id="user",
+      session_id="test_session",
+      live_request_queue=live_queue,
+  )
+
+  event = await agen.__anext__()
+  await agen.aclose()
+
+  assert event.author == "streaming_agent"
+
+  # Verify streaming tools were detected correctly
+  active_tools = captured_context.get("active_streaming_tools", {})
+  assert "raw_streaming_tool" in active_tools
+  assert "wrapped_streaming_tool" in active_tools
+  # Non-streaming tool should not be detected
+  assert "non_streaming_tool" not in active_tools
+
+
 @pytest.mark.asyncio
 async def test_runner_allows_nested_agent_directories(tmp_path, monkeypatch):
   project_root = tmp_path / "workspace"