feat: add tier detection and split gemini model families for rate limiting (#7)

* feat: add tier detection and split gemini model families for rate limiting

- Add automatic tier detection (free vs paid) via loadCodeAssist API
- Split ModelFamily: gemini -> gemini-flash + gemini-pro for independent rate limiting
- Implement paid account prioritization in multi-account rotation
- Add strict account stickiness to preserve context cache
- Remove redundant code in OAuth flow
- Add comprehensive tier priority tests

Fixes #6

* fix: add storage v3 migration for split gemini model families

- Bump storage version to 3 with gemini-flash and gemini-pro keys
- Add v2 interface with legacy single 'gemini' key
- Add migrateV2ToV3 that splits gemini -> gemini-flash + gemini-pro
- Migrations now preserve expired rate limits (cleanup at runtime)
- Add tests for v2->v3 migration and tier preservation

Author: theblazehen

.gitignore

@@ -36,4 +36,3 @@ gemini-debug-*.log
 
 # Log files
 *.log
-CLIProxyAPI

src/antigravity/oauth.ts

@@ -32,6 +32,7 @@ interface AntigravityTokenExchangeSuccess {
   expires: number;
   email?: string;
   projectId: string;
+  tier?: "free" | "paid";
 }
 
 interface AntigravityTokenExchangeFailure {
@@ -72,10 +73,10 @@ function decodeState(state: string): AntigravityAuthState {
 }
 
 /**
- * Automatically discover project ID from Antigravity API.
- * Tries multiple endpoints to find the user's default project.
+ * Automatically discover project ID and account tier from Antigravity API.
+ * Tries multiple endpoints to find the user's default project and tier.
  */
-async function fetchProjectID(accessToken: string): Promise<string> {
+async function fetchAccountInfo(accessToken: string): Promise<{ projectId: string; tier: "free" | "paid" }> {
   const errors: string[] = [];
 
   const loadHeaders: Record<string, string> = {
@@ -110,15 +111,38 @@ async function fetchProjectID(accessToken: string): Promise<string> {
       }
 
       const data = await response.json();
+      let projectId = "";
+      let tier: "free" | "paid" = "free";
+
+      // Extract Project ID
       if (typeof data.cloudaicompanionProject === "string" && data.cloudaicompanionProject) {
-        return data.cloudaicompanionProject;
-      }
-      if (
+        projectId = data.cloudaicompanionProject;
+      } else if (
         data.cloudaicompanionProject &&
         typeof data.cloudaicompanionProject.id === "string" &&
         data.cloudaicompanionProject.id
       ) {
-        return data.cloudaicompanionProject.id;
+        projectId = data.cloudaicompanionProject.id;
+      }
+
+      // Extract Tier
+      // Default to "free" (legacy-tier), check for paid indicators
+      if (Array.isArray(data.allowedTiers)) {
+        const defaultTier = data.allowedTiers.find((t: any) => t.isDefault);
+        if (defaultTier && typeof defaultTier.id === "string") {
+          const tierId = defaultTier.id;
+          // "legacy-tier" is the default free tier. Anything else is likely paid/upgraded.
+          // We can refine this list as we learn more about tier IDs.
+          if (tierId !== "legacy-tier" && !tierId.includes("free") && !tierId.includes("zero")) {
+            tier = "paid";
+          } else if (tierId !== "legacy-tier") {
+             console.debug(`[antigravity] Detected free tier variant: ${tierId}`);
+          }
+        }
+      }
+
+      if (projectId) {
+        return { projectId, tier };
       }
 
       errors.push(`loadCodeAssist missing project id at ${baseEndpoint}`);
@@ -132,9 +156,9 @@ async function fetchProjectID(accessToken: string): Promise<string> {
   }
 
   if (errors.length) {
-    console.warn("Failed to resolve Antigravity project via loadCodeAssist:", errors.join("; "));
+    console.warn("Failed to resolve Antigravity account info via loadCodeAssist:", errors.join("; "));
   }
-  return "";
+  return { projectId: "", tier: "free" };
 }
 
 export async function authorizeAntigravity(projectId = ""): Promise<AntigravityAuthorization> {
@@ -208,11 +232,14 @@ export async function exchangeAntigravity(
       return { type: "failed", error: "Missing refresh token in response" };
     }
 
-    // Auto-discover project ID if not provided
+    // Auto-discover project ID and tier if not provided, or fetch tier if project ID is provided
     let effectiveProjectId = projectId;
+
+    const accountInfo = await fetchAccountInfo(tokenPayload.access_token);
     if (!effectiveProjectId) {
-      effectiveProjectId = await fetchProjectID(tokenPayload.access_token);
+      effectiveProjectId = accountInfo.projectId;
     }
+    const tier = accountInfo.tier;
 
     // Don't embed email in refresh token - we store it separately in accounts.json
     const storedRefresh = `${refreshToken}|${effectiveProjectId || ""}`;
@@ -224,6 +251,7 @@ export async function exchangeAntigravity(
       expires: Date.now() + tokenPayload.expires_in * 1000,
       email: userInfo.email,
       projectId: effectiveProjectId || "",
+      tier,
     };
   } catch (error) {
     return {

src/plugin.ts

@@ -3,7 +3,7 @@ import type { AntigravityTokenExchangeResult } from "./antigravity/oauth";
 import { authorizeAntigravity, exchangeAntigravity } from "./antigravity/oauth";
 import { ANTIGRAVITY_PROVIDER_ID, MAX_ACCOUNTS } from "./constants";
 import { accessTokenExpired, isOAuthAuth, parseRefreshParts, formatMultiAccountRefresh } from "./plugin/auth";
-import { AccountManager } from "./plugin/accounts";
+import { AccountManager, type AccountTier } from "./plugin/accounts";
 import { openBrowser } from "./plugin/browser";
 import { promptProjectId, promptAddAnotherAccount } from "./plugin/cli";
 import { createAntigravityFetch } from "./plugin/fetch-wrapper";
@@ -35,7 +35,7 @@ async function getAuthContext(
 
   const storedAccounts = await loadAccounts();
   const accountManager = new AccountManager(auth, storedAccounts);
-  const account = accountManager.getCurrentOrNextForFamily("gemini");
+  const account = accountManager.getCurrentOrNextForFamily("gemini-flash");
   if (!account) {
     return null;
   }
@@ -107,7 +107,7 @@ function createGoogleSearchTool(getAuth: GetAuth, client: PluginContext["client"
 async function authenticateSingleAccount(
   client: PluginContext["client"],
   isHeadless: boolean,
-): Promise<{ refresh: string; access: string; expires: number; projectId: string; email?: string } | null> {
+): Promise<{ refresh: string; access: string; expires: number; projectId: string; email?: string; tier?: AccountTier } | null> {
   let listener: OAuthListener | null = null;
   if (!isHeadless) {
     try {
@@ -229,6 +229,7 @@ async function authenticateSingleAccount(
     expires: result.expires,
     projectId: result.projectId,
     email: result.email,
+    tier: result.tier as AccountTier,
   };
 }
 
@@ -280,6 +281,7 @@ export const AntigravityOAuthPlugin = async ({ client }: PluginContext): Promise
               expires: number;
               projectId: string;
               email?: string;
+              tier?: AccountTier;
             }> = [];
 
             const firstAccount = await authenticateSingleAccount(client, isHeadless);
@@ -337,11 +339,12 @@ export const AntigravityOAuthPlugin = async ({ client }: PluginContext): Promise
 
             try {
               await saveAccounts({
-                version: 2,
+                version: 3,
                 accounts: accounts.map((acc, index) => ({
                   email: acc.email,
                   refreshToken: acc.refresh,
                   projectId: acc.projectId,
+                  tier: acc.tier,
                   managedProjectId: undefined,
                   addedAt: Date.now(),
                   lastUsed: index === 0 ? Date.now() : 0,

src/plugin/accounts.test.ts

@@ -2,7 +2,7 @@ import { describe, it, expect } from "bun:test";
 import { AccountManager, type ModelFamily } from "./accounts";
 import type { OAuthAuthDetails } from "./types";
 
-const FAMILY: ModelFamily = "gemini";
+const FAMILY: ModelFamily = "gemini-flash";
 
 describe("AccountManager", () => {
   it("should initialize with single account", () => {
@@ -149,7 +149,7 @@ describe("AccountManager", () => {
     const claudeAccount = manager.getCurrentOrNextForFamily("claude");
     expect(claudeAccount).toBeNull();
 
-    const geminiAccount = manager.getCurrentOrNextForFamily("gemini");
+    const geminiAccount = manager.getCurrentOrNextForFamily("gemini-flash");
     expect(geminiAccount).not.toBeNull();
     expect(geminiAccount?.index).toBe(0);
   });

src/plugin/accounts.ts

@@ -5,9 +5,9 @@ import {
   parseRefreshParts,
   formatRefreshParts,
 } from "./auth";
-import { saveAccounts, type AccountStorage, type RateLimitState, type ModelFamily } from "./storage";
+import { saveAccounts, type AccountStorage, type RateLimitState, type ModelFamily, type AccountTier } from "./storage";
 
-export type { ModelFamily } from "./storage";
+export type { ModelFamily, AccountTier } from "./storage";
 
 export interface ManagedAccount {
   index: number;
@@ -17,6 +17,7 @@ export interface ManagedAccount {
   rateLimitResetTimes: RateLimitState;
   lastUsed: number;
   email?: string;
+  tier?: AccountTier;
   lastSwitchReason?: "rate-limit" | "initial" | "rotation";
 }
 
@@ -30,8 +31,11 @@ function clearExpiredRateLimits(account: ManagedAccount): void {
   if (account.rateLimitResetTimes.claude !== undefined && now >= account.rateLimitResetTimes.claude) {
     delete account.rateLimitResetTimes.claude;
   }
-  if (account.rateLimitResetTimes.gemini !== undefined && now >= account.rateLimitResetTimes.gemini) {
-    delete account.rateLimitResetTimes.gemini;
+  if (account.rateLimitResetTimes["gemini-flash"] !== undefined && now >= account.rateLimitResetTimes["gemini-flash"]) {
+    delete account.rateLimitResetTimes["gemini-flash"];
+  }
+  if (account.rateLimitResetTimes["gemini-pro"] !== undefined && now >= account.rateLimitResetTimes["gemini-pro"]) {
+    delete account.rateLimitResetTimes["gemini-pro"];
   }
 }
 
@@ -67,6 +71,7 @@ export class AccountManager {
         rateLimitResetTimes: acc.rateLimitResetTimes ?? {},
         lastUsed: acc.lastUsed,
         email: acc.email,
+        tier: acc.tier,
         lastSwitchReason: acc.lastSwitchReason,
       }));
     } else {
@@ -99,9 +104,10 @@ export class AccountManager {
 
   async save(): Promise<void> {
     const storage: AccountStorage = {
-      version: 2,
+      version: 3,
       accounts: this.accounts.map((acc) => ({
         email: acc.email,
+        tier: acc.tier,
         refreshToken: acc.parts.refreshToken,
         projectId: acc.parts.projectId,
         managedProjectId: acc.parts.managedProjectId,
@@ -133,12 +139,19 @@ export class AccountManager {
   }
 
   getCurrentOrNextForFamily(family: ModelFamily): ManagedAccount | null {
+    this.accounts.forEach(clearExpiredRateLimits);
+
     const current = this.getCurrentAccount();
     if (current) {
-      clearExpiredRateLimits(current);
       if (!isRateLimitedForFamily(current, family)) {
-        current.lastUsed = Date.now();
-        return current;
+        const betterTierAvailable =
+          current.tier !== "paid" &&
+          this.accounts.some((a) => a.tier === "paid" && !isRateLimitedForFamily(a, family));
+
+        if (!betterTierAvailable) {
+          current.lastUsed = Date.now();
+          return current;
+        }
       }
     }
 
@@ -150,16 +163,17 @@ export class AccountManager {
   }
 
   getNextForFamily(family: ModelFamily): ManagedAccount | null {
-    const available = this.accounts.filter((a) => {
-      clearExpiredRateLimits(a);
-      return !isRateLimitedForFamily(a, family);
-    });
+    const available = this.accounts.filter((a) => !isRateLimitedForFamily(a, family));
 
     if (available.length === 0) {
       return null;
     }
 
-    const account = available[this.currentIndex % available.length];
+    // Prioritize paid accounts
+    const paidAvailable = available.filter((a) => a.tier === "paid");
+    const pool = paidAvailable.length > 0 ? paidAvailable : available;
+
+    const account = pool[this.currentIndex % pool.length];
     if (!account) {
       return null;
     }
@@ -195,7 +209,7 @@ export class AccountManager {
     };
   }
 
-  addAccount(parts: RefreshParts, access?: string, expires?: number, email?: string): void {
+  addAccount(parts: RefreshParts, access?: string, expires?: number, email?: string, tier?: AccountTier): void {
     this.accounts.push({
       index: this.accounts.length,
       parts,
@@ -204,6 +218,7 @@ export class AccountManager {
       rateLimitResetTimes: {},
       lastUsed: 0,
       email,
+      tier,
     });
   }
 

src/plugin/cache.ts

@@ -69,7 +69,7 @@ export function clearCachedAuth(refresh?: string): void {
 
 const signatureCache = new Map<string, string>();
 
-export type ModelFamily = "claude" | "gemini";
+export type ModelFamily = "claude" | "gemini-flash" | "gemini-pro";
 
 /**
  * Generates a SHA-256 hash key for a thought block.

src/plugin/fetch-wrapper.ts

@@ -52,7 +52,10 @@ function getModelFamilyFromUrl(urlString: string): ModelFamily {
   if (model && model.includes("claude")) {
     return "claude";
   }
-  return "gemini";
+  if (model && model.includes("flash")) {
+    return "gemini-flash";
+  }
+  return "gemini-pro";
 }
 
 export function sleep(ms: number, signal?: AbortSignal | null): Promise<void> {
@@ -464,6 +467,7 @@ async function handleServerError(
 ): Promise<EndpointLoopResult> {
   const retryAfterMs = 60000;
 
+  // For 500 errors, we use a fixed short retry (1 min) rather than the heavy defaults
   accountManager.markRateLimited(account, retryAfterMs, family);
 
   log.warn(`Account ${account.index + 1}/${accountCount} received ${response.status} error on all endpoints`, {

src/plugin/request.ts

@@ -39,7 +39,11 @@ const MODEL_FALLBACKS: Record<string, string> = {
 };
 
 function getModelFamily(model: string): ModelFamily {
-  return model.includes("claude") ? "claude" : "gemini";
+  return model.includes("claude") 
+    ? "claude" 
+    : model.includes("flash") 
+      ? "gemini-flash" 
+      : "gemini-pro";
 }
 
 export function isGenerativeLanguageRequest(input: RequestInfo): boolean {
@@ -642,7 +646,7 @@ export async function transformAntigravityResponse(
   const contentType = response.headers.get("content-type") ?? "";
   const isJsonResponse = contentType.includes("application/json");
   const isEventStreamResponse = contentType.includes("text/event-stream");
-  const family: ModelFamily = requestedModel ? getModelFamily(requestedModel) : "gemini";
+  const family: ModelFamily = requestedModel ? getModelFamily(requestedModel) : "gemini-flash";
 
   if (!isJsonResponse && !isEventStreamResponse) {
     logAntigravityDebugResponse(debugContext, response, {