ai/src/util/write-to-server-response.ts

AI SDK by Vercel - build apps like ChatGPT, Claude, Gemini, and more with a single interface for any model using the Vercel AI Gateway or go direct to OpenAI, Anthropic, Google, or any other model provider.

import type { ServerResponse } from 'node:http';

type FlushableServerResponse = ServerResponse & {
  flush?: () => void;
};

/**
 * Writes the content of a stream to a server response.
 */
export function writeToServerResponse({
  response,
  status,
  statusText,
  headers,
  stream,
}: {
  response: ServerResponse;
  status?: number;
  statusText?: string;
  headers?: Record<string, string | number | string[]>;
  stream: ReadableStream<Uint8Array>;
}): void {
  const statusCode = status ?? 200;
  if (statusText !== undefined) {
    response.writeHead(statusCode, statusText, headers);
  } else {
    response.writeHead(statusCode, headers);
  }

  const reader = stream.getReader();
  const read = async () => {
    try {
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        // Respect backpressure: if write() returns false, wait for 'drain' event
        const canContinue = response.write(value);
        const flush = (response as FlushableServerResponse).flush;
        if (typeof flush === 'function') {
          flush.call(response);
        }

        if (!canContinue) {
          await new Promise<void>(resolve => {
            response.once('drain', resolve);
          });
        }
      }
    } catch (error) {
      throw error;
    } finally {
      response.end();
    }
  };

  read();
}