Chunked streaming body, parsing out a routing config along the way

chunked-streaming-parser.js

// Chunk size to pull from streams (256KB)
const CHUNK_BYTES = 256 * 1024;

// Buffer overlap size to handle model spanning chunk boundaries
// 50 chars for model value + ~20 chars for "model": " prefix
const OVERLAP_BUFFER_SIZE = 70;

function extractModelFromText(text) {
  const match = text.match(/"model"\s*:\s*"([^"]+)"/);
  return match ? match[1] : null;
}

// Extract model by streaming the body (for larger requests)
// This (potentially) avoids loading the entire body into memory
// Returns { model, bodyStream } where bodyStream is for forwarding
async function extractModelFromStream(request) {
  // create a dedicated *streaming* decoder for this stream
  const textDecoder = new TextDecoder();
  const reader = request.body.getReader({ mode: "byob" });

  const doRead = async () => {
    const { done, value } = await reader.read(
      new Uint8Array(new ArrayBuffer(CHUNK_BYTES)),
      { min: CHUNK_BYTES }
    );
    return {
      done,
      chunk:
        value != null && value.byteLength > 0
          ? new Uint8Array(value.buffer, 0, value.byteLength)
          : null,
    };
  };

  const chunksSoFar = [];
  let buffer = "";
  let model = null;
  while (true) {
    const { done, chunk } = await doRead();

    if (chunk != null) {
      chunksSoFar.push(chunk);
      buffer += textDecoder.decode(chunk, { stream: true });
    }

    model = extractModelFromText(buffer);

    if (done || model) {
      break;
    }

    // Keep only the last OVERLAP_BUFFER_SIZE characters to handle
    // model spanning chunk boundaries, discard the rest to save memory
    if (buffer.length > OVERLAP_BUFFER_SIZE) {
      buffer = buffer.slice(-OVERLAP_BUFFER_SIZE);
    }
  }

  const body = new ReadableStream({
    async pull(controller) {
      if (chunksSoFar.length > 0) {
        controller.enqueue(chunksSoFar.shift());
        return;
      }
      const { done, chunk } = await doRead();
      if (chunk != null) {
        controller.enqueue(chunk);
      }
      if (done) {
        controller.close();
      }
    },
    cancel(reason) {
      return reader.cancel(reason);
    },
  });

  return { model, body };
}

export default {
  async fetch(request, env, ctx) {
    const headers = new Headers(request.headers);
    const requestOverrides = { headers };

    const { model, body } = await extractModelFromStream(request);
    if (model) {
      // modify the request however you want, eg add headers for routing
      headers.set("x-api-model", model);
    }
    requestOverrides.body = body;
    
    // The above will transfer with no content-length and transfer-encoding: chunked
    // If your origin needs content-length then do this instead:

    // const contentLength = parseInt(
    //   request.headers.get("content-length") || "0",
    //   10
    // );
    // requestOverrides.body = contentLength
    //   ? body.pipeThrough(new FixedLengthStream(contentLength))
    //   : body;

    return fetch(new Request(request, requestOverrides));
  }
}