diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli.ts index fb65b523..80b8b6c0 100644 --- a/src/cliproxy/quota/quota-fetcher-gemini-cli.ts +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli.ts @@ -1,1179 +1,27 @@ /** - * Quota Fetcher for Gemini CLI Accounts + * Quota Fetcher for Gemini CLI Accounts (barrel) * * Fetches quota information from Google Cloud Code internal API. * Used for displaying bucket-based quotas grouped by model series. - */ - -import * as fs from 'node:fs'; -import * as path from 'node:path'; -import { getAuthDir } from '../config/config-generator'; -import { getProviderAccounts, getPausedDir, setAccountTier } from '../accounts/account-manager'; -import { getTokenExpiryTimestamp, sanitizeEmail, isTokenExpired } from '../auth/auth-utils'; -import { - buildGeminiCliBucketsFromParsedBuckets, - type GeminiCliParsedBucket, -} from './gemini-cli-quota-normalizer'; -import { mapExternalProviderName } from '../provider-capabilities'; -import { - buildManagementHeaders, - buildProxyUrl, - getProxyTarget, -} from '../proxy/proxy-target-resolver'; -import type { GeminiCliQuotaResult, GeminiCliBucket } from './quota-types'; -import { - buildProviderEntitlementEvidence, - getProviderTierLabel, - isModelCapacityExhausted, - normalizeProviderTierId, -} from '../auth/provider-entitlement-evidence'; -import type { ProviderEntitlementEvidence } from '../auth/provider-entitlement-types'; -import { createLogger } from '../../services/logging'; - -// Diagnostic-only logger: quota fetch progress, upstream HTTP status, and -// recovery hints. accountId is attached as provider context; token values -// are never logged (they live in auth files and are not read into messages). -const logger = createLogger('cliproxy:quota:gemini-cli'); - -/** Google Cloud Code API endpoints */ -const GEMINI_CLI_API_BASE = 'https://cloudcode-pa.googleapis.com'; -const GEMINI_CLI_API_VERSION = 'v1internal'; -const GEMINI_CLI_QUOTA_URL = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:retrieveUserQuota`; -const GEMINI_CLI_CODE_ASSIST_URL = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:loadCodeAssist`; -const GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH = 320; -const GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX = '...[truncated]'; -const GEMINI_CLI_G1_CREDIT_TYPE = 'GOOGLE_ONE_AI'; -const MANAGEMENT_API_TIMEOUT_MS = 5000; -const SECONDARY_REQUEST_TIMEOUT_MS = 2000; - -/** Auth data extracted from Gemini CLI auth file */ -interface GeminiCliAuthData { - accessToken: string; - projectId: string | null; - isExpired: boolean; - expiresAt: string | number | null; -} - -/** Raw bucket from API response */ -interface RawGeminiCliBucket { - model_id?: string; - modelId?: string; - token_type?: string | null; - tokenType?: string | null; - remaining_fraction?: number; - remainingFraction?: number; - remaining_amount?: number; - remainingAmount?: number; - reset_time?: string | null; - resetTime?: string | null; -} - -/** Raw API response structure */ -interface GeminiCliQuotaResponse { - buckets?: RawGeminiCliBucket[]; -} - -interface GeminiCliCredits { - creditType?: string; - credit_type?: string; - creditAmount?: string | number; - credit_amount?: string | number; -} - -interface GeminiCliUserTier { - id?: string; - availableCredits?: GeminiCliCredits[]; - available_credits?: GeminiCliCredits[]; -} - -interface GeminiCliCodeAssistResponse { - currentTier?: GeminiCliUserTier | null; - current_tier?: GeminiCliUserTier | null; - paidTier?: GeminiCliUserTier | null; - paid_tier?: GeminiCliUserTier | null; -} - -interface ParsedGeminiCliErrorBody { - errorCode?: string; - errorDetail?: string; - message?: string; -} - -interface GeminiCliSupplementaryInfo { - tierLabel: string | null; - tierId: string | null; - creditBalance: number | null; - normalizedTier: 'free' | 'pro' | 'ultra' | 'unknown'; -} - -interface ManagementAuthFile { - auth_index?: string | number; - provider?: string; - type?: string; - email?: string; - name?: string; -} - -interface ManagementApiCallResponse { - status_code?: number; - body?: string; -} - -interface ManagedResponse { - status: number; - bodyText: string; - json: unknown; - viaManagement: boolean; -} - -interface ManagedGeminiAuthContext { - authIndexLookupPromise?: Promise; -} - -interface ManagedGeminiAuthLookupResult { - authIndex: string | number | null; - unavailable: boolean; -} - -interface ManagedGeminiRequestResult { - response: ManagedResponse | null; - unavailable: boolean; -} - -function getRemainingTimeoutMs(deadlineMs: number): number { - return Math.max(1, deadlineMs - Date.now()); -} - -/** - * Extract project ID from account field - * Input: "user@example.com (cloudaicompanion-abc-123)" - * Output: "cloudaicompanion-abc-123" - */ -function resolveGeminiCliProjectId(accountField: string): string | null { - const regex = /\(([^()]+)\)/g; - let match: RegExpExecArray | null; - let lastMatch: string | null = null; - while ((match = regex.exec(accountField)) !== null) { - lastMatch = match[1]; - } - return lastMatch; -} - -/** - * Extract access token from Gemini auth file data - * Handles both flat (access_token) and nested (token.access_token) structures - */ -function extractAccessToken(data: Record): string | null { - // Flat structure: { access_token: "..." } - if (typeof data.access_token === 'string') { - return data.access_token; - } - // Nested structure: { token: { access_token: "..." } } - if (data.token && typeof data.token === 'object') { - const token = data.token as Record; - if (typeof token.access_token === 'string') { - return token.access_token; - } - } - return null; -} - -/** - * Extract expiry from Gemini auth file data - * Handles both flat (expired) and nested (token.expiry) structures - */ -function extractExpiry(data: Record): string | number | null { - // Flat structure: { expired: "..." } - if (typeof data.expired === 'string') { - return data.expired; - } - if (typeof data.expired === 'number') { - return data.expired; - } - // Nested structure: { token: { expiry: "..." } } - if (data.token && typeof data.token === 'object') { - const token = data.token as Record; - if (typeof token.expiry === 'string') { - return token.expiry; - } - if (typeof token.expiry === 'number') { - return token.expiry; - } - } - return null; -} - -/** - * Check if file matches Gemini CLI auth file patterns - * Patterns: gemini-*.json OR *-gen-lang-client-*.json OR email@domain.com-*.json with type=gemini - */ -function isGeminiAuthFile(filename: string): boolean { - if (!filename.endsWith('.json')) return false; - // Legacy pattern: gemini-email.json - if (filename.startsWith('gemini-')) return true; - // New pattern: email-gen-lang-client-projectId.json - if (filename.includes('-gen-lang-client-')) return true; - // Check if contains @ (email pattern) - will verify type inside - if (filename.includes('@')) return true; - return false; -} - -function safeParseJson(bodyText: string): unknown { - try { - return JSON.parse(bodyText); - } catch { - return null; - } -} - -async function readManagedResponse( - response: Response, - viaManagement: boolean -): Promise { - const bodyText = await response.text(); - return { - status: response.status, - bodyText, - json: safeParseJson(bodyText), - viaManagement, - }; -} - -function isGeminiAuthFileForAccount(file: ManagementAuthFile, accountId: string): boolean { - const rawProvider = normalizeStringValue(file.provider ?? file.type); - if (!rawProvider || mapExternalProviderName(rawProvider) !== 'gemini') { - return false; - } - - const email = normalizeStringValue(file.email); - const normalizedAccountId = accountId.trim().toLowerCase(); - if (email?.toLowerCase() === normalizedAccountId) { - return true; - } - - const normalizedName = normalizeStringValue(file.name); - if (!normalizedName) { - return false; - } - - const normalizedFileName = normalizedName.toLowerCase(); - const sanitizedAccount = sanitizeEmail(accountId).toLowerCase(); - return ( - normalizedFileName === `gemini-${sanitizedAccount}.json` || - normalizedFileName.startsWith(`${normalizedAccountId}-gen-lang-client-`) || - normalizedFileName.includes(sanitizedAccount) - ); -} - -async function findManagedGeminiAuthIndex( - accountId: string, - timeoutMs: number -): Promise { - const target = getProxyTarget(); - const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), timeoutMs); - - try { - const response = await fetch(buildProxyUrl(target, '/v0/management/auth-files'), { - signal: controller.signal, - headers: buildManagementHeaders(target), - }); - clearTimeout(timeoutId); - - if (!response.ok) { - return { authIndex: null, unavailable: true }; - } - - const data = (await response.json()) as { files?: ManagementAuthFile[] }; - const match = data.files?.find((file) => isGeminiAuthFileForAccount(file, accountId)); - return { authIndex: match?.auth_index ?? null, unavailable: false }; - } catch { - clearTimeout(timeoutId); - return { authIndex: null, unavailable: true }; - } -} - -async function getManagedGeminiAuthIndex( - accountId: string, - timeoutMs: number, - context?: ManagedGeminiAuthContext -): Promise { - if (!context) { - return await findManagedGeminiAuthIndex(accountId, timeoutMs); - } - - context.authIndexLookupPromise ??= findManagedGeminiAuthIndex(accountId, timeoutMs); - return await context.authIndexLookupPromise; -} - -class GeminiManagedAuthUnavailableError extends Error { - constructor() { - super('CLIProxy managed Gemini auth is temporarily unavailable'); - this.name = 'GeminiManagedAuthUnavailableError'; - } -} - -async function performManagedGeminiRequest( - accountId: string, - url: string, - body: string, - timeoutMs: number, - authContext?: ManagedGeminiAuthContext -): Promise { - const deadlineMs = Date.now() + timeoutMs; - const lookupResult = await getManagedGeminiAuthIndex( - accountId, - getRemainingTimeoutMs(deadlineMs), - authContext - ); - if (lookupResult.unavailable) { - return { response: null, unavailable: true }; - } - - const authIndex = lookupResult.authIndex; - if (authIndex === null || authIndex === undefined) { - return { response: null, unavailable: false }; - } - - const target = getProxyTarget(); - const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), getRemainingTimeoutMs(deadlineMs)); - - try { - const response = await fetch(buildProxyUrl(target, '/v0/management/api-call'), { - method: 'POST', - signal: controller.signal, - headers: buildManagementHeaders(target, { - 'Content-Type': 'application/json', - }), - body: JSON.stringify({ - auth_index: authIndex, - method: 'POST', - url, - header: { - Authorization: 'Bearer $TOKEN$', - 'Content-Type': 'application/json', - }, - data: body, - }), - }); - clearTimeout(timeoutId); - - if (!response.ok) { - return { response: null, unavailable: true }; - } - - const apiResponse = (await response.json()) as ManagementApiCallResponse; - const bodyText = typeof apiResponse.body === 'string' ? apiResponse.body : ''; - return { - response: { - status: typeof apiResponse.status_code === 'number' ? apiResponse.status_code : 500, - bodyText, - json: safeParseJson(bodyText), - viaManagement: true, - }, - unavailable: false, - }; - } catch { - clearTimeout(timeoutId); - return { response: null, unavailable: true }; - } -} - -async function performGeminiCliRequest( - accountId: string, - accessToken: string, - url: string, - body: string, - preferManagement = false, - authContext?: ManagedGeminiAuthContext -): Promise { - let managementAttempted = false; - let managementUnavailable = false; - - if (preferManagement) { - managementAttempted = true; - const managedResult = await performManagedGeminiRequest( - accountId, - url, - body, - MANAGEMENT_API_TIMEOUT_MS, - authContext - ); - managementUnavailable = managedResult.unavailable; - if (managedResult.response) { - return managedResult.response; - } - } - - const controller = new AbortController(); - const timeoutId = setTimeout( - () => controller.abort(), - managementAttempted ? SECONDARY_REQUEST_TIMEOUT_MS : MANAGEMENT_API_TIMEOUT_MS - ); - - try { - const response = await fetch(url, { - method: 'POST', - signal: controller.signal, - headers: { - Authorization: `Bearer ${accessToken}`, - 'Content-Type': 'application/json', - }, - body, - }); - clearTimeout(timeoutId); - - const directResult = await readManagedResponse(response, false); - if (directResult.status !== 401) { - return directResult; - } - - if (managementAttempted) { - if (managementUnavailable) { - throw new GeminiManagedAuthUnavailableError(); - } - return directResult; - } - - const managedResult = await performManagedGeminiRequest( - accountId, - url, - body, - SECONDARY_REQUEST_TIMEOUT_MS, - authContext - ); - if (managedResult.response) { - return managedResult.response; - } - if (managedResult.unavailable) { - throw new GeminiManagedAuthUnavailableError(); - } - return directResult; - } catch (error) { - clearTimeout(timeoutId); - throw error; - } -} - -/** - * Read auth data from Gemini CLI auth file - * Supports multiple file naming conventions and JSON structures - */ -function readGeminiCliAuthData(accountId: string): GeminiCliAuthData | null { - const authDirs = [getAuthDir(), getPausedDir()]; - const sanitizedId = sanitizeEmail(accountId); - const expectedFiles = [ - `gemini-${sanitizedId}.json`, // Legacy format - `${accountId}-gen-lang-client-`, // New format prefix (partial match) - ]; - - for (const authDir of authDirs) { - if (!fs.existsSync(authDir)) continue; - - // Try exact legacy match first - const legacyPath = path.join(authDir, expectedFiles[0]); - if (fs.existsSync(legacyPath)) { - try { - const content = fs.readFileSync(legacyPath, 'utf-8'); - const data = JSON.parse(content) as Record; - const accessToken = extractAccessToken(data); - if (accessToken) { - const projectId = - typeof data.project_id === 'string' - ? data.project_id - : resolveGeminiCliProjectId(String(data.account || '')); - const expiry = extractExpiry(data); - - return { - accessToken, - projectId, - isExpired: isTokenExpired(expiry ?? undefined), - expiresAt: expiry, - }; - } - } catch { - // Continue to fallback - } - } - - // Scan directory for matching files - const files = fs.readdirSync(authDir); - for (const file of files) { - if (!isGeminiAuthFile(file)) continue; - - const candidatePath = path.join(authDir, file); - try { - const content = fs.readFileSync(candidatePath, 'utf-8'); - const data = JSON.parse(content) as Record; - - // Check if this file matches our account - const fileEmail = typeof data.email === 'string' ? data.email : null; - const fileType = typeof data.type === 'string' ? data.type : null; - const matchesEmail = fileEmail === accountId; - const matchesFilename = file.startsWith(`${accountId}-`) || file.includes(sanitizedId); - const isGeminiType = fileType === 'gemini' || fileType === 'gemini-cli'; - - // Must match account AND be gemini type (or legacy gemini- prefix) - if ((matchesEmail || matchesFilename) && (isGeminiType || file.startsWith('gemini-'))) { - const accessToken = extractAccessToken(data); - if (accessToken) { - const projectId = - typeof data.project_id === 'string' - ? data.project_id - : resolveGeminiCliProjectId(String(data.account || '')); - const expiry = extractExpiry(data); - - return { - accessToken, - projectId, - isExpired: isTokenExpired(expiry ?? undefined), - expiresAt: expiry, - }; - } - } - } catch { - continue; - } - } - } - - return null; -} - -function normalizeStringValue(value: unknown): string | null { - return typeof value === 'string' && value.trim().length > 0 ? value.trim() : null; -} - -function normalizeNumberValue(value: unknown): number | null { - if (typeof value === 'number' && Number.isFinite(value)) { - return value; - } - if (typeof value === 'string' && value.trim().length > 0) { - const parsed = Number(value); - if (Number.isFinite(parsed)) { - return parsed; - } - } - return null; -} - -function resolveGeminiCliTierId(payload: GeminiCliCodeAssistResponse | null): string | null { - if (!payload) return null; - const currentTier = payload.currentTier ?? payload.current_tier; - const paidTier = payload.paidTier ?? payload.paid_tier; - const rawId = normalizeStringValue(paidTier?.id) ?? normalizeStringValue(currentTier?.id); - return rawId ? rawId.toLowerCase() : null; -} - -function resolveGeminiCliTierLabel(payload: GeminiCliCodeAssistResponse | null): string | null { - const tierId = resolveGeminiCliTierId(payload); - return getProviderTierLabel(tierId); -} - -function resolveGeminiCliCreditBalance(payload: GeminiCliCodeAssistResponse | null): number | null { - if (!payload) return null; - - const paidTier = payload.paidTier ?? payload.paid_tier; - const currentTier = payload.currentTier ?? payload.current_tier; - const tier = paidTier ?? currentTier; - if (!tier) return null; - - const credits = tier.availableCredits ?? tier.available_credits ?? []; - let total = 0; - let found = false; - for (const credit of credits) { - const creditType = normalizeStringValue(credit.creditType ?? credit.credit_type); - if (creditType !== GEMINI_CLI_G1_CREDIT_TYPE) continue; - - const amount = normalizeNumberValue(credit.creditAmount ?? credit.credit_amount); - if (amount !== null) { - total += amount; - found = true; - } - } - - return found ? total : null; -} - -async function fetchGeminiCliSupplementary( - accountId: string, - accessToken: string, - projectId: string, - verbose: boolean, - authContext?: ManagedGeminiAuthContext -): Promise { - const requestBody = JSON.stringify({ - cloudaicompanionProject: projectId, - metadata: { - ideType: 'IDE_UNSPECIFIED', - platform: 'PLATFORM_UNSPECIFIED', - pluginType: 'GEMINI', - duetProject: projectId, - }, - }); - - try { - const response = await performGeminiCliRequest( - accountId, - accessToken, - GEMINI_CLI_CODE_ASSIST_URL, - requestBody, - false, - authContext - ); - - if (response.status !== 200) { - if (verbose) { - const source = response.viaManagement ? 'managed' : 'direct'; - logger.info( - 'gemini_cli.supplementary_metadata_unavailable', - `Gemini CLI supplementary metadata unavailable via ${source}: HTTP ${response.status}`, - { provider: 'gemini', accountId, httpStatus: response.status, source } - ); - } - return { tierLabel: null, tierId: null, creditBalance: null, normalizedTier: 'unknown' }; - } - - const payload = response.json as GeminiCliCodeAssistResponse | null; - return { - tierLabel: resolveGeminiCliTierLabel(payload), - tierId: resolveGeminiCliTierId(payload), - creditBalance: resolveGeminiCliCreditBalance(payload), - normalizedTier: normalizeProviderTierId(resolveGeminiCliTierId(payload)), - }; - } catch (error) { - if (verbose) { - const message = error instanceof Error ? error.message : 'Unknown error'; - logger.info( - 'gemini_cli.supplementary_metadata_skipped', - `Gemini CLI supplementary metadata skipped: ${message}`, - { - provider: 'gemini', - accountId, - err: error instanceof Error ? { name: error.name, message } : { message }, - } - ); - } - return { tierLabel: null, tierId: null, creditBalance: null, normalizedTier: 'unknown' }; - } -} - -function buildGeminiCliFailureResult( - accountId: string, - projectId: string | null, - options: { - error: string; - httpStatus?: number; - errorCode?: string; - errorDetail?: string; - actionHint?: string; - retryable?: boolean; - needsReauth?: boolean; - isForbidden?: boolean; - entitlement?: ProviderEntitlementEvidence; - } -): GeminiCliQuotaResult { - return { - success: false, - buckets: [], - projectId, - tierLabel: null, - tierId: null, - creditBalance: null, - lastUpdated: Date.now(), - accountId, - error: options.error, - httpStatus: options.httpStatus, - errorCode: options.errorCode, - errorDetail: options.errorDetail, - actionHint: options.actionHint, - retryable: options.retryable, - needsReauth: options.needsReauth, - isForbidden: options.isForbidden, - entitlement: options.entitlement, - }; -} - -function sanitizeGeminiCliErrorDetail(bodyText: string): string | undefined { - const trimmed = bodyText.trim(); - if (!trimmed) { - return undefined; - } - - if (/^]+>/.test(trimmed)) { - return '[HTML error response omitted]'; - } - - let sanitized = trimmed - .replace( - /"(access[_-]?token|refresh[_-]?token|authorization|cookie|set-cookie|api[_-]?key|session[_-]?token|token)"\s*:\s*"[^"]*"/gi, - '"$1":"[redacted]"' - ) - .replace(/Bearer\s+[A-Za-z0-9._-]+/g, 'Bearer [redacted]') - .replace(/\s+/g, ' '); - - if (sanitized.length > GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH) { - sanitized = `${sanitized.slice( - 0, - GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH - GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX.length - )}${GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX}`; - } - - return sanitized; -} - -function extractGeminiCliNestedMessage(value: unknown): string | undefined { - if (Array.isArray(value)) { - for (const entry of value) { - const nested = extractGeminiCliNestedMessage(entry); - if (nested) return nested; - } - return undefined; - } - - if (!value || typeof value !== 'object') { - return undefined; - } - - const record = value as Record; - const directMessage = [ - record.message, - record.localizedMessage, - record.description, - record.reason, - record.error, - ].find( - (candidate): candidate is string => typeof candidate === 'string' && candidate.trim().length > 0 - ); - if (directMessage) { - return directMessage; - } - - return undefined; -} - -function parseGeminiCliErrorBody(bodyText: string): ParsedGeminiCliErrorBody { - const trimmed = bodyText.trim(); - if (!trimmed) { - return {}; - } - - const sanitizedDetail = sanitizeGeminiCliErrorDetail(trimmed); - - try { - const parsed = JSON.parse(trimmed) as Record; - const topLevelMessage = [parsed.message, parsed.error].find( - (candidate): candidate is string => - typeof candidate === 'string' && candidate.trim().length > 0 - ); - const topLevelCode = [parsed.code, parsed.status].find( - (candidate): candidate is string => - typeof candidate === 'string' && candidate.trim().length > 0 - ); - - if (parsed.error && typeof parsed.error === 'object') { - const error = parsed.error as Record; - return { - errorCode: - [error.status, error.code, topLevelCode].find( - (candidate): candidate is string => - typeof candidate === 'string' && candidate.trim().length > 0 - ) || undefined, - errorDetail: sanitizedDetail, - message: - [ - error.message, - error.error, - extractGeminiCliNestedMessage(error.details), - topLevelMessage, - ].find( - (candidate): candidate is string => - typeof candidate === 'string' && candidate.trim().length > 0 - ) || undefined, - }; - } - - return { - errorCode: topLevelCode, - errorDetail: sanitizedDetail, - message: - [topLevelMessage, extractGeminiCliNestedMessage(parsed.details)].find( - (candidate): candidate is string => - typeof candidate === 'string' && candidate.trim().length > 0 - ) || undefined, - }; - } catch { - return { - errorDetail: sanitizedDetail, - message: sanitizedDetail === '[HTML error response omitted]' ? undefined : trimmed, - }; - } -} - -function buildGeminiCliForbiddenActionHint(parsed: ParsedGeminiCliErrorBody): string { - const combined = `${parsed.message || ''} ${parsed.errorDetail || ''}`.toLowerCase(); - if (combined.includes('verify') || combined.includes('verification')) { - return 'Complete the Google account verification mentioned above, then retry quota refresh.'; - } - if (combined.includes('project')) { - return 'Confirm this Google project still has Gemini CLI quota access, then retry.'; - } - return 'Check the Google account or workspace access shown above, then retry quota refresh.'; -} - -function buildGeminiCliHttpFailureResult( - accountId: string, - projectId: string | null, - status: number, - bodyText: string -): GeminiCliQuotaResult { - const parsed = parseGeminiCliErrorBody(bodyText); - - if (status === 401) { - return buildGeminiCliFailureResult(accountId, projectId, { - error: parsed.message || 'Token expired or invalid', - httpStatus: 401, - errorCode: parsed.errorCode || 'reauth_required', - errorDetail: parsed.errorDetail, - actionHint: 'Run ccs gemini --auth to reconnect this account.', - needsReauth: true, - retryable: false, - }); - } - - if (status === 403) { - return buildGeminiCliFailureResult(accountId, projectId, { - error: parsed.message || 'Quota access forbidden for this account', - httpStatus: 403, - errorCode: parsed.errorCode || 'quota_api_forbidden', - errorDetail: parsed.errorDetail, - actionHint: buildGeminiCliForbiddenActionHint(parsed), - isForbidden: true, - retryable: false, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'medium', - accessState: 'not_entitled', - capacityState: 'unknown', - }), - }); - } - - if (status === 429) { - if (isModelCapacityExhausted(parsed.message, parsed.errorDetail, parsed.errorCode)) { - return buildGeminiCliFailureResult(accountId, projectId, { - error: parsed.message || 'Model capacity exhausted for this account right now', - httpStatus: 429, - errorCode: 'capacity_exhausted', - errorDetail: parsed.errorDetail, - actionHint: - 'Retry later or switch to another Gemini model. This indicates temporary model capacity, not an authentication failure.', - retryable: true, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'medium', - accessState: 'entitled', - capacityState: 'capacity_exhausted', - notes: 'Upstream returned MODEL_CAPACITY_EXHAUSTED for this model.', - }), - }); - } - - return buildGeminiCliFailureResult(accountId, projectId, { - error: parsed.message || 'Rate limited - try again later', - httpStatus: 429, - errorCode: parsed.errorCode || 'rate_limited', - errorDetail: parsed.errorDetail, - actionHint: 'Retry after a short delay.', - retryable: true, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'low', - accessState: 'unknown', - capacityState: 'rate_limited', - }), - }); - } - - if (status >= 500) { - return buildGeminiCliFailureResult(accountId, projectId, { - error: parsed.message || `Gemini quota service unavailable (HTTP ${status})`, - httpStatus: status, - errorCode: parsed.errorCode || 'provider_unavailable', - errorDetail: parsed.errorDetail, - actionHint: 'Retry later. This looks like a temporary Google upstream problem.', - retryable: true, - }); - } - - return buildGeminiCliFailureResult(accountId, projectId, { - error: parsed.message || `Gemini quota request failed (HTTP ${status})`, - httpStatus: status, - errorCode: parsed.errorCode || 'quota_request_failed', - errorDetail: parsed.errorDetail, - actionHint: 'Inspect the upstream response details and retry if appropriate.', - retryable: false, - }); -} - -/** - * Build GeminiCliBucket array from API response - * Groups buckets by model series and token type - */ -function buildGeminiCliBuckets(rawBuckets: RawGeminiCliBucket[]): GeminiCliBucket[] { - const parsedBuckets = rawBuckets - .map((bucket): GeminiCliParsedBucket | null => { - const modelId = normalizeStringValue(bucket.model_id ?? bucket.modelId); - if (!modelId) return null; - - const tokenType = normalizeStringValue(bucket.token_type ?? bucket.tokenType); - const remainingFractionRaw = normalizeNumberValue( - bucket.remaining_fraction ?? bucket.remainingFraction - ); - const remainingAmount = normalizeNumberValue( - bucket.remaining_amount ?? bucket.remainingAmount - ); - const resetTime = normalizeStringValue(bucket.reset_time ?? bucket.resetTime); - - let fallbackFraction: number | null = null; - if (remainingAmount !== null) { - fallbackFraction = remainingAmount <= 0 ? 0 : null; - } else if (resetTime) { - fallbackFraction = 0; - } - - return { - modelId, - tokenType, - remainingFraction: remainingFractionRaw ?? fallbackFraction ?? 1, - remainingAmount, - resetTime, - }; - }) - .filter((bucket): bucket is GeminiCliParsedBucket => bucket !== null); - - return buildGeminiCliBucketsFromParsedBuckets(parsedBuckets); -} - -/** - * Internal helper: Fetch quota with validated auth data - * Extracted to support auto-refresh retry logic - */ -async function fetchWithAuthData( - authData: GeminiCliAuthData, - accountId: string, - verbose: boolean -): Promise { - if (!authData.projectId) { - const error = 'Cannot resolve project ID from auth file'; - if (verbose) { - logger.error('gemini_cli.missing_project_id', `Error: ${error}`, { - provider: 'gemini', - accountId, - }); - } - return buildGeminiCliFailureResult(accountId, null, { - error, - errorCode: 'missing_project_id', - actionHint: 'Run ccs gemini --auth to reconnect this account and recover the project ID.', - retryable: false, - }); - } - - const authContext: ManagedGeminiAuthContext = {}; - const supplementaryPromise = fetchGeminiCliSupplementary( - accountId, - authData.accessToken, - authData.projectId, - verbose, - authContext - ); - const requestBody = JSON.stringify({ project: authData.projectId }); - - try { - const response = await performGeminiCliRequest( - accountId, - authData.accessToken, - GEMINI_CLI_QUOTA_URL, - requestBody, - authData.isExpired, - authContext - ); - - if (verbose) { - const source = response.viaManagement ? 'managed' : 'direct'; - logger.info( - 'gemini_cli.api_status', - `Gemini CLI API status via ${source}: ${response.status}`, - { provider: 'gemini', accountId, httpStatus: response.status, source } - ); - } - - if (response.status !== 200) { - return buildGeminiCliHttpFailureResult( - accountId, - authData.projectId, - response.status, - response.bodyText - ); - } - - const data = response.json as GeminiCliQuotaResponse | null; - const rawBuckets = data?.buckets || []; - const buckets = buildGeminiCliBuckets(rawBuckets); - const supplementary = await supplementaryPromise; - - if (verbose) { - logger.info('gemini_cli.buckets_found', `Gemini CLI buckets found: ${buckets.length}`, { - provider: 'gemini', - accountId, - bucketCount: buckets.length, - }); - } - - if (supplementary.normalizedTier !== 'unknown') { - setAccountTier('gemini', accountId, supplementary.normalizedTier); - } - - return { - success: true, - buckets, - projectId: authData.projectId, - tierLabel: supplementary.tierLabel, - tierId: supplementary.tierId, - creditBalance: supplementary.creditBalance, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: supplementary.normalizedTier, - rawTierId: supplementary.tierId, - rawTierLabel: supplementary.tierLabel, - source: supplementary.tierId ? 'runtime_api' : 'runtime_inference', - confidence: supplementary.tierId ? 'high' : 'medium', - accessState: 'entitled', - capacityState: 'available', - }), - lastUpdated: Date.now(), - accountId, - }; - } catch (err) { - if (err instanceof GeminiManagedAuthUnavailableError) { - return buildGeminiCliFailureResult(accountId, authData.projectId, { - error: 'Gemini delegated auth refresh is temporarily unavailable', - errorCode: 'managed_auth_unavailable', - errorDetail: err.message, - actionHint: 'Retry later. CLIProxy management could not refresh this Gemini account.', - retryable: true, - }); - } - - const errorMsg = - err instanceof Error && err.name === 'AbortError' - ? 'Request timeout' - : err instanceof Error - ? err.message - : 'Unknown error'; - - if (verbose) { - logger.error('gemini_cli.quota_fetch_error', `Gemini CLI quota error: ${errorMsg}`, { - provider: 'gemini', - accountId, - err: err instanceof Error ? { name: err.name, message: errorMsg } : { message: errorMsg }, - }); - } - - return buildGeminiCliFailureResult(accountId, authData.projectId, { - error: errorMsg, - errorCode: - err instanceof Error && err.name === 'AbortError' ? 'network_timeout' : 'network_error', - actionHint: 'Retry later. This looks temporary.', - retryable: true, - httpStatus: err instanceof Error && err.name === 'AbortError' ? 408 : undefined, - }); - } -} - -/** - * Fetch quota for a single Gemini CLI account * - * @param accountId - Account identifier (email) - * @param verbose - Show detailed diagnostics - * @returns Quota result with buckets and percentages - */ -export async function fetchGeminiCliQuota( - accountId: string, - verbose = false -): Promise { - if (verbose) { - logger.info('gemini_cli.fetch_start', `Fetching Gemini CLI quota for ${accountId}...`, { - provider: 'gemini', - accountId, - }); - } - - const authData = readGeminiCliAuthData(accountId); - if (!authData) { - const error = 'Auth file not found for Gemini account'; - if (verbose) { - logger.error('gemini_cli.auth_file_missing', `Error: ${error}`, { - provider: 'gemini', - accountId, - }); - } - return buildGeminiCliFailureResult(accountId, null, { - error, - errorCode: 'auth_file_missing', - actionHint: 'Run ccs gemini --auth to reconnect this account.', - retryable: false, - }); - } - - if (authData.isExpired && verbose) { - const expiresAt = getTokenExpiryTimestamp(authData.expiresAt); - const expiryLabel = expiresAt ? new Date(expiresAt).toISOString() : 'unknown'; - logger.info( - 'gemini_cli.token_expired', - `Gemini access token is expired (${expiryLabel}); quota requests will defer to managed auth when available.`, - { provider: 'gemini', accountId, tokenExpired: true, expiresAt: expiryLabel } - ); - } - - return await fetchWithAuthData(authData, accountId, verbose); -} - -/** - * Fetch quota for all Gemini CLI accounts + * This file is a thin re-export barrel. The implementation has been split + * into focused submodules under `./quota-fetcher-gemini-cli/`: + * - token-parsing: access token / expiry / project-id extraction + * - auth-file-discovery: on-disk auth-file lookup and parsing + * - managed-request: managed (CLIProxy) + direct upstream HTTP paths + * - supplementary-metadata: tier / credit resolution (loadCodeAssist) + * - error-parsing: error-body sanitization and failure-result builders + * - bucket-building: raw bucket -> normalized GeminiCliBucket array + * - quota-fetcher: top-level fetchGeminiCliQuota orchestration * - * @param verbose - Show detailed diagnostics - * @returns Array of account quotas + * Public surface is preserved exactly: every previously-importable symbol + * keeps its name and signature. Token values are never logged. */ -export async function fetchAllGeminiCliQuotas( - verbose = false -): Promise<{ account: string; quota: GeminiCliQuotaResult }[]> { - const accounts = getProviderAccounts('gemini'); - if (accounts.length === 0) { - return []; - } - - const results = await Promise.all( - accounts.map(async (account) => ({ - account: account.id, - quota: await fetchGeminiCliQuota(account.id, verbose), - })) - ); - - return results; -} - -export const __testExports = { - sanitizeGeminiCliErrorDetail, - extractGeminiCliNestedMessage, - parseGeminiCliErrorBody, - buildGeminiCliForbiddenActionHint, -}; - -// Export for testing -export { resolveGeminiCliProjectId, buildGeminiCliBuckets }; +export { + fetchGeminiCliQuota, + fetchAllGeminiCliQuotas, + resolveGeminiCliProjectId, + buildGeminiCliBuckets, + __testExports, +} from './quota-fetcher-gemini-cli/index'; diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/auth-file-discovery.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/auth-file-discovery.ts new file mode 100644 index 00000000..dfabdc9f --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/auth-file-discovery.ts @@ -0,0 +1,114 @@ +/** + * Auth file discovery for the Gemini CLI quota fetcher (direct-path credentials). + * + * Locates and parses the on-disk Gemini CLI auth file for a given account, + * supporting both the legacy `gemini-.json` filename and the newer + * `-gen-lang-client-.json` pattern. Scans both the active + * auth directory and the paused-account directory. + * + * Returns the access token, project ID, expiry, and expired flag. The live + * token is returned only so the caller can place it in an Authorization header; + * it is never logged by this module or its callers. + */ + +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import { getAuthDir } from '../../config/config-generator'; +import { getPausedDir } from '../../accounts/account-manager'; +import { isTokenExpired } from '../../auth/auth-utils'; +import { sanitizeEmail } from '../../auth/auth-utils'; +import { isGeminiAuthFile } from './managed-request'; +import { extractAccessToken, extractExpiry, resolveGeminiCliProjectId } from './token-parsing'; +import type { GeminiCliAuthData } from './types'; + +/** + * Read auth data from a Gemini CLI auth file on disk. + * + * Resolution order per auth directory: + * 1. Exact legacy match: `gemini-.json` + * 2. Directory scan for files matching {@link isGeminiAuthFile}, filtered + * by account email/filename and Gemini type. + * + * Scans both the active auth dir and the paused-account dir. Returns null if + * no usable auth file (with an access token) is found. + */ +export function readGeminiCliAuthData(accountId: string): GeminiCliAuthData | null { + const authDirs = [getAuthDir(), getPausedDir()]; + const sanitizedId = sanitizeEmail(accountId); + const expectedFiles = [ + `gemini-${sanitizedId}.json`, // Legacy format + `${accountId}-gen-lang-client-`, // New format prefix (partial match) + ]; + + for (const authDir of authDirs) { + if (!fs.existsSync(authDir)) continue; + + // Try exact legacy match first + const legacyPath = path.join(authDir, expectedFiles[0]); + if (fs.existsSync(legacyPath)) { + try { + const content = fs.readFileSync(legacyPath, 'utf-8'); + const data = JSON.parse(content) as Record; + const accessToken = extractAccessToken(data); + if (accessToken) { + const projectId = + typeof data.project_id === 'string' + ? data.project_id + : resolveGeminiCliProjectId(String(data.account || '')); + const expiry = extractExpiry(data); + + return { + accessToken, + projectId, + isExpired: isTokenExpired(expiry ?? undefined), + expiresAt: expiry, + }; + } + } catch { + // Continue to fallback + } + } + + // Scan directory for matching files + const files = fs.readdirSync(authDir); + for (const file of files) { + if (!isGeminiAuthFile(file)) continue; + + const candidatePath = path.join(authDir, file); + try { + const content = fs.readFileSync(candidatePath, 'utf-8'); + const data = JSON.parse(content) as Record; + + // Check if this file matches our account + const fileEmail = typeof data.email === 'string' ? data.email : null; + const fileType = typeof data.type === 'string' ? data.type : null; + const matchesEmail = fileEmail === accountId; + const matchesFilename = file.startsWith(`${accountId}-`) || file.includes(sanitizedId); + const isGeminiType = fileType === 'gemini' || fileType === 'gemini-cli'; + + // Must match account AND be gemini type (or legacy gemini- prefix) + if ((matchesEmail || matchesFilename) && (isGeminiType || file.startsWith('gemini-'))) { + const accessToken = extractAccessToken(data); + if (accessToken) { + const projectId = + typeof data.project_id === 'string' + ? data.project_id + : resolveGeminiCliProjectId(String(data.account || '')); + const expiry = extractExpiry(data); + + return { + accessToken, + projectId, + isExpired: isTokenExpired(expiry ?? undefined), + expiresAt: expiry, + }; + } + } + } catch { + continue; + } + } + } + + return null; +} diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/bucket-building.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/bucket-building.ts new file mode 100644 index 00000000..5d5f5991 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/bucket-building.ts @@ -0,0 +1,62 @@ +/** + * Bucket building for the Gemini CLI quota fetcher. + * + * Translates raw upstream quota buckets (snake_case and camelCase tolerant) + * into the normalized {@link GeminiCliBucket} array grouped by model series + * and token type. Delegates the grouping to the shared + * `gemini-cli-quota-normalizer` so the grouping rules stay in one place. + */ + +import { + buildGeminiCliBucketsFromParsedBuckets, + type GeminiCliParsedBucket, +} from '../gemini-cli-quota-normalizer'; +import type { GeminiCliBucket } from '../quota-types'; +import type { RawGeminiCliBucket } from './types'; +import { normalizeNumberValue, normalizeStringValue } from './shared-utils'; + +/** + * Build a {@link GeminiCliBucket} array from raw upstream quota buckets. + * + * Each raw bucket is normalized into a {@link GeminiCliParsedBucket}: + * - skips buckets with no resolvable model id + * - coalesces remaining_fraction / remaining_amount / reset_time across + * naming variants, with a fallback of `1` (full) when none are present + * but a reset time or non-positive amount implies exhaustion + * Then delegates to {@link buildGeminiCliBucketsFromParsedBuckets} for the + * model-series and token-type grouping. + */ +export function buildGeminiCliBuckets(rawBuckets: RawGeminiCliBucket[]): GeminiCliBucket[] { + const parsedBuckets = rawBuckets + .map((bucket): GeminiCliParsedBucket | null => { + const modelId = normalizeStringValue(bucket.model_id ?? bucket.modelId); + if (!modelId) return null; + + const tokenType = normalizeStringValue(bucket.token_type ?? bucket.tokenType); + const remainingFractionRaw = normalizeNumberValue( + bucket.remaining_fraction ?? bucket.remainingFraction + ); + const remainingAmount = normalizeNumberValue( + bucket.remaining_amount ?? bucket.remainingAmount + ); + const resetTime = normalizeStringValue(bucket.reset_time ?? bucket.resetTime); + + let fallbackFraction: number | null = null; + if (remainingAmount !== null) { + fallbackFraction = remainingAmount <= 0 ? 0 : null; + } else if (resetTime) { + fallbackFraction = 0; + } + + return { + modelId, + tokenType, + remainingFraction: remainingFractionRaw ?? fallbackFraction ?? 1, + remainingAmount, + resetTime, + }; + }) + .filter((bucket): bucket is GeminiCliParsedBucket => bucket !== null); + + return buildGeminiCliBucketsFromParsedBuckets(parsedBuckets); +} diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/constants.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/constants.ts new file mode 100644 index 00000000..22714fec --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/constants.ts @@ -0,0 +1,34 @@ +/** + * Constants for the Gemini CLI quota fetcher submodule. + * + * Google Cloud Code API endpoints, error-detail sanitization limits, and + * upstream request timeouts. Extracted verbatim from the original god file; + * do not change values without coordinating with callers and tests. + */ + +/** Google Cloud Code internal API base URL. */ +export const GEMINI_CLI_API_BASE = 'https://cloudcode-pa.googleapis.com'; + +/** Google Cloud Code API version path segment. */ +export const GEMINI_CLI_API_VERSION = 'v1internal'; + +/** retrieveUserQuota endpoint - returns bucket-based model quotas. */ +export const GEMINI_CLI_QUOTA_URL = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:retrieveUserQuota`; + +/** loadCodeAssist endpoint - returns tier/credit metadata. */ +export const GEMINI_CLI_CODE_ASSIST_URL = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:loadCodeAssist`; + +/** Max characters retained from a sanitized upstream error detail. */ +export const GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH = 320; + +/** Suffix appended when an error detail is truncated. */ +export const GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX = '...[truncated]'; + +/** Credit type identifying Google One AI (paid tier) credits. */ +export const GEMINI_CLI_G1_CREDIT_TYPE = 'GOOGLE_ONE_AI'; + +/** Timeout for the primary (preferred) management API attempt, in ms. */ +export const MANAGEMENT_API_TIMEOUT_MS = 5000; + +/** Timeout for the secondary / fallback upstream request, in ms. */ +export const SECONDARY_REQUEST_TIMEOUT_MS = 2000; diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/error-parsing.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/error-parsing.ts new file mode 100644 index 00000000..fdadbd58 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/error-parsing.ts @@ -0,0 +1,318 @@ +/** + * Error parsing and failure-result builders for the Gemini CLI quota fetcher. + * + * Translates non-200 upstream responses into structured {@link GeminiCliQuotaResult} + * failure payloads with sanitized error details, recovery hints, and provider + * entitlement evidence. Token values in error bodies are always redacted before + * being surfaced (see {@link sanitizeGeminiCliErrorDetail}). + */ + +import { + buildProviderEntitlementEvidence, + isModelCapacityExhausted, +} from '../../auth/provider-entitlement-evidence'; +import type { + GeminiCliFailureResultOptions, + GeminiCliQuotaResult, + ParsedGeminiCliErrorBody, +} from './types'; +import { + GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH, + GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX, +} from './constants'; + +/** + * Build a structured failure {@link GeminiCliQuotaResult} with empty buckets. + * Centralizes the common failure shape so each HTTP-status branch only needs + * to supply its specific error/hint/entitlement fields. + */ +export function buildGeminiCliFailureResult( + accountId: string, + projectId: string | null, + options: GeminiCliFailureResultOptions +): GeminiCliQuotaResult { + return { + success: false, + buckets: [], + projectId, + tierLabel: null, + tierId: null, + creditBalance: null, + lastUpdated: Date.now(), + accountId, + error: options.error, + httpStatus: options.httpStatus, + errorCode: options.errorCode, + errorDetail: options.errorDetail, + actionHint: options.actionHint, + retryable: options.retryable, + needsReauth: options.needsReauth, + isForbidden: options.isForbidden, + entitlement: options.entitlement, + }; +} + +/** + * Sanitize an upstream error body for safe inclusion in a quota result. + * + * - Collapses HTML responses to a placeholder (never leaks provider HTML). + * - Redacts common token/credential/secret field names and `Bearer `. + * - Collapses internal whitespace to single spaces. + * - Truncates to {@link GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH} with a sentinel suffix. + * + * Returns undefined for empty input. Token values are never preserved. + */ +export function sanitizeGeminiCliErrorDetail(bodyText: string): string | undefined { + const trimmed = bodyText.trim(); + if (!trimmed) { + return undefined; + } + + if (/^]+>/.test(trimmed)) { + return '[HTML error response omitted]'; + } + + let sanitized = trimmed + .replace( + /"(access[_-]?token|refresh[_-]?token|authorization|cookie|set-cookie|api[_-]?key|session[_-]?token|token)"\s*:\s*"[^"]*"/gi, + '"$1":"[redacted]"' + ) + .replace(/Bearer\s+[A-Za-z0-9._-]+/g, 'Bearer [redacted]') + .replace(/\s+/g, ' '); + + if (sanitized.length > GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH) { + sanitized = `${sanitized.slice( + 0, + GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH - GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX.length + )}${GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX}`; + } + + return sanitized; +} + +/** + * Recursively extract the first non-empty message-like field from a nested + * error `details` array/object. Looks for `message`, `localizedMessage`, + * `description`, `reason`, and `error` keys at any level. + */ +export function extractGeminiCliNestedMessage(value: unknown): string | undefined { + if (Array.isArray(value)) { + for (const entry of value) { + const nested = extractGeminiCliNestedMessage(entry); + if (nested) return nested; + } + return undefined; + } + + if (!value || typeof value !== 'object') { + return undefined; + } + + const record = value as Record; + const directMessage = [ + record.message, + record.localizedMessage, + record.description, + record.reason, + record.error, + ].find( + (candidate): candidate is string => typeof candidate === 'string' && candidate.trim().length > 0 + ); + if (directMessage) { + return directMessage; + } + + return undefined; +} + +/** + * Parse an upstream error body into a structured {@link ParsedGeminiCliErrorBody}. + * + * Extracts a top-level code/status, a message (looking inside `error` objects + * and nested `details`), and a sanitized error detail. Non-JSON bodies fall + * back to the raw (sanitized) trimmed text as the message. HTML bodies surface + * only as the sanitized detail placeholder, never as the message. + */ +export function parseGeminiCliErrorBody(bodyText: string): ParsedGeminiCliErrorBody { + const trimmed = bodyText.trim(); + if (!trimmed) { + return {}; + } + + const sanitizedDetail = sanitizeGeminiCliErrorDetail(trimmed); + + try { + const parsed = JSON.parse(trimmed) as Record; + const topLevelMessage = [parsed.message, parsed.error].find( + (candidate): candidate is string => + typeof candidate === 'string' && candidate.trim().length > 0 + ); + const topLevelCode = [parsed.code, parsed.status].find( + (candidate): candidate is string => + typeof candidate === 'string' && candidate.trim().length > 0 + ); + + if (parsed.error && typeof parsed.error === 'object') { + const error = parsed.error as Record; + return { + errorCode: + [error.status, error.code, topLevelCode].find( + (candidate): candidate is string => + typeof candidate === 'string' && candidate.trim().length > 0 + ) || undefined, + errorDetail: sanitizedDetail, + message: + [ + error.message, + error.error, + extractGeminiCliNestedMessage(error.details), + topLevelMessage, + ].find( + (candidate): candidate is string => + typeof candidate === 'string' && candidate.trim().length > 0 + ) || undefined, + }; + } + + return { + errorCode: topLevelCode, + errorDetail: sanitizedDetail, + message: + [topLevelMessage, extractGeminiCliNestedMessage(parsed.details)].find( + (candidate): candidate is string => + typeof candidate === 'string' && candidate.trim().length > 0 + ) || undefined, + }; + } catch { + return { + errorDetail: sanitizedDetail, + message: sanitizedDetail === '[HTML error response omitted]' ? undefined : trimmed, + }; + } +} + +/** + * Build a user-facing recovery hint for a 403 (forbidden) upstream response. + * Inspects the parsed message/detail for verification, project, or generic + * access signals and returns the matching recovery instruction. + */ +export function buildGeminiCliForbiddenActionHint(parsed: ParsedGeminiCliErrorBody): string { + const combined = `${parsed.message || ''} ${parsed.errorDetail || ''}`.toLowerCase(); + if (combined.includes('verify') || combined.includes('verification')) { + return 'Complete the Google account verification mentioned above, then retry quota refresh.'; + } + if (combined.includes('project')) { + return 'Confirm this Google project still has Gemini CLI quota access, then retry.'; + } + return 'Check the Google account or workspace access shown above, then retry quota refresh.'; +} + +/** + * Build a structured failure result from an HTTP non-200 upstream response. + * + * Status-specific behavior: + * - 401: marks the result as needsReauth (user must re-run `ccs gemini --auth`) + * - 403: marks forbidden with runtime-inferred not_entitled evidence and a + * context-aware action hint + * - 429: distinguishes MODEL_CAPACITY_EXHAUSTED (entitled but capacity-stressed) + * from generic rate limiting + * - >=500: retryable provider-unavailable result + * - other: generic non-retryable quota_request_failed + */ +export function buildGeminiCliHttpFailureResult( + accountId: string, + projectId: string | null, + status: number, + bodyText: string +): GeminiCliQuotaResult { + const parsed = parseGeminiCliErrorBody(bodyText); + + if (status === 401) { + return buildGeminiCliFailureResult(accountId, projectId, { + error: parsed.message || 'Token expired or invalid', + httpStatus: 401, + errorCode: parsed.errorCode || 'reauth_required', + errorDetail: parsed.errorDetail, + actionHint: 'Run ccs gemini --auth to reconnect this account.', + needsReauth: true, + retryable: false, + }); + } + + if (status === 403) { + return buildGeminiCliFailureResult(accountId, projectId, { + error: parsed.message || 'Quota access forbidden for this account', + httpStatus: 403, + errorCode: parsed.errorCode || 'quota_api_forbidden', + errorDetail: parsed.errorDetail, + actionHint: buildGeminiCliForbiddenActionHint(parsed), + isForbidden: true, + retryable: false, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'medium', + accessState: 'not_entitled', + capacityState: 'unknown', + }), + }); + } + + if (status === 429) { + if (isModelCapacityExhausted(parsed.message, parsed.errorDetail, parsed.errorCode)) { + return buildGeminiCliFailureResult(accountId, projectId, { + error: parsed.message || 'Model capacity exhausted for this account right now', + httpStatus: 429, + errorCode: 'capacity_exhausted', + errorDetail: parsed.errorDetail, + actionHint: + 'Retry later or switch to another Gemini model. This indicates temporary model capacity, not an authentication failure.', + retryable: true, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'medium', + accessState: 'entitled', + capacityState: 'capacity_exhausted', + notes: 'Upstream returned MODEL_CAPACITY_EXHAUSTED for this model.', + }), + }); + } + + return buildGeminiCliFailureResult(accountId, projectId, { + error: parsed.message || 'Rate limited - try again later', + httpStatus: 429, + errorCode: parsed.errorCode || 'rate_limited', + errorDetail: parsed.errorDetail, + actionHint: 'Retry after a short delay.', + retryable: true, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'low', + accessState: 'unknown', + capacityState: 'rate_limited', + }), + }); + } + + if (status >= 500) { + return buildGeminiCliFailureResult(accountId, projectId, { + error: parsed.message || `Gemini quota service unavailable (HTTP ${status})`, + httpStatus: status, + errorCode: parsed.errorCode || 'provider_unavailable', + errorDetail: parsed.errorDetail, + actionHint: 'Retry later. This looks like a temporary Google upstream problem.', + retryable: true, + }); + } + + return buildGeminiCliFailureResult(accountId, projectId, { + error: parsed.message || `Gemini quota request failed (HTTP ${status})`, + httpStatus: status, + errorCode: parsed.errorCode || 'quota_request_failed', + errorDetail: parsed.errorDetail, + actionHint: 'Inspect the upstream response details and retry if appropriate.', + retryable: false, + }); +} diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/index.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/index.ts new file mode 100644 index 00000000..24a6fb0a --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/index.ts @@ -0,0 +1,41 @@ +/** + * Barrel for the Gemini CLI quota fetcher submodule. + * + * Re-exports the original public surface of `quota-fetcher-gemini-cli.ts` + * so the file at the original path can be reduced to a thin re-export + * (preserving import paths and signatures). Submodules are private + * implementation detail; only the symbols below are part of the contract. + */ + +// Public API +export { fetchGeminiCliQuota, fetchAllGeminiCliQuotas } from './quota-fetcher'; + +// Exported helpers (also part of the public surface - used by tests and +// the bucket/grouping normalization tests). +export { resolveGeminiCliProjectId } from './token-parsing'; +export { buildGeminiCliBuckets } from './bucket-building'; + +// Test exports: keep the original `__testExports` bag shape stable so the +// existing test suite (which destructures `__testExports`) keeps working. +export { + sanitizeGeminiCliErrorDetail, + extractGeminiCliNestedMessage, + parseGeminiCliErrorBody, + buildGeminiCliForbiddenActionHint, +} from './error-parsing'; + +// Re-export `__testExports` as a single object to preserve the original +// named-const export shape (`__testExports`). +import { + sanitizeGeminiCliErrorDetail, + extractGeminiCliNestedMessage, + parseGeminiCliErrorBody, + buildGeminiCliForbiddenActionHint, +} from './error-parsing'; + +export const __testExports = { + sanitizeGeminiCliErrorDetail, + extractGeminiCliNestedMessage, + parseGeminiCliErrorBody, + buildGeminiCliForbiddenActionHint, +}; diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/managed-request.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/managed-request.ts new file mode 100644 index 00000000..33234595 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/managed-request.ts @@ -0,0 +1,327 @@ +/** + * Managed and direct HTTP request machinery for the Gemini CLI quota fetcher. + * + * Wraps the two upstream call paths used when fetching Gemini CLI quota and + * supplementary metadata: + * - managed: delegated through the CLIProxy management API (uses $TOKEN$ + * substitution so the local process never holds the live token) + * - direct: bearer-token fetch against the Google Cloud Code endpoint + * + * The preferred path is configurable per call. On a 401 from the direct path, + * the managed path is retried as a delegated-auth refresh fallback. A + * `GeminiManagedAuthUnavailableError` is thrown when managed auth is required + * but unreachable, so callers can surface a retryable failure result. + */ + +import { + buildManagementHeaders, + buildProxyUrl, + getProxyTarget, +} from '../../proxy/proxy-target-resolver'; +import { mapExternalProviderName } from '../../provider-capabilities'; +import { sanitizeEmail } from '../../auth/auth-utils'; +import { MANAGEMENT_API_TIMEOUT_MS, SECONDARY_REQUEST_TIMEOUT_MS } from './constants'; +import { getRemainingTimeoutMs, normalizeStringValue, safeParseJson } from './shared-utils'; +import type { + ManagedGeminiAuthContext, + ManagedGeminiAuthLookupResult, + ManagedGeminiRequestResult, + ManagedResponse, + ManagementApiCallResponse, + ManagementAuthFile, +} from './types'; + +/** + * Thrown when Gemini delegated auth refresh via the CLIProxy management API + * is required but temporarily unreachable. Callers translate this into a + * retryable failure result. + */ +export class GeminiManagedAuthUnavailableError extends Error { + constructor() { + super('CLIProxy managed Gemini auth is temporarily unavailable'); + this.name = 'GeminiManagedAuthUnavailableError'; + } +} + +/** + * Read a fetch Response into the normalized {@link ManagedResponse} shape. + * `viaManagement` marks whether the response came through the managed API so + * downstream log messages can attribute the source correctly. + */ +export async function readManagedResponse( + response: Response, + viaManagement: boolean +): Promise { + const bodyText = await response.text(); + return { + status: response.status, + bodyText, + json: safeParseJson(bodyText), + viaManagement, + }; +} + +/** + * Check whether a filename matches the Gemini CLI auth file naming patterns. + * Recognizes three patterns: + * - legacy: gemini-*.json + * - new: *-gen-lang-client-*.json + * - email: contains "@" (verified against type inside the payload later) + */ +export function isGeminiAuthFile(filename: string): boolean { + if (!filename.endsWith('.json')) return false; + // Legacy pattern: gemini-email.json + if (filename.startsWith('gemini-')) return true; + // New pattern: email-gen-lang-client-projectId.json + if (filename.includes('-gen-lang-client-')) return true; + // Check if contains @ (email pattern) - will verify type inside + if (filename.includes('@')) return true; + return false; +} + +/** + * Determine whether a management-API auth-file descriptor belongs to the + * given Gemini account. Matches on provider/type normalized to "gemini", + * then on email, filename, or sanitized email substring. + */ +export function isGeminiAuthFileForAccount(file: ManagementAuthFile, accountId: string): boolean { + const rawProvider = normalizeStringValue(file.provider ?? file.type); + if (!rawProvider || mapExternalProviderName(rawProvider) !== 'gemini') { + return false; + } + + const email = normalizeStringValue(file.email); + const normalizedAccountId = accountId.trim().toLowerCase(); + if (email?.toLowerCase() === normalizedAccountId) { + return true; + } + + const normalizedName = normalizeStringValue(file.name); + if (!normalizedName) { + return false; + } + + const normalizedFileName = normalizedName.toLowerCase(); + const sanitizedAccount = sanitizeEmail(accountId).toLowerCase(); + return ( + normalizedFileName === `gemini-${sanitizedAccount}.json` || + normalizedFileName.startsWith(`${normalizedAccountId}-gen-lang-client-`) || + normalizedFileName.includes(sanitizedAccount) + ); +} + +/** + * Look up the management-API auth index for a Gemini account. + * Hits `/v0/management/auth-files` and matches the entry for this account. + * Returns `{ unavailable: true }` if the management API is unreachable or + * returns a non-OK response, so callers can fall back to direct auth. + */ +export async function findManagedGeminiAuthIndex( + accountId: string, + timeoutMs: number +): Promise { + const target = getProxyTarget(); + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), timeoutMs); + + try { + const response = await fetch(buildProxyUrl(target, '/v0/management/auth-files'), { + signal: controller.signal, + headers: buildManagementHeaders(target), + }); + clearTimeout(timeoutId); + + if (!response.ok) { + return { authIndex: null, unavailable: true }; + } + + const data = (await response.json()) as { files?: ManagementAuthFile[] }; + const match = data.files?.find((file) => isGeminiAuthFileForAccount(file, accountId)); + return { authIndex: match?.auth_index ?? null, unavailable: false }; + } catch { + clearTimeout(timeoutId); + return { authIndex: null, unavailable: true }; + } +} + +/** + * Look up the management auth index for a Gemini account, deduping concurrent + * lookups for the same account via the shared {@link ManagedGeminiAuthContext}. + * The first caller wins; subsequent callers await the same promise. + */ +export async function getManagedGeminiAuthIndex( + accountId: string, + timeoutMs: number, + context?: ManagedGeminiAuthContext +): Promise { + if (!context) { + return await findManagedGeminiAuthIndex(accountId, timeoutMs); + } + + context.authIndexLookupPromise ??= findManagedGeminiAuthIndex(accountId, timeoutMs); + return await context.authIndexLookupPromise; +} + +/** + * Perform a single upstream request to the Gemini CLI API via the CLIProxy + * management `/v0/management/api-call` endpoint. Uses `$TOKEN$` substitution + * so the live token never leaves the management API. Returns + * `{ unavailable: true }` if the management path is unreachable; returns + * `{ response: null, unavailable: false }` if the request succeeded but no + * matching auth file was found. + */ +export async function performManagedGeminiRequest( + accountId: string, + url: string, + body: string, + timeoutMs: number, + authContext?: ManagedGeminiAuthContext +): Promise { + const deadlineMs = Date.now() + timeoutMs; + const lookupResult = await getManagedGeminiAuthIndex( + accountId, + getRemainingTimeoutMs(deadlineMs), + authContext + ); + if (lookupResult.unavailable) { + return { response: null, unavailable: true }; + } + + const authIndex = lookupResult.authIndex; + if (authIndex === null || authIndex === undefined) { + return { response: null, unavailable: false }; + } + + const target = getProxyTarget(); + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), getRemainingTimeoutMs(deadlineMs)); + + try { + const response = await fetch(buildProxyUrl(target, '/v0/management/api-call'), { + method: 'POST', + signal: controller.signal, + headers: buildManagementHeaders(target, { + 'Content-Type': 'application/json', + }), + body: JSON.stringify({ + auth_index: authIndex, + method: 'POST', + url, + header: { + Authorization: 'Bearer $TOKEN$', + 'Content-Type': 'application/json', + }, + data: body, + }), + }); + clearTimeout(timeoutId); + + if (!response.ok) { + return { response: null, unavailable: true }; + } + + const apiResponse = (await response.json()) as ManagementApiCallResponse; + const bodyText = typeof apiResponse.body === 'string' ? apiResponse.body : ''; + return { + response: { + status: typeof apiResponse.status_code === 'number' ? apiResponse.status_code : 500, + bodyText, + json: safeParseJson(bodyText), + viaManagement: true, + }, + unavailable: false, + }; + } catch { + clearTimeout(timeoutId); + return { response: null, unavailable: true }; + } +} + +/** + * Perform a Gemini CLI upstream request, preferring the managed path when + * requested and falling back to direct bearer-token auth. On a 401 from the + * direct path, retries via managed auth as a delegated-auth refresh; throws + * {@link GeminiManagedAuthUnavailableError} if that retry is unreachable. + * + * @param accountId Account identifier (email), used for managed auth lookup. + * @param accessToken Bearer token for the direct path. Never logged. + * @param url Target Gemini CLI API URL. + * @param body JSON request body string. + * @param preferManagement When true, try the managed path first. + * @param authContext Optional shared context to dedupe auth-index lookups. + */ +export async function performGeminiCliRequest( + accountId: string, + accessToken: string, + url: string, + body: string, + preferManagement = false, + authContext?: ManagedGeminiAuthContext +): Promise { + let managementAttempted = false; + let managementUnavailable = false; + + if (preferManagement) { + managementAttempted = true; + const managedResult = await performManagedGeminiRequest( + accountId, + url, + body, + MANAGEMENT_API_TIMEOUT_MS, + authContext + ); + managementUnavailable = managedResult.unavailable; + if (managedResult.response) { + return managedResult.response; + } + } + + const controller = new AbortController(); + const timeoutId = setTimeout( + () => controller.abort(), + managementAttempted ? SECONDARY_REQUEST_TIMEOUT_MS : MANAGEMENT_API_TIMEOUT_MS + ); + + try { + const response = await fetch(url, { + method: 'POST', + signal: controller.signal, + headers: { + Authorization: `Bearer ${accessToken}`, + 'Content-Type': 'application/json', + }, + body, + }); + clearTimeout(timeoutId); + + const directResult = await readManagedResponse(response, false); + if (directResult.status !== 401) { + return directResult; + } + + if (managementAttempted) { + if (managementUnavailable) { + throw new GeminiManagedAuthUnavailableError(); + } + return directResult; + } + + const managedResult = await performManagedGeminiRequest( + accountId, + url, + body, + SECONDARY_REQUEST_TIMEOUT_MS, + authContext + ); + if (managedResult.response) { + return managedResult.response; + } + if (managedResult.unavailable) { + throw new GeminiManagedAuthUnavailableError(); + } + return directResult; + } catch (error) { + clearTimeout(timeoutId); + throw error; + } +} diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/quota-fetcher.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/quota-fetcher.ts new file mode 100644 index 00000000..b71cc864 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/quota-fetcher.ts @@ -0,0 +1,242 @@ +/** + * Top-level quota fetch orchestration for Gemini CLI accounts. + * + * Coordinates auth-file discovery, the managed/direct upstream quota request, + * supplementary tier/credit metadata, and structured failure-result building. + * Preserves the structured logging from the original god file (events: + * gemini_cli.fetch_start, gemini_cli.auth_file_missing, gemini_cli.token_expired, + * gemini_cli.missing_project_id, gemini_cli.api_status, gemini_cli.buckets_found, + * gemini_cli.quota_fetch_error). Token values are never logged. + */ + +import { getProviderAccounts, setAccountTier } from '../../accounts/account-manager'; +import { getTokenExpiryTimestamp } from '../../auth/auth-utils'; +import { buildProviderEntitlementEvidence } from '../../auth/provider-entitlement-evidence'; +import type { GeminiCliQuotaResult } from '../quota-types'; +import { readGeminiCliAuthData } from './auth-file-discovery'; +import { buildGeminiCliBuckets } from './bucket-building'; +import { buildGeminiCliFailureResult, buildGeminiCliHttpFailureResult } from './error-parsing'; +import { GeminiManagedAuthUnavailableError, performGeminiCliRequest } from './managed-request'; +import { fetchGeminiCliSupplementary } from './supplementary-metadata'; +import { logger } from './shared-utils'; +import { GEMINI_CLI_QUOTA_URL } from './constants'; +import type { GeminiCliAuthData, GeminiCliQuotaResponse, ManagedGeminiAuthContext } from './types'; + +/** + * Internal helper: fetch quota with already-validated auth data. + * + * Extracted to support the auto-refresh retry path: the caller resolves auth + * data once (legacy file or managed), then this function performs the upstream + * quota request and supplementary metadata fetch in parallel. On success it + * persists the resolved tier back to the account via `setAccountTier`. + */ +export async function fetchWithAuthData( + authData: GeminiCliAuthData, + accountId: string, + verbose: boolean +): Promise { + if (!authData.projectId) { + const error = 'Cannot resolve project ID from auth file'; + if (verbose) { + logger.error('gemini_cli.missing_project_id', `Error: ${error}`, { + provider: 'gemini', + accountId, + }); + } + return buildGeminiCliFailureResult(accountId, null, { + error, + errorCode: 'missing_project_id', + actionHint: 'Run ccs gemini --auth to reconnect this account and recover the project ID.', + retryable: false, + }); + } + + const authContext: ManagedGeminiAuthContext = {}; + const supplementaryPromise = fetchGeminiCliSupplementary( + accountId, + authData.accessToken, + authData.projectId, + verbose, + authContext + ); + const requestBody = JSON.stringify({ project: authData.projectId }); + + try { + const response = await performGeminiCliRequest( + accountId, + authData.accessToken, + GEMINI_CLI_QUOTA_URL, + requestBody, + authData.isExpired, + authContext + ); + + if (verbose) { + const source = response.viaManagement ? 'managed' : 'direct'; + logger.info( + 'gemini_cli.api_status', + `Gemini CLI API status via ${source}: ${response.status}`, + { provider: 'gemini', accountId, httpStatus: response.status, source } + ); + } + + if (response.status !== 200) { + return buildGeminiCliHttpFailureResult( + accountId, + authData.projectId, + response.status, + response.bodyText + ); + } + + const data = response.json as GeminiCliQuotaResponse | null; + const rawBuckets = data?.buckets || []; + const buckets = buildGeminiCliBuckets(rawBuckets); + const supplementary = await supplementaryPromise; + + if (verbose) { + logger.info('gemini_cli.buckets_found', `Gemini CLI buckets found: ${buckets.length}`, { + provider: 'gemini', + accountId, + bucketCount: buckets.length, + }); + } + + if (supplementary.normalizedTier !== 'unknown') { + setAccountTier('gemini', accountId, supplementary.normalizedTier); + } + + return { + success: true, + buckets, + projectId: authData.projectId, + tierLabel: supplementary.tierLabel, + tierId: supplementary.tierId, + creditBalance: supplementary.creditBalance, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: supplementary.normalizedTier, + rawTierId: supplementary.tierId, + rawTierLabel: supplementary.tierLabel, + source: supplementary.tierId ? 'runtime_api' : 'runtime_inference', + confidence: supplementary.tierId ? 'high' : 'medium', + accessState: 'entitled', + capacityState: 'available', + }), + lastUpdated: Date.now(), + accountId, + }; + } catch (err) { + if (err instanceof GeminiManagedAuthUnavailableError) { + return buildGeminiCliFailureResult(accountId, authData.projectId, { + error: 'Gemini delegated auth refresh is temporarily unavailable', + errorCode: 'managed_auth_unavailable', + errorDetail: err.message, + actionHint: 'Retry later. CLIProxy management could not refresh this Gemini account.', + retryable: true, + }); + } + + const errorMsg = + err instanceof Error && err.name === 'AbortError' + ? 'Request timeout' + : err instanceof Error + ? err.message + : 'Unknown error'; + + if (verbose) { + logger.error('gemini_cli.quota_fetch_error', `Gemini CLI quota error: ${errorMsg}`, { + provider: 'gemini', + accountId, + err: err instanceof Error ? { name: err.name, message: errorMsg } : { message: errorMsg }, + }); + } + + return buildGeminiCliFailureResult(accountId, authData.projectId, { + error: errorMsg, + errorCode: + err instanceof Error && err.name === 'AbortError' ? 'network_timeout' : 'network_error', + actionHint: 'Retry later. This looks temporary.', + retryable: true, + httpStatus: err instanceof Error && err.name === 'AbortError' ? 408 : undefined, + }); + } +} + +/** + * Fetch quota for a single Gemini CLI account. + * + * Reads the on-disk auth file, emits the structured `gemini_cli.fetch_start` + * and `gemini_cli.auth_file_missing` / `gemini_cli.token_expired` log events + * (gated on `verbose`), and delegates to {@link fetchWithAuthData}. Token + * values are never logged; only the expiry label is surfaced. + * + * @param accountId - Account identifier (email) + * @param verbose - Show detailed diagnostics + * @returns Quota result with buckets, percentages, tier, and entitlement evidence + */ +export async function fetchGeminiCliQuota( + accountId: string, + verbose = false +): Promise { + if (verbose) { + logger.info('gemini_cli.fetch_start', `Fetching Gemini CLI quota for ${accountId}...`, { + provider: 'gemini', + accountId, + }); + } + + const authData = readGeminiCliAuthData(accountId); + if (!authData) { + const error = 'Auth file not found for Gemini account'; + if (verbose) { + logger.error('gemini_cli.auth_file_missing', `Error: ${error}`, { + provider: 'gemini', + accountId, + }); + } + return buildGeminiCliFailureResult(accountId, null, { + error, + errorCode: 'auth_file_missing', + actionHint: 'Run ccs gemini --auth to reconnect this account.', + retryable: false, + }); + } + + if (authData.isExpired && verbose) { + const expiresAt = getTokenExpiryTimestamp(authData.expiresAt); + const expiryLabel = expiresAt ? new Date(expiresAt).toISOString() : 'unknown'; + logger.info( + 'gemini_cli.token_expired', + `Gemini access token is expired (${expiryLabel}); quota requests will defer to managed auth when available.`, + { provider: 'gemini', accountId, tokenExpired: true, expiresAt: expiryLabel } + ); + } + + return await fetchWithAuthData(authData, accountId, verbose); +} + +/** + * Fetch quota for all configured Gemini CLI accounts in parallel. + * + * @param verbose - Show detailed diagnostics (forwarded to each per-account fetch) + * @returns Array of `{ account, quota }` entries, one per active Gemini account. + * Returns an empty array when there are no Gemini accounts configured. + */ +export async function fetchAllGeminiCliQuotas( + verbose = false +): Promise<{ account: string; quota: GeminiCliQuotaResult }[]> { + const accounts = getProviderAccounts('gemini'); + + if (accounts.length === 0) { + return []; + } + + const results = await Promise.all( + accounts.map(async (account) => ({ + account: account.id, + quota: await fetchGeminiCliQuota(account.id, verbose), + })) + ); + + return results; +} diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/shared-utils.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/shared-utils.ts new file mode 100644 index 00000000..f07cf53c --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/shared-utils.ts @@ -0,0 +1,62 @@ +/** + * Shared utilities for the Gemini CLI quota fetcher submodule. + * + * Includes the diagnostic logger (provider context = cliproxy:quota:gemini-cli) + * and small value-normalization helpers used by multiple submodules. Token + * values are never logged here; accountId is attached as provider context + * only. + */ + +import { createLogger } from '../../../services/logging'; + +/** + * Diagnostic-only logger for Gemini CLI quota fetch progress, upstream HTTP + * status, and recovery hints. Token values live in auth files and are never + * read into log messages. + */ +export const logger = createLogger('cliproxy:quota:gemini-cli'); + +/** + * Normalize a raw value into a trimmed non-empty string, or null. + * Returns null for empty strings, non-strings, or whitespace-only input. + */ +export function normalizeStringValue(value: unknown): string | null { + return typeof value === 'string' && value.trim().length > 0 ? value.trim() : null; +} + +/** + * Normalize a raw value into a finite number, or null. + * Accepts actual numbers and numeric strings; rejects NaN/Infinity. + */ +export function normalizeNumberValue(value: unknown): number | null { + if (typeof value === 'number' && Number.isFinite(value)) { + return value; + } + if (typeof value === 'string' && value.trim().length > 0) { + const parsed = Number(value); + if (Number.isFinite(parsed)) { + return parsed; + } + } + return null; +} + +/** + * Best-effort JSON.parse that returns null on failure instead of throwing. + * Used when normalizing upstream response bodies into a `json` field. + */ +export function safeParseJson(bodyText: string): unknown { + try { + return JSON.parse(bodyText); + } catch { + return null; + } +} + +/** + * Compute the remaining milliseconds available before a deadline, clamped + * to a minimum of 1ms so AbortController timeouts are always positive. + */ +export function getRemainingTimeoutMs(deadlineMs: number): number { + return Math.max(1, deadlineMs - Date.now()); +} diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/supplementary-metadata.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/supplementary-metadata.ts new file mode 100644 index 00000000..150ef9db --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/supplementary-metadata.ts @@ -0,0 +1,149 @@ +/** + * Supplementary tier/credit metadata fetcher for the Gemini CLI quota fetcher. + * + * Wraps the `loadCodeAssist` endpoint to resolve the account's tier label, + * tier id, normalized tier (free/pro/ultra/unknown), and Google One AI credit + * balance. Runs alongside the primary quota fetch and shares the same + * managed-auth context so auth-index lookups are deduped. + */ + +import { + getProviderTierLabel, + normalizeProviderTierId, +} from '../../auth/provider-entitlement-evidence'; +import { performGeminiCliRequest } from './managed-request'; +import { logger, normalizeNumberValue, normalizeStringValue } from './shared-utils'; +import { GEMINI_CLI_CODE_ASSIST_URL, GEMINI_CLI_G1_CREDIT_TYPE } from './constants'; +import type { + GeminiCliCodeAssistResponse, + GeminiCliSupplementaryInfo, + ManagedGeminiAuthContext, +} from './types'; + +/** + * Resolve the tier id from a loadCodeAssist response. + * Prefers the paid tier id, then the current tier id. Lowercased. + * Returns null if neither is present. + */ +export function resolveGeminiCliTierId(payload: GeminiCliCodeAssistResponse | null): string | null { + if (!payload) return null; + const currentTier = payload.currentTier ?? payload.current_tier; + const paidTier = payload.paidTier ?? payload.paid_tier; + const rawId = normalizeStringValue(paidTier?.id) ?? normalizeStringValue(currentTier?.id); + return rawId ? rawId.toLowerCase() : null; +} + +/** + * Resolve a human-readable tier label from the loadCodeAssist tier id. + * Returns null when the tier id cannot be mapped to a known label. + */ +export function resolveGeminiCliTierLabel( + payload: GeminiCliCodeAssistResponse | null +): string | null { + const tierId = resolveGeminiCliTierId(payload); + return getProviderTierLabel(tierId); +} + +/** + * Resolve the Google One AI credit balance for the account from the + * loadCodeAssist response. Sums all credits with type `GOOGLE_ONE_AI` on the + * paid tier (preferred) or current tier. Returns null if no matching credits + * are present. + */ +export function resolveGeminiCliCreditBalance( + payload: GeminiCliCodeAssistResponse | null +): number | null { + if (!payload) return null; + + const paidTier = payload.paidTier ?? payload.paid_tier; + const currentTier = payload.currentTier ?? payload.current_tier; + const tier = paidTier ?? currentTier; + if (!tier) return null; + + const credits = tier.availableCredits ?? tier.available_credits ?? []; + let total = 0; + let found = false; + for (const credit of credits) { + const creditType = normalizeStringValue(credit.creditType ?? credit.credit_type); + if (creditType !== GEMINI_CLI_G1_CREDIT_TYPE) continue; + + const amount = normalizeNumberValue(credit.creditAmount ?? credit.credit_amount); + if (amount !== null) { + total += amount; + found = true; + } + } + + return found ? total : null; +} + +/** + * Fetch supplementary tier/credit metadata for a Gemini account via the + * loadCodeAssist endpoint. Never throws: on any failure returns a + * supplementary info with `normalizedTier: 'unknown'` so the primary quota + * fetch can still complete. Diagnostic logging is gated on `verbose` and + * records only accountId, HTTP status, and the source (managed/direct); + * token values are never logged. + */ +export async function fetchGeminiCliSupplementary( + accountId: string, + accessToken: string, + projectId: string, + verbose: boolean, + authContext?: ManagedGeminiAuthContext +): Promise { + const requestBody = JSON.stringify({ + cloudaicompanionProject: projectId, + metadata: { + ideType: 'IDE_UNSPECIFIED', + platform: 'PLATFORM_UNSPECIFIED', + pluginType: 'GEMINI', + duetProject: projectId, + }, + }); + + try { + const response = await performGeminiCliRequest( + accountId, + accessToken, + GEMINI_CLI_CODE_ASSIST_URL, + requestBody, + false, + authContext + ); + + if (response.status !== 200) { + if (verbose) { + const source = response.viaManagement ? 'managed' : 'direct'; + logger.info( + 'gemini_cli.supplementary_metadata_unavailable', + `Gemini CLI supplementary metadata unavailable via ${source}: HTTP ${response.status}`, + { provider: 'gemini', accountId, httpStatus: response.status, source } + ); + } + return { tierLabel: null, tierId: null, creditBalance: null, normalizedTier: 'unknown' }; + } + + const payload = response.json as GeminiCliCodeAssistResponse | null; + return { + tierLabel: resolveGeminiCliTierLabel(payload), + tierId: resolveGeminiCliTierId(payload), + creditBalance: resolveGeminiCliCreditBalance(payload), + normalizedTier: normalizeProviderTierId(resolveGeminiCliTierId(payload)), + }; + } catch (error) { + if (verbose) { + const message = error instanceof Error ? error.message : 'Unknown error'; + logger.info( + 'gemini_cli.supplementary_metadata_skipped', + `Gemini CLI supplementary metadata skipped: ${message}`, + { + provider: 'gemini', + accountId, + err: error instanceof Error ? { name: error.name, message } : { message }, + } + ); + } + return { tierLabel: null, tierId: null, creditBalance: null, normalizedTier: 'unknown' }; + } +} diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/token-parsing.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/token-parsing.ts new file mode 100644 index 00000000..7ee1eaa3 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/token-parsing.ts @@ -0,0 +1,73 @@ +/** + * Token parsing helpers for Gemini CLI auth files. + * + * Extracts access tokens, expiry, and project IDs from the raw auth file + * payload. Gemini auth files come in two structural variants: + * - flat: { access_token, expired, project_id, account } + * - nested:{ token: { access_token, expiry }, project_id, account } + * These helpers handle both without throwing on shape mismatches. + */ + +/** + * Extract the access token from a Gemini auth file payload. + * Handles both flat (`access_token`) and nested (`token.access_token`) shapes. + * Returns null if no usable token is present. + */ +export function extractAccessToken(data: Record): string | null { + // Flat structure: { access_token: "..." } + if (typeof data.access_token === 'string') { + return data.access_token; + } + // Nested structure: { token: { access_token: "..." } } + if (data.token && typeof data.token === 'object') { + const token = data.token as Record; + if (typeof token.access_token === 'string') { + return token.access_token; + } + } + return null; +} + +/** + * Extract the token expiry from a Gemini auth file payload. + * Handles both flat (`expired`) and nested (`token.expiry`) shapes. + * Returns the raw string/number, or null if absent. + */ +export function extractExpiry(data: Record): string | number | null { + // Flat structure: { expired: "..." } + if (typeof data.expired === 'string') { + return data.expired; + } + if (typeof data.expired === 'number') { + return data.expired; + } + // Nested structure: { token: { expiry: "..." } } + if (data.token && typeof data.token === 'object') { + const token = data.token as Record; + if (typeof token.expiry === 'string') { + return token.expiry; + } + if (typeof token.expiry === 'number') { + return token.expiry; + } + } + return null; +} + +/** + * Extract the project ID from an auth file's `account` field. + * Input shape: "user@example.com (cloudaicompanion-abc-123)" + * Returns the last parenthesized segment, or null if no match. + * + * Example: + * "user@example.com (cloudaicompanion-abc-123)" -> "cloudaicompanion-abc-123" + */ +export function resolveGeminiCliProjectId(accountField: string): string | null { + const regex = /\(([^()]+)\)/g; + let match: RegExpExecArray | null; + let lastMatch: string | null = null; + while ((match = regex.exec(accountField)) !== null) { + lastMatch = match[1]; + } + return lastMatch; +} diff --git a/src/cliproxy/quota/quota-fetcher-gemini-cli/types.ts b/src/cliproxy/quota/quota-fetcher-gemini-cli/types.ts new file mode 100644 index 00000000..a55d01cc --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher-gemini-cli/types.ts @@ -0,0 +1,132 @@ +/** + * Shared types for the Gemini CLI quota fetcher submodule. + * + * Extracted from the original quota-fetcher-gemini-cli.ts god file. These + * interfaces describe raw API response shapes, internal parsed structures, + * and managed-auth context used across the submodules. + */ + +import type { GeminiCliBucket, GeminiCliQuotaResult } from '../quota-types'; +import type { ProviderEntitlementEvidence } from '../../auth/provider-entitlement-types'; + +/** Auth data extracted from a Gemini CLI auth file. */ +export interface GeminiCliAuthData { + accessToken: string; + projectId: string | null; + isExpired: boolean; + expiresAt: string | number | null; +} + +/** Raw bucket shape returned by the Gemini CLI quota API. */ +export interface RawGeminiCliBucket { + model_id?: string; + modelId?: string; + token_type?: string | null; + tokenType?: string | null; + remaining_fraction?: number; + remainingFraction?: number; + remaining_amount?: number; + remainingAmount?: number; + reset_time?: string | null; + resetTime?: string | null; +} + +/** Raw quota API response wrapper. */ +export interface GeminiCliQuotaResponse { + buckets?: RawGeminiCliBucket[]; +} + +/** Credit entry inside a tier (supports snake_case and camelCase variants). */ +export interface GeminiCliCredits { + creditType?: string; + credit_type?: string; + creditAmount?: string | number; + credit_amount?: string | number; +} + +/** User tier inside a loadCodeAssist response. */ +export interface GeminiCliUserTier { + id?: string; + availableCredits?: GeminiCliCredits[]; + available_credits?: GeminiCliCredits[]; +} + +/** loadCodeAssist response shape (currentTier + paidTier). */ +export interface GeminiCliCodeAssistResponse { + currentTier?: GeminiCliUserTier | null; + current_tier?: GeminiCliUserTier | null; + paidTier?: GeminiCliUserTier | null; + paid_tier?: GeminiCliUserTier | null; +} + +/** Parsed error body extracted from an upstream non-200 response. */ +export interface ParsedGeminiCliErrorBody { + errorCode?: string; + errorDetail?: string; + message?: string; +} + +/** Supplementary tier/credit info resolved alongside the quota buckets. */ +export interface GeminiCliSupplementaryInfo { + tierLabel: string | null; + tierId: string | null; + creditBalance: number | null; + normalizedTier: 'free' | 'pro' | 'ultra' | 'unknown'; +} + +/** Auth-file entry as returned by the CLIProxy management API. */ +export interface ManagementAuthFile { + auth_index?: string | number; + provider?: string; + type?: string; + email?: string; + name?: string; +} + +/** api-call response envelope from the CLIProxy management endpoint. */ +export interface ManagementApiCallResponse { + status_code?: number; + body?: string; +} + +/** Normalized HTTP response used by both direct and managed code paths. */ +export interface ManagedResponse { + status: number; + bodyText: string; + json: unknown; + viaManagement: boolean; +} + +/** Per-account managed-auth context used to dedupe auth-index lookups. */ +export interface ManagedGeminiAuthContext { + authIndexLookupPromise?: Promise; +} + +/** Result of looking up a Gemini auth file index via management API. */ +export interface ManagedGeminiAuthLookupResult { + authIndex: string | number | null; + unavailable: boolean; +} + +/** Result of performing a managed Gemini upstream request. */ +export interface ManagedGeminiRequestResult { + response: ManagedResponse | null; + unavailable: boolean; +} + +/** Options bag for {@link buildGeminiCliFailureResult}. */ +export interface GeminiCliFailureResultOptions { + error: string; + httpStatus?: number; + errorCode?: string; + errorDetail?: string; + actionHint?: string; + retryable?: boolean; + needsReauth?: boolean; + isForbidden?: boolean; + entitlement?: ProviderEntitlementEvidence; +} + +// Re-export the public result shapes so callers can import everything from +// the barrel without reaching into quota-types directly. +export type { GeminiCliBucket, GeminiCliQuotaResult, ProviderEntitlementEvidence }; diff --git a/src/cliproxy/quota/quota-fetcher.ts b/src/cliproxy/quota/quota-fetcher.ts index 898a741e..6ae95388 100644 --- a/src/cliproxy/quota/quota-fetcher.ts +++ b/src/cliproxy/quota/quota-fetcher.ts @@ -1,1106 +1,29 @@ /** - * Quota Fetcher for Antigravity Accounts + * Quota Fetcher for Antigravity Accounts (barrel). * - * Fetches quota information from Google Cloud Code internal API. + * Fetches quota information from the Google Cloud Code internal API. * Used for displaying remaining quota percentages and reset times. - */ - -import * as fs from 'node:fs'; -import * as path from 'node:path'; -import { getAuthDir } from '../config/config-generator'; -import { CLIProxyProvider } from '../types'; -import { - getProviderAccounts, - getPausedDir, - setAccountTier, - type AccountInfo, - type AccountTier, -} from '../accounts/account-manager'; -import { sanitizeEmail, isTokenExpired } from '../auth/auth-utils'; -import { - buildProviderEntitlementEvidence, - getProviderTierLabel, - normalizeProviderTierId, -} from '../auth/provider-entitlement-evidence'; -import type { ProviderEntitlementEvidence } from '../auth/provider-entitlement-types'; -import { - buildManagementHeaders, - buildProxyUrl, - getProxyTarget, -} from '../proxy/proxy-target-resolver'; -import { createLogger } from '../../services/logging'; - -const logger = createLogger('cliproxy:quota:fetcher'); - -/** Individual model quota info */ -export interface ModelQuota { - /** Model name, e.g., "gemini-3-pro-high" */ - name: string; - /** Display name from API, e.g., "Gemini 3 Pro" */ - displayName?: string; - /** Remaining quota as percentage (0-100) */ - percentage: number; - /** ISO timestamp when quota resets, null if unknown */ - resetTime: string | null; -} - -/** Quota fetch result */ -export interface QuotaResult { - /** Whether fetch succeeded */ - success: boolean; - /** Quota for each available model */ - models: ModelQuota[]; - /** Timestamp of fetch */ - lastUpdated: number; - /** Upstream HTTP status when available */ - httpStatus?: number; - /** Stable machine-readable error code */ - errorCode?: string; - /** Additional provider-specific detail/code from upstream */ - errorDetail?: string; - /** True if account lacks quota access (403) */ - isForbidden?: boolean; - /** Error message if fetch failed */ - error?: string; - /** Provider-specific remediation guidance */ - actionHint?: string; - /** True when the failure is temporary and retrying later may help */ - retryable?: boolean; - /** True if token is expired and needs re-auth */ - isExpired?: boolean; - /** True if token refresh cannot proceed and the account should be re-authenticated */ - needsReauth?: boolean; - /** ISO timestamp when token expires/expired */ - expiresAt?: string; - /** True if account hasn't been activated in official Antigravity app */ - isUnprovisioned?: boolean; - /** Account ID (email) this quota belongs to */ - accountId?: string; - /** GCP project ID for this account */ - projectId?: string; - /** Detected account tier based on model access */ - tier?: AccountTier; - /** Richer provider entitlement evidence derived from live/runtime signals */ - entitlement?: ProviderEntitlementEvidence; -} - -/** Google Cloud Code API endpoints */ -const ANTIGRAVITY_DAILY_API_BASE = 'https://daily-cloudcode-pa.googleapis.com'; -const ANTIGRAVITY_API_BASE = 'https://cloudcode-pa.googleapis.com'; -const ANTIGRAVITY_API_VERSION = 'v1internal'; -const ANTIGRAVITY_LOADCODEASSIST_BASE_URLS = [ - ANTIGRAVITY_DAILY_API_BASE, - ANTIGRAVITY_API_BASE, -] as const; -const MANAGEMENT_API_TIMEOUT_MS = 5000; - -/** Headers for loadCodeAssist (matches current CLIProxyAPIPlus control-plane requests) */ -const LOADCODEASSIST_HEADERS = { - 'Content-Type': 'application/json', - 'User-Agent': 'antigravity/1.21.9 darwin/arm64 google-api-nodejs-client/10.3.0', - 'X-Goog-Api-Client': 'gl-node/22.21.1', -}; - -/** Headers for fetchAvailableModels (matches CLIProxyAPI antigravity_executor.go) */ -const FETCHMODELS_HEADERS = { - 'Content-Type': 'application/json', - 'User-Agent': 'antigravity/1.104.0 darwin/arm64', -}; - -/** Auth file structure */ -interface AntigravityAuthFile { - access_token: string; - refresh_token?: string; - email?: string; - expired?: string; - expires_in?: number; - timestamp?: number; - type?: string; - project_id?: string; -} - -/** Auth data returned from file */ -interface AuthData { - accessToken: string; - refreshToken: string | null; - projectId: string | null; - isExpired: boolean; - expiresAt: string | null; -} - -/** Tier info from loadCodeAssist */ -interface TierInfo { - id?: string; - isDefault?: boolean; -} - -/** loadCodeAssist response */ -interface LoadCodeAssistResponse { - cloudaicompanionProject?: string | { id?: string }; - /** Current tier (may be trial/temporary) */ - currentTier?: TierInfo; - /** Paid tier (reflects actual subscription - takes priority) */ - paidTier?: TierInfo; - /** Array of allowed tiers - use isDefault=true to find active tier (CLIProxyAPIPlus approach) */ - allowedTiers?: TierInfo[]; -} - -/** fetchAvailableModels response model */ -interface AvailableModel { - name?: string; - displayName?: string; - quotaInfo?: { - remainingFraction?: number; - remaining_fraction?: number; - remaining?: number; - resetTime?: string; - reset_time?: string; - }; - quota_info?: { - remainingFraction?: number; - remaining_fraction?: number; - remaining?: number; - resetTime?: string; - reset_time?: string; - }; -} - -/** fetchAvailableModels response */ -interface FetchAvailableModelsResponse { - models?: Record; -} - -interface ManagementAuthFile { - auth_index?: string | number; - provider?: string; - type?: string; - email?: string; - name?: string; -} - -interface ManagementApiCallResponse { - status_code?: number; - body?: string; -} - -interface ManagedResponse { - status: number; - bodyText: string; - json: unknown; - viaManagement: boolean; -} - -interface ProjectLookupResult { - projectId: string | null; - tier?: AccountTier; - rawTierId?: string | null; - rawTierLabel?: string | null; - entitlement?: ProviderEntitlementEvidence; - error?: string; - errorCode?: string; - errorDetail?: string; - actionHint?: string; - retryable?: boolean; - httpStatus?: number; - needsReauth?: boolean; - isUnprovisioned?: boolean; -} - -function safeParseJson(bodyText: string): unknown { - try { - return JSON.parse(bodyText); - } catch { - return null; - } -} - -function normalizeErrorDetail(bodyText: string): string | undefined { - const normalized = bodyText.trim(); - if (!normalized) { - return undefined; - } - if (normalized.length <= 400) { - return normalized; - } - return `${normalized.slice(0, 397)}...`; -} - -function buildAntigravityFailure( - status: number | undefined, - bodyText?: string -): Pick< - QuotaResult, - | 'error' - | 'errorCode' - | 'errorDetail' - | 'actionHint' - | 'retryable' - | 'httpStatus' - | 'needsReauth' - | 'entitlement' -> & { isForbidden?: boolean } { - const detail = normalizeErrorDetail(bodyText || ''); - - if (status === 401) { - return { - httpStatus: 401, - error: 'Token expired or invalid', - errorCode: 'reauth_required', - actionHint: - 'Re-authenticate this account. If CLIProxy is running, retry after the proxy finishes refreshing the token.', - needsReauth: true, - errorDetail: detail, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'medium', - accessState: 'unknown', - capacityState: 'unknown', - }), - }; - } - - if (status === 403) { - return { - httpStatus: 403, - error: 'Access forbidden', - errorCode: 'quota_api_forbidden', - actionHint: 'This account does not have Gemini Code Assist quota access.', - isForbidden: true, - errorDetail: detail, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'medium', - accessState: 'not_entitled', - capacityState: 'unknown', - }), - }; - } - - if (status === 429) { - return { - httpStatus: 429, - error: 'Rate limited - try again later', - errorCode: 'rate_limited', - actionHint: 'Retry later. This looks temporary.', - retryable: true, - errorDetail: detail, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'low', - accessState: 'unknown', - capacityState: 'rate_limited', - }), - }; - } - - if (status === 408) { - return { - httpStatus: 408, - error: 'Request timeout', - errorCode: 'network_timeout', - actionHint: 'Retry later. This looks temporary.', - retryable: true, - errorDetail: detail, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'low', - accessState: 'unknown', - capacityState: 'temporarily_unavailable', - }), - }; - } - - if (typeof status === 'number' && status >= 500) { - return { - httpStatus: status, - error: `API error: ${status}`, - errorCode: 'provider_unavailable', - actionHint: 'Retry later. The provider appears unavailable.', - retryable: true, - errorDetail: detail, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'low', - accessState: 'unknown', - capacityState: 'temporarily_unavailable', - }), - }; - } - - if (typeof status === 'number' && status >= 400) { - return { - httpStatus: status, - error: `API error: ${status}`, - errorCode: 'quota_request_failed', - errorDetail: detail, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'low', - accessState: 'unknown', - capacityState: 'unknown', - }), - }; - } - - return { - error: 'Quota request failed', - errorCode: 'quota_request_failed', - errorDetail: detail, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'low', - accessState: 'unknown', - capacityState: 'unknown', - }), - }; -} - -function mergeAntigravityTierEvidence( - entitlement: ProviderEntitlementEvidence | undefined, - tier: AccountTier, - rawTierId: string | null, - rawTierLabel: string | null -): ProviderEntitlementEvidence | undefined { - if (tier === 'unknown' && !entitlement) { - return undefined; - } - - return buildProviderEntitlementEvidence({ - normalizedTier: tier, - rawTierId, - rawTierLabel, - source: rawTierId ? 'runtime_api' : (entitlement?.source ?? 'runtime_inference'), - confidence: rawTierId ? 'high' : (entitlement?.confidence ?? 'medium'), - accessState: entitlement?.accessState ?? 'unknown', - capacityState: entitlement?.capacityState ?? 'unknown', - notes: entitlement?.notes ?? null, - }); -} - -async function readManagedResponse( - response: Response, - viaManagement: boolean -): Promise { - const bodyText = await response.text(); - return { - status: response.status, - bodyText, - json: safeParseJson(bodyText), - viaManagement, - }; -} - -function isAntigravityAuthFileForAccount(file: ManagementAuthFile, accountId: string): boolean { - const provider = (file.provider || file.type || '').trim().toLowerCase(); - if (provider !== 'antigravity' && provider !== 'agy') { - return false; - } - - const normalizedAccount = accountId.trim().toLowerCase(); - const normalizedEmail = file.email?.trim().toLowerCase(); - if (normalizedEmail && normalizedEmail === normalizedAccount) { - return true; - } - - const normalizedName = file.name?.trim().toLowerCase(); - if (!normalizedName) { - return false; - } - - const sanitizedAccount = sanitizeEmail(accountId).toLowerCase(); - return ( - normalizedName === `antigravity-${sanitizedAccount}.json` || - normalizedName === `agy-${sanitizedAccount}.json` - ); -} - -async function findManagedAntigravityAuthIndex(accountId: string): Promise { - const target = getProxyTarget(); - const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); - - try { - const response = await fetch(buildProxyUrl(target, '/v0/management/auth-files'), { - signal: controller.signal, - headers: buildManagementHeaders(target), - }); - clearTimeout(timeoutId); - - if (!response.ok) { - return null; - } - - const data = (await response.json()) as { files?: ManagementAuthFile[] }; - const match = data.files?.find((file) => isAntigravityAuthFileForAccount(file, accountId)); - return match?.auth_index ?? null; - } catch { - clearTimeout(timeoutId); - return null; - } -} - -async function performManagedAntigravityRequest( - accountId: string, - url: string, - headers: Record, - body: string -): Promise { - const authIndex = await findManagedAntigravityAuthIndex(accountId); - if (authIndex === null || authIndex === undefined) { - return null; - } - - const target = getProxyTarget(); - const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); - - try { - const response = await fetch(buildProxyUrl(target, '/v0/management/api-call'), { - method: 'POST', - signal: controller.signal, - headers: buildManagementHeaders(target, { - 'Content-Type': 'application/json', - }), - body: JSON.stringify({ - auth_index: authIndex, - method: 'POST', - url, - header: { - ...headers, - Authorization: 'Bearer $TOKEN$', - }, - data: body, - }), - }); - clearTimeout(timeoutId); - - if (!response.ok) { - return null; - } - - const apiResponse = (await response.json()) as ManagementApiCallResponse; - const bodyText = typeof apiResponse.body === 'string' ? apiResponse.body : ''; - return { - status: typeof apiResponse.status_code === 'number' ? apiResponse.status_code : 500, - bodyText, - json: safeParseJson(bodyText), - viaManagement: true, - }; - } catch { - clearTimeout(timeoutId); - return null; - } -} - -async function performAntigravityRequest( - accountId: string, - accessToken: string, - url: string, - headers: Record, - body: string -): Promise { - const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); - - try { - const response = await fetch(url, { - method: 'POST', - signal: controller.signal, - headers: { - ...headers, - Authorization: `Bearer ${accessToken}`, - }, - body, - }); - clearTimeout(timeoutId); - - const directResult = await readManagedResponse(response, false); - if (directResult.status !== 401) { - return directResult; - } - - const managedResult = await performManagedAntigravityRequest(accountId, url, headers, body); - return managedResult ?? directResult; - } catch (err) { - clearTimeout(timeoutId); - if (err instanceof Error && err.name === 'AbortError') { - return { - status: 408, - bodyText: '', - json: null, - viaManagement: false, - }; - } - - const message = err instanceof Error ? err.message : 'Unknown error'; - return { - status: 503, - bodyText: message, - json: null, - viaManagement: false, - }; - } -} - -async function performAntigravityRequestWithBaseUrlFallback( - accountId: string, - accessToken: string, - baseUrls: readonly string[], - apiPath: string, - headers: Record, - body: string -): Promise { - let lastResponse: ManagedResponse | null = null; - - for (const baseUrl of baseUrls) { - const response = await performAntigravityRequest( - accountId, - accessToken, - `${baseUrl}/${apiPath}`, - headers, - body - ); - if (response.status >= 200 && response.status < 300) { - return response; - } - lastResponse = response; - } - - return ( - lastResponse ?? { - status: 503, - bodyText: 'No Antigravity API endpoint available', - json: null, - viaManagement: false, - } - ); -} - -/** - * Read auth data from auth file (access token, project_id, expiry status) - */ -function readAuthData(provider: CLIProxyProvider, accountId: string): AuthData | null { - // Check both active and paused auth directories (quota needed for paused accounts too) - const authDirs = [getAuthDir(), getPausedDir()]; - - // Sanitize accountId (email) to match auth file naming: @ and . → _ - const sanitizedId = sanitizeEmail(accountId); - const prefix = provider === 'agy' ? 'antigravity-' : `${provider}-`; - const expectedFile = `${prefix}${sanitizedId}.json`; - - for (const authDir of authDirs) { - if (!fs.existsSync(authDir)) continue; - - const filePath = path.join(authDir, expectedFile); - - // Direct file access (most common case) - if (fs.existsSync(filePath)) { - try { - const content = fs.readFileSync(filePath, 'utf-8'); - const data = JSON.parse(content) as AntigravityAuthFile; - if (!data.access_token) continue; - return { - accessToken: data.access_token, - refreshToken: data.refresh_token || null, - projectId: data.project_id || null, - isExpired: isTokenExpired(data.expired), - expiresAt: data.expired || null, - }; - } catch { - continue; - } - } - - // Fallback: scan directory for matching email in file content - const files = fs.readdirSync(authDir); - for (const file of files) { - if (file.startsWith(prefix) && file.endsWith('.json')) { - const candidatePath = path.join(authDir, file); - try { - const content = fs.readFileSync(candidatePath, 'utf-8'); - const data = JSON.parse(content) as AntigravityAuthFile; - // Match by email field inside the auth file - if (data.email === accountId && data.access_token) { - return { - accessToken: data.access_token, - refreshToken: data.refresh_token || null, - projectId: data.project_id || null, - isExpired: isTokenExpired(data.expired), - expiresAt: data.expired || null, - }; - } - } catch { - continue; - } - } - } - } - - return null; -} - -/** - * Map tier ID string to AccountTier type - * API returns: "g1-ultra-tier", "g1-pro-tier", "standard-tier", etc. - * Priority: ultra > pro > free - */ -/** - * Get project ID and tier via loadCodeAssist endpoint - * Uses paidTier.id for accurate tier detection (g1-ultra-tier, g1-pro-tier) - */ -async function getProjectId(accountId: string, accessToken: string): Promise { - const body = JSON.stringify({ - metadata: { - ide_name: 'antigravity', - ide_type: 'ANTIGRAVITY', - ide_version: '1.21.9', - }, - }); - const response = await performAntigravityRequestWithBaseUrlFallback( - accountId, - accessToken, - ANTIGRAVITY_LOADCODEASSIST_BASE_URLS, - `${ANTIGRAVITY_API_VERSION}:loadCodeAssist`, - LOADCODEASSIST_HEADERS, - body - ); - - if (response.status < 200 || response.status >= 300) { - return { - projectId: null, - ...buildAntigravityFailure(response.status, response.bodyText), - }; - } - - const data = response.json as LoadCodeAssistResponse | null; - if (!data) { - return { - projectId: null, - error: 'Invalid quota response from provider', - errorCode: 'provider_unavailable', - retryable: true, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'low', - accessState: 'unknown', - capacityState: 'temporarily_unavailable', - notes: 'Provider returned a 2xx response with an empty or invalid project payload.', - }), - }; - } - - // Extract project ID from response - let projectId: string | undefined; - if (typeof data.cloudaicompanionProject === 'string') { - projectId = data.cloudaicompanionProject; - } else if (typeof data.cloudaicompanionProject === 'object') { - projectId = data.cloudaicompanionProject?.id; - } - - if (!projectId?.trim()) { - return { - projectId: null, - error: 'Sign in to Antigravity app to activate quota.', - errorCode: 'account_unprovisioned', - actionHint: 'Complete sign-in in the Antigravity app, then retry quota refresh.', - isUnprovisioned: true, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'medium', - accessState: 'unknown', - capacityState: 'unknown', - notes: 'Project provisioning is incomplete for this account.', - }), - }; - } - - // Extract tier - paidTier reflects actual subscription status, takes priority - const rawTierId = (data.paidTier?.id || data.currentTier?.id || '').trim() || null; - const tier = normalizeProviderTierId(rawTierId); - - return { - projectId: projectId.trim(), - tier, - rawTierId, - rawTierLabel: getProviderTierLabel(rawTierId), - }; -} - -/** - * Fetch available models with quota info - * Note: projectId is kept for potential future use but not sent in body - * (CLIProxyAPI sends empty {} body for this endpoint) - */ -async function fetchAvailableModels( - accountId: string, - accessToken: string, - _projectId: string -): Promise { - const url = `${ANTIGRAVITY_API_BASE}/${ANTIGRAVITY_API_VERSION}:fetchAvailableModels`; - const response = await performAntigravityRequest( - accountId, - accessToken, - url, - FETCHMODELS_HEADERS, - JSON.stringify({}) - ); - - if (response.status < 200 || response.status >= 300) { - return { - success: false, - models: [], - lastUpdated: Date.now(), - ...buildAntigravityFailure(response.status, response.bodyText), - }; - } - - const data = response.json as FetchAvailableModelsResponse | null; - if (!data) { - return { - success: false, - models: [], - lastUpdated: Date.now(), - error: 'Invalid quota response from provider', - errorCode: 'provider_unavailable', - retryable: true, - entitlement: buildProviderEntitlementEvidence({ - normalizedTier: 'unknown', - source: 'runtime_inference', - confidence: 'low', - accessState: 'unknown', - capacityState: 'temporarily_unavailable', - notes: 'Provider returned a 2xx response with an empty or invalid quota payload.', - }), - }; - } - - const models: ModelQuota[] = []; - - if (data.models && typeof data.models === 'object') { - for (const [modelId, modelData] of Object.entries(data.models)) { - const quotaInfo = modelData.quotaInfo || modelData.quota_info; - if (!quotaInfo) continue; - - const remaining = - quotaInfo.remainingFraction ?? quotaInfo.remaining_fraction ?? quotaInfo.remaining; - const resetTime = quotaInfo.resetTime || quotaInfo.reset_time || null; - - let percentage: number; - if (typeof remaining === 'number' && isFinite(remaining)) { - percentage = Math.max(0, Math.min(100, Math.round(remaining * 100))); - } else if (resetTime) { - percentage = 0; - } else { - continue; - } - - models.push({ - name: modelId, - displayName: modelData.displayName, - percentage, - resetTime, - }); - } - } - - return { - success: true, - models, - lastUpdated: Date.now(), - }; -} - -/** - * Fetch quota for an Antigravity account * - * @param provider - Provider name (only 'agy' supported) - * @param accountId - Account identifier (email) - * @param verbose - Show detailed diagnostics - * @returns Quota result with models and percentages + * This file is a thin barrel over the focused submodules in ./quota-fetcher/. + * The public API surface (exports, signatures, types, and __testExports) is + * preserved exactly so all existing imports from 'quota-fetcher' keep working. */ -export async function fetchAccountQuota( - provider: CLIProxyProvider, - accountId: string, - verbose = false -): Promise { - if (verbose) - logger.info('quota.fetch.start', 'Fetching quota for account', { provider, accountId }); - // Only Antigravity supports quota fetching - if (provider !== 'agy') { - const error = `Quota not supported for provider: ${provider}`; - if (verbose) logger.warn('quota.fetch.unsupported_provider', error, { provider }); - // Stable machine code so callers branch on a code, not the human string. - // This is "no quota API for this provider", which is healthy — distinct - // from a transient fetch failure or an expired token. - return { - success: false, - models: [], - lastUpdated: Date.now(), - error, - errorCode: 'quota_not_supported', - }; - } +// Public types +export type { ModelQuota, QuotaResult, AllAccountsQuotaResult } from './quota-fetcher/types'; - // Read auth data from auth file (checks both active and paused directories) - const authData = readAuthData(provider, accountId); - if (!authData) { - const error = 'Auth file not found for account'; - if (verbose) logger.warn('quota.fetch.auth_missing', error, { provider, accountId }); - return { - success: false, - models: [], - lastUpdated: Date.now(), - error, - errorCode: 'auth_file_missing', - actionHint: 'Reconnect this account so CCS can read a current auth token.', - }; - } +// Public functions +export { fetchAccountQuota } from './quota-fetcher/account-quota-fetcher'; +export { fetchAllProviderQuotas } from './quota-fetcher/all-accounts-fetcher'; +export { findAvailableAccount } from './quota-fetcher/all-accounts-fetcher'; +export { readProjectIdFromAuthFile } from './quota-fetcher/auth-file-reader'; - const accessToken = authData.accessToken; - if (verbose) { - const expiryState = authData.isExpired - ? 'expired' - : authData.expiresAt - ? `expires ${authData.expiresAt}` - : 'expiry unknown'; - logger.info('quota.fetch.auth_state', `Auth token state: ${expiryState}`, { - provider, - state: expiryState, - }); - } - - // Get project ID and tier - prefer stored project ID, but always call API for tier - let projectId = authData.projectId; - let apiTier: AccountTier = 'unknown'; - let rawTierId: string | null = null; - let rawTierLabel: string | null = null; - - // Always call loadCodeAssist to get accurate tier from API. - // If the file token is stale, the helper retries through CLIProxy management auth. - const lastProjectResult = await getProjectId(accountId, accessToken); - - if (!lastProjectResult.projectId && !projectId) { - const error = lastProjectResult.error || 'Failed to retrieve project ID'; - if (verbose) - logger.warn('quota.fetch.project_lookup_failed', error, { - provider, - errorCode: lastProjectResult.errorCode, - httpStatus: lastProjectResult.httpStatus, - }); - return { - success: false, - models: [], - lastUpdated: Date.now(), - error, - errorCode: lastProjectResult.errorCode, - errorDetail: lastProjectResult.errorDetail, - actionHint: lastProjectResult.actionHint, - retryable: lastProjectResult.retryable, - httpStatus: lastProjectResult.httpStatus, - needsReauth: lastProjectResult.needsReauth, - isUnprovisioned: lastProjectResult.isUnprovisioned, - entitlement: lastProjectResult.entitlement, - isExpired: authData.isExpired, - expiresAt: authData.expiresAt || undefined, - }; - } - - // Use API project ID if available, else fallback to stored - projectId = lastProjectResult.projectId || projectId; - apiTier = lastProjectResult.tier || 'unknown'; - rawTierId = lastProjectResult.rawTierId || null; - rawTierLabel = lastProjectResult.rawTierLabel || null; - - if (verbose) - logger.info('quota.fetch.project_resolved', `Project ID: ${projectId || 'not found'}`, { - provider, - }); - - // Fetch models with quota - const result = await fetchAvailableModels(accountId, accessToken, projectId as string); - - if (verbose) - logger.info('quota.fetch.models', `Models found: ${result.models.length}`, { - provider, - count: result.models.length, - }); - result.accountId = accountId; - result.projectId = projectId || undefined; - - // Determine tier from API response only - if (result.success) { - const finalTier = apiTier !== 'unknown' ? apiTier : 'unknown'; - result.tier = finalTier; - result.entitlement = buildProviderEntitlementEvidence({ - normalizedTier: finalTier, - rawTierId, - rawTierLabel, - source: rawTierId ? 'runtime_api' : 'runtime_inference', - confidence: rawTierId ? 'high' : 'medium', - accessState: 'entitled', - capacityState: 'available', - }); - if (finalTier !== 'unknown') { - setAccountTier(provider, accountId, finalTier); - } - } else { - result.isExpired = authData.isExpired; - result.expiresAt = authData.expiresAt || undefined; - result.entitlement = mergeAntigravityTierEvidence( - result.entitlement, - apiTier, - rawTierId, - rawTierLabel - ); - } - - if (verbose && result.error) { - console.log(`[!] Error: ${result.error}`); - } - - return result; -} - -/** - * Read project ID directly from auth file without making API call - * Used for quick project ID comparison in doctor command - */ -export function readProjectIdFromAuthFile( - provider: CLIProxyProvider, - accountId: string -): string | null { - const authData = readAuthData(provider, accountId); - return authData?.projectId || null; -} - -/** Result for all accounts of a provider */ -export interface AllAccountsQuotaResult { - /** Provider name */ - provider: CLIProxyProvider; - /** Results per account */ - accounts: Array<{ - account: AccountInfo; - quota: QuotaResult; - }>; - /** Accounts grouped by project ID (for detecting shared projects) */ - projectGroups: Record; - /** Timestamp of fetch */ - lastUpdated: number; -} - -/** - * Fetch quota for all accounts of a provider - * Also detects accounts sharing same GCP project (failover won't help) - * - * @param provider - Provider name (only 'agy' supported for quota) - * @param verbose - Show detailed diagnostics - * @returns Results for all accounts with project grouping - */ -export async function fetchAllProviderQuotas( - provider: CLIProxyProvider, - verbose = false -): Promise { - const accounts = getProviderAccounts(provider); - const results: AllAccountsQuotaResult = { - provider, - accounts: [], - projectGroups: {}, - lastUpdated: Date.now(), - }; - - if (accounts.length === 0) { - return results; - } - - // Fetch quota for each account in parallel - const quotaPromises = accounts.map(async (account) => { - const quota = await fetchAccountQuota(provider, account.id, verbose); - - // Read project ID from auth file if not in quota result - let projectId = quota.projectId; - if (!projectId) { - projectId = readProjectIdFromAuthFile(provider, account.id) || undefined; - } - - return { - account, - quota: { ...quota, accountId: account.id, projectId }, - }; - }); - - const quotaResults = await Promise.all(quotaPromises); - - // Build project groups for detecting shared projects - for (const { account, quota } of quotaResults) { - results.accounts.push({ account, quota }); - - if (quota.projectId) { - if (!results.projectGroups[quota.projectId]) { - results.projectGroups[quota.projectId] = []; - } - results.projectGroups[quota.projectId].push(account.id); - } - } - - return results; -} +// Test-only exports (consumed by quota-fetcher-antigravity-failure.test.ts via +// dynamic import of ../quota-fetcher?...). buildAntigravityFailure is the only +// pure helper exercised through this surface; keep it stable. +export { buildAntigravityFailure } from './quota-fetcher/status-classifier'; +import { buildAntigravityFailure } from './quota-fetcher/status-classifier'; export const __testExports = { buildAntigravityFailure, }; - -/** - * Find available account with remaining quota - * Used by preflight check for auto-switching - * - * @param provider - Provider name - * @param excludeAccountId - Account to exclude (current exhausted account) - * @param verbose - Show detailed diagnostics - * @returns Account with available quota, or null if none available - */ -export async function findAvailableAccount( - provider: CLIProxyProvider, - excludeAccountId?: string, - verbose = false -): Promise<{ account: AccountInfo; quota: QuotaResult } | null> { - const allQuotas = await fetchAllProviderQuotas(provider, verbose); - - // Get excluded account's project ID to avoid switching to same-project accounts - const excludedProjectId = allQuotas.accounts.find((a) => a.account.id === excludeAccountId)?.quota - .projectId; - - for (const { account, quota } of allQuotas.accounts) { - // Skip excluded account - if (excludeAccountId && account.id === excludeAccountId) { - continue; - } - - // Skip failed quota fetches - if (!quota.success) { - continue; - } - - // Skip accounts sharing same GCP project (quota is pooled) - if (excludedProjectId && quota.projectId === excludedProjectId) { - continue; - } - - // Check if any model has remaining quota (> 5% to avoid edge cases) - const hasQuota = quota.models.some((m) => m.percentage > 5); - if (hasQuota) { - return { account, quota }; - } - } - - return null; -} diff --git a/src/cliproxy/quota/quota-fetcher/account-quota-fetcher.ts b/src/cliproxy/quota/quota-fetcher/account-quota-fetcher.ts new file mode 100644 index 00000000..b72c72f4 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/account-quota-fetcher.ts @@ -0,0 +1,175 @@ +/** + * fetchAccountQuota: top-level Antigravity account quota orchestrator. + * + * Reads the local auth file, calls loadCodeAssist (project + tier), then + * fetchAvailableModels. Merges the results into a QuotaResult, attaching + * entitlement evidence and persisting the resolved tier back to the account + * manager. Preserves the structured createLogger('cliproxy:quota:fetcher') + * logging from P3 (quota.fetch.start / auth_state / project_resolved / models). + */ + +import type { CLIProxyProvider } from '../../types'; +import type { AccountTier } from '../../accounts/account-manager'; +import { setAccountTier } from '../../accounts/account-manager'; +import { buildProviderEntitlementEvidence } from '../../auth/provider-entitlement-evidence'; +import { createLogger } from '../../../services/logging'; + +import { readAuthData } from './auth-file-reader'; +import { fetchAvailableModels } from './available-models-fetcher'; +import { getProjectId } from './project-lookup'; +import { mergeAntigravityTierEvidence } from './status-classifier'; +import type { QuotaResult } from './types'; + +const logger = createLogger('cliproxy:quota:fetcher'); + +/** + * Fetch quota for an Antigravity account. + * + * @param provider - Provider name (only 'agy' supported) + * @param accountId - Account identifier (email) + * @param verbose - Show detailed diagnostics + * @returns Quota result with models and percentages + */ +export async function fetchAccountQuota( + provider: CLIProxyProvider, + accountId: string, + verbose = false +): Promise { + if (verbose) + logger.info('quota.fetch.start', 'Fetching quota for account', { provider, accountId }); + + // Only Antigravity supports quota fetching + if (provider !== 'agy') { + const error = `Quota not supported for provider: ${provider}`; + if (verbose) logger.warn('quota.fetch.unsupported_provider', error, { provider }); + // Stable machine code so callers branch on a code, not the human string. + // This is "no quota API for this provider", which is healthy — distinct + // from a transient fetch failure or an expired token. + return { + success: false, + models: [], + lastUpdated: Date.now(), + error, + errorCode: 'quota_not_supported', + }; + } + + // Read auth data from auth file (checks both active and paused directories) + const authData = readAuthData(provider, accountId); + if (!authData) { + const error = 'Auth file not found for account'; + if (verbose) logger.warn('quota.fetch.auth_missing', error, { provider, accountId }); + return { + success: false, + models: [], + lastUpdated: Date.now(), + error, + errorCode: 'auth_file_missing', + actionHint: 'Reconnect this account so CCS can read a current auth token.', + }; + } + + const accessToken = authData.accessToken; + if (verbose) { + const expiryState = authData.isExpired + ? 'expired' + : authData.expiresAt + ? `expires ${authData.expiresAt}` + : 'expiry unknown'; + logger.info('quota.fetch.auth_state', `Auth token state: ${expiryState}`, { + provider, + state: expiryState, + }); + } + + // Get project ID and tier - prefer stored project ID, but always call API for tier + let projectId = authData.projectId; + let apiTier: AccountTier = 'unknown'; + let rawTierId: string | null = null; + let rawTierLabel: string | null = null; + + // Always call loadCodeAssist to get accurate tier from API. + // If the file token is stale, the helper retries through CLIProxy management auth. + const lastProjectResult = await getProjectId(accountId, accessToken); + + if (!lastProjectResult.projectId && !projectId) { + const error = lastProjectResult.error || 'Failed to retrieve project ID'; + if (verbose) + logger.warn('quota.fetch.project_lookup_failed', error, { + provider, + errorCode: lastProjectResult.errorCode, + httpStatus: lastProjectResult.httpStatus, + }); + return { + success: false, + models: [], + lastUpdated: Date.now(), + error, + errorCode: lastProjectResult.errorCode, + errorDetail: lastProjectResult.errorDetail, + actionHint: lastProjectResult.actionHint, + retryable: lastProjectResult.retryable, + httpStatus: lastProjectResult.httpStatus, + needsReauth: lastProjectResult.needsReauth, + isUnprovisioned: lastProjectResult.isUnprovisioned, + entitlement: lastProjectResult.entitlement, + isExpired: authData.isExpired, + expiresAt: authData.expiresAt || undefined, + }; + } + + // Use API project ID if available, else fallback to stored + projectId = lastProjectResult.projectId || projectId; + apiTier = lastProjectResult.tier || 'unknown'; + rawTierId = lastProjectResult.rawTierId || null; + rawTierLabel = lastProjectResult.rawTierLabel || null; + + if (verbose) + logger.info('quota.fetch.project_resolved', `Project ID: ${projectId || 'not found'}`, { + provider, + }); + + // Fetch models with quota + const result = await fetchAvailableModels(accountId, accessToken, projectId as string); + + if (verbose) + logger.info('quota.fetch.models', `Models found: ${result.models.length}`, { + provider, + count: result.models.length, + }); + result.accountId = accountId; + result.projectId = projectId || undefined; + + // Determine tier from API response only + if (result.success) { + const finalTier = apiTier !== 'unknown' ? apiTier : 'unknown'; + result.tier = finalTier; + result.entitlement = buildProviderEntitlementEvidence({ + normalizedTier: finalTier, + rawTierId, + rawTierLabel, + source: rawTierId ? 'runtime_api' : 'runtime_inference', + confidence: rawTierId ? 'high' : 'medium', + accessState: 'entitled', + capacityState: 'available', + }); + if (finalTier !== 'unknown') { + setAccountTier(provider, accountId, finalTier); + } + } else { + result.isExpired = authData.isExpired; + result.expiresAt = authData.expiresAt || undefined; + result.entitlement = mergeAntigravityTierEvidence( + result.entitlement, + apiTier, + rawTierId, + rawTierLabel + ); + } + + if (verbose && result.error) { + console.log(`[!] Error: ${result.error}`); + } + + return result; +} diff --git a/src/cliproxy/quota/quota-fetcher/all-accounts-fetcher.ts b/src/cliproxy/quota/quota-fetcher/all-accounts-fetcher.ts new file mode 100644 index 00000000..7da48dcb --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/all-accounts-fetcher.ts @@ -0,0 +1,119 @@ +/** + * fetchAllProviderQuotas and findAvailableAccount. + * + * fetchAllProviderQuotas fans quota fetches out across all accounts of a + * provider in parallel and groups them by GCP project id (accounts that share + * a project pool quota together, so failover between them won't help). + * findAvailableAccount wraps that to pick the first account that still has + * remaining quota (used by the auto-switch preflight check). + */ + +import type { CLIProxyProvider } from '../../types'; +import { getProviderAccounts, type AccountInfo } from '../../accounts/account-manager'; + +import { fetchAccountQuota } from './account-quota-fetcher'; +import { readProjectIdFromAuthFile } from './auth-file-reader'; +import type { AllAccountsQuotaResult, QuotaResult } from './types'; + +/** + * Fetch quota for all accounts of a provider. + * Also detects accounts sharing the same GCP project (failover won't help). + * + * @param provider - Provider name (only 'agy' supported for quota) + * @param verbose - Show detailed diagnostics + * @returns Results for all accounts with project grouping + */ +export async function fetchAllProviderQuotas( + provider: CLIProxyProvider, + verbose = false +): Promise { + const accounts = getProviderAccounts(provider); + const results: AllAccountsQuotaResult = { + provider, + accounts: [], + projectGroups: {}, + lastUpdated: Date.now(), + }; + + if (accounts.length === 0) { + return results; + } + + // Fetch quota for each account in parallel + const quotaPromises = accounts.map(async (account) => { + const quota = await fetchAccountQuota(provider, account.id, verbose); + + // Read project ID from auth file if not in quota result + let projectId = quota.projectId; + if (!projectId) { + projectId = readProjectIdFromAuthFile(provider, account.id) || undefined; + } + + return { + account, + quota: { ...quota, accountId: account.id, projectId }, + }; + }); + + const quotaResults = await Promise.all(quotaPromises); + + // Build project groups for detecting shared projects + for (const { account, quota } of quotaResults) { + results.accounts.push({ account, quota }); + + if (quota.projectId) { + if (!results.projectGroups[quota.projectId]) { + results.projectGroups[quota.projectId] = []; + } + results.projectGroups[quota.projectId].push(account.id); + } + } + + return results; +} + +/** + * Find an available account with remaining quota. + * Used by preflight check for auto-switching. + * + * @param provider - Provider name + * @param excludeAccountId - Account to exclude (current exhausted account) + * @param verbose - Show detailed diagnostics + * @returns Account with available quota, or null if none available + */ +export async function findAvailableAccount( + provider: CLIProxyProvider, + excludeAccountId?: string, + verbose = false +): Promise<{ account: AccountInfo; quota: QuotaResult } | null> { + const allQuotas = await fetchAllProviderQuotas(provider, verbose); + + // Get excluded account's project ID to avoid switching to same-project accounts + const excludedProjectId = allQuotas.accounts.find((a) => a.account.id === excludeAccountId)?.quota + .projectId; + + for (const { account, quota } of allQuotas.accounts) { + // Skip excluded account + if (excludeAccountId && account.id === excludeAccountId) { + continue; + } + + // Skip failed quota fetches + if (!quota.success) { + continue; + } + + // Skip accounts sharing the same GCP project (quota is pooled) + if (excludedProjectId && quota.projectId === excludedProjectId) { + continue; + } + + // Check if any model has remaining quota (> 5% to avoid edge cases) + const hasQuota = quota.models.some((m) => m.percentage > 5); + if (hasQuota) { + return { account, quota }; + } + } + + return null; +} diff --git a/src/cliproxy/quota/quota-fetcher/auth-file-reader.ts b/src/cliproxy/quota/quota-fetcher/auth-file-reader.ts new file mode 100644 index 00000000..86928043 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/auth-file-reader.ts @@ -0,0 +1,93 @@ +/** + * Auth file reader for Antigravity quota fetching. + * + * Reads the local Antigravity auth file (active or paused directory) and + * extracts the access token, refresh token, project id, and expiry state. + * Falls back to scanning the directory and matching by the embedded email + * field when the canonical sanitized filename is not present. + */ + +import * as fs from 'node:fs'; +import * as path from 'node:path'; + +import { getAuthDir } from '../../config/config-generator'; +import type { CLIProxyProvider } from '../../types'; +import { isTokenExpired, sanitizeEmail } from '../../auth/auth-utils'; +import { getPausedDir } from '../../accounts/account-manager'; +import type { AntigravityAuthFile, AuthData } from './types'; + +/** + * Read auth data from the auth file (access token, project_id, expiry state). + * Checks both active and paused auth directories (quota is needed for paused + * accounts too). + */ +export function readAuthData(provider: CLIProxyProvider, accountId: string): AuthData | null { + const authDirs = [getAuthDir(), getPausedDir()]; + + // Sanitize accountId (email) to match auth file naming: @ and . → _ + const sanitizedId = sanitizeEmail(accountId); + const prefix = provider === 'agy' ? 'antigravity-' : `${provider}-`; + const expectedFile = `${prefix}${sanitizedId}.json`; + + for (const authDir of authDirs) { + if (!fs.existsSync(authDir)) continue; + + const filePath = path.join(authDir, expectedFile); + + // Direct file access (most common case) + if (fs.existsSync(filePath)) { + try { + const content = fs.readFileSync(filePath, 'utf-8'); + const data = JSON.parse(content) as AntigravityAuthFile; + if (!data.access_token) continue; + return { + accessToken: data.access_token, + refreshToken: data.refresh_token || null, + projectId: data.project_id || null, + isExpired: isTokenExpired(data.expired), + expiresAt: data.expired || null, + }; + } catch { + continue; + } + } + + // Fallback: scan directory for matching email in file content + const files = fs.readdirSync(authDir); + for (const file of files) { + if (file.startsWith(prefix) && file.endsWith('.json')) { + const candidatePath = path.join(authDir, file); + try { + const content = fs.readFileSync(candidatePath, 'utf-8'); + const data = JSON.parse(content) as AntigravityAuthFile; + // Match by email field inside the auth file + if (data.email === accountId && data.access_token) { + return { + accessToken: data.access_token, + refreshToken: data.refresh_token || null, + projectId: data.project_id || null, + isExpired: isTokenExpired(data.expired), + expiresAt: data.expired || null, + }; + } + } catch { + continue; + } + } + } + } + + return null; +} + +/** + * Read project ID directly from auth file without making an API call. + * Used for quick project ID comparison in the doctor command. + */ +export function readProjectIdFromAuthFile( + provider: CLIProxyProvider, + accountId: string +): string | null { + const authData = readAuthData(provider, accountId); + return authData?.projectId || null; +} diff --git a/src/cliproxy/quota/quota-fetcher/available-models-fetcher.ts b/src/cliproxy/quota/quota-fetcher/available-models-fetcher.ts new file mode 100644 index 00000000..1bc80b19 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/available-models-fetcher.ts @@ -0,0 +1,99 @@ +/** + * fetchAvailableModels call for Antigravity quota. + * + * Fetches the model -> remaining-fraction map from the Cloud Code internal + * API and projects it into ModelQuota[] percentages (0-100). The projectId + * is intentionally NOT sent in the body (CLIProxyAPI sends an empty {} body + * for this endpoint); it is accepted only for symmetry with the project + * lookup flow. + */ + +import { buildProviderEntitlementEvidence } from '../../auth/provider-entitlement-evidence'; +import { ANTIGRAVITY_API_BASE, ANTIGRAVITY_API_VERSION, FETCHMODELS_HEADERS } from './constants'; +import { performAntigravityRequest } from './http-client'; +import { buildAntigravityFailure } from './status-classifier'; +import type { FetchAvailableModelsResponse, ModelQuota, QuotaResult } from './types'; + +/** + * Fetch available models with quota info. + * Note: projectId is kept for potential future use but not sent in body + * (CLIProxyAPI sends empty {} body for this endpoint). + */ +export async function fetchAvailableModels( + accountId: string, + accessToken: string, + _projectId: string +): Promise { + const url = `${ANTIGRAVITY_API_BASE}/${ANTIGRAVITY_API_VERSION}:fetchAvailableModels`; + const response = await performAntigravityRequest( + accountId, + accessToken, + url, + FETCHMODELS_HEADERS, + JSON.stringify({}) + ); + + if (response.status < 200 || response.status >= 300) { + return { + success: false, + models: [], + lastUpdated: Date.now(), + ...buildAntigravityFailure(response.status, response.bodyText), + }; + } + + const data = response.json as FetchAvailableModelsResponse | null; + if (!data) { + return { + success: false, + models: [], + lastUpdated: Date.now(), + error: 'Invalid quota response from provider', + errorCode: 'provider_unavailable', + retryable: true, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'low', + accessState: 'unknown', + capacityState: 'temporarily_unavailable', + notes: 'Provider returned a 2xx response with an empty or invalid quota payload.', + }), + }; + } + + const models: ModelQuota[] = []; + + if (data.models && typeof data.models === 'object') { + for (const [modelId, modelData] of Object.entries(data.models)) { + const quotaInfo = modelData.quotaInfo || modelData.quota_info; + if (!quotaInfo) continue; + + const remaining = + quotaInfo.remainingFraction ?? quotaInfo.remaining_fraction ?? quotaInfo.remaining; + const resetTime = quotaInfo.resetTime || quotaInfo.reset_time || null; + + let percentage: number; + if (typeof remaining === 'number' && isFinite(remaining)) { + percentage = Math.max(0, Math.min(100, Math.round(remaining * 100))); + } else if (resetTime) { + percentage = 0; + } else { + continue; + } + + models.push({ + name: modelId, + displayName: modelData.displayName, + percentage, + resetTime, + }); + } + } + + return { + success: true, + models, + lastUpdated: Date.now(), + }; +} diff --git a/src/cliproxy/quota/quota-fetcher/constants.ts b/src/cliproxy/quota/quota-fetcher/constants.ts new file mode 100644 index 00000000..bf377fbf --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/constants.ts @@ -0,0 +1,30 @@ +/** + * Constants for the Antigravity quota fetcher. + * + * Google Cloud Code internal API endpoints, fixed headers used by the + * CLIProxyAPIPlus control-plane requests, and the shared timeout applied to + * every Antigravity management API call. + */ + +/** Google Cloud Code API endpoints */ +export const ANTIGRAVITY_DAILY_API_BASE = 'https://daily-cloudcode-pa.googleapis.com'; +export const ANTIGRAVITY_API_BASE = 'https://cloudcode-pa.googleapis.com'; +export const ANTIGRAVITY_API_VERSION = 'v1internal'; +export const ANTIGRAVITY_LOADCODEASSIST_BASE_URLS = [ + ANTIGRAVITY_DAILY_API_BASE, + ANTIGRAVITY_API_BASE, +] as const; +export const MANAGEMENT_API_TIMEOUT_MS = 5000; + +/** Headers for loadCodeAssist (matches current CLIProxyAPIPlus control-plane requests) */ +export const LOADCODEASSIST_HEADERS = { + 'Content-Type': 'application/json', + 'User-Agent': 'antigravity/1.21.9 darwin/arm64 google-api-nodejs-client/10.3.0', + 'X-Goog-Api-Client': 'gl-node/22.21.1', +}; + +/** Headers for fetchAvailableModels (matches CLIProxyAPI antigravity_executor.go) */ +export const FETCHMODELS_HEADERS = { + 'Content-Type': 'application/json', + 'User-Agent': 'antigravity/1.104.0 darwin/arm64', +}; diff --git a/src/cliproxy/quota/quota-fetcher/http-client.ts b/src/cliproxy/quota/quota-fetcher/http-client.ts new file mode 100644 index 00000000..e0b48dac --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/http-client.ts @@ -0,0 +1,248 @@ +/** + * HTTP transport for Antigravity quota requests. + * + * Wraps fetch() with three fallback strategies: + * 1. Direct call with the local access token. + * 2. If that returns 401 (token rejected), retry through CLIProxy management + * auth using the proxy's stored token. + * 3. For loadCodeAssist, fall back across the daily then prod Cloud Code hosts. + * + * Every call is bounded by MANAGEMENT_API_TIMEOUT_MS via an AbortController. + * Network errors become synthetic 503 responses; abort timeouts become 408. + */ + +import { sanitizeEmail } from '../../auth/auth-utils'; +import { + buildManagementHeaders, + buildProxyUrl, + getProxyTarget, +} from '../../proxy/proxy-target-resolver'; +import { MANAGEMENT_API_TIMEOUT_MS } from './constants'; +import type { ManagedResponse, ManagementApiCallResponse, ManagementAuthFile } from './types'; + +/** Best-effort JSON.parse; returns null on failure. */ +export function safeParseJson(bodyText: string): unknown { + try { + return JSON.parse(bodyText); + } catch { + return null; + } +} + +/** Read the response body once and return a normalized ManagedResponse. */ +async function readManagedResponse( + response: Response, + viaManagement: boolean +): Promise { + const bodyText = await response.text(); + return { + status: response.status, + bodyText, + json: safeParseJson(bodyText), + viaManagement, + }; +} + +/** Does this management-API auth file belong to the given Antigravity account? */ +function isAntigravityAuthFileForAccount(file: ManagementAuthFile, accountId: string): boolean { + const provider = (file.provider || file.type || '').trim().toLowerCase(); + if (provider !== 'antigravity' && provider !== 'agy') { + return false; + } + + const normalizedAccount = accountId.trim().toLowerCase(); + const normalizedEmail = file.email?.trim().toLowerCase(); + if (normalizedEmail && normalizedEmail === normalizedAccount) { + return true; + } + + const normalizedName = file.name?.trim().toLowerCase(); + if (!normalizedName) { + return false; + } + + const sanitizedAccount = sanitizeEmail(accountId).toLowerCase(); + return ( + normalizedName === `antigravity-${sanitizedAccount}.json` || + normalizedName === `agy-${sanitizedAccount}.json` + ); +} + +/** Ask CLIProxy management API for the auth_index of the Antigravity account. */ +async function findManagedAntigravityAuthIndex(accountId: string): Promise { + const target = getProxyTarget(); + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); + + try { + const response = await fetch(buildProxyUrl(target, '/v0/management/auth-files'), { + signal: controller.signal, + headers: buildManagementHeaders(target), + }); + clearTimeout(timeoutId); + + if (!response.ok) { + return null; + } + + const data = (await response.json()) as { files?: ManagementAuthFile[] }; + const match = data.files?.find((file) => isAntigravityAuthFileForAccount(file, accountId)); + return match?.auth_index ?? null; + } catch { + clearTimeout(timeoutId); + return null; + } +} + +/** + * Run a request through the CLIProxy management api-call endpoint using the + * proxy's stored token (substituted server-side as $TOKEN$). Returns null if + * the proxy can't handle the request. + */ +async function performManagedAntigravityRequest( + accountId: string, + url: string, + headers: Record, + body: string +): Promise { + const authIndex = await findManagedAntigravityAuthIndex(accountId); + if (authIndex === null || authIndex === undefined) { + return null; + } + + const target = getProxyTarget(); + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); + + try { + const response = await fetch(buildProxyUrl(target, '/v0/management/api-call'), { + method: 'POST', + signal: controller.signal, + headers: buildManagementHeaders(target, { + 'Content-Type': 'application/json', + }), + body: JSON.stringify({ + auth_index: authIndex, + method: 'POST', + url, + header: { + ...headers, + Authorization: 'Bearer $TOKEN$', + }, + data: body, + }), + }); + clearTimeout(timeoutId); + + if (!response.ok) { + return null; + } + + const apiResponse = (await response.json()) as ManagementApiCallResponse; + const bodyText = typeof apiResponse.body === 'string' ? apiResponse.body : ''; + return { + status: typeof apiResponse.status_code === 'number' ? apiResponse.status_code : 500, + bodyText, + json: safeParseJson(bodyText), + viaManagement: true, + }; + } catch { + clearTimeout(timeoutId); + return null; + } +} + +/** + * Perform a single Antigravity POST. Tries direct with the local access token + * first; on 401, retries through CLIProxy management auth. Network errors map + * to synthetic 503 responses so the caller's status-based classifier still + * works; abort timeouts map to 408. + */ +export async function performAntigravityRequest( + accountId: string, + accessToken: string, + url: string, + headers: Record, + body: string +): Promise { + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); + + try { + const response = await fetch(url, { + method: 'POST', + signal: controller.signal, + headers: { + ...headers, + Authorization: `Bearer ${accessToken}`, + }, + body, + }); + clearTimeout(timeoutId); + + const directResult = await readManagedResponse(response, false); + if (directResult.status !== 401) { + return directResult; + } + + const managedResult = await performManagedAntigravityRequest(accountId, url, headers, body); + return managedResult ?? directResult; + } catch (err) { + clearTimeout(timeoutId); + if (err instanceof Error && err.name === 'AbortError') { + return { + status: 408, + bodyText: '', + json: null, + viaManagement: false, + }; + } + + const message = err instanceof Error ? err.message : 'Unknown error'; + return { + status: 503, + bodyText: message, + json: null, + viaManagement: false, + }; + } +} + +/** + * Run a request against each base URL in order, returning the first 2xx + * response. If none succeed, return the last failure (or a synthetic 503 if + * no URLs were attempted). + */ +export async function performAntigravityRequestWithBaseUrlFallback( + accountId: string, + accessToken: string, + baseUrls: readonly string[], + apiPath: string, + headers: Record, + body: string +): Promise { + let lastResponse: ManagedResponse | null = null; + + for (const baseUrl of baseUrls) { + const response = await performAntigravityRequest( + accountId, + accessToken, + `${baseUrl}/${apiPath}`, + headers, + body + ); + if (response.status >= 200 && response.status < 300) { + return response; + } + lastResponse = response; + } + + return ( + lastResponse ?? { + status: 503, + bodyText: 'No Antigravity API endpoint available', + json: null, + viaManagement: false, + } + ); +} diff --git a/src/cliproxy/quota/quota-fetcher/project-lookup.ts b/src/cliproxy/quota/quota-fetcher/project-lookup.ts new file mode 100644 index 00000000..b2c51c7b --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/project-lookup.ts @@ -0,0 +1,110 @@ +/** + * Antigravity loadCodeAssist project + tier lookup. + * + * Calls loadCodeAssist against the daily then prod Cloud Code hosts to resolve + * the GCP project id and the account tier (paidTier.id takes priority over + * currentTier.id). Returns a structured ProjectLookupResult that the top-level + * fetchAccountQuota merges into a QuotaResult. + */ + +import { + buildProviderEntitlementEvidence, + getProviderTierLabel, + normalizeProviderTierId, +} from '../../auth/provider-entitlement-evidence'; +import { + ANTIGRAVITY_API_VERSION, + ANTIGRAVITY_LOADCODEASSIST_BASE_URLS, + LOADCODEASSIST_HEADERS, +} from './constants'; +import { performAntigravityRequestWithBaseUrlFallback } from './http-client'; +import { buildAntigravityFailure } from './status-classifier'; +import type { LoadCodeAssistResponse, ProjectLookupResult } from './types'; + +/** + * Get project ID and tier via loadCodeAssist endpoint. + * Uses paidTier.id for accurate tier detection (g1-ultra-tier, g1-pro-tier). + * Falls back across the daily then prod Cloud Code hosts. + */ +export async function getProjectId( + accountId: string, + accessToken: string +): Promise { + const body = JSON.stringify({ + metadata: { + ide_name: 'antigravity', + ide_type: 'ANTIGRAVITY', + ide_version: '1.21.9', + }, + }); + const response = await performAntigravityRequestWithBaseUrlFallback( + accountId, + accessToken, + ANTIGRAVITY_LOADCODEASSIST_BASE_URLS, + `${ANTIGRAVITY_API_VERSION}:loadCodeAssist`, + LOADCODEASSIST_HEADERS, + body + ); + + if (response.status < 200 || response.status >= 300) { + return { + projectId: null, + ...buildAntigravityFailure(response.status, response.bodyText), + }; + } + + const data = response.json as LoadCodeAssistResponse | null; + if (!data) { + return { + projectId: null, + error: 'Invalid quota response from provider', + errorCode: 'provider_unavailable', + retryable: true, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'low', + accessState: 'unknown', + capacityState: 'temporarily_unavailable', + notes: 'Provider returned a 2xx response with an empty or invalid project payload.', + }), + }; + } + + // Extract project ID from response + let projectId: string | undefined; + if (typeof data.cloudaicompanionProject === 'string') { + projectId = data.cloudaicompanionProject; + } else if (typeof data.cloudaicompanionProject === 'object') { + projectId = data.cloudaicompanionProject?.id; + } + + if (!projectId?.trim()) { + return { + projectId: null, + error: 'Sign in to Antigravity app to activate quota.', + errorCode: 'account_unprovisioned', + actionHint: 'Complete sign-in in the Antigravity app, then retry quota refresh.', + isUnprovisioned: true, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'medium', + accessState: 'unknown', + capacityState: 'unknown', + notes: 'Project provisioning is incomplete for this account.', + }), + }; + } + + // Extract tier - paidTier reflects actual subscription status, takes priority + const rawTierId = (data.paidTier?.id || data.currentTier?.id || '').trim() || null; + const tier = normalizeProviderTierId(rawTierId); + + return { + projectId: projectId.trim(), + tier, + rawTierId, + rawTierLabel: getProviderTierLabel(rawTierId), + }; +} diff --git a/src/cliproxy/quota/quota-fetcher/status-classifier.ts b/src/cliproxy/quota/quota-fetcher/status-classifier.ts new file mode 100644 index 00000000..eb520583 --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/status-classifier.ts @@ -0,0 +1,195 @@ +/** + * Status classifier for Antigravity quota fetch failures. + * + * Maps an upstream HTTP status code (plus optional response body) into a stable + * QuotaResult failure fragment: error code, action hint, retryability, and + * provider entitlement evidence. Also merges tier evidence from a successful + * project lookup with entitlement evidence derived from a failed models fetch. + */ + +import type { AccountTier } from '../../accounts/account-manager'; +import { buildProviderEntitlementEvidence } from '../../auth/provider-entitlement-evidence'; +import type { ProviderEntitlementEvidence } from '../../auth/provider-entitlement-types'; +import type { QuotaResult } from './types'; + +/** Trim and cap upstream error bodies to keep payloads small. */ +export function normalizeErrorDetail(bodyText: string): string | undefined { + const normalized = bodyText.trim(); + if (!normalized) { + return undefined; + } + if (normalized.length <= 400) { + return normalized; + } + return `${normalized.slice(0, 397)}...`; +} + +/** + * Build the failure fragment for an Antigravity quota request. The returned + * object is spread into a QuotaResult by callers. Status 401/403/429/408/5xx + * each have their own stable error code and capacity/access state signal. + */ +export function buildAntigravityFailure( + status: number | undefined, + bodyText?: string +): Pick< + QuotaResult, + | 'error' + | 'errorCode' + | 'errorDetail' + | 'actionHint' + | 'retryable' + | 'httpStatus' + | 'needsReauth' + | 'entitlement' +> & { isForbidden?: boolean } { + const detail = normalizeErrorDetail(bodyText || ''); + + if (status === 401) { + return { + httpStatus: 401, + error: 'Token expired or invalid', + errorCode: 'reauth_required', + actionHint: + 'Re-authenticate this account. If CLIProxy is running, retry after the proxy finishes refreshing the token.', + needsReauth: true, + errorDetail: detail, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'medium', + accessState: 'unknown', + capacityState: 'unknown', + }), + }; + } + + if (status === 403) { + return { + httpStatus: 403, + error: 'Access forbidden', + errorCode: 'quota_api_forbidden', + actionHint: 'This account does not have Gemini Code Assist quota access.', + isForbidden: true, + errorDetail: detail, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'medium', + accessState: 'not_entitled', + capacityState: 'unknown', + }), + }; + } + + if (status === 429) { + return { + httpStatus: 429, + error: 'Rate limited - try again later', + errorCode: 'rate_limited', + actionHint: 'Retry later. This looks temporary.', + retryable: true, + errorDetail: detail, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'low', + accessState: 'unknown', + capacityState: 'rate_limited', + }), + }; + } + + if (status === 408) { + return { + httpStatus: 408, + error: 'Request timeout', + errorCode: 'network_timeout', + actionHint: 'Retry later. This looks temporary.', + retryable: true, + errorDetail: detail, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'low', + accessState: 'unknown', + capacityState: 'temporarily_unavailable', + }), + }; + } + + if (typeof status === 'number' && status >= 500) { + return { + httpStatus: status, + error: `API error: ${status}`, + errorCode: 'provider_unavailable', + actionHint: 'Retry later. The provider appears unavailable.', + retryable: true, + errorDetail: detail, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'low', + accessState: 'unknown', + capacityState: 'temporarily_unavailable', + }), + }; + } + + if (typeof status === 'number' && status >= 400) { + return { + httpStatus: status, + error: `API error: ${status}`, + errorCode: 'quota_request_failed', + errorDetail: detail, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'low', + accessState: 'unknown', + capacityState: 'unknown', + }), + }; + } + + return { + error: 'Quota request failed', + errorCode: 'quota_request_failed', + errorDetail: detail, + entitlement: buildProviderEntitlementEvidence({ + normalizedTier: 'unknown', + source: 'runtime_inference', + confidence: 'low', + accessState: 'unknown', + capacityState: 'unknown', + }), + }; +} + +/** + * Merge entitlement evidence from a successful project lookup with evidence + * from a failed models fetch. A known tier id from the project lookup always + * wins (runtime_api source, high confidence); otherwise we fall back to the + * pre-existing evidence or build a fresh runtime_inference record. + */ +export function mergeAntigravityTierEvidence( + entitlement: ProviderEntitlementEvidence | undefined, + tier: AccountTier, + rawTierId: string | null, + rawTierLabel: string | null +): ProviderEntitlementEvidence | undefined { + if (tier === 'unknown' && !entitlement) { + return undefined; + } + + return buildProviderEntitlementEvidence({ + normalizedTier: tier, + rawTierId, + rawTierLabel, + source: rawTierId ? 'runtime_api' : (entitlement?.source ?? 'runtime_inference'), + confidence: rawTierId ? 'high' : (entitlement?.confidence ?? 'medium'), + accessState: entitlement?.accessState ?? 'unknown', + capacityState: entitlement?.capacityState ?? 'unknown', + notes: entitlement?.notes ?? null, + }); +} diff --git a/src/cliproxy/quota/quota-fetcher/types.ts b/src/cliproxy/quota/quota-fetcher/types.ts new file mode 100644 index 00000000..f7d34f5b --- /dev/null +++ b/src/cliproxy/quota/quota-fetcher/types.ts @@ -0,0 +1,180 @@ +/** + * Shared types for the Antigravity quota fetcher. + * + * Public types (ModelQuota, QuotaResult, AllAccountsQuotaResult) are re-exported + * from the barrel at ../quota-fetcher.ts so existing import paths keep working. + */ + +import type { CLIProxyProvider } from '../../types'; +import type { AccountInfo, AccountTier } from '../../accounts/account-manager'; +import type { ProviderEntitlementEvidence } from '../../auth/provider-entitlement-types'; + +/** Individual model quota info */ +export interface ModelQuota { + /** Model name, e.g., "gemini-3-pro-high" */ + name: string; + /** Display name from API, e.g., "Gemini 3 Pro" */ + displayName?: string; + /** Remaining quota as percentage (0-100) */ + percentage: number; + /** ISO timestamp when quota resets, null if unknown */ + resetTime: string | null; +} + +/** Quota fetch result */ +export interface QuotaResult { + /** Whether fetch succeeded */ + success: boolean; + /** Quota for each available model */ + models: ModelQuota[]; + /** Timestamp of fetch */ + lastUpdated: number; + /** Upstream HTTP status when available */ + httpStatus?: number; + /** Stable machine-readable error code */ + errorCode?: string; + /** Additional provider-specific detail/code from upstream */ + errorDetail?: string; + /** True if account lacks quota access (403) */ + isForbidden?: boolean; + /** Error message if fetch failed */ + error?: string; + /** Provider-specific remediation guidance */ + actionHint?: string; + /** True when the failure is temporary and retrying later may help */ + retryable?: boolean; + /** True if token is expired and needs re-auth */ + isExpired?: boolean; + /** True if token refresh cannot proceed and the account should be re-authenticated */ + needsReauth?: boolean; + /** ISO timestamp when token expires/expired */ + expiresAt?: string; + /** True if account hasn't been activated in official Antigravity app */ + isUnprovisioned?: boolean; + /** Account ID (email) this quota belongs to */ + accountId?: string; + /** GCP project ID for this account */ + projectId?: string; + /** Detected account tier based on model access */ + tier?: AccountTier; + /** Richer provider entitlement evidence derived from live/runtime signals */ + entitlement?: ProviderEntitlementEvidence; +} + +/** Result for all accounts of a provider */ +export interface AllAccountsQuotaResult { + /** Provider name */ + provider: CLIProxyProvider; + /** Results per account */ + accounts: Array<{ + account: AccountInfo; + quota: QuotaResult; + }>; + /** Accounts grouped by project ID (for detecting shared projects) */ + projectGroups: Record; + /** Timestamp of fetch */ + lastUpdated: number; +} + +// --------------------------------------------------------------------------- +// Internal types (not part of the public surface) +// --------------------------------------------------------------------------- + +/** Auth file structure on disk for Antigravity accounts */ +export interface AntigravityAuthFile { + access_token: string; + refresh_token?: string; + email?: string; + expired?: string; + expires_in?: number; + timestamp?: number; + type?: string; + project_id?: string; +} + +/** Auth data returned from file */ +export interface AuthData { + accessToken: string; + refreshToken: string | null; + projectId: string | null; + isExpired: boolean; + expiresAt: string | null; +} + +/** Tier info from loadCodeAssist */ +export interface TierInfo { + id?: string; + isDefault?: boolean; +} + +/** loadCodeAssist response */ +export interface LoadCodeAssistResponse { + cloudaicompanionProject?: string | { id?: string }; + /** Current tier (may be trial/temporary) */ + currentTier?: TierInfo; + /** Paid tier (reflects actual subscription - takes priority) */ + paidTier?: TierInfo; + /** Array of allowed tiers - use isDefault=true to find active tier (CLIProxyAPIPlus approach) */ + allowedTiers?: TierInfo[]; +} + +/** fetchAvailableModels response model */ +export interface AvailableModel { + name?: string; + displayName?: string; + quotaInfo?: { + remainingFraction?: number; + remaining_fraction?: number; + remaining?: number; + resetTime?: string; + reset_time?: string; + }; + quota_info?: { + remainingFraction?: number; + remaining_fraction?: number; + remaining?: number; + resetTime?: string; + reset_time?: string; + }; +} + +/** fetchAvailableModels response */ +export interface FetchAvailableModelsResponse { + models?: Record; +} + +export interface ManagementAuthFile { + auth_index?: string | number; + provider?: string; + type?: string; + email?: string; + name?: string; +} + +export interface ManagementApiCallResponse { + status_code?: number; + body?: string; +} + +export interface ManagedResponse { + status: number; + bodyText: string; + json: unknown; + viaManagement: boolean; +} + +export interface ProjectLookupResult { + projectId: string | null; + tier?: AccountTier; + rawTierId?: string | null; + rawTierLabel?: string | null; + entitlement?: ProviderEntitlementEvidence; + error?: string; + errorCode?: string; + errorDetail?: string; + actionHint?: string; + retryable?: boolean; + httpStatus?: number; + needsReauth?: boolean; + isUnprovisioned?: boolean; +}