@timber-js/app 0.2.0-alpha.19 → 0.2.0-alpha.191

package/src/server/html-injector-core.ts

@@ -0,0 +1,495 @@
+/**
+ * HTML injector core — pure stateful helpers for streaming HTML post-processing.
+ *
+ * These helpers contain the actual byte/string logic for the held-flush
+ * streaming pipeline (buffered transforms, suffix moving, head injection,
+ * RSC flight interleaving). They are stream-shape agnostic: they know
+ * nothing about Web `TransformStream` or Node.js `Transform` and contain
+ * no `setImmediate` scheduling or controller plumbing.
+ *
+ * The Web Stream wrapper (`html-injectors.ts`) and the Node Stream wrapper
+ * (`node-stream-transforms.ts`) compose these helpers behind whichever
+ * stream shape their target runtime needs:
+ *
+ *   - On Cloudflare/edge: Web `TransformStream` wraps the helpers.
+ *   - On Node.js/Bun: Node `Transform` wraps the helpers (faster — C++ streams).
+ *
+ * Keeping this module split per-shape (rather than one file exporting both)
+ * preserves the build-time tree-shake: Cloudflare bundles never import
+ * `node:stream` even transitively. Only the helpers below are shared.
+ *
+ * **Pure functions only.** Anything that touches a stream API, schedules
+ * a tick, or owns a controller belongs in the wrapper, not here.
+ *
+ * Design docs: 02-rendering-pipeline.md §"Streaming Constraints", 18-build-system.md §"Entry Files"
+ */
+
+import { createMachine, type Machine } from '../utils/state-machine.js';
+import { flightChunkScript } from './flight-scripts.js';
+import {
+  flightInjectionTransitions,
+  isHtmlDone,
+  isPullDone,
+  type FlightInjectionState,
+  type FlightInjectionEvent,
+} from './flight-injection-state.js';
+import { withTimeout, RenderTimeoutError } from './render-timeout.js';
+
+// ─── Encoders / Constants ────────────────────────────────────────────────────
+
+const encoder = new TextEncoder();
+
+/** Closing tags React Fizz emits inside the shell chunk. */
+export const SUFFIX = '</body></html>';
+/** UTF-8 bytes of {@link SUFFIX}. Cached so wrappers don't re-encode per request. */
+export const SUFFIX_BYTES: Uint8Array = encoder.encode(SUFFIX);
+
+/** Encode a UTF-8 string. Convenience for wrappers that need a single source of truth. */
+export function encodeUtf8(text: string): Uint8Array {
+  return encoder.encode(text);
+}
+
+// ─── Buffered Aggregator ─────────────────────────────────────────────────────
+
+/**
+ * Pure buffer state for the buffered transform.
+ *
+ * Both Web and Node wrappers want the same behaviour: collect chunks that
+ * arrive in the same event-loop tick, decide whether the buffer has grown
+ * beyond `maxBufferByteLength` and must flush synchronously, then emit a
+ * single concatenated chunk on the next tick.
+ *
+ * The only thing the two wrappers disagree on is *who schedules the tick*
+ * (`setImmediate` callback vs Promise-based) and *how to push the merged
+ * chunk* (`controller.enqueue` vs `transform.push`). This class owns
+ * everything else.
+ */
+export class BufferAggregator {
+  private chunks: Uint8Array[] = [];
+  private byteLength = 0;
+
+  constructor(private readonly maxBufferByteLength: number = Infinity) {}
+
+  /**
+   * Append a chunk. Returns `true` if the caller should flush synchronously
+   * right now (the buffer has reached the byte cap), `false` if a deferred
+   * tick-end flush is sufficient.
+   */
+  append(chunk: Uint8Array): boolean {
+    this.chunks.push(chunk);
+    this.byteLength += chunk.byteLength;
+    return this.byteLength >= this.maxBufferByteLength;
+  }
+
+  /** Drain the buffer into a single concatenated chunk, or `null` if empty. */
+  drain(): Uint8Array | null {
+    if (this.chunks.length === 0) return null;
+
+    const merged = new Uint8Array(this.byteLength);
+    let offset = 0;
+    for (const chunk of this.chunks) {
+      merged.set(chunk, offset);
+      offset += chunk.byteLength;
+    }
+    this.chunks = [];
+    this.byteLength = 0;
+    return merged;
+  }
+
+  get bufferedByteLength(): number {
+    return this.byteLength;
+  }
+
+  get isEmpty(): boolean {
+    return this.chunks.length === 0;
+  }
+}
+
+// ─── Suffix Mover ────────────────────────────────────────────────────────────
+
+/**
+ * Result of feeding one chunk through the suffix-mover state.
+ *
+ * - `passthrough` — already found and removed the suffix on a previous chunk.
+ *   Caller emits the chunk unchanged.
+ * - `noSuffix` — current chunk does not contain the suffix. When `emitSafe`
+ *   is set, the caller must emit those bytes (the safe head minus the
+ *   retained tail). When `emitSafe` is null, the entire chunk was absorbed
+ *   into the tail buffer and there is nothing to emit yet.
+ * - `suffixFound` — caller emits `before` and `after` (skipping the suffix
+ *   itself). The suffix will be re-emitted later by {@link suffixFlush}.
+ */
+export type SuffixChunkResult =
+  | { kind: 'passthrough' }
+  | { kind: 'noSuffix'; emitSafe?: Uint8Array }
+  | { kind: 'suffixFound'; before: Uint8Array | null; after: Uint8Array | null };
+
+/**
+ * State for the suffix mover. Works entirely in bytes — no TextDecoder
+ * needed because `</body></html>` is pure ASCII and its byte values are
+ * all < 0x80, so the pattern cannot appear inside multi-byte UTF-8
+ * sequences. This avoids the surrogate-pair and codepoint-corruption
+ * hazards that a text-based search + re-encode would introduce.
+ *
+ * We retain a trailing byte buffer (length `SUFFIX_BYTES.length - 1`)
+ * across chunks so the suffix is detected even when split across chunk
+ * boundaries — the same tail-buffer approach the head injector uses.
+ */
+export interface SuffixState {
+  found: boolean;
+  /** Trailing bytes from the previous chunk for cross-boundary search. */
+  tail: Uint8Array | null;
+  readonly tailLen: number;
+}
+
+export function createSuffixState(): SuffixState {
+  return {
+    found: false,
+    tail: null,
+    tailLen: SUFFIX_BYTES.byteLength - 1,
+  };
+}
+
+/** Find `needle` in `haystack` starting from byte offset `start`. Returns -1 if not found. */
+function indexOfBytes(haystack: Uint8Array, needle: Uint8Array, start = 0): number {
+  const hLen = haystack.byteLength;
+  const nLen = needle.byteLength;
+  if (nLen === 0) return start;
+  const last = hLen - nLen;
+  outer: for (let i = start; i <= last; i++) {
+    for (let j = 0; j < nLen; j++) {
+      if (haystack[i + j] !== needle[j]) continue outer;
+    }
+    return i;
+  }
+  return -1;
+}
+
+/** Concatenate two Uint8Arrays into one. */
+function concatBytes(a: Uint8Array, b: Uint8Array): Uint8Array {
+  const result = new Uint8Array(a.byteLength + b.byteLength);
+  result.set(a, 0);
+  result.set(b, a.byteLength);
+  return result;
+}
+
+/**
+ * Search a chunk for {@link SUFFIX_BYTES}. If found, mark state and return
+ * the surrounding bytes (the caller emits them; the suffix is held for the
+ * end of the stream). A trailing byte tail is retained across calls so the
+ * suffix is detected even when split across chunk boundaries.
+ *
+ * All output is raw byte slices — no text decoding or re-encoding — so
+ * multi-byte UTF-8 codepoints and astral characters are never corrupted.
+ */
+export function processSuffixChunk(state: SuffixState, chunk: Uint8Array): SuffixChunkResult {
+  if (state.found) return { kind: 'passthrough' };
+
+  // Combine tail (if any) with the new chunk for searching.
+  const combined = state.tail ? concatBytes(state.tail, chunk) : chunk;
+
+  const idx = indexOfBytes(combined, SUFFIX_BYTES);
+  if (idx !== -1) {
+    state.found = true;
+    state.tail = null;
+    const beforeEnd = idx;
+    const afterStart = idx + SUFFIX_BYTES.byteLength;
+    return {
+      kind: 'suffixFound',
+      before: beforeEnd > 0 ? combined.slice(0, beforeEnd) : null,
+      after: afterStart < combined.byteLength ? combined.slice(afterStart) : null,
+    };
+  }
+
+  // Suffix not yet seen — emit everything except the last tailLen bytes,
+  // which might be the start of the suffix split across chunks.
+  const safeEnd = Math.max(0, combined.byteLength - state.tailLen);
+  state.tail = combined.slice(safeEnd);
+
+  if (safeEnd === 0) return { kind: 'noSuffix' };
+  return { kind: 'noSuffix', emitSafe: combined.slice(0, safeEnd) };
+}
+
+/**
+ * Final emit for the suffix mover. Returns any leftover tail bytes
+ * followed by the suffix. Always appends the suffix — even if none was
+ * ever observed in the input — so output is well-formed.
+ */
+export function suffixFlush(state: SuffixState): Uint8Array {
+  if (state.tail && state.tail.byteLength > 0) {
+    const result = concatBytes(state.tail, SUFFIX_BYTES);
+    state.tail = null;
+    return result;
+  }
+  state.tail = null;
+  return SUFFIX_BYTES;
+}
+
+// ─── Tag Injector (head / body) ──────────────────────────────────────────────
+
+/**
+ * Pure state for `injectHead`-style transforms: scan for a target tag,
+ * inject `content` once at the matching position, then pass everything
+ * else through. We keep a small trailing string buffer (just under the
+ * length of the target tag) so the tag can be detected even when split
+ * across chunk boundaries.
+ *
+ * The `decoder` is a long-lived `TextDecoder` reused across every chunk
+ * with `{ stream: true }`. This is required so a UTF-8 codepoint that
+ * spans two chunks (e.g. the bytes of `é` arriving in two halves) is
+ * reassembled correctly. A fresh decoder per chunk would drop or
+ * replace the partial codepoint and corrupt streamed HTML.
+ */
+export interface InjectorState {
+  injected: boolean;
+  tail: string;
+  readonly content: string;
+  readonly targetTag: string;
+  readonly position: 'before' | 'after';
+  readonly tailLen: number;
+  readonly decoder: TextDecoder;
+}
+
+export interface InjectorOptions {
+  content: string;
+  targetTag: string;
+  position?: 'before' | 'after';
+}
+
+export function createInjectorState(options: InjectorOptions): InjectorState {
+  return {
+    injected: false,
+    tail: '',
+    content: options.content,
+    targetTag: options.targetTag,
+    position: options.position ?? 'before',
+    tailLen: options.targetTag.length - 1,
+    decoder: new TextDecoder('utf-8'),
+  };
+}
+
+/**
+ * Result of feeding one chunk through {@link processInjectorChunk}.
+ *
+ * - `passthrough` — already injected. Emit `chunk` unchanged.
+ * - `emit` — emit the returned bytes (or nothing if `null`).
+ *
+ * Note: when not yet injected and the target tag is not present, the
+ * helper still emits the *safe head* of the buffer — everything except
+ * the trailing `tailLen` characters that might be the start of the
+ * target tag spilled into the next chunk. The unsafe tail is held in
+ * `state.tail` for the next call.
+ */
+export type InjectorChunkResult =
+  | { kind: 'passthrough' }
+  | { kind: 'emit'; bytes: Uint8Array | null };
+
+export function processInjectorChunk(state: InjectorState, chunk: Uint8Array): InjectorChunkResult {
+  if (state.injected) return { kind: 'passthrough' };
+
+  // Reuse the long-lived decoder so split UTF-8 codepoints are preserved
+  // across chunk boundaries. See comment on InjectorState.decoder.
+  const decoded = state.decoder.decode(chunk, { stream: true });
+  const text = state.tail + decoded;
+  const tagIndex = text.indexOf(state.targetTag);
+
+  if (tagIndex !== -1) {
+    const splitPoint = state.position === 'before' ? tagIndex : tagIndex + state.targetTag.length;
+    const before = text.slice(0, splitPoint);
+    const after = text.slice(splitPoint);
+    state.injected = true;
+    state.tail = '';
+    return { kind: 'emit', bytes: encoder.encode(before + state.content + after) };
+  }
+
+  // Tag not yet seen — flush everything except the last tailLen chars,
+  // which might be the start of the target tag split across chunks.
+  const safeEnd = Math.max(0, text.length - state.tailLen);
+  const safe = safeEnd > 0 ? text.slice(0, safeEnd) : '';
+  state.tail = text.slice(safeEnd);
+  return { kind: 'emit', bytes: safe ? encoder.encode(safe) : null };
+}
+
+/**
+ * Final emit for the injector. If the target tag never appeared, the
+ * leftover tail bytes are returned so they aren't lost. (The injection
+ * itself is silently skipped — same behaviour as before extraction.)
+ */
+export function injectorFlush(state: InjectorState): Uint8Array | null {
+  if (!state.injected && state.tail) return encoder.encode(state.tail);
+  return null;
+}
+
+// ─── Flight Injection Core ───────────────────────────────────────────────────
+
+/**
+ * Pure async core of the RSC flight-injection transform.
+ *
+ * Both wrappers (Web `TransformStream` and Node `Transform`) implement the
+ * exact same protocol around an RSC reader:
+ *
+ *   1. On the first HTML chunk, send `FIRST_CHUNK` and start the pull loop.
+ *   2. The pull loop reads from `rscReader` (timeout-guarded), wraps each
+ *      chunk in a `<script>` via `flightChunkScript`, and pushes it onto a
+ *      pending queue. Between reads, it yields with a single-shot
+ *      `setImmediate` so HTML chunks get event-loop priority — and after
+ *      yielding, it drains the queue (via the wrapper's `onYieldDrain`)
+ *      so flight data reaches the client at shell-flush time.
+ *   3. After every emitted HTML chunk, the wrapper drains the queue.
+ *   4. On `flush()` the wrapper sends `HTML_DONE` and awaits `pullPromise`.
+ *      No polling — just `.then()`.
+ *
+ * The wrappers only differ in:
+ *   - How they push bytes (controller.enqueue vs transform.push).
+ *   - How they signal a terminal error (controller.error vs transform.destroy).
+ *
+ * This class owns the state machine, the pending queue, and the pull
+ * loop itself. The wrapper supplies the drain callback and consumes the
+ * pending queue + machine state to decide what to push.
+ *
+ * **No polling.** See design/02-rendering-pipeline.md §"No Polling".
+ * **Always timeout-guarded reads.** Per design/02 §"Streaming Constraints".
+ */
+export class FlightInjectorCore {
+  readonly machine: Machine<FlightInjectionState, FlightInjectionEvent>;
+  /** Script chunks waiting to be drained between HTML chunks. */
+  readonly pending: Uint8Array[] = [];
+
+  private readonly rscReader: ReadableStreamDefaultReader<Uint8Array>;
+  private readonly timeoutMs: number;
+  private readonly errorSignal?: AbortSignal;
+  private readonly decoder = new TextDecoder('utf-8', { fatal: true });
+  private pullPromise: Promise<void> | null = null;
+  private _aborted = false;
+
+  constructor(
+    rscStream: ReadableStream<Uint8Array>,
+    renderTimeoutMs?: number,
+    errorSignal?: AbortSignal
+  ) {
+    this.rscReader = rscStream.getReader();
+    this.timeoutMs = (renderTimeoutMs ?? 30_000) || 120_000;
+    this.errorSignal = errorSignal;
+    this.machine = createMachine<FlightInjectionState, FlightInjectionEvent>({
+      initial: { phase: 'init' },
+      transitions: flightInjectionTransitions,
+    });
+  }
+
+  /** Tell the machine the first HTML chunk has arrived. Idempotent-ish: caller checks. */
+  notifyFirstChunk(): void {
+    this.machine.send({ type: 'FIRST_CHUNK' });
+  }
+
+  /** Tell the machine HTML output is finished — pull loop will stop yielding. */
+  notifyHtmlDone(): void {
+    this.machine.send({ type: 'HTML_DONE' });
+  }
+
+  get isInit(): boolean {
+    return this.machine.state.phase === 'init';
+  }
+
+  get isPullDone(): boolean {
+    return isPullDone(this.machine.state);
+  }
+
+  get hasPullPromise(): boolean {
+    return this.pullPromise !== null;
+  }
+
+  /**
+   * Get (or start) the pull-loop promise. The wrapper calls this from
+   * `flush()` and chains `.then()` onto it — never `await` without an
+   * existing finish callback (we don't want to block transform()).
+   *
+   * `onYieldDrain` is invoked between RSC reads while HTML is still
+   * streaming. It must drain the pending queue into the output stream.
+   */
+  ensurePullLoop(onYieldDrain: () => void): Promise<void> {
+    if (!this.pullPromise) {
+      this.pullPromise = this.runPullLoop(onYieldDrain);
+    }
+    return this.pullPromise;
+  }
+
+  /**
+   * Get the current terminal error if the machine is in the error phase,
+   * else `null`. Wrappers use this after a drain to decide whether to
+   * `controller.error()` / `transform.destroy()`.
+   */
+  get terminalError(): unknown {
+    return this.machine.state.phase === 'error' ? this.machine.state.error : null;
+  }
+
+  /**
+   * Abort the pull loop and release the RSC reader. Called by stream
+   * wrappers when the HTML pipeline is torn down (Node `destroy()`,
+   * Web `cancel()`). Safe to call multiple times.
+   */
+  abort(reason?: unknown): void {
+    if (this._aborted) return;
+    this._aborted = true;
+    this.rscReader.cancel(reason).catch(() => {});
+    this.pending.length = 0;
+  }
+
+  get isAborted(): boolean {
+    return this._aborted;
+  }
+
+  private async runPullLoop(onYieldDrain: () => void): Promise<void> {
+    // Yield once so the first HTML shell chunk flows through the wrapper
+    // before we start reading RSC data.
+    await new Promise<void>((r) => setImmediate(r));
+
+    try {
+      for (;;) {
+        if (this._aborted) return;
+
+        // Guard each read with a timeout so a permanently hung RSC stream
+        // eventually aborts. See design/02 §"Streaming Constraints".
+        const readPromise = this.rscReader.read();
+        const { done, value } = await withTimeout(
+          readPromise,
+          this.timeoutMs,
+          'RSC stream read timed out'
+        );
+
+        if (done) {
+          this.machine.send({ type: 'PULL_DONE' });
+          return;
+        }
+
+        // If an unhandled RSC error was detected, stop injecting Flight
+        // data. Continue reading (to prevent tee buffer divergence) but
+        // discard the chunks so corrupted data never reaches the client.
+        // See TIM-938.
+        if (this.errorSignal?.aborted) {
+          continue;
+        }
+
+        const decoded = this.decoder.decode(value, { stream: true });
+        this.pending.push(encoder.encode(flightChunkScript(decoded)));
+
+        // Yield between reads so HTML chunks get priority — but only while
+        // HTML is still streaming. Once flush() fires (HTML_DONE), drain
+        // without yielding so remaining RSC data doesn't stall.
+        if (!isHtmlDone(this.machine.state)) {
+          await new Promise<void>((r) => setImmediate(r));
+          // After yielding, drain. Without this, hydration blocks waiting
+          // for flight data stuck in pending[]. Safe because we're between
+          // event-loop ticks (no transform() call mid-execution → no
+          // mid-tag risk; the buffered transform upstream coalesced HTML
+          // chunks to coherent fragments).
+          if (this.pending.length > 0) onYieldDrain();
+        }
+      }
+    } catch (err) {
+      if (err instanceof RenderTimeoutError) {
+        this.rscReader.cancel(err).catch(() => {});
+      }
+      this.machine.send({ type: 'PULL_ERROR', error: err });
+    }
+  }
+}