diff --git a/.github/workflows/selected-upstream-regression.yml b/.github/workflows/selected-upstream-regression.yml index 264c4fd96272..048a86e8608d 100644 --- a/.github/workflows/selected-upstream-regression.yml +++ b/.github/workflows/selected-upstream-regression.yml @@ -86,7 +86,14 @@ jobs: tests/unit/http-client-abort-guard-empty-response-16347.test.ts \ tests/unit/release-version-workspace-consistency.test.ts \ tests/unit/server-owned-tool-loop-flag.test.ts \ - tests/unit/bug-8020-codex-peek-body-timeout.test.ts + tests/unit/bug-8020-codex-peek-body-timeout.test.ts \ + tests/unit/antigravity-adaptive-health.test.ts \ + tests/unit/antigravity-timeout-retry.test.ts \ + tests/unit/antigravity-timing.test.ts \ + tests/unit/preferred-connection-failover.test.ts \ + tests/unit/codex-same-account-transport-retry-9708.test.ts \ + tests/integration/antigravity-resilience-settings.test.ts \ + tests/integration/antigravity-resilience-routing.test.ts - name: UI provider availability (Vitest / jsdom) run: ./node_modules/.bin/vitest run tests/unit/ui/resilience-connections.test.tsx --maxWorkers=2 - name: Typecheck core diff --git a/docs/audits/2026-10-24-antigravity-resilience/README.md b/docs/audits/2026-10-24-antigravity-resilience/README.md new file mode 100644 index 000000000000..789a26a1d668 --- /dev/null +++ b/docs/audits/2026-10-24-antigravity-resilience/README.md @@ -0,0 +1,89 @@ +# Antigravity resilience overhaul audit + +This branch starts at production commit `44fdae6b4b9d63fa41f62a1941e3ff2f366743f6`. +It is a code and static review only. The production host was not built, started, or +used for tests, and no upstream model calls were made. + +## Root cause and baseline + +The request loop classified every Antigravity 504 as a pre-response timeout in one +branch, but the generic same-account transport retry remained reachable in other +paths. Its transport status set included 504, so a confirmed header timeout could +send the same account again before account rotation. The change requires both +`errorType=upstream_timeout` and +`errorCode=ANTIGRAVITY_PRE_RESPONSE_TIMEOUT` for the specialized path, excludes +confirmed Antigravity timeouts from same-account retry, and keeps 502, 503 and +socket-reset handling separate. + +The owner supplied the baseline for the last 24 hours: about 2,824 Gemini 3.8 +records, 261 HTTP 504, 576 HTTP 499 and 1,874 HTTP 200. Prefix summaries were +04572930: 303 requests (97 504, 106 499, 93 429, 7 200), 96c44b1c: 2,099 +(100 504, 375 499, 1,612 200), and 42684230: 232 (59 504, 57 499, 116 200). +The prefixes are not account numbers. The 113 same-account repeats reported by +the owner were treated as the regression signal. 499 cancellations are excluded +from health denominators and do not mark accounts; 429 remains quota/rate-limit +state. + +## Implemented behavior + +`open-sse/services/antigravityHealth.ts` provides a bounded process-local +connection + normalized-model rolling window. It requires eight completed samples, +uses a 50% timeout ratio or three consecutive confirmed timeouts, quarantines for +30 seconds with exponential backoff capped at five minutes, admits one half-open +probe, and requires two sequential successes for recovery. 499, quota, auth and +local rejection outcomes are ignored. Entries and history are bounded, and a +generation fence prevents a stale concurrent result from undoing a newer state. +Selection applies health after existing quota, terminal, model and forced-connection +rules. If every eligible pair is quarantined, one real demand-driven request is +allowed as an escape probe. The feature is off by default. + +The connections API exposes only bounded model, state, counts, ratios, cooldown, +backoff and inflight fields. It does not expose credentials, upstream text, URLs or +PII. The dashboard connection detail renders these fields. The API requires the +existing management authentication policy and access-token write scope. + +`open-sse/services/antigravityTiming.ts` records optional phase telemetry with a +bounded hashed account label and bounded event count. Idle duration, inflight +count, header latency and first meaningful Gemini delta are measured from actual +phase boundaries. DNS, connect and TLS phases are explicitly `null` when the +available fetch API cannot measure them. Timing is disabled by default and emits no +request body, bearer token, URL or user content. Existing dispatcher keepalive and +fresh-socket recovery paths remain in control; no synthetic paid probes were added. + +Preferred failover defaults remain exported as 5,000 and 12,000 ms for compatibility. +The runtime settings object adds bounded, persisted values (500..120,000 ms), +adaptive-health and timing switches. `/api/resilience` validates the patch, requires +management auth, writes the settings atomically with the existing revision conflict +guard, and reads settings for each new request. The dashboard exposes the four +values. + +## Tradeoffs and operator knobs + +The health tracker is process-local, so a restart clears observations and does not +create a migration. This avoids persisting ephemeral per-account state or secrets. +The default feature gate is conservative. Operators can enable adaptive health and +timing telemetry from Settings → Resilience. Preferred timeout values are clamped to +500..120,000 ms. Recovery uses real traffic only, so an idle account may remain +quarantined until the next request reaches its half-open window. + +## Verification + +Added focused unit coverage for exact timeout classification, forced/pinned behavior, +bounded transport retry, model normalization, low-sample protection, 499 exclusion, +rolling expiry, backoff, half-open concurrency, recovery hysteresis, bounded state, +stream completion and timing lifecycle. Added the two test files to +`.github/workflows/selected-upstream-regression.yml`. + +Only static checks were performed in this worktree: source inspection and +`git diff --check`. CI has not run yet. This report intentionally does not claim +typecheck, lint, unit, integration, build, browser or deployment verification. + +## Remaining uncertainties and limitations + +The owner-provided aggregate counts do not identify Antigravity accounts, and this +branch cannot establish the Gemini 3.8 account #2 idle-restart diagnosis from +production. The telemetry will distinguish idle, concurrency, header and first-delta +phases after CI and an approved deployment. DNS/connect/TLS remain null unless the +selected dispatcher exposes safe phase hooks. The adaptive tracker is not shared +between multiple worker processes. No deployment, merge or production setting change +is included in this PR. diff --git a/open-sse/executors/antigravity.ts b/open-sse/executors/antigravity.ts index 7847a7c237cf..0e1812299ab3 100644 --- a/open-sse/executors/antigravity.ts +++ b/open-sse/executors/antigravity.ts @@ -1,4 +1,5 @@ import crypto, { randomUUID } from "crypto"; +import { isAntigravityPreResponseTimeout } from "../services/antigravityTimeout.ts"; import { BaseExecutor, mergeUpstreamExtraHeaders, @@ -565,8 +566,7 @@ export class AntigravityExecutor extends BaseExecutor { (status === HTTP_STATUS.RATE_LIMITED || status === HTTP_STATUS.NOT_FOUND || status === HTTP_STATUS.BAD_GATEWAY || - status === HTTP_STATUS.SERVICE_UNAVAILABLE || - status === HTTP_STATUS.GATEWAY_TIMEOUT) && + status === HTTP_STATUS.SERVICE_UNAVAILABLE) && urlIndex + 1 < this.getFallbackCount() ); } @@ -1197,6 +1197,7 @@ export class AntigravityExecutor extends BaseExecutor { try { result = await this.executeOnce(input, candidate, physicalSendCounter); } catch (error) { + if (isAntigravityPreResponseTimeout(error)) throw error; const outcome = handleAntigravityFallbackChainError( input, error, @@ -1338,13 +1339,14 @@ export class AntigravityExecutor extends BaseExecutor { if (signal?.aborted || isAbortError(error)) { throw signal?.reason ?? error; } + if (isAntigravityPreResponseTimeout(error)) throw error; lastError = error; l.error( "TELEMETRY", - `[Antigravity] Network/Fetch Error - URL: ${url}, Model: ${model}, Error: ${error instanceof Error ? error.message : String(error)}` + "[Antigravity] Network/fetch failure before executor result" ); if (urlIndex + 1 < fallbackCount) { - l.debug("RETRY", `Error on ${url}, trying fallback ${urlIndex + 1}`); + l.debug("RETRY", `Transport failure, trying endpoint ${urlIndex + 1}`); continue; } throw error; @@ -1658,6 +1660,7 @@ export class AntigravityExecutor extends BaseExecutor { switchAuth: decision.kind === "short_cooldown_switch_auth", }; } catch (error) { + if (isAntigravityPreResponseTimeout(error)) throw error; if (signal?.aborted || isAbortError(error)) { throw signal?.reason ?? error; } diff --git a/open-sse/executors/antigravity/executeAttempt.ts b/open-sse/executors/antigravity/executeAttempt.ts index 276412efaaea..4a8c4988c394 100644 --- a/open-sse/executors/antigravity/executeAttempt.ts +++ b/open-sse/executors/antigravity/executeAttempt.ts @@ -27,6 +27,8 @@ import { type SsePassthroughResult, } from "./streamingPassthrough.ts"; import type { AntigravityCredentials } from "../antigravity.ts"; +import { isAntigravityPreResponseTimeout } from "../../services/antigravityTimeout.ts"; +import { timeAntigravityFetch } from "../../services/antigravityTiming.ts"; const LONG_RETRY_THRESHOLD_MS = 60_000; const CREDITS_EXHAUSTED_TTL_MS = 5 * 60 * 60 * 1000; // 5 hours @@ -90,22 +92,12 @@ 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}`); + constructor(timeoutMs: number) { + super(`Antigravity upstream did not return response headers within ${timeoutMs}ms`); 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; -} - export function isAbortError(error: unknown): boolean { return ( (error instanceof DOMException && error.name === "AbortError") || @@ -123,6 +115,14 @@ export async function fetchAntigravityWithReadinessTimeout( url: string, init: RequestInit, timeoutMs = STREAM_READINESS_TIMEOUT_MS +): Promise { + return timeAntigravityFetch(() => fetchWithReadinessTimeout(url, init, timeoutMs), init.signal); +} + +async function fetchWithReadinessTimeout( + url: string, + init: RequestInit, + timeoutMs = STREAM_READINESS_TIMEOUT_MS ): Promise { const boundedTimeoutMs = Math.max(0, Math.floor(timeoutMs)); if (boundedTimeoutMs <= 0) { @@ -131,7 +131,7 @@ export async function fetchAntigravityWithReadinessTimeout( const timeoutController = new AbortController(); let timeoutId: ReturnType | null = setTimeout(() => { - timeoutController.abort(new AntigravityPreResponseTimeoutError(boundedTimeoutMs, url)); + timeoutController.abort(new AntigravityPreResponseTimeoutError(boundedTimeoutMs)); }, boundedTimeoutMs); const existingSignal = init.signal instanceof AbortSignal ? init.signal : null; @@ -142,6 +142,10 @@ export async function fetchAntigravityWithReadinessTimeout( try { return await fetch(url, { ...init, signal: combinedSignal }); } catch (error) { + // The caller's cancellation wins; never relabel it as an account timeout. + if (existingSignal?.aborted && combinedSignal.reason !== timeoutController.signal.reason) { + throw existingSignal.reason; + } if ( timeoutController.signal.aborted && isAntigravityPreResponseTimeout(timeoutController.signal.reason) @@ -261,11 +265,6 @@ function serializeAntigravityRequest( return applyFingerprint(provider, { ...headers }, serializedBody); } -function getRequestTargetModel(body: Record): string { - const target = body.model; - return typeof target === "string" && target.length > 0 ? target : "unknown"; -} - /** Sanitize unsupported schema metadata, then apply credits-first injection for one attempt. */ export function finalizeAntigravityRequestBody( transformed: Record, @@ -347,7 +346,7 @@ export async function sendAntigravityRequest( log.debug( "TELEMETRY", - `[Antigravity] Execute - URL: ${url}, Model: ${model}, Target: ${getRequestTargetModel(transformedBody)}, RetryAttempt: ${retryAttempt}` + `[Antigravity] Execute attempt ${Math.min(100, retryAttempt)}` ); // Dump outgoing headers (mask Authorization) and envelope shape for debugging. @@ -361,7 +360,7 @@ export async function sendAntigravityRequest( const physicalSendOrdinal = ++physicalSendCounter.value; log.debug( "TELEMETRY", - `[Antigravity] PhysicalSend - RequestId: ${correlationId ?? "none"}, URL: ${url}, Model: ${model}, PhysicalSend: ${physicalSendOrdinal}, RetryAttempt: ${retryAttempt}` + `[Antigravity] PhysicalSend: ${physicalSendOrdinal}, RetryAttempt: ${retryAttempt}` ); let response = await fetchAntigravityWithReadinessTimeout(url, { method: "POST", @@ -379,7 +378,7 @@ export async function sendAntigravityRequest( const retryPhysicalSendOrdinal = ++physicalSendCounter.value; log.debug( "TELEMETRY", - `[Antigravity] PhysicalSend - RequestId: ${correlationId ?? "none"}, URL: ${url}, Model: ${model}, PhysicalSend: ${retryPhysicalSendOrdinal}, RetryAttempt: ${retryAttempt}, Cause: x-goog-user-project-403` + `[Antigravity] PhysicalSend: ${retryPhysicalSendOrdinal}, Cause: x-goog-user-project-403` ); response = await fetchAntigravityWithReadinessTimeout(url, { method: "POST", @@ -394,7 +393,7 @@ export async function sendAntigravityRequest( if (!response.ok) { log.warn( "TELEMETRY", - `[Antigravity] Error Response - URL: ${url}, Status: ${response.status}, Model: ${model}` + `[Antigravity] Error Response: ${response.status}` ); if (response.status === HTTP_STATUS.NOT_FOUND) { // The backend may have shipped/renamed models the synced catalog does not @@ -450,7 +449,7 @@ export async function tryCreditsRetry( const creditsPhysicalSendOrdinal = ++physicalSendCounter.value; log.debug( "TELEMETRY", - `[Antigravity] PhysicalSend - RequestId: ${correlationId ?? "none"}, URL: ${url}, PhysicalSend: ${creditsPhysicalSendOrdinal}, Cause: google-one-ai-credits-retry` + `[Antigravity] PhysicalSend: ${creditsPhysicalSendOrdinal}, Cause: google-one-ai-credits-retry` ); const creditsResp = await fetchAntigravityWithReadinessTimeout(url, { method: "POST", @@ -492,6 +491,7 @@ export async function tryCreditsRetry( markCreditsExhausted(accountId); return null; } catch (creditsErr) { + if (isAntigravityPreResponseTimeout(creditsErr)) throw creditsErr; if (signal?.aborted || isAbortError(creditsErr)) { throw signal?.reason ?? creditsErr; } diff --git a/open-sse/executors/antigravity/proFallbackChain.ts b/open-sse/executors/antigravity/proFallbackChain.ts index ec5b710c1f85..a32c0769d9ad 100644 --- a/open-sse/executors/antigravity/proFallbackChain.ts +++ b/open-sse/executors/antigravity/proFallbackChain.ts @@ -4,6 +4,7 @@ // Extracted from antigravity.ts (file-size cap) -- mirrors the existing // antigravity/sseCollect.ts submodule pattern. import type { ExecuteInput } from "../base.ts"; +import { isAntigravityPreResponseTimeout } from "../../services/antigravityTimeout.ts"; /** Shape of one execute()/executeOnce() result (kept local to avoid importing the class). */ export type AntigravityExecuteResult = { @@ -42,7 +43,7 @@ export function handleAntigravityFallbackChainError( resolvedUpstreamId: string ): AntigravityFallbackChainErrorOutcome { // Abort signal (user disconnect) — propagate immediately, do not retry. - if (isAntigravityAbortError(input, error)) { + if (isAntigravityAbortError(input, error) || isAntigravityPreResponseTimeout(error)) { return { action: "throw", error }; } if (i < chain.length - 1) { diff --git a/open-sse/handlers/chatCore/nonStreamingProviderLeg.ts b/open-sse/handlers/chatCore/nonStreamingProviderLeg.ts index efe732f7ee1c..b10d90558439 100644 --- a/open-sse/handlers/chatCore/nonStreamingProviderLeg.ts +++ b/open-sse/handlers/chatCore/nonStreamingProviderLeg.ts @@ -25,6 +25,8 @@ import { extractUsageFromResponse } from "../usageExtractor.ts"; import { sanitizeUsagePayloadForRequest } from "../../utils/usageTracking.ts"; import { createErrorResult, formatProviderError } from "../../utils/error.ts"; import { isLocalStreamLifecycleError } from "@/shared/utils/circuitBreaker"; +import { isAntigravityPreResponseTimeout } from "../../services/antigravityTimeout.ts"; +import { ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE } from "../../config/constants.ts"; import { unwrapClinepassEnvelope } from "../../utils/clinepassEnvelope.ts"; import { unwrapClineNonStreamingEnvelope } from "./clineResponseEnvelope.ts"; import { @@ -513,7 +515,18 @@ export async function runNonStreamingProviderLeg( connectionId, model: currentModel, }); - const errorResult = legError(failureStatus, failureMessage, error); + // Preserve the PRIVATE exact pre-response timeout marker across the + // non-streaming leg. The status-only legError used to discard it, so the + // outer router saw an indistinguishable generic 504 and could not score + // account+model health or attribute the account switch correctly. + const isHeaderTimeout = + (provider === "antigravity" || provider === "agy") && + isAntigravityPreResponseTimeout(error); + const errorResult = legError( + failureStatus, failureMessage, error, null, + isHeaderTimeout ? ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE : undefined, + isHeaderTimeout ? "upstream_timeout" : undefined + ); return { kind: "error", result: errorResult as ChatCoreErrorResult, diff --git a/open-sse/handlers/chatCore/upstreamTimeouts.ts b/open-sse/handlers/chatCore/upstreamTimeouts.ts index 2b5bdbed5e57..43c13d693d07 100644 --- a/open-sse/handlers/chatCore/upstreamTimeouts.ts +++ b/open-sse/handlers/chatCore/upstreamTimeouts.ts @@ -1,5 +1,6 @@ import { EXECUTOR_CONTRACT_VIOLATION_CODE, + ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, FETCH_TIMEOUT_MS, HTTP_STATUS, } from "../../config/constants.ts"; @@ -47,6 +48,12 @@ export function createUpstreamStartTimeoutError( `Upstream request did not return response headers after ${timeoutMs}ms (${provider}/${model})` ); err.name = "TimeoutError"; + // The outer per-request header timer is another producer of the SAME + // Antigravity pre-response condition as the executor's own readiness + // timer. Carry a private marker through the leg, never a raw URL/token. + if (provider === "antigravity" || provider === "agy") { + (err as Error & { code?: string }).code = ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE; + } return err; } diff --git a/open-sse/services/antigravityHealth.ts b/open-sse/services/antigravityHealth.ts new file mode 100644 index 000000000000..87c668b63601 --- /dev/null +++ b/open-sse/services/antigravityHealth.ts @@ -0,0 +1,189 @@ +import { resolveAntigravityModelId } from "../config/antigravityModelAliases.ts"; + +export type AntigravityHealthOutcome = "success" | "timeout" | "transport" | "ignored"; +type Sample = { at: number; outcome: Exclude }; +type Entry = { + connectionId: string; + model: string; + samples: Sample[]; + consecutiveTimeouts: number; + until: number; + backoff: number; + recoverySuccesses: number; + generation: number; + probe: boolean; + inflight: number; + touched: number; +}; + +export const ANTIGRAVITY_HEALTH_POLICY = Object.freeze({ + minSamples: 8, + maxSamples: 32, + windowMs: 300_000, + timeoutRatio: 0.5, + consecutiveTimeouts: 3, + baseCooldownMs: 30_000, + maxCooldownMs: 300_000, + recoverySuccesses: 2, + maxEntries: 1000, + idleTtlMs: 1_800_000, +}); + +export function normalizeAntigravityHealthModel(model: string): string { + const id = model.trim().replace(/^(?:antigravity|agy)\//i, "").toLowerCase(); + // Never retain arbitrary request content in diagnostic state. + return /^[a-z0-9_.:-]{1,160}$/.test(id) ? resolveAntigravityModelId(id) : "unknown"; +} + +export interface AntigravityHealthAttempt { + readonly probe: boolean; + finish(outcome: AntigravityHealthOutcome): void; +} + +/** Process-local, bounded, demand-driven state. No background model calls. */ +export class AntigravityHealthTracker { + private entries = new Map(); + constructor(private now: () => number = Date.now) {} + + private key(connectionId: string, model: string): string { + return JSON.stringify([connectionId, normalizeAntigravityHealthModel(model)]); + } + + private trim(entry: Entry): void { + entry.samples = entry.samples.filter((sample) => sample.at > this.now() - ANTIGRAVITY_HEALTH_POLICY.windowMs); + if (entry.samples.length === 0) entry.consecutiveTimeouts = 0; + } + + private get(connectionId: string, model: string): Entry | null { + const key = this.key(connectionId, model); + const existing = this.entries.get(key); + if (existing) { + this.trim(existing); + return existing; + } + for (const [id, entry] of this.entries) { + if (entry.inflight === 0 && this.now() - entry.touched > ANTIGRAVITY_HEALTH_POLICY.idleTtlMs) { + this.entries.delete(id); + } + } + if (this.entries.size >= ANTIGRAVITY_HEALTH_POLICY.maxEntries) { + // Never evict an in-flight probe: it would allow a second probe for the same pair. + const idle = [...this.entries].filter(([, entry]) => entry.inflight === 0) + .sort((a, b) => a[1].touched - b[1].touched)[0]; + if (!idle) return null; + this.entries.delete(idle[0]); + } + const entry: Entry = { + connectionId, model: normalizeAntigravityHealthModel(model), samples: [], + consecutiveTimeouts: 0, until: 0, backoff: 0, recoverySuccesses: 0, + generation: 0, probe: false, inflight: 0, touched: this.now(), + }; + this.entries.set(key, entry); + return entry; + } + + /** Apply only to a pool that already passed quota, terminal, model and pin rules. */ + candidates(connections: T[], model: string, forced = false): T[] { + const healthy: T[] = []; + const due: T[] = []; + const open: T[] = []; + for (const connection of connections) { + const entry = this.entries.get(this.key(connection.id, model)); + if (!entry || entry.backoff === 0) healthy.push(connection); + else if (!entry.probe) (entry.until <= this.now() ? due : open).push(connection); + } + if (forced) return [...healthy, ...due, ...open].slice(0, 1); + // Include due probes so priority/round-robin can eventually test recovery. + if (healthy.length || due.length) return [...healthy, ...due]; + // Availability escape hatch: one real request to the earliest quarantined pair. + // begin() enforces a single probe even when selectors race. + return open.sort((a, b) => this.entries.get(this.key(a.id, model))!.until - + this.entries.get(this.key(b.id, model))!.until).slice(0, 1); + } + + remainingCooldownMs(connectionId: string, model: string): number { + return Math.max(0, (this.entries.get(this.key(connectionId, model))?.until ?? 0) - this.now()); + } + + begin(connectionId: string, model: string, allowEarlyProbe = false): AntigravityHealthAttempt | null { + const entry = this.get(connectionId, model); + // Capacity overflow fails open without allocating unbounded state. + if (!entry) return { probe: false, finish() {} }; + const probe = entry.backoff > 0; + if (entry.probe) return null; + if (probe && entry.until > this.now() && !allowEarlyProbe) return null; + if (probe) entry.probe = true; + entry.inflight++; + entry.touched = this.now(); + const generation = entry.generation; + let finished = false; + return { + probe, + finish: (outcome) => { + if (finished) return; + finished = true; + entry.inflight--; + entry.touched = this.now(); + // An older concurrent success must not undo a newer quarantine. + if (generation !== entry.generation) return; + if (probe) entry.probe = false; + if (outcome === "ignored") return; // 499, quota, auth, local rejection + this.trim(entry); + entry.samples.push({ at: this.now(), outcome }); + entry.samples = entry.samples.slice(-ANTIGRAVITY_HEALTH_POLICY.maxSamples); + entry.consecutiveTimeouts = outcome === "timeout" ? entry.consecutiveTimeouts + 1 : 0; + if (probe) { + if (outcome === "success") { + entry.recoverySuccesses++; + if (entry.recoverySuccesses >= ANTIGRAVITY_HEALTH_POLICY.recoverySuccesses) { + entry.backoff = 0; + entry.until = 0; + entry.samples = []; + entry.consecutiveTimeouts = 0; + entry.generation++; + } else entry.until = this.now(); + } else this.quarantine(entry); + return; + } + const timeouts = entry.samples.filter((sample) => sample.outcome === "timeout").length; + if (entry.samples.length >= ANTIGRAVITY_HEALTH_POLICY.minSamples && + (timeouts / entry.samples.length >= ANTIGRAVITY_HEALTH_POLICY.timeoutRatio || + entry.consecutiveTimeouts >= ANTIGRAVITY_HEALTH_POLICY.consecutiveTimeouts)) { + this.quarantine(entry); + } + }, + }; + } + + private quarantine(entry: Entry): void { + entry.backoff = Math.min(5, entry.backoff + 1); + entry.until = this.now() + Math.min(ANTIGRAVITY_HEALTH_POLICY.maxCooldownMs, + ANTIGRAVITY_HEALTH_POLICY.baseCooldownMs * 2 ** (entry.backoff - 1)); + entry.recoverySuccesses = 0; + entry.generation++; + entry.probe = false; + } + + snapshot(connectionId?: string) { + return [...this.entries.values()].filter((entry) => !connectionId || entry.connectionId === connectionId) + .filter((entry) => entry.inflight > 0 || this.now() - entry.touched <= ANTIGRAVITY_HEALTH_POLICY.idleTtlMs) + .map((entry) => { + this.trim(entry); + const timeouts = entry.samples.filter((sample) => sample.outcome === "timeout").length; + return { + connectionId: entry.connectionId, model: entry.model, + state: entry.probe ? "probing" : entry.backoff === 0 ? "healthy" : entry.until > this.now() ? "quarantined" : "half_open", + reason: entry.backoff > 0 ? "pre_response_timeout_rate" : null, + samples: entry.samples.length, timeouts, + transportFailures: entry.samples.filter((sample) => sample.outcome === "transport").length, + timeoutRatio: entry.samples.length ? timeouts / entry.samples.length : 0, + consecutiveTimeouts: entry.consecutiveTimeouts, + cooldownRemainingMs: Math.max(0, entry.until - this.now()), + backoffLevel: entry.backoff, recoverySuccesses: entry.recoverySuccesses, + inflight: entry.inflight, + }; + }); + } +} + +export const antigravityHealth = new AntigravityHealthTracker(); diff --git a/open-sse/services/antigravityTimeout.ts b/open-sse/services/antigravityTimeout.ts new file mode 100644 index 000000000000..b0e6a163b30c --- /dev/null +++ b/open-sse/services/antigravityTimeout.ts @@ -0,0 +1,14 @@ +import { ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE } from "../config/constants.ts"; + +/** Exact phase evidence. A generic 504 or BodyTimeoutError is insufficient. */ +export function isAntigravityPreResponseTimeout(error: unknown): boolean { + return ( + error !== null && + typeof error === "object" && + (error as { code?: unknown }).code === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE + ); +} + +export function isAntigravityProvider(provider: string | null | undefined): boolean { + return provider === "antigravity" || provider === "agy"; +} diff --git a/open-sse/services/antigravityTiming.ts b/open-sse/services/antigravityTiming.ts new file mode 100644 index 000000000000..6ac0ddecd084 --- /dev/null +++ b/open-sse/services/antigravityTiming.ts @@ -0,0 +1,171 @@ +import { AsyncLocalStorage } from "node:async_hooks"; +import { createHash, randomUUID } from "node:crypto"; +import { normalizeAntigravityHealthModel } from "./antigravityHealth.ts"; +import { isAntigravityPreResponseTimeout } from "./antigravityTimeout.ts"; + +type TimingContext = { + connectionId: string; + model: string; + requestId: string; + sends: number; + emit: (event: Record) => void; +}; +const context = new AsyncLocalStorage(); +const MAX_ACCOUNTS = 1000; +const MAX_EVENTS_PER_ATTEMPT = 24; +type AccountTiming = { inflight: number; lastFinishedAt: number | null }; + +/** Only allow actual Gemini output to set first delta; pings/errors/usage do not. */ +export function hasAntigravityMeaningfulDelta(frame: string): boolean { + for (const line of frame.split(/\r?\n/)) { + if (!line.startsWith("data:")) continue; + try { + const parsed = JSON.parse(line.slice(5).trim()) as Record; + const response = (parsed.response ?? parsed) as Record; + if (!Array.isArray(response.candidates)) continue; + for (const candidate of response.candidates) { + const parts = candidate?.content?.parts; + if (!Array.isArray(parts)) continue; + if (parts.some((part) => (typeof part?.text === "string" && part.text.length > 0) || + (typeof part?.functionCall?.name === "string" && part.functionCall.name.length > 0) || + (typeof part?.inlineData?.data === "string" && part.inlineData.data.length > 0) || + part?.executableCode || part?.codeExecutionResult)) return true; + } + } catch { /* Partial, invalid or terminal frame: no measured delta. */ } + } + return false; +} + +export class AntigravityTimingTracker { + private accounts = new Map(); + constructor(private now: () => number = () => performance.now()) {} + + async fetch(send: () => Promise, signal?: AbortSignal | null): Promise { + const scope = context.getStore(); + if (!scope) return send(); + let account = this.accounts.get(scope.connectionId); + if (!account) { + if (this.accounts.size >= MAX_ACCOUNTS) { + const idle = [...this.accounts].find(([, item]) => item.inflight === 0); + if (!idle) return send(); + this.accounts.delete(idle[0]); + } + account = { inflight: 0, lastFinishedAt: null }; + this.accounts.set(scope.connectionId, account); + } + const started = this.now(); + const idleDurationMs = account.inflight > 0 ? 0 : account.lastFinishedAt === null + ? null : Math.max(0, started - account.lastFinishedAt); + const concurrentInflight = ++account.inflight; + const physicalSend = ++scope.sends; + let headersAt: number | null = null; + let firstDeltaAt: number | null = null; + let finished = false; + const event = (phase: string, extra: Record = {}) => { + if (physicalSend > MAX_EVENTS_PER_ATTEMPT / 3) return; + // Correlation ids stay log fields, never metric labels. No URLs or payloads. + const fields = { + event: "antigravity_timing", phase, requestId: scope.requestId, + account: createHash("sha256").update(scope.connectionId).digest("hex").slice(0, 16), + model: normalizeAntigravityHealthModel(scope.model), physicalSend, + idleDurationMs, concurrentInflight, + dnsMs: null, connectMs: null, tlsMs: null, + headerLatencyMs: headersAt === null ? null : headersAt - started, + firstMeaningfulDeltaFromHeadersMs: firstDeltaAt === null || headersAt === null ? null : firstDeltaAt - headersAt, + firstMeaningfulDeltaFromSendMs: firstDeltaAt === null ? null : firstDeltaAt - started, + elapsedMs: Math.max(0, this.now() - started), ...extra, + }; + try { scope.emit(fields); } catch { /* Diagnostics cannot change request behavior. */ } + }; + const finish = (phase: string) => { + if (finished) return; + finished = true; + account!.inflight--; + account!.lastFinishedAt = this.now(); + signal?.removeEventListener("abort", onAbort); + event(phase); + }; + const abortPhase = () => (signal?.reason as { name?: string } | undefined)?.name === "TimeoutError" + ? "upstream_deadline" : "caller_abort"; + const onAbort = () => finish(abortPhase()); + signal?.addEventListener("abort", onAbort, { once: true }); + try { + const response = await send(); + headersAt = this.now(); + event("headers", { status: response.status }); + if (signal?.aborted) finish(abortPhase()); + if (!response.body) { + finish("complete"); + return response; + } + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let pending = ""; + let skipOversizedFrame = false; + const inspect = (bytes: Uint8Array) => { + if (finished || firstDeltaAt !== null) return; + // Never retain an unbounded chunk/frame, including large image deltas. + if (bytes.byteLength > 65536) { pending = ""; skipOversizedFrame = true; return; } + pending += decoder.decode(bytes, { stream: true }); + for (;;) { + const boundary = pending.search(/\r?\n\r?\n/); + if (boundary < 0) break; + const frame = pending.slice(0, boundary); + pending = pending.slice(boundary).replace(/^\r?\n\r?\n/, ""); + if (!skipOversizedFrame && hasAntigravityMeaningfulDelta(frame)) { + firstDeltaAt = this.now(); + event("first_meaningful_delta"); + pending = ""; + return; + } + skipOversizedFrame = false; + } + if (pending.length > 65536) { pending = ""; skipOversizedFrame = true; } + }; + const body = new ReadableStream({ + async pull(controller) { + try { + const chunk = await reader.read(); + if (chunk.done) { + finish("complete"); + reader.releaseLock(); + controller.close(); + } else { + inspect(chunk.value); + controller.enqueue(chunk.value); + } + } catch (error) { + finish("body_error"); + reader.releaseLock(); + controller.error(error); + } + }, + async cancel(reason) { + finish("cancelled"); + try { await reader.cancel(reason); } finally { reader.releaseLock(); } + }, + }); + return new Response(body, { + status: response.status, statusText: response.statusText, headers: response.headers, + }); + } catch (error) { + finish(isAntigravityPreResponseTimeout(error) ? "pre_response_timeout" : signal?.aborted ? abortPhase() : "fetch_failed"); + throw error; + } + } +} + +const tracker = new AntigravityTimingTracker(); +export const timeAntigravityFetch = (send: () => Promise, signal?: AbortSignal | null) => + tracker.fetch(send, signal); + +export function withAntigravityTiming( + options: { enabled: boolean; connectionId: string; model: string; correlationId?: string | null; + emit: TimingContext["emit"] }, + fn: () => T +): T { + if (!options.enabled) return fn(); + const requestId = /^[0-9a-f]{8}-[0-9a-f-]{27}$/i.test(options.correlationId ?? "") + ? options.correlationId! : randomUUID(); + return context.run({ ...options, requestId, sends: 0 }, fn); +} diff --git a/src/app/(dashboard)/dashboard/resilience/connections/components/ConnectionDetail.tsx b/src/app/(dashboard)/dashboard/resilience/connections/components/ConnectionDetail.tsx index 653a346ce8e6..ed835fb5991d 100644 --- a/src/app/(dashboard)/dashboard/resilience/connections/components/ConnectionDetail.tsx +++ b/src/app/(dashboard)/dashboard/resilience/connections/components/ConnectionDetail.tsx @@ -101,6 +101,19 @@ export default function ConnectionDetail({ )}

{t("detail.lockouts")}

+ {(connection.antigravityHealth?.length ?? 0) > 0 && ( +
+

{typeof t.has === "function" && t.has("detail.antigravityHealth" as never) ? t("detail.antigravityHealth" as never) : "Antigravity model health"}

+ {connection.antigravityHealth?.map((health) => ( +
+
{health.model}: {health.state}
+
{health.timeouts}/{health.samples} confirmed header timeouts ({Math.round(health.timeoutRatio * 100)}%)
+
{health.reason ?? "healthy"}; {formatRemaining(health.cooldownRemainingMs)}; {health.inflight} inflight
+
{health.recoverySuccesses}/2 recovery successes
+
+ ))} +
+ )} {connection.lockouts.length > 0 ? (
    {connection.lockouts.map((l, i) => ( diff --git a/src/app/(dashboard)/dashboard/settings/components/AntigravityResilienceCard.tsx b/src/app/(dashboard)/dashboard/settings/components/AntigravityResilienceCard.tsx new file mode 100644 index 000000000000..171ddd64e789 --- /dev/null +++ b/src/app/(dashboard)/dashboard/settings/components/AntigravityResilienceCard.tsx @@ -0,0 +1,49 @@ +"use client"; + +import { useState } from "react"; +import { useTranslations } from "next-intl"; +import { Button, Card } from "@/shared/components"; +import type { AntigravityResilienceSettings } from "@/lib/resilience/antigravitySettings"; +import { BooleanField, NumberField } from "./ResilienceFields"; + +export default function AntigravityResilienceCard({ value, saving, onSave }: { + value: AntigravityResilienceSettings; + saving: boolean; + onSave: (next: AntigravityResilienceSettings) => Promise; +}) { + const t = useTranslations("settings"); + const common = useTranslations("common"); + const [draft, setDraft] = useState(value); + const [previous, setPrevious] = useState(value); + if (previous !== value) { setPrevious(value); setDraft(value); } + const tx = (key: string, fallback: string) => typeof t.has === "function" && t.has(key as never) ? t(key as never) : fallback; + + return ( + +

    {tx("antigravityResilienceTitle", "Antigravity resilience")}

    +

    + {tx("antigravityResilienceDescription", "Preferred account timeouts apply when another eligible account is available. Changes apply to new requests.")} +

    +
    + setDraft({ ...draft, preferredTimeoutMs })} /> + setDraft({ ...draft, gemini38PreferredTimeoutMs })} /> + setDraft({ ...draft, adaptiveHealthEnabled })} /> + setDraft({ ...draft, timingTelemetryEnabled })} /> +
    +
    + + +
    +     + ); +} diff --git a/src/app/(dashboard)/dashboard/settings/components/ResilienceTab.tsx b/src/app/(dashboard)/dashboard/settings/components/ResilienceTab.tsx index 7c066a7d88b7..8af345b8cd61 100644 --- a/src/app/(dashboard)/dashboard/settings/components/ResilienceTab.tsx +++ b/src/app/(dashboard)/dashboard/settings/components/ResilienceTab.tsx @@ -7,6 +7,8 @@ import { useTranslations } from "next-intl"; import AutoDisableCard from "./AutoDisableCard"; import ModelLockoutCard from "./ModelLockoutCard"; import { NumberField, BooleanField } from "./ResilienceFields"; +import AntigravityResilienceCard from "./AntigravityResilienceCard"; +import type { AntigravityResilienceSettings } from "@/lib/resilience/antigravitySettings"; type RequestQueueSettings = { autoEnableApiKeyProviders: boolean; @@ -61,6 +63,7 @@ export type ProviderCooldownSettings = { }; type ResilienceResponse = { + antigravity?: AntigravityResilienceSettings; requestQueue: RequestQueueSettings; connectionCooldown: { oauth: ConnectionCooldownProfileSettings; @@ -79,6 +82,7 @@ type ResilienceResponse = { function toResilienceResponse(json: ResilienceResponse): ResilienceResponse { return { + antigravity: json.antigravity, requestQueue: json.requestQueue, connectionCooldown: json.connectionCooldown, providerBreaker: json.providerBreaker, @@ -1217,6 +1221,10 @@ export default function ResilienceTab() { return (
    + {data.antigravity && ( + savePatch("antigravity", { antigravity })} /> + )} : {}; + const timeout = (key: "preferredTimeoutMs" | "gemini38PreferredTimeoutMs") => { + const candidate = record[key]; + return typeof candidate === "number" && Number.isInteger(candidate) && candidate >= 500 && candidate <= 120000 + ? candidate : fallback[key]; + }; + return { + preferredTimeoutMs: timeout("preferredTimeoutMs"), + gemini38PreferredTimeoutMs: timeout("gemini38PreferredTimeoutMs"), + adaptiveHealthEnabled: typeof record.adaptiveHealthEnabled === "boolean" + ? record.adaptiveHealthEnabled : fallback.adaptiveHealthEnabled, + timingTelemetryEnabled: typeof record.timingTelemetryEnabled === "boolean" + ? record.timingTelemetryEnabled : fallback.timingTelemetryEnabled, + }; +} + +export function resolveAntigravitySettings(settings: Record): AntigravityResilienceSettings { + const resilience = settings.resilienceSettings as { antigravity?: unknown } | undefined; + return normalizeAntigravitySettings(resilience?.antigravity); +} diff --git a/src/lib/resilience/settings.ts b/src/lib/resilience/settings.ts index a76de7b65a56..faefca9fb51e 100644 --- a/src/lib/resilience/settings.ts +++ b/src/lib/resilience/settings.ts @@ -5,6 +5,7 @@ import { } from "@omniroute/open-sse/config/constants"; import type { JsonRecord, ResilienceSettings, ResilienceSettingsPatch } from "./settings/types"; +import { DEFAULT_ANTIGRAVITY_RESILIENCE, normalizeAntigravitySettings } from "./antigravitySettings"; import { asRecord, toInteger, @@ -65,6 +66,7 @@ export const DEFAULT_REQUEST_QUEUE_MAX_DEPTH = (() => { })(); export const DEFAULT_RESILIENCE_SETTINGS: ResilienceSettings = { + antigravity: DEFAULT_ANTIGRAVITY_RESILIENCE, requestQueue: { autoEnableApiKeyProviders: true, requestsPerMinute: DEFAULT_API_LIMITS.requestsPerMinute, @@ -285,6 +287,7 @@ function buildLegacyFallback(settings: JsonRecord): ResilienceSettings { DEFAULT_RESILIENCE_SETTINGS.waitForCooldown.budgetMs ), }, + antigravity: DEFAULT_ANTIGRAVITY_RESILIENCE, comboCooldownWait: DEFAULT_RESILIENCE_SETTINGS.comboCooldownWait, quotaShareConcurrencyLimit: DEFAULT_RESILIENCE_SETTINGS.quotaShareConcurrencyLimit, providerCooldown: DEFAULT_RESILIENCE_SETTINGS.providerCooldown, @@ -326,6 +329,7 @@ export function resolveResilienceSettings( const fallback = buildLegacyFallback(record); return { + antigravity: normalizeAntigravitySettings(current.antigravity), requestQueue: normalizeRequestQueueSettings(current.requestQueue, fallback.requestQueue), connectionCooldown: { oauth: normalizeConnectionCooldownProfile( @@ -387,6 +391,7 @@ export function mergeResilienceSettings( updates: ResilienceSettingsPatch ): ResilienceSettings { return { + antigravity: normalizeAntigravitySettings(updates.antigravity, current.antigravity), requestQueue: normalizeRequestQueueSettings(updates.requestQueue, current.requestQueue), connectionCooldown: { oauth: normalizeConnectionCooldownProfile( diff --git a/src/lib/resilience/settings/types.ts b/src/lib/resilience/settings/types.ts index 028cf113b7c2..23d7a5b71a2c 100644 --- a/src/lib/resilience/settings/types.ts +++ b/src/lib/resilience/settings/types.ts @@ -236,6 +236,7 @@ export interface StreamThroughputWatchdogSettings { } export interface ResilienceSettings { + antigravity: import("../antigravitySettings").AntigravityResilienceSettings; requestQueue: RequestQueueSettings; connectionCooldown: Record; providerBreaker: Record; @@ -250,6 +251,7 @@ export interface ResilienceSettings { } export interface ResilienceSettingsPatch { + antigravity?: Partial; requestQueue?: Partial; connectionCooldown?: Partial>>; providerBreaker?: Partial>>; diff --git a/src/server/authz/accessScopes.ts b/src/server/authz/accessScopes.ts index f5422a746b38..2cd46ae9ccd3 100644 --- a/src/server/authz/accessScopes.ts +++ b/src/server/authz/accessScopes.ts @@ -32,6 +32,7 @@ export const ADMIN_SCOPE_PREFIXES: readonly string[] = [ /** Require `admin` only for mutating methods; GET/HEAD under these stay `read`. */ export const ADMIN_MUTATION_PREFIXES: readonly string[] = [ + "/api/resilience", // persisted routing timeouts and model-health policy "/api/providers", // POST add provider / rotate key = admin; GET status = read "/api/cli-tools/apply", // writes config onto the host filesystem ]; diff --git a/src/shared/validation/schemas/settings.ts b/src/shared/validation/schemas/settings.ts index 106a14e68a24..d14949e1efa5 100644 --- a/src/shared/validation/schemas/settings.ts +++ b/src/shared/validation/schemas/settings.ts @@ -160,6 +160,12 @@ export const credentialHealthCheckSettingsSchema = z export const updateResilienceSchema = z .object({ + antigravity: z.object({ + preferredTimeoutMs: z.number().int().min(500).max(120000).optional(), + gemini38PreferredTimeoutMs: z.number().int().min(500).max(120000).optional(), + adaptiveHealthEnabled: z.boolean().optional(), + timingTelemetryEnabled: z.boolean().optional(), + }).strict().optional(), requestQueue: requestQueueSettingsSchema.optional(), connectionCooldown: z .object({ @@ -223,7 +229,8 @@ export const updateResilienceSchema = z !value.profiles && !value.defaults && !value.providerQuotaOverrides && - !value.credentialHealthCheck + !value.credentialHealthCheck && + !value.antigravity ) { ctx.addIssue({ code: z.ZodIssueCode.custom, diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 0aee8b5665de..01ebeb94114b 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -83,6 +83,11 @@ import { import { dispatchChatWithAffinityEviction } from "./chatDispatch"; import { resolveInternalUpstreamStartTimeoutMs } from "../services/internalUpstreamStartTimeout"; import { resolvePreferredConnectionUpstreamStartTimeoutMs } from "../services/preferredConnectionFailover"; +import { getSettings } from "@/lib/db/settings"; +import { resolveAntigravitySettings, DEFAULT_ANTIGRAVITY_RESILIENCE } from "@/lib/resilience/antigravitySettings"; +import { antigravityHealth } from "@omniroute/open-sse/services/antigravityHealth"; +import { withAntigravityTiming } from "@omniroute/open-sse/services/antigravityTiming"; +import { bindAntigravityHealthAttempt, holdAntigravityHealthThroughResponse } from "../services/antigravityHealthLifecycle"; import { getCachedSettings, getCombosCacheVersion } from "@/lib/db/readCache"; import { comboCheckProvider, ghComboGate } from "./chat/githubLiveCatalogFilter.ts"; import { comboTargetPassesKeyModelPolicy } from "./chat/comboTargetKeyPolicy.ts"; @@ -1688,6 +1693,11 @@ async function handleSingleModelChat( // original early-EOF 502 so an exhausted sibling pool surfaces it verbatim (combo detection). let earlyEofOriginal: Response | null = null; const sameAccountTransportRetries = new Map(); + // Snapshot once per NEW request, bypassing inherited combo settings snapshots. + const antigravitySettings = provider === "antigravity" || provider === "agy" + ? resolveAntigravitySettings(await getSettings()) : DEFAULT_ANTIGRAVITY_RESILIENCE; + const timedOutConnectionIds = new Set(); + let lastAntigravityTimeoutResponse: Response | null = null; let accountFallbacks = 0; let selectedConnectionPlan: string | null = null; let sameAccountRetries = 0; @@ -1716,7 +1726,9 @@ async function handleSingleModelChat( const agy = agyLease.startAntigravityLeaseRequest(provider, runtimeOptions.correlationId); requestAttemptLoop: while (true) { - const excludedConnectionIds = new Set(agy.on ? agy.attempted : []); + const excludedConnectionIds = new Set([ + ...(agy.on ? agy.attempted : []), ...timedOutConnectionIds, + ]); let lastError = requestRetryLastError; let lastStatus = requestRetryLastStatus; let lastCooldownMs = requestRetryLastCooldownMs; @@ -1724,6 +1736,7 @@ async function handleSingleModelChat( initialPreselectedCredentials = null; while (true) { + if (requestSignal?.aborted) return errorResponse(499, "Request aborted"); const credentials = preselectedCredentials && excludedConnectionIds.size === 0 && !agy.on ? preselectedCredentials @@ -1792,6 +1805,7 @@ async function handleSingleModelChat( "allExpired" in credentials || !credentials.connectionId ) { + if (lastAntigravityTimeoutResponse) return lastAntigravityTimeoutResponse; if (earlyEofOriginal) return earlyEofOriginal; if (!credentials?.allRateLimited && agy.earliestRetryHintAtMs !== null) return agyLease.buildAntigravityPoolBusyResponse(agy.earliestRetryHintAtMs); @@ -2024,11 +2038,30 @@ async function handleSingleModelChat( reselectPoolMember?: () => Promise; } = { proxy: null }; const proxyStartTime = Date.now(); + const healthTicket = antigravitySettings.adaptiveHealthEnabled + ? antigravityHealth.begin(credentials.connectionId, effectiveModel, + hasForcedConnection || credentials.antigravityHealthFallback === true) + : undefined; + if (healthTicket === null) { + releaseOAuthSession(); + agyLease.release(leaseId); + if (hasForcedConnection) return errorResponse(503, "Antigravity model recovery probe already in flight"); + excludedConnectionIds.add(credentials.connectionId); + continue; + } + const antigravityHealthAttempt = healthTicket + ? bindAntigravityHealthAttempt(healthTicket, requestSignal) : undefined; // 4. Execute chat via core after breaker gate checks (with optional TLS tracking) if (telemetry) telemetry.startPhase("connect"); let execution: Awaited>; try { - execution = await dispatchChatWithAffinityEviction( + execution = await withAntigravityTiming({ + enabled: antigravitySettings.timingTelemetryEnabled, + connectionId: credentials.connectionId, + model: effectiveModel, + correlationId: runtimeOptions.correlationId, + emit: (event) => log.info("ANTIGRAVITY_TIMING", "Antigravity upstream phase", event), + }, () => dispatchChatWithAffinityEviction( { bypassCircuitBreaker: forceLiveComboTest || hasForcedConnection, breaker, @@ -2074,23 +2107,38 @@ async function handleSingleModelChat( provider, credentials, requestUpstreamStartTimeoutMs, - effectiveModel + effectiveModel, + antigravitySettings ), forcedConnectionId: hasForcedConnection ? forcedConnectionId : null, // #14116 }, runtimeOptions - ); + )); } catch (error) { + antigravityHealthAttempt?.finish("ignored"); releaseOAuthSession(); agyLease.release(leaseId); throw error; } if (telemetry) telemetry.endPhase(); if ("localResourcePressureResult" in execution) { + antigravityHealthAttempt?.finish("ignored"); + releaseOAuthSession(); agyLease.release(leaseId); return execution.localResourcePressureResult.response; } const { result, tlsFingerprintUsed } = execution; + if (antigravityHealthAttempt) { + if (requestSignal?.aborted || result.status === 499) antigravityHealthAttempt.finish("ignored"); + else if (result.success && requestBody.stream === true) { + result.response = holdAntigravityHealthThroughResponse(result.response, antigravityHealthAttempt); + } else if (result.success) antigravityHealthAttempt.finish("success"); + else antigravityHealthAttempt.finish( + result.status === 504 && result.errorType === "upstream_timeout" && + result.errorCode === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE + ? "timeout" : result.status === 502 || result.status === 503 ? "transport" : "ignored" + ); + } if (!result.success) releaseOAuthSession(); // Hand the lease to the SSE body's terminal lifecycle; anything else frees it now. if (result.success && agyLease.isStreamingAntigravityResponse(result.response)) @@ -2294,13 +2342,15 @@ async function handleSingleModelChat( return withRoutingTelemetry(result.response, credentials?.connectionId); } - const isAntigravityPreResponseTimeout = - provider === "antigravity" && - result.status === HTTP_STATUS.GATEWAY_TIMEOUT && - (result.errorType === "upstream_timeout" || - result.errorCode === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE); + const isAntigravityGatewayTimeout = + (provider === "antigravity" || provider === "agy") && + result.status === HTTP_STATUS.GATEWAY_TIMEOUT; - if (isAntigravityPreResponseTimeout) { + if (isAntigravityGatewayTimeout) { + // Persist across cooldown-loop restarts; one logical request never sends + // again on an account that already spent its header/504 waiting budget. + timedOutConnectionIds.add(credentials.connectionId); + lastAntigravityTimeoutResponse = withRoutingTelemetry(result.response, credentials.connectionId); const { shouldFallback, cooldownMs } = await markAccountUnavailable( credentials.connectionId, result.status, @@ -2308,13 +2358,17 @@ async function handleSingleModelChat( provider, model, providerProfile, - buildExhaustionOptions(runtimeOptions.correlationId ?? null, { isCombo }) + buildExhaustionOptions(runtimeOptions.correlationId ?? null, { + isCombo, + structuredError: { code: result.errorCode, type: result.errorType }, + healthAttempt: antigravityHealthAttempt, + }) ); if (shouldFallback && !hasForcedConnection) { log.warn( "AUTH", - `Antigravity connection ${accountId}... timed out before response headers, trying fallback connection` + "Antigravity upstream timeout, trying a distinct eligible connection" ); if (Number.isFinite(cooldownMs) && cooldownMs > 0) { lastCooldownMs = cooldownMs; @@ -2333,7 +2387,9 @@ async function handleSingleModelChat( // best-effort: selection also excludes this connection for the current retry. } } - noteAccountFallback("pre_response_timeout"); + noteAccountFallback(result.errorCode === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE + ? "pre_response_timeout" : "http_504"); + lastAntigravityTimeoutResponse = withRoutingTelemetry(result.response, credentials.connectionId); excludedConnectionIds.add(credentials.connectionId); lastError = result.error; lastStatus = result.status; @@ -2519,6 +2575,7 @@ async function handleSingleModelChat( !comboName && !forceLiveComboTest && shouldRetrySameAccountTransport({ + provider, status: result.status, errorText: errorStr, errorCode: result.errorCode, diff --git a/src/sse/services/antigravityHealthLifecycle.ts b/src/sse/services/antigravityHealthLifecycle.ts new file mode 100644 index 000000000000..56ec77b76d43 --- /dev/null +++ b/src/sse/services/antigravityHealthLifecycle.ts @@ -0,0 +1,58 @@ +import type { AntigravityHealthAttempt } from "@omniroute/open-sse/services/antigravityHealth"; +import { createStreamContentWatcher } from "@omniroute/open-sse/utils/streamReadiness"; + +export function bindAntigravityHealthAttempt( + attempt: AntigravityHealthAttempt, + signal?: AbortSignal | null +): AntigravityHealthAttempt { + const onAbort = () => finish("ignored"); + const finish: AntigravityHealthAttempt["finish"] = (outcome) => { + signal?.removeEventListener("abort", onAbort); + attempt.finish(signal?.aborted ? "ignored" : outcome); + }; + if (signal?.aborted) finish("ignored"); + else signal?.addEventListener("abort", onAbort, { once: true }); + return { probe: attempt.probe, finish }; +} + +/** Count a successful stream only when it completes, never at HTTP 200 headers. */ +export function holdAntigravityHealthThroughResponse( + response: Response, + attempt: AntigravityHealthAttempt +): Response { + if (!response.body) { + attempt.finish("success"); + return response; + } + const reader = response.body.getReader(); + const watcher = createStreamContentWatcher(); + const decoder = new TextDecoder(); + const body = new ReadableStream({ + async pull(controller) { + try { + const chunk = await reader.read(); + if (chunk.done) { + watcher.note(decoder.decode()); + watcher.finish(); + attempt.finish(watcher.sawContent() && !watcher.sawError() ? "success" : "ignored"); + controller.close(); + reader.releaseLock(); + } else { + watcher.note(decoder.decode(chunk.value, { stream: true })); + controller.enqueue(chunk.value); + } + } catch (error) { + attempt.finish("ignored"); + controller.error(error); + reader.releaseLock(); + } + }, + async cancel(reason) { + attempt.finish("ignored"); + try { await reader.cancel(reason); } finally { reader.releaseLock(); } + }, + }); + return new Response(body, { + status: response.status, statusText: response.statusText, headers: response.headers, + }); +} diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 20da6a7795c9..b5c619efc173 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -97,6 +97,9 @@ import { isQuotaPreflightEnabled, } from "@omniroute/open-sse/services/quotaPreflight.ts"; import { resolveResilienceSettings } from "@/lib/resilience/settings"; +import { resolveAntigravitySettings } from "@/lib/resilience/antigravitySettings"; +import { antigravityHealth, type AntigravityHealthAttempt } from "@omniroute/open-sse/services/antigravityHealth"; +import { isAntigravityProvider } from "@omniroute/open-sse/services/antigravityTimeout"; import { resolveModelLockoutSettings } from "@/lib/resilience/modelLockoutSettings"; import { buildMixedAvailabilityError, @@ -1077,6 +1080,7 @@ async function materializeConnection( requestedModel?: string | null; selectedByApiKeyPreference?: boolean; preferredFallbackAvailable?: boolean; + antigravityHealthFallback?: boolean; } = {} ) { const providerSpecificData = await hydrateConnectionProviderSpecificData(connection); @@ -1979,7 +1983,21 @@ export async function getProviderCredentials( : null; } - const orderedConnections = [...leasePolicy.connections].sort((a, b) => { + const healthEnabled = isAntigravityProvider(provider) && requestedModel && + resolveAntigravitySettings(settings).adaptiveHealthEnabled; + const healthCandidates = healthEnabled + ? antigravityHealth.candidates(leasePolicy.connections, requestedModel, Boolean(forcedConnectionId)) + : leasePolicy.connections; + if (healthCandidates.length === 0) { + return { + allRateLimited: true, + retryAfter: new Date(Date.now() + 1000).toISOString(), + retryAfterHuman: "1s", + lastError: "Antigravity model recovery probe already in flight", + lastErrorCode: 503, + }; + } + const orderedConnections = [...healthCandidates].sort((a, b) => { if (a.authType !== "oauth" || b.authType !== "oauth") return 0; const priorityDelta = (a.priority || 999) - (b.priority || 999); if (priorityDelta !== 0) return priorityDelta; @@ -2252,6 +2270,9 @@ export async function getProviderCredentials( selectNextLeaseCandidate, selectedByApiKeyPreference, preferredFallbackAvailable, + requestedModel, + antigravityHealthFallback: Boolean(healthEnabled && + antigravityHealth.remainingCooldownMs(connection.id, requestedModel) > 0), ...probeStamp, }); } @@ -2303,6 +2324,8 @@ export async function getProviderCredentials( requestedModel, selectedByApiKeyPreference, preferredFallbackAvailable, + antigravityHealthFallback: Boolean(healthEnabled && + antigravityHealth.remainingCooldownMs(connection.id, requestedModel) > 0), }); } finally { selectionLock?.release(); @@ -2721,11 +2744,15 @@ export function buildExhaustionOptions( /** Caller is the combo engine — it records its own model-level lockouts. */ isCombo?: boolean; headers?: Headers | Record | null; + structuredError?: { code?: string | null; type?: string | null } | null; + healthAttempt?: AntigravityHealthAttempt; } = {} ): { persistUnavailableState?: boolean; isCombo?: boolean; headers?: Headers | Record | null; + structuredError?: { code?: string | null; type?: string | null } | null; + healthAttempt?: AntigravityHealthAttempt; correlationId: string | null; } { return { ...rest, correlationId }; @@ -2747,8 +2774,23 @@ export async function markAccountUnavailable( correlationId?: string | null; streamOutputEmitted?: boolean; structuredError?: { code?: string | null; type?: string | null } | null; + healthAttempt?: AntigravityHealthAttempt; } = {} ) { + // Cancellation never belongs to account health, quota or cooldown state. + if (status === 499) { + options.healthAttempt?.finish("ignored"); + return { shouldFallback: false, cooldownMs: 0 }; + } + if (isAntigravityProvider(provider) && model && status === 504 && + options.structuredError?.type === "upstream_timeout" && + options.structuredError?.code === "ANTIGRAVITY_PRE_RESPONSE_TIMEOUT" && + (options.healthAttempt || resolveAntigravitySettings(await getCachedSettings()).adaptiveHealthEnabled)) { + // The handler owns an exactly-once ticket. Other callers without a ticket + // still rotate, but cannot invent an unpaired health observation here. + options.healthAttempt?.finish("timeout"); + return { shouldFallback: true, cooldownMs: antigravityHealth.remainingCooldownMs(connectionId, model) }; + } const currentMutex = markMutexes.get(connectionId) || Promise.resolve(); let resolveMutex: (() => void) | undefined; markMutexes.set( diff --git a/src/sse/services/preferredConnectionFailover.ts b/src/sse/services/preferredConnectionFailover.ts index f2f8c83d9838..08f5e83e37e6 100644 --- a/src/sse/services/preferredConnectionFailover.ts +++ b/src/sse/services/preferredConnectionFailover.ts @@ -1,5 +1,8 @@ -export const ANTIGRAVITY_PREFERRED_FAILOVER_START_TIMEOUT_MS = 5_000; -export const ANTIGRAVITY_GEMINI_38_PREFERRED_FAILOVER_START_TIMEOUT_MS = 12_000; +import { DEFAULT_ANTIGRAVITY_RESILIENCE, type AntigravityResilienceSettings } from "@/lib/resilience/antigravitySettings"; +import { normalizeAntigravityHealthModel } from "@omniroute/open-sse/services/antigravityHealth"; + +export const ANTIGRAVITY_PREFERRED_FAILOVER_START_TIMEOUT_MS = DEFAULT_ANTIGRAVITY_RESILIENCE.preferredTimeoutMs; +export const ANTIGRAVITY_GEMINI_38_PREFERRED_FAILOVER_START_TIMEOUT_MS = DEFAULT_ANTIGRAVITY_RESILIENCE.gemini38PreferredTimeoutMs; type PreferredCredentialState = { selectedByApiKeyPreference?: boolean; @@ -15,7 +18,8 @@ export function resolvePreferredConnectionUpstreamStartTimeoutMs( provider: string, credentials: PreferredCredentialState | null | undefined, requestHintMs?: number, - modelId?: string + modelId?: string, + settings: AntigravityResilienceSettings = DEFAULT_ANTIGRAVITY_RESILIENCE ): number | undefined { const requestHint = typeof requestHintMs === "number" && Number.isFinite(requestHintMs) && requestHintMs > 0 @@ -26,9 +30,9 @@ export function resolvePreferredConnectionUpstreamStartTimeoutMs( isAntigravity && credentials?.selectedByApiKeyPreference === true && credentials?.preferredFallbackAvailable === true - ? /^(?:antigravity\/)?gemini-3\.8-flash-tiered$/i.test((modelId ?? "").trim()) - ? ANTIGRAVITY_GEMINI_38_PREFERRED_FAILOVER_START_TIMEOUT_MS - : ANTIGRAVITY_PREFERRED_FAILOVER_START_TIMEOUT_MS + ? /^gemini-3\.8-flash(?:-tiered|-high|-medium|-low)?$/i.test(normalizeAntigravityHealthModel(modelId ?? "")) + ? settings.gemini38PreferredTimeoutMs + : settings.preferredTimeoutMs : undefined; if (preferredFailoverCap === undefined) return requestHint; diff --git a/src/sse/services/sameAccountTransportRetry.ts b/src/sse/services/sameAccountTransportRetry.ts index 4c82d94765a9..240381d6c5aa 100644 --- a/src/sse/services/sameAccountTransportRetry.ts +++ b/src/sse/services/sameAccountTransportRetry.ts @@ -36,14 +36,23 @@ export function isRetryablePreOutputTransportError( status: unknown, errorText: string | null | undefined, errorCode?: string | null, - errorType?: string | null + errorType?: string | null, + provider?: string | null ): boolean { + // Timeouts already spent their waiting budget. They must reach account fallback, + // including callers outside chat.ts's specialized Antigravity branch. + if (errorCode === "ANTIGRAVITY_PRE_RESPONSE_TIMEOUT" || errorType === "upstream_timeout") { + return false; + } + if ((provider === "antigravity" || provider === "agy") && Number(status) === 504) { + return false; + } if (errorType && NON_RETRYABLE_ERROR_TYPES.has(errorType)) return false; if (errorCode && String(errorCode).startsWith("LEASE_")) return false; const text = String(errorText || ""); const numericStatus = Number(status); - if (numericStatus === 429 || numericStatus === 401 || numericStatus === 400) return false; + if (numericStatus === 499 || numericStatus === 429 || numericStatus === 401 || numericStatus === 400) return false; if (/quota (threshold|exhausted)|credits exhausted/i.test(text)) return false; if (/invalid_request|prompt is too long|context.?length|unsupported model/i.test(text)) { return false; @@ -75,6 +84,7 @@ export function shouldRetrySameAccountTransport(options: { attempt: number; hasForcedConnection?: boolean; hasEmittedOutput?: boolean; + provider?: string | null; }): boolean { if (options.hasForcedConnection) return false; if (options.hasEmittedOutput) return false; @@ -83,7 +93,8 @@ export function shouldRetrySameAccountTransport(options: { options.status, options.errorText, options.errorCode, - options.errorType + options.errorType, + options.provider ); } diff --git a/src/types/resilience.ts b/src/types/resilience.ts index 2797a4efc5f6..98f425b021e7 100644 --- a/src/types/resilience.ts +++ b/src/types/resilience.ts @@ -29,6 +29,21 @@ export interface RotationAccountState { } export interface ConnectionState { + /** Process-local model health; bounded, structured and free of upstream messages. */ + antigravityHealth?: Array<{ + model: string; + state: string; + reason: string | null; + samples: number; + timeouts: number; + transportFailures: number; + timeoutRatio: number; + consecutiveTimeouts: number; + cooldownRemainingMs: number; + backoffLevel: number; + recoverySuccesses: number; + inflight: number; + }>; id: string; provider: string; name: string | null; diff --git a/tests/integration/antigravity-resilience-routing.test.ts b/tests/integration/antigravity-resilience-routing.test.ts new file mode 100644 index 000000000000..0631425259e1 --- /dev/null +++ b/tests/integration/antigravity-resilience-routing.test.ts @@ -0,0 +1,150 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { NextRequest } from "next/server"; +import { createChatPipelineHarness } from "./_chatPipelineHarness.ts"; +import { createManagementSessionHeaders } from "../helpers/managementSession.ts"; + +const harness = await createChatPipelineHarness("antigravity-resilience-routing"); +const { buildRequest, handleChat, settingsDb, resetStorage } = harness; +const providers = await import("../../src/lib/db/providers.ts"); +const auth = await import("../../src/sse/services/auth.ts"); +const { antigravityHealth } = await import("../../open-sse/services/antigravityHealth.ts"); +const { seedAntigravityIdeVersionCache, clearAntigravityVersionCaches } = await import("../../open-sse/services/antigravityVersion.ts"); +const { GET: stateGET } = await import("../../src/app/api/resilience/connections/route.ts"); +const MODEL = "gemini-3.8-flash-tiered"; +const TIMEOUT_CODE = "ANTIGRAVITY_PRE_RESPONSE_TIMEOUT"; + +test.beforeEach(async () => { + await resetStorage(); + seedAntigravityIdeVersionCache("2.1.1"); + await settingsDb.updateSettings({ + requestRetry: 0, + resilienceSettings: { antigravity: { adaptiveHealthEnabled: true }, waitForCooldown: { enabled: false } }, + }); +}); +test.afterEach(() => clearAntigravityVersionCaches()); +test.after(async () => harness.cleanup()); + +async function seed(priority: number) { + return providers.createProviderConnection({ + provider: "antigravity", authType: "oauth", name: `fixture-${priority}`, + accessToken: `fixture-token-${priority}`, refreshToken: `fixture-refresh-${priority}`, + projectId: "fixture-project", expiresAt: new Date(Date.now() + 3_600_000).toISOString(), + isActive: true, testStatus: "active", priority, + }); +} +function request(headers: Record = {}) { + return buildRequest({ headers, body: { + model: `antigravity/${MODEL}`, stream: false, + messages: [{ role: "user", content: `fixture-${crypto.randomUUID()}` }], + } }); +} +function success() { + return new Response('data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"fixture success"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}}\n\n', { + headers: { "content-type": "text/event-stream" }, + }); +} +function confirmedTimeout(): never { + throw Object.assign(new Error("fixture header timeout"), { name: "TimeoutError", code: TIMEOUT_CODE }); +} +function quarantine(id: string) { + for (let i = 0; i < 8; i++) antigravityHealth.begin(id, MODEL)!.finish("timeout"); +} + +test("confirmed pre-response 504 sends once on first account then immediately serves a distinct sibling", async () => { + const first = await seed(1); + const second = await seed(2); + const sends: string[] = []; + globalThis.fetch = async (_input, init) => { + const token = new Headers(init?.headers).get("authorization") || ""; + sends.push(token); + if (token.endsWith("fixture-token-1")) confirmedTimeout(); + return success(); + }; + const response = await handleChat(request()); + assert.equal(response.status, 200); + await response.text(); + assert.deepEqual(sends, ["Bearer fixture-token-1", "Bearer fixture-token-2"]); + assert.equal(response.headers.get("X-OmniRoute-Selected-Connection-Id"), second.id); + assert.equal(response.headers.get("X-OmniRoute-Same-Account-Retries"), null); + assert.equal(response.headers.get("X-OmniRoute-Account-Fallback-Reasons"), "pre_response_timeout"); + assert.equal(antigravityHealth.snapshot(first.id)[0].timeouts, 1); + const stored = await providers.getProviderConnectionById(first.id); + assert.equal(stored?.rateLimitedUntil ?? null, null, "model health must not cool the whole connection"); +}); + +test("exhausted pool cannot revisit either timed-out account via outer cooldown retry", async () => { + await seed(1); + await seed(2); + await settingsDb.updateSettings({ requestRetry: 3, resilienceSettings: { + antigravity: { adaptiveHealthEnabled: true }, + waitForCooldown: { enabled: true, maxRetries: 3, maxRetryWaitSec: 1, budgetMs: 3000 }, + } }); + let sends = 0; + globalThis.fetch = async () => { sends++; return confirmedTimeout(); }; + const response = await handleChat(request()); + assert.equal(response.status, 504); + assert.equal(sends, 2); +}); + +test("explicit connection pin surfaces original 504 and never rotates", async () => { + const first = await seed(1); + await seed(2); + let sends = 0; + globalThis.fetch = async () => { sends++; return confirmedTimeout(); }; + const response = await handleChat(request({ "x-omniroute-connection": first.id })); + assert.equal(response.status, 504); + const failure = await response.json(); + assert.equal(failure?.error?.code, TIMEOUT_CODE, "public timeout identifier stays from the sanitized safe code vocabulary"); + assert.equal(antigravityHealth.snapshot(first.id)[0].timeouts, 1, "private timeout marker survives non-streaming dispatch for quarantine"); + assert.equal(sends, 1); +}); + +test("quarantine is model scoped, quota/allowlist constrained, with all-quarantined fallback", async () => { + const first = await seed(1); + const second = await seed(2); + quarantine(first.id); + const selected = await auth.getProviderCredentials("antigravity", null, null, MODEL); + assert.equal(selected?.connectionId, second.id); + const otherModel = await auth.getProviderCredentials("antigravity", null, [first.id], "gemini-3.8-flash-low"); + assert.equal(otherModel?.connectionId, first.id); + quarantine(second.id); + const escape = await auth.getProviderCredentials("antigravity", null, [second.id], MODEL); + assert.equal(escape?.connectionId, second.id); + assert.equal(escape?.antigravityHealthFallback, true); + await providers.updateProviderConnection(second.id, { rateLimitedUntil: new Date(Date.now() + 60000).toISOString() }); + const quotaBlocked = await auth.getProviderCredentials("antigravity", null, [second.id], MODEL); + assert.equal(quotaBlocked?.connectionId, undefined, "health escape cannot override connection cooldown"); +}); + +test("499 never counts as a fault or triggers sibling rotation or persistent cooldown", async () => { + const first = await seed(1); + await seed(2); + let sends = 0; + globalThis.fetch = async () => { sends++; throw new DOMException("fixture client abort", "AbortError"); }; + const response = await handleChat(request()); + assert.equal(response.status, 499); + assert.equal(sends, 1); + assert.equal(antigravityHealth.snapshot(first.id)[0].samples, 0); + assert.equal(antigravityHealth.snapshot(first.id)[0].inflight, 0); + assert.equal(((await providers.getProviderConnectionById(first.id))?.rateLimitedUntil ?? null), null); +}); + +test("state API returns structured bounded model health without credentials or upstream text", async () => { + const connection = await seed(1); + quarantine(connection.id); + const response = await stateGET(new NextRequest("http://localhost/api/resilience/connections?provider=antigravity", { + headers: await createManagementSessionHeaders(), + })); + assert.equal(response.status, 200); + const body = await response.json(); + const state = body.connections.find((entry: { id: string }) => entry.id === connection.id).antigravityHealth[0]; + assert.equal(state.model, MODEL); + assert.equal(state.state, "quarantined"); + assert.equal(state.reason, "pre_response_timeout_rate"); + assert.equal(state.samples, 8); + const serialized = JSON.stringify(body); + for (const forbidden of ["fixture-token", "fixture-refresh", "fixture-project", "Authorization", "accessToken"]) { + assert.equal(serialized.includes(forbidden), false); + } +}); diff --git a/tests/integration/antigravity-resilience-settings.test.ts b/tests/integration/antigravity-resilience-settings.test.ts new file mode 100644 index 000000000000..6a22dc4f6d85 --- /dev/null +++ b/tests/integration/antigravity-resilience-settings.test.ts @@ -0,0 +1,86 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { makeManagementSessionRequest } from "../helpers/managementSession.ts"; +import { getDbInstance, resetDbInstance } from "../../src/lib/db/core.ts"; +import { getSettings, getSettingsRevision, updateSettings, SettingsRevisionConflictError } from "../../src/lib/db/settings.ts"; +import { invalidateDbCache } from "../../src/lib/db/readCache.ts"; +import { createAccessToken } from "../../src/lib/db/accessTokens.ts"; +import { resolveAntigravitySettings } from "../../src/lib/resilience/antigravitySettings.ts"; +import { resolvePreferredConnectionUpstreamStartTimeoutMs } from "../../src/sse/services/preferredConnectionFailover.ts"; +import { GET, PATCH } from "../../src/app/api/resilience/route.ts"; + +const url = "http://localhost/api/resilience"; +const preferred = { selectedByApiKeyPreference: true, preferredFallbackAvailable: true }; +test.beforeEach(() => { + getDbInstance().prepare("DELETE FROM key_value WHERE namespace = 'settings'").run(); + invalidateDbCache(); +}); +test.after(() => resetDbInstance()); + +async function patch(antigravity: Record) { + return PATCH(await makeManagementSessionRequest(url, { method: "PATCH", body: { antigravity } })); +} + +test("hot persisted timeout updates affect new request snapshots without a restart", async () => { + const before = resolveAntigravitySettings(await getSettings()); + assert.equal(resolvePreferredConnectionUpstreamStartTimeoutMs("agy", preferred, undefined, "gemini-3.8-flash", before), 12000); + assert.equal((await patch({ preferredTimeoutMs: 7000, gemini38PreferredTimeoutMs: 19000 })).status, 200); + const after = resolveAntigravitySettings(await getSettings()); + assert.equal(resolvePreferredConnectionUpstreamStartTimeoutMs("agy", preferred, undefined, "gemini-3.8-flash-high", after), 19000); + assert.equal(resolvePreferredConnectionUpstreamStartTimeoutMs("agy", preferred, undefined, "claude-sonnet-4-6", after), 7000); + assert.equal(resolvePreferredConnectionUpstreamStartTimeoutMs("agy", preferred, 1000, "gemini-3.8-flash-low", after), 1000); + assert.equal(before.gemini38PreferredTimeoutMs, 12000, "in-flight request snapshot remains stable"); + const response = await GET(await makeManagementSessionRequest(url)); + assert.equal((await response.json()).antigravity.gemini38PreferredTimeoutMs, 19000); + resetDbInstance(); + invalidateDbCache(); + assert.equal(resolveAntigravitySettings(await getSettings()).gemini38PreferredTimeoutMs, 19000, "persists across DB reopen"); +}); + +test("all fields validate before the atomic write; invalid paired values cannot partially persist", async () => { + const before = await getSettingsRevision(); + for (const value of [499, 120001, 12.5, "5000", null]) { + const response = await patch({ preferredTimeoutMs: 8000, gemini38PreferredTimeoutMs: value }); + assert.equal(response.status, 400); + assert.equal(await getSettingsRevision(), before); + assert.equal(resolveAntigravitySettings(await getSettings()).preferredTimeoutMs, 5000); + } + assert.equal((await patch({ preferredTimeoutMs: 500, gemini38PreferredTimeoutMs: 120000 })).status, 200); + assert.equal((await patch({ unknownKnob: 10 })).status, 400); +}); + +test("mutation requires admin-grade credentials even when login is disabled", async () => { + await updateSettings({ requireLogin: false }); + const anonymous = await PATCH(new Request(url, { + method: "PATCH", body: JSON.stringify({ antigravity: { preferredTimeoutMs: 8000 } }), + headers: { "content-type": "application/json" }, + })); + assert.equal(anonymous.status, 401); + const token = createAccessToken({ name: "write-only-fixture", scope: "write" }); + const denied = await PATCH(new Request(url, { + method: "PATCH", body: JSON.stringify({ antigravity: { preferredTimeoutMs: 8000 } }), + headers: { authorization: `Bearer ${token.secret}`, "content-type": "application/json" }, + })); + assert.equal(denied.status, 403); + assert.equal(resolveAntigravitySettings(await getSettings()).preferredTimeoutMs, 5000); + assert.equal(JSON.stringify(await denied.json()).includes("at /"), false); +}); + +test("revision conflict preserves the newer complete settings row", async () => { + const revision = await getSettingsRevision(); + await updateSettings({ resilienceSettings: { antigravity: { preferredTimeoutMs: 9000 } } }); + await assert.rejects(updateSettings({ resilienceSettings: { antigravity: { preferredTimeoutMs: 6000 } } }, { + expectedRevision: revision, + }), SettingsRevisionConflictError); + assert.equal(resolveAntigravitySettings(await getSettings()).preferredTimeoutMs, 9000); +}); + +test("partial updates preserve other timeout and gate values", async () => { + assert.equal((await patch({ preferredTimeoutMs: 6000, adaptiveHealthEnabled: true })).status, 200); + assert.equal((await patch({ timingTelemetryEnabled: true })).status, 200); + const settings = resolveAntigravitySettings(await getSettings()); + assert.deepEqual(settings, { + preferredTimeoutMs: 6000, gemini38PreferredTimeoutMs: 12000, + adaptiveHealthEnabled: true, timingTelemetryEnabled: true, + }); +}); diff --git a/tests/unit/antigravity-adaptive-health.test.ts b/tests/unit/antigravity-adaptive-health.test.ts new file mode 100644 index 000000000000..3a6eba26c8bf --- /dev/null +++ b/tests/unit/antigravity-adaptive-health.test.ts @@ -0,0 +1,148 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { + AntigravityHealthTracker, + ANTIGRAVITY_HEALTH_POLICY, + normalizeAntigravityHealthModel, +} from "../../open-sse/services/antigravityHealth.ts"; +import { + bindAntigravityHealthAttempt, + holdAntigravityHealthThroughResponse, +} from "../../src/sse/services/antigravityHealthLifecycle.ts"; + +const MODEL = "gemini-3.8-flash-high"; +function timeoutRun(tracker: AntigravityHealthTracker, id = "a") { + for (let i = 0; i < 8; i++) tracker.begin(id, MODEL)!.finish("timeout"); +} + +test("normalization shares aliases, without sharing distinct models or accounts", () => { + const tracker = new AntigravityHealthTracker(() => 1000); + assert.equal(normalizeAntigravityHealthModel(" AGY/Gemini-3.8-Flash "), MODEL); + timeoutRun(tracker); + assert.deepEqual(tracker.candidates([{ id: "a" }, { id: "b" }], MODEL), [{ id: "b" }]); + assert.deepEqual(tracker.candidates([{ id: "a" }], "gemini-3.8-flash-low"), [{ id: "a" }]); +}); + +test("low samples never quarantine, cancellations and quota do not enter denominator", () => { + const tracker = new AntigravityHealthTracker(() => 1000); + for (let i = 0; i < 7; i++) tracker.begin("a", MODEL)!.finish("timeout"); + for (let i = 0; i < 20; i++) tracker.begin("a", MODEL)!.finish("ignored"); + assert.equal(tracker.snapshot()[0].samples, 7); + assert.equal(tracker.snapshot()[0].state, "healthy"); + tracker.begin("a", MODEL)!.finish("success"); + assert.equal(tracker.snapshot()[0].state, "quarantined"); +}); + +test("transport failures stay separate from confirmed header timeouts", () => { + const tracker = new AntigravityHealthTracker(() => 1000); + for (let i = 0; i < 12; i++) tracker.begin("a", MODEL)!.finish("transport"); + assert.equal(tracker.snapshot()[0].state, "healthy"); + assert.equal(tracker.snapshot()[0].timeouts, 0); + assert.equal(tracker.snapshot()[0].transportFailures, 12); +}); + +test("window expiry and bounded history prevent stale samples from quarantining", () => { + let now = 1000; + const tracker = new AntigravityHealthTracker(() => now); + for (let i = 0; i < 7; i++) tracker.begin("a", MODEL)!.finish("timeout"); + now += ANTIGRAVITY_HEALTH_POLICY.windowMs + 1; + tracker.begin("a", MODEL)!.finish("timeout"); + assert.equal(tracker.snapshot()[0].samples, 1); + assert.equal(tracker.snapshot()[0].state, "healthy"); + for (let i = 0; i < 50; i++) tracker.begin("b", MODEL)!.finish("success"); + assert.equal(tracker.snapshot("b")[0].samples, ANTIGRAVITY_HEALTH_POLICY.maxSamples); +}); + +test("exactly one half-open probe, two successes to recover, stale successes fenced", async () => { + let now = 1000; + const tracker = new AntigravityHealthTracker(() => now); + const staleSuccess = tracker.begin("a", MODEL)!; + timeoutRun(tracker); + staleSuccess.finish("success"); + assert.equal(tracker.snapshot()[0].state, "quarantined"); + assert.equal(tracker.begin("a", MODEL), null); + now += 30_001; + const tickets = await Promise.all(Array.from({ length: 10 }, async () => tracker.begin("a", MODEL))); + assert.equal(tickets.filter(Boolean).length, 1); + tickets.find(Boolean)!.finish("success"); + assert.equal(tracker.snapshot()[0].state, "half_open"); + assert.equal(tracker.snapshot()[0].recoverySuccesses, 1); + const second = tracker.begin("a", MODEL)!; + second.finish("success"); + second.finish("timeout"); // duplicate result/callback cannot reopen + assert.equal(tracker.snapshot()[0].state, "healthy"); + assert.equal(tracker.snapshot()[0].samples, 0); +}); + +test("all quarantined pool admits one bounded real-request escape, respecting supplied pool", () => { + const tracker = new AntigravityHealthTracker(() => 1000); + timeoutRun(tracker, "a"); + timeoutRun(tracker, "b"); + const candidates = tracker.candidates([{ id: "a" }, { id: "b" }], MODEL); + assert.equal(candidates.length, 1); + const probe = tracker.begin(candidates[0].id, MODEL, true)!; + assert.equal(probe.probe, true); + assert.equal(tracker.begin(candidates[0].id, MODEL, true), null); + assert.deepEqual(tracker.candidates([], MODEL), []); + probe.finish("ignored"); +}); + +test("failed probes increase finite cooldown, capped at five minutes", () => { + let now = 1000; + const tracker = new AntigravityHealthTracker(() => now); + timeoutRun(tracker); + for (let i = 0; i < 10; i++) { + const remaining = tracker.snapshot()[0].cooldownRemainingMs; + assert.ok(remaining > 0 && remaining <= 300_000); + now += remaining + 1; + tracker.begin("a", MODEL)!.finish("timeout"); + } + assert.equal(tracker.snapshot()[0].cooldownRemainingMs, 300_000); +}); + +test("caller abort releases probe without failure or success double count", () => { + const tracker = new AntigravityHealthTracker(() => 1000); + timeoutRun(tracker); + const signal = new AbortController(); + const attempt = bindAntigravityHealthAttempt(tracker.begin("a", MODEL, true)!, signal.signal); + signal.abort(new DOMException("caller cancelled", "AbortError")); + attempt.finish("timeout"); + attempt.finish("success"); + const state = tracker.snapshot()[0]; + assert.equal(state.samples, 8); + assert.equal(state.inflight, 0); + assert.equal(state.backoffLevel, 1); + assert.notEqual(state.state, "probing"); +}); + +test("cancelled response stream never counts as a healthy completed request", async () => { + const tracker = new AntigravityHealthTracker(() => 1000); + const ticket = tracker.begin("a", MODEL)!; + const response = holdAntigravityHealthThroughResponse(new Response(new ReadableStream()), ticket); + await response.body!.cancel(); + ticket.finish("success"); + assert.equal(tracker.snapshot()[0].samples, 0); + assert.equal(tracker.snapshot()[0].inflight, 0); +}); + +test("finished useful SSE counts once, error-only SSE counts zero", async () => { + const tracker = new AntigravityHealthTracker(() => 1000); + for (const [id, data] of [ + ["success", 'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n'], + ["error", 'data: {"error":{"message":"upstream failed"}}\n\n'], + ]) { + const response = holdAntigravityHealthThroughResponse(new Response(data), tracker.begin(id, MODEL)!); + await response.text(); + } + assert.equal(tracker.snapshot("success")[0].samples, 1); + assert.equal(tracker.snapshot("error")[0].samples, 0); +}); + +test("bounded registry never evicts in-flight probes", () => { + const tracker = new AntigravityHealthTracker(() => 1000); + const held = tracker.begin("held", MODEL)!; + for (let i = 0; i < 1100; i++) tracker.begin(`account-${i}`, MODEL)!.finish("success"); + assert.ok(tracker.snapshot().length <= ANTIGRAVITY_HEALTH_POLICY.maxEntries); + assert.equal(tracker.snapshot("held")[0].inflight, 1); + held.finish("success"); +}); diff --git a/tests/unit/antigravity-timeout-retry.test.ts b/tests/unit/antigravity-timeout-retry.test.ts new file mode 100644 index 000000000000..3fd945a70814 --- /dev/null +++ b/tests/unit/antigravity-timeout-retry.test.ts @@ -0,0 +1,77 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { shouldRetrySameAccountTransport } from "../../src/sse/services/sameAccountTransportRetry.ts"; +import { isAntigravityPreResponseTimeout } from "../../open-sse/services/antigravityTimeout.ts"; +import { createUpstreamStartTimeoutError } from "../../open-sse/handlers/chatCore/upstreamTimeouts.ts"; +import { handleAntigravityFallbackChainError } from "../../open-sse/executors/antigravity/proFallbackChain.ts"; +import { AntigravityExecutor } from "../../open-sse/executors/antigravity.ts"; +import { seedAntigravityIdeVersionCache, clearAntigravityVersionCaches } from "../../open-sse/services/antigravityVersion.ts"; + +const code = "ANTIGRAVITY_PRE_RESPONSE_TIMEOUT"; +test("confirmed timeout is never a same-account transport retry, even with reset text", () => { + for (const provider of ["antigravity", "agy", undefined]) { + assert.equal(shouldRetrySameAccountTransport({ + provider, status: 504, errorCode: code, errorType: "upstream_timeout", + errorText: "connection reset", attempt: 0, + }), false); + } + assert.equal(shouldRetrySameAccountTransport({ provider: "agy", status: 504, attempt: 0 }), false); + assert.equal(isAntigravityPreResponseTimeout({ code }), true); + assert.equal(isAntigravityPreResponseTimeout({ status: 504, type: "upstream_timeout" }), false); +}); + +test("outer pre-header timeout carries Antigravity-only marker without tagging other providers", () => { + const antigravity = createUpstreamStartTimeoutError(12000, "antigravity", "gemini-3.8-flash-high"); + assert.equal(isAntigravityPreResponseTimeout(antigravity), true); + assert.equal((antigravity as Error & { code?: string }).code, code); + const other = createUpstreamStartTimeoutError(12000, "codex", "gpt-test"); + assert.equal(isAntigravityPreResponseTimeout(other), false); +}); + +test("transient 502, 503 and socket resets retain bounded same-account retry", () => { + for (const status of [502, 503]) { + assert.equal(shouldRetrySameAccountTransport({ provider: "antigravity", status, attempt: 0 }), true); + assert.equal(shouldRetrySameAccountTransport({ provider: "antigravity", status, attempt: 1 }), false); + } + assert.equal(shouldRetrySameAccountTransport({ status: 500, errorText: "ECONNRESET", attempt: 0 }), true); + for (const status of [499, 429, 401, 400]) { + assert.equal(shouldRetrySameAccountTransport({ status, attempt: 0 }), false); + } +}); + +test("forced connections and emitted output prohibit retries", () => { + assert.equal(shouldRetrySameAccountTransport({ status: 503, attempt: 0, hasForcedConnection: true }), false); + assert.equal(shouldRetrySameAccountTransport({ status: 502, attempt: 0, hasEmittedOutput: true }), false); +}); + +test("Pro candidate fallback preserves the exact timeout instead of hiding or replaying it", () => { + const error = Object.assign(new Error("headers timeout"), { code }); + const outcome = handleAntigravityFallbackChainError({ + model: "gemini-3.1-pro-low", body: {}, stream: true, credentials: {}, + }, error, "gemini-3.1-pro-low", 0, ["gemini-3.1-pro-low", "gemini-3-pro-low"], null, "gemini-3.1-pro-low"); + assert.equal(outcome.action, "throw"); + if (outcome.action === "throw") assert.equal(outcome.error, error); +}); + +test("executor sends a confirmed header timeout only once across all URLs and Pro ids", async () => { + seedAntigravityIdeVersionCache("2.1.1"); + const original = globalThis.fetch; + let sends = 0; + globalThis.fetch = async () => { + sends++; + throw Object.assign(new Error("headers timeout"), { name: "TimeoutError", code }); + }; + try { + const executor = new AntigravityExecutor(); + await assert.rejects(executor.execute({ + model: "gemini-3.1-pro-low", body: { messages: [{ role: "user", content: "fixture" }] }, + stream: true, credentials: { accessToken: "fixture", projectId: "fixture-project" }, + }), (error: unknown) => isAntigravityPreResponseTimeout(error)); + assert.equal(sends, 1); + assert.equal(executor.shouldRetry(504, 0), false); + assert.equal(executor.shouldRetry(503, 0), executor.getFallbackCount() > 1); + } finally { + globalThis.fetch = original; + clearAntigravityVersionCaches(); + } +}); diff --git a/tests/unit/antigravity-timing.test.ts b/tests/unit/antigravity-timing.test.ts new file mode 100644 index 000000000000..c1bba574c3ee --- /dev/null +++ b/tests/unit/antigravity-timing.test.ts @@ -0,0 +1,102 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { AntigravityTimingTracker, withAntigravityTiming, hasAntigravityMeaningfulDelta } from "../../open-sse/services/antigravityTiming.ts"; + +const text = 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"fixture"}]}}]}}\n\n'; +const encoder = new TextEncoder(); +const requestId = "24a00000-0000-4000-a000-000000000024"; + +test("pings, roles, usage and errors are not meaningful deltas", () => { + for (const frame of [": ping\n\n", 'data: {"usageMetadata":{"totalTokenCount":1}}\n\n', + 'data: {"response":{"candidates":[{"content":{"role":"model","parts":[]}}]}}\n\n', + 'data: {"error":{"message":"failure"}}\n\n', "data: [DONE]\n\n"]) { + assert.equal(hasAntigravityMeaningfulDelta(frame), false); + } + assert.equal(hasAntigravityMeaningfulDelta(text), true); + assert.equal(hasAntigravityMeaningfulDelta('data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"tool"}}]}}]}\n\n'), true); +}); + +test("measures actual header and first delta phases with null unavailable network phases", async () => { + let now = 100; + const events: Record[] = []; + const tracker = new AntigravityTimingTracker(() => now); + let upstream: ReadableStreamDefaultController; + const response = await withAntigravityTiming({ + enabled: true, connectionId: "fixture-account", model: "agy/gemini-3.8-flash", + correlationId: requestId, emit: (event) => events.push(event), + }, () => tracker.fetch(async () => { + now = 125; + return new Response(new ReadableStream({ start(controller) { upstream = controller; } })); + })); + assert.equal(events[0].headerLatencyMs, 25); + assert.equal(events[0].idleDurationMs, null); + assert.equal(events[0].dnsMs, null); + assert.equal(events[0].connectMs, null); + assert.equal(events[0].tlsMs, null); + now = 140; + upstream!.enqueue(encoder.encode(": ping\n\n")); + const reading = response.text(); + now = 170; + upstream!.enqueue(encoder.encode(text.slice(0, 20))); + upstream!.enqueue(encoder.encode(text.slice(20))); + upstream!.close(); + assert.equal(await reading, ": ping\n\n" + text); + const first = events.find((event) => event.phase === "first_meaningful_delta")!; + assert.equal(first.firstMeaningfulDeltaFromHeadersMs, 45); + assert.equal(first.firstMeaningfulDeltaFromSendMs, 70); + assert.equal(first.requestId, requestId); + assert.equal(first.model, "gemini-3.8-flash-high"); + assert.equal(events.filter((event) => event.phase === "complete").length, 1); + const serialized = JSON.stringify(events); + for (const forbidden of ["fixture-account", "fixture\"", "Bearer", "https://", "parts", "candidates"]) { + assert.equal(serialized.includes(forbidden), false); + } +}); + +test("idle duration and inflight count are per account and end on actual lifecycle", async () => { + let now = 10; + const tracker = new AntigravityTimingTracker(() => now); + const events: Record[] = []; + const send = (id = "a") => withAntigravityTiming({ + enabled: true, connectionId: id, model: "gemini-3.8-flash-high", emit: (event) => events.push(event), + }, () => tracker.fetch(async () => new Response(new ReadableStream()))); + const first = await send(); + now = 20; + const concurrent = await send(); + assert.equal(events[1].concurrentInflight, 2); + assert.equal(events[1].idleDurationMs, 0); + await first.body!.cancel(); + await concurrent.body!.cancel(); + now = 90; + const afterIdle = await send(); + assert.equal(events.findLast((event) => event.phase === "headers")!.idleDurationMs, 70); + const other = await send("b"); + assert.equal(events.findLast((event) => event.phase === "headers")!.idleDurationMs, null); + await afterIdle.body!.cancel(); + await other.body!.cancel(); +}); + +test("feature off does not wrap or consume the response, and logs nothing", async () => { + const tracker = new AntigravityTimingTracker(); + const original = new Response(text); + const events: unknown[] = []; + const response = await withAntigravityTiming({ + enabled: false, connectionId: "a", model: "m", emit: (event) => events.push(event), + }, () => tracker.fetch(async () => original)); + assert.equal(response, original); + assert.equal(response.bodyUsed, false); + assert.deepEqual(events, []); +}); + +test("failed fetch never fabricates header or first delta latency", async () => { + let now = 0; + const tracker = new AntigravityTimingTracker(() => now); + const events: Record[] = []; + await assert.rejects(withAntigravityTiming({ + enabled: true, connectionId: "a", model: "m", correlationId: "Bearer secret", emit: (event) => events.push(event), + }, () => tracker.fetch(async () => { now = 500; throw new Error("secret upstream URL"); }))); + assert.equal(events[0].headerLatencyMs, null); + assert.equal(events[0].firstMeaningfulDeltaFromSendMs, null); + assert.equal(events[0].elapsedMs, 500); + assert.equal(JSON.stringify(events).includes("secret"), false); +});