@kamel-ahmed/proxy-claude 1.0.0

package/src/account-manager/strategies/trackers/token-bucket-tracker.js

@@ -0,0 +1,121 @@
+/**
+ * Token Bucket Tracker
+ *
+ * Client-side rate limiting using the token bucket algorithm.
+ * Each account has a bucket of tokens that regenerate over time.
+ * Requests consume tokens; accounts without tokens are deprioritized.
+ */
+
+// Default configuration (matches opencode-antigravity-auth)
+const DEFAULT_CONFIG = {
+    maxTokens: 50,        // Maximum token capacity
+    tokensPerMinute: 6,   // Regeneration rate
+    initialTokens: 50     // Starting tokens
+};
+
+export class TokenBucketTracker {
+    #buckets = new Map(); // email -> { tokens, lastUpdated }
+    #config;
+
+    /**
+     * Create a new TokenBucketTracker
+     * @param {Object} config - Token bucket configuration
+     */
+    constructor(config = {}) {
+        this.#config = { ...DEFAULT_CONFIG, ...config };
+    }
+
+    /**
+     * Get the current token count for an account
+     * @param {string} email - Account email
+     * @returns {number} Current token count (with regeneration applied)
+     */
+    getTokens(email) {
+        const bucket = this.#buckets.get(email);
+        if (!bucket) {
+            return this.#config.initialTokens;
+        }
+
+        // Apply token regeneration based on time elapsed
+        const now = Date.now();
+        const minutesElapsed = (now - bucket.lastUpdated) / (1000 * 60);
+        const regenerated = minutesElapsed * this.#config.tokensPerMinute;
+        const currentTokens = Math.min(
+            this.#config.maxTokens,
+            bucket.tokens + regenerated
+        );
+
+        return currentTokens;
+    }
+
+    /**
+     * Check if an account has tokens available
+     * @param {string} email - Account email
+     * @returns {boolean} True if account has at least 1 token
+     */
+    hasTokens(email) {
+        return this.getTokens(email) >= 1;
+    }
+
+    /**
+     * Consume a token from an account's bucket
+     * @param {string} email - Account email
+     * @returns {boolean} True if token was consumed, false if no tokens available
+     */
+    consume(email) {
+        const currentTokens = this.getTokens(email);
+        if (currentTokens < 1) {
+            return false;
+        }
+
+        this.#buckets.set(email, {
+            tokens: currentTokens - 1,
+            lastUpdated: Date.now()
+        });
+        return true;
+    }
+
+    /**
+     * Refund a token to an account's bucket (e.g., on request failure before processing)
+     * @param {string} email - Account email
+     */
+    refund(email) {
+        const currentTokens = this.getTokens(email);
+        const newTokens = Math.min(
+            this.#config.maxTokens,
+            currentTokens + 1
+        );
+        this.#buckets.set(email, {
+            tokens: newTokens,
+            lastUpdated: Date.now()
+        });
+    }
+
+    /**
+     * Get the maximum token capacity
+     * @returns {number} Maximum tokens per bucket
+     */
+    getMaxTokens() {
+        return this.#config.maxTokens;
+    }
+
+    /**
+     * Reset the bucket for an account
+     * @param {string} email - Account email
+     */
+    reset(email) {
+        this.#buckets.set(email, {
+            tokens: this.#config.initialTokens,
+            lastUpdated: Date.now()
+        });
+    }
+
+    /**
+     * Clear all tracked buckets
+     */
+    clear() {
+        this.#buckets.clear();
+    }
+}
+
+export default TokenBucketTracker;

package/src/auth/database.js

@@ -0,0 +1,169 @@
+/**
+ * SQLite Database Access Module
+ * Provides cross-platform database operations for Antigravity state.
+ *
+ * Uses better-sqlite3 for:
+ * - Windows compatibility (no CLI dependency)
+ * - Native performance
+ * - Synchronous API (simple error handling)
+ *
+ * Includes auto-rebuild capability for handling Node.js version updates
+ * that cause native module incompatibility.
+ */
+
+import { createRequire } from 'module';
+import { ANTIGRAVITY_DB_PATH } from '../constants.js';
+import { isModuleVersionError, attemptAutoRebuild, clearRequireCache } from '../utils/native-module-helper.js';
+import { logger } from '../utils/logger.js';
+import { NativeModuleError } from '../errors.js';
+
+const require = createRequire(import.meta.url);
+
+// Lazy-loaded Database constructor
+let Database = null;
+let moduleLoadError = null;
+
+/**
+ * Load the better-sqlite3 module with auto-rebuild on version mismatch
+ * Uses synchronous require to maintain API compatibility
+ * @returns {Function} The Database constructor
+ * @throws {Error} If module cannot be loaded even after rebuild
+ */
+function loadDatabaseModule() {
+    // Return cached module if already loaded
+    if (Database) return Database;
+
+    // Re-throw cached error if previous load failed permanently
+    if (moduleLoadError) throw moduleLoadError;
+
+    try {
+        Database = require('better-sqlite3');
+        return Database;
+    } catch (error) {
+        if (isModuleVersionError(error)) {
+            logger.warn('[Database] Native module version mismatch detected');
+
+            if (attemptAutoRebuild(error)) {
+                // Clear require cache and retry
+                try {
+                    const resolvedPath = require.resolve('better-sqlite3');
+                    // Clear the module and all its dependencies from cache
+                    clearRequireCache(resolvedPath, require.cache);
+
+                    Database = require('better-sqlite3');
+                    logger.success('[Database] Module reloaded successfully after rebuild');
+                    return Database;
+                } catch (retryError) {
+                    // Rebuild succeeded but reload failed - user needs to restart
+                    moduleLoadError = new NativeModuleError(
+                        'Native module rebuild completed. Please restart the server to apply the fix.',
+                        true,  // rebuildSucceeded
+                        true   // restartRequired
+                    );
+                    logger.info('[Database] Rebuild succeeded - server restart required');
+                    throw moduleLoadError;
+                }
+            } else {
+                moduleLoadError = new NativeModuleError(
+                    'Failed to auto-rebuild native module. Please run manually:\n' +
+                    '  npm rebuild better-sqlite3\n' +
+                    'Or if using npx, find the package location in the error and run:\n' +
+                    '  cd /path/to/better-sqlite3 && npm rebuild',
+                    false,  // rebuildSucceeded
+                    false   // restartRequired
+                );
+                throw moduleLoadError;
+            }
+        }
+
+        // Non-version-mismatch error, just throw it
+        throw error;
+    }
+}
+
+/**
+ * Query Antigravity database for authentication status
+ * @param {string} [dbPath] - Optional custom database path
+ * @returns {Object} Parsed auth data with apiKey, email, name, etc.
+ * @throws {Error} If database doesn't exist, query fails, or no auth status found
+ */
+export function getAuthStatus(dbPath = ANTIGRAVITY_DB_PATH) {
+    const Db = loadDatabaseModule();
+    let db;
+    try {
+        // Open database in read-only mode
+        db = new Db(dbPath, {
+            readonly: true,
+            fileMustExist: true
+        });
+
+        // Prepare and execute query
+        const stmt = db.prepare(
+            "SELECT value FROM ItemTable WHERE key = 'antigravityAuthStatus'"
+        );
+        const row = stmt.get();
+
+        if (!row || !row.value) {
+            throw new Error('No auth status found in database');
+        }
+
+        // Parse JSON value
+        const authData = JSON.parse(row.value);
+
+        if (!authData.apiKey) {
+            throw new Error('Auth data missing apiKey field');
+        }
+
+        return authData;
+    } catch (error) {
+        // Enhance error messages for common issues
+        if (error.code === 'SQLITE_CANTOPEN') {
+            throw new Error(
+                `Database not found at ${dbPath}. ` +
+                'Make sure Antigravity is installed and you are logged in.'
+            );
+        }
+        // Re-throw with context if not already our error
+        if (error.message.includes('No auth status') || error.message.includes('missing apiKey')) {
+            throw error;
+        }
+        // Re-throw native module errors from loadDatabaseModule without wrapping
+        if (error instanceof NativeModuleError) {
+            throw error;
+        }
+        throw new Error(`Failed to read Antigravity database: ${error.message}`);
+    } finally {
+        // Always close database connection
+        if (db) {
+            db.close();
+        }
+    }
+}
+
+/**
+ * Check if database exists and is accessible
+ * @param {string} [dbPath] - Optional custom database path
+ * @returns {boolean} True if database exists and can be opened
+ */
+export function isDatabaseAccessible(dbPath = ANTIGRAVITY_DB_PATH) {
+    let db;
+    try {
+        const Db = loadDatabaseModule();
+        db = new Db(dbPath, {
+            readonly: true,
+            fileMustExist: true
+        });
+        return true;
+    } catch {
+        return false;
+    } finally {
+        if (db) {
+            db.close();
+        }
+    }
+}
+
+export default {
+    getAuthStatus,
+    isDatabaseAccessible
+};

package/src/auth/oauth.js

@@ -0,0 +1,419 @@
+/**
+ * Google OAuth with PKCE for Antigravity
+ *
+ * Implements the same OAuth flow as opencode-antigravity-auth
+ * to obtain refresh tokens for multiple Google accounts.
+ * Uses a local callback server to automatically capture the auth code.
+ */
+
+import crypto from 'crypto';
+import http from 'http';
+import {
+    ANTIGRAVITY_ENDPOINT_FALLBACKS,
+    LOAD_CODE_ASSIST_HEADERS,
+    OAUTH_CONFIG,
+    OAUTH_REDIRECT_URI
+} from '../constants.js';
+import { logger } from '../utils/logger.js';
+import { onboardUser, getDefaultTierId } from '../account-manager/onboarding.js';
+
+/**
+ * Generate PKCE code verifier and challenge
+ */
+function generatePKCE() {
+    const verifier = crypto.randomBytes(32).toString('base64url');
+    const challenge = crypto
+        .createHash('sha256')
+        .update(verifier)
+        .digest('base64url');
+    return { verifier, challenge };
+}
+
+/**
+ * Generate authorization URL for Google OAuth
+ * Returns the URL and the PKCE verifier (needed for token exchange)
+ *
+ * @param {string} [customRedirectUri] - Optional custom redirect URI (e.g. for WebUI)
+ * @returns {{url: string, verifier: string, state: string}} Auth URL and PKCE data
+ */
+export function getAuthorizationUrl(customRedirectUri = null) {
+    const { verifier, challenge } = generatePKCE();
+    const state = crypto.randomBytes(16).toString('hex');
+
+    const params = new URLSearchParams({
+        client_id: OAUTH_CONFIG.clientId,
+        redirect_uri: customRedirectUri || OAUTH_REDIRECT_URI,
+        response_type: 'code',
+        scope: OAUTH_CONFIG.scopes.join(' '),
+        access_type: 'offline',
+        prompt: 'consent',
+        code_challenge: challenge,
+        code_challenge_method: 'S256',
+        state: state
+    });
+
+    return {
+        url: `${OAUTH_CONFIG.authUrl}?${params.toString()}`,
+        verifier,
+        state
+    };
+}
+
+/**
+ * Extract authorization code and state from user input.
+ * User can paste either:
+ * - Full callback URL: http://localhost:51121/oauth-callback?code=xxx&state=xxx
+ * - Just the code parameter: 4/0xxx...
+ *
+ * @param {string} input - User input (URL or code)
+ * @returns {{code: string, state: string|null}} Extracted code and optional state
+ */
+export function extractCodeFromInput(input) {
+    if (!input || typeof input !== 'string') {
+        throw new Error('No input provided');
+    }
+
+    const trimmed = input.trim();
+
+    // Check if it looks like a URL
+    if (trimmed.startsWith('http://') || trimmed.startsWith('https://')) {
+        try {
+            const url = new URL(trimmed);
+            const code = url.searchParams.get('code');
+            const state = url.searchParams.get('state');
+            const error = url.searchParams.get('error');
+
+            if (error) {
+                throw new Error(`OAuth error: ${error}`);
+            }
+
+            if (!code) {
+                throw new Error('No authorization code found in URL');
+            }
+
+            return { code, state };
+        } catch (e) {
+            if (e.message.includes('OAuth error') || e.message.includes('No authorization code')) {
+                throw e;
+            }
+            throw new Error('Invalid URL format');
+        }
+    }
+
+    // Assume it's a raw code
+    // Google auth codes typically start with "4/" and are long
+    if (trimmed.length < 10) {
+        throw new Error('Input is too short to be a valid authorization code');
+    }
+
+    return { code: trimmed, state: null };
+}
+
+/**
+ * Start a local server to receive the OAuth callback
+ * Returns a promise that resolves with the authorization code
+ *
+ * @param {string} expectedState - Expected state parameter for CSRF protection
+ * @param {number} timeoutMs - Timeout in milliseconds (default 120000)
+ * @returns {Promise<string>} Authorization code from OAuth callback
+ */
+export function startCallbackServer(expectedState, timeoutMs = 120000) {
+    return new Promise((resolve, reject) => {
+        const server = http.createServer((req, res) => {
+            const url = new URL(req.url, `http://localhost:${OAUTH_CONFIG.callbackPort}`);
+
+            if (url.pathname !== '/oauth-callback') {
+                res.writeHead(404);
+                res.end('Not found');
+                return;
+            }
+
+            const code = url.searchParams.get('code');
+            const state = url.searchParams.get('state');
+            const error = url.searchParams.get('error');
+
+            if (error) {
+                res.writeHead(400, { 'Content-Type': 'text/html; charset=utf-8' });
+                res.end(`
+                    <html>
+                    <head><meta charset="UTF-8"><title>Authentication Failed</title></head>
+                    <body style="font-family: system-ui; padding: 40px; text-align: center;">
+                        <h1 style="color: #dc3545;">❌ Authentication Failed</h1>
+                        <p>Error: ${error}</p>
+                        <p>You can close this window.</p>
+                    </body>
+                    </html>
+                `);
+                server.close();
+                reject(new Error(`OAuth error: ${error}`));
+                return;
+            }
+
+            if (state !== expectedState) {
+                res.writeHead(400, { 'Content-Type': 'text/html; charset=utf-8' });
+                res.end(`
+                    <html>
+                    <head><meta charset="UTF-8"><title>Authentication Failed</title></head>
+                    <body style="font-family: system-ui; padding: 40px; text-align: center;">
+                        <h1 style="color: #dc3545;">❌ Authentication Failed</h1>
+                        <p>State mismatch - possible CSRF attack.</p>
+                        <p>You can close this window.</p>
+                    </body>
+                    </html>
+                `);
+                server.close();
+                reject(new Error('State mismatch'));
+                return;
+            }
+
+            if (!code) {
+                res.writeHead(400, { 'Content-Type': 'text/html; charset=utf-8' });
+                res.end(`
+                    <html>
+                    <head><meta charset="UTF-8"><title>Authentication Failed</title></head>
+                    <body style="font-family: system-ui; padding: 40px; text-align: center;">
+                        <h1 style="color: #dc3545;">❌ Authentication Failed</h1>
+                        <p>No authorization code received.</p>
+                        <p>You can close this window.</p>
+                    </body>
+                    </html>
+                `);
+                server.close();
+                reject(new Error('No authorization code'));
+                return;
+            }
+
+            // Success!
+            res.writeHead(200, { 'Content-Type': 'text/html; charset=utf-8' });
+            res.end(`
+                <html>
+                <head><meta charset="UTF-8"><title>Authentication Successful</title></head>
+                <body style="font-family: system-ui; padding: 40px; text-align: center;">
+                    <h1 style="color: #28a745;">✅ Authentication Successful!</h1>
+                    <p>You can close this window and return to the terminal.</p>
+                    <script>setTimeout(() => window.close(), 2000);</script>
+                </body>
+                </html>
+            `);
+
+            server.close();
+            resolve(code);
+        });
+
+        server.on('error', (err) => {
+            if (err.code === 'EADDRINUSE') {
+                reject(new Error(`Port ${OAUTH_CONFIG.callbackPort} is already in use. Close any other OAuth flows and try again.`));
+            } else {
+                reject(err);
+            }
+        });
+
+        server.listen(OAUTH_CONFIG.callbackPort, () => {
+            logger.info(`[OAuth] Callback server listening on port ${OAUTH_CONFIG.callbackPort}`);
+        });
+
+        // Timeout after specified duration
+        setTimeout(() => {
+            server.close();
+            reject(new Error('OAuth callback timeout - no response received'));
+        }, timeoutMs);
+    });
+}
+
+/**
+ * Exchange authorization code for tokens
+ *
+ * @param {string} code - Authorization code from OAuth callback
+ * @param {string} verifier - PKCE code verifier
+ * @returns {Promise<{accessToken: string, refreshToken: string, expiresIn: number}>} OAuth tokens
+ */
+export async function exchangeCode(code, verifier) {
+    const response = await fetch(OAUTH_CONFIG.tokenUrl, {
+        method: 'POST',
+        headers: {
+            'Content-Type': 'application/x-www-form-urlencoded'
+        },
+        body: new URLSearchParams({
+            client_id: OAUTH_CONFIG.clientId,
+            client_secret: OAUTH_CONFIG.clientSecret,
+            code: code,
+            code_verifier: verifier,
+            grant_type: 'authorization_code',
+            redirect_uri: OAUTH_REDIRECT_URI
+        })
+    });
+
+    if (!response.ok) {
+        const error = await response.text();
+        logger.error(`[OAuth] Token exchange failed: ${response.status} ${error}`);
+        throw new Error(`Token exchange failed: ${error}`);
+    }
+
+    const tokens = await response.json();
+
+    if (!tokens.access_token) {
+        logger.error('[OAuth] No access token in response:', tokens);
+        throw new Error('No access token received');
+    }
+
+    logger.info(`[OAuth] Token exchange successful, access_token length: ${tokens.access_token?.length}`);
+
+    return {
+        accessToken: tokens.access_token,
+        refreshToken: tokens.refresh_token,
+        expiresIn: tokens.expires_in
+    };
+}
+
+/**
+ * Refresh access token using refresh token
+ *
+ * @param {string} refreshToken - OAuth refresh token
+ * @returns {Promise<{accessToken: string, expiresIn: number}>} New access token
+ */
+export async function refreshAccessToken(refreshToken) {
+    const response = await fetch(OAUTH_CONFIG.tokenUrl, {
+        method: 'POST',
+        headers: {
+            'Content-Type': 'application/x-www-form-urlencoded'
+        },
+        body: new URLSearchParams({
+            client_id: OAUTH_CONFIG.clientId,
+            client_secret: OAUTH_CONFIG.clientSecret,
+            refresh_token: refreshToken,
+            grant_type: 'refresh_token'
+        })
+    });
+
+    if (!response.ok) {
+        const error = await response.text();
+        throw new Error(`Token refresh failed: ${error}`);
+    }
+
+    const tokens = await response.json();
+    return {
+        accessToken: tokens.access_token,
+        expiresIn: tokens.expires_in
+    };
+}
+
+/**
+ * Get user email from access token
+ *
+ * @param {string} accessToken - OAuth access token
+ * @returns {Promise<string>} User's email address
+ */
+export async function getUserEmail(accessToken) {
+    const response = await fetch(OAUTH_CONFIG.userInfoUrl, {
+        headers: {
+            'Authorization': `Bearer ${accessToken}`
+        }
+    });
+
+    if (!response.ok) {
+        const errorText = await response.text();
+        logger.error(`[OAuth] getUserEmail failed: ${response.status} ${errorText}`);
+        throw new Error(`Failed to get user info: ${response.status}`);
+    }
+
+    const userInfo = await response.json();
+    return userInfo.email;
+}
+
+/**
+ * Discover project ID for the authenticated user
+ *
+ * @param {string} accessToken - OAuth access token
+ * @returns {Promise<string|null>} Project ID or null if not found
+ */
+export async function discoverProjectId(accessToken) {
+    let loadCodeAssistData = null;
+
+    for (const endpoint of ANTIGRAVITY_ENDPOINT_FALLBACKS) {
+        try {
+            const response = await fetch(`${endpoint}/v1internal:loadCodeAssist`, {
+                method: 'POST',
+                headers: {
+                    'Authorization': `Bearer ${accessToken}`,
+                    'Content-Type': 'application/json',
+                    ...LOAD_CODE_ASSIST_HEADERS
+                },
+                body: JSON.stringify({
+                    metadata: {
+                        ideType: 'IDE_UNSPECIFIED',
+                        platform: 'PLATFORM_UNSPECIFIED',
+                        pluginType: 'GEMINI'
+                    }
+                })
+            });
+
+            if (!response.ok) continue;
+
+            const data = await response.json();
+            loadCodeAssistData = data;
+
+            if (typeof data.cloudaicompanionProject === 'string') {
+                return data.cloudaicompanionProject;
+            }
+            if (data.cloudaicompanionProject?.id) {
+                return data.cloudaicompanionProject.id;
+            }
+
+            // No project found - try to onboard
+            logger.info('[OAuth] No project in loadCodeAssist response, attempting onboardUser...');
+            break;
+        } catch (error) {
+            logger.warn(`[OAuth] Project discovery failed at ${endpoint}:`, error.message);
+        }
+    }
+
+    // Try onboarding if we got a response but no project
+    if (loadCodeAssistData) {
+        const tierId = getDefaultTierId(loadCodeAssistData.allowedTiers) || 'FREE';
+        logger.info(`[OAuth] Onboarding user with tier: ${tierId}`);
+
+        const onboardedProject = await onboardUser(accessToken, tierId);
+        if (onboardedProject) {
+            logger.success(`[OAuth] Successfully onboarded, project: ${onboardedProject}`);
+            return onboardedProject;
+        }
+    }
+
+    return null;
+}
+
+/**
+ * Complete OAuth flow: exchange code and get all account info
+ *
+ * @param {string} code - Authorization code from OAuth callback
+ * @param {string} verifier - PKCE code verifier
+ * @returns {Promise<{email: string, refreshToken: string, accessToken: string, projectId: string|null}>} Complete account info
+ */
+export async function completeOAuthFlow(code, verifier) {
+    // Exchange code for tokens
+    const tokens = await exchangeCode(code, verifier);
+
+    // Get user email
+    const email = await getUserEmail(tokens.accessToken);
+
+    // Discover project ID
+    const projectId = await discoverProjectId(tokens.accessToken);
+
+    return {
+        email,
+        refreshToken: tokens.refreshToken,
+        accessToken: tokens.accessToken,
+        projectId
+    };
+}
+
+export default {
+    getAuthorizationUrl,
+    extractCodeFromInput,
+    startCallbackServer,
+    exchangeCode,
+    refreshAccessToken,
+    getUserEmail,
+    discoverProjectId,
+    completeOAuthFlow
+};