commons-proxy 2.0.0

package/src/account-manager/index.js

@@ -0,0 +1,462 @@
+/**
+ * Account Manager
+ * Manages multiple Antigravity accounts with configurable selection strategies,
+ * 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,
+    getRateLimitInfo as getLimitInfo,
+    getConsecutiveFailures as getFailures,
+    resetConsecutiveFailures as resetFailures,
+    incrementConsecutiveFailures as incrementFailures,
+    markAccountCoolingDown as markCoolingDown,
+    isAccountCoolingDown as checkCoolingDown,
+    clearAccountCooldown as clearCooldown,
+    getCooldownRemaining as getCooldownMs,
+    CooldownReason
+} from './rate-limits.js';
+import {
+    getTokenForAccount as fetchToken,
+    getProjectForAccount as fetchProject,
+    clearProjectCache as clearProject,
+    clearTokenCache as clearToken
+} from './credentials.js';
+import { createStrategy, getStrategyLabel, DEFAULT_STRATEGY } from './strategies/index.js';
+import { logger } from '../utils/logger.js';
+import { config } from '../config.js';
+
+export class AccountManager {
+    #accounts = [];
+    #currentIndex = 0;
+    #configPath;
+    #settings = {};
+    #initialized = false;
+    #strategy = null;
+    #strategyName = DEFAULT_STRATEGY;
+
+    // Per-account caches
+    #tokenCache = new Map(); // email -> { token, extractedAt }
+    #projectCache = new Map(); // email -> projectId
+
+    constructor(configPath = ACCOUNT_CONFIG_PATH, strategyName = null) {
+        this.#configPath = configPath;
+        // Strategy name can be set at construction or later via initialize
+        if (strategyName) {
+            this.#strategyName = strategyName;
+        }
+    }
+
+    /**
+     * Initialize the account manager by loading config
+     * @param {string} [strategyOverride] - Override strategy name (from CLI flag or env var)
+     */
+    async initialize(strategyOverride = null) {
+        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;
+        }
+
+        // Determine strategy: CLI override > env var > config file > default
+        const configStrategy = config?.accountSelection?.strategy;
+        const envStrategy = process.env.ACCOUNT_STRATEGY;
+        this.#strategyName = strategyOverride || envStrategy || configStrategy || this.#strategyName;
+
+        // Create the strategy instance
+        const strategyConfig = config?.accountSelection || {};
+        this.#strategy = createStrategy(this.#strategyName, strategyConfig);
+        logger.info(`[AccountManager] Using ${getStrategyLabel(this.#strategyName)} selection strategy`);
+
+        // Clear any expired rate limits
+        this.clearExpiredLimits();
+
+        this.#initialized = true;
+    }
+
+    /**
+     * Reload accounts from disk (force re-initialization)
+     * Useful when accounts.json is modified externally (e.g., by WebUI)
+     */
+    async reload() {
+        this.#initialized = false;
+        await this.initialize();
+        logger.info('[AccountManager] Accounts reloaded from disk');
+    }
+
+    /**
+     * Get the number of accounts
+     * @returns {number} Number of configured accounts
+     */
+    getAccountCount() {
+        return this.#accounts.length;
+    }
+
+    /**
+     * Check if all accounts are rate-limited
+     * @param {string} [modelId] - Optional model ID
+     * @returns {boolean} True if all accounts are rate-limited
+     */
+    isAllRateLimited(modelId = null) {
+        return checkAllRateLimited(this.#accounts, modelId);
+    }
+
+    /**
+     * Get list of available (non-rate-limited, non-invalid) accounts
+     * @param {string} [modelId] - Optional model ID
+     * @returns {Array<Object>} Array of available account objects
+     */
+    getAvailableAccounts(modelId = null) {
+        return getAvailable(this.#accounts, modelId);
+    }
+
+    /**
+     * 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);
+    }
+
+    /**
+     * Select an account using the configured strategy.
+     * This is the main method to use for account selection.
+     * @param {string} [modelId] - Model ID for the request
+     * @param {Object} [options] - Additional options
+     * @param {string} [options.sessionId] - Session ID for cache continuity
+     * @returns {{account: Object|null, waitMs: number}} Account to use and optional wait time
+     */
+    selectAccount(modelId = null, options = {}) {
+        if (!this.#strategy) {
+            throw new Error('AccountManager not initialized. Call initialize() first.');
+        }
+
+        const result = this.#strategy.selectAccount(this.#accounts, modelId, {
+            currentIndex: this.#currentIndex,
+            onSave: () => this.saveToDisk(),
+            ...options
+        });
+
+        this.#currentIndex = result.index;
+        return { account: result.account, waitMs: result.waitMs || 0 };
+    }
+
+    /**
+     * Notify the strategy of a successful request
+     * @param {Object} account - The account that was used
+     * @param {string} modelId - The model ID that was used
+     */
+    notifySuccess(account, modelId) {
+        if (this.#strategy) {
+            this.#strategy.onSuccess(account, modelId);
+        }
+        // Reset consecutive failures on success (matches opencode-cloudcode-auth)
+        if (account?.email) {
+            resetFailures(this.#accounts, account.email);
+        }
+    }
+
+    /**
+     * Notify the strategy of a rate limit
+     * @param {Object} account - The account that was rate-limited
+     * @param {string} modelId - The model ID that was rate-limited
+     */
+    notifyRateLimit(account, modelId) {
+        if (this.#strategy) {
+            this.#strategy.onRateLimit(account, modelId);
+        }
+    }
+
+    /**
+     * Notify the strategy of a failure
+     * @param {Object} account - The account that failed
+     * @param {string} modelId - The model ID that failed
+     */
+    notifyFailure(account, modelId) {
+        if (this.#strategy) {
+            this.#strategy.onFailure(account, modelId);
+        }
+    }
+
+    /**
+     * Get the consecutive failure count for an account
+     * Used for progressive backoff calculation
+     * @param {string} email - Account email
+     * @returns {number} Number of consecutive failures
+     */
+    getConsecutiveFailures(email) {
+        return getFailures(this.#accounts, email);
+    }
+
+    /**
+     * Increment the consecutive failure count without marking as rate limited
+     * Used for quick retries to track failures while staying on same account
+     * @param {string} email - Account email
+     * @returns {number} New consecutive failure count
+     */
+    incrementConsecutiveFailures(email) {
+        return incrementFailures(this.#accounts, email);
+    }
+
+    /**
+     * Get the current strategy name
+     * @returns {string} Strategy name
+     */
+    getStrategyName() {
+        return this.#strategyName;
+    }
+
+    /**
+     * Get the strategy display label
+     * @returns {string} Strategy display label
+     */
+    getStrategyLabel() {
+        return getStrategyLabel(this.#strategyName);
+    }
+
+    /**
+     * Get the health tracker from the current strategy (if available)
+     * Used by handlers for consecutive failure tracking
+     * Only available when using hybrid strategy
+     * @returns {Object|null} Health tracker instance or null if not available
+     */
+    getHealthTracker() {
+        if (this.#strategy && typeof this.#strategy.getHealthTracker === 'function') {
+            return this.#strategy.getHealthTracker();
+        }
+        return null;
+    }
+
+    /**
+     * 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)
+     * @param {string} [modelId] - Optional model ID to mark specific limit
+     */
+    markRateLimited(email, resetMs = null, modelId = null) {
+        markLimited(this.#accounts, email, resetMs, modelId);
+        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
+     * @param {string} [modelId] - Optional model ID
+     * @returns {number} Wait time in milliseconds
+     */
+    getMinWaitTimeMs(modelId = null) {
+        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);
+    }
+
+    // ============================================================================
+    // Cooldown Methods (matches opencode-cloudcode-auth)
+    // ============================================================================
+
+    /**
+     * Mark an account as cooling down for a specified duration
+     * Used for temporary backoff separate from rate limits
+     * @param {string} email - Email of the account
+     * @param {number} cooldownMs - Duration of cooldown in milliseconds
+     * @param {string} [reason] - Reason for the cooldown (use CooldownReason constants)
+     */
+    markAccountCoolingDown(email, cooldownMs, reason = CooldownReason.RATE_LIMIT) {
+        markCoolingDown(this.#accounts, email, cooldownMs, reason);
+    }
+
+    /**
+     * Check if an account is currently cooling down
+     * @param {string} email - Email of the account
+     * @returns {boolean} True if account is cooling down
+     */
+    isAccountCoolingDown(email) {
+        const account = this.#accounts.find(a => a.email === email);
+        return account ? checkCoolingDown(account) : false;
+    }
+
+    /**
+     * Clear the cooldown for an account
+     * @param {string} email - Email of the account
+     */
+    clearAccountCooldown(email) {
+        const account = this.#accounts.find(a => a.email === email);
+        if (account) {
+            clearCooldown(account);
+        }
+    }
+
+    /**
+     * Get time remaining until cooldown expires for an account
+     * @param {string} email - Email of the account
+     * @returns {number} Milliseconds until cooldown expires, 0 if not cooling down
+     */
+    getCooldownRemaining(email) {
+        const account = this.#accounts.find(a => a.email === email);
+        return account ? getCooldownMs(account) : 0;
+    }
+
+    /**
+     * 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) {
+        // Pass onSave callback to persist managedProjectId in refresh token
+        return fetchProject(account, token, this.#projectCache, () => this.saveToDisk());
+    }
+
+    /**
+     * 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 invalid = this.getInvalidAccounts();
+
+        // Count accounts that have any active model-specific rate limits
+        const rateLimited = this.#accounts.filter(a => {
+            if (!a.modelRateLimits) return false;
+            return Object.values(a.modelRateLimits).some(
+                limit => limit.isRateLimited && limit.resetTime > Date.now()
+            );
+        });
+
+        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,
+                enabled: a.enabled !== false,  // Default to true if undefined
+                projectId: a.projectId || null,
+                modelRateLimits: a.modelRateLimits || {},
+                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;
+    }
+}
+
+// Re-export CooldownReason for use by handlers
+export { CooldownReason };
+
+export default AccountManager;

package/src/account-manager/onboarding.js

@@ -0,0 +1,112 @@
+/**
+ * User Onboarding
+ *
+ * Handles provisioning of managed projects for accounts that don't have one.
+ */
+
+import {
+    ONBOARD_USER_ENDPOINTS,
+    ANTIGRAVITY_HEADERS
+} from '../constants.js';
+import { logger } from '../utils/logger.js';
+import { sleep } from '../utils/helpers.js';
+
+/**
+ * Get the default tier ID from allowed tiers list
+ *
+ * @param {Array} allowedTiers - List of allowed tiers from loadCodeAssist
+ * @returns {string|undefined} Default tier ID
+ */
+export function getDefaultTierId(allowedTiers) {
+    if (!allowedTiers || allowedTiers.length === 0) {
+        return undefined;
+    }
+
+    // Find the tier marked as default
+    for (const tier of allowedTiers) {
+        if (tier?.isDefault) {
+            return tier.id;
+        }
+    }
+
+    // Fall back to first tier
+    return allowedTiers[0]?.id;
+}
+
+/**
+ * Onboard a user to get a managed project
+ *
+ * @param {string} token - OAuth access token
+ * @param {string} tierId - Tier ID (raw API value, e.g., 'free-tier', 'standard-tier', 'g1-pro-tier')
+ * @param {string} [projectId] - Optional GCP project ID (required for non-free tiers)
+ * @param {number} [maxAttempts=10] - Maximum polling attempts
+ * @param {number} [delayMs=5000] - Delay between polling attempts
+ * @returns {Promise<string|null>} Managed project ID or null if failed
+ */
+export async function onboardUser(token, tierId, projectId = undefined, maxAttempts = 10, delayMs = 5000) {
+    const metadata = {
+        ideType: 'IDE_UNSPECIFIED',
+        platform: 'PLATFORM_UNSPECIFIED',
+        pluginType: 'GEMINI'
+    };
+
+    if (projectId) {
+        metadata.duetProject = projectId;
+    }
+
+    const requestBody = {
+        tierId,
+        metadata
+    };
+    // Note: Do NOT add cloudaicompanionProject to requestBody
+    // Reference implementation only sets metadata.duetProject, not the body field
+    // Adding cloudaicompanionProject causes 400 errors for auto-provisioned tiers (g1-pro, g1-ultra)
+
+    logger.debug(`[Onboarding] Starting onboard with tierId: ${tierId}, projectId: ${projectId}`);
+
+    for (const endpoint of ONBOARD_USER_ENDPOINTS) {
+        for (let attempt = 0; attempt < maxAttempts; attempt++) {
+            try {
+                const response = await fetch(`${endpoint}/v1internal:onboardUser`, {
+                    method: 'POST',
+                    headers: {
+                        'Authorization': `Bearer ${token}`,
+                        'Content-Type': 'application/json',
+                        ...ANTIGRAVITY_HEADERS
+                    },
+                    body: JSON.stringify(requestBody)
+                });
+
+                if (!response.ok) {
+                    const errorText = await response.text();
+                    logger.warn(`[Onboarding] onboardUser failed at ${endpoint}: ${response.status} - ${errorText}`);
+                    break; // Try next endpoint
+                }
+
+                const data = await response.json();
+                logger.debug(`[Onboarding] onboardUser response (attempt ${attempt + 1}):`, JSON.stringify(data));
+
+                // Check if onboarding is complete
+                const managedProjectId = data.response?.cloudaicompanionProject?.id;
+                if (data.done && managedProjectId) {
+                    return managedProjectId;
+                }
+                if (data.done && projectId) {
+                    return projectId;
+                }
+
+                // Not done yet, wait and retry
+                if (attempt < maxAttempts - 1) {
+                    logger.debug(`[Onboarding] onboardUser not complete, waiting ${delayMs}ms...`);
+                    await sleep(delayMs);
+                }
+            } catch (error) {
+                logger.warn(`[Onboarding] onboardUser error at ${endpoint}:`, error.message);
+                break; // Try next endpoint
+            }
+        }
+    }
+
+    logger.warn(`[Onboarding] All onboarding attempts failed for tierId: ${tierId}`);
+    return null;
+}