From 7fabbf10f2c0813686cb842290baefda790a65e4 Mon Sep 17 00:00:00 2001 From: CTWalk <100585900+CTWalk@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:35:51 +0800 Subject: [PATCH] fix: cancel the invocation when a sync run() generator is closed early Closing the generator returned by the synchronous Runner.run() did not stop the invocation it started. The agent kept running on the background event loop and could append further events to the session after Generator.close() returned. Runner._cleanup_root_task() documents that the root task must be cancelled when the caller stops iterating early, and test_run_async_teardown_on_aclose pins that behavior for the async entrance. The sync wrapper starts run_async() in a background thread but never propagated the foreground generator's close to that task. Hand the background event loop and task to the foreground generator, tell normal queue exhaustion apart from an early exit, and on early exit only cancel the background task -- which unwinds through the existing aclosing(...) and so reuses run_async()'s own _cleanup_root_task() teardown. Treat the resulting CancelledError as expected thread teardown, and join the thread on both paths so close() does not return while the invocation is still alive. Adds test_run_teardown_on_close, the sync counterpart of the existing test_run_async_teardown_on_aclose. --- src/google/adk/runners.py | 45 ++++++++++++++---- tests/unittests/test_runners.py | 82 +++++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+), 9 deletions(-) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index d8c4e6a04e..77385b63c4 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -1057,8 +1057,18 @@ def run( """ run_config = run_config or RunConfig() event_queue: queue.Queue[Event | None] = queue.Queue() + # Handle to the background invocation, so that closing this generator early + # can cancel it instead of leaking a running task. See + # `_cleanup_root_task()` for the equivalent guarantee on `run_async()`. + invocation_handle: queue.Queue[ + tuple[asyncio.AbstractEventLoop, asyncio.Task[Any] | None] + ] = queue.Queue(maxsize=1) + caller_closed_early = False async def _invoke_run_async() -> None: + invocation_handle.put( + (asyncio.get_running_loop(), asyncio.current_task()) + ) try: async with aclosing( self.run_async( @@ -1077,21 +1087,38 @@ async def _invoke_run_async() -> None: def _asyncio_thread_main() -> None: try: asyncio.run(_invoke_run_async()) + except asyncio.CancelledError: + if not caller_closed_early: + raise finally: event_queue.put(None) thread = create_thread(target=_asyncio_thread_main) thread.start() - # consumes and re-yield the events from background thread. - while True: - event = event_queue.get() - if event is None: - break - else: - yield event - - thread.join() + exhausted = False + try: + # consumes and re-yield the events from background thread. + while True: + event = event_queue.get() + if event is None: + exhausted = True + break + else: + yield event + finally: + if not exhausted: + # The caller stopped iterating early, so cancel the invocation before + # it can run further tools or append more events to the session. + caller_closed_early = True + loop, task = invocation_handle.get() + if task is not None: + try: + loop.call_soon_threadsafe(task.cancel) + except RuntimeError: + # The background loop already finished; nothing to cancel. + pass + thread.join() async def run_async( self, diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index b76953eab6..185e0601d4 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -2084,6 +2084,88 @@ async def _run_async_impl( assert was_cancelled["value"] is True +def test_run_teardown_on_close(): + """Closing the sync run() generator should cancel the running agent task.""" + import asyncio + import threading + import time + + session_service = InMemorySessionService() + + release_second = threading.Event() + was_cancelled = {"value": False} + completed = {"value": False} + + class CancellingAgent(BaseAgent): + + async def _run_async_impl( + self, invocation_context: InvocationContext + ) -> AsyncGenerator[Event, None]: + try: + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text="First response")] + ), + ) + # Bounded wait so a broken teardown fails instead of hanging. + deadline = time.monotonic() + 5.0 + while not release_second.is_set() and time.monotonic() < deadline: + await asyncio.sleep(0.01) + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text="Second response")] + ), + ) + completed["value"] = True + except (asyncio.CancelledError, GeneratorExit): + was_cancelled["value"] = True + raise + + runner = Runner( + app_name=TEST_APP_ID, + agent=CancellingAgent(name="cancel_agent"), + session_service=session_service, + artifact_service=InMemoryArtifactService(), + auto_create_session=True, + ) + + # Given a sync run stream + stream = runner.run( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=types.Content(role="user", parts=[types.Part(text="hello")]), + ) + + # When the client reads the first event and then calls close() + event = next(stream) + assert event.content.parts[0].text == "First response" + + stream.close() + release_second.set() + + # Then the running agent was cancelled before it could do further work + assert was_cancelled["value"] is True + assert completed["value"] is False + + # And no later event was appended to the session. + session = asyncio.run( + session_service.get_session( + app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + ) + texts = [ + part.text + for session_event in session.events + if session_event.content + for part in session_event.content.parts + ] + assert texts == ["hello", "First response"] + + @pytest.mark.asyncio async def test_run_live_passes_get_session_config(): """run_live should forward RunConfig.get_session_config to get_session."""