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
This commit is contained in:
@@ -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`);
|
||||
|
||||
99
src/sse/services/antigravityQuota.js
Normal file
99
src/sse/services/antigravityQuota.js
Normal file
@@ -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;
|
||||
}
|
||||
@@ -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));
|
||||
|
||||
Reference in New Issue
Block a user