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:
@@ -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.
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user