@kamel-ahmed/proxy-claude 1.0.0

package/src/account-manager/credentials.js

@@ -0,0 +1,243 @@
+/**
+ * Credentials Management
+ *
+ * Handles OAuth token handling and project discovery.
+ */
+
+import {
+    ANTIGRAVITY_DB_PATH,
+    TOKEN_REFRESH_INTERVAL_MS,
+    LOAD_CODE_ASSIST_ENDPOINTS,
+    LOAD_CODE_ASSIST_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';
+import { isNetworkError } from '../utils/helpers.js';
+import { onboardUser, getDefaultTierId } from './onboarding.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) {
+            // Check if it's a transient network error
+            if (isNetworkError(error)) {
+                logger.warn(`[AccountManager] Failed to refresh token for ${account.email} due to network error: ${error.message}`);
+                // Do NOT mark as invalid, just throw so caller knows it failed
+                throw new Error(`AUTH_NETWORK_ERROR: ${error.message}`);
+            }
+
+            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) {
+    let lastError = null;
+    let gotSuccessfulResponse = false;
+    let loadCodeAssistData = null;
+
+    for (const endpoint of LOAD_CODE_ASSIST_ENDPOINTS) {
+        try {
+            const response = await fetch(`${endpoint}/v1internal:loadCodeAssist`, {
+                method: 'POST',
+                headers: {
+                    'Authorization': `Bearer ${token}`,
+                    'Content-Type': 'application/json',
+                    ...LOAD_CODE_ASSIST_HEADERS
+                },
+                body: JSON.stringify({
+                    metadata: {
+                        ideType: 'IDE_UNSPECIFIED',
+                        platform: 'PLATFORM_UNSPECIFIED',
+                        pluginType: 'GEMINI',
+                        duetProject: DEFAULT_PROJECT_ID
+                    }
+                })
+            });
+
+            if (!response.ok) {
+                const errorText = await response.text();
+                lastError = `${response.status} - ${errorText}`;
+                logger.debug(`[AccountManager] loadCodeAssist failed at ${endpoint}: ${lastError}`);
+                continue;
+            }
+
+            const data = await response.json();
+            gotSuccessfulResponse = true;
+            loadCodeAssistData = data;
+
+            logger.debug(`[AccountManager] loadCodeAssist response from ${endpoint}:`, JSON.stringify(data));
+
+            if (typeof data.cloudaicompanionProject === 'string') {
+                logger.success(`[AccountManager] Discovered project: ${data.cloudaicompanionProject}`);
+                return data.cloudaicompanionProject;
+            }
+            if (data.cloudaicompanionProject?.id) {
+                logger.success(`[AccountManager] Discovered project: ${data.cloudaicompanionProject.id}`);
+                return data.cloudaicompanionProject.id;
+            }
+
+            // No project found - log tier data and try to onboard the user
+            logger.info(`[AccountManager] No project in loadCodeAssist response, attempting onboardUser...`);
+            logger.debug(`[AccountManager] Tier data for onboarding: paidTier=${JSON.stringify(data.paidTier)}, currentTier=${JSON.stringify(data.currentTier)}, allowedTiers=${JSON.stringify(data.allowedTiers?.map(t => ({ id: t?.id, isDefault: t?.isDefault })))}`);
+            break;
+        } catch (error) {
+            lastError = error.message;
+            logger.debug(`[AccountManager] loadCodeAssist error at ${endpoint}:`, error.message);
+        }
+    }
+
+    // If we got a successful response but no project, try onboarding
+    if (gotSuccessfulResponse && loadCodeAssistData) {
+        // Priority: paidTier > currentTier > allowedTiers (consistent with model-api.js)
+        let tierId = null;
+        let tierSource = null;
+
+        if (loadCodeAssistData.paidTier?.id) {
+            tierId = loadCodeAssistData.paidTier.id;
+            tierSource = 'paidTier';
+        } else if (loadCodeAssistData.currentTier?.id) {
+            tierId = loadCodeAssistData.currentTier.id;
+            tierSource = 'currentTier';
+        } else {
+            tierId = getDefaultTierId(loadCodeAssistData.allowedTiers);
+            tierSource = 'allowedTiers';
+        }
+
+        tierId = tierId || 'free-tier';
+        logger.info(`[AccountManager] Onboarding user with tier: ${tierId} (source: ${tierSource})`);
+
+        // Check if this is a free tier (raw API values contain 'free')
+        const isFree = tierId.toLowerCase().includes('free');
+
+        // For non-free tiers, pass DEFAULT_PROJECT_ID as the GCP project
+        // The API requires a project for paid tier onboarding
+        const onboardedProject = await onboardUser(
+            token,
+            tierId,
+            isFree ? null : DEFAULT_PROJECT_ID
+        );
+        if (onboardedProject) {
+            logger.success(`[AccountManager] Successfully onboarded, project: ${onboardedProject}`);
+            return onboardedProject;
+        }
+
+        logger.warn(`[AccountManager] Onboarding failed, using default project: ${DEFAULT_PROJECT_ID}`);
+    }
+
+    // Only warn if all endpoints failed with errors (not just missing project)
+    if (!gotSuccessfulResponse) {
+        logger.warn(`[AccountManager] loadCodeAssist failed for all endpoints: ${lastError}`);
+    }
+    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,380 @@
+/**
+ * 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
+} 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);
+        }
+    }
+
+    /**
+     * 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 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);
+    }
+
+    /**
+     * 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 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;
+    }
+}
+
+export default AccountManager;