fix(sse): grace period before finalizing a client disconnect as 499 (#9653)

A client that closes its connection right after reading a fully-completed
SSE stream can race OmniRoute's own completion bookkeeping: the bytes
already reached the client, but the transform stream's own completion
callback (onStreamComplete, which flips streamCompletionRecorded) hasn't
finished bubbling up when the disconnect handler fires, so the request gets
persisted as a false 499 with zero token usage even though it delivered its
full response.

Confirmed live on real traffic before this fix: a request whose server log
showed "disconnect: request_signal_aborted" at 18236ms was persisted with
status 200 and full token usage (82814/1292) once the grace period let the
real completion win the race, matching what the client actually received.

createClientDisconnectGraceHandler (new leaf in
streamFailureFinalization.ts) polls isStreamCompletionRecorded() for up to
STREAM_DISCONNECT_GRACE_PERIOD_MS (default 10s, env-configurable, 0
disables) before finalizing as a failure. If a real completion lands within
the window, handleStreamFailure's own guard is a no-op and the genuine 200
stands.

Covered by tests/unit/stream-disconnect-grace-period-9653.test.ts (fake-timer
driven: already-recorded completion short-circuits, disabled-grace-period
finalizes immediately, a completion landing mid-window skips finalize
entirely, and no completion ever landing finalizes once the deadline
passes).

(cherry picked from commit 5d0fe28c42)

open-sse/config/constants.ts

@@ -18,6 +18,15 @@ export const FETCH_TIMEOUT_MS = upstreamTimeouts.fetchTimeoutMs;
 // idle for this duration. Override with STREAM_IDLE_TIMEOUT_MS env var.
 export const STREAM_IDLE_TIMEOUT_MS = upstreamTimeouts.streamIdleTimeoutMs;
 
+// Grace period (ms) a client-disconnect finalization waits for the stream's own
+// completion bookkeeping to land before persisting a 499. See #9653 — a client
+// that closes right after reading a fully-completed SSE stream can otherwise
+// race OmniRoute's own completion callback, resulting in a false 499 with zero
+// token usage for a request that actually delivered its full response. Set
+// STREAM_DISCONNECT_GRACE_PERIOD_MS=0 to disable and restore the old
+// immediate-fail behavior.
+export const STREAM_DISCONNECT_GRACE_PERIOD_MS = upstreamTimeouts.streamDisconnectGracePeriodMs;
+
 // Timeout for the first non-ping SSE event. Inherits REQUEST_TIMEOUT_MS when
 // set, unless STREAM_READINESS_TIMEOUT_MS is specified directly. This must stay
 // conservative for large prompts and slow first-byte reasoning providers.

open-sse/handlers/chatCore.ts

@@ -171,6 +171,7 @@ import {
   ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE,
   STREAM_RECOVERY,
   DEFAULT_MAX_TOKENS,
+  STREAM_DISCONNECT_GRACE_PERIOD_MS,
 } from "../config/constants.ts";
 import { createRecoverableStream, makeContinuationBody } from "../services/streamRecovery.ts";
 import {
@@ -4917,13 +4918,20 @@ export async function handleChatCore({
   });
   const handleStreamFailure = streamFailureFinalizers.handleStreamFailure;
   onPipelineStreamError = streamFailureFinalizers.onPipelineStreamError;
-  onClientDisconnectFinalize = (event) =>
-    handleStreamFailure({
-      status: 499,
-      message: `Client disconnected: ${event.reason}`,
-      code: "client_disconnected",
-      type: "client_disconnected",
-    });
+  // #9653: gives a genuine, race-delayed completion a chance to land (see
+  // createClientDisconnectGraceHandler's doc comment) before persisting a false
+  // 499/0-tokens for a request that actually delivered its full response.
+  onClientDisconnectFinalize = streamFailure.createClientDisconnectGraceHandler({
+    isStreamCompletionRecorded: () => streamCompletionRecorded,
+    gracePeriodMs: STREAM_DISCONNECT_GRACE_PERIOD_MS,
+    finalize: (event) =>
+      handleStreamFailure({
+        status: 499,
+        message: `Client disconnected: ${event.reason}`,
+        code: "client_disconnected",
+        type: "client_disconnected",
+      }),
+  });
 
   // For providers using Responses API format, translate stream back to openai (Chat Completions) format
   // UNLESS client is Droid CLI which expects openai-responses format back

open-sse/utils/streamFailureFinalization.ts

@@ -29,6 +29,59 @@ export type PipelineStreamErrorHandler = (event: {
   statusCode: number;
 }) => boolean;
 
+export type ClientDisconnectEvent = { reason: string; duration: number };
+
+/**
+ * #9653: a client that closes its connection right after reading a fully-completed
+ * SSE stream can race the stream's own completion bookkeeping — the bytes already
+ * reached the client, but the transform stream's completion callback (which flips
+ * `isStreamCompletionRecorded()` to true) hasn't finished bubbling up yet when the
+ * disconnect handler fires. Persisting immediately in that case records a false
+ * 499 with zero token usage for a request that actually delivered its full response.
+ *
+ * This wraps a disconnect finalizer with a grace period: instead of finalizing
+ * immediately, poll `isStreamCompletionRecorded()` until it flips true (a real
+ * completion landed — nothing more to do) or the deadline passes (genuinely gone —
+ * finalize as a 499 same as before). Pass `gracePeriodMs <= 0` to disable and
+ * finalize immediately, matching the pre-#9653 behavior.
+ */
+export function createClientDisconnectGraceHandler({
+  isStreamCompletionRecorded,
+  gracePeriodMs,
+  finalize,
+  pollIntervalMs = 250,
+  setTimeoutFn = setTimeout,
+}: {
+  isStreamCompletionRecorded: () => boolean;
+  gracePeriodMs: number;
+  finalize: (event: ClientDisconnectEvent) => unknown;
+  pollIntervalMs?: number;
+  setTimeoutFn?: (callback: () => void, ms: number) => unknown;
+}): (event: ClientDisconnectEvent) => boolean {
+  return (event) => {
+    if (isStreamCompletionRecorded()) return true;
+    if (gracePeriodMs <= 0) {
+      finalize(event);
+      return true;
+    }
+
+    const deadline = Date.now() + gracePeriodMs;
+    const poll = () => {
+      if (isStreamCompletionRecorded()) return;
+      if (Date.now() >= deadline) {
+        finalize(event);
+        return;
+      }
+      setTimeoutFn(poll, pollIntervalMs);
+    };
+    setTimeoutFn(poll, pollIntervalMs);
+
+    // Claim "handled" immediately so the caller's own immediate-finalize fallback
+    // doesn't fire while the grace-period poll is still pending.
+    return true;
+  };
+}
+
 export function finalizeStreamRequestLog({
   pendingRequestId,
   model,
@@ -107,9 +160,7 @@ export function createStreamFailureFinalizers({
     const message = failure.message || "Upstream stream error";
     const code = failure.code || failure.type || String(status);
     const classification =
-      failure.code || failure.type
-        ? { code: failure.code, type: failure.type }
-        : undefined;
+      failure.code || failure.type ? { code: failure.code, type: failure.type } : undefined;
 
     if (!isFailureCompletionRecorded()) {
       const errorBody = buildErrorBody(status, message, undefined, classification);

src/shared/utils/runtimeTimeouts.ts

@@ -27,6 +27,14 @@ export const DEFAULT_API_BRIDGE_SERVER_SOCKET_TIMEOUT_MS = 0;
 // idle-pool window, mirroring the API bridge server's pattern.
 export const DEFAULT_MAIN_SERVER_KEEPALIVE_TIMEOUT_MS = 65_000;
 export const DEFAULT_MAIN_SERVER_HEADERS_TIMEOUT_MS = 66_000;
+// A client that closes its connection right after reading a fully-completed
+// SSE stream can race OmniRoute's own completion bookkeeping (#9653): the
+// bytes already reached the client, but the disconnect handler can fire
+// before the stream's own completion callback finishes recording it,
+// persisting a false 499 with zero token usage. Before committing to that
+// failure, wait this long for the real completion to land. Set to 0 to
+// disable and restore the old immediate-fail behavior.
+export const DEFAULT_STREAM_DISCONNECT_GRACE_PERIOD_MS = 10_000;
 
 function hasEnvValue(env: EnvSource, name: string): boolean {
   const raw = env[name];
@@ -43,6 +51,7 @@ export type UpstreamTimeoutConfig = {
   fetchBodyTimeoutMs: number;
   fetchConnectTimeoutMs: number;
   fetchKeepAliveTimeoutMs: number;
+  streamDisconnectGracePeriodMs: number;
 };
 
 export type TlsClientTimeoutConfig = {
@@ -136,6 +145,15 @@ export function getUpstreamTimeoutConfig(
       logger,
     }
   );
+  const streamDisconnectGracePeriodMs = readTimeoutMs(
+    env,
+    "STREAM_DISCONNECT_GRACE_PERIOD_MS",
+    DEFAULT_STREAM_DISCONNECT_GRACE_PERIOD_MS,
+    {
+      allowZero: true,
+      logger,
+    }
+  );
 
   return {
     fetchTimeoutMs,
@@ -143,6 +161,7 @@ export function getUpstreamTimeoutConfig(
     streamReadinessTimeoutMs,
     streamReadinessMaxTimeoutMs,
     sseHeartbeatIntervalMs,
+    streamDisconnectGracePeriodMs,
     fetchHeadersTimeoutMs: readTimeoutMs(env, "FETCH_HEADERS_TIMEOUT_MS", fetchTimeoutMs, {
       allowZero: true,
       logger,

tests/unit/stream-disconnect-grace-period-9653.test.ts

@@ -0,0 +1,144 @@
+/**
+ * Regression tests for #9653 — a client that closes its connection right after
+ * reading a fully-completed SSE stream could race OmniRoute's own completion
+ * bookkeeping, persisting a false 499/0-tokens for a request that actually
+ * delivered its full response. createClientDisconnectGraceHandler gives a
+ * delayed completion a grace period to land before finalizing as a failure.
+ */
+import test from "node:test";
+import assert from "node:assert/strict";
+
+import { createClientDisconnectGraceHandler } from "../../open-sse/utils/streamFailureFinalization.ts";
+
+/** Fake scheduler: setTimeoutFn calls are queued instead of run on a real clock;
+ * the test advances them one at a time by calling `runNext()`. */
+function createFakeScheduler() {
+  const queue: Array<() => void> = [];
+  return {
+    setTimeoutFn: (callback: () => void) => {
+      queue.push(callback);
+      return 0;
+    },
+    runNext: () => {
+      const callback = queue.shift();
+      if (!callback) throw new Error("no pending timer to run");
+      callback();
+    },
+    pending: () => queue.length,
+  };
+}
+
+test("#9653: a completion already recorded before disconnect finalizes immediately, never calls finalize", () => {
+  const scheduler = createFakeScheduler();
+  let finalizeCalls = 0;
+  const handler = createClientDisconnectGraceHandler({
+    isStreamCompletionRecorded: () => true,
+    gracePeriodMs: 10_000,
+    finalize: () => {
+      finalizeCalls++;
+    },
+    setTimeoutFn: scheduler.setTimeoutFn,
+  });
+
+  const result = handler({ reason: "request_signal_aborted", duration: 100 });
+
+  assert.equal(result, true);
+  assert.equal(finalizeCalls, 0);
+  assert.equal(scheduler.pending(), 0);
+});
+
+test("#9653: gracePeriodMs <= 0 finalizes immediately (old pre-#9653 behavior)", () => {
+  const scheduler = createFakeScheduler();
+  let finalizeCalls = 0;
+  const handler = createClientDisconnectGraceHandler({
+    isStreamCompletionRecorded: () => false,
+    gracePeriodMs: 0,
+    finalize: () => {
+      finalizeCalls++;
+    },
+    setTimeoutFn: scheduler.setTimeoutFn,
+  });
+
+  const result = handler({ reason: "request_signal_aborted", duration: 100 });
+
+  assert.equal(result, true);
+  assert.equal(finalizeCalls, 1);
+  assert.equal(scheduler.pending(), 0);
+});
+
+test("#9653: a real completion landing during the grace period skips finalize entirely", () => {
+  const scheduler = createFakeScheduler();
+  let finalizeCalls = 0;
+  let completed = false;
+  const realNow = Date.now;
+  let fakeNow = 0;
+  Date.now = () => fakeNow;
+  try {
+    const handler = createClientDisconnectGraceHandler({
+      isStreamCompletionRecorded: () => completed,
+      gracePeriodMs: 1000,
+      finalize: () => {
+        finalizeCalls++;
+      },
+      pollIntervalMs: 100,
+      setTimeoutFn: scheduler.setTimeoutFn,
+    });
+
+    const result = handler({ reason: "request_signal_aborted", duration: 100 });
+    assert.equal(result, true, "claims handled immediately, poll is pending");
+    assert.equal(finalizeCalls, 0);
+    assert.equal(scheduler.pending(), 1);
+
+    // The real completion lands mid-grace-period.
+    completed = true;
+    fakeNow += 100;
+    scheduler.runNext();
+
+    assert.equal(finalizeCalls, 0, "completion landed — must not finalize as a failure");
+    assert.equal(scheduler.pending(), 0, "poll must stop once completion is recorded");
+  } finally {
+    Date.now = realNow;
+  }
+});
+
+test("#9653: no completion ever lands — finalizes as a failure once the deadline passes", () => {
+  const scheduler = createFakeScheduler();
+  let finalizeCalls = 0;
+  let finalizedEvent: unknown = null;
+  const realNow = Date.now;
+  let fakeNow = 0;
+  Date.now = () => fakeNow;
+  try {
+    const handler = createClientDisconnectGraceHandler({
+      isStreamCompletionRecorded: () => false,
+      gracePeriodMs: 1000,
+      finalize: (event) => {
+        finalizeCalls++;
+        finalizedEvent = event;
+      },
+      pollIntervalMs: 250,
+      setTimeoutFn: scheduler.setTimeoutFn,
+    });
+
+    const event = { reason: "request_signal_aborted", duration: 100 };
+    const result = handler(event);
+    assert.equal(result, true);
+    assert.equal(finalizeCalls, 0);
+
+    // Advance through the grace period in 250ms polling steps; completion never lands.
+    for (let i = 0; i < 3; i++) {
+      fakeNow += 250;
+      scheduler.runNext();
+      assert.equal(finalizeCalls, 0, `must not finalize before the deadline (step ${i})`);
+    }
+
+    fakeNow += 250; // now at 1000ms, deadline reached
+    scheduler.runNext();
+
+    assert.equal(finalizeCalls, 1, "must finalize as a failure once the deadline passes");
+    assert.deepEqual(finalizedEvent, event);
+    assert.equal(scheduler.pending(), 0, "must stop polling after finalizing");
+  } finally {
+    Date.now = realNow;
+  }
+});