@kamel-ahmed/proxy-claude 1.0.0

package/src/account-manager/onboarding.js

@@ -0,0 +1,117 @@
+/**
+ * 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 = null, maxAttempts = 10, delayMs = 5000) {
+    const metadata = {
+        ideType: 'IDE_UNSPECIFIED',
+        platform: 'PLATFORM_UNSPECIFIED',
+        pluginType: 'GEMINI'
+    };
+
+    if (projectId) {
+        metadata.duetProject = projectId;
+    }
+
+    const requestBody = {
+        tierId,
+        metadata
+    };
+
+    // Check if this is a free tier (handles raw API values like 'free-tier')
+    const isFree = tierId.toLowerCase().includes('free');
+
+    // Non-free tiers require a cloudaicompanionProject
+    if (!isFree && projectId) {
+        requestBody.cloudaicompanionProject = projectId;
+    }
+
+    logger.debug(`[Onboarding] Starting onboard with tierId: ${tierId}, projectId: ${projectId}, isFree: ${isFree}`);
+
+    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;
+}

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

@@ -0,0 +1,237 @@
+/**
+ * Rate Limit Management
+ *
+ * Handles rate limit tracking and state management for accounts.
+ * All rate limits are model-specific.
+ */
+
+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 for a specific model
+ *
+ * @param {Array} accounts - Array of account objects
+ * @param {string} modelId - Model ID to check rate limits for
+ * @returns {boolean} True if all accounts are rate-limited
+ */
+export function isAllRateLimited(accounts, modelId) {
+    if (accounts.length === 0) return true;
+    if (!modelId) return false; // No model specified = not rate limited
+
+    return accounts.every(acc => {
+        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 limit = modelLimits[modelId];
+        return limit && limit.isRateLimited && limit.resetTime > Date.now();
+    });
+}
+
+/**
+ * Get list of available (non-rate-limited, non-invalid) accounts for a model
+ *
+ * @param {Array} accounts - Array of account objects
+ * @param {string} [modelId] - Model ID to filter by
+ * @returns {Array} Array of available account objects
+ */
+export function getAvailableAccounts(accounts, modelId = null) {
+    return accounts.filter(acc => {
+        if (acc.isInvalid) return false;
+
+        // WebUI: Skip disabled accounts
+        if (acc.enabled === false) return false;
+
+        if (modelId && acc.modelRateLimits && acc.modelRateLimits[modelId]) {
+            const limit = acc.modelRateLimits[modelId];
+            if (limit.isRateLimited && limit.resetTime > Date.now()) {
+                return false;
+            }
+        }
+
+        return true;
+    });
+}
+
+/**
+ * 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.modelRateLimits) {
+            for (const [modelId, limit] of Object.entries(account.modelRateLimits)) {
+                if (limit.isRateLimited && limit.resetTime <= now) {
+                    limit.isRateLimited = false;
+                    limit.resetTime = null;
+                    cleared++;
+                    logger.success(`[AccountManager] Rate limit expired for: ${account.email} (model: ${modelId})`);
+                }
+            }
+        }
+    }
+
+    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) {
+        if (account.modelRateLimits) {
+            for (const key of Object.keys(account.modelRateLimits)) {
+                account.modelRateLimits[key] = { isRateLimited: false, resetTime: null };
+            }
+        }
+    }
+    logger.warn('[AccountManager] Reset all rate limits for optimistic retry');
+}
+
+/**
+ * Mark an account as rate-limited for a specific model
+ *
+ * @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 (from API)
+ * @param {string} modelId - Model ID to mark rate limit for
+ * @returns {boolean} True if account was found and marked
+ */
+export function markRateLimited(accounts, email, resetMs = null, modelId) {
+    const account = accounts.find(a => a.email === email);
+    if (!account) return false;
+
+    // Store the ACTUAL reset time from the API
+    // This is used to decide whether to wait (short) or switch accounts (long)
+    const actualResetMs = (resetMs && resetMs > 0) ? resetMs : DEFAULT_COOLDOWN_MS;
+
+    if (!account.modelRateLimits) {
+        account.modelRateLimits = {};
+    }
+
+    account.modelRateLimits[modelId] = {
+        isRateLimited: true,
+        resetTime: Date.now() + actualResetMs,  // Actual reset time for decisions
+        actualResetMs: actualResetMs             // Original duration from API
+    };
+
+    // Log appropriately based on duration
+    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;
+}
+
+/**
+ * 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 for a model
+ *
+ * @param {Array} accounts - Array of account objects
+ * @param {string} modelId - Model ID to check
+ * @returns {number} Wait time in milliseconds
+ */
+export function getMinWaitTimeMs(accounts, modelId) {
+    if (!isAllRateLimited(accounts, modelId)) return 0;
+
+    const now = Date.now();
+    let minWait = Infinity;
+    let soonestAccount = null;
+
+    for (const account of accounts) {
+        if (modelId && account.modelRateLimits && account.modelRateLimits[modelId]) {
+            const limit = account.modelRateLimits[modelId];
+            if (limit.isRateLimited && limit.resetTime) {
+                const wait = limit.resetTime - 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;
+}
+
+/**
+ * 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
+    };
+}

package/src/account-manager/storage.js

@@ -0,0 +1,136 @@
+/**
+ * Account Storage
+ *
+ * Handles loading and saving account configuration to disk.
+ */
+
+import { readFile, writeFile, mkdir, access } from 'fs/promises';
+import { constants as fsConstants } from 'fs';
+import { dirname } from 'path';
+import { ACCOUNT_CONFIG_PATH } from '../constants.js';
+import { getAuthStatus } from '../auth/database.js';
+import { logger } from '../utils/logger.js';
+
+/**
+ * Load accounts from the config file
+ *
+ * @param {string} configPath - Path to the config file
+ * @returns {Promise<{accounts: Array, settings: Object, activeIndex: number}>}
+ */
+export async function loadAccounts(configPath = ACCOUNT_CONFIG_PATH) {
+    try {
+        // Check if config file exists using async access
+        await access(configPath, fsConstants.F_OK);
+        const configData = await readFile(configPath, 'utf-8');
+        const config = JSON.parse(configData);
+
+        const accounts = (config.accounts || []).map(acc => ({
+            ...acc,
+            lastUsed: acc.lastUsed || null,
+            enabled: acc.enabled !== false, // Default to true if not specified
+            // Reset invalid flag on startup - give accounts a fresh chance to refresh
+            isInvalid: false,
+            invalidReason: null,
+            modelRateLimits: acc.modelRateLimits || {},
+            // New fields for subscription and quota tracking
+            subscription: acc.subscription || { tier: 'unknown', projectId: null, detectedAt: null },
+            quota: acc.quota || { models: {}, lastChecked: null }
+        }));
+
+        const settings = config.settings || {};
+        let activeIndex = config.activeIndex || 0;
+
+        // Clamp activeIndex to valid range
+        if (activeIndex >= accounts.length) {
+            activeIndex = 0;
+        }
+
+        logger.info(`[AccountManager] Loaded ${accounts.length} account(s) from config`);
+
+        return { accounts, settings, activeIndex };
+    } catch (error) {
+        if (error.code === 'ENOENT') {
+            // No config file - return empty
+            logger.info('[AccountManager] No config file found. Using Antigravity database (single account mode)');
+        } else {
+            logger.error('[AccountManager] Failed to load config:', error.message);
+        }
+        return { accounts: [], settings: {}, activeIndex: 0 };
+    }
+}
+
+/**
+ * Load the default account from Antigravity's database
+ *
+ * @param {string} dbPath - Optional path to the database
+ * @returns {{accounts: Array, tokenCache: Map}}
+ */
+export function loadDefaultAccount(dbPath) {
+    try {
+        const authData = getAuthStatus(dbPath);
+        if (authData?.apiKey) {
+            const account = {
+                email: authData.email || 'default@antigravity',
+                source: 'database',
+                lastUsed: null,
+                modelRateLimits: {}
+            };
+
+            const tokenCache = new Map();
+            tokenCache.set(account.email, {
+                token: authData.apiKey,
+                extractedAt: Date.now()
+            });
+
+            logger.info(`[AccountManager] Loaded default account: ${account.email}`);
+
+            return { accounts: [account], tokenCache };
+        }
+    } catch (error) {
+        logger.error('[AccountManager] Failed to load default account:', error.message);
+    }
+
+    return { accounts: [], tokenCache: new Map() };
+}
+
+/**
+ * Save account configuration to disk
+ *
+ * @param {string} configPath - Path to the config file
+ * @param {Array} accounts - Array of account objects
+ * @param {Object} settings - Settings object
+ * @param {number} activeIndex - Current active account index
+ */
+export async function saveAccounts(configPath, accounts, settings, activeIndex) {
+    try {
+        // Ensure directory exists
+        const dir = dirname(configPath);
+        await mkdir(dir, { recursive: true });
+
+        const config = {
+            accounts: accounts.map(acc => ({
+                email: acc.email,
+                source: acc.source,
+                enabled: acc.enabled !== false, // Persist enabled state
+                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,
+                isInvalid: acc.isInvalid || false,
+                invalidReason: acc.invalidReason || null,
+                modelRateLimits: acc.modelRateLimits || {},
+                lastUsed: acc.lastUsed,
+                // Persist subscription and quota data
+                subscription: acc.subscription || { tier: 'unknown', projectId: null, detectedAt: null },
+                quota: acc.quota || { models: {}, lastChecked: null }
+            })),
+            settings: settings,
+            activeIndex: activeIndex
+        };
+
+        await writeFile(configPath, JSON.stringify(config, null, 2));
+    } catch (error) {
+        logger.error('[AccountManager] Failed to save config:', error.message);
+    }
+}

package/src/account-manager/strategies/base-strategy.js

@@ -0,0 +1,104 @@
+/**
+ * Base Strategy
+ *
+ * Abstract base class defining the interface for account selection strategies.
+ * All strategies must implement the selectAccount method.
+ */
+
+/**
+ * @typedef {Object} SelectionResult
+ * @property {Object|null} account - The selected account or null if none available
+ * @property {number} index - The index of the selected account
+ * @property {number} [waitMs] - Optional wait time before account becomes available
+ */
+
+export class BaseStrategy {
+    /**
+     * Create a new BaseStrategy
+     * @param {Object} config - Strategy configuration
+     */
+    constructor(config = {}) {
+        if (new.target === BaseStrategy) {
+            throw new Error('BaseStrategy is abstract and cannot be instantiated directly');
+        }
+        this.config = config;
+    }
+
+    /**
+     * Select an account for a request
+     * @param {Array} accounts - Array of account objects
+     * @param {string} modelId - The model ID for the request
+     * @param {Object} options - Additional options
+     * @param {number} options.currentIndex - Current account index
+     * @param {string} [options.sessionId] - Session ID for cache continuity
+     * @param {Function} [options.onSave] - Callback to save changes
+     * @returns {SelectionResult} The selected account and index
+     */
+    selectAccount(accounts, modelId, options = {}) {
+        throw new Error('selectAccount must be implemented by subclass');
+    }
+
+    /**
+     * Called after a successful request
+     * @param {Object} account - The account that was used
+     * @param {string} modelId - The model ID that was used
+     */
+    onSuccess(account, modelId) {
+        // Default: no-op, override in subclass if needed
+    }
+
+    /**
+     * Called when a request is rate-limited
+     * @param {Object} account - The account that was rate-limited
+     * @param {string} modelId - The model ID that was rate-limited
+     */
+    onRateLimit(account, modelId) {
+        // Default: no-op, override in subclass if needed
+    }
+
+    /**
+     * Called when a request fails (non-rate-limit error)
+     * @param {Object} account - The account that failed
+     * @param {string} modelId - The model ID that failed
+     */
+    onFailure(account, modelId) {
+        // Default: no-op, override in subclass if needed
+    }
+
+    /**
+     * Check if an account is usable for a specific model
+     * @param {Object} account - Account object
+     * @param {string} modelId - Model ID to check
+     * @returns {boolean} True if account is usable
+     */
+    isAccountUsable(account, modelId) {
+        if (!account || account.isInvalid) return false;
+
+        // Skip disabled accounts
+        if (account.enabled === false) return false;
+
+        // Check model-specific rate limit
+        if (modelId && account.modelRateLimits && account.modelRateLimits[modelId]) {
+            const limit = account.modelRateLimits[modelId];
+            if (limit.isRateLimited && limit.resetTime > Date.now()) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    /**
+     * Get all usable accounts for a model
+     * @param {Array} accounts - Array of account objects
+     * @param {string} modelId - Model ID to check
+     * @returns {Array} Array of usable accounts with their original indices
+     */
+    getUsableAccounts(accounts, modelId) {
+        return accounts
+            .map((account, index) => ({ account, index }))
+            .filter(({ account }) => this.isAccountUsable(account, modelId));
+    }
+}
+
+export default BaseStrategy;