relay-eager.ts

/**
 * Eager bounded single-reader SSE relay (#314 mitigation, WP2).
 *
 * Replaces the tee()+background-inspection passthrough shape on runtimes where
 * the Bun#32111 async-pull cancel fix is present (src/lib/bun-stream-caps.ts):
 * ONE eager producer loop reads upstream, feeds every chunk through the shared
 * SSE inspector (terminal outcome, quota, request log, context cache), and
 * enqueues it into a byte-bounded client queue. When the queue is full the
 * producer pauses — no unbounded tee branch queue can build up behind a slow
 * client.
 *
 * Honesty caveats (audit M5): full leak relief additionally assumes the
 * runtime carries the Bun#29831 fetch receive-backpressure fix and that Bun's
 * native Response sink pull-paces a JS ReadableStream. Neither is provable in
 * bun:test (a JS reader always paces); both remain "awaiting Windows user
 * verification".
 *
 * #44 cancel semantics: after client cancel the relay keeps reading upstream in
 * DISCARD-DRAIN mode (inspection only) until a terminal is seen or the bounded
 * drain window (ms/bytes) expires — a genuinely reached terminal records as
 * completed/failed, never downgraded to cancel. Only when no terminal arrives
 * within bounds does onClientCancel fire. This bounds today's unbounded tee
 * drain; the tradeoff is that client-cancel log finalization may be delayed by
 * up to the drain window.
 */

import { buildFailedTailPayload } from "./relay";
import {
  nextSseBlock,
  replaceSseDataPayload,
  sseDataPayload,
  type SsePayloadRewrite,
} from "./sse-payload-rewrite";
import type { TranslatorBudget } from "../lib/translator-budget";

export type EagerRelayHooks = {
  /** Feed one upstream chunk through SSE inspection (createSseInspector.feed). */
  inspectChunk: (chunk: Uint8Array) => void;
  /**
   * Optional inline client-facing payload rewrite, framed to complete SSE
   * blocks inside the single reader. This is what lets win32 rewrite traffic
   * (image_gen restore, item-id repair) use this relay instead of the
   * Bun#32111-unsafe tee()+JS-pull chain (#864).
   */
  rewritePayload?: SsePayloadRewrite;
  /** Flush inspection at upstream end (createSseInspector.finish). */
  finishInspection: () => void;
  /** Drop inspector-owned frame/item state during producer teardown. */
  disposeInspection?: () => void;
  /** True once inspection has reported a protocol terminal (inspector.reported). */
  sawTerminal: () => boolean;
  /** Record a synthetic terminal (caller decides incomplete vs failed-502). */
  onSynthetic: (kind: "incomplete" | "failed") => void;
  /** Client cancelled and NO terminal arrived within the drain bounds. */
  onClientCancel: () => void;
  /** Exactly once, after the producer fully stops (unregisterTurn parity). */
  onDone: () => void;
};

export type EagerRelayOptions = {
  /** Optional trailer emitted after rewrite flush (e.g. data: [DONE]). */
  trailer?: string | (() => string | undefined);
  /** Bounded client queue in bytes; producer pauses above it. Default 8 MiB. */
  maxQueueBytes?: number;
  /** Transient-budget owner for the inline-rewrite frame buffer. */
  rewriteBudget?: TranslatorBudget;
  /** Post-cancel discard-drain wall-clock bound. Default 15 000 ms. */
  postCancelDrainMs?: number;
  /** Post-cancel discard-drain byte bound. Default 32 MiB. */
  postCancelDrainBytes?: number;
  /** Injectable clock for tests. */
  now?: () => number;
};

const DEFAULT_MAX_QUEUE_BYTES = 8 * 1024 * 1024;
const DEFAULT_DRAIN_MS = 15_000;
const DEFAULT_DRAIN_BYTES = 32 * 1024 * 1024;

/**
 * Relay `body` to the returned stream with eager bounded reading and inline
 * inspection. `upstream` is aborted on cancel-drain expiry and observed for
 * shutdown teardown (its abort wakes a paused producer and suppresses
 * synthetic terminals — audit M3).
 */
export function relaySseEagerBounded(
  body: ReadableStream<Uint8Array>,
  upstream: AbortController,
  hooks: EagerRelayHooks,
  opts?: EagerRelayOptions,
): ReadableStream<Uint8Array> {
  const maxQueueBytes = opts?.maxQueueBytes ?? DEFAULT_MAX_QUEUE_BYTES;
  const drainMs = opts?.postCancelDrainMs ?? DEFAULT_DRAIN_MS;
  const drainBytes = opts?.postCancelDrainBytes ?? DEFAULT_DRAIN_BYTES;
  const now = opts?.now ?? Date.now;

  const reader = body.getReader();
  const rewrite = hooks.rewritePayload;
  const rewriteDecoder = rewrite ? new TextDecoder() : null;
  const rewriteEncoder = rewrite ? new TextEncoder() : null;
  const rewriteBudget = opts?.rewriteBudget;
  let frameBuffer = "";
  let frameBufferBytes = 0;
  /** Frame complete SSE blocks and rewrite each block's data payload in place. */
  const rewriteOutbound = (value: Uint8Array): Uint8Array => {
    let out = "";
    const fragment = rewriteDecoder!.decode(value, { stream: true });
    if (rewriteBudget) {
      const nextBytes = frameBufferBytes + rewriteEncoder!.encode(fragment).byteLength;
      const reservation = rewriteBudget.reserveTransient(nextBytes, { kind: "live_transient" });
      try {
        frameBuffer += fragment;
        reservation.commitRetained();
        rewriteBudget.releaseRetained(frameBufferBytes, { kind: "live_transient" });
        frameBufferBytes = nextBytes;
      } catch (error) {
        reservation.release();
        throw error;
      }
    } else {
      frameBuffer += fragment;
      frameBufferBytes += value.byteLength;
    }
    for (;;) {
      const next = nextSseBlock(frameBuffer);
      if (!next) break;
      const payload = sseDataPayload(next.block);
      const rewrittenPayload = payload === null ? null : rewrite!(payload);
      frameBuffer = next.rest;
      // Returning null from rewrite drops the entire SSE event.
      if (payload !== null && rewrittenPayload === null) continue;
      // Replace only on an actual change: replaceSseDataPayload collapses
      // multi-data-line events and normalizes newline style even when the
      // payload is identical, which corrupts valid streams.
      const block = payload !== null && rewrittenPayload !== null && rewrittenPayload !== payload
        ? replaceSseDataPayload(next.block, rewrittenPayload)
        : next.block;
      out += block + next.delimiter;
    }
    if (rewriteBudget) {
      const remaining = rewriteEncoder!.encode(frameBuffer).byteLength;
      rewriteBudget.releaseRetained(frameBufferBytes - remaining, { kind: "live_transient" });
      frameBufferBytes = remaining;
    } else {
      frameBufferBytes = rewriteEncoder!.encode(frameBuffer).byteLength;
    }
    return rewriteEncoder!.encode(out);
  };
  /** Flush any trailing partial block at upstream end (rewrite applied, matching the pull relay). */
  const flushRewriteTail = (): Uint8Array => {
    if (!rewrite) return new Uint8Array(0);
    // Decoder-flushed bytes logically follow everything already decoded.
    let tail = frameBuffer + rewriteDecoder!.decode();
    const payload = sseDataPayload(tail);
    if (payload !== null) {
      const rewrittenPayload = rewrite(payload);
      if (rewrittenPayload === null) {
        frameBuffer = "";
        if (rewriteBudget && frameBufferBytes > 0) {
          rewriteBudget.releaseRetained(frameBufferBytes, { kind: "live_transient" });
        }
        frameBufferBytes = 0;
        return rewriteEncoder!.encode("");
      }
      if (rewrittenPayload !== payload) tail = replaceSseDataPayload(tail, rewrittenPayload);
    }
    frameBuffer = "";
    if (rewriteBudget && frameBufferBytes > 0) {
      rewriteBudget.releaseRetained(frameBufferBytes, { kind: "live_transient" });
    }
    frameBufferBytes = 0;
    return rewriteEncoder!.encode(tail);
  };
  let queuedBytes = 0;
  let cancelled = false;
  let done = false;
  // Pause gate: resolved by client pull, client cancel, or upstream abort so a
  // paused producer ALWAYS resumes (audit blocker 2 — no deadlock; onDone and
  // turn unregistration stay reachable, drainAndShutdown never hangs).
  let wake: (() => void) | null = null;
  const wakeUp = () => { const w = wake; wake = null; w?.(); };
  const paused = () => new Promise<void>(resolve => { wake = resolve; });
  upstream.signal.addEventListener("abort", wakeUp, { once: true });

  let controllerRef: ReadableStreamDefaultController<Uint8Array> | null = null;
  let doneFired = false;
  let drainTimer: ReturnType<typeof setTimeout> | null = null;
  const fireDone = () => {
    if (doneFired) return;
    doneFired = true;
    if (drainTimer) { clearTimeout(drainTimer); drainTimer = null; }
    try { hooks.onDone(); } catch { /* lifecycle callbacks must not break teardown */ }
  };
  // A silent upstream after cancel would park the drain loop in reader.read();
  // the wall-clock bound must fire regardless, so cancel arms a hard timer that
  // aborts upstream at the deadline (the abort wakes the read).
  const armDrainTimer = () => {
    if (drainTimer) return;
    drainTimer = setTimeout(() => {
      drainTimer = null;
      upstream.abort(new Error("post-cancel drain window expired"));
    }, drainMs);
    (drainTimer as { unref?: () => void }).unref?.();
  };

  const producer = async () => {
    let syntheticKind: "incomplete" | "failed" | null = null;
    // reader.read() is not intrinsically tied to the upstream AbortController
    // (a fetch body usually rejects on abort, but that coupling is the fetch
    // implementation's, not the stream's), so abort must break a parked read on
    // a silent upstream. Cancelling the reader does that: the pending read
    // settles and the loop observes the abort. This is deliberately NOT a
    // shared `Promise.race([reader.read(), aborted])` companion — racing every
    // read against one never-settled promise retains a reaction per chunk, and
    // that is the exact retention class relay.ts avoids at its own drain.
    const wakeParkedRead = () => { reader.cancel(upstream.signal.reason).catch(() => {}); };
    if (upstream.signal.aborted) wakeParkedRead();
    else upstream.signal.addEventListener("abort", wakeParkedRead, { once: true });
    try {
      for (;;) {
        const result = await reader.read();
        const { done: upstreamDone, value } = result;
        // A chunk that already settled is INSPECTED before abort is honored. A read
        // can settle with a real chunk in the same tick the signal fires (post-cancel
        // drain: the terminal frame arrives, then the drain timer aborts upstream).
        // Checking the signal first discarded that frame, so the terminal was never
        // recorded and the turn was accounted as a plain cancel.
        if (!upstreamDone && value !== undefined) hooks.inspectChunk(value);
        if (upstream.signal.aborted) break;
        if (upstreamDone) {
          hooks.finishInspection();
          if (rewrite) {
            const tail = flushRewriteTail();
            if (tail.byteLength > 0 && !cancelled) {
              queuedBytes += tail.byteLength;
              try { controllerRef?.enqueue(tail); } catch { /* client already gone */ }
            }
            const trailer = typeof opts?.trailer === "function" ? opts.trailer() : opts?.trailer;
            if (trailer && !cancelled) {
              const trailerBytes = new TextEncoder().encode(trailer);
              queuedBytes += trailerBytes.byteLength;
              try { controllerRef?.enqueue(trailerBytes); } catch { /* client already gone */ }
            }
          }
          if (!hooks.sawTerminal() && !cancelled && !upstream.signal.aborted) {
            syntheticKind = "incomplete";
          }
          break;
        }
        if (cancelled) {
          // Discard-drain: inspection only, nothing queued. Stop at terminal
          // or when the bounded window expires.
          drainedBytes += value.byteLength;
          if (hooks.sawTerminal() || drainedBytes >= drainBytes || now() >= drainDeadline) {
            break;
          }
          continue;
        }
        const outbound = rewrite ? rewriteOutbound(value) : value;
        if (outbound.byteLength === 0) continue;
        queuedBytes += outbound.byteLength;
        try {
          controllerRef?.enqueue(outbound);
        } catch {
          // Controller already torn down (client went away without cancel()).
          cancelled = true;
          drainDeadline = now() + drainMs;
          armDrainTimer();
          continue;
        }
        while (queuedBytes > maxQueueBytes && !cancelled && !upstream.signal.aborted) {
          await paused();
        }
      }
    } catch (err) {
      // Upstream read failure. Distinguish genuine mid-stream reset from
      // abort-driven teardown (shutdown/cancel-expiry) — audit M3.
      if (!hooks.sawTerminal() && !cancelled && !upstream.signal.aborted) {
        // Serializing `err` can run user-defined accessors (Error.message
        // getters, toString) that re-entrantly cancel the client or abort the
        // upstream. Build the tail FIRST, then re-check eligibility before
        // committing to the synthetic terminal (adversarial review blocker).
        const tail = new TextEncoder().encode(
          `\n\nevent: response.failed\ndata: ${buildFailedTailPayload(err)}\n\ndata: [DONE]\n\n`,
        );
        if (!hooks.sawTerminal() && !cancelled && !upstream.signal.aborted) {
          syntheticKind = "failed";
          queuedBytes += tail.byteLength;
          try { controllerRef?.enqueue(tail); } catch { /* client already torn down */ }
          try { controllerRef?.close(); } catch { /* client already torn down */ }
        }
      }
    } finally {
      // Release any retained rewrite-buffer bytes on every teardown path
      // (error, cancel, upstream abort) — consumption/EOF release alone
      // leaves them charged.
      if (rewriteBudget && frameBufferBytes > 0) {
        try { rewriteBudget.releaseRetained(frameBufferBytes, { kind: "live_transient" }); } catch { /* teardown must not throw */ }
        frameBufferBytes = 0;
      }
      if (syntheticKind) hooks.onSynthetic(syntheticKind);
      if (cancelled && !hooks.sawTerminal()) {
        hooks.onClientCancel();
      }
      if (cancelled || upstream.signal.aborted || syntheticKind === "failed") {
        upstream.abort();
        reader.cancel().catch(() => {});
      }
      if (!cancelled) {
        try { controllerRef?.close(); } catch { /* already closed/errored */ }
      }
      try { hooks.disposeInspection?.(); } catch { /* inspection teardown must not block lifecycle cleanup */ }
      fireDone();
    }
  };

  let drainedBytes = 0;
  let drainDeadline = Number.POSITIVE_INFINITY;

  return new ReadableStream<Uint8Array>({
    start(controller) {
      controllerRef = controller;
      void producer();
    },
    pull() {
      // The client consumed from the queue; approximate accounting: reset on
      // pull below cap. desiredSize reflects internal queue in chunks, not
      // bytes, so we track bytes ourselves and drain optimistically.
      queuedBytes = 0;
      wakeUp();
    },
    cancel() {
      cancelled = true;
      drainDeadline = now() + drainMs;
      armDrainTimer();
      wakeUp();
    },
  });
}