Drain live AI streams until the body closes

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<Uint8Array>,
+  signal: AbortSignal | undefined,
+  capMs: number,
+): Promise<boolean> {
+  let timer: ReturnType<typeof setTimeout> | 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<T>(
   response: Response,
@@ -82,21 +124,47 @@ export async function readLiveResponse<T>(
     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;

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<void>(() => {}),
+      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);
+  });
 });