feat: Add Graceful Plugin Shutdown to Runner

This change introduces a shutdown lifecycle hook for plugins. The `PluginManager` now has an `async def shutdown()` method that will call `await plugin.shutdown()` on any registered plugins that implement the method. This is called from `Runner.close()`, allowing plugins to perform cleanup tasks like flushing logs or closing connections when the runner instance is being closed. This improves the reliability of plugins that perform background operations.

PiperOrigin-RevId: 831037737

Author: google-genai-bot

src/google/adk/plugins/base_plugin.py

@@ -187,6 +187,14 @@ async def after_run_callback(
     """
     pass
 
+  async def close(self) -> None:
+    """Method executed when the runner is closed.
+
+    This method is used for cleanup tasks such as closing network connections
+    or releasing resources.
+    """
+    pass
+
   async def before_agent_callback(
       self, *, agent: BaseAgent, callback_context: CallbackContext
   ) -> Optional[types.Content]:

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

@@ -455,7 +455,7 @@ async def _log(self, data: dict):
     self._background_tasks.add(task)
     task.add_done_callback(self._background_tasks.discard)
 
-  async def shutdown(self):
+  async def close(self):
     """Flushes pending logs and closes client."""
     # 1. Wait for pending background logs (best effort, 2s timeout)
     if self._background_tasks:

src/google/adk/plugins/plugin_manager.py

@@ -14,7 +14,9 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
+import sys
 from typing import Any
 from typing import List
 from typing import Literal
@@ -70,13 +72,19 @@ class PluginManager:
   tool calls, or model requests.
   """
 
-  def __init__(self, plugins: Optional[List[BasePlugin]] = None):
+  def __init__(
+      self,
+      plugins: Optional[List[BasePlugin]] = None,
+      close_timeout: float = 5.0,
+  ):
     """Initializes the plugin service.
 
     Args:
       plugins: An optional list of plugins to register upon initialization.
+      close_timeout: The timeout in seconds for each plugin's close method.
     """
     self.plugins: List[BasePlugin] = []
+    self._close_timeout = close_timeout
     if plugins:
       for plugin in plugins:
         self.register_plugin(plugin)
@@ -297,3 +305,43 @@ async def _run_callbacks(
         raise RuntimeError(error_message) from e
 
     return None
+
+  async def close(self) -> None:
+    """Calls the close method on all registered plugins concurrently.
+
+    Raises:
+      RuntimeError: If one or more plugins failed to close, containing
+        details of all failures.
+    """
+    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.
+    for plugin in self.plugins:
+      try:
+        if sys.version_info >= (3, 11):
+          async with asyncio.timeout(self._close_timeout):
+            await plugin.close()
+        else:
+          # For Python < 3.11, we use wait_for which creates a new task.
+          # This might still cause issues with task-local contexts, but
+          # asyncio.timeout is not available.
+          await asyncio.wait_for(plugin.close(), timeout=self._close_timeout)
+      except Exception as e:
+        exceptions[plugin.name] = e
+        if isinstance(e, (asyncio.TimeoutError, asyncio.CancelledError)):
+          logger.warning(
+              "Timeout/Cancelled while closing plugin: %s", plugin.name
+          )
+        else:
+          logger.error(
+              "Error during close of plugin %s: %s",
+              plugin.name,
+              e,
+              exc_info=e,
+          )
+
+    if exceptions:
+      error_summary = ", ".join(
+          f"'{name}': {type(exc).__name__}" for name, exc in exceptions.items()
+      )
+      raise RuntimeError(f"Failed to close plugins: {error_summary}")

src/google/adk/runners.py

@@ -115,6 +115,7 @@ def __init__(
       session_service: BaseSessionService,
       memory_service: Optional[BaseMemoryService] = None,
       credential_service: Optional[BaseCredentialService] = None,
+      plugin_close_timeout: float = 5.0,
   ):
     """Initializes the Runner.
 
@@ -134,6 +135,7 @@ def __init__(
         session_service: The session service for the runner.
         memory_service: The memory service for the runner.
         credential_service: The credential service for the runner.
+        plugin_close_timeout: The timeout in seconds for plugin close methods.
 
     Raises:
         ValueError: If `app` is provided along with `app_name` or `plugins`, or
@@ -151,7 +153,9 @@ def __init__(
     self.session_service = session_service
     self.memory_service = memory_service
     self.credential_service = credential_service
-    self.plugin_manager = PluginManager(plugins=plugins)
+    self.plugin_manager = PluginManager(
+        plugins=plugins, close_timeout=plugin_close_timeout
+    )
     (
         self._agent_origin_app_name,
         self._agent_origin_dir,
@@ -1297,8 +1301,16 @@ async def _cleanup_toolsets(self, toolsets_to_close: set[BaseToolset]):
 
   async def close(self):
     """Closes the runner."""
+    logger.info('Closing runner...')
+    # Close Toolsets
     await self._cleanup_toolsets(self._collect_toolset(self.agent))
 
+    # Close Plugins
+    if self.plugin_manager:
+      await self.plugin_manager.close()
+
+    logger.info('Runner closed.')
+
   async def __aenter__(self):
     """Async context manager entry."""
     return self
@@ -1329,13 +1341,17 @@ def __init__(
       app_name: Optional[str] = None,
       plugins: Optional[list[BasePlugin]] = None,
       app: Optional[App] = None,
+      plugin_close_timeout: float = 5.0,
   ):
     """Initializes the InMemoryRunner.
 
     Args:
         agent: The root agent to run.
         app_name: The application name of the runner. Defaults to
           'InMemoryRunner'.
+        plugins: Optional list of plugins for the runner.
+        app: Optional App instance.
+        plugin_close_timeout: The timeout in seconds for plugin close methods.
     """
     if app is None and app_name is None:
       app_name = 'InMemoryRunner'
@@ -1347,4 +1363,5 @@ def __init__(
         app=app,
         session_service=InMemorySessionService(),
         memory_service=InMemoryMemoryService(),
+        plugin_close_timeout=plugin_close_timeout,
     )

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

@@ -513,10 +513,8 @@ async def fake_append_rows_with_error(requests, **kwargs):
     mock_write_client.append_rows.assert_called_once()
 
   @pytest.mark.asyncio
-  async def test_shutdown(
-      self, bq_plugin_inst, mock_bq_client, mock_write_client
-  ):
-    await bq_plugin_inst.shutdown()
+  async def test_close(self, bq_plugin_inst, mock_bq_client, mock_write_client):
+    await bq_plugin_inst.close()
     mock_write_client.transport.close.assert_called_once()
     mock_bq_client.close.assert_called_once()
 

tests/unittests/plugins/test_plugin_manager.py

@@ -16,6 +16,8 @@
 
 from __future__ import annotations
 
+import asyncio
+from unittest.mock import AsyncMock
 from unittest.mock import Mock
 
 from google.adk.models.llm_response import LlmResponse
@@ -267,3 +269,51 @@ async def test_all_callbacks_are_supported(
       "on_model_error_callback",
   ]
   assert set(plugin1.call_log) == set(expected_callbacks)
+
+
+@pytest.mark.asyncio
+async def test_close_calls_plugin_close(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that close calls the close method on registered plugins."""
+  plugin1.close = AsyncMock()
+  service.register_plugin(plugin1)
+
+  await service.close()
+
+  plugin1.close.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_close_raises_runtime_error_on_plugin_exception(
+    service: PluginManager, plugin1: TestPlugin
+):
+  """Tests that close raises a RuntimeError if a plugin's close fails."""
+  plugin1.close = AsyncMock(side_effect=ValueError("Shutdown error"))
+  service.register_plugin(plugin1)
+
+  with pytest.raises(
+      RuntimeError, match="Failed to close plugins: 'plugin1': ValueError"
+  ):
+    await service.close()
+
+  plugin1.close.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_close_with_timeout(plugin1: TestPlugin):
+  """Tests that close respects the timeout and raises on failure."""
+  service = PluginManager(close_timeout=0.1)
+
+  async def slow_close():
+    await asyncio.sleep(0.2)
+
+  plugin1.close = slow_close
+  service.register_plugin(plugin1)
+
+  with pytest.raises(RuntimeError) as excinfo:
+    await service.close()
+
+  assert "Failed to close plugins: 'plugin1': TimeoutError" in str(
+      excinfo.value
+  )

tests/unittests/test_runners.py

@@ -15,6 +15,7 @@
 from pathlib import Path
 import textwrap
 from typing import Optional
+from unittest.mock import AsyncMock
 
 from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.context_cache_config import ContextCacheConfig
@@ -562,6 +563,29 @@ async def test_runner_modifies_event_after_execution(self):
 
     assert modified_event_message == MockPlugin.ON_EVENT_CALLBACK_MSG
 
+  @pytest.mark.asyncio
+  async def test_runner_close_calls_plugin_close(self):
+    """Test that runner.close() calls plugin manager close."""
+    # Mock the plugin manager's close method
+    self.runner.plugin_manager.close = AsyncMock()
+
+    await self.runner.close()
+
+    self.runner.plugin_manager.close.assert_awaited_once()
+
+  @pytest.mark.asyncio
+  async def test_runner_passes_plugin_close_timeout(self):
+    """Test that runner passes plugin_close_timeout to PluginManager."""
+    runner = Runner(
+        app_name="test_app",
+        agent=MockLlmAgent("test_agent"),
+        session_service=self.session_service,
+        artifact_service=self.artifact_service,
+        plugins=[self.plugin],
+        plugin_close_timeout=10.0,
+    )
+    assert runner.plugin_manager._close_timeout == 10.0
+
   def test_runner_init_raises_error_with_app_and_app_name_and_agent(self):
     """Test that ValueError is raised when app, app_name and agent are provided."""
     with pytest.raises(