From 1a3db1efae5c758483a7af9e3cfb01a7a64015df Mon Sep 17 00:00:00 2001 From: Nim Date: Thu, 27 Aug 2026 16:14:15 +0700 Subject: [PATCH] feat(antigravity): quota-aware routing with reset-aware fallback On a 409/429 from Antigravity, fetch live quota to learn the exact per-model resetAt instead of guessing a backoff, then skip only the exhausted account/model pair until that time. - antigravityQuota.js: in-memory quota cache, coalesced concurrent refreshes, 30s throttle per connection (applied to failures too), keeps known cache when upstream returns 401/403 error payloads - auth.js: pre-filter exhausted account/model pairs; report the earliest quota reset when every account is blocked; skip the 30-minute cooldown cap so the upstream resetAt is not truncated - chat.js: antigravity 409/429 falls back on the RAM cache only, no persistent modelLock_* for this path - Logs identify accounts by id prefix, never email or name Closes #3561 --- src/sse/handlers/chat.js | 19 +- src/sse/services/antigravityQuota.js | 99 +++++++++++ src/sse/services/auth.js | 29 +++- tests/unit/antigravity-quota-routing.test.js | 172 +++++++++++++++++++ 4 files changed, 314 insertions(+), 5 deletions(-) create mode 100644 src/sse/services/antigravityQuota.js create mode 100644 tests/unit/antigravity-quota-routing.test.js diff --git a/src/sse/handlers/chat.js b/src/sse/handlers/chat.js index 88353c50..dc7f98b1 100644 --- a/src/sse/handlers/chat.js +++ b/src/sse/handlers/chat.js @@ -7,6 +7,7 @@ import { extractApiKey, isValidApiKey, } from "../services/auth.js"; +import { handleAntigravityQuotaError } from "../services/antigravityQuota.js"; import { getSettings } from "@/lib/localDb"; import { getModelInfo, getComboModels } from "../services/model.js"; import { handleChatCore } from "open-sse/handlers/chatCore.js"; @@ -300,8 +301,22 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re if (result.success) return result.response; - // Mark account unavailable (auto-calculates cooldown with exponential backoff, or precise resetsAtMs) - const { shouldFallback } = await markAccountUnavailable(credentials.connectionId, result.status, result.error, provider, model, result.resetsAtMs); + // Antigravity 409/429: refresh live quota to get exact resetAt before locking + let quotaResetMs = null; + let resetsAtMs = result.resetsAtMs; + if (provider === "antigravity" && (result.status === 409 || result.status === 429)) { + quotaResetMs = await handleAntigravityQuotaError( + credentials.connectionId, result.status, model, + refreshedCredentials.accessToken, credentials.providerSpecificData + ); + if (quotaResetMs) resetsAtMs = quotaResetMs; + } + + // Exhausted Antigravity model is blocked only in RAM cache until upstream resetAt. + // Do not persist a modelLock_* for this path. + const shouldFallback = provider === "antigravity" && quotaResetMs + ? true + : (await markAccountUnavailable(credentials.connectionId, result.status, result.error, provider, model, resetsAtMs)).shouldFallback; if (shouldFallback) { log.warn("FALLBACK", `⇄ ACC:${credentials.connectionName} UNAVAILABLE (${result.status}) → NEXT ACCOUNT`); diff --git a/src/sse/services/antigravityQuota.js b/src/sse/services/antigravityQuota.js new file mode 100644 index 00000000..e991614f --- /dev/null +++ b/src/sse/services/antigravityQuota.js @@ -0,0 +1,99 @@ +/** + * Antigravity live quota cache — in-memory, refreshed on demand. + * Used by auth.js pre-filter to skip accounts with exhausted model quota. + * Also triggered by 409/429 error handler to sync exact resetAt from upstream. + */ + +import { resolveConnectionProxyConfig } from "@/lib/network/connectionProxy"; +import { getAntigravityUsage } from "open-sse/services/usage/google.js"; +import * as log from "../utils/logger.js"; + +// In-memory cache: connectionId → { [modelId]: { remainingPercentage, resetAt } } +const quotaCache = new Map(); +// Track last refresh per connection to avoid hammering +const lastRefreshAt = new Map(); +// In-flight refresh promises — dedup concurrent 409/429 bursts +const inflightRefresh = new Map(); + +const MIN_REFRESH_INTERVAL_MS = 30_000; // 30s between refreshes per connection + +/** + * Get the quota cache (read-only reference for auth.js pre-filter). + */ +export function getAntigravityQuotaCache() { + return quotaCache; +} + +/** + * Refresh quota for a single antigravity connection from upstream API. + * Updates in-memory cache only. Cache expiry is the upstream model resetAt. + * @returns {object|null} quotas map or null on failure + */ +export async function refreshAntigravityQuota(connectionId, accessToken, providerSpecificData) { + const now = Date.now(); + // Coalesce concurrent refreshes before applying the interval gate. + const inflight = inflightRefresh.get(connectionId); + if (inflight) return inflight; + + const lastRefresh = lastRefreshAt.get(connectionId) || 0; + if (now - lastRefresh < MIN_REFRESH_INTERVAL_MS) { + log.debug("AG_QUOTA", `${connectionId.slice(0, 8)} | skip refresh (${Math.round((now - lastRefresh) / 1000)}s ago)`); + return quotaCache.get(connectionId) || null; + } + + // Record every attempt so failed quota calls cannot amplify an upstream 429 burst. + lastRefreshAt.set(connectionId, now); + const promise = _doRefresh(connectionId, accessToken, providerSpecificData, now); + inflightRefresh.set(connectionId, promise); + try { + return await promise; + } finally { + inflightRefresh.delete(connectionId); + } +} + +async function _doRefresh(connectionId, accessToken, providerSpecificData, now) { + try { + const proxyCfg = await resolveConnectionProxyConfig(providerSpecificData || {}); + const proxyOptions = { + connectionProxyEnabled: proxyCfg.connectionProxyEnabled === true, + connectionProxyUrl: proxyCfg.connectionProxyUrl || "", + connectionNoProxy: proxyCfg.connectionNoProxy || "", + vercelRelayUrl: proxyCfg.vercelRelayUrl || "", + strictProxy: proxyCfg.strictProxy === true, + }; + + const usage = await getAntigravityUsage(accessToken, providerSpecificData, proxyOptions); + // 401/403 usage responses can contain an empty quotas object plus message. + // Preserve known cache instead of replacing it with an upstream error response. + if (!usage?.quotas || usage.message) return null; + + // Update in-memory cache. Caller logs CACHE_BLOCK only if requested model is exhausted. + quotaCache.set(connectionId, usage.quotas); + + return usage.quotas; + } catch (e) { + log.warn("AG_QUOTA", `${connectionId.slice(0, 8)} | refresh failed: ${e.message}`); + return null; + } +} + +/** + * Handle Antigravity 409/429 — refresh RAM cache and return model resetAt when exhausted. + * Called from chat handler error path. + * @returns {number|null} resetAt timestamp ms (for resetsAtMs passthrough) or null + */ +export async function handleAntigravityQuotaError(connectionId, status, model, accessToken, providerSpecificData) { + log.info("AG_QUOTA", `${connectionId.slice(0, 8)} | ${status} on ${model} — refreshing quota`); + + // Throttle applies to error paths too: one quota request per account/30s. + // The first 409/429 populates cache; concurrent or repeated errors reuse it. + const quota = (await refreshAntigravityQuota(connectionId, accessToken, providerSpecificData))?.[model]; + if (!quota || quota.remainingPercentage > 0 || !quota.resetAt) return null; + + const resetMs = new Date(quota.resetAt).getTime(); + if (resetMs <= Date.now()) return null; + + log.warn("AG_QUOTA", `${connectionId.slice(0, 8)} | UPSTREAM_${status} ${model} — quota exhausted; CACHE_BLOCK until ${quota.resetAt}`); + return resetMs; +} diff --git a/src/sse/services/auth.js b/src/sse/services/auth.js index feaaa2ab..d85c4901 100644 --- a/src/sse/services/auth.js +++ b/src/sse/services/auth.js @@ -3,6 +3,7 @@ import { resolveConnectionProxyConfig, pickProxyPoolId } from "@/lib/network/con import { formatRetryAfter, checkFallbackError, isModelLockActive, buildModelLockUpdate, getEarliestModelLockUntil } from "open-sse/services/accountFallback.js"; import { MAX_RATE_LIMIT_COOLDOWN_MS } from "open-sse/config/errorConfig.js"; import { resolveProviderId, FREE_PROVIDERS } from "@/shared/constants/providers.js"; +import { getAntigravityQuotaCache } from "./antigravityQuota.js"; import * as log from "../utils/logger.js"; // Mutex to prevent race conditions during account selection @@ -76,10 +77,23 @@ export async function getProviderCredentials(provider, excludeConnectionIds = nu return null; } - // Filter out model-locked and excluded connections + // Antigravity quota cache is lazy: only populated after that account returns 409/429. + const isAntigravity = providerId === "antigravity"; + const antigravityQuotaCache = isAntigravity && model ? getAntigravityQuotaCache() : null; + + // Filter out model-locked, excluded, and Antigravity quota-exhausted connections. const availableConnections = connections.filter(c => { if (excludeSet.has(c.id)) return false; if (isModelLockActive(c, model)) return false; + // Antigravity: skip if live quota exhausted for this model + if (isAntigravity && model && antigravityQuotaCache) { + const quota = antigravityQuotaCache.get(c.id)?.[model]; + if (quota && quota.remainingPercentage <= 0 && quota.resetAt && new Date(quota.resetAt).getTime() > Date.now()) { + const account = c.id?.slice(0, 8) || "unknown"; + log.info("AG_QUOTA", `${account} | CACHE_BLOCK ${model} — skip upstream until ${quota.resetAt}`); + return false; + } + } return true; }); @@ -94,9 +108,15 @@ export async function getProviderCredentials(provider, excludeConnectionIds = nu }); if (availableConnections.length === 0) { - // Find earliest lock expiry across all connections for retry timing + // Find earliest persistent lock or lazy Antigravity quota-cache reset for retry timing. const lockedConns = connections.filter(c => isModelLockActive(c, model)); const expiries = lockedConns.map(c => getEarliestModelLockUntil(c)).filter(Boolean); + if (isAntigravity && model && antigravityQuotaCache) { + connections.forEach((c) => { + const resetAt = antigravityQuotaCache.get(c.id)?.[model]?.resetAt; + if (resetAt && new Date(resetAt).getTime() > Date.now()) expiries.push(resetAt); + }); + } const earliest = expiries.sort()[0] || null; if (earliest) { const earliestConn = lockedConns[0]; @@ -233,7 +253,10 @@ export async function markAccountUnavailable(connectionId, status, errorText, pr newBackoffLevel = 0; } else if (resetsAtMs && resetsAtMs > Date.now()) { shouldFallback = true; - cooldownMs = Math.min(resetsAtMs - Date.now(), MAX_RATE_LIMIT_COOLDOWN_MS); + // Antigravity quota API provides exact per-model resetAt. Do not truncate it. + cooldownMs = resolveProviderId(provider) === "antigravity" + ? resetsAtMs - Date.now() + : Math.min(resetsAtMs - Date.now(), MAX_RATE_LIMIT_COOLDOWN_MS); newBackoffLevel = 0; } else { ({ shouldFallback, cooldownMs, newBackoffLevel } = checkFallbackError(status, errorText, backoffLevel)); diff --git a/tests/unit/antigravity-quota-routing.test.js b/tests/unit/antigravity-quota-routing.test.js new file mode 100644 index 00000000..c7778a3c --- /dev/null +++ b/tests/unit/antigravity-quota-routing.test.js @@ -0,0 +1,172 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + getProviderConnections: vi.fn(), + getSettings: vi.fn(), + resolveConnectionProxyConfig: vi.fn(), + getAntigravityUsage: vi.fn(), +})); + +vi.mock("@/lib/localDb", () => ({ + getProviderConnections: mocks.getProviderConnections, + getSettings: mocks.getSettings, + getProxyPools: vi.fn(), + validateApiKey: vi.fn(), + updateProviderConnection: vi.fn(), +})); +vi.mock("@/lib/network/connectionProxy", () => ({ + resolveConnectionProxyConfig: mocks.resolveConnectionProxyConfig, + pickProxyPoolId: vi.fn(), +})); +vi.mock("@/shared/constants/providers.js", () => ({ + FREE_PROVIDERS: {}, + resolveProviderId: (provider) => provider, +})); +vi.mock("open-sse/services/usage/google.js", () => ({ + getAntigravityUsage: mocks.getAntigravityUsage, +})); +vi.mock("@/sse/utils/logger.js", () => ({ debug: vi.fn(), info: vi.fn(), warn: vi.fn() })); + +const { getAntigravityQuotaCache, handleAntigravityQuotaError, refreshAntigravityQuota } = await import("@/sse/services/antigravityQuota.js"); +const { getProviderCredentials } = await import("@/sse/services/auth.js"); + +const MODEL = "claude-opus-4-6-thinking"; +const FUTURE_RESET = "2026-09-01T00:00:00.000Z"; + +beforeEach(() => { + vi.clearAllMocks(); + getAntigravityQuotaCache().clear(); + mocks.resolveConnectionProxyConfig.mockResolvedValue({}); + mocks.getSettings.mockResolvedValue({}); +}); + +describe("Antigravity quota-aware routing", () => { + it("records exhausted upstream quota after 429 and returns its exact reset time", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + mocks.getAntigravityUsage.mockResolvedValue({ quotas: { + [MODEL]: { remainingPercentage: 0, resetAt: FUTURE_RESET }, + } }); + + try { + await expect(handleAntigravityQuotaError("ag-a", 429, MODEL, "token", {})) + .resolves.toBe(Date.parse(FUTURE_RESET)); + expect(getAntigravityQuotaCache().get("ag-a")[MODEL]).toEqual({ + remainingPercentage: 0, + resetAt: FUTURE_RESET, + }); + } finally { + vi.useRealTimers(); + } + }); + + it("skips exhausted account/model and selects the next account", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + mocks.getProviderConnections.mockResolvedValue([ + { id: "ag-a", email: "a@example.com", isActive: true }, + { id: "ag-b", email: "b@example.com", isActive: true }, + ]); + getAntigravityQuotaCache().set("ag-a", { + [MODEL]: { remainingPercentage: 0, resetAt: FUTURE_RESET }, + }); + + try { + await expect(getProviderCredentials("antigravity", null, MODEL)).resolves.toMatchObject({ + connectionId: "ag-b", + connectionName: "b@example.com", + }); + } finally { + vi.useRealTimers(); + } + }); + + it("reports retry time when every account is cache-blocked", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + mocks.getProviderConnections.mockResolvedValue([{ id: "ag-a", email: "a@example.com", isActive: true }]); + getAntigravityQuotaCache().set("ag-a", { + [MODEL]: { remainingPercentage: 0, resetAt: FUTURE_RESET }, + }); + + try { + await expect(getProviderCredentials("antigravity", null, MODEL)).resolves.toMatchObject({ + allRateLimited: true, + retryAfter: FUTURE_RESET, + }); + } finally { + vi.useRealTimers(); + } + }); + + it("lets account back into rotation once reset time has passed", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-09-01T00:00:01.000Z")); + mocks.getProviderConnections.mockResolvedValue([{ id: "ag-a", email: "a@example.com", isActive: true }]); + getAntigravityQuotaCache().set("ag-a", { + [MODEL]: { remainingPercentage: 0, resetAt: FUTURE_RESET }, + }); + + try { + await expect(getProviderCredentials("antigravity", null, MODEL)).resolves.toMatchObject({ + connectionId: "ag-a", + connectionName: "a@example.com", + }); + } finally { + vi.useRealTimers(); + } + }); + + it("coalesces concurrent quota refreshes for one account", async () => { + let resolveUsage; + mocks.getAntigravityUsage.mockReturnValue(new Promise(resolve => { resolveUsage = resolve; })); + + const first = refreshAntigravityQuota("ag-concurrent", "token", {}); + const second = refreshAntigravityQuota("ag-concurrent", "token", {}); + resolveUsage({ quotas: { [MODEL]: { remainingPercentage: 0, resetAt: FUTURE_RESET } } }); + + await expect(Promise.all([first, second])).resolves.toEqual([ + { [MODEL]: { remainingPercentage: 0, resetAt: FUTURE_RESET } }, + { [MODEL]: { remainingPercentage: 0, resetAt: FUTURE_RESET } }, + ]); + expect(mocks.getAntigravityUsage).toHaveBeenCalledTimes(1); + }); + + it("preserves strict proxy policy for usage refresh", async () => { + mocks.resolveConnectionProxyConfig.mockResolvedValue({ strictProxy: true }); + mocks.getAntigravityUsage.mockResolvedValue({ quotas: {} }); + + await refreshAntigravityQuota("ag-strict-proxy", "token", {}); + + expect(mocks.getAntigravityUsage).toHaveBeenCalledWith("token", {}, expect.objectContaining({ + strictProxy: true, + })); + }); + + it("keeps known cache when quota endpoint returns an error payload", async () => { + const cached = { [MODEL]: { remainingPercentage: 0, resetAt: FUTURE_RESET } }; + getAntigravityQuotaCache().set("ag-error-response", cached); + mocks.getAntigravityUsage.mockResolvedValue({ message: "Unauthorized", quotas: {} }); + + await expect(refreshAntigravityQuota("ag-error-response", "token", {})).resolves.toBeNull(); + expect(getAntigravityQuotaCache().get("ag-error-response")).toBe(cached); + }); + + it("throttles failed refresh attempts for 30 seconds", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-08-26T00:00:00.000Z")); + mocks.getAntigravityUsage.mockRejectedValue(new Error("usage unavailable")); + + try { + await refreshAntigravityQuota("ag-failed-refresh", "token", {}); + await refreshAntigravityQuota("ag-failed-refresh", "token", {}); + expect(mocks.getAntigravityUsage).toHaveBeenCalledTimes(1); + + await vi.advanceTimersByTimeAsync(30_000); + await refreshAntigravityQuota("ag-failed-refresh", "token", {}); + expect(mocks.getAntigravityUsage).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); +});