fix(quota): use API tier detection instead of model-based heuristics

Tier detection now uses paidTier/currentTier from loadCodeAssist API
response instead of inferring from model names. This is the correct
approach - all Antigravity accounts have access to same models
regardless of tier.

- Add TierInfo interface and tier fields to LoadCodeAssistResponse
- Add mapTierString() helper for API → AccountTier conversion
- Update getProjectId() to extract and return tier
- Update fetchAccountQuota() to use API tier
- Remove flawed detectTier() function
This commit is contained in:
kaitranntt
2026-01-06 14:41:25 -05:00
parent 0af185f6a0
commit aad0d44069
+69 -67
View File
@@ -111,42 +111,18 @@ interface TokenRefreshResponse {
error_description?: string; error_description?: string;
} }
/** /** Tier info from loadCodeAssist */
* Detect account tier from quota API model list. interface TierInfo {
* Ultra accounts have access to experimental/preview models. id?: string;
* Pro accounts have standard model access.
* Free/unknown accounts have limited model access.
*/
export function detectTier(quotaResult: QuotaResult): AccountTier {
if (!quotaResult.success || quotaResult.models.length === 0) {
return 'unknown';
}
const modelNames = quotaResult.models.map((m) => m.name.toLowerCase());
// Ultra indicators: experimental, preview, ultra in name
// Note: 2.5-pro is a standard pro model, not ultra
const ultraIndicators = ['ultra', 'experimental', 'preview', '3-ultra'];
const hasUltra = modelNames.some((name) =>
ultraIndicators.some((indicator) => name.includes(indicator))
);
if (hasUltra) return 'ultra';
// Pro indicators: gemini-2.x, gemini-3, pro models (includes 2.5-pro)
const proIndicators = ['gemini-2', 'gemini-3', '-pro'];
const hasPro = modelNames.some((name) =>
proIndicators.some((indicator) => name.includes(indicator))
);
if (hasPro) return 'pro';
return 'free';
} }
/** loadCodeAssist response */ /** loadCodeAssist response */
interface LoadCodeAssistResponse { interface LoadCodeAssistResponse {
cloudaicompanionProject?: string | { id?: string }; cloudaicompanionProject?: string | { id?: string };
/** Current tier (may be trial/temporary) */
currentTier?: TierInfo;
/** Paid tier (reflects actual subscription - takes priority) */
paidTier?: TierInfo;
} }
/** fetchAvailableModels response model */ /** fetchAvailableModels response model */
@@ -308,11 +284,27 @@ function readAuthData(provider: CLIProxyProvider, accountId: string): AuthData |
} }
/** /**
* Get project ID via loadCodeAssist endpoint * Map API tier string to AccountTier type
* API returns: "FREE", "PRO", "ULTRA" (or variants like "pro", "ultra")
*/ */
async function getProjectId( function mapTierString(tierStr: string | undefined): AccountTier {
accessToken: string if (!tierStr) return 'unknown';
): Promise<{ projectId: string | null; error?: string; isUnprovisioned?: boolean }> { const normalized = tierStr.toLowerCase();
if (normalized.includes('ultra')) return 'ultra';
if (normalized.includes('pro')) return 'pro';
if (normalized.includes('free')) return 'free';
return 'unknown';
}
/**
* Get project ID and subscription tier via loadCodeAssist endpoint
*/
async function getProjectId(accessToken: string): Promise<{
projectId: string | null;
tier?: AccountTier;
error?: string;
isUnprovisioned?: boolean;
}> {
const url = `${ANTIGRAVITY_API_BASE}/${ANTIGRAVITY_API_VERSION}:loadCodeAssist`; const url = `${ANTIGRAVITY_API_BASE}/${ANTIGRAVITY_API_VERSION}:loadCodeAssist`;
const controller = new AbortController(); const controller = new AbortController();
@@ -367,7 +359,11 @@ async function getProjectId(
}; };
} }
return { projectId: projectId.trim() }; // Extract tier: priority paidTier > currentTier (paid reflects actual subscription)
const tierStr = data.paidTier?.id || data.currentTier?.id;
const tier = mapTierString(tierStr);
return { projectId: projectId.trim(), tier };
} catch (err) { } catch (err) {
clearTimeout(timeoutId); clearTimeout(timeoutId);
if (err instanceof Error && err.name === 'AbortError') { if (err instanceof Error && err.name === 'AbortError') {
@@ -529,57 +525,63 @@ export async function fetchAccountQuota(
} }
} }
// Get project ID - prefer stored value, fallback to API call // Get project ID and tier - prefer stored project ID, but always call API for tier
let projectId = authData.projectId; let projectId = authData.projectId;
if (!projectId) { let apiTier: AccountTier = 'unknown';
let lastProjectResult = await getProjectId(accessToken);
if (!lastProjectResult.projectId) { // Always call loadCodeAssist to get accurate tier from API
// If project ID fetch fails, it might be token issue - try refresh if we haven't let lastProjectResult = await getProjectId(accessToken);
if (authData.refreshToken && accessToken === authData.accessToken) {
const refreshResult = await refreshAccessToken(authData.refreshToken); if (!lastProjectResult.projectId && !projectId) {
if (refreshResult.accessToken) { // If project ID fetch fails, it might be token issue - try refresh if we haven't
accessToken = refreshResult.accessToken; if (authData.refreshToken && accessToken === authData.accessToken) {
lastProjectResult = await getProjectId(accessToken); const refreshResult = await refreshAccessToken(authData.refreshToken);
} if (refreshResult.accessToken) {
} accessToken = refreshResult.accessToken;
if (!lastProjectResult.projectId) { lastProjectResult = await getProjectId(accessToken);
return {
success: false,
models: [],
lastUpdated: Date.now(),
error: lastProjectResult.error || 'Failed to retrieve project ID',
isUnprovisioned: lastProjectResult.isUnprovisioned,
};
} }
} }
projectId = lastProjectResult.projectId; if (!lastProjectResult.projectId) {
return {
success: false,
models: [],
lastUpdated: Date.now(),
error: lastProjectResult.error || 'Failed to retrieve project ID',
isUnprovisioned: lastProjectResult.isUnprovisioned,
};
}
} }
// Use API project ID if available, else fallback to stored
projectId = lastProjectResult.projectId || projectId;
apiTier = lastProjectResult.tier || 'unknown';
// Fetch models with quota // Fetch models with quota
const result = await fetchAvailableModels(accessToken, projectId); const result = await fetchAvailableModels(accessToken, projectId as string);
// If quota fetch fails with auth error and we haven't refreshed yet, try refresh // If quota fetch fails with auth error and we haven't refreshed yet, try refresh
if (!result.success && result.error?.includes('expired') && authData.refreshToken) { if (!result.success && result.error?.includes('expired') && authData.refreshToken) {
const refreshResult = await refreshAccessToken(authData.refreshToken); const refreshResult = await refreshAccessToken(authData.refreshToken);
if (refreshResult.accessToken) { if (refreshResult.accessToken) {
const retryResult = await fetchAvailableModels(refreshResult.accessToken, projectId); const retryResult = await fetchAvailableModels(
// Detect and persist tier for retry result refreshResult.accessToken,
projectId as string
);
// Use API tier (from loadCodeAssist) instead of model-based detection
if (retryResult.success) { if (retryResult.success) {
const tier = detectTier(retryResult); retryResult.tier = apiTier;
retryResult.tier = tier;
retryResult.accountId = accountId; retryResult.accountId = accountId;
setAccountTier(provider, accountId, tier); setAccountTier(provider, accountId, apiTier);
} }
return retryResult; return retryResult;
} }
} }
// Detect and persist tier for successful result // Use API tier (from loadCodeAssist) instead of model-based detection
if (result.success) { if (result.success) {
const tier = detectTier(result); result.tier = apiTier;
result.tier = tier;
result.accountId = accountId; result.accountId = accountId;
setAccountTier(provider, accountId, tier); setAccountTier(provider, accountId, apiTier);
} }
return result; return result;