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.
2026-07-31 04:12:10 +03:00 · 2026-06-28 19:05:27 +03:00
parent 6143ffc021
commit 6f150cfd2c
2 changed files with 249 additions and 13 deletions
--- 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<string>, 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<string, number>();
@@ -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<string, unknown>).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<string, KeyHealth>
      | 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,
--- a/tests/unit/auth-antigravity-account-retry-v2.test.ts
+++ 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);
+});