antigravity-claude-proxy 1.0.0

package/src/server.js

@@ -0,0 +1,517 @@
+/**
+ * Express Server - Anthropic-compatible API
+ * Proxies to Google Cloud Code via Antigravity
+ * Supports multi-account load balancing
+ */
+
+import express from 'express';
+import cors from 'cors';
+import { sendMessage, sendMessageStream, listModels, getModelQuotas } from './cloudcode-client.js';
+import { forceRefresh } from './token-extractor.js';
+import { REQUEST_BODY_LIMIT } from './constants.js';
+import { AccountManager } from './account-manager.js';
+import { formatDuration } from './utils/helpers.js';
+
+const app = express();
+
+// Initialize account manager (will be fully initialized on first request or startup)
+const accountManager = new AccountManager();
+
+// Track initialization status
+let isInitialized = false;
+let initError = null;
+let initPromise = null;
+
+/**
+ * Ensure account manager is initialized (with race condition protection)
+ */
+async function ensureInitialized() {
+    if (isInitialized) return;
+
+    // If initialization is already in progress, wait for it
+    if (initPromise) return initPromise;
+
+    initPromise = (async () => {
+        try {
+            await accountManager.initialize();
+            isInitialized = true;
+            const status = accountManager.getStatus();
+            console.log(`[Server] Account pool initialized: ${status.summary}`);
+        } catch (error) {
+            initError = error;
+            initPromise = null; // Allow retry on failure
+            console.error('[Server] Failed to initialize account manager:', error.message);
+            throw error;
+        }
+    })();
+
+    return initPromise;
+}
+
+// Middleware
+app.use(cors());
+app.use(express.json({ limit: REQUEST_BODY_LIMIT }));
+
+/**
+ * Parse error message to extract error type, status code, and user-friendly message
+ */
+function parseError(error) {
+    let errorType = 'api_error';
+    let statusCode = 500;
+    let errorMessage = error.message;
+
+    if (error.message.includes('401') || error.message.includes('UNAUTHENTICATED')) {
+        errorType = 'authentication_error';
+        statusCode = 401;
+        errorMessage = 'Authentication failed. Make sure Antigravity is running with a valid token.';
+    } else if (error.message.includes('429') || error.message.includes('RESOURCE_EXHAUSTED') || error.message.includes('QUOTA_EXHAUSTED')) {
+        errorType = 'invalid_request_error';  // Use invalid_request_error to force client to purge/stop
+        statusCode = 400;  // Use 400 to ensure client does not retry (429 and 529 trigger retries)
+
+        // Try to extract the quota reset time from the error
+        const resetMatch = error.message.match(/quota will reset after (\d+h\d+m\d+s|\d+m\d+s|\d+s)/i);
+        const modelMatch = error.message.match(/"model":\s*"([^"]+)"/);
+        const model = modelMatch ? modelMatch[1] : 'the model';
+
+        if (resetMatch) {
+            errorMessage = `You have exhausted your capacity on ${model}. Quota will reset after ${resetMatch[1]}.`;
+        } else {
+            errorMessage = `You have exhausted your capacity on ${model}. Please wait for your quota to reset.`;
+        }
+    } else if (error.message.includes('invalid_request_error') || error.message.includes('INVALID_ARGUMENT')) {
+        errorType = 'invalid_request_error';
+        statusCode = 400;
+        const msgMatch = error.message.match(/"message":"([^"]+)"/);
+        if (msgMatch) errorMessage = msgMatch[1];
+    } else if (error.message.includes('All endpoints failed')) {
+        errorType = 'api_error';
+        statusCode = 503;
+        errorMessage = 'Unable to connect to Claude API. Check that Antigravity is running.';
+    } else if (error.message.includes('PERMISSION_DENIED')) {
+        errorType = 'permission_error';
+        statusCode = 403;
+        errorMessage = 'Permission denied. Check your Antigravity license.';
+    }
+
+    return { errorType, statusCode, errorMessage };
+}
+
+// Request logging middleware
+app.use((req, res, next) => {
+    console.log(`[${new Date().toISOString()}] ${req.method} ${req.path}`);
+    next();
+});
+
+/**
+ * Health check endpoint
+ */
+app.get('/health', async (req, res) => {
+    try {
+        await ensureInitialized();
+        const status = accountManager.getStatus();
+
+        res.json({
+            status: 'ok',
+            accounts: status.summary,
+            available: status.available,
+            rateLimited: status.rateLimited,
+            invalid: status.invalid,
+            timestamp: new Date().toISOString()
+        });
+    } catch (error) {
+        res.status(503).json({
+            status: 'error',
+            error: error.message,
+            timestamp: new Date().toISOString()
+        });
+    }
+});
+
+/**
+ * Account limits endpoint - fetch quota/limits for all accounts × all models
+ * Returns a table showing remaining quota and reset time for each combination
+ * Use ?format=table for ASCII table output, default is JSON
+ */
+app.get('/account-limits', async (req, res) => {
+    try {
+        await ensureInitialized();
+        const allAccounts = accountManager.getAllAccounts();
+        const format = req.query.format || 'json';
+
+        // Fetch quotas for each account in parallel
+        const results = await Promise.allSettled(
+            allAccounts.map(async (account) => {
+                // Skip invalid accounts
+                if (account.isInvalid) {
+                    return {
+                        email: account.email,
+                        status: 'invalid',
+                        error: account.invalidReason,
+                        models: {}
+                    };
+                }
+
+                try {
+                    const token = await accountManager.getTokenForAccount(account);
+                    const quotas = await getModelQuotas(token);
+
+                    return {
+                        email: account.email,
+                        status: 'ok',
+                        models: quotas
+                    };
+                } catch (error) {
+                    return {
+                        email: account.email,
+                        status: 'error',
+                        error: error.message,
+                        models: {}
+                    };
+                }
+            })
+        );
+
+        // Process results
+        const accountLimits = results.map((result, index) => {
+            if (result.status === 'fulfilled') {
+                return result.value;
+            } else {
+                return {
+                    email: allAccounts[index].email,
+                    status: 'error',
+                    error: result.reason?.message || 'Unknown error',
+                    models: {}
+                };
+            }
+        });
+
+        // Collect all unique model IDs
+        const allModelIds = new Set();
+        for (const account of accountLimits) {
+            for (const modelId of Object.keys(account.models || {})) {
+                allModelIds.add(modelId);
+            }
+        }
+
+        const sortedModels = Array.from(allModelIds).filter(m => m.includes('claude')).sort();
+
+        // Return ASCII table format
+        if (format === 'table') {
+            res.setHeader('Content-Type', 'text/plain; charset=utf-8');
+
+            // Build table
+            const lines = [];
+            const timestamp = new Date().toLocaleString();
+            lines.push(`Account Limits (${timestamp})`);
+
+            // Get account status info
+            const status = accountManager.getStatus();
+            lines.push(`Accounts: ${status.total} total, ${status.available} available, ${status.rateLimited} rate-limited, ${status.invalid} invalid`);
+            lines.push('');
+
+            // Table 1: Account status
+            const accColWidth = 25;
+            const statusColWidth = 15;
+            const lastUsedColWidth = 25;
+            const resetColWidth = 25;
+
+            let accHeader = 'Account'.padEnd(accColWidth) + 'Status'.padEnd(statusColWidth) + 'Last Used'.padEnd(lastUsedColWidth) + 'Quota Reset';
+            lines.push(accHeader);
+            lines.push('─'.repeat(accColWidth + statusColWidth + lastUsedColWidth + resetColWidth));
+
+            for (const acc of status.accounts) {
+                const shortEmail = acc.email.split('@')[0].slice(0, 22);
+                const lastUsed = acc.lastUsed ? new Date(acc.lastUsed).toLocaleString() : 'never';
+
+                // Get status and error from accountLimits
+                const accLimit = accountLimits.find(a => a.email === acc.email);
+                let accStatus;
+                if (acc.isInvalid) {
+                    accStatus = 'invalid';
+                } else if (acc.isRateLimited) {
+                    const remaining = acc.rateLimitResetTime ? acc.rateLimitResetTime - Date.now() : 0;
+                    accStatus = remaining > 0 ? `limited (${formatDuration(remaining)})` : 'rate-limited';
+                } else {
+                    accStatus = accLimit?.status || 'ok';
+                }
+
+                // Get reset time from quota API
+                const claudeModel = sortedModels.find(m => m.includes('claude'));
+                const quota = claudeModel && accLimit?.models?.[claudeModel];
+                const resetTime = quota?.resetTime
+                    ? new Date(quota.resetTime).toLocaleString()
+                    : '-';
+
+                let row = shortEmail.padEnd(accColWidth) + accStatus.padEnd(statusColWidth) + lastUsed.padEnd(lastUsedColWidth) + resetTime;
+
+                // Add error on next line if present
+                if (accLimit?.error) {
+                    lines.push(row);
+                    lines.push('  └─ ' + accLimit.error);
+                } else {
+                    lines.push(row);
+                }
+            }
+            lines.push('');
+
+            // Calculate column widths
+            const modelColWidth = Math.max(25, ...sortedModels.map(m => m.length)) + 2;
+            const accountColWidth = 22;
+
+            // Header row
+            let header = 'Model'.padEnd(modelColWidth);
+            for (const acc of accountLimits) {
+                const shortEmail = acc.email.split('@')[0].slice(0, 18);
+                header += shortEmail.padEnd(accountColWidth);
+            }
+            lines.push(header);
+            lines.push('─'.repeat(modelColWidth + accountLimits.length * accountColWidth));
+
+            // Data rows
+            for (const modelId of sortedModels) {
+                let row = modelId.padEnd(modelColWidth);
+                for (const acc of accountLimits) {
+                    const quota = acc.models?.[modelId];
+                    let cell;
+                    if (acc.status !== 'ok') {
+                        cell = `[${acc.status}]`;
+                    } else if (!quota) {
+                        cell = '-';
+                    } else if (quota.remainingFraction === null) {
+                        cell = '0% (exhausted)';
+                    } else {
+                        const pct = Math.round(quota.remainingFraction * 100);
+                        cell = `${pct}%`;
+                    }
+                    row += cell.padEnd(accountColWidth);
+                }
+                lines.push(row);
+            }
+
+            return res.send(lines.join('\n'));
+        }
+
+        // Default: JSON format
+        res.json({
+            timestamp: new Date().toLocaleString(),
+            totalAccounts: allAccounts.length,
+            models: sortedModels,
+            accounts: accountLimits.map(acc => ({
+                email: acc.email,
+                status: acc.status,
+                error: acc.error || null,
+                limits: Object.fromEntries(
+                    sortedModels.map(modelId => {
+                        const quota = acc.models?.[modelId];
+                        if (!quota) {
+                            return [modelId, null];
+                        }
+                        return [modelId, {
+                            remaining: quota.remainingFraction !== null
+                                ? `${Math.round(quota.remainingFraction * 100)}%`
+                                : 'N/A',
+                            remainingFraction: quota.remainingFraction,
+                            resetTime: quota.resetTime || null
+                        }];
+                    })
+                )
+            }))
+        });
+    } catch (error) {
+        res.status(500).json({
+            status: 'error',
+            error: error.message
+        });
+    }
+});
+
+/**
+ * Force token refresh endpoint
+ */
+app.post('/refresh-token', async (req, res) => {
+    try {
+        await ensureInitialized();
+        // Clear all caches
+        accountManager.clearTokenCache();
+        accountManager.clearProjectCache();
+        // Force refresh default token
+        const token = await forceRefresh();
+        res.json({
+            status: 'ok',
+            message: 'Token caches cleared and refreshed',
+            tokenPrefix: token.substring(0, 10) + '...'
+        });
+    } catch (error) {
+        res.status(500).json({
+            status: 'error',
+            error: error.message
+        });
+    }
+});
+
+/**
+ * List models endpoint (OpenAI-compatible format)
+ */
+app.get('/v1/models', (req, res) => {
+    res.json(listModels());
+});
+
+/**
+ * Main messages endpoint - Anthropic Messages API compatible
+ */
+app.post('/v1/messages', async (req, res) => {
+    try {
+        // Ensure account manager is initialized
+        await ensureInitialized();
+
+        // Optimistic Retry: If ALL accounts are rate-limited, reset them to force a fresh check.
+        // If we have some available accounts, we try them first.
+        if (accountManager.isAllRateLimited()) {
+            console.log('[Server] All accounts rate-limited. Resetting state for optimistic retry.');
+            accountManager.resetAllRateLimits();
+        }
+
+        const {
+            model,
+            messages,
+            max_tokens,
+            stream,
+            system,
+            tools,
+            tool_choice,
+            thinking,
+            top_p,
+            top_k,
+            temperature
+        } = req.body;
+
+        // Validate required fields
+        if (!messages || !Array.isArray(messages)) {
+            return res.status(400).json({
+                type: 'error',
+                error: {
+                    type: 'invalid_request_error',
+                    message: 'messages is required and must be an array'
+                }
+            });
+        }
+
+        // Build the request object
+        const request = {
+            model: model || 'claude-3-5-sonnet-20241022',
+            messages,
+            max_tokens: max_tokens || 4096,
+            stream,
+            system,
+            tools,
+            tool_choice,
+            thinking,
+            top_p,
+            top_k,
+            temperature
+        };
+
+        console.log(`[API] Request for model: ${request.model}, stream: ${!!stream}`);
+
+        // Debug: Log message structure to diagnose tool_use/tool_result ordering
+        if (process.env.DEBUG) {
+            console.log('[API] Message structure:');
+            messages.forEach((msg, i) => {
+                const contentTypes = Array.isArray(msg.content)
+                    ? msg.content.map(c => c.type || 'text').join(', ')
+                    : (typeof msg.content === 'string' ? 'text' : 'unknown');
+                console.log(`  [${i}] ${msg.role}: ${contentTypes}`);
+            });
+        }
+
+        if (stream) {
+            // Handle streaming response
+            res.setHeader('Content-Type', 'text/event-stream');
+            res.setHeader('Cache-Control', 'no-cache');
+            res.setHeader('Connection', 'keep-alive');
+            res.setHeader('X-Accel-Buffering', 'no');
+
+            // Flush headers immediately to start the stream
+            res.flushHeaders();
+
+            try {
+                // Use the streaming generator with account manager
+                for await (const event of sendMessageStream(request, accountManager)) {
+                    res.write(`event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`);
+                    // Flush after each event for real-time streaming
+                    if (res.flush) res.flush();
+                }
+                res.end();
+
+            } catch (streamError) {
+                console.error('[API] Stream error:', streamError);
+
+                const { errorType, errorMessage } = parseError(streamError);
+
+                res.write(`event: error\ndata: ${JSON.stringify({
+                    type: 'error',
+                    error: { type: errorType, message: errorMessage }
+                })}\n\n`);
+                res.end();
+            }
+
+        } else {
+            // Handle non-streaming response
+            const response = await sendMessage(request, accountManager);
+            res.json(response);
+        }
+
+    } catch (error) {
+        console.error('[API] Error:', error);
+
+        let { errorType, statusCode, errorMessage } = parseError(error);
+
+        // For auth errors, try to refresh token
+        if (errorType === 'authentication_error') {
+            console.log('[API] Token might be expired, attempting refresh...');
+            try {
+                accountManager.clearProjectCache();
+                accountManager.clearTokenCache();
+                await forceRefresh();
+                errorMessage = 'Token was expired and has been refreshed. Please retry your request.';
+            } catch (refreshError) {
+                errorMessage = 'Could not refresh token. Make sure Antigravity is running.';
+            }
+        }
+
+        console.log(`[API] Returning error response: ${statusCode} ${errorType} - ${errorMessage}`);
+
+        // Check if headers have already been sent (for streaming that failed mid-way)
+        if (res.headersSent) {
+            console.log('[API] Headers already sent, writing error as SSE event');
+            res.write(`event: error\ndata: ${JSON.stringify({
+                type: 'error',
+                error: { type: errorType, message: errorMessage }
+            })}\n\n`);
+            res.end();
+        } else {
+            res.status(statusCode).json({
+                type: 'error',
+                error: {
+                    type: errorType,
+                    message: errorMessage
+                }
+            });
+        }
+    }
+});
+
+/**
+ * Catch-all for unsupported endpoints
+ */
+app.use('*', (req, res) => {
+    res.status(404).json({
+        type: 'error',
+        error: {
+            type: 'not_found_error',
+            message: `Endpoint ${req.method} ${req.originalUrl} not found`
+        }
+    });
+});
+
+export default app;

package/src/token-extractor.js

@@ -0,0 +1,146 @@
+/**
+ * Token Extractor Module
+ * Extracts OAuth tokens from Antigravity's SQLite database
+ *
+ * The database is automatically updated by Antigravity when tokens refresh,
+ * so this approach doesn't require any manual intervention.
+ */
+
+import { execSync } from 'child_process';
+import {
+    TOKEN_REFRESH_INTERVAL_MS,
+    ANTIGRAVITY_AUTH_PORT,
+    ANTIGRAVITY_DB_PATH
+} from './constants.js';
+
+// Cache for the extracted token
+let cachedToken = null;
+let tokenExtractedAt = null;
+
+/**
+ * Extract token from Antigravity's SQLite database
+ * This is the preferred method as the DB is auto-updated
+ */
+function extractTokenFromDB() {
+    try {
+        const result = execSync(
+            `sqlite3 "${ANTIGRAVITY_DB_PATH}" "SELECT value FROM ItemTable WHERE key = 'antigravityAuthStatus';"`,
+            { encoding: 'utf-8', timeout: 5000 }
+        );
+
+        if (!result || !result.trim()) {
+            throw new Error('No auth status found in database');
+        }
+
+        const authData = JSON.parse(result.trim());
+        return {
+            apiKey: authData.apiKey,
+            name: authData.name,
+            email: authData.email,
+            // Include other fields we might need
+            ...authData
+        };
+    } catch (error) {
+        console.error('[Token] Database extraction failed:', error.message);
+        throw error;
+    }
+}
+
+/**
+ * Extract the chat params from Antigravity's HTML page (fallback method)
+ */
+async function extractChatParams() {
+    try {
+        const response = await fetch(`http://127.0.0.1:${ANTIGRAVITY_AUTH_PORT}/`);
+        const html = await response.text();
+
+        // Find the base64-encoded chatParams in the HTML
+        const match = html.match(/window\.chatParams\s*=\s*'([^']+)'/);
+        if (!match) {
+            throw new Error('Could not find chatParams in Antigravity page');
+        }
+
+        // Decode base64
+        const base64Data = match[1];
+        const jsonString = Buffer.from(base64Data, 'base64').toString('utf-8');
+        const config = JSON.parse(jsonString);
+
+        return config;
+    } catch (error) {
+        if (error.code === 'ECONNREFUSED') {
+            throw new Error(
+                `Cannot connect to Antigravity on port ${ANTIGRAVITY_AUTH_PORT}. ` +
+                'Make sure Antigravity is running.'
+            );
+        }
+        throw error;
+    }
+}
+
+/**
+ * Get fresh token data - tries DB first, falls back to HTML page
+ */
+async function getTokenData() {
+    // Try database first (preferred - always has fresh token)
+    try {
+        const dbData = extractTokenFromDB();
+        if (dbData?.apiKey) {
+            console.log('[Token] Got fresh token from SQLite database');
+            return dbData;
+        }
+    } catch (err) {
+        console.log('[Token] DB extraction failed, trying HTML page...');
+    }
+
+    // Fallback to HTML page
+    try {
+        const pageData = await extractChatParams();
+        if (pageData?.apiKey) {
+            console.log('[Token] Got token from HTML page (may be stale)');
+            return pageData;
+        }
+    } catch (err) {
+        console.log('[Token] HTML page extraction failed:', err.message);
+    }
+
+    throw new Error(
+        'Could not extract token from Antigravity. ' +
+        'Make sure Antigravity is running and you are logged in.'
+    );
+}
+
+/**
+ * Check if the cached token needs refresh
+ */
+function needsRefresh() {
+    if (!cachedToken || !tokenExtractedAt) {
+        return true;
+    }
+    return Date.now() - tokenExtractedAt > TOKEN_REFRESH_INTERVAL_MS;
+}
+
+/**
+ * Get the current OAuth token (with caching)
+ */
+export async function getToken() {
+    if (needsRefresh()) {
+        const data = await getTokenData();
+        cachedToken = data.apiKey;
+        tokenExtractedAt = Date.now();
+    }
+    return cachedToken;
+}
+
+/**
+ * Force refresh the token (useful if requests start failing)
+ */
+export async function forceRefresh() {
+    cachedToken = null;
+    tokenExtractedAt = null;
+    return getToken();
+}
+
+export default {
+    getToken,
+    forceRefresh
+};

package/src/utils/helpers.js

@@ -0,0 +1,33 @@
+/**
+ * Shared Utility Functions
+ *
+ * General-purpose helper functions used across multiple modules.
+ */
+
+/**
+ * Format duration in milliseconds to human-readable string
+ * @param {number} ms - Duration in milliseconds
+ * @returns {string} Human-readable duration (e.g., "1h23m45s")
+ */
+export function formatDuration(ms) {
+    const seconds = Math.floor(ms / 1000);
+    const hours = Math.floor(seconds / 3600);
+    const minutes = Math.floor((seconds % 3600) / 60);
+    const secs = seconds % 60;
+
+    if (hours > 0) {
+        return `${hours}h${minutes}m${secs}s`;
+    } else if (minutes > 0) {
+        return `${minutes}m${secs}s`;
+    }
+    return `${secs}s`;
+}
+
+/**
+ * Sleep for specified milliseconds
+ * @param {number} ms - Duration to sleep in milliseconds
+ * @returns {Promise<void>} Resolves after the specified duration
+ */
+export function sleep(ms) {
+    return new Promise(resolve => setTimeout(resolve, ms));
+}