import axios, { AxiosError, AxiosInstance } from 'axios'; import { QuotaData, Tier, LoadProjectResponse, QuotaApiResponse } from './types'; // Constants const QUOTA_API_URL = 'https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels'; const CLOUD_CODE_BASE_URL = 'https://cloudcode-pa.googleapis.com'; const USER_AGENT = 'antigravity/1.11.3 Darwin/arm64'; // Keeping the same UA as source // Service Class export class QuotaService { private static createClient(timeoutSecs: number = 15): AxiosInstance { return axios.create({ timeout: timeoutSecs * 1000, headers: { 'User-Agent': USER_AGENT, 'Content-Type': 'application/json', }, }); } /** * Fetch Project ID and Subscription Type */ private static async fetchProjectId( accessToken: string, email: string, ): Promise<[string | undefined, string | undefined]> { const client = this.createClient(); const meta = { metadata: { ideType: 'ANTIGRAVITY' } }; try { const res = await client.post( `${CLOUD_CODE_BASE_URL}/v1internal:loadCodeAssist`, meta, { headers: { Authorization: `Bearer ${accessToken}`, 'User-Agent': 'antigravity/windows/amd64', }, }, ); if (res.status >= 200 && res.status < 300) { const data = res.data; const projectId = data.cloudaicompanionProject; // Core logic: Preferentially get subscription ID from paid_tier const subscriptionTier = data.paidTier?.id || data.currentTier?.id; if (subscriptionTier) { console.log(`📊 [${email}] Subscription Identified: ${subscriptionTier}`); } return [projectId, subscriptionTier]; } else { console.warn(`⚠️ [${email}] loadCodeAssist failed: Status: ${res.status}`); } } catch (error: any) { console.error(`❌ [${email}] loadCodeAssist Network Error: ${error.message}`); } return [undefined, undefined]; } /** * Unified entry point for querying account quota */ public static async fetchQuota(accessToken: string, email: string) { return this.fetchQuotaInner(accessToken, email); } /** * Logic for querying account quota (Inner) */ private static async fetchQuotaInner( accessToken: string, email: string, ): Promise<{ quotaData: QuotaData; projectId?: string }> { // 1. Get Project ID and Subscription Type const [projectId, subscriptionTier] = await this.fetchProjectId(accessToken, email); const finalProjectId = projectId; const client = this.createClient(); const payload = { project: finalProjectId }; const url = QUOTA_API_URL; const maxRetries = 3; console.log(`Sending quota request to ${url}`); for (let attempt = 1; attempt <= maxRetries; attempt++) { try { const response = await client.post(url, payload, { headers: { Authorization: `Bearer ${accessToken}`, 'User-Agent': USER_AGENT, }, }); const quotaResponse = response.data; const quotaData: QuotaData = { models: {}, isForbidden: false, subscriptionTier: subscriptionTier, }; console.log(`Quota API returned ${Object.keys(quotaResponse.models || {}).length} models:`); if (quotaResponse.models) { for (const [name, info] of Object.entries(quotaResponse.models)) { console.log(` - ${name}`); if (info.quotaInfo) { const fraction = info.quotaInfo.remainingFraction ?? 0; const percentage = Math.floor(fraction * 100); const resetTime = info.quotaInfo.resetTime || ''; // Only save models we care about if (name.includes('gemini') || name.includes('claude')) { quotaData.models[name] = { percentage, resetTime }; } } } } return { quotaData, projectId: projectId }; } catch (error: any) { if (axios.isAxiosError(error)) { const status = error.response?.status; let text = ''; try { text = JSON.stringify(error.response?.data || ''); } catch { text = '[Unable to serialize response data]'; } // ✅ Handle 403 Forbidden specifically - return immediately, do not retry if (status === 403) { console.warn(`Account no permission (403 Forbidden), marked as forbidden`); return { quotaData: { models: {}, isForbidden: true, subscriptionTier: subscriptionTier, }, projectId: projectId, }; } console.warn(`API Error: ${status} - ${text} (Attempt ${attempt}/${maxRetries})`); } else { console.warn(`Request Failed: ${error.message} (Attempt ${attempt}/${maxRetries})`); } if (attempt < maxRetries) { await new Promise((r) => setTimeout(r, 1000)); } else { throw new Error(`Quota query failed: ${error.message}`); } } } throw new Error('Unknown error in fetchQuota'); // Should not reach here } }