fix: cancel the invocation when a sync run() generator is closed early - #6564
Open
CTWalk wants to merge 2 commits into
Open
fix: cancel the invocation when a sync run() generator is closed early#6564CTWalk wants to merge 2 commits into
CTWalk wants to merge 2 commits into
Conversation
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.
|
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Component:
core—Runner.run()insrc/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 twinalready cancels correctly on
aclose(); the sync wrapper never propagatesclose to the background task. This cancels it on early exit only, reusing
run_async()'s existing teardown, and adds the sync counterpart of the testthat pins the async behavior.
Describe the bug
Closing the generator returned by the synchronous
Runner.run()does not stopthe invocation it started. The agent keeps running on the background event loop
and can append further events to the session after
Generator.close()hasreturned.
Runner._cleanup_root_task()documents the intended behavior: when the callerstops iterating early, the root task must be cancelled to avoid a leaked task.
test_run_async_teardown_on_aclosepins that for the async entrance. The syncwrapper starts
run_async()in a background thread but never propagates theforeground 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.mdsuggests. Happy to open aseparate issue first if you would prefer that.
Steps to reproduce
Keyless — no model or network call.
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 alreadyreturned.
Observed behavior (before this change)
Expected behavior (after this change)
stream.close()cancels the underlying invocation before any further agent workor session append, matching
run_async().aclose().Root cause
Runner.run()runs_invoke_run_async()in a background event-loop thread andconsumes an event queue in the foreground generator.
Generator.close()raisesGeneratorExitat the foregroundyield; the frame exits without cancellingthe 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 insrc/:through a one-item queue;
early exit;
loop.call_soon_threadsafe(task.cancel)— which unwinds through the existingaclosing(...)and so reusesrun_async()'s own_cleanup_root_task()teardown;
CancelledErroris treated as expected thread teardown —without this,
threading's default excepthook prints aCancelledErrortraceback to stderr on every early close; and
close()does not return while theinvocation 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 thecancelled root task). An agent that swallows
CancelledErrorand keepsrunning will therefore delay
close(), exactly as it already delaysaclose().I measured the cases that note implies, on the same pin:
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_closeintests/unittests/test_runners.py, thesync counterpart of the existing
test_run_async_teardown_on_aclose. Itconsumes 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.pymakes it fail on the cancellation assertion, so it fails ontoday's
mainand passes with this change.Ran locally against
f4e7233469e3595336dfb0d84c281b2f6245ce4c: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/unittestsrun also passes: 9260 passed against 9259 onunmodified
main— the delta is exactly the added test, with an identical setof 15 pre-existing failures caused by optional dependencies missing from my
local environment. I did not run the
toxmatrix, nortests/unittests/evaluationandtests/unittests/optimization, which do notcollect 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:Console output before the change:
Console output after the change:
The relevant lines are
cancelledflipping toTrueand'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, andnone of them are the component this PR belongs to.
disappearing across the sync thread boundary, not early-close teardown.
run_async()from outside the agent #4796 requests an external "stop generating" handle forrun_async(). Thischange does not add a new API; it makes the existing generator-close path
behave as documented.
context-detach error on early
run_async()close, intelemetry/_instrumentation.py. No file, symbol, or behavior is shared withthis PR.
Environment
google-adk2.6.1, editable checkout atf4e7233469e3595336dfb0d84c281b2f6245ce4cdeterministic agent and make no model or network call.