Repository navigation
[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
Description
Activity
Confirmed on main (f866804), and I opened #1621 for the propagation half. The public cancel API stays with #1341, and #1434 doesn't touch
cache(), so they compose.Repro: a
TestLlmthat returnsFlowable.never().doOnCancel(...), subscribe throughrunner.runAsync(...)orBaseLlmFlow.run(ctx), dispose after the model subscribes, anddoOnCancelnever fires. The mechanism differs from a framework subscriber holding the connection open:FlowableCachehas no path that cancels upstream at all.CacheSubscription.cancel()only removes the subscriber, so after a disconnect the step runs on with zero subscribers, and a tool call the model emits afterwards still executes.#1621 keeps
cache()and puts atakeUntilin front of it that fires when a subscriber cancels, so thetoList()side still gets the buffer and completion. I'll rebase onto #1565'srunSteponce it lands. Eight of nine tests pass there; the late-subscriber test needs updating for itsdefer.Thanks for pointing at the exact line, that saved digging. Would you rather take #1621 from here?
Hi @caps-xia, thank you for taking time and reporting this issue.
@sangkyoonnam, thank you for your contribution and We appreciate you taking the time to submit this pull request. I have reproduced this behavior and through local testing that your fix addresses it.
Currently, the issue and the associated PR are under review by our team and we will keep you updated if any additional information is required. Thank you.
Please keep #1621 — you're 8/9 tests in, and takeUntil-in-front-of-cache is exactly the right shape: late-subscriber replay stays intact while cancellation finally reaches upstream.
Your mechanism write-up also matches our production incident precisely: after the disconnect the step kept running with zero subscribers, and tool rounds the model emitted afterwards still executed. That's why we run a cooperative cancel flag (REST → per-run flag → checked at chunk/round boundaries) alongside — the two approaches compose.
I'd rather contribute validation than take over: happy to run #1621 against our multi-round tool-loop workloads (disconnect mid-stream, dispose during tool execution, late-subscriber replay) and report results here.
Validated #1621 (head a63b9cd) in our production environment — a multi-agent log-troubleshooting workload with multi-round tool loops over an OpenAI-compatible gateway (GLM), served via SSE.
Setup: built the PR branch locally, installed
google-adk 1.11.1-SNAPSHOT, switched our agent platform onto it, and ran our full regression suite first — 311 tests green, no behavioral drift.Scenario 1 — client disconnects mid-stream (the original incident path):
Started a streaming troubleshooting run, let it emit its first tool round (a member-lookup call), waited for the model to keep streaming, then closed the browser tab outright — no stop button, no cooperative cancel, just the client vanishing.
Result: the disconnect is caught on the model stream's own delivery path (the SSE write fails on the next chunk, and the trace shows
takeUntil → cacheexactly where the fix sits). After the broken pipe, zero further log activity — no additional tool rounds, no follow-up model calls, nothing running in the background. On the unfixed version this exact scenario kept generating for 3+ minutes and even executed tool rounds the model emitted after the disconnect.Scenario 2 — cooperative cancellation still works alongside:
Our platform also runs an application-level cancel (REST stop → per-run flag → checked at chunk/round boundaries). With #1621 underneath, a stop request still closes the run out gracefully with partial results preserved (825 chars of reasoning kept as a partial result) — the two mechanisms compose cleanly, no interference.
Takeaway: both halves of the problem in this issue are now covered — dispose-based cancellation propagates to the model stream (#1621), and our cooperative flag remains the application-level stop entry point. Thanks @sangkyoonnam — the takeUntil-in-front-of-cache shape held up exactly as designed under real workload.
Thanks for running it on the real workload, that's the evidence I couldn't produce here. I'll link this from #1621 so the reviewers see it.
Describe the Bug
When a client disconnects while an agent is streaming (e.g. the SSE consumer cancels its subscription to
Runner's event flow), the underlying model call keeps running to completion. Tokens continue to be billed with no consumer attached. In one production incident the model kept generating for 3+ minutes after the client went away (observed via gateway-side logs and token metering).The cause appears to be
BaseLlmFlow.javaline 530 (main at the time of writing):With RxJava
cache(), the upstream connection is only disposed when all subscribers of the cached flow cancel. As long as any subscriber attached by the framework itself (event handling/aggregation within the invocation) is alive, a user-side cancellation cannot reach the model stream.On top of that, there is no public API to cancel an in-flight run: neither
RunnernorInvocationContextnor the session exposes a "stop generating" handle. The only way to stop a long generation is to let it finish.To Reproduce
Expected behavior
Cancelling the downstream subscription (or any public cancellation entry point) must reach the model stream: the in-flight LLM request is disposed (or aborted at the next chunk/round boundary), and no further rounds are started.
Environment
cache()behavior is in core (BaseLlmFlow) and should affect any transportAdditional context
In production we could only implement "stop generating" with a cooperative cancellation protocol outside the framework: a per-run cancel flag (REST endpoint sets it) that the model-stream loop checks at every chunk and at tool-round boundaries, plus
setCancellablehooks for the current round. It works, but it is entirely application-side — the framework neither propagates the cancellation nor offers a canonical way to express it.Suggested directions (happy to discuss or contribute):
cache()with cancellation-propagating semantics (e.g.publish().refCount()-style sharing, or a scoped subscription that disposes the source when the invocation ends/cancels), or