fix: accurate quota reporting with project ID and improved rate limit handling

- Pass project ID to fetchAvailableModels for accurate per-project quota
- Treat missing remainingFraction with resetTime as 0% (exhausted)
- Fix double-escaped regex in rate-limit-parser.js (\\d -> \d)
- Use ANTIGRAVITY_HEADERS for loadCodeAssist consistency
- Store actual reset time from API instead of capping at default
- Add getRateLimitInfo() for detailed rate limit state
- Handle disabled accounts in rate limit checks

Fixes issue where free tier accounts showed 100% quota but were actually exhausted.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Badri Narayanan S
2026-01-15 16:18:13 +05:30
parent a20cba90ee
commit 77363c679e
8 changed files with 252 additions and 163 deletions

View File

@@ -14,7 +14,8 @@ import {
resetAllRateLimits as resetLimits, resetAllRateLimits as resetLimits,
markRateLimited as markLimited, markRateLimited as markLimited,
markInvalid as markAccountInvalid, markInvalid as markAccountInvalid,
getMinWaitTimeMs as getMinWait getMinWaitTimeMs as getMinWait,
getRateLimitInfo as getLimitInfo
} from './rate-limits.js'; } from './rate-limits.js';
import { import {
getTokenForAccount as fetchToken, getTokenForAccount as fetchToken,
@@ -214,6 +215,16 @@ export class AccountManager {
return getMinWait(this.#accounts, modelId); return getMinWait(this.#accounts, modelId);
} }
/**
* Get rate limit info for a specific account and model
* @param {string} email - Email of the account
* @param {string} modelId - Model ID to check
* @returns {{isRateLimited: boolean, actualResetMs: number|null, waitMs: number}} Rate limit info
*/
getRateLimitInfo(email, modelId) {
return getLimitInfo(this.#accounts, email, modelId);
}
/** /**
* Get OAuth token for an account * Get OAuth token for an account
* @param {Object} account - Account object with email and credentials * @param {Object} account - Account object with email and credentials

View File

@@ -22,6 +22,7 @@ export function isAllRateLimited(accounts, modelId) {
return accounts.every(acc => { return accounts.every(acc => {
if (acc.isInvalid) return true; // Invalid accounts count as unavailable if (acc.isInvalid) return true; // Invalid accounts count as unavailable
if (acc.enabled === false) return true; // Disabled accounts count as unavailable
const modelLimits = acc.modelRateLimits || {}; const modelLimits = acc.modelRateLimits || {};
const limit = modelLimits[modelId]; const limit = modelLimits[modelId];
return limit && limit.isRateLimited && limit.resetTime > Date.now(); return limit && limit.isRateLimited && limit.resetTime > Date.now();
@@ -118,18 +119,9 @@ export function markRateLimited(accounts, email, resetMs = null, modelId) {
const account = accounts.find(a => a.email === email); const account = accounts.find(a => a.email === email);
if (!account) return false; if (!account) return false;
// Use configured cooldown as the maximum wait time // Store the ACTUAL reset time from the API
// If API returns a reset time, cap it at DEFAULT_COOLDOWN_MS // This is used to decide whether to wait (short) or switch accounts (long)
// If API doesn't return a reset time, use DEFAULT_COOLDOWN_MS const actualResetMs = (resetMs && resetMs > 0) ? resetMs : DEFAULT_COOLDOWN_MS;
let cooldownMs;
if (resetMs && resetMs > 0) {
// API provided a reset time - cap it at configured maximum
cooldownMs = Math.min(resetMs, DEFAULT_COOLDOWN_MS);
} else {
// No reset time from API - use configured default
cooldownMs = DEFAULT_COOLDOWN_MS;
}
const resetTime = Date.now() + cooldownMs;
if (!account.modelRateLimits) { if (!account.modelRateLimits) {
account.modelRateLimits = {}; account.modelRateLimits = {};
@@ -137,12 +129,20 @@ export function markRateLimited(accounts, email, resetMs = null, modelId) {
account.modelRateLimits[modelId] = { account.modelRateLimits[modelId] = {
isRateLimited: true, isRateLimited: true,
resetTime: resetTime resetTime: Date.now() + actualResetMs, // Actual reset time for decisions
actualResetMs: actualResetMs // Original duration from API
}; };
logger.warn( // Log appropriately based on duration
`[AccountManager] Rate limited: ${email} (model: ${modelId}). Available in ${formatDuration(cooldownMs)}` if (actualResetMs > DEFAULT_COOLDOWN_MS) {
); logger.warn(
`[AccountManager] Quota exhausted: ${email} (model: ${modelId}). Resets in ${formatDuration(actualResetMs)}`
);
} else {
logger.warn(
`[AccountManager] Rate limited: ${email} (model: ${modelId}). Available in ${formatDuration(actualResetMs)}`
);
}
return true; return true;
} }
@@ -209,3 +209,29 @@ export function getMinWaitTimeMs(accounts, modelId) {
return minWait === Infinity ? DEFAULT_COOLDOWN_MS : minWait; return minWait === Infinity ? DEFAULT_COOLDOWN_MS : minWait;
} }
/**
* Get the rate limit info for a specific account and model
* Returns the actual reset time from API, not capped
*
* @param {Array} accounts - Array of account objects
* @param {string} email - Email of the account
* @param {string} modelId - Model ID to check
* @returns {{isRateLimited: boolean, actualResetMs: number|null, waitMs: number}} Rate limit info
*/
export function getRateLimitInfo(accounts, email, modelId) {
const account = accounts.find(a => a.email === email);
if (!account || !account.modelRateLimits || !account.modelRateLimits[modelId]) {
return { isRateLimited: false, actualResetMs: null, waitMs: 0 };
}
const limit = account.modelRateLimits[modelId];
const now = Date.now();
const waitMs = limit.resetTime ? Math.max(0, limit.resetTime - now) : 0;
return {
isRateLimited: limit.isRateLimited && waitMs > 0,
actualResetMs: limit.actualResetMs || null,
waitMs
};
}

View File

@@ -9,6 +9,7 @@ import {
ANTIGRAVITY_ENDPOINT_FALLBACKS, ANTIGRAVITY_ENDPOINT_FALLBACKS,
MAX_RETRIES, MAX_RETRIES,
MAX_WAIT_BEFORE_ERROR_MS, MAX_WAIT_BEFORE_ERROR_MS,
DEFAULT_COOLDOWN_MS,
isThinkingModel isThinkingModel
} from '../constants.js'; } from '../constants.js';
import { convertGoogleToAnthropic } from '../format/index.js'; import { convertGoogleToAnthropic } from '../format/index.js';
@@ -39,67 +40,56 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab
// Retry loop with account failover // Retry loop with account failover
// Ensure we try at least as many times as there are accounts to cycle through everyone // Ensure we try at least as many times as there are accounts to cycle through everyone
// +1 to ensure we hit the "all accounts rate-limited" check at the start of the next loop
const maxAttempts = Math.max(MAX_RETRIES, accountManager.getAccountCount() + 1); const maxAttempts = Math.max(MAX_RETRIES, accountManager.getAccountCount() + 1);
for (let attempt = 0; attempt < maxAttempts; attempt++) { for (let attempt = 0; attempt < maxAttempts; attempt++) {
// Use sticky account selection for cache continuity // Clear any expired rate limits before picking
const { account: stickyAccount, waitMs } = accountManager.pickStickyAccount(model); accountManager.clearExpiredLimits();
let account = stickyAccount;
// Handle waiting for sticky account // Get available accounts for this model
if (!account && waitMs > 0) { const availableAccounts = accountManager.getAvailableAccounts(model);
logger.info(`[CloudCode] Waiting ${formatDuration(waitMs)} for sticky account...`);
await sleep(waitMs);
accountManager.clearExpiredLimits();
account = accountManager.getCurrentStickyAccount(model);
}
// Handle all accounts rate-limited // If no accounts available, check if we should wait or throw error
if (!account) { if (availableAccounts.length === 0) {
if (accountManager.isAllRateLimited(model)) { if (accountManager.isAllRateLimited(model)) {
const allWaitMs = accountManager.getMinWaitTimeMs(model); const minWaitMs = accountManager.getMinWaitTimeMs(model);
const resetTime = new Date(Date.now() + allWaitMs).toISOString(); const resetTime = new Date(Date.now() + minWaitMs).toISOString();
// If wait time is too long (> 2 minutes), throw error immediately // If wait time is too long (> 2 minutes), throw error immediately
if (allWaitMs > MAX_WAIT_BEFORE_ERROR_MS) { if (minWaitMs > MAX_WAIT_BEFORE_ERROR_MS) {
throw new Error( throw new Error(
`RESOURCE_EXHAUSTED: Rate limited on ${model}. Quota will reset after ${formatDuration(allWaitMs)}. Next available: ${resetTime}` `RESOURCE_EXHAUSTED: Rate limited on ${model}. Quota will reset after ${formatDuration(minWaitMs)}. Next available: ${resetTime}`
); );
} }
// Wait for reset (applies to both single and multi-account modes) // Wait for shortest reset time
const accountCount = accountManager.getAccountCount(); const accountCount = accountManager.getAccountCount();
logger.warn(`[CloudCode] All ${accountCount} account(s) rate-limited. Waiting ${formatDuration(allWaitMs)}...`); logger.warn(`[CloudCode] All ${accountCount} account(s) rate-limited. Waiting ${formatDuration(minWaitMs)}...`);
await sleep(allWaitMs); await sleep(minWaitMs + 500); // Add 500ms buffer
// Add small buffer after waiting to ensure rate limits have truly expired
await sleep(500);
accountManager.clearExpiredLimits(); accountManager.clearExpiredLimits();
account = accountManager.pickNext(model); continue; // Retry the loop
// If still no account after waiting, try optimistic reset
// This handles cases where the API rate limit is transient
if (!account) {
logger.warn('[CloudCode] No account available after wait, attempting optimistic reset...');
accountManager.resetAllRateLimits();
account = accountManager.pickNext(model);
}
} }
if (!account) { // Check if fallback is enabled and available
// Check if fallback is enabled and available if (fallbackEnabled) {
if (fallbackEnabled) { const fallbackModel = getFallbackModel(model);
const fallbackModel = getFallbackModel(model); if (fallbackModel) {
if (fallbackModel) { logger.warn(`[CloudCode] All accounts exhausted for ${model}. Attempting fallback to ${fallbackModel}`);
logger.warn(`[CloudCode] All accounts exhausted for ${model}. Attempting fallback to ${fallbackModel}`); const fallbackRequest = { ...anthropicRequest, model: fallbackModel };
// Retry with fallback model return await sendMessage(fallbackRequest, accountManager, false);
const fallbackRequest = { ...anthropicRequest, model: fallbackModel };
return await sendMessage(fallbackRequest, accountManager, false); // Disable fallback for recursive call
}
} }
throw new Error('No accounts available');
} }
throw new Error('No accounts available');
}
// Pick sticky account (prefers current for cache continuity)
let account = accountManager.getCurrentStickyAccount(model);
if (!account) {
account = accountManager.pickNext(model);
}
if (!account) {
continue; // Shouldn't happen, but safety check
} }
try { try {
@@ -112,6 +102,8 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab
// Try each endpoint // Try each endpoint
let lastError = null; let lastError = null;
let retriedOnce = false; // Track if we've already retried for short rate limit
for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) { for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) {
try { try {
const url = isThinking const url = isThinking
@@ -137,14 +129,51 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab
} }
if (response.status === 429) { if (response.status === 429) {
// Rate limited on this endpoint - try next endpoint first (DAILY → PROD)
logger.debug(`[CloudCode] Rate limited at ${endpoint}, trying next endpoint...`);
const resetMs = parseResetTime(response, errorText); const resetMs = parseResetTime(response, errorText);
// Keep minimum reset time across all 429 responses
if (!lastError?.is429 || (resetMs && (!lastError.resetMs || resetMs < lastError.resetMs))) { // Decision: wait and retry OR switch account
lastError = { is429: true, response, errorText, resetMs }; if (resetMs && resetMs > DEFAULT_COOLDOWN_MS) {
// Long-term quota exhaustion (> 10s) - switch to next account
logger.info(`[CloudCode] Quota exhausted for ${account.email} (${formatDuration(resetMs)}), switching account...`);
accountManager.markRateLimited(account.email, resetMs, model);
throw new Error(`QUOTA_EXHAUSTED: ${errorText}`);
} else {
// Short-term rate limit (<= 10s) - wait and retry once
const waitMs = resetMs || DEFAULT_COOLDOWN_MS;
if (!retriedOnce) {
retriedOnce = true;
logger.info(`[CloudCode] Short rate limit (${formatDuration(waitMs)}), waiting and retrying...`);
await sleep(waitMs);
// Retry same endpoint
const retryResponse = await fetch(url, {
method: 'POST',
headers: buildHeaders(token, model, isThinking ? 'text/event-stream' : 'application/json'),
body: JSON.stringify(payload)
});
if (retryResponse.ok) {
// Process retry response
if (isThinking) {
return await parseThinkingSSEResponse(retryResponse, anthropicRequest.model);
}
const data = await retryResponse.json();
logger.debug('[CloudCode] Response received after retry');
return convertGoogleToAnthropic(data, anthropicRequest.model);
}
// Retry also failed - parse new reset time
const retryErrorText = await retryResponse.text();
const retryResetMs = parseResetTime(retryResponse, retryErrorText);
logger.warn(`[CloudCode] Retry also failed, marking and switching...`);
accountManager.markRateLimited(account.email, retryResetMs || waitMs, model);
throw new Error(`RATE_LIMITED_AFTER_RETRY: ${retryErrorText}`);
} else {
// Already retried once, mark and switch
accountManager.markRateLimited(account.email, waitMs, model);
throw new Error(`RATE_LIMITED: ${errorText}`);
}
} }
continue;
} }
if (response.status >= 400) { if (response.status >= 400) {
@@ -179,7 +208,6 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab
// If all endpoints failed for this account // If all endpoints failed for this account
if (lastError) { if (lastError) {
// If all endpoints returned 429, mark account as rate-limited
if (lastError.is429) { if (lastError.is429) {
logger.warn(`[CloudCode] All endpoints rate-limited for ${account.email}`); logger.warn(`[CloudCode] All endpoints rate-limited for ${account.email}`);
accountManager.markRateLimited(account.email, lastError.resetMs, model); accountManager.markRateLimited(account.email, lastError.resetMs, model);
@@ -199,18 +227,17 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab
logger.warn(`[CloudCode] Account ${account.email} has invalid credentials, trying next...`); logger.warn(`[CloudCode] Account ${account.email} has invalid credentials, trying next...`);
continue; continue;
} }
// Non-rate-limit error: throw immediately // Handle 5xx errors
// UNLESS it's a 500 error, then we treat it as a "soft" failure for this account and try the next one
if (error.message.includes('API error 5') || error.message.includes('500') || error.message.includes('503')) { if (error.message.includes('API error 5') || error.message.includes('500') || error.message.includes('503')) {
logger.warn(`[CloudCode] Account ${account.email} failed with 5xx error, trying next...`); logger.warn(`[CloudCode] Account ${account.email} failed with 5xx error, trying next...`);
accountManager.pickNext(model); // Force advance to next account accountManager.pickNext(model);
continue; continue;
} }
if (isNetworkError(error)) { if (isNetworkError(error)) {
logger.warn(`[CloudCode] Network error for ${account.email}, trying next account... (${error.message})`); logger.warn(`[CloudCode] Network error for ${account.email}, trying next account... (${error.message})`);
await sleep(1000); // Brief pause before retry await sleep(1000);
accountManager.pickNext(model); // Advance to next account accountManager.pickNext(model);
continue; continue;
} }
@@ -224,7 +251,7 @@ export async function sendMessage(anthropicRequest, accountManager, fallbackEnab
if (fallbackModel) { if (fallbackModel) {
logger.warn(`[CloudCode] All retries exhausted for ${model}. Attempting fallback to ${fallbackModel}`); logger.warn(`[CloudCode] All retries exhausted for ${model}. Attempting fallback to ${fallbackModel}`);
const fallbackRequest = { ...anthropicRequest, model: fallbackModel }; const fallbackRequest = { ...anthropicRequest, model: fallbackModel };
return await sendMessage(fallbackRequest, accountManager, false); // Disable fallback for recursive call return await sendMessage(fallbackRequest, accountManager, false);
} }
} }

View File

@@ -57,22 +57,26 @@ export async function listModels(token) {
* Returns model quotas including remaining fraction and reset time * Returns model quotas including remaining fraction and reset time
* *
* @param {string} token - OAuth access token * @param {string} token - OAuth access token
* @param {string} [projectId] - Optional project ID for accurate quota info
* @returns {Promise<Object>} Raw response from fetchAvailableModels API * @returns {Promise<Object>} Raw response from fetchAvailableModels API
*/ */
export async function fetchAvailableModels(token) { export async function fetchAvailableModels(token, projectId = null) {
const headers = { const headers = {
'Authorization': `Bearer ${token}`, 'Authorization': `Bearer ${token}`,
'Content-Type': 'application/json', 'Content-Type': 'application/json',
...ANTIGRAVITY_HEADERS ...ANTIGRAVITY_HEADERS
}; };
// Include project ID in body for accurate quota info (per Quotio implementation)
const body = projectId ? { project: projectId } : {};
for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) { for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) {
try { try {
const url = `${endpoint}/v1internal:fetchAvailableModels`; const url = `${endpoint}/v1internal:fetchAvailableModels`;
const response = await fetch(url, { const response = await fetch(url, {
method: 'POST', method: 'POST',
headers, headers,
body: JSON.stringify({}) body: JSON.stringify(body)
}); });
if (!response.ok) { if (!response.ok) {
@@ -95,10 +99,11 @@ export async function fetchAvailableModels(token) {
* Extracts quota info (remaining fraction and reset time) for each model * Extracts quota info (remaining fraction and reset time) for each model
* *
* @param {string} token - OAuth access token * @param {string} token - OAuth access token
* @param {string} [projectId] - Optional project ID for accurate quota info
* @returns {Promise<Object>} Map of modelId -> { remainingFraction, resetTime } * @returns {Promise<Object>} Map of modelId -> { remainingFraction, resetTime }
*/ */
export async function getModelQuotas(token) { export async function getModelQuotas(token, projectId = null) {
const data = await fetchAvailableModels(token); const data = await fetchAvailableModels(token, projectId);
if (!data || !data.models) return {}; if (!data || !data.models) return {};
const quotas = {}; const quotas = {};
@@ -108,7 +113,8 @@ export async function getModelQuotas(token) {
if (modelData.quotaInfo) { if (modelData.quotaInfo) {
quotas[modelId] = { quotas[modelId] = {
remainingFraction: modelData.quotaInfo.remainingFraction ?? null, // When remainingFraction is missing but resetTime is present, quota is exhausted (0%)
remainingFraction: modelData.quotaInfo.remainingFraction ?? (modelData.quotaInfo.resetTime ? 0 : null),
resetTime: modelData.quotaInfo.resetTime ?? null resetTime: modelData.quotaInfo.resetTime ?? null
}; };
} }

View File

@@ -78,7 +78,7 @@ export function parseResetTime(responseOrError, errorText = '') {
// Try to extract "quotaResetDelay" first (e.g. "754.431528ms" or "1.5s") // Try to extract "quotaResetDelay" first (e.g. "754.431528ms" or "1.5s")
// This is Google's preferred format for rate limit reset delay // This is Google's preferred format for rate limit reset delay
const quotaDelayMatch = msg.match(/quotaResetDelay[:\s"]+(\\d+(?:\\.\\d+)?)(ms|s)/i); const quotaDelayMatch = msg.match(/quotaResetDelay[:\s"]+(\d+(?:\.\d+)?)(ms|s)/i);
if (quotaDelayMatch) { if (quotaDelayMatch) {
const value = parseFloat(quotaDelayMatch[1]); const value = parseFloat(quotaDelayMatch[1]);
const unit = quotaDelayMatch[2].toLowerCase(); const unit = quotaDelayMatch[2].toLowerCase();
@@ -103,7 +103,7 @@ export function parseResetTime(responseOrError, errorText = '') {
// Try to extract "retry-after-ms" or "retryDelay" - check seconds format first (e.g. "7739.23s") // Try to extract "retry-after-ms" or "retryDelay" - check seconds format first (e.g. "7739.23s")
// Added stricter regex to avoid partial matches // Added stricter regex to avoid partial matches
if (!resetMs) { if (!resetMs) {
const secMatch = msg.match(/(?:retry[-_]?after[-_]?ms|retryDelay)[:\s"]+([\\d\\.]+)(?:s\b|s")/i); const secMatch = msg.match(/(?:retry[-_]?after[-_]?ms|retryDelay)[:\s"]+([\d.]+)(?:s\b|s")/i);
if (secMatch) { if (secMatch) {
resetMs = Math.ceil(parseFloat(secMatch[1]) * 1000); resetMs = Math.ceil(parseFloat(secMatch[1]) * 1000);
logger.debug(`[CloudCode] Parsed retry seconds from body (precise): ${resetMs}ms`); logger.debug(`[CloudCode] Parsed retry seconds from body (precise): ${resetMs}ms`);

View File

@@ -9,7 +9,8 @@ import {
ANTIGRAVITY_ENDPOINT_FALLBACKS, ANTIGRAVITY_ENDPOINT_FALLBACKS,
MAX_RETRIES, MAX_RETRIES,
MAX_EMPTY_RESPONSE_RETRIES, MAX_EMPTY_RESPONSE_RETRIES,
MAX_WAIT_BEFORE_ERROR_MS MAX_WAIT_BEFORE_ERROR_MS,
DEFAULT_COOLDOWN_MS
} from '../constants.js'; } from '../constants.js';
import { isRateLimitError, isAuthError, isEmptyResponseError } from '../errors.js'; import { isRateLimitError, isAuthError, isEmptyResponseError } from '../errors.js';
import { formatDuration, sleep, isNetworkError } from '../utils/helpers.js'; import { formatDuration, sleep, isNetworkError } from '../utils/helpers.js';
@@ -38,68 +39,57 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
// Retry loop with account failover // Retry loop with account failover
// Ensure we try at least as many times as there are accounts to cycle through everyone // Ensure we try at least as many times as there are accounts to cycle through everyone
// +1 to ensure we hit the "all accounts rate-limited" check at the start of the next loop
const maxAttempts = Math.max(MAX_RETRIES, accountManager.getAccountCount() + 1); const maxAttempts = Math.max(MAX_RETRIES, accountManager.getAccountCount() + 1);
for (let attempt = 0; attempt < maxAttempts; attempt++) { for (let attempt = 0; attempt < maxAttempts; attempt++) {
// Use sticky account selection for cache continuity // Clear any expired rate limits before picking
const { account: stickyAccount, waitMs } = accountManager.pickStickyAccount(model); accountManager.clearExpiredLimits();
let account = stickyAccount;
// Handle waiting for sticky account // Get available accounts for this model
if (!account && waitMs > 0) { const availableAccounts = accountManager.getAvailableAccounts(model);
logger.info(`[CloudCode] Waiting ${formatDuration(waitMs)} for sticky account...`);
await sleep(waitMs);
accountManager.clearExpiredLimits();
account = accountManager.getCurrentStickyAccount(model);
}
// Handle all accounts rate-limited // If no accounts available, check if we should wait or throw error
if (!account) { if (availableAccounts.length === 0) {
if (accountManager.isAllRateLimited(model)) { if (accountManager.isAllRateLimited(model)) {
const allWaitMs = accountManager.getMinWaitTimeMs(model); const minWaitMs = accountManager.getMinWaitTimeMs(model);
const resetTime = new Date(Date.now() + allWaitMs).toISOString(); const resetTime = new Date(Date.now() + minWaitMs).toISOString();
// If wait time is too long (> 2 minutes), throw error immediately // If wait time is too long (> 2 minutes), throw error immediately
if (allWaitMs > MAX_WAIT_BEFORE_ERROR_MS) { if (minWaitMs > MAX_WAIT_BEFORE_ERROR_MS) {
throw new Error( throw new Error(
`RESOURCE_EXHAUSTED: Rate limited on ${model}. Quota will reset after ${formatDuration(allWaitMs)}. Next available: ${resetTime}` `RESOURCE_EXHAUSTED: Rate limited on ${model}. Quota will reset after ${formatDuration(minWaitMs)}. Next available: ${resetTime}`
); );
} }
// Wait for reset (applies to both single and multi-account modes) // Wait for shortest reset time
const accountCount = accountManager.getAccountCount(); const accountCount = accountManager.getAccountCount();
logger.warn(`[CloudCode] All ${accountCount} account(s) rate-limited. Waiting ${formatDuration(allWaitMs)}...`); logger.warn(`[CloudCode] All ${accountCount} account(s) rate-limited. Waiting ${formatDuration(minWaitMs)}...`);
await sleep(allWaitMs); await sleep(minWaitMs + 500); // Add 500ms buffer
// Add small buffer after waiting to ensure rate limits have truly expired
await sleep(500);
accountManager.clearExpiredLimits(); accountManager.clearExpiredLimits();
account = accountManager.pickNext(model); continue; // Retry the loop
// If still no account after waiting, try optimistic reset
// This handles cases where the API rate limit is transient
if (!account) {
logger.warn('[CloudCode] No account available after wait, attempting optimistic reset...');
accountManager.resetAllRateLimits();
account = accountManager.pickNext(model);
}
} }
if (!account) { // Check if fallback is enabled and available
// Check if fallback is enabled and available if (fallbackEnabled) {
if (fallbackEnabled) { const fallbackModel = getFallbackModel(model);
const fallbackModel = getFallbackModel(model); if (fallbackModel) {
if (fallbackModel) { logger.warn(`[CloudCode] All accounts exhausted for ${model}. Attempting fallback to ${fallbackModel} (streaming)`);
logger.warn(`[CloudCode] All accounts exhausted for ${model}. Attempting fallback to ${fallbackModel} (streaming)`); const fallbackRequest = { ...anthropicRequest, model: fallbackModel };
// Retry with fallback model yield* sendMessageStream(fallbackRequest, accountManager, false);
const fallbackRequest = { ...anthropicRequest, model: fallbackModel }; return;
yield* sendMessageStream(fallbackRequest, accountManager, false); // Disable fallback for recursive call
return;
}
} }
throw new Error('No accounts available');
} }
throw new Error('No accounts available');
}
// Pick sticky account (prefers current for cache continuity)
let account = accountManager.getCurrentStickyAccount(model);
if (!account) {
account = accountManager.pickNext(model);
}
if (!account) {
continue; // Shouldn't happen, but safety check
} }
try { try {
@@ -112,6 +102,8 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
// Try each endpoint for streaming // Try each endpoint for streaming
let lastError = null; let lastError = null;
let retriedOnce = false; // Track if we've already retried for short rate limit
for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) { for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) {
try { try {
const url = `${endpoint}/v1internal:streamGenerateContent?alt=sse`; const url = `${endpoint}/v1internal:streamGenerateContent?alt=sse`;
@@ -134,14 +126,48 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
} }
if (response.status === 429) { if (response.status === 429) {
// Rate limited on this endpoint - try next endpoint first (DAILY → PROD)
logger.debug(`[CloudCode] Stream rate limited at ${endpoint}, trying next endpoint...`);
const resetMs = parseResetTime(response, errorText); const resetMs = parseResetTime(response, errorText);
// Keep minimum reset time across all 429 responses
if (!lastError?.is429 || (resetMs && (!lastError.resetMs || resetMs < lastError.resetMs))) { // Decision: wait and retry OR switch account
lastError = { is429: true, response, errorText, resetMs }; if (resetMs && resetMs > DEFAULT_COOLDOWN_MS) {
// Long-term quota exhaustion (> 10s) - switch to next account
logger.info(`[CloudCode] Quota exhausted for ${account.email} (${formatDuration(resetMs)}), switching account...`);
accountManager.markRateLimited(account.email, resetMs, model);
throw new Error(`QUOTA_EXHAUSTED: ${errorText}`);
} else {
// Short-term rate limit (<= 10s) - wait and retry once
const waitMs = resetMs || DEFAULT_COOLDOWN_MS;
if (!retriedOnce) {
retriedOnce = true;
logger.info(`[CloudCode] Short rate limit (${formatDuration(waitMs)}), waiting and retrying...`);
await sleep(waitMs);
// Retry same endpoint
const retryResponse = await fetch(url, {
method: 'POST',
headers: buildHeaders(token, model, 'text/event-stream'),
body: JSON.stringify(payload)
});
if (retryResponse.ok) {
// Stream the retry response
yield* streamSSEResponse(retryResponse, anthropicRequest.model);
logger.debug('[CloudCode] Stream completed after retry');
return;
}
// Retry also failed - parse new reset time
const retryErrorText = await retryResponse.text();
const retryResetMs = parseResetTime(retryResponse, retryErrorText);
logger.warn(`[CloudCode] Retry also failed, marking and switching...`);
accountManager.markRateLimited(account.email, retryResetMs || waitMs, model);
throw new Error(`RATE_LIMITED_AFTER_RETRY: ${retryErrorText}`);
} else {
// Already retried once, mark and switch
accountManager.markRateLimited(account.email, waitMs, model);
throw new Error(`RATE_LIMITED: ${errorText}`);
}
} }
continue;
} }
lastError = new Error(`API error ${response.status}: ${errorText}`); lastError = new Error(`API error ${response.status}: ${errorText}`);
@@ -156,7 +182,6 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
} }
// Stream the response with retry logic for empty responses // Stream the response with retry logic for empty responses
// Uses a for-loop for clearer retry semantics
let currentResponse = response; let currentResponse = response;
for (let emptyRetries = 0; emptyRetries <= MAX_EMPTY_RESPONSE_RETRIES; emptyRetries++) { for (let emptyRetries = 0; emptyRetries <= MAX_EMPTY_RESPONSE_RETRIES; emptyRetries++) {
@@ -207,28 +232,22 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
throw new Error(`401 AUTH_INVALID during retry: ${retryErrorText}`); throw new Error(`401 AUTH_INVALID during retry: ${retryErrorText}`);
} }
// For 5xx errors, don't pass to streamer - just continue to next retry // For 5xx errors, continue retrying
if (currentResponse.status >= 500) { if (currentResponse.status >= 500) {
logger.warn(`[CloudCode] Retry got ${currentResponse.status}, will retry...`); logger.warn(`[CloudCode] Retry got ${currentResponse.status}, will retry...`);
// Don't continue here - let the loop increment and refetch
// Set currentResponse to null to force refetch at loop start
emptyRetries--; // Compensate for loop increment since we didn't actually try
await sleep(1000); await sleep(1000);
// Refetch immediately for 5xx
currentResponse = await fetch(url, { currentResponse = await fetch(url, {
method: 'POST', method: 'POST',
headers: buildHeaders(token, model, 'text/event-stream'), headers: buildHeaders(token, model, 'text/event-stream'),
body: JSON.stringify(payload) body: JSON.stringify(payload)
}); });
if (currentResponse.ok) { if (currentResponse.ok) {
continue; // Try streaming with new response continue;
} }
// If still failing, let it fall through to throw
} }
throw new Error(`Empty response retry failed: ${currentResponse.status} - ${retryErrorText}`); throw new Error(`Empty response retry failed: ${currentResponse.status} - ${retryErrorText}`);
} }
// Response is OK, loop will continue to try streamSSEResponse
} }
} }
@@ -237,7 +256,7 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
throw endpointError; // Re-throw to trigger account switch throw endpointError; // Re-throw to trigger account switch
} }
if (isEmptyResponseError(endpointError)) { if (isEmptyResponseError(endpointError)) {
throw endpointError; // Re-throw empty response errors to outer handler throw endpointError;
} }
logger.warn(`[CloudCode] Stream error at ${endpoint}:`, endpointError.message); logger.warn(`[CloudCode] Stream error at ${endpoint}:`, endpointError.message);
lastError = endpointError; lastError = endpointError;
@@ -246,7 +265,6 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
// If all endpoints failed for this account // If all endpoints failed for this account
if (lastError) { if (lastError) {
// If all endpoints returned 429, mark account as rate-limited
if (lastError.is429) { if (lastError.is429) {
logger.warn(`[CloudCode] All endpoints rate-limited for ${account.email}`); logger.warn(`[CloudCode] All endpoints rate-limited for ${account.email}`);
accountManager.markRateLimited(account.email, lastError.resetMs, model); accountManager.markRateLimited(account.email, lastError.resetMs, model);
@@ -266,18 +284,17 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
logger.warn(`[CloudCode] Account ${account.email} has invalid credentials, trying next...`); logger.warn(`[CloudCode] Account ${account.email} has invalid credentials, trying next...`);
continue; continue;
} }
// Non-rate-limit error: throw immediately // Handle 5xx errors
// UNLESS it's a 500 error, then we treat it as a "soft" failure for this account and try the next one
if (error.message.includes('API error 5') || error.message.includes('500') || error.message.includes('503')) { if (error.message.includes('API error 5') || error.message.includes('500') || error.message.includes('503')) {
logger.warn(`[CloudCode] Account ${account.email} failed with 5xx stream error, trying next...`); logger.warn(`[CloudCode] Account ${account.email} failed with 5xx stream error, trying next...`);
accountManager.pickNext(model); // Force advance to next account accountManager.pickNext(model);
continue; continue;
} }
if (isNetworkError(error)) { if (isNetworkError(error)) {
logger.warn(`[CloudCode] Network error for ${account.email} (stream), trying next account... (${error.message})`); logger.warn(`[CloudCode] Network error for ${account.email} (stream), trying next account... (${error.message})`);
await sleep(1000); // Brief pause before retry await sleep(1000);
accountManager.pickNext(model); // Advance to next account accountManager.pickNext(model);
continue; continue;
} }
@@ -291,7 +308,7 @@ export async function* sendMessageStream(anthropicRequest, accountManager, fallb
if (fallbackModel) { if (fallbackModel) {
logger.warn(`[CloudCode] All retries exhausted for ${model}. Attempting fallback to ${fallbackModel} (streaming)`); logger.warn(`[CloudCode] All retries exhausted for ${model}. Attempting fallback to ${fallbackModel} (streaming)`);
const fallbackRequest = { ...anthropicRequest, model: fallbackModel }; const fallbackRequest = { ...anthropicRequest, model: fallbackModel };
yield* sendMessageStream(fallbackRequest, accountManager, false); // Disable fallback for recursive call yield* sendMessageStream(fallbackRequest, accountManager, false);
return; return;
} }
} }

View File

@@ -69,15 +69,16 @@ export const ONBOARD_USER_ENDPOINTS = ANTIGRAVITY_ENDPOINT_FALLBACKS;
// Hybrid headers specifically for loadCodeAssist // Hybrid headers specifically for loadCodeAssist
// Uses google-api-nodejs-client User-Agent (required for project discovery on some accounts) // Uses google-api-nodejs-client User-Agent (required for project discovery on some accounts)
export const LOAD_CODE_ASSIST_HEADERS = { // export const LOAD_CODE_ASSIST_HEADERS = {
'User-Agent': 'google-api-nodejs-client/9.15.1', // 'User-Agent': 'google-api-nodejs-client/9.15.1',
'X-Goog-Api-Client': 'google-cloud-sdk vscode_cloudshelleditor/0.1', // 'X-Goog-Api-Client': 'google-cloud-sdk vscode_cloudshelleditor/0.1',
'Client-Metadata': JSON.stringify({ // 'Client-Metadata': JSON.stringify({
ideType: 'IDE_UNSPECIFIED', // ideType: 'IDE_UNSPECIFIED',
platform: 'PLATFORM_UNSPECIFIED', // platform: 'PLATFORM_UNSPECIFIED',
pluginType: 'GEMINI' // pluginType: 'GEMINI'
}) // })
}; // };
export const LOAD_CODE_ASSIST_HEADERS = ANTIGRAVITY_HEADERS;
// Default project ID if none can be discovered // Default project ID if none can be discovered
export const DEFAULT_PROJECT_ID = 'rising-fact-p41fc'; export const DEFAULT_PROJECT_ID = 'rising-fact-p41fc';

View File

@@ -214,7 +214,8 @@ app.get('/health', async (req, res) => {
try { try {
const token = await accountManager.getTokenForAccount(account); const token = await accountManager.getTokenForAccount(account);
const quotas = await getModelQuotas(token); const projectId = account.subscription?.projectId || null;
const quotas = await getModelQuotas(token, projectId);
// Format quotas for readability // Format quotas for readability
const formattedQuotas = {}; const formattedQuotas = {};
@@ -309,11 +310,11 @@ app.get('/account-limits', async (req, res) => {
try { try {
const token = await accountManager.getTokenForAccount(account); const token = await accountManager.getTokenForAccount(account);
// Fetch both quotas and subscription tier in parallel // Fetch subscription tier first to get project ID
const [quotas, subscription] = await Promise.all([ const subscription = await getSubscriptionTier(token);
getModelQuotas(token),
getSubscriptionTier(token) // Then fetch quotas with project ID for accurate quota info
]); const quotas = await getModelQuotas(token, subscription.projectId);
// Update account object with fresh data // Update account object with fresh data
account.subscription = { account.subscription = {