fix(flows): propagate downstream cancellation to the model stream

[sangkyoonnam]

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

• Related: [core] Client disconnects don't cancel the model stream (per-step flow is cached) — and there is no public API to cancel an in-flight run #1618
• Related: Feature Request: Run Cancellation / Abort Support for Java ADK #1341

This covers the propagation half of #1618; a public cancel API (#1341) stays separate. Rebased onto #1565, which moved the changed line into runStepAndCheckContinue; the swap is the same there.

Problem:
Disposing the Flowable from Runner.runAsync doesn't stop the model call. BaseLlmFlow.run shares each step as runOneStep(...).cache(), and FlowableCache never cancels its upstream: CacheSubscription.cancel() only drops the subscriber from the list. Once the consumer disposes, the step runs on with zero subscribers, including the after-model callbacks and the step's tool calls. I measured the last part: a function call the model emits after the dispose still runs the tool.

Solution:
Keep cache() and put a takeUntil in front of it: step.takeUntil(cancelled.toFlowable()).cache().doOnCancel(cancelled::onComplete), one CompletableSubject per step. A cancel completes the signal, takeUntil cancels runOneStep and completes the cached stream. What run() relies on from cache() is unchanged: unbounded upstream demand, terminal delivered as soon as a subscriber drains the buffer, buffer kept for the toList() subscriber that arrives after the step finished.

After the change a disconnect cancels the step's pending reactive work: the in-flight model Flowable, after-model callbacks not yet run, tool calls not yet started. A tool already executing has its Single disposed; whether its work stops is up to the tool. A model error arriving after the cancel can go to RxJavaPlugins.onError instead of sitting in the cache for nobody, which is RxJava's normal post-cancel routing. Whether the HTTP request stops depends on the model client. For Gemini streaming, this PR schedules the response stream's close on cancel, including a stream that opens afterwards; the close can wait behind a blocked read.

ParallelAgent escalation now cancels an active sibling's model stream, with a regression test. Outside durable resumable mode, LoopAgent escalation also cancels its active sub-agent's remaining step. #1434's cooperative CancellationToken cuts the stream inside callLlm and doesn't touch cache(), so the two compose.

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

BaseLlmFlowCancelTest (new, 9 tests): cancel reaches the model through BaseLlmFlow.run and through InMemoryRunner.runAsync, a dispose before a function call leaves the tool unrun, a cancel inside onSubscribe doesn't leave the model streaming (this one already passes on the rebased base, since #1565 checks for the cancel before subscribing to the step), a ParallelAgent sibling escalation cancels the model; cache() semantics kept under bounded demand (run(ctx).test(1) gets the final event and completion, and a model error after a buffered event) and a new run that completes when run() is subscribed again after a cancel; a control run still runs the tool. Without the BaseLlmFlow change 5 of the 9 fail. ./mvnw -pl core test: 2342 tests, 0 failures, 0 errors, 0 skipped.

Manual End-to-End (E2E) Tests:

Not run. The dev server's /run_sse disposes on onCompletion/onTimeout and the A2A executor's cancel() disposes its CompositeDisposable; with this change both reach the model. I read those paths, I haven't exercised them.

Checklist

• I have read the CONTRIBUTING.md document.
• My pull request contains a single commit.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules.

Additional context

runLive doesn't use cache() and isn't touched. Nothing in core subscribes to one run() result twice at the same time, so "a subscriber cancels" and "the last subscriber cancels" are the same event here.

I wrote this with Claude Code and reviewed and ran it myself.

[hemasekhar-p]

Hi @sangkyoonnam, thank you for your contribution. We appreciate you taking the time to submit this pull request. As per contribution policy, PR should consists of a single commit. Could you please change your commits accordingly?

[sangkyoonnam]

Done. The branch is one commit again, a63b9cd, rebased onto main at 9debe96. The second commit was a merge from main I clicked through on GitHub right after opening the PR; the diff is unchanged.

[sangkyoonnam]

The reporter validated this head (a63b9cd) in production: 311 regression tests green, a mid-stream client disconnect left no further model or tool activity, and their cooperative cancel flag still works alongside. Details in #1618 (comment).

[hemasekhar-p]

@sangkyoonnam,thank you for the update and I noticed there is nothing in the merge commit. To proceed with the review, could you please squash them into a single commit?

[sangkyoonnam]

Squashed again: the branch is back to the single commit a63b9cd. The extra commit was an empty merge from main made with the Update branch button; I won't use it on this PR again.

[hemasekhar-p]

@sangkyoonnam, thanks for the quick update. Currently this PR is under review by our team, we will keep you posted if any additional information is required. thank you.

[MiloszSobczyk]

Hello, thank you for the PR.

I checked the change against FlowableCache, FlowableTakeUntil, and doOnCancel. doOnCancel marks itself cancelled before it runs the callback, so the cancelling subscriber never sees the forced onComplete, and run() does not start the next step. I also ran the new tests: 9/9 pass, and 6/9 fail with the one-line change reverted, as you describe. This also matches Python ADK, which closes the generator chain with Aclosing.

One thing outside this diff, please: Gemini builds the SSE stream as Flowable.fromFuture(streamFuture).flatMapIterable(iterable -> iterable) and never calls ResponseStream.close(). Before this change, the detached cache() read every stream to the end, so the HTTP connection was released. With this change, a cancel stops reading partway and nothing closes the stream, leaving the response open until it is garbage collected. That happens on every client disconnect in /run_sse, which is the case this PR fixes. Could you please close the stream on cancel—for example, using Flowable.using(..., ResponseStream::close)—either here or in a follow-up PR? Note that closing the reader can block while another thread is inside readLine().

Could you please also list the visible behavior changes in the commit message so users can find them in the release notes:

• In-flight tool calls are disposed on cancel. With ToolExecutionMode.PARALLEL_SUBSCRIBE the tool runs on Schedulers.io() by default, so its thread is interrupted.
• An error that arrives after a cancel goes to RxJavaPlugins.onError. Without a global error handler this becomes an UndeliverableException, which crashes Android apps.
• An escalation in ParallelAgent or LoopAgent now stops the sibling sub-agents' model calls and tools, which used to run to the end.

[sangkyoonnam]

Both are in the amended commit: the stream now closes on cancel, and the commit message lists the visible changes. I did the close here rather than in a follow-up, since this PR is what leaves the stream half-read on a disconnect.

Gemini reads the stream through Flowable.using, and the close runs on Schedulers.io(), so the cancelling thread never waits on the readLine() lock. It also closes a stream whose future completes after the cancel, and a continuation response the SDK opens while an iterator call is already in flight. That second case is new with the SDK continuation that landed on main this week.

The two cancel tests in GeminiTest and the continuation test in GeminiContinuationTest fail with the old read path, and the continuation test also fails with my first version of this fix, which only closed the stream once. A read that completes or fails without cancellation is left alone; the extra cleanup only runs on cancel. I also narrowed three of your bullets after checking them against the code: interruption depends on the scheduler and the disposing thread, not every error after a cancel reaches RxJavaPlugins.onError, and LoopAgent cancels the active sub-agent's remaining work outside durable resumable mode.

Still one commit, rebased onto main. ./mvnw -pl core test passes: 2360 run, 0 failures, 0 errors, 24 skipped.

[sangkyoonnam]

Shortened both Javadocs and replaced the timed wait with assertThat(tool.ran.getCount()).isEqualTo(1) in dd914fc.

[sangkyoonnam]

Rebased the single commit onto main at 0ee0d4b as 542ad25. The conflict was with #1565: cancellableCache now wraps the step in runStepAndCheckContinue. Apart from that relocation and the test update below, the patch is unchanged.

Since #1565, run() is deferred, so subscribing again after a cancel starts a new run. run_subscribedAgainAfterCancel_startsNewRunAndCompletes now checks that the second subscription makes a new model request and completes; with .cache() in place of cancellableCache, its cancellation assertion still fails. All 2,342 core tests pass on JDK 17. PR description updated to match.