Skip to content

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

Open
CTWalk wants to merge 2 commits into
google:mainfrom
CTWalk:fix/sync-run-close-cancels-invocation
Open

fix: cancel the invocation when a sync run() generator is closed early#6564
CTWalk wants to merge 2 commits into
google:mainfrom
CTWalk:fix/sync-run-close-cancels-invocation

Conversation

@CTWalk

@CTWalk CTWalk commented Aug 3, 2026

Copy link
Copy Markdown

Component: coreRunner.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 resulting CancelledError is treated as expected thread teardown —
    without this, threading's default excepthook prints a CancelledError
    traceback to stderr on every early close; and
  • 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 the cases that note implies, on the same pin:

close() while the agent is mid-flight      -> returned in 0.4 ms, agent cancelled
stream dropped, then gc.collect()          -> returned in 30.7 ms, agent cancelled
full consumption (unchanged path)          -> same events, 1.7 ms

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.

As a sanity check that the test is not vacuous: keeping the test and reverting
only runners.py makes it fail on the cancellation assertion, so it fails on
today's main and passes with this change.

Ran locally against f4e7233469e3595336dfb0d84c281b2f6245ce4c:

pytest tests/unittests/test_runners.py -q            -> 78 passed
pytest tests/unittests/test_runners.py -q -n auto    -> 78 passed
new test, 10 consecutive runs                        -> no flakes
pre-commit run --files <the two changed files>       -> all hooks passed
mypy vs unmodified main, using the baseline/PR diff  -> no new errors
  the CI type-check job itself performs

Because the change is threading and cancellation, I also ran the reproducer
above on both ends of the CI matrix — CPython 3.10 and 3.14 — with the same
post-fix output. Development host was macOS 15.7.4 / CPython 3.11.14.

The full tests/unittests run also passes: 9260 passed against 9259 on
unmodified main — the delta is exactly the added test, with an identical set
of 15 pre-existing failures caused by optional dependencies missing from my
local environment. I did not run the tox matrix, nor
tests/unittests/evaluation and tests/unittests/optimization, which do not
collect locally for the same reason; I am not claiming CI coverage for those.

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.

Environment

  • ADK version: google-adk 2.6.1, editable checkout at
    f4e7233469e3595336dfb0d84c281b2f6245ce4c
  • 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.

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.
@google-cla

google-cla Bot commented Aug 3, 2026

Copy link
Copy Markdown

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.

@adk-bot adk-bot added the core [Component] This issue is related to the core interface and implementation label Aug 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core [Component] This issue is related to the core interface and implementation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants