diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 299c6d9a65c..3b2020c74ee 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -915,8 +915,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, @@ -946,17 +955,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 3465d82d4c7..0ad4e8ca169 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -484,6 +484,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 @@ -2110,6 +2199,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."""