fix(tools): don't close parent's plugins from AgentTool's sub-Runner

AgentTool.run_async spins up a sub-Runner per call that, with the
default include_plugins=True, shares the parent runner's plugin
instances. On exit the sub-Runner's close() tore down those shared
plugins, so long-lived plugins (e.g. observability exporters) hit the
5s plugin_close_timeout and surfaced as
"RuntimeError: Failed to close plugins: 'X': TimeoutError" on the
parent's tool call.

Add PluginManager.set_skip_closing_plugins(value); when set to True,
close() is a no-op. AgentTool calls it on the sub-Runner's
plugin_manager after construction (only when include_plugins=True) so
inherited plugins are no longer closed by the sub-Runner.

Change-Id: I94f81a33e7f6acef728855eaa9237a93a17b66ec
This commit is contained in:
thacht
2026-05-28 18:10:50 +00:00
committed by George Weale
parent e896c620f6
commit 2a68c4e746
4 changed files with 133 additions and 0 deletions
+24
View File
@@ -85,10 +85,25 @@ class PluginManager:
"""
self.plugins: List[BasePlugin] = []
self._close_timeout = close_timeout
self._skip_closing_plugins = False
if plugins:
for plugin in plugins:
self.register_plugin(plugin)
def set_skip_closing_plugins(self, value: bool) -> None:
"""Controls whether `close()` will tear down the registered plugins.
Set to True when the plugins are owned by another component (e.g. a parent
`Runner` whose plugin list this manager is sharing). When set, subsequent
calls to `close()` become a no-op so the shared plugins are not torn down
while still in use.
Args:
value: True to skip closing the plugins; False (default) to close them
normally.
"""
self._skip_closing_plugins = value
def register_plugin(self, plugin: BasePlugin) -> None:
"""Registers a new plugin.
@@ -309,10 +324,19 @@ class PluginManager:
async def close(self) -> None:
"""Calls the close method on all registered plugins concurrently.
If this manager was constructed with `skip_closing_plugins=True`, this
method is a no-op so plugins owned by another component (e.g. a parent
`Runner`) are not torn down while still in use.
Raises:
RuntimeError: If one or more plugins failed to close, containing
details of all failures.
"""
if self._skip_closing_plugins:
logger.debug(
"Skipping plugin close; plugins are owned by another component."
)
return
exceptions = {}
# We iterate sequentially to avoid creating new tasks which can cause issues
# with some libraries (like anyio/mcp) that rely on task-local context.
+6
View File
@@ -249,6 +249,12 @@ class AgentTool(BaseTool):
credential_service=tool_context._invocation_context.credential_service,
plugins=plugins,
)
# When plugins are inherited from the parent runner, the parent still owns
# them; tell the sub-Runner's plugin manager to skip closing them on exit
# so shared plugins (e.g. observability exporters) are not torn down while
# the parent is still using them.
if self.include_plugins:
runner.plugin_manager.set_skip_closing_plugins(True)
state_dict = {
k: v
@@ -317,3 +317,49 @@ async def test_close_with_timeout(plugin1: TestPlugin):
assert "Failed to close plugins: 'plugin1': TimeoutError" in str(
excinfo.value
)
@pytest.mark.asyncio
async def test_close_is_noop_after_set_skip_closing_plugins(
plugin1: TestPlugin,
):
"""Tests that close is a no-op after set_skip_closing_plugins(True)."""
plugin1.close = AsyncMock()
service = PluginManager(plugins=[plugin1])
service.set_skip_closing_plugins(True)
await service.close()
plugin1.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_set_skip_closing_plugins_bypasses_close_timeout(
plugin1: TestPlugin,
):
"""Tests that set_skip_closing_plugins(True) bypasses close timeout."""
async def slow_close():
await asyncio.sleep(10) # Would otherwise time out.
plugin1.close = slow_close
service = PluginManager(plugins=[plugin1], close_timeout=0.01)
service.set_skip_closing_plugins(True)
# Should return immediately without raising.
await service.close()
@pytest.mark.asyncio
async def test_set_skip_closing_plugins_false_reverts_to_closing(
plugin1: TestPlugin,
):
"""Tests that set_skip_closing_plugins(False) re-enables closing."""
plugin1.close = AsyncMock()
service = PluginManager(plugins=[plugin1])
service.set_skip_closing_plugins(True)
service.set_skip_closing_plugins(False)
await service.close()
plugin1.close.assert_awaited_once()
+57
View File
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
from typing import Any
from typing import Optional
@@ -30,6 +31,7 @@ from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.plugins.base_plugin import BasePlugin
from google.adk.plugins.plugin_manager import PluginManager
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.tools.agent_tool import AgentTool
from google.adk.tools.tool_context import ToolContext
@@ -727,6 +729,61 @@ def test_include_plugins_false():
assert tracking_plugin.before_agent_calls == 1
@pytest.mark.asyncio
async def test_include_plugins_true_sub_runner_does_not_close_parent_plugins():
"""Sub-Runner must not close plugins owned by the parent runner."""
class SlowClosePlugin(BasePlugin):
def __init__(self, name: str):
super().__init__(name)
self.close_calls = 0
async def close(self):
self.close_calls += 1
# Would otherwise blow past the sub-Runner's plugin_close_timeout.
await asyncio.sleep(10)
parent_plugin = SlowClosePlugin(name='parent_plugin')
mock_model = testing_utils.MockModel.create(
responses=[function_call_no_schema, 'response1', 'response2']
)
tool_agent = Agent(name='tool_agent', model=mock_model)
root_agent = Agent(
name='root_agent',
model=mock_model,
tools=[AgentTool(agent=tool_agent, include_plugins=True)],
)
runner = Runner(
app_name='test_app',
agent=root_agent,
artifact_service=InMemoryArtifactService(),
session_service=InMemorySessionService(),
memory_service=InMemoryMemoryService(),
plugins=[parent_plugin],
# Tight timeout amplifies the bug if it regresses; with the fix, the
# sub-Runner's close skips the parent's plugins entirely.
plugin_close_timeout=0.01,
)
session = await runner.session_service.create_session(
app_name='test_app', user_id='test_user'
)
# Must not raise RuntimeError("Failed to close plugins: ...") from the
# sub-Runner closing the parent's slow-to-close plugin.
async for _ in runner.run_async(
user_id=session.user_id,
session_id=session.id,
new_message=testing_utils.get_user_content('test1'),
):
pass
# The sub-Runner must not have closed the parent's plugin.
assert parent_plugin.close_calls == 0
def test_agent_tool_description_with_input_schema():
"""Test that agent description is propagated when using input_schema."""