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:
Nim committed 2026-08-27 16:18:08 +07:00
1 parent f0a6d35818
commit 1a3db1efae
4 files changed
+314 -5

No files matched your search

+99
View 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;
}