| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import { createLogger } from '@automaker/utils'; |
| import * as fs from 'fs'; |
| import * as path from 'path'; |
| import * as os from 'os'; |
| import { execFileSync } from 'child_process'; |
|
|
| const logger = createLogger('GeminiUsage'); |
|
|
| |
| const QUOTA_API_URL = 'https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota'; |
|
|
| |
| const CODE_ASSIST_URL = 'https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist'; |
|
|
| |
| const GOOGLE_TOKEN_URL = 'https://oauth2.googleapis.com/token'; |
|
|
| |
| const FETCH_TIMEOUT_MS = 10_000; |
|
|
| |
| const CREDENTIALS_CACHE_TTL_MS = 5 * 60 * 1000; |
|
|
| export interface GeminiQuotaBucket { |
| |
| modelId: string; |
| |
| remainingFraction: number; |
| |
| resetTime: string; |
| } |
|
|
| |
| export interface GeminiTierQuota { |
| |
| usedPercent: number; |
| |
| remainingPercent: number; |
| |
| resetText?: string; |
| |
| resetTime?: string; |
| } |
|
|
| export interface GeminiUsageData { |
| |
| authenticated: boolean; |
| |
| authMethod: 'cli_login' | 'api_key' | 'none'; |
| |
| usedPercent: number; |
| |
| remainingPercent: number; |
| |
| resetText?: string; |
| |
| resetTime?: string; |
| |
| constrainedModel?: string; |
| |
| flashQuota?: GeminiTierQuota; |
| |
| proQuota?: GeminiTierQuota; |
| |
| quotaBuckets?: GeminiQuotaBucket[]; |
| |
| lastUpdated: string; |
| |
| error?: string; |
| } |
|
|
| interface OAuthCredentials { |
| access_token?: string; |
| id_token?: string; |
| refresh_token?: string; |
| token_type?: string; |
| expiry_date?: number; |
| client_id?: string; |
| client_secret?: string; |
| } |
|
|
| interface OAuthClientCredentials { |
| clientId: string; |
| clientSecret: string; |
| } |
|
|
| interface QuotaResponse { |
| |
| buckets?: Array<{ |
| modelId?: string; |
| remainingFraction?: number; |
| resetTime?: string; |
| tokenType?: string; |
| }>; |
| |
| quotaBuckets?: Array<{ |
| modelId?: string; |
| remainingFraction?: number; |
| resetTime?: string; |
| tokenType?: string; |
| }>; |
| } |
|
|
| |
| |
| |
| |
| |
| |
| export class GeminiUsageService { |
| private cachedCredentials: OAuthCredentials | null = null; |
| private cachedCredentialsAt: number | null = null; |
| private cachedClientCredentials: OAuthClientCredentials | null = null; |
| private credentialsPath: string; |
| |
| private loadedCredentialsPath: string | null = null; |
|
|
| constructor() { |
| |
| this.credentialsPath = path.join(os.homedir(), '.gemini', 'oauth_creds.json'); |
| } |
|
|
| |
| |
| |
| async isAvailable(): Promise<boolean> { |
| const creds = await this.loadCredentials(); |
| return Boolean(creds?.access_token || creds?.refresh_token); |
| } |
|
|
| |
| |
| |
| async fetchUsageData(): Promise<GeminiUsageData> { |
| logger.info('[fetchUsageData] Starting...'); |
|
|
| const creds = await this.loadCredentials(); |
|
|
| if (!creds || (!creds.access_token && !creds.refresh_token)) { |
| logger.info('[fetchUsageData] No credentials found'); |
| return { |
| authenticated: false, |
| authMethod: 'none', |
| usedPercent: 0, |
| remainingPercent: 100, |
| lastUpdated: new Date().toISOString(), |
| error: 'Not authenticated. Run "gemini auth login" to authenticate.', |
| }; |
| } |
|
|
| try { |
| |
| const accessToken = await this.getValidAccessToken(creds); |
|
|
| if (!accessToken) { |
| return { |
| authenticated: false, |
| authMethod: 'none', |
| usedPercent: 0, |
| remainingPercent: 100, |
| lastUpdated: new Date().toISOString(), |
| error: 'Failed to obtain access token. Try running "gemini auth login" again.', |
| }; |
| } |
|
|
| |
| |
| let projectId: string | undefined; |
| try { |
| const codeAssistResponse = await fetch(CODE_ASSIST_URL, { |
| method: 'POST', |
| headers: { |
| Authorization: `Bearer ${accessToken}`, |
| 'Content-Type': 'application/json', |
| }, |
| body: JSON.stringify({}), |
| signal: AbortSignal.timeout(FETCH_TIMEOUT_MS), |
| }); |
|
|
| if (codeAssistResponse.ok) { |
| const codeAssistData = (await codeAssistResponse.json()) as { |
| cloudaicompanionProject?: string; |
| currentTier?: { id?: string; name?: string }; |
| }; |
| projectId = codeAssistData.cloudaicompanionProject; |
| logger.debug('[fetchUsageData] Got project ID:', projectId); |
| } |
| } catch (e) { |
| logger.debug('[fetchUsageData] Failed to get project ID:', e); |
| } |
|
|
| |
| |
| const response = await fetch(QUOTA_API_URL, { |
| method: 'POST', |
| headers: { |
| Authorization: `Bearer ${accessToken}`, |
| 'Content-Type': 'application/json', |
| }, |
| body: JSON.stringify(projectId ? { project: projectId } : {}), |
| signal: AbortSignal.timeout(FETCH_TIMEOUT_MS), |
| }); |
|
|
| if (!response.ok) { |
| const errorText = await response.text().catch(() => ''); |
| logger.error('[fetchUsageData] Quota API error:', response.status, errorText); |
|
|
| |
| return { |
| authenticated: true, |
| authMethod: 'cli_login', |
| usedPercent: 0, |
| remainingPercent: 100, |
| lastUpdated: new Date().toISOString(), |
| error: `Quota API unavailable (${response.status})`, |
| }; |
| } |
|
|
| const data = (await response.json()) as QuotaResponse; |
|
|
| |
| const apiBuckets = data.buckets || data.quotaBuckets; |
|
|
| logger.debug('[fetchUsageData] Raw buckets:', JSON.stringify(apiBuckets)); |
|
|
| if (!apiBuckets || apiBuckets.length === 0) { |
| return { |
| authenticated: true, |
| authMethod: 'cli_login', |
| usedPercent: 0, |
| remainingPercent: 100, |
| lastUpdated: new Date().toISOString(), |
| }; |
| } |
|
|
| |
| |
| |
| let flashLowestRemaining = 1.0; |
| let flashResetTime: string | undefined; |
| let hasFlashModels = false; |
| let proLowestRemaining = 1.0; |
| let proResetTime: string | undefined; |
| let hasProModels = false; |
| let overallLowestRemaining = 1.0; |
| let constrainedModel: string | undefined; |
| let overallResetTime: string | undefined; |
|
|
| const quotaBuckets: GeminiQuotaBucket[] = apiBuckets.map((bucket) => { |
| const remaining = bucket.remainingFraction ?? 1.0; |
| const modelId = bucket.modelId?.toLowerCase() || ''; |
|
|
| |
| if (remaining < overallLowestRemaining) { |
| overallLowestRemaining = remaining; |
| constrainedModel = bucket.modelId; |
| overallResetTime = bucket.resetTime; |
| } |
|
|
| |
| if (modelId.includes('flash')) { |
| hasFlashModels = true; |
| if (remaining < flashLowestRemaining) { |
| flashLowestRemaining = remaining; |
| flashResetTime = bucket.resetTime; |
| } |
| |
| if (!flashResetTime && bucket.resetTime) { |
| flashResetTime = bucket.resetTime; |
| } |
| } else if (modelId.includes('pro')) { |
| hasProModels = true; |
| if (remaining < proLowestRemaining) { |
| proLowestRemaining = remaining; |
| proResetTime = bucket.resetTime; |
| } |
| |
| if (!proResetTime && bucket.resetTime) { |
| proResetTime = bucket.resetTime; |
| } |
| } |
|
|
| return { |
| modelId: bucket.modelId || 'unknown', |
| remainingFraction: remaining, |
| resetTime: bucket.resetTime || '', |
| }; |
| }); |
|
|
| const usedPercent = Math.round((1 - overallLowestRemaining) * 100); |
| const remainingPercent = Math.round(overallLowestRemaining * 100); |
|
|
| |
| const flashQuota: GeminiTierQuota | undefined = hasFlashModels |
| ? { |
| usedPercent: Math.round((1 - flashLowestRemaining) * 100), |
| remainingPercent: Math.round(flashLowestRemaining * 100), |
| resetText: flashResetTime ? this.formatResetTime(flashResetTime) : undefined, |
| resetTime: flashResetTime, |
| } |
| : undefined; |
|
|
| const proQuota: GeminiTierQuota | undefined = hasProModels |
| ? { |
| usedPercent: Math.round((1 - proLowestRemaining) * 100), |
| remainingPercent: Math.round(proLowestRemaining * 100), |
| resetText: proResetTime ? this.formatResetTime(proResetTime) : undefined, |
| resetTime: proResetTime, |
| } |
| : undefined; |
|
|
| return { |
| authenticated: true, |
| authMethod: 'cli_login', |
| usedPercent, |
| remainingPercent, |
| resetText: overallResetTime ? this.formatResetTime(overallResetTime) : undefined, |
| resetTime: overallResetTime, |
| constrainedModel, |
| flashQuota, |
| proQuota, |
| quotaBuckets, |
| lastUpdated: new Date().toISOString(), |
| }; |
| } catch (error) { |
| const errorMsg = error instanceof Error ? error.message : 'Unknown error'; |
| logger.error('[fetchUsageData] Error:', errorMsg); |
|
|
| return { |
| authenticated: true, |
| authMethod: 'cli_login', |
| usedPercent: 0, |
| remainingPercent: 100, |
| lastUpdated: new Date().toISOString(), |
| error: `Failed to fetch quota: ${errorMsg}`, |
| }; |
| } |
| } |
|
|
| |
| |
| |
| |
| private async loadCredentials(): Promise<OAuthCredentials | null> { |
| |
| if (this.cachedCredentials && this.cachedCredentialsAt) { |
| const now = Date.now(); |
| const cacheAge = now - this.cachedCredentialsAt; |
|
|
| if (cacheAge < CREDENTIALS_CACHE_TTL_MS) { |
| |
| const sourcePath = this.loadedCredentialsPath || this.credentialsPath; |
| try { |
| const stat = fs.statSync(sourcePath); |
| if (stat.mtimeMs <= this.cachedCredentialsAt) { |
| |
| return this.cachedCredentials; |
| } |
| |
| logger.debug('[loadCredentials] File modified since cache, re-reading'); |
| } catch { |
| |
| return this.cachedCredentials; |
| } |
| } else { |
| |
| logger.debug('[loadCredentials] Cache TTL expired, re-reading'); |
| } |
|
|
| |
| this.cachedCredentials = null; |
| this.cachedCredentialsAt = null; |
| } |
|
|
| |
| const rawPaths = [ |
| this.credentialsPath, |
| path.join(os.homedir(), '.config', 'gemini', 'oauth_creds.json'), |
| ]; |
| const possiblePaths = [...new Set(rawPaths)]; |
|
|
| for (const credPath of possiblePaths) { |
| try { |
| if (fs.existsSync(credPath)) { |
| const content = fs.readFileSync(credPath, 'utf8'); |
| const creds = JSON.parse(content); |
|
|
| |
| if (creds.access_token || creds.refresh_token) { |
| this.cachedCredentials = creds; |
| this.cachedCredentialsAt = Date.now(); |
| this.loadedCredentialsPath = credPath; |
| logger.info('[loadCredentials] Loaded from:', credPath); |
| return creds; |
| } |
|
|
| |
| if (creds.web?.client_id || creds.installed?.client_id) { |
| const clientCreds = creds.web || creds.installed; |
| this.cachedCredentials = { |
| client_id: clientCreds.client_id, |
| client_secret: clientCreds.client_secret, |
| }; |
| this.cachedCredentialsAt = Date.now(); |
| this.loadedCredentialsPath = credPath; |
| return this.cachedCredentials; |
| } |
| } |
| } catch (error) { |
| logger.debug('[loadCredentials] Failed to load from', credPath, error); |
| } |
| } |
|
|
| return null; |
| } |
|
|
| |
| |
| |
| private findGeminiBinaryPath(): string | null { |
| |
| const whichCmd = process.platform === 'win32' ? 'where' : 'which'; |
| try { |
| const whichResult = execFileSync(whichCmd, ['gemini'], { |
| encoding: 'utf8', |
| timeout: 5000, |
| stdio: ['pipe', 'pipe', 'pipe'], |
| }).trim(); |
| |
| const firstLine = whichResult.split('\n')[0]?.trim(); |
| if (firstLine && fs.existsSync(firstLine)) { |
| return firstLine; |
| } |
| } catch { |
| |
| } |
|
|
| |
| const possiblePaths = [ |
| |
| path.join(os.homedir(), '.npm-global', 'bin', 'gemini'), |
| '/usr/local/bin/gemini', |
| '/usr/bin/gemini', |
| |
| '/opt/homebrew/bin/gemini', |
| '/usr/local/opt/gemini/bin/gemini', |
| |
| path.join(os.homedir(), '.nvm', 'versions', 'node'), |
| path.join(os.homedir(), '.fnm', 'node-versions'), |
| |
| path.join(os.homedir(), 'AppData', 'Roaming', 'npm', 'gemini.cmd'), |
| path.join(os.homedir(), 'AppData', 'Roaming', 'npm', 'gemini'), |
| ]; |
|
|
| for (const p of possiblePaths) { |
| if (fs.existsSync(p)) { |
| return p; |
| } |
| } |
|
|
| return null; |
| } |
|
|
| |
| |
| |
| |
| private extractOAuthClientCredentials(): OAuthClientCredentials | null { |
| if (this.cachedClientCredentials) { |
| return this.cachedClientCredentials; |
| } |
|
|
| const geminiBinary = this.findGeminiBinaryPath(); |
| if (!geminiBinary) { |
| logger.debug('[extractOAuthClientCredentials] Gemini binary not found'); |
| return null; |
| } |
|
|
| |
| let resolvedPath = geminiBinary; |
| try { |
| resolvedPath = fs.realpathSync(geminiBinary); |
| } catch { |
| |
| } |
|
|
| const baseDir = path.dirname(resolvedPath); |
| logger.debug('[extractOAuthClientCredentials] Base dir:', baseDir); |
|
|
| |
| |
| const possibleOAuth2Paths = [ |
| |
| path.join( |
| baseDir, |
| '..', |
| 'lib', |
| 'node_modules', |
| '@google', |
| 'gemini-cli', |
| 'dist', |
| 'src', |
| 'code_assist', |
| 'oauth2.js' |
| ), |
| path.join( |
| baseDir, |
| '..', |
| 'lib', |
| 'node_modules', |
| '@google', |
| 'gemini-cli-core', |
| 'dist', |
| 'src', |
| 'code_assist', |
| 'oauth2.js' |
| ), |
| |
| path.join( |
| baseDir, |
| '..', |
| 'libexec', |
| 'lib', |
| 'node_modules', |
| '@google', |
| 'gemini-cli', |
| 'dist', |
| 'src', |
| 'code_assist', |
| 'oauth2.js' |
| ), |
| path.join( |
| baseDir, |
| '..', |
| 'libexec', |
| 'lib', |
| 'node_modules', |
| '@google', |
| 'gemini-cli-core', |
| 'dist', |
| 'src', |
| 'code_assist', |
| 'oauth2.js' |
| ), |
| |
| path.join(baseDir, '..', 'gemini-cli-core', 'dist', 'src', 'code_assist', 'oauth2.js'), |
| path.join(baseDir, '..', 'gemini-cli', 'dist', 'src', 'code_assist', 'oauth2.js'), |
| |
| path.join( |
| baseDir, |
| '..', |
| '..', |
| 'lib', |
| 'node_modules', |
| '@google', |
| 'gemini-cli', |
| 'dist', |
| 'src', |
| 'code_assist', |
| 'oauth2.js' |
| ), |
| path.join( |
| baseDir, |
| '..', |
| '..', |
| 'lib', |
| 'node_modules', |
| '@google', |
| 'gemini-cli-core', |
| 'dist', |
| 'src', |
| 'code_assist', |
| 'oauth2.js' |
| ), |
| ]; |
|
|
| for (const oauth2Path of possibleOAuth2Paths) { |
| try { |
| const normalizedPath = path.normalize(oauth2Path); |
| if (fs.existsSync(normalizedPath)) { |
| logger.debug('[extractOAuthClientCredentials] Found oauth2.js at:', normalizedPath); |
| const content = fs.readFileSync(normalizedPath, 'utf8'); |
| const creds = this.parseOAuthCredentialsFromSource(content); |
| if (creds) { |
| this.cachedClientCredentials = creds; |
| logger.info('[extractOAuthClientCredentials] Extracted credentials from CLI'); |
| return creds; |
| } |
| } |
| } catch (error) { |
| logger.debug('[extractOAuthClientCredentials] Failed to read', oauth2Path, error); |
| } |
| } |
|
|
| |
| if (process.platform !== 'win32') { |
| try { |
| const searchBase = path.resolve(baseDir, '..'); |
| const searchResult = execFileSync( |
| 'find', |
| [searchBase, '-name', 'oauth2.js', '-path', '*gemini*', '-path', '*code_assist*'], |
| { encoding: 'utf8', timeout: 5000, stdio: ['pipe', 'pipe', 'pipe'] } |
| ) |
| .trim() |
| .split('\n')[0]; |
|
|
| if (searchResult && fs.existsSync(searchResult)) { |
| logger.debug('[extractOAuthClientCredentials] Found via search:', searchResult); |
| const content = fs.readFileSync(searchResult, 'utf8'); |
| const creds = this.parseOAuthCredentialsFromSource(content); |
| if (creds) { |
| this.cachedClientCredentials = creds; |
| logger.info( |
| '[extractOAuthClientCredentials] Extracted credentials from CLI (via search)' |
| ); |
| return creds; |
| } |
| } |
| } catch { |
| |
| } |
| } |
|
|
| logger.warn('[extractOAuthClientCredentials] Could not extract credentials from CLI'); |
| return null; |
| } |
|
|
| |
| |
| |
| private parseOAuthCredentialsFromSource(content: string): OAuthClientCredentials | null { |
| |
| |
| const clientIdPatterns = [ |
| /OAUTH_CLIENT_ID\s*=\s*["']([^"']+)["']/, |
| /clientId\s*[:=]\s*["']([^"']+)["']/, |
| /client_id\s*[:=]\s*["']([^"']+)["']/, |
| /"clientId"\s*:\s*["']([^"']+)["']/, |
| ]; |
|
|
| const clientSecretPatterns = [ |
| /OAUTH_CLIENT_SECRET\s*=\s*["']([^"']+)["']/, |
| /clientSecret\s*[:=]\s*["']([^"']+)["']/, |
| /client_secret\s*[:=]\s*["']([^"']+)["']/, |
| /"clientSecret"\s*:\s*["']([^"']+)["']/, |
| ]; |
|
|
| let clientId: string | null = null; |
| let clientSecret: string | null = null; |
|
|
| for (const pattern of clientIdPatterns) { |
| const match = content.match(pattern); |
| if (match && match[1]) { |
| clientId = match[1]; |
| break; |
| } |
| } |
|
|
| for (const pattern of clientSecretPatterns) { |
| const match = content.match(pattern); |
| if (match && match[1]) { |
| clientSecret = match[1]; |
| break; |
| } |
| } |
|
|
| if (clientId && clientSecret) { |
| logger.debug('[parseOAuthCredentialsFromSource] Found client credentials'); |
| return { clientId, clientSecret }; |
| } |
|
|
| return null; |
| } |
|
|
| |
| |
| |
| private async getValidAccessToken(creds: OAuthCredentials): Promise<string | null> { |
| |
| if (creds.access_token && creds.expiry_date) { |
| const now = Date.now(); |
| if (creds.expiry_date > now + 5 * 60 * 1000) { |
| logger.debug('[getValidAccessToken] Using existing token (not expired)'); |
| return creds.access_token; |
| } |
| } |
|
|
| |
| if (creds.refresh_token) { |
| |
| const extractedCreds = this.extractOAuthClientCredentials(); |
|
|
| |
| const clientId = extractedCreds?.clientId || creds.client_id; |
| const clientSecret = extractedCreds?.clientSecret || creds.client_secret; |
|
|
| if (!clientId || !clientSecret) { |
| logger.error('[getValidAccessToken] No client credentials available for token refresh'); |
| |
| return creds.access_token || null; |
| } |
|
|
| try { |
| logger.debug('[getValidAccessToken] Refreshing token...'); |
| const response = await fetch(GOOGLE_TOKEN_URL, { |
| method: 'POST', |
| headers: { |
| 'Content-Type': 'application/x-www-form-urlencoded', |
| }, |
| body: new URLSearchParams({ |
| client_id: clientId, |
| client_secret: clientSecret, |
| refresh_token: creds.refresh_token, |
| grant_type: 'refresh_token', |
| }), |
| signal: AbortSignal.timeout(FETCH_TIMEOUT_MS), |
| }); |
|
|
| if (response.ok) { |
| const data = (await response.json()) as { access_token?: string; expires_in?: number }; |
| const newAccessToken = data.access_token; |
| const expiresIn = data.expires_in || 3600; |
|
|
| if (newAccessToken) { |
| logger.info('[getValidAccessToken] Token refreshed successfully'); |
|
|
| |
| this.cachedCredentials = { |
| ...creds, |
| access_token: newAccessToken, |
| expiry_date: Date.now() + expiresIn * 1000, |
| }; |
| this.cachedCredentialsAt = Date.now(); |
|
|
| |
| const writePath = this.loadedCredentialsPath || this.credentialsPath; |
| try { |
| fs.writeFileSync(writePath, JSON.stringify(this.cachedCredentials, null, 2)); |
| } catch (e) { |
| logger.debug('[getValidAccessToken] Could not save refreshed token:', e); |
| } |
|
|
| return newAccessToken; |
| } |
| } else { |
| const errorText = await response.text().catch(() => ''); |
| logger.error('[getValidAccessToken] Token refresh failed:', response.status, errorText); |
| } |
| } catch (error) { |
| logger.error('[getValidAccessToken] Token refresh error:', error); |
| } |
| } |
|
|
| |
| return creds.access_token || null; |
| } |
|
|
| |
| |
| |
| private formatResetTime(isoTime: string): string { |
| try { |
| const resetDate = new Date(isoTime); |
| const now = new Date(); |
| const diff = resetDate.getTime() - now.getTime(); |
|
|
| if (diff < 0) { |
| return 'Resetting soon'; |
| } |
|
|
| const minutes = Math.floor(diff / 60000); |
| const hours = Math.floor(minutes / 60); |
|
|
| if (hours > 0) { |
| const remainingMins = minutes % 60; |
| return remainingMins > 0 ? `Resets in ${hours}h ${remainingMins}m` : `Resets in ${hours}h`; |
| } |
|
|
| return `Resets in ${minutes}m`; |
| } catch { |
| return ''; |
| } |
| } |
|
|
| |
| |
| |
| clearCache(): void { |
| this.cachedCredentials = null; |
| this.cachedCredentialsAt = null; |
| this.cachedClientCredentials = null; |
| } |
| } |
|
|
| |
| let usageServiceInstance: GeminiUsageService | null = null; |
|
|
| |
| |
| |
| export function getGeminiUsageService(): GeminiUsageService { |
| if (!usageServiceInstance) { |
| usageServiceInstance = new GeminiUsageService(); |
| } |
| return usageServiceInstance; |
| } |
|
|