Drain live AI streams until the body closes (#204)

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: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Gáborjáni-Szabó Ádám <Programcode@users.noreply.github.com>

Author: Programcode

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