gitfor/platform

import { randomId } from "../crypto";
import { run } from "../db";
import { recordUsage } from "./record";

/**
 * Wraps a body so `onDone(bytes)` runs exactly once, when it has fully streamed or the client cancelled
 * (bytes already delivered are still metered).
 */
export function countingBody(
  body: ReadableStream<Uint8Array> | null,
  onDone: (bytes: number) => void,
  onChunk?: (chunk: Uint8Array) => void,
): ReadableStream<Uint8Array> | null {
  if (!body) {
    onDone(0);
    return null;
  }
  const reader = body.getReader();
  let bytes = 0;
  let done = false;
  const finish = () => {
    if (done) return;
    done = true;
    try {
      onDone(bytes);
    } catch (error) {
      console.warn("metering callback failed", error);
    }
  };
  return new ReadableStream<Uint8Array>({
    async pull(controller) {
      try {
        const next = await reader.read();
        if (next.done) {
          finish();
          controller.close();
          return;
        }
        bytes += next.value.byteLength;
        onChunk?.(next.value);