From d86fd3cd2a4d23573d65731c5d4d4e5da3228906 Mon Sep 17 00:00:00 2001 From: Markus Hartung Date: Fri, 7 Aug 2026 02:04:58 +0200 Subject: [PATCH] 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 5d0fe28c4246518f7c7b588795a6da1573b5df16) --- open-sse/config/constants.ts | 9 ++ open-sse/handlers/chatCore.ts | 22 ++- open-sse/utils/streamFailureFinalization.ts | 57 ++++++- src/shared/utils/runtimeTimeouts.ts | 19 +++ ...tream-disconnect-grace-period-9653.test.ts | 144 ++++++++++++++++++ 5 files changed, 241 insertions(+), 10 deletions(-) create mode 100644 tests/unit/stream-disconnect-grace-period-9653.test.ts diff --git a/open-sse/config/constants.ts b/open-sse/config/constants.ts index ad648f39bf..1aae014a7a 100644 --- a/open-sse/config/constants.ts +++ b/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. diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 770e373030..3e7f6cab2f 100644 --- a/open-sse/handlers/chatCore.ts +++ b/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 diff --git a/open-sse/utils/streamFailureFinalization.ts b/open-sse/utils/streamFailureFinalization.ts index dfd4def6be..7d4e57ffba 100644 --- a/open-sse/utils/streamFailureFinalization.ts +++ b/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); diff --git a/src/shared/utils/runtimeTimeouts.ts b/src/shared/utils/runtimeTimeouts.ts index cd148d9177..667fa82b64 100644 --- a/src/shared/utils/runtimeTimeouts.ts +++ b/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, diff --git a/tests/unit/stream-disconnect-grace-period-9653.test.ts b/tests/unit/stream-disconnect-grace-period-9653.test.ts new file mode 100644 index 0000000000..93a7e10eac --- /dev/null +++ b/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; + } +});