fix: cancel the invocation when a sync run() generator is closed early

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,

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