diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 9ac294332d1..52f908a0270 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -20,6 +20,7 @@ import logging from pathlib import Path import queue +import threading from types import TracebackType from typing import Any from typing import AsyncGenerator @@ -31,6 +32,7 @@ from typing import Optional from typing import TYPE_CHECKING import warnings +import weakref from google.genai import types from opentelemetry import context @@ -77,6 +79,37 @@ _EventQueueItem = tuple[object, asyncio.Event | None] +# Toolsets are often shared across Runners (parallel AgentTool, two Runners +# on the same agent). Holders are weak, so a Runner dropped without close() +# does not pin the toolset open for the process lifetime. +_toolset_holders: weakref.WeakKeyDictionary[ + BaseToolset, weakref.WeakSet[Runner] +] = weakref.WeakKeyDictionary() +_toolset_holders_lock = threading.Lock() + + +def _register_toolset_runner(toolset: BaseToolset, runner: Runner) -> None: + with _toolset_holders_lock: + holders = _toolset_holders.get(toolset) + if holders is None: + holders = weakref.WeakSet() + _toolset_holders[toolset] = holders + holders.add(runner) + + +def _unregister_toolset_runner(toolset: BaseToolset, runner: Runner) -> bool: + """Returns True if no live runner still holds the toolset.""" + with _toolset_holders_lock: + holders = _toolset_holders.get(toolset) + if holders is None: + return True + holders.discard(runner) + if holders: + return False + _toolset_holders.pop(toolset, None) + return True + + # Silence unused warning. # tracer is imported for backwards compatibility, to avoid breaking change in the API. _ = tracer @@ -292,6 +325,12 @@ def __init__( self._app_name_alignment_hint: Optional[str] = None self._enforce_app_name_alignment() self._warn_uncached_agent_transfer() + self._held_toolsets: set[BaseToolset] = set() + self._toolsets_released = False + if isinstance(self.agent, BaseAgent): + self._held_toolsets = self._collect_toolset(self.agent) + for toolset in self._held_toolsets: + _register_toolset_runner(toolset, self) def _require_root_agent(self) -> BaseAgent: """Returns the root as an agent for agent-only execution paths.""" @@ -2140,9 +2179,17 @@ async def _cleanup_toolsets( async def close(self) -> None: """Closes the runner.""" logger.info('Closing runner...') - # Close Toolsets - if isinstance(self.agent, BaseAgent): - await self._cleanup_toolsets(self._collect_toolset(self.agent)) + # Close Toolsets only when this runner was the last live holder. + if isinstance(self.agent, BaseAgent) and not self._toolsets_released: + self._toolsets_released = True + collected = self._collect_toolset(self.agent) + to_close = { + toolset + for toolset in collected | self._held_toolsets + if _unregister_toolset_runner(toolset, self) + } + self._held_toolsets = set() + await self._cleanup_toolsets(to_close) # Close Plugins if self.plugin_manager: diff --git a/src/google/adk/tools/agent_tool.py b/src/google/adk/tools/agent_tool.py index cb7ae50760c..08d5ed4ed48 100644 --- a/src/google/adk/tools/agent_tool.py +++ b/src/google/adk/tools/agent_tool.py @@ -317,27 +317,30 @@ async def run_async( last_content = None last_error_message = None last_grounding_metadata = None - async with Aclosing( - runner.run_async( - user_id=session.user_id, - session_id=session.id, - new_message=content, - run_config=nested_run_config, - ) - ) as agen: - async for event in agen: - # Forward state delta to parent session. - if event.actions.state_delta: - tool_context.state.update(event.actions.state_delta) - if event.error_message: - last_error_message = event.error_message - if event.content: - last_content = event.content - last_grounding_metadata = event.grounding_metadata - - # Clean up runner resources (especially MCP sessions) - # to avoid "Attempted to exit cancel scope in a different task" errors - await runner.close() + try: + async with Aclosing( + runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=content, + run_config=nested_run_config, + ) + ) as agen: + async for event in agen: + # Forward state delta to parent session. + if event.actions.state_delta: + tool_context.state.update(event.actions.state_delta) + if event.error_message: + last_error_message = event.error_message + if event.content: + last_content = event.content + last_grounding_metadata = event.grounding_metadata + finally: + # Same-task close so MCP cancel scopes unwind on this task even + # when the nested run raises. Runner.close only closes a toolset + # once no other live Runner holds it, so this is safe when the + # toolset is shared with the parent. + await runner.close() if last_content is None or last_content.parts is None: return last_error_message or '' diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index 4fa4df522b6..9ea9d524896 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -14,6 +14,7 @@ import asyncio from contextlib import aclosing +import gc import importlib import logging from pathlib import Path @@ -60,6 +61,21 @@ TEST_SESSION_ID = "test_session" +class CountingToolset(BaseToolset): + """Toolset that records how many times close() ran.""" + + def __init__(self): + super().__init__() + self.close_count = 0 + + async def get_tools(self, readonly_context=None): + del readonly_context + return [] + + async def close(self) -> None: + self.close_count += 1 + + class MockAgent(BaseAgent): """Mock agent for unit testing.""" @@ -1274,6 +1290,88 @@ async def close(self) -> None: assert toolset.close_cancelled is False assert toolset.close_finished.is_set() + @pytest.mark.asyncio + async def test_shared_toolset_stays_open_until_last_runner_closes(self): + """A toolset shared by two Runners must not close with the first one.""" + + toolset = CountingToolset() + agent = LlmAgent( + name="shared_agent", model="gemini-1.5-pro", tools=[toolset] + ) + runner_one = Runner( + app_name="test_app", + agent=agent, + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + runner_two = Runner( + app_name="test_app", + agent=agent, + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + + await runner_one.close() + assert toolset.close_count == 0 + await runner_two.close() + assert toolset.close_count == 1 + + @pytest.mark.asyncio + async def test_single_runner_still_closes_its_toolset(self): + toolset = CountingToolset() + runner = Runner( + app_name="test_app", + agent=LlmAgent( + name="solo_agent", model="gemini-1.5-pro", tools=[toolset] + ), + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + + await runner.close() + assert toolset.close_count == 1 + await runner.close() + assert toolset.close_count == 1 + + @pytest.mark.asyncio + async def test_dropped_runner_does_not_pin_shared_toolset(self): + toolset = CountingToolset() + agent = LlmAgent( + name="shared_agent", model="gemini-1.5-pro", tools=[toolset] + ) + dropped = Runner( + app_name="test_app", + agent=agent, + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + surviving = Runner( + app_name="test_app", + agent=agent, + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + del dropped + gc.collect() + + await surviving.close() + assert toolset.close_count == 1 + + @pytest.mark.asyncio + async def test_late_added_toolset_is_closed(self): + toolset = CountingToolset() + agent = LlmAgent(name="late_agent", model="gemini-1.5-pro", tools=[]) + runner = Runner( + app_name="test_app", + agent=agent, + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + agent.tools.append(toolset) + + await runner.close() + assert toolset.close_count == 1 + @pytest.mark.asyncio async def test_runner_passes_plugin_close_timeout(self): """Test that runner passes plugin_close_timeout to PluginManager.""" diff --git a/tests/unittests/tools/test_agent_tool.py b/tests/unittests/tools/test_agent_tool.py index ab34045da08..1d334df1387 100644 --- a/tests/unittests/tools/test_agent_tool.py +++ b/tests/unittests/tools/test_agent_tool.py @@ -39,6 +39,7 @@ import google.adk.runners as _runners_module from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.tools.agent_tool import AgentTool +from google.adk.tools.base_toolset import BaseToolset from google.adk.tools.tool_context import ToolContext from google.adk.utils._schema_utils import validate_node_data from google.adk.utils.variant_utils import GoogleLLMVariant @@ -50,6 +51,37 @@ from .. import testing_utils + +class CountingToolset(BaseToolset): + """Toolset that records how many times close() ran.""" + + def __init__(self): + super().__init__() + self.close_count = 0 + + async def get_tools(self, readonly_context=None): + del readonly_context + return [] + + async def close(self): + self.close_count += 1 + + +async def _tool_context_for(root_agent: Agent) -> ToolContext: + session_service = InMemorySessionService() + session = await session_service.create_session( + app_name='test_app', user_id='test_user' + ) + invocation_context = InvocationContext( + invocation_id='invocation_id', + agent=root_agent, + session=session, + session_service=session_service, + plugin_manager=PluginManager(), + ) + return ToolContext(invocation_context=invocation_context) + + function_call_custom = Part.from_function_call( name='tool_agent', args={'custom_input': 'test1'} ) @@ -921,6 +953,75 @@ async def close(self): assert parent_plugin.close_calls == 0 +@mark.asyncio +async def test_parallel_agent_tools_do_not_close_shared_toolset_early(): + toolset = CountingToolset() + started = asyncio.Event() + release = asyncio.Event() + + async def _hold(_callback_context): + started.set() + await release.wait() + return None + + mock_model = testing_utils.MockModel.create(responses=['ok-a', 'ok-b']) + agent_a = Agent(name='tool_a', model=mock_model, tools=[toolset]) + agent_b = Agent( + name='tool_b', + model=mock_model, + tools=[toolset], + before_agent_callback=_hold, + ) + root_agent = Agent(name='root_agent', model='test-model') + tool_context = await _tool_context_for(root_agent) + + slow = asyncio.create_task( + AgentTool(agent=agent_b).run_async( + args={'request': 'b'}, tool_context=tool_context + ) + ) + await started.wait() + await AgentTool(agent=agent_a).run_async( + args={'request': 'a'}, tool_context=tool_context + ) + assert toolset.close_count == 0 + release.set() + await slow + assert toolset.close_count == 1 + + +@mark.asyncio +async def test_failed_agent_tool_still_closes_shared_toolset(): + async def _boom(_callback_context): + raise RuntimeError('boom') + + toolset = CountingToolset() + boom_agent = Agent( + name='boom_agent', + model='test-model', + tools=[toolset], + before_agent_callback=_boom, + ) + ok_agent = Agent( + name='ok_agent', + model=testing_utils.MockModel.create(responses=['ok']), + tools=[toolset], + ) + root_agent = Agent(name='root_agent', model='test-model') + tool_context = await _tool_context_for(root_agent) + + with pytest.raises(RuntimeError, match='boom'): + await AgentTool(agent=boom_agent).run_async( + args={'request': 'x'}, tool_context=tool_context + ) + assert toolset.close_count == 1 + + await AgentTool(agent=ok_agent).run_async( + args={'request': 'y'}, tool_context=tool_context + ) + assert toolset.close_count == 2 + + def test_agent_tool_description_with_input_schema(): """Test that agent description is propagated when using input_schema."""