From ec6abf401019c39e8e1a8d1b2c7d5cf5e8c7ac56 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Wed, 14 Jan 2026 20:59:49 -0800 Subject: [PATCH] fix: Use canonical tools to find streaming tools and use tool.name to register them Original codes use tool.__name__ to register streaming tools, this is problemetic, it only works with python function passed as tools directly, if they are wrapped in FunctionTool, then FunctionTool doesn't have "__name__" property. canonical_tools wrap python function in FunctionTool uniformly thus we can use tool.name uniformly Co-authored-by: Xiang (Sean) Zhou PiperOrigin-RevId: 856472936 --- src/google/adk/runners.py | 17 +++++-- tests/unittests/test_runners.py | 84 +++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 5 deletions(-) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 97eb85df..64343302 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -1015,12 +1015,15 @@ class Runner: # 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 @@ class Runner: 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 ) diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index bb44ce73..c876bff5 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -23,6 +23,7 @@ from unittest.mock import AsyncMock 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.plugins.base_plugin import BasePlugin 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"