fix(server): bound client-facing SSE frame retention

src/server/relay.ts

@@ -11,11 +11,16 @@ import {
   type RequestLogContext,
   type RequestLogEntry,
 } from "./request-log";
+import {
+  BoundedSseFrameBuffer,
+  joinSseFrameBytes,
+  MAX_CLIENT_SSE_FRAME_BYTES,
+} from "./sse-frame-buffer";
 
 const nativePassthroughSseResponses = new WeakSet<Response>();
 const eagerRelaySseResponses = new WeakSet<Response>();
 
-export const MAX_INSPECTION_SSE_FRAME_BYTES = 4 * 1024 * 1024;
+export const MAX_INSPECTION_SSE_FRAME_BYTES = MAX_CLIENT_SSE_FRAME_BYTES;
 export const MAX_COMPLETED_OUTPUT_ITEMS = 256;
 export const MAX_COMPLETED_OUTPUT_ITEM_SOURCE_BYTES = 8 * 1024 * 1024;
 export const MAX_TAIL_ERROR_MESSAGE_CHARS = 512;
@@ -104,30 +109,29 @@ export type SseTerminalOutputBoundary = {
 
 /**
  * Frame-aware client output boundary shared by both native Responses relays.
- * It buffers only the current incomplete SSE block, forwards complete blocks
- * through the first Responses terminal, and drops every later block/byte.
+ * It buffers only the current incomplete SSE block under the same hard byte
+ * cap as inspection, forwards complete blocks through the first Responses
+ * terminal, and drops every later block/byte.
  */
 export function createSseTerminalOutputBoundary(): SseTerminalOutputBoundary {
-  let decoder: TextDecoder | null = new TextDecoder();
-  const encoder = new TextEncoder();
-  let buffer = "";
+  const decoder = new TextDecoder();
+  const framer = new BoundedSseFrameBuffer(MAX_INSPECTION_SSE_FRAME_BYTES);
   let terminal = false;
   let done = false;
   let disposed = false;
 
-  const process = (flush: boolean): Uint8Array => {
-    if (disposed || terminal) return new Uint8Array(0);
-    let output = "";
+  const processFrames = (
+    frames: ReturnType<BoundedSseFrameBuffer["feed"]>,
+  ): Uint8Array => {
+    if (disposed || terminal || frames.length === 0) return new Uint8Array(0);
+    const output: Uint8Array[] = [];
     let responsesTerminal = false;
-    for (;;) {
-      const next = nextSseBlock(buffer);
-      if (!next) break;
-      buffer = next.rest;
-      const payload = sseDataPayload(next.block);
-      if (!responsesTerminal) output += next.block + next.delimiter;
+    for (const frame of frames) {
+      const payload = sseDataPayload(decoder.decode(frame.block));
+      if (!responsesTerminal) output.push(frame.block, frame.delimiter);
       if (payload === "[DONE]") {
         done = true;
-        if (responsesTerminal) output += next.block + next.delimiter;
+        if (responsesTerminal) output.push(frame.block, frame.delimiter);
         continue;
       }
       if (!responsesTerminal && payload && terminalStatusFromSsePayload(payload)) {
@@ -136,33 +140,26 @@ export function createSseTerminalOutputBoundary(): SseTerminalOutputBoundary {
     }
     if (responsesTerminal) {
       terminal = true;
-      buffer = "";
-    }
-    if (flush && !terminal && buffer.length > 0) {
-      output += buffer;
-      buffer = "";
+      framer.dispose();
     }
-    return encoder.encode(output);
+    return joinSseFrameBytes(output);
   };
 
   return {
     feed(chunk) {
       if (disposed || terminal) return new Uint8Array(0);
-      buffer += decoder!.decode(chunk, { stream: true });
-      return process(false);
+      return processFrames(framer.feed(chunk));
     },
     finish() {
       if (disposed || terminal) return new Uint8Array(0);
-      buffer += decoder!.decode();
-      return process(true);
+      return framer.finish();
     },
     terminalSeen: () => terminal,
     doneSeen: () => done,
     dispose() {
       if (disposed) return;
       disposed = true;
-      decoder = null;
-      buffer = "";
+      framer.dispose();
     },
   };
 }

src/server/sse-frame-buffer.ts

@@ -0,0 +1,209 @@
+export const MAX_CLIENT_SSE_FRAME_BYTES = 4 * 1024 * 1024;
+
+export class SseFrameTooLargeError extends Error {
+  readonly maxBytes: number;
+
+  constructor(maxBytes: number) {
+    super(`upstream SSE frame exceeded ${maxBytes} bytes`);
+    this.name = "SseFrameTooLargeError";
+    this.maxBytes = maxBytes;
+  }
+}
+
+export type BoundedSseFrame = {
+  block: Uint8Array;
+  delimiter: Uint8Array;
+};
+
+function delimiterLengthAt(
+  index: number,
+  length: number,
+  byteAt: (index: number) => number,
+): number | 0 | undefined {
+  const first = byteAt(index);
+  if (first === 10) {
+    if (index + 1 >= length) return undefined;
+    const second = byteAt(index + 1);
+    if (second === 10) return 2;
+    if (second !== 13) return 0;
+    if (index + 2 >= length) return undefined;
+    return byteAt(index + 2) === 10 ? 3 : 0;
+  }
+  if (first !== 13) return 0;
+  if (index + 1 >= length) return undefined;
+  if (byteAt(index + 1) !== 10) return 0;
+  if (index + 2 >= length) return undefined;
+  const third = byteAt(index + 2);
+  if (third === 10) return 3;
+  if (third !== 13) return 0;
+  if (index + 3 >= length) return undefined;
+  return byteAt(index + 3) === 10 ? 4 : 0;
+}
+
+function copyRange(
+  start: number,
+  end: number,
+  tailLength: number,
+  previousTail: Uint8Array,
+  chunk: Uint8Array,
+): Uint8Array {
+  const out = new Uint8Array(end - start);
+  for (let index = start; index < end; index += 1) {
+    out[index - start] = index < tailLength
+      ? previousTail[index]!
+      : chunk[index - tailLength]!;
+  }
+  return out;
+}
+
+/**
+ * Byte-bounded SSE block framer for client-facing protocol paths.
+ *
+ * The delimiter scanner works on raw bytes, so fragmented UTF-8 cannot change
+ * accounting and a hostile upstream cannot grow an unterminated JS string
+ * without limit. Candidate bytes live in one geometrically grown buffer rather
+ * than one allocation per upstream chunk, bounding both bytes and object count.
+ * Complete blocks are returned without their delimiter; the exact delimiter
+ * bytes are returned separately so callers can relay bytes unchanged.
+ */
+export class BoundedSseFrameBuffer {
+  private readonly maxFrameBytes: number;
+  private delimiterTail: Uint8Array = new Uint8Array(0);
+  private candidate: Uint8Array = new Uint8Array(0);
+  private candidateBytes = 0;
+  private disposed = false;
+
+  constructor(maxFrameBytes = MAX_CLIENT_SSE_FRAME_BYTES) {
+    if (!Number.isSafeInteger(maxFrameBytes) || maxFrameBytes <= 0) {
+      throw new RangeError("maxFrameBytes must be a positive safe integer");
+    }
+    this.maxFrameBytes = maxFrameBytes;
+  }
+
+  private clear(): void {
+    this.delimiterTail = new Uint8Array(0);
+    this.candidate = new Uint8Array(0);
+    this.candidateBytes = 0;
+  }
+
+  private ensureCapacity(requiredBytes: number): void {
+    if (this.candidate.byteLength >= requiredBytes) return;
+    let capacity = this.candidate.byteLength === 0
+      ? Math.min(this.maxFrameBytes, Math.max(requiredBytes, 4096))
+      : this.candidate.byteLength;
+    while (capacity < requiredBytes) {
+      capacity = Math.min(this.maxFrameBytes, Math.max(requiredBytes, capacity * 2));
+    }
+    const grown = new Uint8Array(capacity);
+    if (this.candidateBytes > 0) {
+      grown.set(this.candidate.subarray(0, this.candidateBytes));
+    }
+    this.candidate = grown;
+  }
+
+  private retain(slice: Uint8Array): void {
+    if (slice.byteLength === 0) return;
+    const nextBytes = this.candidateBytes + slice.byteLength;
+    if (nextBytes > this.maxFrameBytes) {
+      this.clear();
+      this.disposed = true;
+      throw new SseFrameTooLargeError(this.maxFrameBytes);
+    }
+    this.ensureCapacity(nextBytes);
+    this.candidate.set(slice, this.candidateBytes);
+    this.candidateBytes = nextBytes;
+  }
+
+  private takeCandidate(): Uint8Array {
+    if (this.candidateBytes === 0) return new Uint8Array(0);
+    const block = this.candidate.slice(0, this.candidateBytes);
+    this.candidate = new Uint8Array(0);
+    this.candidateBytes = 0;
+    return block;
+  }
+
+  feed(chunk: Uint8Array): BoundedSseFrame[] {
+    if (this.disposed) return [];
+    if (chunk.byteLength === 0) return [];
+
+    const frames: BoundedSseFrame[] = [];
+    const previousTail = this.delimiterTail;
+    this.delimiterTail = new Uint8Array(0);
+    const tailLength = previousTail.byteLength;
+    const totalLength = tailLength + chunk.byteLength;
+    const byteAt = (index: number): number => index < tailLength
+      ? previousTail[index]!
+      : chunk[index - tailLength]!;
+    const retainRange = (start: number, end: number): void => {
+      if (end <= start) return;
+      if (start < tailLength) {
+        this.retain(previousTail.subarray(start, Math.min(end, tailLength)));
+      }
+      if (end > tailLength) {
+        this.retain(chunk.subarray(Math.max(0, start - tailLength), end - tailLength));
+      }
+    };
+
+    let index = 0;
+    let retainedThrough = 0;
+    while (index < totalLength) {
+      const delimiterLength = delimiterLengthAt(index, totalLength, byteAt);
+      if (delimiterLength === undefined) break;
+      if (delimiterLength > 0) {
+        retainRange(retainedThrough, index);
+        const block = this.takeCandidate();
+        const delimiter = copyRange(
+          index,
+          index + delimiterLength,
+          tailLength,
+          previousTail,
+          chunk,
+        );
+        frames.push({ block, delimiter });
+        index += delimiterLength;
+        retainedThrough = index;
+        continue;
+      }
+      index += 1;
+    }
+
+    retainRange(retainedThrough, index);
+    if (index < totalLength) {
+      this.delimiterTail = copyRange(index, totalLength, tailLength, previousTail, chunk);
+    }
+    return frames;
+  }
+
+  /** Return the final unterminated block bytes and release all retained state. */
+  finish(): Uint8Array {
+    if (this.disposed) return new Uint8Array(0);
+    try {
+      this.retain(this.delimiterTail);
+      this.delimiterTail = new Uint8Array(0);
+      return this.takeCandidate();
+    } finally {
+      this.clear();
+      this.disposed = true;
+    }
+  }
+
+  dispose(): void {
+    if (this.disposed) return;
+    this.clear();
+    this.disposed = true;
+  }
+}
+
+export function joinSseFrameBytes(parts: readonly Uint8Array[]): Uint8Array {
+  let byteLength = 0;
+  for (const part of parts) byteLength += part.byteLength;
+  if (byteLength === 0) return new Uint8Array(0);
+  if (parts.length === 1 && parts[0]!.byteLength === byteLength) return parts[0]!;
+  const joined = new Uint8Array(byteLength);
+  let offset = 0;
+  for (const part of parts) {
+    joined.set(part, offset);
+    offset += part.byteLength;
+  }
+  return joined;
+}

src/server/ws-bridge.ts

@@ -6,6 +6,7 @@ import { headersForCodexAuthContext } from "../codex/auth-context";
 import type { ResponsesTerminalStatus } from "../bridge";
 import type { DataPlaneAdmission } from "./auth-cors";
 import type { AdmissionLease, AdmissionReservation } from "../lib/admission";
+import { BoundedSseFrameBuffer } from "./sse-frame-buffer";
 
 const OPEN = 1;
 type ResponsesTerminalReporter = (status: ResponsesTerminalStatus) => void;
@@ -163,15 +164,6 @@ function parseSseBlock(block: string): string | null {
   return data.length > 0 ? data.join("\n") : null;
 }
 
-function nextSseBlock(buffer: string): { block: string; rest: string } | null {
-  const match = buffer.match(/\r?\n\r?\n/);
-  if (!match || match.index === undefined) return null;
-  return {
-    block: buffer.slice(0, match.index),
-    rest: buffer.slice(match.index + match[0].length),
-  };
-}
-
 function payloadType(payload: string): string | null {
   try {
     const json = JSON.parse(payload) as { type?: unknown };
@@ -231,7 +223,7 @@ export async function pumpResponsesSseToWebSocket(
   ws.data.cancel = cancel;
 
   const decoder = new TextDecoder();
-  let buffer = "";
+  const framer = new BoundedSseFrameBuffer();
   let terminalSeen = false;
 
   const handlePayload = (payload: string): boolean => {
@@ -266,29 +258,28 @@ export async function pumpResponsesSseToWebSocket(
     while (!terminalSeen) {
       const { done, value } = await reader.read();
       if (done) break;
-      buffer += decoder.decode(value, { stream: true });
-      let next: { block: string; rest: string } | null;
-      while ((next = nextSseBlock(buffer))) {
-        buffer = next.rest;
-        const payload = parseSseBlock(next.block);
+      for (const frame of framer.feed(value)) {
+        const payload = parseSseBlock(decoder.decode(frame.block));
         if (payload && handlePayload(payload)) break;
       }
     }
-    buffer += decoder.decode();
-    if (!terminalSeen && buffer.trim()) {
-      const payload = parseSseBlock(buffer);
+    const tail = framer.finish();
+    if (!terminalSeen && tail.byteLength > 0) {
+      const payload = parseSseBlock(decoder.decode(tail));
       if (payload) handlePayload(payload);
     }
     if (!terminalSeen && isCurrent() && !clientCancelled) {
       reportTerminal("incomplete");
       sendProtocolError(ws, 502, "Upstream stream ended before response terminal event");
     }
   } catch (err) {
+    framer.dispose();
     if (!terminalSeen && isCurrent() && ws.readyState === OPEN) {
       if (!(err instanceof WsSendDroppedError)) reportTerminal("incomplete");
       sendProtocolError(ws, 502, err instanceof Error ? err.message : String(err));
     }
   } finally {
+    framer.dispose();
     if (ws.data.cancel === cancel) ws.data.cancel = undefined;
   }
 }