From c9e52397a56a28c98ca4ef7c3b825eaa41cce39b Mon Sep 17 00:00:00 2001 From: Diego Rodrigues de Sa e Souza <8016841+diegosouzapw@users.noreply.github.com> Date: Thu, 21 May 2026 23:17:04 -0300 Subject: [PATCH] fix(antigravity): fail over stalled sessions before response headers (port #2464 to v3.8.2) (#2537) Integrated into release/v3.8.2 --- open-sse/config/constants.ts | 5 ++ open-sse/executors/antigravity.ts | 74 +++++++++++++++++-- open-sse/handlers/chatCore.ts | 43 ++++++++++- open-sse/utils/error.ts | 2 + src/sse/handlers/chat.ts | 65 +++++++++++++++- tests/unit/error-message-sanitization.test.ts | 7 ++ tests/unit/executor-antigravity.test.ts | 49 +++++++++++- 7 files changed, 231 insertions(+), 14 deletions(-) diff --git a/open-sse/config/constants.ts b/open-sse/config/constants.ts index 8d3b1b979d..d45ad56a96 100644 --- a/open-sse/config/constants.ts +++ b/open-sse/config/constants.ts @@ -22,6 +22,11 @@ export const STREAM_IDLE_TIMEOUT_MS = upstreamTimeouts.streamIdleTimeoutMs; // first token, while dead 200 OK streams fail fast enough for combo fallback. export const STREAM_READINESS_TIMEOUT_MS = upstreamTimeouts.streamReadinessTimeoutMs; +// Error code used when an upstream Antigravity request stalls before response +// headers are returned. Keep it shared so executor, core normalization and +// account fallback detection cannot drift. +export const ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE = "ANTIGRAVITY_PRE_RESPONSE_TIMEOUT"; + // Heartbeat interval for synthetic SSE keepalive emission toward the downstream // client (Capy, Claude Code, OpenAI SDK, etc). Keeps strict proxies from // dropping the connection during long upstream thinking phases. Set to 0 to diff --git a/open-sse/executors/antigravity.ts b/open-sse/executors/antigravity.ts index 28cba9ec38..81bc0a138e 100644 --- a/open-sse/executors/antigravity.ts +++ b/open-sse/executors/antigravity.ts @@ -1,13 +1,20 @@ import crypto, { randomUUID } from "crypto"; import { BaseExecutor, + mergeAbortSignals, mergeUpstreamExtraHeaders, type ExecuteInput, type ExecutorLog, type ProviderCredentials, } from "./base.ts"; import { applyFingerprint, isCliCompatEnabled } from "../config/cliFingerprints.ts"; -import { PROVIDERS, OAUTH_ENDPOINTS, HTTP_STATUS } from "../config/constants.ts"; +import { + PROVIDERS, + OAUTH_ENDPOINTS, + HTTP_STATUS, + STREAM_READINESS_TIMEOUT_MS, + ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, +} from "../config/constants.ts"; import { scrubProxyAndFingerprintHeaders } from "../services/antigravityHeaderScrub.ts"; import { antigravityNativeOAuthUserAgent, @@ -44,7 +51,6 @@ import { const MAX_RETRY_AFTER_MS = 60_000; const LONG_RETRY_THRESHOLD_MS = 60_000; const CREDITS_EXHAUSTED_TTL_MS = 5 * 60 * 60 * 1000; // 5 hours - const BARE_PRO_IDS = new Set(["gemini-3.1-pro"]); interface AntigravityContent { @@ -130,6 +136,26 @@ type AntigravityRequestEnvelope = Record & { enabledCreditTypes?: string[]; }; +class AntigravityPreResponseTimeoutError extends Error { + code = ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE; + status = HTTP_STATUS.GATEWAY_TIMEOUT; + + constructor(timeoutMs: number, url: string) { + super(`Antigravity upstream did not return response headers within ${timeoutMs}ms: ${url}`); + this.name = "TimeoutError"; + } +} + +function getAbortErrorCode(error: unknown): string | null { + if (!error || typeof error !== "object") return null; + const value = (error as { code?: unknown }).code; + return typeof value === "string" ? value : null; +} + +function isAntigravityPreResponseTimeout(error: unknown): boolean { + return getAbortErrorCode(error) === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE; +} + /** * Per-account GOOGLE_ONE_AI credits-exhausted tracker. * Key: accountId (OAuth subject / email). Value: expiry timestamp. @@ -777,6 +803,44 @@ export class AntigravityExecutor extends BaseExecutor { const creditsMode = getCreditsMode(); const useCreditsFirst = shouldUseCreditsFirst(credentials?.accessToken || "", creditsMode); + const fetchWithReadinessTimeout = async ( + url: string, + init: RequestInit, + timeoutMs = STREAM_READINESS_TIMEOUT_MS + ): Promise => { + const boundedTimeoutMs = Math.max(0, Math.floor(timeoutMs)); + if (boundedTimeoutMs <= 0) { + return fetch(url, init); + } + + const timeoutController = new AbortController(); + let timeoutId: ReturnType | null = setTimeout(() => { + timeoutController.abort(new AntigravityPreResponseTimeoutError(boundedTimeoutMs, url)); + }, boundedTimeoutMs); + + const existingSignal = init.signal instanceof AbortSignal ? init.signal : null; + const combinedSignal = existingSignal + ? mergeAbortSignals(existingSignal, timeoutController.signal) + : timeoutController.signal; + + try { + return await fetch(url, { ...init, signal: combinedSignal }); + } catch (error) { + if ( + timeoutController.signal.aborted && + isAntigravityPreResponseTimeout(timeoutController.signal.reason) + ) { + throw timeoutController.signal.reason; + } + throw error; + } finally { + if (timeoutId) { + clearTimeout(timeoutId); + timeoutId = null; + } + } + }; + for (let urlIndex = 0; urlIndex < fallbackCount; urlIndex++) { const url = this.buildUrl(model, upstreamStream, urlIndex); const headers = this.buildHeaders(credentials, upstreamStream); @@ -852,7 +916,7 @@ export class AntigravityExecutor extends BaseExecutor { ); } - let response = await fetch(url, { + let response = await fetchWithReadinessTimeout(url, { method: "POST", headers: finalHeaders, body: getChunkedOrFixedBody(serializedRequest.bodyString, stream), @@ -864,7 +928,7 @@ export class AntigravityExecutor extends BaseExecutor { const retryHeaders = { ...finalHeaders }; removeHeaderCaseInsensitive(retryHeaders, "x-goog-user-project"); log?.debug?.("RETRY", "403 with x-goog-user-project, retrying once without it"); - response = await fetch(url, { + response = await fetchWithReadinessTimeout(url, { method: "POST", headers: retryHeaders, body: getChunkedOrFixedBody(serializedRequest.bodyString, stream), @@ -934,7 +998,7 @@ export class AntigravityExecutor extends BaseExecutor { ); const finalCreditsHeaders = serializedCreditsRequest.headers; try { - const creditsResp = await fetch(url, { + const creditsResp = await fetchWithReadinessTimeout(url, { method: "POST", headers: finalCreditsHeaders, body: getChunkedOrFixedBody(serializedCreditsRequest.bodyString, stream), diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index f63b474d81..2a189975ec 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -40,6 +40,7 @@ import { SSE_HEARTBEAT_INTERVAL_MS, STREAM_IDLE_TIMEOUT_MS, STREAM_READINESS_TIMEOUT_MS, + ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, } from "../config/constants.ts"; import { classifyProviderError, @@ -894,11 +895,19 @@ function isSemaphoreCapacityError(error: unknown): error is Error & { code: stri ); } -function createStreamingErrorResult(statusCode: number, message: string, code?: string) { +function createStreamingErrorResult( + statusCode: number, + message: string, + code?: string, + type?: string +) { const errorBody = buildErrorBody(statusCode, message); if (code) { errorBody.error.code = code; } + if (type) { + errorBody.error.type = type; + } const body = `data: ${JSON.stringify(errorBody)}\n\ndata: [DONE]\n\n`; @@ -918,6 +927,12 @@ function createStreamingErrorResult(statusCode: number, message: string, code?: }; } +function getUpstreamErrorIdentifier(error: unknown): string | undefined { + if (!error || typeof error !== "object") return undefined; + const value = (error as { code?: unknown }).code; + return typeof value === "string" && value.length > 0 ? value : undefined; +} + function wrapReadableStreamWithFinalize( readable: ReadableStream, finalize: () => void @@ -3665,6 +3680,9 @@ export async function handleChatCore({ error.name === "AbortError" ? "Request aborted" : formatProviderError(error, provider, model, failureStatus); + const upstreamErrorCode = getUpstreamErrorIdentifier(error); + const upstreamErrorType = + upstreamErrorCode === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE ? "upstream_timeout" : undefined; appendRequestLog({ model, provider, @@ -3685,10 +3703,29 @@ export async function handleChatCore({ } persistFailureUsage( failureStatus, - error instanceof Error && error.name ? error.name : "upstream_error" + upstreamErrorCode || (error instanceof Error && error.name ? error.name : "upstream_error") ); console.log(`${COLORS.red}[ERROR] ${failureMessage}${COLORS.reset}`); - return createErrorResult(failureStatus, failureMessage); + if (stream && upstreamErrorCode) { + const result = createStreamingErrorResult( + failureStatus, + failureMessage, + upstreamErrorCode, + upstreamErrorType + ); + return { + ...result, + errorType: upstreamErrorType, + errorCode: upstreamErrorCode, + }; + } + return createErrorResult( + failureStatus, + failureMessage, + null, + upstreamErrorCode, + upstreamErrorType + ); } // We need to peek at the error text if it's 400 for Qwen let upstreamErrorParsed = false; diff --git a/open-sse/utils/error.ts b/open-sse/utils/error.ts index 91a886a86a..9cf1bf10ef 100644 --- a/open-sse/utils/error.ts +++ b/open-sse/utils/error.ts @@ -317,6 +317,7 @@ export function createErrorResult( status: number; error: string; errorType?: string; + errorCode?: string; response: Response; retryAfterMs?: number; } = { @@ -324,6 +325,7 @@ export function createErrorResult( status: statusCode, error: body.error.message, errorType, + errorCode, response: new Response(JSON.stringify(body), { status: statusCode, headers: { "Content-Type": "application/json" }, diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 24435e23c7..7c31d8507f 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -19,7 +19,10 @@ import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; import { handleComboChat } from "@omniroute/open-sse/services/combo.ts"; import { resolveComboConfig } from "@omniroute/open-sse/services/comboConfig.ts"; import { injectHandoffIntoBody } from "@omniroute/open-sse/services/contextHandoff.ts"; -import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; +import { + HTTP_STATUS, + ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, +} from "@omniroute/open-sse/config/constants.ts"; import { getTargetFormat } from "@omniroute/open-sse/services/provider.ts"; import { getModelTargetFormat, @@ -29,7 +32,12 @@ import type { AutoVariant } from "@omniroute/open-sse/services/autoCombo/autoPre import * as log from "../utils/logger"; import { checkAndRefreshToken } from "../services/tokenRefresh"; import { deleteHandoff, getHandoff } from "@/lib/db/contextHandoffs"; -import { getCachedSettings, getCombos } from "@/lib/localDb"; +import { + deleteSessionAccountAffinity, + getCachedSettings, + getCombos, + getSessionAccountAffinity, +} from "@/lib/localDb"; import { ensureOpenAIStoreSessionFallback, isOpenAIResponsesStoreEnabled, @@ -961,8 +969,57 @@ async function handleSingleModelChat( } if (result.errorType === "stream_timeout" || result.errorType === "stream_early_eof") { - // Stream readiness timeout is an upstream stall, not an account/quota failure. - // Do NOT mark the account as unavailable or trip the circuit breaker. + // Stream readiness timeout is an upstream stall after an HTTP response was received, + // not an account/quota failure. Do NOT mark the account unavailable here. + return result.response; + } + + const isAntigravityPreResponseTimeout = + provider === "antigravity" && + result.status === HTTP_STATUS.GATEWAY_TIMEOUT && + (result.errorType === "upstream_timeout" || + result.errorCode === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE); + + if (isAntigravityPreResponseTimeout) { + const { shouldFallback, cooldownMs } = await markAccountUnavailable( + credentials.connectionId, + result.status, + result.error || ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, + provider, + model, + providerProfile + ); + + if (shouldFallback && !hasForcedConnection) { + log.warn( + "AUTH", + `Antigravity connection ${accountId}... timed out before response headers, trying fallback connection` + ); + if (Number.isFinite(cooldownMs) && cooldownMs > 0) { + lastCooldownMs = cooldownMs; + requestRetryLastCooldownMs = cooldownMs; + } + if (runtimeOptions.sessionAffinityKey) { + try { + const affinity = getSessionAccountAffinity( + runtimeOptions.sessionAffinityKey, + provider + ); + if (affinity?.connectionId === credentials.connectionId) { + deleteSessionAccountAffinity(runtimeOptions.sessionAffinityKey, provider); + } + } catch { + // best-effort: selection also excludes this connection for the current retry. + } + } + excludedConnectionIds.add(credentials.connectionId); + lastError = result.error; + lastStatus = result.status; + requestRetryLastError = result.error; + requestRetryLastStatus = result.status; + continue; + } + return result.response; } diff --git a/tests/unit/error-message-sanitization.test.ts b/tests/unit/error-message-sanitization.test.ts index 657a8aefd6..6c4a1d7192 100644 --- a/tests/unit/error-message-sanitization.test.ts +++ b/tests/unit/error-message-sanitization.test.ts @@ -335,6 +335,13 @@ test("createErrorResult — response body excludes upstream_details when not pro assert.ok(!("upstream_details" in body), "upstream_details must be absent when not provided"); }); +test("createErrorResult — exposes error code/type on the result object", async () => { + const { createErrorResult } = await import("../../open-sse/utils/error.ts"); + const result = createErrorResult(504, "upstream timeout", null, "UPSTREAM_TIMEOUT", "timeout"); + assert.equal(result.errorCode, "UPSTREAM_TIMEOUT"); + assert.equal(result.errorType, "timeout"); +}); + test("regression: upstream_details never contains stack trace text", async () => { const { createErrorResult } = await import("../../open-sse/utils/error.ts"); const upstream = { error: { message: "err" }, stack: "Error\n at /abs/path.ts:1:2" }; diff --git a/tests/unit/executor-antigravity.test.ts b/tests/unit/executor-antigravity.test.ts index 51f9dae5f5..399a49985a 100644 --- a/tests/unit/executor-antigravity.test.ts +++ b/tests/unit/executor-antigravity.test.ts @@ -588,6 +588,51 @@ test("AntigravityExecutor.execute embeds retryAfterMs when the upstream asks for } }); +test("AntigravityExecutor.execute tags pre-response stalls with a fallbackable timeout code", async () => { + const executor = new AntigravityExecutor(); + const originalFetch = globalThis.fetch; + const originalSetTimeout = globalThis.setTimeout; + seedAntigravityVersionCache("2026.04.17-test"); + + globalThis.fetch = async (_url, init) => { + await new Promise((_resolve, reject) => { + const signal = init?.signal as AbortSignal | undefined; + if (signal?.aborted) { + reject(signal.reason); + return; + } + signal?.addEventListener("abort", () => reject(signal.reason), { once: true }); + }); + throw new Error("unreachable"); + }; + globalThis.setTimeout = ((callback) => { + (callback as () => void)(); + return 0; + }) as typeof setTimeout; + + try { + await assert.rejects( + () => + executor.execute({ + model: "antigravity/gemini-2.5-flash", + body: { request: { contents: [] } }, + stream: true, + credentials: { accessToken: "token", projectId: "project-1" }, + log: { debug() {}, warn() {}, error() {} }, + }), + (error: unknown) => { + assert.equal((error as { code?: string }).code, "ANTIGRAVITY_PRE_RESPONSE_TIMEOUT"); + assert.equal((error as { name?: string }).name, "TimeoutError"); + assert.match((error as Error).message, /did not return response headers/); + return true; + } + ); + } finally { + globalThis.fetch = originalFetch; + globalThis.setTimeout = originalSetTimeout; + } +}); + test("AntigravityExecutor.execute applies CLI fingerprint when enabled", async () => { const executor = new AntigravityExecutor(); const originalFetch = globalThis.fetch; @@ -598,9 +643,9 @@ test("AntigravityExecutor.execute applies CLI fingerprint when enabled", async ( const headers = init?.headers as Record; const parsedBody = JSON.parse(String(init?.body)); - assert.equal( + assert.match( headers["User-Agent"], - "Antigravity/2026.04.17-test (Macintosh; Intel Mac OS X 10_15_7) Chrome/132.0.6834.160 Electron/39.2.3" + /^Antigravity\/2026\.04\.17-test \(.+\) Chrome\/132\.0\.6834\.160 Electron\/39\.2\.3$/ ); assert.equal(headers["x-client-name"], "antigravity"); assert.equal(headers["x-client-version"], "2026.04.17-test");