From 48bce8d69a0492f7b9876708a07ad71e6511407e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 22 Oct 2026 23:05:46 +0000 Subject: [PATCH] Drain live AI streams until the body closes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit readLiveResponse stopped at the result event and always cancelled the reader. The edge function closes the body a moment later, so that cancel could abort the fetch and drop the response. Keep reading until the stream ends, or for two seconds, and cancel only if it is still open. The answer is returned as soon as result arrives. Co-authored-by: Gáborjáni-Szabó Ádám --- src/lib/ai/live.ts | 76 +++++++++++++++++++++++++++++++++++++--- src/test/ai/live.test.ts | 73 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 143 insertions(+), 6 deletions(-) diff --git a/src/lib/ai/live.ts b/src/lib/ai/live.ts index faf9187c8..4e5fd5be4 100644 --- a/src/lib/ai/live.ts +++ b/src/lib/ai/live.ts @@ -17,14 +17,56 @@ export interface LiveGeneration { const DRAFT_INTERVAL_MS = 50; +/** + * After `result`, keep reading only this long. The function closes the body + * a moment later; waiting for that close avoids cancelling the fetch, and + * the cap stops a body that never closes from holding the reader open. + */ +export const LIVE_STREAM_DRAIN_MS = 2_000; + function isPhase(value: unknown): value is LivePhase { return value === 'research' || value === 'writing' || value === 'checking'; } +/** + * Reads until the body ends, the caller aborts, the read fails, or `capMs` + * elapses. True means the body ended, so the caller must not cancel it. + */ +async function drainUntilEnd( + reader: ReadableStreamDefaultReader, + signal: AbortSignal | undefined, + capMs: number, +): Promise { + let timer: ReturnType | undefined; + const cap = new Promise<'cap'>((resolve) => { + timer = setTimeout(() => resolve('cap'), capMs); + }); + try { + while (!signal?.aborted) { + const read = reader.read().then( + (chunk) => ({ kind: 'chunk' as const, chunk }), + (): { kind: 'error' } => ({ kind: 'error' }), + ); + const winner = await Promise.race([ + read, + cap.then((): { kind: 'cap' } => ({ kind: 'cap' })), + ]); + if (winner.kind !== 'chunk') return false; + if (winner.chunk.done) return true; + } + return false; + } finally { + if (timer) clearTimeout(timer); + } +} + /** * Reads a successful `ai` response. A JSON body (a server from before live * generation, or a cached reply) is returned as is; an event stream is * followed to its `result`, whose status decides success or failure. + * The answer is returned as soon as `result` arrives. Reading continues + * until the body ends so the fetch is not cancelled in the gap before the + * server closes it. */ export async function readLiveResponse( response: Response, @@ -82,21 +124,47 @@ export async function readLiveResponse( reader.cancel(signal?.reason).catch((): undefined => undefined); }; signal?.addEventListener('abort', stopReading, { once: true }); + + // Cancel only while the body is still open. Cancelling after `result`, + // before the server closes, aborts the fetch and drops the body. + let streamEnded = false; + let draining = false; + const release = () => { + signal?.removeEventListener('abort', stopReading); + if (timer) clearTimeout(timer); + timer = null; + if (!streamEnded) reader.cancel().catch((): undefined => undefined); + }; + try { while (!result) { if (signal?.aborted) throw new DOMException('The request was canceled.', 'AbortError'); const { done, value } = await reader.read(); if (signal?.aborted) throw new DOMException('The request was canceled.', 'AbortError'); - if (done) break; + if (done) { + streamEnded = true; + break; + } const parsed = consumeSSEChunk(buffer, decoder.decode(value, { stream: true })); buffer = parsed.buffer; handle(parsed.events); } if (!result) handle(consumeSSEChunk('', `${buffer}${decoder.decode()}\n\n`).events); + + // Hand the answer back before the tail, so a slow close cannot delay the UI. + if (result && !streamEnded) { + if (timer) clearTimeout(timer); + timer = null; + draining = true; + void drainUntilEnd(reader, signal, LIVE_STREAM_DRAIN_MS) + .then((ended) => { + streamEnded = ended; + }) + .catch((): undefined => undefined) + .finally(release); + } } finally { - signal?.removeEventListener('abort', stopReading); - if (timer) clearTimeout(timer); - reader.cancel().catch((): undefined => undefined); + if (!draining) release(); } const final = result as { status: number; body: unknown } | null; diff --git a/src/test/ai/live.test.ts b/src/test/ai/live.test.ts index 8f6a16a59..d9331bcb6 100644 --- a/src/test/ai/live.test.ts +++ b/src/test/ai/live.test.ts @@ -1,6 +1,6 @@ -import { describe, expect, it, vi } from 'vitest'; +import { afterEach, describe, expect, it, vi } from 'vitest'; import { FunctionsFetchError, FunctionsHttpError } from '@supabase/supabase-js'; -import { readLiveResponse } from '@/lib/ai/live'; +import { LIVE_STREAM_DRAIN_MS, readLiveResponse } from '@/lib/ai/live'; function eventStream(chunks: string[]): Response { const encoder = new TextEncoder(); @@ -102,4 +102,73 @@ describe('readLiveResponse', () => { const error = await readLiveResponse(eventStream([event('delta', { text: '{"a": "b' })]), { onDraft: () => {} }); expect(error).toBeInstanceOf(FunctionsFetchError); }); + + it('cancels immediately when the caller aborts', async () => { + let cancelCount = 0; + const encoder = new TextEncoder(); + const response = new Response(new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(event('delta', { text: 'partial' }))); + }, + // The body stays open, as a fetch can after the request signal aborts. + pull: () => new Promise(() => {}), + cancel() { cancelCount += 1; }, + }), { headers: { 'Content-Type': 'text/event-stream' } }); + + const controller = new AbortController(); + const reading = readLiveResponse(response, { onDraft: () => {} }, controller.signal); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(cancelCount).toBe(0); + controller.abort(); + // The abort listener cancels before the read loop resumes, not after the drain cap. + expect(cancelCount).toBe(1); + await expect(reading).rejects.toMatchObject({ name: 'AbortError' }); + expect(cancelCount).toBe(1); + }); +}); + +describe('draining the body after result', () => { + afterEach(() => { + vi.useRealTimers(); + }); + + function openResultStream(closeAfterMs: number | null, onCancel: () => void): Response { + const encoder = new TextEncoder(); + return new Response(new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(event('result', { status: 200, body: { output: { ok: true } } }))); + if (closeAfterMs !== null) setTimeout(() => controller.close(), closeAfterMs); + }, + cancel: onCancel, + }), { headers: { 'Content-Type': 'text/event-stream' } }); + } + + it('returns the result before the body closes and does not cancel once it ends', async () => { + vi.useFakeTimers(); + let cancelled = false; + const pending = readLiveResponse(openResultStream(400, () => { cancelled = true; }), { onDraft: () => {} }); + + await expect(pending).resolves.toEqual({ output: { ok: true } }); + expect(cancelled).toBe(false); + + await vi.advanceTimersByTimeAsync(400); + expect(cancelled).toBe(false); + // The drain cap must not cancel a body that already ended. + await vi.advanceTimersByTimeAsync(LIVE_STREAM_DRAIN_MS); + expect(cancelled).toBe(false); + }); + + it('cancels when the body stays open through the drain cap', async () => { + vi.useFakeTimers(); + let cancelled = false; + const pending = readLiveResponse(openResultStream(null, () => { cancelled = true; }), { onDraft: () => {} }); + + await expect(pending).resolves.toEqual({ output: { ok: true } }); + expect(cancelled).toBe(false); + + await vi.advanceTimersByTimeAsync(LIVE_STREAM_DRAIN_MS - 1); + expect(cancelled).toBe(false); + await vi.advanceTimersByTimeAsync(1); + expect(cancelled).toBe(true); + }); });