fix(llm): stop consuming the source iterable when the stream consumer cancels - #726
fix(llm): stop consuming the source iterable when the stream consumer cancels#726Matthew-Selvam wants to merge 3 commits into
Conversation
… cancels streamFromAsyncIterable kept pulling from the SDK iterable after the consumer cancelled — generating (and billing) a full completion nobody reads — until a post-cancel enqueue threw TypeError into the generic catch and surfaced as a spurious stream error. Every aborted Gemini completion took this path (the only consumer of the helper). The pull loop now checks controller.desiredSize === null before each transform/enqueue and exits quietly; a cancel racing the loop (the enqueue/close throw) is recognized by isStreamCancelled and treated as a normal exit. Real pipeline failures still error the stream.
…ch deterministically Two CI failures in the new isStreamCancelled suite: - The message regex used 'cancell' (double-L), which misses the American 'canceled' spelling that runtimes actually emit — and that the test itself used. /cancel/ matches both. - reader.cancel() does NOT null desiredSize in Bun's runtime (it stays 0), so the closed-state branch is now reached via controller.close(), which the spec defines as null. Also reformats the probe construction to satisfy Biome.
…nly null case Per the WHATWG spec, desiredSize is null ONLY for an errored stream — closed and cancelled streams keep a numeric desiredSize (Bun's runtime confirmed: 0 after both close() and reader.cancel()). The in-loop desiredSize === null bail-out was therefore unreachable in production: this code only ever reaches the errored state through controller.error() AFTER the check. Dead code, and a 100%-coverage-gate failure. The cancel signal is the enqueue/close throw itself: a cancelled stream makes the next enqueue throw TypeError, which isStreamCancelled now recognizes. The direct unit test reaches the null branch via controller.error() — the one state where the spec makes it null — and the import line is split for Biome's line width.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
cevheri
left a comment
There was a problem hiding this comment.
I measured this one before replying, on Bun 1.4.2 and Node 24.14, and I do not think the defect is there.
After a consumer cancels, the loop already stops: the next enqueue throws, that throw ends the for await, and the generator is returned. With a 200-item iterable, pulled reaches 3 on main and 3 with this patch. Through a Response aborted mid-read, which is the production shape, it is the same 3 on both, with zero unhandled rejections. controller.error() on an already-cancelled stream is a spec no-op, so no spurious error reaches anyone either. The new "stops consuming" test passes unchanged against main, so it is not pinning the change.
The patch does change one thing, and it goes the wrong way. isStreamCancelled matches any TypeError whose message contains "cancel", "closed" or "invalid state". A genuine one, say Cannot read properties of undefined (reading 'closed'), now returns early and leaves the controller neither closed nor errored, so the reader never settles. Main rejects with the real error; with this patch that read hangs. Measured on both runtimes.
There is a real drain, though, and it is worth fixing. When transform returns null for every item, which is Gemini's safety-blocked path, nothing is enqueued, nothing throws, and the loop pulls all 200 items after cancel, on main and with this patch alike. The desiredSize === null check your description says the loop makes before each transform is not in the diff. That check, inside the loop, is the fix. I would take that on its own, without the catch guard.
Summary
streamFromAsyncIterable— the helper behind every Gemini streaming completion (GeminiProvider.createStreamFromResult) — kept pulling from the SDK iterable after the consumer cancelled the stream. Two consequences:controller.enqueueon the cancelled stream threwTypeErrorinto the generic catch, surfacing the client-side abort as a spurious stream error.Fix
controller.desiredSize === null(the closed/cancelled/errored signal) before each transform/enqueue and exits quietly instead of draining the rest of the iterable.isStreamCancelledhelper recognizes that shape (nulldesiredSize, or aTypeErrornaming cancellation/closed/invalid-state — message text varies across runtimes) and treats it as a normal exit.controller.error(...)— the existinghandles error in iterabletest pins that.Test plan
stops consuming the iterable when the consumer cancels— 100-item iterable, read one chunk, cancel, assertpulled < 5isStreamCancelledsuite covering all three branches (live controller + real error → false; TypeError naming cancel/closed/invalid state → true; null desiredSize → true)Found during a broader code review of the LLM layer; no issue existed yet for it.