antigravity-claude-proxy 1.1.5 → 1.2.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "antigravity-claude-proxy",
-  "version": "1.1.5",
+  "version": "1.2.0",
   "description": "Proxy server to use Antigravity's Claude models with Claude Code CLI",
   "main": "src/index.js",
   "type": "module",
@@ -14,11 +14,11 @@
   "scripts": {
     "start": "node src/index.js",
     "dev": "node --watch src/index.js",
-    "accounts": "node src/accounts-cli.js",
-    "accounts:add": "node src/accounts-cli.js add",
-    "accounts:list": "node src/accounts-cli.js list",
-    "accounts:remove": "node src/accounts-cli.js remove",
-    "accounts:verify": "node src/accounts-cli.js verify",
+    "accounts": "node src/cli/accounts.js",
+    "accounts:add": "node src/cli/accounts.js add",
+    "accounts:list": "node src/cli/accounts.js list",
+    "accounts:remove": "node src/cli/accounts.js remove",
+    "accounts:verify": "node src/cli/accounts.js verify",
     "test": "node tests/run-all.cjs",
     "test:signatures": "node tests/test-thinking-signatures.cjs",
     "test:multiturn": "node tests/test-multiturn-thinking-tools.cjs",

package/src/account-manager/credentials.js

@@ -0,0 +1,171 @@
+/**
+ * Credentials Management
+ *
+ * Handles OAuth token handling and project discovery.
+ */
+
+import {
+    ANTIGRAVITY_DB_PATH,
+    TOKEN_REFRESH_INTERVAL_MS,
+    ANTIGRAVITY_ENDPOINT_FALLBACKS,
+    ANTIGRAVITY_HEADERS,
+    DEFAULT_PROJECT_ID
+} from '../constants.js';
+import { refreshAccessToken } from '../auth/oauth.js';
+import { getAuthStatus } from '../auth/database.js';
+import { logger } from '../utils/logger.js';
+
+/**
+ * Get OAuth token for an account
+ *
+ * @param {Object} account - Account object with email and credentials
+ * @param {Map} tokenCache - Token cache map
+ * @param {Function} onInvalid - Callback when account is invalid (email, reason)
+ * @param {Function} onSave - Callback to save changes
+ * @returns {Promise<string>} OAuth access token
+ * @throws {Error} If token refresh fails
+ */
+export async function getTokenForAccount(account, tokenCache, onInvalid, onSave) {
+    // Check cache first
+    const cached = 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;
+                if (onSave) await onSave();
+            }
+            logger.success(`[AccountManager] Refreshed OAuth token for: ${account.email}`);
+        } catch (error) {
+            logger.error(`[AccountManager] Failed to refresh token for ${account.email}:`, error.message);
+            // Mark account as invalid (credentials need re-auth)
+            if (onInvalid) onInvalid(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
+    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
+ * @param {Map} projectCache - Project cache map
+ * @returns {Promise<string>} Project ID
+ */
+export async function getProjectForAccount(account, token, projectCache) {
+    // Check cache first
+    const cached = projectCache.get(account.email);
+    if (cached) {
+        return cached;
+    }
+
+    // OAuth or manual accounts may have projectId specified
+    if (account.projectId) {
+        projectCache.set(account.email, account.projectId);
+        return account.projectId;
+    }
+
+    // Discover project via loadCodeAssist API
+    const project = await discoverProject(token);
+    projectCache.set(account.email, project);
+    return project;
+}
+
+/**
+ * Discover project ID via Cloud Code API
+ *
+ * @param {string} token - OAuth access token
+ * @returns {Promise<string>} Project ID
+ */
+export async function 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) {
+            logger.warn(`[AccountManager] Project discovery failed at ${endpoint}:`, error.message);
+        }
+    }
+
+    logger.info(`[AccountManager] Using default project: ${DEFAULT_PROJECT_ID}`);
+    return DEFAULT_PROJECT_ID;
+}
+
+/**
+ * Clear project cache for an account
+ *
+ * @param {Map} projectCache - Project cache map
+ * @param {string|null} email - Email to clear cache for, or null to clear all
+ */
+export function clearProjectCache(projectCache, email = null) {
+    if (email) {
+        projectCache.delete(email);
+    } else {
+        projectCache.clear();
+    }
+}
+
+/**
+ * Clear token cache for an account
+ *
+ * @param {Map} tokenCache - Token cache map
+ * @param {string|null} email - Email to clear cache for, or null to clear all
+ */
+export function clearTokenCache(tokenCache, email = null) {
+    if (email) {
+        tokenCache.delete(email);
+    } else {
+        tokenCache.clear();
+    }
+}

package/src/account-manager/index.js

@@ -0,0 +1,293 @@
+/**
+ * Account Manager
+ * Manages multiple Antigravity accounts with sticky selection,
+ * automatic failover, and smart cooldown for rate-limited accounts.
+ */
+
+import { ACCOUNT_CONFIG_PATH } from '../constants.js';
+import { loadAccounts, loadDefaultAccount, saveAccounts } from './storage.js';
+import {
+    isAllRateLimited as checkAllRateLimited,
+    getAvailableAccounts as getAvailable,
+    getInvalidAccounts as getInvalid,
+    clearExpiredLimits as clearLimits,
+    resetAllRateLimits as resetLimits,
+    markRateLimited as markLimited,
+    markInvalid as markAccountInvalid,
+    getMinWaitTimeMs as getMinWait
+} from './rate-limits.js';
+import {
+    getTokenForAccount as fetchToken,
+    getProjectForAccount as fetchProject,
+    clearProjectCache as clearProject,
+    clearTokenCache as clearToken
+} from './credentials.js';
+import {
+    pickNext as selectNext,
+    getCurrentStickyAccount as getSticky,
+    shouldWaitForCurrentAccount as shouldWait,
+    pickStickyAccount as selectSticky
+} from './selection.js';
+import { logger } from '../utils/logger.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;
+
+        const { accounts, settings, activeIndex } = await loadAccounts(this.#configPath);
+
+        this.#accounts = accounts;
+        this.#settings = settings;
+        this.#currentIndex = activeIndex;
+
+        // If config exists but has no accounts, fall back to Antigravity database
+        if (this.#accounts.length === 0) {
+            logger.warn('[AccountManager] No accounts in config. Falling back to Antigravity database');
+            const { accounts: defaultAccounts, tokenCache } = loadDefaultAccount();
+            this.#accounts = defaultAccounts;
+            this.#tokenCache = tokenCache;
+        }
+
+        // Clear any expired rate limits
+        this.clearExpiredLimits();
+
+        this.#initialized = true;
+    }
+
+    /**
+     * 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() {
+        return checkAllRateLimited(this.#accounts);
+    }
+
+    /**
+     * Get list of available (non-rate-limited, non-invalid) accounts
+     * @returns {Array<Object>} Array of available account objects
+     */
+    getAvailableAccounts() {
+        return getAvailable(this.#accounts);
+    }
+
+    /**
+     * Get list of invalid accounts
+     * @returns {Array<Object>} Array of invalid account objects
+     */
+    getInvalidAccounts() {
+        return getInvalid(this.#accounts);
+    }
+
+    /**
+     * Clear expired rate limits
+     * @returns {number} Number of rate limits cleared
+     */
+    clearExpiredLimits() {
+        const cleared = clearLimits(this.#accounts);
+        if (cleared > 0) {
+            this.saveToDisk();
+        }
+        return cleared;
+    }
+
+    /**
+     * Clear all rate limits to force a fresh check
+     * (Optimistic retry strategy)
+     * @returns {void}
+     */
+    resetAllRateLimits() {
+        resetLimits(this.#accounts);
+    }
+
+    /**
+     * 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() {
+        const { account, newIndex } = selectNext(this.#accounts, this.#currentIndex, () => this.saveToDisk());
+        this.#currentIndex = newIndex;
+        return account;
+    }
+
+    /**
+     * 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() {
+        const { account, newIndex } = getSticky(this.#accounts, this.#currentIndex, () => this.saveToDisk());
+        this.#currentIndex = newIndex;
+        return account;
+    }
+
+    /**
+     * 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() {
+        return shouldWait(this.#accounts, this.#currentIndex);
+    }
+
+    /**
+     * 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() {
+        const { account, waitMs, newIndex } = selectSticky(this.#accounts, this.#currentIndex, () => this.saveToDisk());
+        this.#currentIndex = newIndex;
+        return { account, waitMs };
+    }
+
+    /**
+     * 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) {
+        markLimited(this.#accounts, email, resetMs, this.#settings);
+        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') {
+        markAccountInvalid(this.#accounts, email, reason);
+        this.saveToDisk();
+    }
+
+    /**
+     * Get the minimum wait time until any account becomes available
+     * @returns {number} Wait time in milliseconds
+     */
+    getMinWaitTimeMs() {
+        return getMinWait(this.#accounts);
+    }
+
+    /**
+     * Get OAuth token for an account
+     * @param {Object} account - Account object with email and credentials
+     * @returns {Promise<string>} OAuth access token
+     * @throws {Error} If token refresh fails
+     */
+    async getTokenForAccount(account) {
+        return fetchToken(
+            account,
+            this.#tokenCache,
+            (email, reason) => this.markInvalid(email, reason),
+            () => this.saveToDisk()
+        );
+    }
+
+    /**
+     * Get project ID for an account
+     * @param {Object} account - Account object
+     * @param {string} token - OAuth access token
+     * @returns {Promise<string>} Project ID
+     */
+    async getProjectForAccount(account, token) {
+        return fetchProject(account, token, this.#projectCache);
+    }
+
+    /**
+     * 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) {
+        clearProject(this.#projectCache, email);
+    }
+
+    /**
+     * 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) {
+        clearToken(this.#tokenCache, email);
+    }
+
+    /**
+     * Save current state to disk (async)
+     * @returns {Promise<void>}
+     */
+    async saveToDisk() {
+        await saveAccounts(this.#configPath, this.#accounts, this.#settings, this.#currentIndex);
+    }
+
+    /**
+     * 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<Object>} Array of account objects
+     */
+    getAllAccounts() {
+        return this.#accounts;
+    }
+}
+
+export default AccountManager;

package/src/account-manager/rate-limits.js

@@ -0,0 +1,157 @@
+/**
+ * Rate Limit Management
+ *
+ * Handles rate limit tracking and state management for accounts.
+ */
+
+import { DEFAULT_COOLDOWN_MS } from '../constants.js';
+import { formatDuration } from '../utils/helpers.js';
+import { logger } from '../utils/logger.js';
+
+/**
+ * Check if all accounts are rate-limited
+ *
+ * @param {Array} accounts - Array of account objects
+ * @returns {boolean} True if all accounts are rate-limited
+ */
+export function isAllRateLimited(accounts) {
+    if (accounts.length === 0) return true;
+    return accounts.every(acc => acc.isRateLimited);
+}
+
+/**
+ * Get list of available (non-rate-limited, non-invalid) accounts
+ *
+ * @param {Array} accounts - Array of account objects
+ * @returns {Array} Array of available account objects
+ */
+export function getAvailableAccounts(accounts) {
+    return accounts.filter(acc => !acc.isRateLimited && !acc.isInvalid);
+}
+
+/**
+ * Get list of invalid accounts
+ *
+ * @param {Array} accounts - Array of account objects
+ * @returns {Array} Array of invalid account objects
+ */
+export function getInvalidAccounts(accounts) {
+    return accounts.filter(acc => acc.isInvalid);
+}
+
+/**
+ * Clear expired rate limits
+ *
+ * @param {Array} accounts - Array of account objects
+ * @returns {number} Number of rate limits cleared
+ */
+export function clearExpiredLimits(accounts) {
+    const now = Date.now();
+    let cleared = 0;
+
+    for (const account of accounts) {
+        if (account.isRateLimited && account.rateLimitResetTime && account.rateLimitResetTime <= now) {
+            account.rateLimitResetTime = null;
+            cleared++;
+            logger.success(`[AccountManager] Rate limit expired for: ${account.email}`);
+        }
+    }
+
+    return cleared;
+}
+
+/**
+ * Clear all rate limits to force a fresh check (optimistic retry strategy)
+ *
+ * @param {Array} accounts - Array of account objects
+ */
+export function resetAllRateLimits(accounts) {
+    for (const account of accounts) {
+        account.isRateLimited = false;
+        account.rateLimitResetTime = null;
+    }
+    logger.warn('[AccountManager] Reset all rate limits for optimistic retry');
+}
+
+/**
+ * Mark an account as rate-limited
+ *
+ * @param {Array} accounts - Array of account objects
+ * @param {string} email - Email of the account to mark
+ * @param {number|null} resetMs - Time in ms until rate limit resets (optional)
+ * @param {Object} settings - Settings object with cooldownDurationMs
+ * @returns {boolean} True if account was found and marked
+ */
+export function markRateLimited(accounts, email, resetMs = null, settings = {}) {
+    const account = accounts.find(a => a.email === email);
+    if (!account) return false;
+
+    account.isRateLimited = true;
+    const cooldownMs = resetMs || settings.cooldownDurationMs || DEFAULT_COOLDOWN_MS;
+    account.rateLimitResetTime = Date.now() + cooldownMs;
+
+    logger.warn(
+        `[AccountManager] Rate limited: ${email}. Available in ${formatDuration(cooldownMs)}`
+    );
+
+    return true;
+}
+
+/**
+ * Mark an account as invalid (credentials need re-authentication)
+ *
+ * @param {Array} accounts - Array of account objects
+ * @param {string} email - Email of the account to mark
+ * @param {string} reason - Reason for marking as invalid
+ * @returns {boolean} True if account was found and marked
+ */
+export function markInvalid(accounts, email, reason = 'Unknown error') {
+    const account = accounts.find(a => a.email === email);
+    if (!account) return false;
+
+    account.isInvalid = true;
+    account.invalidReason = reason;
+    account.invalidAt = Date.now();
+
+    logger.error(
+        `[AccountManager] ⚠ Account INVALID: ${email}`
+    );
+    logger.error(
+        `[AccountManager]   Reason: ${reason}`
+    );
+    logger.error(
+        `[AccountManager]   Run 'npm run accounts' to re-authenticate this account`
+    );
+
+    return true;
+}
+
+/**
+ * Get the minimum wait time until any account becomes available
+ *
+ * @param {Array} accounts - Array of account objects
+ * @returns {number} Wait time in milliseconds
+ */
+export function getMinWaitTimeMs(accounts) {
+    if (!isAllRateLimited(accounts)) return 0;
+
+    const now = Date.now();
+    let minWait = Infinity;
+    let soonestAccount = null;
+
+    for (const account of accounts) {
+        if (account.rateLimitResetTime) {
+            const wait = account.rateLimitResetTime - now;
+            if (wait > 0 && wait < minWait) {
+                minWait = wait;
+                soonestAccount = account;
+            }
+        }
+    }
+
+    if (soonestAccount) {
+        logger.info(`[AccountManager] Shortest wait: ${formatDuration(minWait)} (account: ${soonestAccount.email})`);
+    }
+
+    return minWait === Infinity ? DEFAULT_COOLDOWN_MS : minWait;
+}