@ai-sdk/provider-utils/src/read-response-with-size-limit.ts

import { cancelResponseBody } from './cancel-response-body';
import { DownloadError } from './download-error';

/**
 * Default maximum download size: 2 GiB.
 *
 * `fetch().arrayBuffer()` has ~2x peak memory overhead (undici buffers the
 * body internally, then creates the JS ArrayBuffer), so very large downloads
 * risk exceeding the default V8 heap limit on 64-bit systems and terminating
 * the process with an out-of-memory error.
 *
 * Setting this limit converts an unrecoverable OOM crash into a catchable
 * `DownloadError`.
 */
export const DEFAULT_MAX_DOWNLOAD_SIZE = 2 * 1024 * 1024 * 1024;

/**
 * Reads a fetch Response body with a size limit to prevent memory exhaustion.
 *
 * Checks the Content-Length header for early rejection, then reads the body
 * incrementally via ReadableStream and aborts with a DownloadError when the
 * limit is exceeded.
 *
 * @param response - The fetch Response to read.
 * @param url - The URL being downloaded (used in error messages).
 * @param maxBytes - Maximum allowed bytes. Defaults to DEFAULT_MAX_DOWNLOAD_SIZE.
 * @param onChunk - Optional. Receives each chunk instead of the chunks being
 * collected, so the body is not held in memory, and the returned array is empty.
 * The chunk may be a view on a larger buffer, copy what you keep.
 * @returns A Uint8Array containing the response body.
 * @throws DownloadError if the response exceeds maxBytes.
 */
export async function readResponseWithSizeLimit({
  response,
  url,
  maxBytes = DEFAULT_MAX_DOWNLOAD_SIZE,
  onChunk,
}: {
  response: Response;
  url: string;
  maxBytes?: number;
  onChunk?: (chunk: Uint8Array) => void;
}): Promise<Uint8Array> {
  // Early rejection based on Content-Length header
  const contentLength = response.headers.get('content-length');
  if (contentLength != null) {
    const length = parseInt(contentLength, 10);
    if (!isNaN(length) && length > maxBytes) {
      // Cancel the body so the underlying connection is released back to the
      // pool instead of being left open until the socket is exhausted.
      await cancelResponseBody(response);
      throw new DownloadError({
        url,
        message: `Download of ${url} exceeded maximum size of ${maxBytes} bytes (Content-Length: ${length}).`,
      });
    }
  }

  const body = response.body;

  // Handle missing body (empty responses)
  if (body == null) {
    return new Uint8Array(0);
  }

  const reader = body.getReader();
  const chunks: Uint8Array[] = [];
  let totalBytes = 0;

  try {
    while (true) {
      const { done, value } = await reader.read();

      if (done) {
        break;
      }

      totalBytes += value.length;

      if (totalBytes > maxBytes) {
        throw new DownloadError({
          url,
          message: `Download of ${url} exceeded maximum size of ${maxBytes} bytes.`,
        });
      }

      if (onChunk != null) {
        onChunk(value);
        continue;
      }

      // A small chunk can keep a much larger backing buffer alive. Copy such
      // views before retaining them. Unlike Uint8Array, Node Buffer's slice()
      // returns another view, so use the Uint8Array constructor to copy.
      chunks.push(
        value.byteLength < value.buffer.byteLength
          ? new Uint8Array(value)
          : value,
      );
    }
  } finally {
    try {
      await reader.cancel();
    } catch {
      // Ignore cancel errors so the original rejection is preserved.
    } finally {
      reader.releaseLock();
    }
  }

  // Concatenate chunks into a single Uint8Array
  const result = new Uint8Array(onChunk != null ? 0 : totalBytes);
  let offset = 0;
  for (const chunk of chunks) {
    result.set(chunk, offset);
    offset += chunk.length;
  }

  return result;
}