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); + }); });