fix: graceful shutdown drain + passthrough tee relay (Bun#32111)

Graceful shutdown (#30):
- Track active HTTP/WS turns via Set<AbortController>
- On SIGINT/SIGTERM/api-stop: reject new /v1/responses with 503,
  drain up to shutdownTimeoutMs (default 5s), then abort remaining
- trackStreamLifetime() wraps Response bodies for stream-bounded
  turn lifetime (not handleResponses-bounded)
- WebSocket turns tracked at pump start/end

Bun#32111 workaround (#31):
- Passthrough SSE uses body.tee() + native relay — branch[0]
  returned directly to Response (Bun native, no JS Sink.write),
  branch[1] consumed in background for terminal-outcome inspection
- Avoids the async-pull segfault on Windows where
  controller.enqueue() hits a freed buffer after client abort

Closes #30, closes #31.

Author: lidge-jun

src/cli.ts

@@ -6,7 +6,7 @@ import { restoreLegacyOpenaiHistory } from "./codex-history-provider";
 import { codexAutoStartEnabled, getConfigDir, loadConfig, readPid, removePid, saveConfig, writePid } from "./config";
 import { findAvailablePort } from "./ports";
 import { serviceCommand, stopServiceIfInstalled, uninstallServiceIfInstalled } from "./service";
-import { startServer } from "./server";
+import { drainAndShutdown, startServer } from "./server";
 import { maybeShowStarPrompt } from "./star-prompt";
 
 const args = process.argv.slice(2);
@@ -200,14 +200,15 @@ async function handleStart(options: { block?: boolean } = {}) {
   const server = startServer(port);
   writePid(process.pid);
 
+  const config = loadConfig();
   const shutdown = () => {
     console.log("\n🛑 Shutting down opencodex proxy...");
-    server.stop(true);
-    removePid();
-    // Under the service (OCX_SERVICE), a restart re-injects on start — don't churn Codex config.
-    // `ocx service stop/uninstall` restore explicitly.
-    if (!process.env.OCX_SERVICE) { try { restoreNativeCodex(); } catch { /* best-effort restore */ } }
-    process.exit(0);
+    void (async () => {
+      await drainAndShutdown(server, config.shutdownTimeoutMs ?? 5000);
+      removePid();
+      if (!process.env.OCX_SERVICE) { try { restoreNativeCodex(); } catch { /* best-effort restore */ } }
+      process.exit(0);
+    })();
   };
 
   process.on("SIGINT", shutdown);

src/server.ts

@@ -57,6 +57,66 @@ import {
 } from "./codex-routing";
 import { registerCodexWebSocket, unregisterCodexWebSocket } from "./codex-websocket-registry";
 
+// ---------------------------------------------------------------------------
+// Active turn tracking + graceful shutdown drain
+// ---------------------------------------------------------------------------
+
+const activeTurns = new Set<AbortController>();
+let draining = false;
+
+export function registerTurn(ac: AbortController): void { activeTurns.add(ac); }
+export function unregisterTurn(ac: AbortController): void { activeTurns.delete(ac); }
+export function isDraining(): boolean { return draining; }
+export function getActiveTurnCount(): number { return activeTurns.size; }
+
+export function trackStreamLifetime(
+  body: ReadableStream<Uint8Array>,
+  ac: AbortController,
+): ReadableStream<Uint8Array> {
+  registerTurn(ac);
+  const reader = body.getReader();
+  return new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      try {
+        const { done, value } = await reader.read();
+        if (done) { unregisterTurn(ac); controller.close(); return; }
+        controller.enqueue(value);
+      } catch (err) {
+        unregisterTurn(ac);
+        try { controller.error(err); } catch { /* already closed */ }
+      }
+    },
+    cancel(reason) {
+      unregisterTurn(ac);
+      ac.abort(reason);
+      reader.cancel(reason).catch(() => {});
+    },
+  });
+}
+
+let _serverRef: ReturnType<typeof Bun.serve> | undefined;
+
+export async function drainAndShutdown(
+  server: ReturnType<typeof Bun.serve> | undefined,
+  timeoutMs: number,
+): Promise<void> {
+  const s = server ?? _serverRef;
+  draining = true;
+  const deadline = Date.now() + timeoutMs;
+  while (activeTurns.size > 0 && Date.now() < deadline) {
+    await Bun.sleep(100);
+  }
+  if (activeTurns.size > 0) {
+    console.log(`⚠️  Aborting ${activeTurns.size} in-flight turn(s) after ${timeoutMs}ms deadline`);
+    for (const ac of activeTurns) {
+      ac.abort(new Error("server shutdown"));
+    }
+    activeTurns.clear();
+  }
+  s?.stop(true);
+  draining = false;
+}
+
 // Single source of truth = package.json (../ from src/), so /healthz + the GUI badge match the
 // installed npm version instead of a stale hardcode.
 const VERSION = (() => {
@@ -316,15 +376,28 @@ async function handleResponses(
       }
     }
 
-    const body = isEventStream
-      ? relaySseWithHeartbeat(
-          upstreamResponse.body,
-          upstream,
-          15_000,
-          terminalBodyWillRecord && recordTerminalOutcomes ? terminalRecorder : undefined,
-        )
-      : relayWithAbort(upstreamResponse.body, upstream);
-    return new Response(body, {
+    // Bun#32111 workaround: passthrough SSE uses tee()+native relay to avoid the
+    // async-pull segfault on Windows. Branch[0] goes directly to the Response (Bun
+    // native relay, never enters JS Sink.write); branch[1] is consumed in the
+    // background for terminal-outcome/quota inspection only.
+    if (isEventStream && upstreamResponse.body) {
+      const [nativeBody, inspectBody] = upstreamResponse.body.tee();
+      const turnAc = new AbortController();
+      const trackedNative = trackStreamLifetime(nativeBody, turnAc);
+      if (terminalBodyWillRecord && recordTerminalOutcomes && terminalRecorder) {
+        consumeForInspection(inspectBody, terminalRecorder);
+      } else {
+        inspectBody.cancel().catch(() => {});
+      }
+      return new Response(trackedNative, {
+        status: upstreamResponse.status,
+        headers,
+      });
+    }
+    const body = relayWithAbort(upstreamResponse.body, upstream);
+    const turnAc = new AbortController();
+    const tracked = body ? trackStreamLifetime(body, turnAc) : null;
+    return new Response(tracked, {
       status: upstreamResponse.status,
       headers,
     });
@@ -391,7 +464,9 @@ async function handleResponses(
         hideThinkingSummary: parsed.options.hideThinkingSummary,
       },
     );
-    return new Response(sseStream, {
+    const bridgeTurnAc = new AbortController();
+    const trackedSse = trackStreamLifetime(sseStream, bridgeTurnAc);
+    return new Response(trackedSse, {
       headers: { "Content-Type": "text/event-stream", "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no" },
     });
   }
@@ -608,6 +683,54 @@ export function relaySseWithHeartbeat(
   });
 }
 
+/**
+ * Background-consume an SSE stream purely for terminal-outcome inspection (quota tracking).
+ * Does not produce output; safe to ignore errors (the client-facing stream is separate).
+ */
+function consumeForInspection(
+  body: ReadableStream<Uint8Array>,
+  onTerminal: (status: ResponsesTerminalStatus) => void,
+): void {
+  const reader = body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+  let reported = false;
+  const pump = async () => {
+    try {
+      for (;;) {
+        const { done, value } = await reader.read();
+        if (done) {
+          buffer += decoder.decode();
+          if (buffer.trim() && !reported) {
+            const payload = sseDataPayload(buffer);
+            if (payload) {
+              const status = terminalStatusFromSsePayload(payload);
+              if (status) { reported = true; onTerminal(status); }
+            }
+          }
+          if (!reported) onTerminal("incomplete");
+          return;
+        }
+        buffer += decoder.decode(value, { stream: true });
+        let next: { block: string; rest: string } | null;
+        while ((next = nextSseBlock(buffer))) {
+          buffer = next.rest;
+          if (!reported) {
+            const payload = sseDataPayload(next.block);
+            if (payload) {
+              const status = terminalStatusFromSsePayload(payload);
+              if (status) { reported = true; onTerminal(status); }
+            }
+          }
+        }
+      }
+    } catch {
+      if (!reported) onTerminal("incomplete");
+    }
+  };
+  pump();
+}
+
 /**
  * Bun's fetch auto-decompresses the response body but leaves the upstream `content-encoding`
  * (and a now-stale `content-length`) on `response.headers`. Relaying those with the already-decoded
@@ -967,7 +1090,10 @@ async function handleManagementAPI(req: Request, url: URL, config: OcxConfig): P
     const { stopServiceIfInstalled } = await import("./service");
     stopServiceIfInstalled();
     restoreNativeCodex();
-    setTimeout(() => process.exit(0), 200);
+    setTimeout(async () => {
+      await drainAndShutdown(undefined, config.shutdownTimeoutMs ?? 5000);
+      process.exit(0);
+    }, 200);
     return jsonResponse({ success: true, message: "Proxy stopping, native Codex restored." });
   }
 
@@ -1026,6 +1152,9 @@ export function startServer(port?: number) {
       // Responses WebSocket (phase 120.2). Codex upgrades the same /v1/responses path; auth is
       // handshake-time only, so capture inbound headers and thread them into the pipeline.
       if (url.pathname === "/v1/responses" && req.headers.get("upgrade")?.toLowerCase() === "websocket") {
+        if (draining) {
+          return new Response("Service shutting down", { status: 503, headers: { ...corsHeaders(), "Retry-After": "5" } });
+        }
         const apiAuthError = requireApiAuth(req, config, "data-plane");
         if (apiAuthError) return apiAuthError;
         if (!isLocalOrigin(req)) {
@@ -1090,6 +1219,9 @@ export function startServer(port?: number) {
       }
 
       if (url.pathname === "/v1/responses" && req.method === "POST") {
+        if (draining) {
+          return new Response("Service shutting down", { status: 503, headers: { ...corsHeaders(), "Retry-After": "5" } });
+        }
         const apiAuthError = requireApiAuth(req, config, "data-plane");
         if (apiAuthError) return apiAuthError;
         if (!isLocalOrigin(req)) {
@@ -1151,6 +1283,8 @@ export function startServer(port?: number) {
 
         const payload: Record<string, unknown> = { ...frame };
         delete payload.type;
+        const wsTurnAc = new AbortController();
+        registerTurn(wsTurnAc);
         void (async () => {
           const logCtx = { model: "unknown", provider: "unknown" };
           const fwd = new Headers({ "content-type": "application/json" });
@@ -1192,6 +1326,7 @@ export function startServer(port?: number) {
               /* socket already gone or send dropped */
             }
           } finally {
+            unregisterTurn(wsTurnAc);
             if (ws.data.cancel === cancelTurn) ws.data.cancel = undefined;
           }
         })();
@@ -1203,6 +1338,8 @@ export function startServer(port?: number) {
     },
   });
 
+  _serverRef = server;
+
   console.log(`🚀 opencodex proxy running on http://localhost:${listenPort}`);
   console.log(`   POST /v1/responses → provider translation`);
   console.log(`   GET  /healthz      → health check`);

src/types.ts

@@ -183,6 +183,8 @@ export interface OcxConfig {
   stallTimeoutSec?: number;
   /** Connect timeout (ms) for upstream fetch — covers DNS, TCP, TLS, and response header. Default 30000. */
   connectTimeoutMs?: number;
+  /** Graceful shutdown drain timeout (ms). Active turns are aborted after this deadline. Default 5000. */
+  shutdownTimeoutMs?: number;
   /** Advertise supports_websockets so Codex opens the WS endpoint. Default false; set true to opt in. */
   websockets?: boolean;
   /** Auto-start/sync the proxy from the Codex shim before launching Codex. Default true. */

tests/shutdown-drain.test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, test } from "bun:test";
+import {
+  registerTurn,
+  unregisterTurn,
+  isDraining,
+  getActiveTurnCount,
+  trackStreamLifetime,
+} from "../src/server";
+
+describe("active turn tracking", () => {
+  test("register/unregister tracks active turns", () => {
+    const ac1 = new AbortController();
+    const ac2 = new AbortController();
+    const before = getActiveTurnCount();
+    registerTurn(ac1);
+    registerTurn(ac2);
+    expect(getActiveTurnCount()).toBe(before + 2);
+    unregisterTurn(ac1);
+    expect(getActiveTurnCount()).toBe(before + 1);
+    unregisterTurn(ac2);
+    expect(getActiveTurnCount()).toBe(before);
+  });
+
+  test("isDraining() is false by default", () => {
+    expect(isDraining()).toBe(false);
+  });
+});
+
+describe("trackStreamLifetime", () => {
+  test("registers on start and unregisters on stream close", async () => {
+    const enc = new TextEncoder();
+    const chunks = [enc.encode("hello"), enc.encode("world")];
+    let i = 0;
+    const source = new ReadableStream<Uint8Array>({
+      pull(controller) {
+        if (i < chunks.length) controller.enqueue(chunks[i++]);
+        else controller.close();
+      },
+    });
+    const ac = new AbortController();
+    const before = getActiveTurnCount();
+    const tracked = trackStreamLifetime(source, ac);
+    expect(getActiveTurnCount()).toBe(before + 1);
+
+    const reader = tracked.getReader();
+    const dec = new TextDecoder();
+    let text = "";
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      text += dec.decode(value, { stream: true });
+    }
+    expect(text).toBe("helloworld");
+    expect(getActiveTurnCount()).toBe(before);
+  });
+
+  test("unregisters on cancel", async () => {
+    const source = new ReadableStream<Uint8Array>({
+      pull() {
+        // never closes — simulate long stream
+      },
+    });
+    const ac = new AbortController();
+    const before = getActiveTurnCount();
+    const tracked = trackStreamLifetime(source, ac);
+    expect(getActiveTurnCount()).toBe(before + 1);
+
+    await tracked.cancel("test cancel");
+    expect(getActiveTurnCount()).toBe(before);
+    expect(ac.signal.aborted).toBe(true);
+  });
+});