fix: abort websocket turns before upstream headers

Author: lidge-jun

src/server.ts

@@ -100,7 +100,7 @@ async function handleResponses(
   req: Request,
   config: OcxConfig,
   logCtx: { model: string; provider: string },
-  options: { forceEmptyResponseId?: boolean } = {},
+  options: { forceEmptyResponseId?: boolean; abortSignal?: AbortSignal } = {},
 ): Promise<Response> {
   let body: unknown;
   try {
@@ -161,6 +161,7 @@ async function handleResponses(
     // consumer's cancel to a signalled fetch, so we pass the signal and relay through relayWithAbort,
     // whose cancel() aborts the upstream — preventing leaked connections (RC2, passthrough path).
     const upstream = new AbortController();
+    linkAbortSignal(upstream, options.abortSignal);
     let upstreamResponse: Response;
     try {
       upstreamResponse = await fetch(request.url, {
@@ -199,6 +200,7 @@ async function handleResponses(
   // Abort the upstream fetch if the client (Codex) disconnects mid-stream, so a cancelled turn does
   // not leak the upstream connection or keep draining tokens. The bridge's cancel() fires upstream.abort() (RC2).
   const upstream = new AbortController();
+  linkAbortSignal(upstream, options.abortSignal);
   let upstreamResponse: Response;
   try {
     upstreamResponse = await fetch(request.url, {
@@ -259,6 +261,15 @@ async function handleResponses(
   return formatErrorResponse(500, "internal_error", "Non-streaming not supported by this adapter");
 }
 
+export function linkAbortSignal(upstream: AbortController, signal?: AbortSignal): void {
+  if (!signal) return;
+  if (signal.aborted) {
+    upstream.abort(signal.reason);
+    return;
+  }
+  signal.addEventListener("abort", () => upstream.abort(signal.reason), { once: true });
+}
+
 const requestLog: { timestamp: number; model: string; provider: string; status: number; durationMs: number }[] = [];
 const MAX_LOG_SIZE = 200;
 
@@ -620,12 +631,18 @@ export function startServer(port?: number) {
         const turnId = (ws.data.turnId ?? 0) + 1;
         ws.data.turnId = turnId;
         const isCurrent = () => ws.data.turnId === turnId;
+        const turnAbort = new AbortController();
+        const cancelTurn = () => {
+          turnAbort.abort("websocket turn superseded or closed");
+        };
+        ws.data.cancel = cancelTurn;
 
         if (frame.generate === false) {
           for (const payload of buildWarmupCompletionFrames(frame)) {
             if (!isCurrent()) return;
             sendTextFrame(ws, payload);
           }
+          if (ws.data.cancel === cancelTurn) ws.data.cancel = undefined;
           return;
         }
 
@@ -641,7 +658,10 @@ export function startServer(port?: number) {
             body: JSON.stringify({ ...payload, stream: true }),
           });
           try {
-            const response = await handleResponses(req, config, logCtx, { forceEmptyResponseId: true });
+            const response = await handleResponses(req, config, logCtx, {
+              forceEmptyResponseId: true,
+              abortSignal: turnAbort.signal,
+            });
             await sendResponseToWebSocket(ws, response, isCurrent);
           } catch (err) {
             if (!isCurrent()) return;
@@ -653,6 +673,8 @@ export function startServer(port?: number) {
             } catch {
               /* socket already gone or send dropped */
             }
+          } finally {
+            if (ws.data.cancel === cancelTurn) ws.data.cancel = undefined;
           }
         })();
       },

src/ws-bridge.ts

@@ -38,7 +38,12 @@ export function safeResponseHeaders(headers: Headers): Record<string, string> {
   const out: Record<string, string> = {};
   for (const [name, value] of headers) {
     const lower = name.toLowerCase();
-    if (SAFE_RESPONSE_HEADER_EXACT.has(lower) || lower.startsWith("x-ratelimit-")) {
+    if (
+      SAFE_RESPONSE_HEADER_EXACT.has(lower) ||
+      lower.startsWith("x-ratelimit-") ||
+      /^x-codex(?:-[a-z0-9-]+)?-(primary|secondary)-(used-percent|window-minutes|reset-at)$/.test(lower) ||
+      /^x-codex(?:-[a-z0-9-]+)?-limit-name$/.test(lower)
+    ) {
       out[lower] = value;
     }
   }
@@ -240,7 +245,10 @@ export async function sendResponseToWebSocket(
   response: Response,
   isCurrent: () => boolean,
 ): Promise<void> {
-  if (!isCurrent()) return;
+  if (!isCurrent()) {
+    await response.body?.cancel().catch(() => {});
+    return;
+  }
 
   if (!response.ok) {
     const text = await response.text().catch(() => "");
@@ -273,7 +281,10 @@ export async function sendResponseToWebSocket(
   }
 
   const { prefix, stream } = await readBoundedPrefix(response.body);
-  if (!isCurrent()) return;
+  if (!isCurrent()) {
+    await stream.cancel().catch(() => {});
+    return;
+  }
   if (looksLikeSse(prefix)) {
     await pumpResponsesSseToWebSocket(ws, stream, { isCurrent });
     return;

tests/passthrough-abort.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test";
-import { relayWithAbort } from "../src/server";
+import { linkAbortSignal, relayWithAbort } from "../src/server";
 
 function streamFromChunks(chunks: Uint8Array[]): ReadableStream<Uint8Array> {
   let i = 0;
@@ -45,4 +45,14 @@ describe("passthrough relayWithAbort (RC2, passthrough path)", () => {
     expect(relayWithAbort(null, ac)).toBeNull();
     expect(ac.signal.aborted).toBe(false);
   });
+
+  test("turn-level abort signal aborts the upstream fetch before headers arrive", () => {
+    const upstream = new AbortController();
+    const turn = new AbortController();
+    linkAbortSignal(upstream, turn.signal);
+    expect(upstream.signal.aborted).toBe(false);
+    turn.abort("replacement turn");
+    expect(upstream.signal.aborted).toBe(true);
+    expect(upstream.signal.reason).toBe("replacement turn");
+  });
 });

tests/ws-endpoint.test.ts

@@ -193,10 +193,18 @@ describe("WS endpoint re-framer (120/132)", () => {
       "set-cookie": "secret=1",
       "x-ratelimit-remaining": "4",
       "x-codex-turn-state": "state",
+      "x-codex-primary-used-percent": "100.0",
+      "x-codex-primary-window-minutes": "15",
+      "x-codex-secondary-primary-reset-at": "1781928000",
+      "x-codex-secondary-limit-name": "Secondary",
     }));
     expect(outbound).toEqual({
       "retry-after": "2",
       "x-codex-turn-state": "state",
+      "x-codex-primary-used-percent": "100.0",
+      "x-codex-primary-window-minutes": "15",
+      "x-codex-secondary-primary-reset-at": "1781928000",
+      "x-codex-secondary-limit-name": "Secondary",
       "x-ratelimit-remaining": "4",
     });
   });
@@ -223,6 +231,24 @@ describe("WS endpoint re-framer (120/132)", () => {
     expect(JSON.parse(sent[0]).type).toBe("response.completed");
   });
 
+  test("cancels a stale successful response body before pumping", async () => {
+    const { ws, sent } = mockWs();
+    let cancelled = false;
+    const body = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(new TextEncoder().encode(
+          'data: {"type":"response.completed","response":{"id":"stale"}}\n\n',
+        ));
+      },
+      cancel() { cancelled = true; },
+    });
+    await sendResponseToWebSocket(ws, new Response(body, {
+      headers: { "content-type": "text/event-stream" },
+    }), () => false);
+    expect(sent).toEqual([]);
+    expect(cancelled).toBe(true);
+  });
+
   test("sniffs mislabelled SSE responses", async () => {
     const { ws, sent } = mockWs();
     await sendResponseToWebSocket(ws, new Response(