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."""