feat(proxy): improve parity, rate limiting, and token scheduling

src/services/GoogleAPIService.ts

@@ -1,6 +1,11 @@
 import { ConfigManager } from '../ipc/config/manager';
 import { ProxyAgent } from 'undici';
 import { logger } from '../utils/logger';
+import {
+  buildUserAgent,
+  FALLBACK_VERSION,
+  resolveLocalInstalledVersion,
+} from '@/server/modules/proxy/request-user-agent';
 
 // --- Constants & Config ---
 
@@ -52,13 +57,23 @@ export interface UserInfo {
 }
 
 export interface QuotaData {
-  models: Record<
-    string,
-    {
-      percentage: number;
-      resetTime: string;
-    }
-  >;
+  models: Record<string, ModelQuotaInfo>;
+  model_forwarding_rules?: Record<string, string>;
+  subscription_tier?: string;
+  is_forbidden?: boolean;
+}
+
+export interface ModelQuotaInfo {
+  percentage: number;
+  resetTime: string;
+  display_name?: string;
+  supports_images?: boolean;
+  supports_thinking?: boolean;
+  thinking_budget?: number;
+  recommended?: boolean;
+  max_tokens?: number;
+  max_output_tokens?: number;
+  supported_mime_types?: Record<string, boolean>;
 }
 
 // Internal types for API parsing
@@ -67,10 +82,138 @@ interface ModelInfoRaw {
     remainingFraction?: number;
     resetTime?: string;
   };
+  displayName?: string;
+  supportsImages?: boolean;
+  supportsThinking?: boolean;
+  thinkingBudget?: number;
+  recommended?: boolean;
+  maxTokens?: number;
+  maxOutputTokens?: number;
+  supportedMimeTypes?: Record<string, boolean>;
+}
+
+interface DeprecatedModelInfoRaw {
+  newModelId?: string;
+}
+
+interface IneligibleTierRaw {
+  reasonCode?: string;
+}
+
+interface TierRaw {
+  is_default?: boolean;
+  id?: string;
+  quotaTier?: string;
+  name?: string;
+  slug?: string;
 }
 
 interface LoadProjectResponse {
   cloudaicompanionProject?: string;
+  currentTier?: TierRaw;
+  paidTier?: TierRaw;
+  allowedTiers?: TierRaw[];
+  ineligibleTiers?: IneligibleTierRaw[];
+}
+
+interface FetchModelsResponse {
+  models?: Record<string, ModelInfoRaw>;
+  deprecatedModelIds?: Record<string, DeprecatedModelInfoRaw>;
+}
+
+interface ProjectContext {
+  projectId?: string;
+  subscriptionTier?: string;
+}
+
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms);
+  });
+}
+
+function buildInternalApiHeaders(accessToken: string): Record<string, string> {
+  const discoveryVersion = resolveLocalInstalledVersion() ?? FALLBACK_VERSION;
+  return {
+    Authorization: `Bearer ${accessToken}`,
+    'User-Agent': buildUserAgent(discoveryVersion),
+    'Content-Type': 'application/json',
+  };
+}
+
+function resolveSubscriptionTier(payload: LoadProjectResponse): string | undefined {
+  const paidTier = payload.paidTier;
+  if (paidTier?.name && paidTier.name.trim() !== '') {
+    return paidTier.name;
+  }
+  if (paidTier?.id && paidTier.id.trim() !== '') {
+    return paidTier.id;
+  }
+
+  const ineligible = Array.isArray(payload.ineligibleTiers) && payload.ineligibleTiers.length > 0;
+  if (!ineligible) {
+    const currentTier = payload.currentTier;
+    if (currentTier?.name && currentTier.name.trim() !== '') {
+      return currentTier.name;
+    }
+    if (currentTier?.id && currentTier.id.trim() !== '') {
+      return currentTier.id;
+    }
+  }
+
+  if (Array.isArray(payload.allowedTiers)) {
+    const preferredAllowedTier =
+      payload.allowedTiers.find((tier) => tier.is_default === true) ?? payload.allowedTiers[0];
+    if (preferredAllowedTier?.name && preferredAllowedTier.name.trim() !== '') {
+      return ineligible ? `${preferredAllowedTier.name} (Restricted)` : preferredAllowedTier.name;
+    }
+    if (preferredAllowedTier?.id && preferredAllowedTier.id.trim() !== '') {
+      return ineligible ? `${preferredAllowedTier.id} (Restricted)` : preferredAllowedTier.id;
+    }
+  }
+
+  return undefined;
+}
+
+function isTrackedModel(modelName: string): boolean {
+  return /^(gemini|claude|gpt|image|imagen)/i.test(modelName);
+}
+
+function toModelQuotaInfo(modelName: string, info: ModelInfoRaw): ModelQuotaInfo | null {
+  if (!isTrackedModel(modelName) || !info.quotaInfo) {
+    return null;
+  }
+
+  const fraction = info.quotaInfo.remainingFraction ?? 0;
+  return {
+    percentage: Math.floor(fraction * 100),
+    resetTime: info.quotaInfo.resetTime || '',
+    display_name: info.displayName,
+    supports_images: info.supportsImages,
+    supports_thinking: info.supportsThinking,
+    thinking_budget: info.thinkingBudget,
+    recommended: info.recommended,
+    max_tokens: info.maxTokens,
+    max_output_tokens: info.maxOutputTokens,
+    supported_mime_types: info.supportedMimeTypes,
+  };
+}
+
+function toModelForwardingRules(
+  deprecatedModelIds: FetchModelsResponse['deprecatedModelIds'],
+): Record<string, string> | undefined {
+  if (!deprecatedModelIds || Object.keys(deprecatedModelIds).length === 0) {
+    return undefined;
+  }
+
+  const forwardingRules: Record<string, string> = {};
+  for (const [oldModelId, deprecatedInfo] of Object.entries(deprecatedModelIds)) {
+    if (typeof deprecatedInfo.newModelId === 'string' && deprecatedInfo.newModelId !== '') {
+      forwardingRules[oldModelId] = deprecatedInfo.newModelId;
+    }
+  }
+
+  return Object.keys(forwardingRules).length > 0 ? forwardingRules : undefined;
 }
 
 // --- Service Implementation ---
@@ -218,70 +361,71 @@ export class GoogleAPIService {
     return response.json() as Promise<UserInfo>;
   }
 
-  /**
-   * Core logic: Fetches the internal project ID needed for quota checks.
-   */
-  public static async fetchProjectId(accessToken: string): Promise<string | null> {
+  public static async fetchProjectContext(accessToken: string): Promise<ProjectContext> {
     const body = {
       metadata: { ideType: 'ANTIGRAVITY' },
     };
 
-    // Simple retry logic (2 attempts)
+    let projectId: string | undefined;
+    let subscriptionTier: string | undefined;
+
     for (let i = 0; i < 2; i++) {
       try {
         const response = await fetch(URLS.LOAD_PROJECT, {
           method: 'POST',
-          headers: {
-            Authorization: `Bearer ${accessToken}`,
-            'User-Agent': USER_AGENT,
-            'Content-Type': 'application/json',
-          },
+          headers: buildInternalApiHeaders(accessToken),
           body: JSON.stringify(body),
+          signal: createTimeoutSignal(REQUEST_TIMEOUT_MS),
           ...this.getFetchOptions(),
         });
 
         if (response.ok) {
           const data = (await response.json()) as LoadProjectResponse;
-          if (data.cloudaicompanionProject) {
-            return data.cloudaicompanionProject;
+          if (typeof data.cloudaicompanionProject === 'string') {
+            projectId = data.cloudaicompanionProject;
           }
+          subscriptionTier = resolveSubscriptionTier(data);
+          break;
         }
-      } catch (e) {
-        logger.warn(`[GoogleAPIService] Failed to fetch project ID (Attempt ${i + 1}):`, e);
-        await new Promise((r) => setTimeout(r, 500)); // Sleep 500ms
+      } catch (error) {
+        logger.warn(`[GoogleAPIService] Failed to fetch project ID (Attempt ${i + 1}):`, error);
+        await sleep(500);
       }
     }
-    return null;
+
+    return {
+      projectId,
+      subscriptionTier,
+    };
+  }
+
+  public static async fetchProjectId(accessToken: string): Promise<string | null> {
+    const context = await this.fetchProjectContext(accessToken);
+    return context.projectId ?? null;
   }
 
   /**
    * Core logic: Fetches detailed model quota information.
    */
   static async fetchQuota(accessToken: string): Promise<QuotaData> {
-    // 1. Get Project ID (Critical)
-    const projectId = await this.fetchProjectId(accessToken);
+    const { projectId, subscriptionTier } = await this.fetchProjectContext(accessToken);
 
-    // 2. Build Payload
     const payload: Record<string, unknown> = {};
     if (projectId) {
       payload['project'] = projectId;
     }
 
-    // 3. Send Request with Retries
     const maxRetries = 3;
     let lastError: Error | null = null;
-    const fetchOptions = this.getFetchOptions(); // Reuse options
+    const fetchOptions = this.getFetchOptions();
 
     for (let attempt = 1; attempt <= maxRetries; attempt++) {
       try {
         const response = await fetch(URLS.QUOTA, {
           method: 'POST',
-          headers: {
-            Authorization: `Bearer ${accessToken}`,
-            'User-Agent': USER_AGENT,
-            'Content-Type': 'application/json',
-          },
+          headers: buildInternalApiHeaders(accessToken),
           body: JSON.stringify(payload),
+          signal: createTimeoutSignal(REQUEST_TIMEOUT_MS),
           ...fetchOptions,
         });
 
@@ -302,39 +446,42 @@ export class GoogleAPIService {
           );
 
           if (attempt < maxRetries) {
-            await new Promise((r) => setTimeout(r, 1000));
+            await sleep(1000);
             continue;
           } else {
             throw new Error(errorMsg);
           }
         }
 
-        const data = (await response.json()) as { models: Record<string, ModelInfoRaw> };
-        const result: QuotaData = { models: {} };
+        const data = (await response.json()) as FetchModelsResponse;
+        const result: QuotaData = {
+          models: {},
+          subscription_tier: subscriptionTier,
+          is_forbidden: false,
+        };
 
-        // Parse relevant models
-        for (const [name, info] of Object.entries(data.models || {})) {
-          if (info.quotaInfo) {
-            const fraction = info.quotaInfo.remainingFraction ?? 0;
-            const percentage = Math.floor(fraction * 100);
-            const resetTime = info.quotaInfo.resetTime || '';
-
-            if (name.toLowerCase().includes('gemini') || name.toLowerCase().includes('claude')) {
-              result.models[name] = { percentage, resetTime };
-            }
+        for (const [modelName, modelInfoRaw] of Object.entries(data.models || {})) {
+          const modelQuota = toModelQuotaInfo(modelName, modelInfoRaw);
+          if (modelQuota) {
+            result.models[modelName] = modelQuota;
           }
         }
 
+        const modelForwardingRules = toModelForwardingRules(data.deprecatedModelIds);
+        if (modelForwardingRules) {
+          result.model_forwarding_rules = modelForwardingRules;
+        }
+
         return result;
-      } catch (e: unknown) {
-        if (e instanceof Error) {
+      } catch (error: unknown) {
+        if (error instanceof Error) {
           logger.warn(
-            `[GoogleAPIService] Request failed: ${e.message} (Attempt ${attempt}/${maxRetries})`,
+            `[GoogleAPIService] Request failed: ${error.message} (Attempt ${attempt}/${maxRetries})`,
           );
         }
-        lastError = e instanceof Error ? e : new Error(String(e));
+        lastError = error instanceof Error ? error : new Error(String(error));
         if (attempt < maxRetries) {
-          await new Promise((r) => setTimeout(r, 1000));
+          await sleep(1000);
         }
       }
     }