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

[CTWalk]

Component: core — Runner.run() in src/google/adk/runners.py.

Summary: closing the generator returned by the synchronous Runner.run()
does not stop the invocation behind it, so the agent keeps running and can
append events to the session after close() has returned. The async twin
already cancels correctly on aclose(); the sync wrapper never propagates
close to the background task. This cancels it on early exit only, reusing
run_async()'s existing teardown, and adds the sync counterpart of the test
that pins the async behavior.

Describe the bug

Closing the generator returned by the synchronous Runner.run() does not stop
the invocation it started. The agent keeps running on the background event loop
and can append further events to the session after Generator.close() has
returned.

Runner._cleanup_root_task() documents the intended behavior: when the caller
stops iterating early, the root task must be cancelled to avoid a leaked task.
test_run_async_teardown_on_aclose pins that for the async entrance. The sync
wrapper starts run_async() in a background thread but never propagates the
foreground generator's close to that task.

I did not find an existing issue for this, so I have followed the bug-template
structure in this description as CONTRIBUTING.md suggests. Happy to open a
separate issue first if you would prefer that.

Steps to reproduce

Keyless — no model or network call.

import asyncio
import threading
from typing import AsyncGenerator

from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.events.event import Event
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types

release_second = threading.Event()
completed = threading.Event()
cancelled = threading.Event()


class TwoEventAgent(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")]),
      )
      while not release_second.is_set():
        await asyncio.sleep(0.001)
      yield Event(
          invocation_id=invocation_context.invocation_id,
          author=self.name,
          content=types.Content(role="model", parts=[types.Part(text="second")]),
      )
      completed.set()
    except (asyncio.CancelledError, GeneratorExit):
      cancelled.set()
      raise


session_service = InMemorySessionService()
runner = Runner(
    app_name="close_repro",
    agent=TwoEventAgent(name="two_events"),
    session_service=session_service,
    auto_create_session=True,
)

stream = runner.run(
    user_id="user",
    session_id="session",
    new_message=types.Content(role="user", parts=[types.Part(text="go")]),
)
assert next(stream).content.parts[0].text == "first"
stream.close()
release_second.set()
completed.wait(timeout=5)

session = asyncio.run(
    session_service.get_session(
        app_name="close_repro", user_id="user", session_id="session"
    )
)
persisted = [
    event.content.parts[0].text
    for event in session.events
    if event.content and event.content.parts
]
print(f"completed={completed.is_set()}")
print(f"cancelled={cancelled.is_set()}")
print(f"persisted={persisted}")

The agent blocks on a thread event before its second yield, so the ordering is
deterministic: the second event can only be produced after close() has already
returned.

Observed behavior (before this change)

completed=True
cancelled=False
persisted=['go', 'first', 'second']

Expected behavior (after this change)

stream.close() cancels the underlying invocation before any further agent work
or session append, matching run_async().aclose().

completed=False
cancelled=True
persisted=['go', 'first']

Root cause

Runner.run() runs _invoke_run_async() in a background event-loop thread and
consumes an event queue in the foreground generator. Generator.close() raises
GeneratorExit at the foreground yield; the frame exits without cancelling
the background task and without joining the thread. The background invocation
stays free to run tools, emit events, and mutate session history.

What this change does

Runner.run() is the only function changed in src/:

• the background coroutine hands its event loop and task to the foreground
  through a one-item queue;
• the foreground consumer loop distinguishes normal queue exhaustion from an
  early exit;
• on early exit only, it cancels the background task with
  loop.call_soon_threadsafe(task.cancel) — which unwinds through the existing
  aclosing(...) and so reuses run_async()'s own _cleanup_root_task()
  teardown;
• the thread is joined on both paths, so close() does not return while the
  invocation is still alive.

No public signature, dependency, documentation, or unrelated error behavior
changes. Normal full-consumption runs take the same path as before.

Behavioral note for reviewers: close() now blocks until teardown completes,
which is the same contract run_async().aclose() already has (it awaits the
cancelled root task). An agent that swallows CancelledError and keeps
running will therefore delay close(), exactly as it already delays aclose().
I measured both entrances against an agent that holds the signal for 1s: sync
close() blocked 1002.9 ms, async aclose() blocked 1002.4 ms.

One asymmetry worth naming, since the two paths stop the agent by different
means: teardown reaches the agent as CancelledError on run().close() and as
GeneratorExit on run_async().aclose(). The observable contract is the same —
the invocation stops and appends nothing further — and the existing
test_run_async_teardown_on_aclose already treats the two as one case by
catching (asyncio.CancelledError, GeneratorExit) together, which the new sync
tests mirror. Happy to unify the signal type instead if you would prefer that.

I measured the cases that note implies, on the same pin:

close() while the agent is mid-flight      -> returned in 0.3 ms, agent cancelled
stream dropped, then gc.collect()          -> returned in 40.4 ms, agent cancelled
full consumption (unchanged path)          -> same events, agent runs to completion

The abandoned-stream case is worth calling out: before this change, dropping
the last reference to a partially consumed stream left the invocation running
to completion in a detached thread. After it, the generator's own finalizer
cancels the invocation and returns promptly, so the abandonment path stops
leaking as well.

Testing plan

Unit tests

Added test_run_teardown_on_close in tests/unittests/test_runners.py, the
sync counterpart of the existing test_run_async_teardown_on_aclose. It
consumes the first event, closes the stream, and asserts that the agent was
cancelled, did not complete, and appended no later event to the session. The
agent's wait is bounded, so a broken teardown fails the test rather than
hanging it.

A second test, test_run_close_cancels_agent_parked_without_timers, pins the
part the first one cannot see. Its agent parks on a future and schedules no
timer, so nothing wakes the background event loop on its own. Cancelling the
task without call_soon_threadsafe sets the flag but never wakes the selector,
and close() blocks forever; a polling agent hides this, because its own
sleeps keep waking the loop. The test calls close() on a helper thread and
asserts it returns, so that failure mode fails the test instead of hanging it,
and it releases the agent in a finally so the non-daemon runner thread always
exits.

As a sanity check that the tests are not vacuous: keeping them and reverting
only runners.py makes them fail on the cancellation assertion, so they fail
on today's main and pass with this change. I also checked them against the
plausible wrong fixes — dropping the thread.join(), cancelling without
call_soon_threadsafe, and not cancelling at all — and each is caught.

Ran locally against 989721746aba65e90f644e51606375699175f709:

pytest tests/unittests/test_runners.py -q            -> 91 passed
  (89 on unmodified main; the delta is exactly the two added tests)
new tests, 10 consecutive runs                       -> no flakes
pre-commit run --files <the two changed files>       -> all hooks passed

The runner-adjacent sweep (-k "runner or run_" across tests/unittests)
passes as well: 866 passed, with one failure —
test_import_loading.py::test_entry_point_loads_only_allowlisted_packages[runner]
— that reproduces identically on unmodified main, so it is pre-existing and
not caused by this change. Development host was macOS 15.7.4 / CPython 3.11.14.

Scope of what I did not run, so this is not read as a CI claim: the tox
matrix, other Python versions, and tests/unittests/evaluation and
tests/unittests/optimization, which do not collect in my local environment
for want of optional dependencies.

Manual E2E (Runner)

Runner setup and agent definition: the reproducer in "Steps to reproduce"
above — an in-memory session service, a deterministic two-event BaseAgent,
and Runner.run(). Command:

python repro_run_close.py

Console output before the change:

completed=True
cancelled=False
persisted=['go', 'first', 'second']

Console output after the change:

completed=False
cancelled=True
persisted=['go', 'first']

The relevant lines are cancelled flipping to True and 'second'
disappearing from the persisted session events: the invocation stops at close
instead of running on and appending.

Related, not duplicates

This change is confined to core; none of the items below overlap with it, and
none of them are the component this PR belongs to.

• Runner.run() silently terminates the event generator when run_async raises — exception reaches stderr but no signal to caller #5567 (closed by its author as filed in error) concerns exceptions
  disappearing across the sync thread boundary, not early-close teardown.
• "Stop generating" — ability to stop run_async() from outside the agent #4796 requests an external "stop generating" handle for run_async(). This
  change does not add a new API; it makes the existing generator-close path
  behave as documented.
• b55000d9 ("fix: re-raise agent errors from the synchronous Runner.run()")
  landed upstream after this PR was filed and rewrote the same function. This
  branch is rebased onto it, and the two compose rather than conflict: that
  change forwards an agent failure to the calling thread, this one cancels the
  invocation when the caller stops iterating. The rebase made this patch
  smaller — its except BaseException handler already absorbs the
  CancelledError raised by teardown, so the explicit suppression this PR
  originally carried is gone. test_run_reraises_agent_error,
  test_run_yields_events_before_reraising_agent_error, and
  test_run_reports_agent_cancellation_as_runtime_error all still pass.
• fix: avoid OTel "Failed to detach context" error on early run_async close #6560 (open) proposes a fix in a different component — an OpenTelemetry
  context-detach error on early run_async() close, in
  telemetry/_instrumentation.py. No file, symbol, or behavior is shared with
  this PR.

Environment

• ADK version: google-adk 2.6.1, editable checkout at
  989721746aba65e90f644e51606375699175f709
• Python: 3.11.14
• OS: macOS 15.7.4
• LiteLLM: no. Model: none — the reproducer and the new test use a
  deterministic agent and make no model or network call.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[CTWalk]

Rebased onto main (9897217) — the branch was conflicting and is now clean.

The conflict was with b55000d ("fix: re-raise agent errors from the synchronous Runner.run()"), which rewrote the same function after this PR was filed. The two changes compose rather than overlap: that one forwards an agent failure to the calling thread, this one cancels the invocation when the caller stops iterating early.

The rebase made this patch smaller. That change's except BaseException handler already absorbs the CancelledError raised by teardown, so the explicit suppression this PR originally carried is gone. What remains is the invocation handle plus the try/finally that cancels on early exit.

Re-verified on the new base:

• the new test_run_teardown_on_close still fails on unmodified main (on the cancellation assertion) and passes with the change, so it is not vacuous;
• test_run_reraises_agent_error, test_run_yields_events_before_reraising_agent_error, and test_run_reports_agent_cancellation_as_runtime_error all still pass;
• tests/unittests/test_runners.py → 90 passed, against 89 on unmodified main; the delta is exactly the added test;
• 10 consecutive runs of the new test, no flakes; pre-commit on both changed files, all hooks pass.

The PR description has been updated to match the rebased patch, including replacing two verification claims that no longer applied to this base.

[CTWalk]

Added a second test, and updated the description to match.

test_run_close_cancels_agent_parked_without_timers covers what the first test structurally cannot see. Its agent parks on a future and schedules no timer, so nothing wakes the background event loop on its own. Cancelling the task without call_soon_threadsafe sets the flag but never wakes the selector, and close() blocks forever — a polling agent hides this, because its own asyncio.sleep() calls keep waking the loop regardless.

I checked the tests against the plausible wrong fixes rather than only against the bug:

variant	caught
drop the thread.join()	yes
task.cancel() without call_soon_threadsafe	yes — only by the new test
never cancel (no-op fix)	yes
cancel unconditionally, dropping if not exhausted	no — and deliberately so: once the run is exhausted the loop is closed, call_soon_threadsafe raises RuntimeError and is caught, so that variant is equivalent. The guard is clarity, not correctness.

The new test calls close() on a daemon helper thread with a bounded wait, so a teardown that never returns fails the test rather than hanging it, and it releases the agent in a finally so the non-daemon runner thread always exits.

Also now stated explicitly in the description: teardown reaches the agent as CancelledError on run().close() and as GeneratorExit on run_async().aclose(). The observable contract is identical, and the pre-existing test_run_async_teardown_on_aclose already catches the two together — but I would rather name it than have it found in review. Happy to unify the signal type if you prefer.

tests/unittests/test_runners.py is now 91 passed, against 89 on unmodified main; the delta is exactly the two added tests. src/ is unchanged from the previous push.