From 100026d8029afc31459b0db5982544c813482d0c Mon Sep 17 00:00:00 2001 From: CTWalk 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 | 47 +++++++-- tests/unittests/test_runners.py | 171 ++++++++++++++++++++++++++++++++ 2 files changed, 207 insertions(+), 11 deletions(-) diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 2a1ed913222..632a09129ca 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -1136,8 +1136,17 @@ def run( """ run_config = run_config or RunConfig() event_queue: queue.Queue[Event | BaseException | 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) async def _invoke_run_async() -> None: + invocation_handle.put( + (asyncio.get_running_loop(), asyncio.current_task()) + ) async with aclosing( self.run_async( user_id=user_id, @@ -1167,17 +1176,33 @@ def _asyncio_thread_main() -> None: # consumes and re-yield the events from background thread. agent_error: BaseException | None = None - while True: - item = event_queue.get() - if item is None: - break - elif isinstance(item, BaseException): - agent_error = item - break - else: - yield item - - thread.join() + exhausted = False + try: + while True: + item = event_queue.get() + if item is None: + exhausted = True + break + elif isinstance(item, BaseException): + agent_error = item + exhausted = True + break + else: + yield item + 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. The + # cancellation unwinds through `aclosing(...)` above, reusing + # `run_async()`'s own `_cleanup_root_task()` teardown. + 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() if isinstance(agent_error, Exception): raise agent_error if agent_error is not None: diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index ebfc0018254..9947bed4338 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -738,6 +738,95 @@ async def test_run_live_auto_create_session(): assert session is not None +def test_run_close_cancels_agent_parked_without_timers(): + """close() must wake the background loop, not just flag the task. + + The agent parks on a future and schedules no timer, so nothing wakes the + event loop on its own. Cancelling the task from this thread without + `call_soon_threadsafe` sets the flag but never wakes the selector, and + close() blocks forever. A polling agent hides that, because its own sleeps + keep waking the loop. + """ + import asyncio + import threading + + session_service = InMemorySessionService() + + parked = threading.Event() + handle = {} + was_cancelled = {"value": False} + + class ParkingAgent(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")] + ), + ) + loop = asyncio.get_running_loop() + handle["loop"] = loop + handle["future"] = loop.create_future() + parked.set() + await handle["future"] + except (asyncio.CancelledError, GeneratorExit): + was_cancelled["value"] = True + raise + + runner = Runner( + app_name=TEST_APP_ID, + agent=ParkingAgent(name="parking_agent"), + session_service=session_service, + artifact_service=InMemoryArtifactService(), + auto_create_session=True, + ) + + # Given a sync run stream whose agent is parked with no pending timer + stream = runner.run( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=types.Content(role="user", parts=[types.Part(text="hello")]), + ) + event = next(stream) + assert event.content.parts[0].text == "First response" + assert parked.wait(timeout=5.0) is True + + # When the client closes the stream, from a helper thread so that a + # teardown that never returns fails this test instead of hanging it + closed = threading.Event() + + def _close_stream() -> None: + stream.close() + closed.set() + + closer = threading.Thread(target=_close_stream, daemon=True) + closer.start() + + try: + # Then close() returns instead of blocking on an unwoken event loop + assert closed.wait(timeout=10.0) is True + assert was_cancelled["value"] is True + finally: + # Release the agent either way, so the non-daemon runner thread always + # exits and cannot wedge the rest of the suite. + loop = handle.get("loop") + future = handle.get("future") + if loop is not None and future is not None: + try: + loop.call_soon_threadsafe( + lambda: future.cancel() if not future.done() else None + ) + except RuntimeError: + # Teardown already finished and closed the loop; nothing to release. + pass + closer.join(timeout=5.0) + + def test_run_passes_state_delta(): """run should forward state_delta down to run_async.""" import asyncio @@ -2364,6 +2453,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."""