/** * Account Manager * Manages multiple Antigravity accounts with sticky selection, * automatic failover, and smart cooldown for rate-limited accounts. */ import { readFile, writeFile, mkdir, access } from 'fs/promises'; import { constants as fsConstants } from 'fs'; import { dirname } from 'path'; import { ACCOUNT_CONFIG_PATH, ANTIGRAVITY_DB_PATH, DEFAULT_COOLDOWN_MS, TOKEN_REFRESH_INTERVAL_MS, ANTIGRAVITY_ENDPOINT_FALLBACKS, ANTIGRAVITY_HEADERS, DEFAULT_PROJECT_ID, MAX_WAIT_BEFORE_ERROR_MS } from './constants.js'; import { refreshAccessToken } from './oauth.js'; import { formatDuration } from './utils/helpers.js'; import { getAuthStatus } from './db/database.js'; export class AccountManager { #accounts = []; #currentIndex = 0; #configPath; #settings = {}; #initialized = false; // Per-account caches #tokenCache = new Map(); // email -> { token, extractedAt } #projectCache = new Map(); // email -> projectId constructor(configPath = ACCOUNT_CONFIG_PATH) { this.#configPath = configPath; } /** * Initialize the account manager by loading config */ async initialize() { if (this.#initialized) return; try { // Check if config file exists using async access await access(this.#configPath, fsConstants.F_OK); const configData = await readFile(this.#configPath, 'utf-8'); const config = JSON.parse(configData); this.#accounts = (config.accounts || []).map(acc => ({ ...acc, isRateLimited: acc.isRateLimited || false, rateLimitResetTime: acc.rateLimitResetTime || null, lastUsed: acc.lastUsed || null })); this.#settings = config.settings || {}; this.#currentIndex = config.activeIndex || 0; // Clamp currentIndex to valid range if (this.#currentIndex >= this.#accounts.length) { this.#currentIndex = 0; } console.log(`[AccountManager] Loaded ${this.#accounts.length} account(s) from config`); // If config exists but has no accounts, fall back to Antigravity database if (this.#accounts.length === 0) { console.log('[AccountManager] No accounts in config. Falling back to Antigravity database'); await this.#loadDefaultAccount(); } } catch (error) { if (error.code === 'ENOENT') { // No config file - use single account from Antigravity database console.log('[AccountManager] No config file found. Using Antigravity database (single account mode)'); } else { console.error('[AccountManager] Failed to load config:', error.message); } // Fall back to default account await this.#loadDefaultAccount(); } // Clear any expired rate limits this.clearExpiredLimits(); this.#initialized = true; } /** * Load the default account from Antigravity's database */ async #loadDefaultAccount() { try { const authData = getAuthStatus(); if (authData?.apiKey) { this.#accounts = [{ email: authData.email || 'default@antigravity', source: 'database', isRateLimited: false, rateLimitResetTime: null, lastUsed: null }]; // Pre-cache the token this.#tokenCache.set(this.#accounts[0].email, { token: authData.apiKey, extractedAt: Date.now() }); console.log(`[AccountManager] Loaded default account: ${this.#accounts[0].email}`); } } catch (error) { console.error('[AccountManager] Failed to load default account:', error.message); // Create empty account list - will fail on first request this.#accounts = []; } } /** * Get the number of accounts * @returns {number} Number of configured accounts */ getAccountCount() { return this.#accounts.length; } /** * Check if all accounts are rate-limited * @returns {boolean} True if all accounts are rate-limited */ isAllRateLimited() { if (this.#accounts.length === 0) return true; return this.#accounts.every(acc => acc.isRateLimited); } /** * Get list of available (non-rate-limited, non-invalid) accounts * @returns {Array} Array of available account objects */ getAvailableAccounts() { return this.#accounts.filter(acc => !acc.isRateLimited && !acc.isInvalid); } /** * Get list of invalid accounts * @returns {Array} Array of invalid account objects */ getInvalidAccounts() { return this.#accounts.filter(acc => acc.isInvalid); } /** * Clear expired rate limits * @returns {number} Number of rate limits cleared */ clearExpiredLimits() { const now = Date.now(); let cleared = 0; for (const account of this.#accounts) { if (account.isRateLimited && account.rateLimitResetTime && account.rateLimitResetTime <= now) { account.isRateLimited = false; account.rateLimitResetTime = null; cleared++; console.log(`[AccountManager] Rate limit expired for: ${account.email}`); } } if (cleared > 0) { this.saveToDisk(); } return cleared; } /** * Clear all rate limits to force a fresh check * (Optimistic retry strategy) * @returns {void} */ resetAllRateLimits() { for (const account of this.#accounts) { account.isRateLimited = false; // distinct from "clearing" expired limits, we blindly reset here // we keep the time? User said "clear isRateLimited value, and rateLimitResetTime" // So we clear both. account.rateLimitResetTime = null; } console.log('[AccountManager] Reset all rate limits for optimistic retry'); } /** * Pick the next available account (fallback when current is unavailable). * Sets activeIndex to the selected account's index. * @returns {Object|null} The next available account or null if none available */ pickNext() { this.clearExpiredLimits(); const available = this.getAvailableAccounts(); if (available.length === 0) { return null; } // Clamp index to valid range if (this.#currentIndex >= this.#accounts.length) { this.#currentIndex = 0; } // Find next available account starting from index AFTER current for (let i = 1; i <= this.#accounts.length; i++) { const idx = (this.#currentIndex + i) % this.#accounts.length; const account = this.#accounts[idx]; if (!account.isRateLimited && !account.isInvalid) { // Set activeIndex to this account (not +1) this.#currentIndex = idx; account.lastUsed = Date.now(); const position = idx + 1; const total = this.#accounts.length; console.log(`[AccountManager] Using account: ${account.email} (${position}/${total})`); // Persist the change (don't await to avoid blocking) this.saveToDisk(); return account; } } return null; } /** * Get the current account without advancing the index (sticky selection). * Used for cache continuity - sticks to the same account until rate-limited. * @returns {Object|null} The current account or null if unavailable/rate-limited */ getCurrentStickyAccount() { this.clearExpiredLimits(); if (this.#accounts.length === 0) { return null; } // Clamp index to valid range if (this.#currentIndex >= this.#accounts.length) { this.#currentIndex = 0; } // Get current account directly (activeIndex = current account) const account = this.#accounts[this.#currentIndex]; // Return if available if (account && !account.isRateLimited && !account.isInvalid) { account.lastUsed = Date.now(); // Persist the change (don't await to avoid blocking) this.saveToDisk(); return account; } return null; } /** * Check if we should wait for the current account's rate limit to reset. * Used for sticky account selection - wait if rate limit is short (≤ threshold). * @returns {{shouldWait: boolean, waitMs: number, account: Object|null}} */ shouldWaitForCurrentAccount() { if (this.#accounts.length === 0) { return { shouldWait: false, waitMs: 0, account: null }; } // Clamp index to valid range if (this.#currentIndex >= this.#accounts.length) { this.#currentIndex = 0; } // Get current account directly (activeIndex = current account) const account = this.#accounts[this.#currentIndex]; if (!account || account.isInvalid) { return { shouldWait: false, waitMs: 0, account: null }; } if (account.isRateLimited && account.rateLimitResetTime) { const waitMs = account.rateLimitResetTime - Date.now(); // If wait time is within threshold, recommend waiting if (waitMs > 0 && waitMs <= MAX_WAIT_BEFORE_ERROR_MS) { return { shouldWait: true, waitMs, account }; } } return { shouldWait: false, waitMs: 0, account }; } /** * Pick an account with sticky selection preference. * Prefers the current account for cache continuity, only switches when: * - Current account is rate-limited for > 2 minutes * - Current account is invalid * @returns {{account: Object|null, waitMs: number}} Account to use and optional wait time */ pickStickyAccount() { // First try to get the current sticky account const stickyAccount = this.getCurrentStickyAccount(); if (stickyAccount) { return { account: stickyAccount, waitMs: 0 }; } // Check if we should wait for current account const waitInfo = this.shouldWaitForCurrentAccount(); if (waitInfo.shouldWait) { console.log(`[AccountManager] Waiting ${formatDuration(waitInfo.waitMs)} for sticky account: ${waitInfo.account.email}`); return { account: null, waitMs: waitInfo.waitMs }; } // Current account unavailable for too long, switch to next available const nextAccount = this.pickNext(); if (nextAccount) { console.log(`[AccountManager] Switched to new account for cache: ${nextAccount.email}`); } return { account: nextAccount, waitMs: 0 }; } /** * Mark an account as rate-limited * @param {string} email - Email of the account to mark * @param {number|null} resetMs - Time in ms until rate limit resets (optional) */ markRateLimited(email, resetMs = null) { const account = this.#accounts.find(a => a.email === email); if (!account) return; account.isRateLimited = true; const cooldownMs = resetMs || this.#settings.cooldownDurationMs || DEFAULT_COOLDOWN_MS; account.rateLimitResetTime = Date.now() + cooldownMs; console.log( `[AccountManager] Rate limited: ${email}. Available in ${formatDuration(cooldownMs)}` ); this.saveToDisk(); } /** * Mark an account as invalid (credentials need re-authentication) * @param {string} email - Email of the account to mark * @param {string} reason - Reason for marking as invalid */ markInvalid(email, reason = 'Unknown error') { const account = this.#accounts.find(a => a.email === email); if (!account) return; account.isInvalid = true; account.invalidReason = reason; account.invalidAt = Date.now(); console.log( `[AccountManager] ⚠ Account INVALID: ${email}` ); console.log( `[AccountManager] Reason: ${reason}` ); console.log( `[AccountManager] Run 'npm run accounts' to re-authenticate this account` ); this.saveToDisk(); } /** * Get the minimum wait time until any account becomes available * @returns {number} Wait time in milliseconds */ getMinWaitTimeMs() { if (!this.isAllRateLimited()) return 0; const now = Date.now(); let minWait = Infinity; let soonestAccount = null; for (const account of this.#accounts) { if (account.rateLimitResetTime) { const wait = account.rateLimitResetTime - now; if (wait > 0 && wait < minWait) { minWait = wait; soonestAccount = account; } } } if (soonestAccount) { console.log(`[AccountManager] Shortest wait: ${formatDuration(minWait)} (account: ${soonestAccount.email})`); } return minWait === Infinity ? DEFAULT_COOLDOWN_MS : minWait; } /** * Get OAuth token for an account * @param {Object} account - Account object with email and credentials * @returns {Promise} OAuth access token * @throws {Error} If token refresh fails */ async getTokenForAccount(account) { // Check cache first const cached = this.#tokenCache.get(account.email); if (cached && (Date.now() - cached.extractedAt) < TOKEN_REFRESH_INTERVAL_MS) { return cached.token; } // Get fresh token based on source let token; if (account.source === 'oauth' && account.refreshToken) { // OAuth account - use refresh token to get new access token try { const tokens = await refreshAccessToken(account.refreshToken); token = tokens.accessToken; // Clear invalid flag on success if (account.isInvalid) { account.isInvalid = false; account.invalidReason = null; await this.saveToDisk(); } console.log(`[AccountManager] Refreshed OAuth token for: ${account.email}`); } catch (error) { console.error(`[AccountManager] Failed to refresh token for ${account.email}:`, error.message); // Mark account as invalid (credentials need re-auth) this.markInvalid(account.email, error.message); throw new Error(`AUTH_INVALID: ${account.email}: ${error.message}`); } } else if (account.source === 'manual' && account.apiKey) { token = account.apiKey; } else { // Extract from database const dbPath = account.dbPath || ANTIGRAVITY_DB_PATH; const authData = getAuthStatus(dbPath); token = authData.apiKey; } // Cache the token this.#tokenCache.set(account.email, { token, extractedAt: Date.now() }); return token; } /** * Get project ID for an account * @param {Object} account - Account object * @param {string} token - OAuth access token * @returns {Promise} Project ID */ async getProjectForAccount(account, token) { // Check cache first const cached = this.#projectCache.get(account.email); if (cached) { return cached; } // OAuth or manual accounts may have projectId specified if (account.projectId) { this.#projectCache.set(account.email, account.projectId); return account.projectId; } // Discover project via loadCodeAssist API const project = await this.#discoverProject(token); this.#projectCache.set(account.email, project); return project; } /** * Discover project ID via Cloud Code API */ async #discoverProject(token) { for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) { try { const response = await fetch(`${endpoint}/v1internal:loadCodeAssist`, { method: 'POST', headers: { 'Authorization': `Bearer ${token}`, 'Content-Type': 'application/json', ...ANTIGRAVITY_HEADERS }, body: JSON.stringify({ metadata: { ideType: 'IDE_UNSPECIFIED', platform: 'PLATFORM_UNSPECIFIED', pluginType: 'GEMINI' } }) }); if (!response.ok) continue; const data = await response.json(); if (typeof data.cloudaicompanionProject === 'string') { return data.cloudaicompanionProject; } if (data.cloudaicompanionProject?.id) { return data.cloudaicompanionProject.id; } } catch (error) { console.log(`[AccountManager] Project discovery failed at ${endpoint}:`, error.message); } } console.log(`[AccountManager] Using default project: ${DEFAULT_PROJECT_ID}`); return DEFAULT_PROJECT_ID; } /** * Clear project cache for an account (useful on auth errors) * @param {string|null} email - Email to clear cache for, or null to clear all */ clearProjectCache(email = null) { if (email) { this.#projectCache.delete(email); } else { this.#projectCache.clear(); } } /** * Clear token cache for an account (useful on auth errors) * @param {string|null} email - Email to clear cache for, or null to clear all */ clearTokenCache(email = null) { if (email) { this.#tokenCache.delete(email); } else { this.#tokenCache.clear(); } } /** * Save current state to disk (async) * @returns {Promise} */ async saveToDisk() { try { // Ensure directory exists const dir = dirname(this.#configPath); await mkdir(dir, { recursive: true }); const config = { accounts: this.#accounts.map(acc => ({ email: acc.email, source: acc.source, dbPath: acc.dbPath || null, refreshToken: acc.source === 'oauth' ? acc.refreshToken : undefined, apiKey: acc.source === 'manual' ? acc.apiKey : undefined, projectId: acc.projectId || undefined, addedAt: acc.addedAt || undefined, isRateLimited: acc.isRateLimited, rateLimitResetTime: acc.rateLimitResetTime, isInvalid: acc.isInvalid || false, invalidReason: acc.invalidReason || null, lastUsed: acc.lastUsed })), settings: this.#settings, activeIndex: this.#currentIndex }; await writeFile(this.#configPath, JSON.stringify(config, null, 2)); } catch (error) { console.error('[AccountManager] Failed to save config:', error.message); } } /** * Get status object for logging/API * @returns {{accounts: Array, settings: Object}} Status object with accounts and settings */ getStatus() { const available = this.getAvailableAccounts(); const rateLimited = this.#accounts.filter(a => a.isRateLimited); const invalid = this.getInvalidAccounts(); return { total: this.#accounts.length, available: available.length, rateLimited: rateLimited.length, invalid: invalid.length, summary: `${this.#accounts.length} total, ${available.length} available, ${rateLimited.length} rate-limited, ${invalid.length} invalid`, accounts: this.#accounts.map(a => ({ email: a.email, source: a.source, isRateLimited: a.isRateLimited, rateLimitResetTime: a.rateLimitResetTime, isInvalid: a.isInvalid || false, invalidReason: a.invalidReason || null, lastUsed: a.lastUsed })) }; } /** * Get settings * @returns {Object} Current settings object */ getSettings() { return { ...this.#settings }; } /** * Get all accounts (internal use for quota fetching) * Returns the full account objects including credentials * @returns {Array} Array of account objects */ getAllAccounts() { return this.#accounts; } } export default AccountManager;