From 6f150cfd2c00eeebbd3036fd2c9591f1c0d0116a Mon Sep 17 00:00:00 2001 From: Ardem2025 Date: Sun, 28 Jun 2026 19:05:27 +0300 Subject: [PATCH] fix(antigravity): retry excluded accounts via fallback LRU (#5222) Antigravity: retry excluded accounts via fallback LRU + family-inferred 429 cooldown. Test strengthened into a real LRU regression guard; cooldown constant extracted. Integrated into release/v3.8.40. --- src/sse/services/auth.ts | 98 +++++++++-- .../auth-antigravity-account-retry-v2.test.ts | 164 ++++++++++++++++++ 2 files changed, 249 insertions(+), 13 deletions(-) create mode 100644 tests/unit/auth-antigravity-account-retry-v2.test.ts diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 9ca5f83b74..26bf73f61a 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -17,6 +17,7 @@ import { getQuotaWindowStatus, isAccountQuotaExhausted, } from "@/domain/quotaCache"; +import { getQuotaScopeLabelForProvider } from "@omniroute/open-sse/services/antigravityQuotaFamily.ts"; import { isAccountUnavailable, getUnavailableUntil, @@ -125,6 +126,10 @@ interface CooldownInspectionState { const MIN_QUOTA_THRESHOLD_PERCENT = 1; const MAX_QUOTA_THRESHOLD_PERCENT = 100; const NON_RETRYABLE_MODEL_LOCKOUT_REASONS = new Set(["not_found", "not_found_local"]); +// Antigravity Gemini family 429 with no parseable upstream hint: seed the backoff at +// this base. Real upstream Retry-After hints still win — they flow through +// `exactCooldownMs` (usedUpstreamRetryHint), not this base. (#5222) +const ANTIGRAVITY_FAMILY_INFERRED_BASE_COOLDOWN_MS = 30_000; function asRecord(value: unknown): JsonRecord { return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; @@ -908,6 +913,14 @@ function normalizeExcludedConnectionIds( return normalized; } +function formatConnectionPrefixesForLog(ids: Iterable, max = 6): string { + const prefixes = Array.from(ids) + .filter((id) => typeof id === "string" && id.length > 0) + .slice(0, max) + .map((id) => `${id.slice(0, 8)}...`); + return prefixes.length > 0 ? prefixes.join(",") : "none"; +} + function buildQuotaPreflightRateLimitedResult( provider: string, blockedByPreflight: Array<{ @@ -1097,12 +1110,29 @@ export async function getProviderCredentials( if (forcedConnectionId) { connections = connections.filter((conn) => conn.id === forcedConnectionId); } + const activeConnectionsCount = connections.length; + const rawConnectionsCount = connectionsRaw.length; + const blockedByForcedConnection = forcedConnectionId + ? rawConnectionsCount - connections.length + : 0; + const blockedByAllowedConnections = + allowedConnections && allowedConnections.length > 0 + ? Math.max(0, rawConnectionsCount - connections.length - blockedByForcedConnection) + : 0; + const forcedIdForLog = forcedConnectionId ? `${forcedConnectionId.slice(0, 8)}...` : "none"; + log.debug( "AUTH", - `${provider} | total connections: ${connections.length}, excludeIds: ${ - excludedConnectionIds.size > 0 ? Array.from(excludedConnectionIds).join(",") : "none" - }, forcedId: ${forcedConnectionId || "none"}` + `${provider} | active=${activeConnectionsCount}, excluded=${excludedConnectionIds.size} (${formatConnectionPrefixesForLog(excludedConnectionIds)}), forcedId=${forcedIdForLog}, blocked_forced=${blockedByForcedConnection}, blocked_allowed=${blockedByAllowedConnections}` ); + if (provider === "antigravity" && (forcedConnectionId || allowedConnections?.length)) { + const reasons: string[] = []; + if (forcedConnectionId) reasons.push(`forcedConnectionId kept ${connections.length}`); + if (allowedConnections?.length) { + reasons.push(`allowedConnections=${allowedConnections.length}`); + } + log.info("AUTH", `${provider} selection constrained: ${reasons.join(", ")}`); + } if (connections.length === 0) { // Check all connections (including inactive) to see if rate limited @@ -1149,7 +1179,10 @@ export async function getProviderCredentials( // the dashboard sees a misleading "bad_request" code. const terminalConnections = allConnections.filter(isTerminalConnectionStatus); if (terminalConnections.length === allConnections.length) { - const syntheticFallback = await maybeSyntheticNoAuthFallback(resolvedId, excludedConnectionIds); + const syntheticFallback = await maybeSyntheticNoAuthFallback( + resolvedId, + excludedConnectionIds + ); if (syntheticFallback) return syntheticFallback; const statusCounts = new Map(); @@ -1166,7 +1199,10 @@ export async function getProviderCredentials( }; } } - const syntheticFallback = await maybeSyntheticNoAuthFallback(resolvedId, excludedConnectionIds); + const syntheticFallback = await maybeSyntheticNoAuthFallback( + resolvedId, + excludedConnectionIds + ); if (syntheticFallback) return syntheticFallback; log.warn("AUTH", `No credentials for ${provider}`); return null; @@ -1195,6 +1231,8 @@ export async function getProviderCredentials( } } + let modelLockedCount = 0; + let familyLockedCount = 0; // Filter out unavailable accounts and excluded connection const availableConnections = connections.filter((c) => { if (excludedConnectionIds.has(c.id)) return false; @@ -1205,8 +1243,18 @@ export async function getProviderCredentials( if (!allowRateLimitedConnections && isAccountUnavailable(c.rateLimitedUntil)) return false; if (isTerminalConnectionStatus(c)) return false; if (provider === "codex" && isCodexScopeUnavailable(c, requestedModel)) return false; - // Per-model lockout: if this specific model is locked on this connection, skip it - if (requestedModel && isModelLocked(provider, c.id, requestedModel)) return false; + // Per-model lockout: if this specific model/family is locked on this connection, skip it + if (requestedModel && isModelLocked(provider, c.id, requestedModel)) { + if ( + provider === "antigravity" && + getQuotaScopeLabelForProvider(provider, requestedModel) === "family" + ) { + familyLockedCount += 1; + } else { + modelLockedCount += 1; + } + return false; + } } return true; }); @@ -1215,6 +1263,12 @@ export async function getProviderCredentials( "AUTH", `${provider} | available: ${availableConnections.length}/${connections.length}` ); + if (provider === "antigravity") { + log.info( + "AUTH", + `${provider} selection candidates model=${requestedModel || "none"}: active=${activeConnectionsCount}, excluded=${excludedConnectionIds.size}, modelLocked=${modelLockedCount}, familyLocked=${familyLockedCount}, eligible=${availableConnections.length}` + ); + } connections.forEach((c) => { const excluded = excludedConnectionIds.has(c.id); const rateLimited = isAccountUnavailable(c.rateLimitedUntil); @@ -1337,7 +1391,10 @@ export async function getProviderCredentials( cooldownModel: allBlockedByModelCooldown ? requestedModel : null, }; } - const syntheticFallback = await maybeSyntheticNoAuthFallback(resolvedId, excludedConnectionIds); + const syntheticFallback = await maybeSyntheticNoAuthFallback( + resolvedId, + excludedConnectionIds + ); if (syntheticFallback) return syntheticFallback; log.warn("AUTH", `${provider} | all ${connections.length} accounts unavailable`); return null; @@ -1461,9 +1518,9 @@ export async function getProviderCredentials( } else if (strategy === "round-robin") { const stickyLimit = toNumber((settings as Record).stickyRoundRobinLimit, 3); - // If excluding an account (fallback scenario), skip sticky logic and go straight to LRU - // This prevents the system from getting stuck on a failed account - const isFallbackScenario = excludeConnectionId !== null; + // If excluding account(s) (fallback scenario), skip sticky logic and go straight to LRU. + // This prevents same-model retries from getting stuck on a failed account. + const isFallbackScenario = excludeConnectionId !== null || excludedConnectionIds.size > 0; if (!isFallbackScenario) { // Sort by lastUsed (most recent first) to find current candidate @@ -1533,7 +1590,7 @@ export async function getProviderCredentials( connection = sortedByOldest[0]; log.info( "AUTH", - `${provider} round-robin: FALLBACK MODE - excluded ${excludeConnectionId?.slice(0, 8)}..., picked LRU ${connection.id?.slice(0, 8)}...` + `${provider} round-robin: FALLBACK MODE - excluded_count=${excludedConnectionIds.size} excluded=${formatConnectionPrefixesForLog(excludedConnectionIds)} picked_lru=${connection.id?.slice(0, 8)}...` ); // Update lastUsedAt and reset count to 1 (await to ensure persistence) @@ -1591,6 +1648,13 @@ export async function getProviderCredentials( connection = orderedConnections[0]; } + if (provider === "antigravity" && connection) { + log.info( + "AUTH", + `${provider} selected account=${connection.id?.slice(0, 8)}... eligible=${orderedConnections.length} excluded=${excludedConnectionIds.size}` + ); + } + const apiKeyHealth = connection.providerSpecificData?.apiKeyHealth as | Record | undefined; @@ -1932,6 +1996,11 @@ export async function markAccountUnavailable( : status === 429 ? "rate_limited" : "server_error"; + const quotaScope = getQuotaScopeLabelForProvider(provider, model); + const antigravityFamilyInferredBaseCooldownMs = + provider === "antigravity" && quotaScope === "family" && status === 429 + ? ANTIGRAVITY_FAMILY_INFERRED_BASE_COOLDOWN_MS + : null; const lockout = recordModelLockoutFailure( provider, connectionId, @@ -1940,7 +2009,10 @@ export async function markAccountUnavailable( status, status === 404 ? (effectiveProviderProfile?.baseCooldownMs ?? COOLDOWN_MS.notFoundLocal) - : (fallbackResult.baseCooldownMs ?? effectiveProviderProfile?.baseCooldownMs ?? 0), + : (antigravityFamilyInferredBaseCooldownMs ?? + fallbackResult.baseCooldownMs ?? + effectiveProviderProfile?.baseCooldownMs ?? + 0), effectiveProviderProfile, { ...modelLockoutOptions, diff --git a/tests/unit/auth-antigravity-account-retry-v2.test.ts b/tests/unit/auth-antigravity-account-retry-v2.test.ts new file mode 100644 index 0000000000..e26a0625cc --- /dev/null +++ b/tests/unit/auth-antigravity-account-retry-v2.test.ts @@ -0,0 +1,164 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-auth-ag-retry-v2-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const settingsDb = await import("../../src/lib/db/settings.ts"); +const auth = await import("../../src/sse/services/auth.ts"); + +type CreatedConnection = { id: string }; + +function connectionId(connection: unknown): string { + assert.ok(connection && typeof connection === "object" && "id" in connection); + const id = (connection as CreatedConnection).id; + assert.equal(typeof id, "string"); + return id; +} + +async function resetStorage() { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.after(() => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("round-robin same-model retry treats multi-exclude as fallback LRU and skips all excluded accounts", async () => { + await resetStorage(); + + const first = await providersDb.createProviderConnection({ + provider: "antigravity", + authType: "oauth", + email: "first@example.test", + accessToken: "tok-first", + isActive: true, + testStatus: "active", + priority: 1, + }); + const second = await providersDb.createProviderConnection({ + provider: "antigravity", + authType: "oauth", + email: "second@example.test", + accessToken: "tok-second", + isActive: true, + testStatus: "active", + priority: 2, + }); + // Two NON-excluded eligible accounts with diverging lastUsedAt so sticky + // (most-recently-used) and fallback LRU (least-recently-used) pick different + // accounts — this is what makes the test discriminate the + // `excludedConnectionIds.size > 0` fallback branch. `recent` is the sticky + // pick; `stale` is the LRU pick. + const recent = await providersDb.createProviderConnection({ + provider: "antigravity", + authType: "oauth", + email: "recent@example.test", + accessToken: "tok-recent", + isActive: true, + testStatus: "active", + priority: 3, + }); + const stale = await providersDb.createProviderConnection({ + provider: "antigravity", + authType: "oauth", + email: "stale@example.test", + accessToken: "tok-stale", + isActive: true, + testStatus: "active", + priority: 4, + }); + + const firstId = connectionId(first); + const secondId = connectionId(second); + const recentId = connectionId(recent); + const staleId = connectionId(stale); + + await providersDb.updateProviderConnection(firstId, { + lastUsedAt: new Date(Date.now() - 5_000).toISOString(), + consecutiveUseCount: 1, + }); + await providersDb.updateProviderConnection(secondId, { + lastUsedAt: new Date(Date.now() - 4_000).toISOString(), + consecutiveUseCount: 1, + }); + // `recent` is the most-recently-used eligible account: WITHOUT the fallback + // change, sticky routing (count 1 < limit 3) would stay on it. + await providersDb.updateProviderConnection(recentId, { + lastUsedAt: new Date(Date.now() - 1_000).toISOString(), + consecutiveUseCount: 1, + }); + // `stale` is the least-recently-used eligible account: the fallback LRU branch + // must pick this one. + await providersDb.updateProviderConnection(staleId, { + lastUsedAt: new Date(Date.now() - 90_000).toISOString(), + consecutiveUseCount: 1, + }); + + await settingsDb.updateSettings({ fallbackStrategy: "round-robin", stickyRoundRobinLimit: 3 }); + + const selected = await auth.getProviderCredentials("antigravity", null, null, "gemini-3-pro", { + excludeConnectionIds: [firstId, secondId], + }); + + assert.ok(selected, "expected an eligible non-excluded Antigravity account"); + // Excluded accounts must never be selected. + assert.notEqual(selected.connectionId, firstId); + assert.notEqual(selected.connectionId, secondId); + // The fallback (excludedConnectionIds.size > 0) must route to the LRU account + // (`stale`), NOT the sticky most-recently-used one (`recent`). Without the + // `excludedConnectionIds.size > 0` fallback trigger this assertion gets + // `recentId` and fails. + assert.equal(selected.connectionId, staleId); +}); + +test("Antigravity inferred Gemini family cooldown starts around 30s when no upstream hint exists", async () => { + await resetStorage(); + + const conn = await providersDb.createProviderConnection({ + provider: "antigravity", + authType: "oauth", + email: "quota@example.test", + accessToken: "tok-quota", + isActive: true, + testStatus: "active", + }); + + const before = Date.now(); + const result = await auth.markAccountUnavailable( + connectionId(conn), + 429, + "RESOURCE_EXHAUSTED: Resource has been exhausted (queries per minute limit was reached)", + "antigravity", + "gemini-3-pro" + ); + const elapsedAllowanceMs = Date.now() - before; + + assert.equal(result.shouldFallback, true); + assert.ok( + result.cooldownMs >= 30_000 - elapsedAllowanceMs - 500, + `expected inferred cooldown near 30s+, got ${result.cooldownMs}` + ); + assert.ok( + result.cooldownMs <= 65_000, + `expected bounded initial cooldown, got ${result.cooldownMs}` + ); + + const otherGemini = await auth.getProviderCredentials( + "antigravity", + null, + null, + "gemini-2.5-pro" + ); + assert.ok(otherGemini && "allRateLimited" in otherGemini && otherGemini.allRateLimited); + assert.equal(otherGemini.cooldownScope, "model"); + assert.equal(otherGemini.lastErrorCode, 429); +});