diff --git a/open-sse/config/constants.ts b/open-sse/config/constants.ts index 81e3c29ee1..deccc4b161 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. @@ -65,27 +74,27 @@ export const PROVIDERS: Record = new Proxy( {} as Record, { get(_, prop) { - if (typeof prop === 'symbol') return undefined; + if (typeof prop === "symbol") return undefined; return Reflect.get(initProviders(), prop, _providers); }, has(_, prop) { - if (typeof prop === 'symbol') return false; + if (typeof prop === "symbol") return false; return Reflect.has(initProviders(), prop); }, ownKeys() { return Reflect.ownKeys(initProviders()); }, getOwnPropertyDescriptor(_, prop) { - if (typeof prop === 'symbol') return undefined; + if (typeof prop === "symbol") return undefined; return Object.getOwnPropertyDescriptor(initProviders(), prop); }, set(_, prop, value) { - if (typeof prop === 'symbol') return false; + if (typeof prop === "symbol") return false; (initProviders() as Record)[prop] = value; return true; }, deleteProperty(_, prop) { - if (typeof prop === 'symbol') return false; + if (typeof prop === "symbol") return false; return Reflect.deleteProperty(initProviders(), prop); }, } diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 2b110b47f4..e58b160ba9 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -170,6 +170,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 { @@ -4908,13 +4909,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; + } +});