@iislee/opencodex 2.11.0 → 2.31.0

package/src/server/relay.ts

@@ -6,17 +6,21 @@ import {
   addFinalRequestLog,
   httpStatusForRequestLogTerminal,
   inspectResponseLogJson,
-  inspectResponseLogSsePayload,
   inspectResponseLogSsePayloadParsed,
   recordFirstOutput,
   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;
@@ -95,6 +99,74 @@ export function buildFailedTailPayload(err: unknown): string {
   });
 }
 
+export type SseTerminalOutputBoundary = {
+  feed(chunk: Uint8Array): Uint8Array;
+  finish(): Uint8Array;
+  terminalSeen(): boolean;
+  doneSeen(): boolean;
+  dispose(): void;
+};
+
+/**
+ * Frame-aware client output boundary shared by both native Responses relays.
+ * 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 {
+  const decoder = new TextDecoder();
+  const framer = new BoundedSseFrameBuffer(MAX_INSPECTION_SSE_FRAME_BYTES);
+  let terminal = false;
+  let done = false;
+  let disposed = false;
+
+  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 frame of frames) {
+      const payload = sseDataPayload(decoder.decode(frame.block));
+      const isDone = payload === "[DONE]";
+      // Preserve every frame through the first Responses terminal. A [DONE]
+      // frame is also preserved when it immediately follows that terminal in
+      // the same upstream chunk; every later non-DONE frame is dropped.
+      if (!responsesTerminal || isDone) output.push(frame.block, frame.delimiter);
+      if (isDone) {
+        done = true;
+        continue;
+      }
+      if (!responsesTerminal && payload && terminalStatusFromSsePayload(payload)) {
+        responsesTerminal = true;
+      }
+    }
+    if (responsesTerminal) {
+      terminal = true;
+      framer.dispose();
+    }
+    return joinSseFrameBytes(output);
+  };
+
+  return {
+    feed(chunk) {
+      if (disposed || terminal) return new Uint8Array(0);
+      return processFrames(framer.feed(chunk));
+    },
+    finish() {
+      if (disposed || terminal) return new Uint8Array(0);
+      return framer.finish();
+    },
+    terminalSeen: () => terminal,
+    doneSeen: () => done,
+    dispose() {
+      if (disposed) return;
+      disposed = true;
+      framer.dispose();
+    },
+  };
+}
+
 /**
  * Relay a passthrough SSE body like relayWithAbort, but convert a MID-STREAM failure (upstream
  * reset after headers) into a clean terminal: any partial block is closed off, then a synthetic
@@ -110,18 +182,64 @@ export function relaySseWithFailedTail(
 ): ReadableStream<Uint8Array> {
   const reader = body.getReader();
   const encoder = new TextEncoder();
+  const terminalBoundary = createSseTerminalOutputBoundary();
+  let closed = false;
+  const relayChunk = (
+    controller: ReadableStreamDefaultController<Uint8Array>,
+    value: Uint8Array,
+  ): "terminal" | "output" | "buffered" => {
+    const outbound = terminalBoundary.feed(value);
+    if (outbound.byteLength > 0) controller.enqueue(outbound);
+    if (!terminalBoundary.terminalSeen()) return outbound.byteLength > 0 ? "output" : "buffered";
+
+    // A Responses terminal frame is the protocol boundary. Some compatible
+    // gateways leave the HTTP connection open after response.completed, which
+    // otherwise leaves Codex waiting forever even though the model turn is done.
+    // Preserve through the terminal block only, add the conventional sentinel
+    // when there was no real [DONE] data event, then stop reading upstream.
+    if (!terminalBoundary.doneSeen()) {
+      controller.enqueue(encoder.encode("data: [DONE]\n\n"));
+    }
+    closed = true;
+    controller.close();
+    const reason = "Responses terminal event received";
+    // Notify the tee inspection branch as well. It has already received the
+    // same terminal-bearing upstream chunk, so its bounded drain records the
+    // real terminal and then releases the turn/upstream keep-alive connection.
+    onClientGone?.(reason);
+    reader.cancel(reason).catch(() => {});
+    terminalBoundary.dispose();
+    return "terminal";
+  };
   return new ReadableStream<Uint8Array>({
     async pull(controller) {
       try {
-        const { done, value } = await reader.read();
-        if (done) {
-          controller.close();
-          return;
+        for (;;) {
+          const { done, value } = await reader.read();
+          if (done) {
+            const tail = terminalBoundary.finish();
+            if (tail.byteLength > 0) controller.enqueue(tail);
+            terminalBoundary.dispose();
+            controller.close();
+            return;
+          }
+          const result = relayChunk(controller, value);
+          if (result !== "buffered") return;
         }
-        controller.enqueue(value);
       } catch (err) {
+        let partial: Uint8Array = new Uint8Array(0);
+        try {
+          partial = terminalBoundary.finish();
+        } catch {
+          // A near-cap ambiguous delimiter tail may itself overflow at EOF.
+          // Preserve the original read/framing failure and continue emitting
+          // the bounded failed tail instead of letting cleanup throw again.
+        }
+        terminalBoundary.dispose();
+        if (closed) return;
         const payload = buildFailedTailPayload(err);
         try {
+          if (partial.byteLength > 0) controller.enqueue(partial);
           // Leading blank line terminates a partial SSE block so the failed frame parses cleanly.
           controller.enqueue(encoder.encode(`\n\nevent: response.failed\ndata: ${payload}\n\ndata: [DONE]\n\n`));
           controller.close();
@@ -130,6 +248,7 @@ export function relaySseWithFailedTail(
       }
     },
     cancel(reason) {
+      terminalBoundary.dispose();
       if (onClientGone) onClientGone(reason);
       else upstream.abort(reason);
       reader.cancel(reason).catch(() => {});
@@ -137,11 +256,12 @@ export function relaySseWithFailedTail(
   });
 }
 
-export function nextSseBlock(buffer: string): { block: string; rest: string } | null {
+export function nextSseBlock(buffer: string): { block: string; delimiter: 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),
+    delimiter: match[0],
     rest: buffer.slice(match.index + match[0].length),
   };
 }
@@ -245,53 +365,52 @@ export function trackSseForRequestLog(
   onFirstOutput?: () => void,
 ): ReadableStream<Uint8Array> {
   const reader = body.getReader();
-  const decoder = new TextDecoder();
-  let buffer = "";
   let terminalReported = false;
-  const reportFirstOutput = createFirstOutputReporter(onFirstOutput);
+  let cancelled = false;
 
   const reportTerminal = (status: ResponsesTerminalStatus) => {
     if (terminalReported) return;
     terminalReported = true;
     onTerminal(status);
   };
-
-  const inspectPayload = (payload: string | null) => {
-    if (!payload) return;
-    if (logCtx) inspectResponseLogSsePayload(logCtx, payload);
-    reportFirstOutput.payload(payload);
-    const status = terminalStatusFromSsePayload(payload);
-    if (status) reportTerminal(status);
-  };
-
-  const inspectChunk = (value: Uint8Array) => {
-    buffer += decoder.decode(value, { stream: true });
-    let next: { block: string; rest: string } | null;
-    while ((next = nextSseBlock(buffer))) {
-      buffer = next.rest;
-      inspectPayload(sseDataPayload(next.block));
-    }
-  };
+  // Reuse the byte-bounded inspector so translated responses cannot retain an
+  // unterminated upstream frame or parse the same event once per observer.
+  const inspector = createSseInspector({
+    onTerminal: reportTerminal,
+    logCtx,
+    onFirstOutput,
+  });
 
   return new ReadableStream<Uint8Array>({
     async pull(controller) {
       try {
         const { done, value } = await reader.read();
         if (done) {
-          buffer += decoder.decode();
-          if (buffer.trim()) inspectPayload(sseDataPayload(buffer));
-          if (!terminalReported) reportTerminal("incomplete");
+          if (!cancelled) {
+            inspector.finish();
+            if (!terminalReported) reportTerminal("incomplete");
+          }
+          inspector.dispose();
           controller.close();
           return;
         }
-        inspectChunk(value);
+        inspector.feed(value);
         controller.enqueue(value);
       } catch (err) {
-        if (!terminalReported) reportTerminal("incomplete");
+        // The upstream read rejected: the 200 body died mid-flight. Client
+        // cancellation is the caller's separate 499 path, so a cancel-drained
+        // pending read (cancelled=true) must not carry the truncation marker.
+        if (!cancelled && !terminalReported && logCtx?.activeAttempt) {
+          logCtx.activeAttempt.streamAborted = true;
+        }
+        if (!cancelled && !terminalReported) reportTerminal("incomplete");
+        inspector.dispose();
         try { controller.error(err); } catch { /* already torn down */ }
       }
     },
     cancel(reason) {
+      cancelled = true;
+      inspector.dispose();
       onCancel();
       reader.cancel(reason).catch(() => {});
     },
@@ -402,38 +521,23 @@ export function relaySseWithHeartbeat(
 ): ReadableStream<Uint8Array> | null {
   if (!body) return null;
   const reader = body.getReader();
-  const decoder = new TextDecoder();
   const heartbeat = new TextEncoder().encode(": opencodex keepalive\n\n");
   let timer: ReturnType<typeof setInterval> | undefined;
   let closed = false;
   let clientCancelled = false;
   let terminalReported = false;
-  let buffer = "";
 
   const reportTerminal = (status: ResponsesTerminalStatus) => {
     if (terminalReported || clientCancelled || closed) return;
     terminalReported = true;
     onTerminal?.(status);
   };
-
-  const inspectPayload = (payload: string | null) => {
-    if (!payload) return;
-    const status = terminalStatusFromSsePayload(payload);
-    if (status) reportTerminal(status);
-  };
-
-  const inspectChunk = (value: Uint8Array) => {
-    buffer += decoder.decode(value, { stream: true });
-    let next: { block: string; rest: string } | null;
-    while ((next = nextSseBlock(buffer))) {
-      buffer = next.rest;
-      inspectPayload(sseDataPayload(next.block));
-    }
-  };
+  const inspector = createSseInspector({ onTerminal: reportTerminal });
 
   const cleanup = () => {
     if (closed) return;
     closed = true;
+    inspector.dispose();
     if (timer) clearInterval(timer);
     timer = undefined;
     options?.onDone?.();
@@ -455,14 +559,13 @@ export function relaySseWithHeartbeat(
       try {
         const { done, value } = await reader.read();
         if (done) {
-          buffer += decoder.decode();
-          if (buffer.trim()) inspectPayload(sseDataPayload(buffer));
+          inspector.finish();
           if (!terminalReported && !clientCancelled) reportTerminal("incomplete");
           cleanup();
           controller.close();
           return;
         }
-        inspectChunk(value);
+        inspector.feed(value);
         controller.enqueue(value);
       } catch (err) {
         if (!clientCancelled) reportTerminal("incomplete");
@@ -500,7 +603,20 @@ export type SseInspectorHandlers = {
   onTerminal?: (status: ResponsesTerminalStatus, httpStatusOverride?: number) => void;
   logCtx?: RequestLogContext;
   onCompletedResponse?: (response: { id?: unknown; output?: unknown; status?: unknown }) => void;
+  /**
+   * Every parsed SSE payload, delivered BEFORE any onCompletedResponse derived from that same
+   * payload. A caller that must decide on the whole turn -- not just its terminal snapshot --
+   * needs to see the incremental events, because a stream can announce an item and then close
+   * with an empty `output`.
+   */
+  onParsedPayload?: (payload: unknown) => void;
   onFirstOutput?: () => void;
+  /**
+   * Provider-scoped compatibility: persist the completed snapshot under the
+   * first response id exposed to the client when an upstream changes ids
+   * between `response.created` and `response.completed`.
+   */
+  pinCompletedResponseIdToFirstSeen?: boolean;
 };
 
 type CompletedOutputItem = { item: unknown; sourceBytes: number };
@@ -530,17 +646,6 @@ function delimiterLengthAt(
   return byteAt(index + 3) === 10 ? 4 : 0;
 }
 
-function joinedBytes(slices: readonly Uint8Array[], byteLength: number): Uint8Array {
-  if (slices.length === 1 && slices[0]!.byteLength === byteLength) return slices[0]!;
-  const joined = new Uint8Array(byteLength);
-  let offset = 0;
-  for (const slice of slices) {
-    joined.set(slice, offset);
-    offset += slice.byteLength;
-  }
-  return joined;
-}
-
 /**
  * Per-chunk SSE inspection state machine shared by consumeForInspection,
  * consumeForResponseLogMetadata, and the eager bounded relay (relay-eager.ts).
@@ -561,8 +666,8 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
   let reported = false;
   let sawTerminal = false;
   let disposed = false;
-  let delimiterTail = new Uint8Array(0);
-  let candidateSlices: Uint8Array[] = [];
+  let delimiterTail: Uint8Array = new Uint8Array(0);
+  let candidate: Uint8Array = new Uint8Array(0);
   let candidateBytes = 0;
   let discardingOversizedFrame = false;
   const reportFirstOutput = createFirstOutputReporter(handlers.onFirstOutput);
@@ -572,10 +677,11 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
     : null;
   let aggregateItemBytes = 0;
   let reconstructionTainted = false;
+  let firstResponseId: string | undefined;
 
   const clearFrameState = (): void => {
     delimiterTail = new Uint8Array(0);
-    candidateSlices = [];
+    candidate = new Uint8Array(0);
     candidateBytes = 0;
     discardingOversizedFrame = false;
   };
@@ -592,6 +698,31 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
     decoder = null;
     clearFrameState();
     clearCompletedItems();
+    firstResponseId = undefined;
+  };
+
+  const ensureCandidateCapacity = (requiredBytes: number): void => {
+    if (candidate.byteLength >= requiredBytes) return;
+    let capacity = candidate.byteLength === 0
+      ? Math.min(MAX_INSPECTION_SSE_FRAME_BYTES, Math.max(requiredBytes, 4096))
+      : candidate.byteLength;
+    while (capacity < requiredBytes) {
+      capacity = Math.min(
+        MAX_INSPECTION_SSE_FRAME_BYTES,
+        Math.max(requiredBytes, capacity * 2),
+      );
+    }
+    const grown = new Uint8Array(capacity);
+    if (candidateBytes > 0) grown.set(candidate.subarray(0, candidateBytes));
+    candidate = grown;
+  };
+
+  const takeCandidate = (): Uint8Array => {
+    if (candidateBytes === 0) return new Uint8Array(0);
+    const frame = candidate.slice(0, candidateBytes);
+    candidate = new Uint8Array(0);
+    candidateBytes = 0;
+    return frame;
   };
 
   const retainCandidateSlice = (slice: Uint8Array): void => {
@@ -602,7 +733,7 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
       Math.min(nextBytes, MAX_INSPECTION_SSE_FRAME_BYTES),
     );
     if (nextBytes > MAX_INSPECTION_SSE_FRAME_BYTES) {
-      candidateSlices = [];
+      candidate = new Uint8Array(0);
       candidateBytes = 0;
       discardingOversizedFrame = true;
       inspectionCounters.frameCapOverflows += 1;
@@ -612,10 +743,8 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
       reconstructionTainted = true;
       return;
     }
-    // `subarray()` aliases the upstream chunk's backing buffer. Copy only the
-    // live candidate bytes so a tiny trailing frame cannot pin a multi-MiB
-    // chunk whose preceding frames have already been consumed.
-    candidateSlices.push(slice.slice());
+    ensureCandidateCapacity(nextBytes);
+    candidate.set(slice, candidateBytes);
     candidateBytes = nextBytes;
   };
 
@@ -664,6 +793,11 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
     if (!reported && handlers.logCtx) {
       inspectResponseLogSsePayloadParsed(handlers.logCtx, payload, parsed);
     }
+    // Before any terminal handling: a consumer deciding on the whole turn must observe this
+    // payload even when the terminal snapshot that follows no longer mentions it.
+    if (handlers.onParsedPayload && parsed !== undefined) {
+      try { handlers.onParsedPayload(parsed); } catch { /* inspection must never throw into the pump */ }
+    }
     reportFirstOutput.parsed(parsed);
     const status = terminalStatusFromParsed(parsed);
     if (status) sawTerminal = true;
@@ -686,6 +820,17 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
       const parsedEvent = parsed && typeof parsed === "object" && !Array.isArray(parsed)
         ? parsed as ParsedSseEvent
         : null;
+      const responseRecord = parsedEvent
+        && typeof parsedEvent.response === "object"
+        && parsedEvent.response !== null
+        && !Array.isArray(parsedEvent.response)
+        ? parsedEvent.response as { id?: unknown }
+        : null;
+      if (handlers.pinCompletedResponseIdToFirstSeen
+        && responseRecord
+        && typeof responseRecord.id === "string") {
+        firstResponseId ??= responseRecord.id;
+      }
       const doneItem = parsedEvent?.type === "response.output_item.done" ? parsedEvent.item : undefined;
       if (parsedEvent
         && doneItem !== undefined
@@ -700,6 +845,11 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
 
       let response = completedResponseFromParsedEvent(parsedEvent);
       if (response) {
+        if (handlers.pinCompletedResponseIdToFirstSeen
+          && firstResponseId !== undefined
+          && response.id !== firstResponseId) {
+          response = { ...response, id: firstResponseId };
+        }
         // Authoritative output is a NON-EMPTY ARRAY only. Anything else
         // (missing, null, scalar, object) keeps the historical backfill
         // behavior so a malformed terminal cannot reach rememberResponseState
@@ -735,9 +885,7 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
       return;
     }
     const sourceBytes = candidateBytes;
-    const frame = joinedBytes(candidateSlices, sourceBytes);
-    candidateSlices = [];
-    candidateBytes = 0;
+    const frame = takeCandidate();
     if (reported && !handlers.onCompletedResponse) return;
     const decoded = decoder!.decode(frame);
     scanPayload(sseDataPayload(decoded), sourceBytes);
@@ -794,7 +942,7 @@ export function createSseInspector(handlers: SseInspectorHandlers): SseInspector
         delimiterTail = new Uint8Array(0);
         if (!discardingOversizedFrame && candidateBytes > 0 && !reported) {
           const sourceBytes = candidateBytes;
-          const decoded = decoder!.decode(joinedBytes(candidateSlices, sourceBytes));
+          const decoded = decoder!.decode(takeCandidate());
           scanPayload(decoded.trim() ? sseDataPayload(decoded) : null, sourceBytes);
         }
       } finally {
@@ -815,6 +963,10 @@ export type InspectionConsumerOptions = {
   drainBounds?: Partial<InspectionDrainBounds>;
   upstream?: AbortController;
   now?: () => number;
+  /** Forward provider-scoped response-id pinning to the owned inspector. */
+  pinCompletedResponseIdToFirstSeen?: boolean;
+  /** Observe every parsed SSE payload on the inspection side; see SseInspectorHandlers. */
+  onParsedPayload?: (payload: unknown) => void;
   /** Test seam for proving both public consumers dispose their owned inspector. */
   inspectorFactory?: (handlers: SseInspectorHandlers) => SseInspector;
 };
@@ -970,7 +1122,9 @@ export function consumeForInspection(
     onTerminal,
     logCtx,
     onCompletedResponse,
+    onParsedPayload: options?.onParsedPayload,
     onFirstOutput,
+    pinCompletedResponseIdToFirstSeen: options?.pinCompletedResponseIdToFirstSeen,
   });
   startBoundedInspectionPump({
     ...options,
@@ -993,6 +1147,10 @@ export function consumeForInspection(
         if (logCtx) {
           logCtx.transportPhase = "mid_stream";
           logCtx.terminalSource = "synthetic";
+          // A truncated 200 body must not meter as a success the client never
+          // received; the router's equivalent turn carries 502 + streamAborted
+          // (codex-router #139).
+          if (logCtx.activeAttempt) logCtx.activeAttempt.streamAborted = true;
         }
         onTerminal("failed", 502);
       }
@@ -1015,7 +1173,9 @@ export function consumeForResponseLogMetadata(
   const inspector = (options?.inspectorFactory ?? createSseInspector)({
     logCtx,
     onCompletedResponse,
+    onParsedPayload: options?.onParsedPayload,
     onFirstOutput,
+    pinCompletedResponseIdToFirstSeen: options?.pinCompletedResponseIdToFirstSeen,
   });
   startBoundedInspectionPump({ ...options, reader, inspector, signal, onDone });
 }