From fd0d50f0b0f116ccd18d1e3a8a4517640d508c6c Mon Sep 17 00:00:00 2001 From: Sehlani042 <257166922+Sehlani042@users.noreply.github.com> Date: Fri, 3 Jul 2026 10:54:18 +0800 Subject: [PATCH 1/2] fix(runner): shield toolset cleanup from cancellation --- src/google/adk/runners.py | 34 +++++++++++++++++++++--- tests/unittests/test_runners.py | 47 +++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 9802801dcad..e90620b3355 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -2095,10 +2095,12 @@ async def _cleanup_toolsets(self, toolsets_to_close: set[BaseToolset]): # This maintains the same task context throughout cleanup for toolset in toolsets_to_close: + cleanup_task = asyncio.create_task( + asyncio.wait_for(toolset.close(), timeout=10.0) + ) try: logger.info('Closing toolset: %s', type(toolset).__name__) - # Use asyncio.wait_for to add timeout protection - await asyncio.wait_for(toolset.close(), timeout=10.0) + await asyncio.shield(cleanup_task) logger.info('Successfully closed toolset: %s', type(toolset).__name__) except asyncio.TimeoutError: logger.warning('Toolset %s cleanup timed out', type(toolset).__name__) @@ -2113,8 +2115,34 @@ async def _cleanup_toolsets(self, toolsets_to_close: set[BaseToolset]): # improved context propagation across task boundaries, and better cancellation # handling prevent the cross-task cancel scope violation. logger.warning( - 'Toolset %s cleanup cancelled: %s', type(toolset).__name__, e + 'Toolset %s cleanup cancellation requested: %s', + type(toolset).__name__, + e, ) + try: + await cleanup_task + logger.info( + 'Successfully closed toolset after cancellation request: %s', + type(toolset).__name__, + ) + except asyncio.TimeoutError: + cleanup_task.cancel() + logger.warning( + 'Toolset %s cleanup timed out after cancellation request', + type(toolset).__name__, + ) + except asyncio.CancelledError as close_cancelled: + logger.warning( + 'Toolset %s cleanup cancelled: %s', + type(toolset).__name__, + close_cancelled, + ) + except Exception as close_error: + logger.error( + 'Error closing toolset %s after cancellation request: %s', + type(toolset).__name__, + close_error, + ) except Exception as e: logger.error('Error closing toolset %s: %s', type(toolset).__name__, e) diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index 925769951d8..1fdd604b1e4 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -36,6 +36,7 @@ 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.base_toolset import BaseToolset from google.genai import types import pytest @@ -1120,6 +1121,52 @@ async def test_runner_close_calls_plugin_close(self): self.runner.plugin_manager.close.assert_awaited_once() + @pytest.mark.asyncio + async def test_runner_close_does_not_cancel_toolset_cleanup(self): + """Caller cancellation should not cancel an in-flight toolset close.""" + import asyncio + + class SlowCloseToolset(BaseToolset): + + def __init__(self): + super().__init__() + self.close_started = asyncio.Event() + self.close_finished = asyncio.Event() + self.close_cancelled = False + + async def get_tools(self, readonly_context=None): + del readonly_context + return [] + + async def close(self) -> None: + self.close_started.set() + try: + await asyncio.sleep(0.05) + self.close_finished.set() + except asyncio.CancelledError: + self.close_cancelled = True + raise + + toolset = SlowCloseToolset() + runner = Runner( + app_name="test_app", + agent=LlmAgent( + name="test_agent", model="gemini-1.5-pro", tools=[toolset] + ), + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + + close_task = asyncio.create_task(runner.close()) + await toolset.close_started.wait() + close_task.cancel() + + await close_task + + assert close_task.cancelled() is False + assert toolset.close_cancelled is False + assert toolset.close_finished.is_set() + @pytest.mark.asyncio async def test_runner_passes_plugin_close_timeout(self): """Test that runner passes plugin_close_timeout to PluginManager.""" From f8fd10af81fe91266285ba56734d99f22b4ecb8f Mon Sep 17 00:00:00 2001 From: Sehlani042 <257166922+Sehlani042@users.noreply.github.com> Date: Fri, 3 Jul 2026 14:06:03 +0800 Subject: [PATCH 2/2] chore: retrigger checks after CLA