Shade/packages/shade-files/src/client/inline-threshold.ts

/**
 * Decide whether a write should travel inline (base64 inside the RPC
 * envelope) or via a dedicated `@shade/transfer` stream.
 *
 * Threshold: 256 KiB plaintext. Anything strictly above goes through the
 * stream path; anything ≤ goes inline.
 *
 * For known-size inputs (`Uint8Array`, `Blob`, `File`) the decision is
 * cheap: compare `byteLength`/`size` against the threshold.
 *
 * For `ReadableStream` we cannot know the size up front. We pull chunks
 * into a temporary buffer until we either:
 *  - hit EOF before the threshold → inline (we have all the bytes)
 *  - cross the threshold → streams (the buffered prefix + the rest of the
 *    stream are returned via a fresh `ReadableStream` so the caller can
 *    feed it to `shade.upload`).
 */
/** Plaintext size at which inline transitions to streams. */
export const INLINE_THRESHOLD = 256 * 1024;

export type InlineDecision =
  | {
      kind: 'inline';
      bytes: Uint8Array;
      contentType?: string;
    }
  | {
      kind: 'streams';
      stream: ReadableStream<Uint8Array>;
      /** Plaintext size when known (Blob/File). undefined for raw streams. */
      size?: number;
      contentType?: string;
    };

/**
 * Public input shape for `FileClient.write`. Runtime discriminator helper.
 */
export type WriteSource =
  | Uint8Array
  | Blob
  | File
  | ReadableStream<Uint8Array>
  | { stream: ReadableStream<Uint8Array>; size: number; contentType?: string };

export async function decideInline(input: WriteSource): Promise<InlineDecision> {
  // 1. Uint8Array — direct size check
  if (input instanceof Uint8Array) {
    if (input.byteLength <= INLINE_THRESHOLD) {
      return { kind: 'inline', bytes: input };
    }
    return {
      kind: 'streams',
      stream: uint8ArrayToStream(input),
      size: input.byteLength,
    };
  }

  // 2. Blob / File — known size (Blob.type is a string; File extends Blob)
  if (typeof Blob !== 'undefined' && input instanceof Blob) {
    const blob: Blob = input;
    const contentType = blob.type === '' ? undefined : blob.type;
    if (blob.size <= INLINE_THRESHOLD) {
      const bytes = new Uint8Array(await blob.arrayBuffer());
      return contentType !== undefined
        ? { kind: 'inline', bytes, contentType }
        : { kind: 'inline', bytes };
    }
    return contentType !== undefined
      ? { kind: 'streams', stream: blob.stream(), size: blob.size, contentType }
      : { kind: 'streams', stream: blob.stream(), size: blob.size };
  }

  // 3. Pre-wrapped { stream, size } — trust caller's declared size
  if (
    typeof input === 'object' &&
    input !== null &&
    'stream' in input &&
    'size' in input &&
    (input as { stream: unknown }).stream instanceof ReadableStream
  ) {
    const wrapped = input as { stream: ReadableStream<Uint8Array>; size: number; contentType?: string };
    if (wrapped.size <= INLINE_THRESHOLD) {
      const bytes = await drainStream(wrapped.stream, wrapped.size);
      const contentType = wrapped.contentType;
      return contentType !== undefined
        ? { kind: 'inline', bytes, contentType }
        : { kind: 'inline', bytes };
    }
    const contentType = wrapped.contentType;
    return contentType !== undefined
      ? { kind: 'streams', stream: wrapped.stream, size: wrapped.size, contentType }
      : { kind: 'streams', stream: wrapped.stream, size: wrapped.size };
  }

  // 4. Bare ReadableStream — peek until threshold or EOF
  if (input instanceof ReadableStream) {
    return await peekStream(input);
  }

  throw new TypeError(
    `decideInline: unsupported input type ${Object.prototype.toString.call(input)}`,
  );
}

/**
 * Drain a stream into a Uint8Array, with a soft cap at `expected` + slack.
 * Used for the inline path when the caller declared a size.
 */
async function drainStream(
  stream: ReadableStream<Uint8Array>,
  expected: number,
): Promise<Uint8Array> {
  const reader = stream.getReader();
  const chunks: Uint8Array[] = [];
  let total = 0;
  try {
    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      if (value === undefined) continue;
      chunks.push(value);
      total += value.byteLength;
      if (total > expected + INLINE_THRESHOLD) {
        throw new Error(
          `decideInline: stream produced more bytes (${total}) than declared size (${expected})`,
        );
      }
    }
  } finally {
    reader.releaseLock();
  }
  return concat(chunks, total);
}

/**
 * Peek a `ReadableStream` of unknown length: buffer up to `INLINE_THRESHOLD + 1`
 * bytes. If EOF first, return inline. Otherwise reconstruct a stream that
 * yields the buffered prefix followed by the remainder.
 */
async function peekStream(stream: ReadableStream<Uint8Array>): Promise<InlineDecision> {
  const reader = stream.getReader();
  const buffered: Uint8Array[] = [];
  let total = 0;
  try {
    while (total <= INLINE_THRESHOLD) {
      const { value, done } = await reader.read();
      if (done) {
        reader.releaseLock();
        return { kind: 'inline', bytes: concat(buffered, total) };
      }
      if (value === undefined) continue;
      buffered.push(value);
      total += value.byteLength;
    }
    // We have at least INLINE_THRESHOLD + 1 bytes buffered. Promote to streams.
    const reconstructed = reconstructStream(buffered, reader);
    return { kind: 'streams', stream: reconstructed };
  } catch (err) {
    reader.releaseLock();
    throw err;
  }
}

/**
 * Structural mirror of WHATWG `ReadableStreamDefaultReader<Uint8Array>`.
 *
 * The disjoint union shape with `value?: T | undefined` is the lowest
 * common denominator across every lib environment we care about:
 * - `bun-types` emits `{ done: true; value?: undefined }`
 * - `lib.dom`   emits `{ done: true; value?: T }`
 * - `node:stream/web` emits the union form
 *
 * `value?: T | undefined` is assignable from all three. A flat
 * `{ value?: T; done: boolean }` is rejected by
 * `exactOptionalPropertyTypes` because the present branches require
 * `value: T`. Defining it as an explicit union avoids the trap.
 */
type MinimalReadResult<T> =
  | { done: false; value: T }
  | { done: true; value?: T | undefined };

interface MinimalReader<T> {
  read(): Promise<MinimalReadResult<T>>;
  cancel(reason?: unknown): Promise<void>;
  releaseLock(): void;
}

function reconstructStream(
  prefix: Uint8Array[],
  reader: MinimalReader<Uint8Array>,
): ReadableStream<Uint8Array> {
  let prefixIdx = 0;
  return new ReadableStream<Uint8Array>({
    async pull(controller) {
      if (prefixIdx < prefix.length) {
        controller.enqueue(prefix[prefixIdx]!);
        prefixIdx++;
        return;
      }
      const { value, done } = await reader.read();
      if (done) {
        controller.close();
        reader.releaseLock();
        return;
      }
      if (value !== undefined) controller.enqueue(value);
    },
    async cancel(reason) {
      try {
        await reader.cancel(reason);
      } finally {
        reader.releaseLock();
      }
    },
  });
}

function uint8ArrayToStream(bytes: Uint8Array): ReadableStream<Uint8Array> {
  return new ReadableStream<Uint8Array>({
    start(controller) {
      controller.enqueue(bytes);
      controller.close();
    },
  });
}

function concat(chunks: Uint8Array[], total: number): Uint8Array {
  const out = new Uint8Array(total);
  let offset = 0;
  for (const c of chunks) {
    out.set(c, offset);
    offset += c.byteLength;
  }
  return out;
}