fix: time out stalled streams

Author: victor0602

src/client/http.ts

@@ -18,6 +18,84 @@ export interface RequestOpts {
   authStyle?: 'bearer' | 'x-api-key';
 }
 
+function timeoutError(message: string): DOMException {
+  return new DOMException(message, 'TimeoutError');
+}
+
+function withIdleTimeout(
+  response: Response,
+  timeoutMs: number,
+  abortController: AbortController,
+): Response {
+  if (!response.body) return response;
+
+  const reader = response.body.getReader();
+  let released = false;
+
+  const releaseReader = (): void => {
+    if (released) return;
+    released = true;
+    reader.releaseLock();
+  };
+
+  const cancelRequest = async (reason?: unknown): Promise<void> => {
+    if (!abortController.signal.aborted) abortController.abort(reason);
+    try {
+      await reader.cancel(reason);
+    } finally {
+      releaseReader();
+    }
+  };
+
+  const body = new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      let timer: ReturnType<typeof setTimeout> | undefined;
+
+      try {
+        const result = await new Promise<Awaited<ReturnType<typeof reader.read>>>((resolve, reject) => {
+          timer = setTimeout(() => {
+            const error = timeoutError(`Stream received no data for ${timeoutMs}ms.`);
+            abortController.abort(error);
+            reject(error);
+          }, timeoutMs);
+
+          reader.read().then(resolve, reject);
+        });
+
+        if (result.done) {
+          releaseReader();
+          controller.close();
+        } else {
+          controller.enqueue(result.value);
+        }
+      } catch (error) {
+        await cancelRequest(error).catch(() => {});
+        controller.error(error);
+      } finally {
+        if (timer) clearTimeout(timer);
+      }
+    },
+    cancel(reason) {
+      return cancelRequest(reason);
+    },
+  });
+
+  const timedResponse = new Response(body, {
+    status: response.status,
+    statusText: response.statusText,
+    headers: response.headers,
+  });
+
+  // Response's constructor does not carry fetch metadata across to a wrapped body.
+  Object.defineProperties(timedResponse, {
+    url: { value: response.url },
+    redirected: { value: response.redirected },
+    type: { value: response.type },
+  });
+
+  return timedResponse;
+}
+
 export async function request(config: Config, opts: RequestOpts): Promise<Response> {
   const isFormData = typeof FormData !== 'undefined' && opts.body instanceof FormData;
 
@@ -54,16 +132,32 @@ export async function request(config: Config, opts: RequestOpts): Promise<Respon
 
   const timeoutMs = (opts.timeout ?? config.timeout) * 1000;
 
-  const res = await fetch(opts.url, {
-    method: opts.method ?? 'GET',
-    headers,
-    body: opts.body
-      ? isFormData
-        ? (opts.body as FormData)
-        : JSON.stringify(opts.body)
-      : undefined,
-    signal: opts.stream ? undefined : AbortSignal.timeout(timeoutMs),
-  });
+  const streamAbortController = opts.stream ? new AbortController() : undefined;
+  const headerTimeout = streamAbortController
+    ? setTimeout(() => {
+        streamAbortController.abort(timeoutError(`Request headers were not received within ${timeoutMs}ms.`));
+      }, timeoutMs)
+    : undefined;
+
+  let res: Response;
+  try {
+    res = await fetch(opts.url, {
+      method: opts.method ?? 'GET',
+      headers,
+      body: opts.body
+        ? isFormData
+          ? (opts.body as FormData)
+          : JSON.stringify(opts.body)
+        : undefined,
+      signal: streamAbortController?.signal ?? AbortSignal.timeout(timeoutMs),
+    });
+  } finally {
+    if (headerTimeout) clearTimeout(headerTimeout);
+  }
+
+  if (streamAbortController) {
+    res = withIdleTimeout(res, timeoutMs, streamAbortController);
+  }
 
   if (config.verbose) {
     process.stderr.write(`< ${res.status} ${res.statusText}\n`);

test/client/http.test.ts

@@ -1,5 +1,5 @@
-import { describe, it, expect, afterEach } from 'bun:test';
-import { requestJson } from '../../src/client/http';
+import { describe, it, expect, afterEach, spyOn } from 'bun:test';
+import { request, requestJson } from '../../src/client/http';
 import { CLI_VERSION } from '../../src/version';
 import { createMockServer, jsonResponse, type MockServer } from '../helpers/mock-server';
 import type { Config } from '../../src/config/schema';
@@ -93,4 +93,158 @@ describe('HTTP client', () => {
       requestJson(config, { url: `${server.url}/v1/test` }),
     ).rejects.toThrow('Rate limit');
   });
+
+  it('aborts a streaming request when response headers stall', async () => {
+    let aborted = false;
+    const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((
+      (_input: string | URL | Request, init?: RequestInit) =>
+        new Promise<Response>((_resolve, reject) => {
+          init?.signal?.addEventListener('abort', () => {
+            aborted = true;
+            reject(init.signal?.reason);
+          }, { once: true });
+        })
+    ) as unknown as typeof fetch);
+
+    try {
+      const config = makeConfig('https://example.com');
+      await expect(request(config, {
+        url: 'https://example.com/stream',
+        stream: true,
+        timeout: 0.02,
+        noAuth: true,
+      })).rejects.toMatchObject({ name: 'TimeoutError' });
+      expect(aborted).toBe(true);
+    } finally {
+      fetchSpy.mockRestore();
+    }
+  });
+
+  it('aborts and cancels a streaming response when its body stalls', async () => {
+    let aborted = false;
+    let cancelled = false;
+    const encoder = new TextEncoder();
+    const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((
+      (_input: string | URL | Request, init?: RequestInit) => {
+        init?.signal?.addEventListener('abort', () => {
+          aborted = true;
+        }, { once: true });
+
+        const body = new ReadableStream<Uint8Array>({
+          start(controller) {
+            controller.enqueue(encoder.encode('data: first\n\n'));
+          },
+          cancel() {
+            cancelled = true;
+          },
+        });
+        return Promise.resolve(new Response(body));
+      }
+    ) as unknown as typeof fetch);
+
+    try {
+      const config = makeConfig('https://example.com');
+      const response = await request(config, {
+        url: 'https://example.com/stream',
+        stream: true,
+        timeout: 0.02,
+        noAuth: true,
+      });
+      const reader = response.body!.getReader();
+
+      expect(new TextDecoder().decode((await reader.read()).value)).toBe('data: first\n\n');
+      await expect(reader.read()).rejects.toMatchObject({ name: 'TimeoutError' });
+      expect(aborted).toBe(true);
+      expect(cancelled).toBe(true);
+    } finally {
+      fetchSpy.mockRestore();
+    }
+  });
+
+  it('allows an active stream to outlive a single timeout interval', async () => {
+    let aborted = false;
+    const encoder = new TextEncoder();
+    const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((
+      (_input: string | URL | Request, init?: RequestInit) => {
+        init?.signal?.addEventListener('abort', () => {
+          aborted = true;
+        }, { once: true });
+
+        let interval: ReturnType<typeof setInterval> | undefined;
+        const body = new ReadableStream<Uint8Array>({
+          start(controller) {
+            let chunk = 0;
+            interval = setInterval(() => {
+              controller.enqueue(encoder.encode(String(chunk)));
+              chunk += 1;
+              if (chunk === 6) {
+                clearInterval(interval);
+                controller.close();
+              }
+            }, 10);
+          },
+          cancel() {
+            if (interval) clearInterval(interval);
+          },
+        });
+        return Promise.resolve(new Response(body));
+      }
+    ) as unknown as typeof fetch);
+
+    try {
+      const config = makeConfig('https://example.com');
+      const response = await request(config, {
+        url: 'https://example.com/stream',
+        stream: true,
+        timeout: 0.03,
+        noAuth: true,
+      });
+      const chunks: string[] = [];
+
+      for await (const chunk of response.body!) {
+        chunks.push(new TextDecoder().decode(chunk));
+      }
+
+      expect(chunks).toEqual(['0', '1', '2', '3', '4', '5']);
+      expect(aborted).toBe(false);
+    } finally {
+      fetchSpy.mockRestore();
+    }
+  });
+
+  it('cancels the network request when the response consumer cancels', async () => {
+    let aborted = false;
+    let cancelled = false;
+    const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((
+      (_input: string | URL | Request, init?: RequestInit) => {
+        init?.signal?.addEventListener('abort', () => {
+          aborted = true;
+        }, { once: true });
+
+        const body = new ReadableStream<Uint8Array>({
+          cancel() {
+            cancelled = true;
+          },
+        });
+        return Promise.resolve(new Response(body));
+      }
+    ) as unknown as typeof fetch);
+
+    try {
+      const config = makeConfig('https://example.com');
+      const response = await request(config, {
+        url: 'https://example.com/stream',
+        stream: true,
+        timeout: 1,
+        noAuth: true,
+      });
+
+      await response.body!.cancel('consumer stopped');
+
+      expect(aborted).toBe(true);
+      expect(cancelled).toBe(true);
+    } finally {
+      fetchSpy.mockRestore();
+    }
+  });
 });