stream.ts

export async function fetchStream(
  url: RequestInfo | URL,
  handler: (text: string) => any,
): Promise<Response> {
  const response = await fetch(url);
  (async () => {
    const reader = response.body!.getReader();
    const decoder = new TextDecoder();
    let buffer = "";
    while (true) {
      const { done, value } = await reader.read();
      if (done) {
        if (buffer) handler(buffer);
        break;
      }
      buffer += decoder.decode(value, { stream: true });
      const lines = buffer.split("\n");
      buffer = lines.pop()!;
      for (const line of lines) handler(line);
    }
  })();
  return response;
}

export function fetchJsonStream<T = any>(
  url: RequestInfo | URL,
  handler: (value: T) => any,
): Promise<Response> {
  return fetchStream(url, (text) => handler(JSON.parse(text)));
}

export interface StreamedResponse<T = string> extends Response {
  send(msg: T): void;
  close(): void;
}

export function responseStream(
  init: ResponseInit = { headers: { "Content-Type": "text/event-stream" } },
): StreamedResponse {
  let send: (msg: string) => void = undefined!;
  let close: () => void = undefined!;
  const stream = new ReadableStream({
    start(controller) {
      const encoder = new TextEncoder();
      send = (msg) => controller.enqueue(encoder.encode(`${msg}\n`));
      close = () => controller.close();
    },
  });
  return Object.assign(new Response(stream, init), { send, close });
}

export function responseJsonStream<T = any>(
  init: ResponseInit = { headers: { "Content-Type": "application/jsonl" } },
): StreamedResponse<T> {
  const response = responseStream(init);
  const { send: sendText } = response;
  const send = (value: T) => sendText(JSON.stringify(value));
  return Object.assign(response, { send });
}