|
|
const redisClient = require('../models/redis') |
|
|
const { v4: uuidv4 } = require('uuid') |
|
|
const crypto = require('crypto') |
|
|
const config = require('../../config/config') |
|
|
const logger = require('../utils/logger') |
|
|
const { OAuth2Client } = require('google-auth-library') |
|
|
const { maskToken } = require('../utils/tokenMask') |
|
|
const ProxyHelper = require('../utils/proxyHelper') |
|
|
const { |
|
|
logRefreshStart, |
|
|
logRefreshSuccess, |
|
|
logRefreshError, |
|
|
logTokenUsage, |
|
|
logRefreshSkipped |
|
|
} = require('../utils/tokenRefreshLogger') |
|
|
const tokenRefreshService = require('./tokenRefreshService') |
|
|
const LRUCache = require('../utils/lruCache') |
|
|
|
|
|
|
|
|
const OAUTH_CLIENT_ID = '681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com' |
|
|
const OAUTH_CLIENT_SECRET = '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' |
|
|
const OAUTH_SCOPES = ['https://www.googleapis.com/auth/cloud-platform'] |
|
|
|
|
|
|
|
|
const ALGORITHM = 'aes-256-cbc' |
|
|
const ENCRYPTION_SALT = 'gemini-account-salt' |
|
|
const IV_LENGTH = 16 |
|
|
|
|
|
|
|
|
|
|
|
let _encryptionKeyCache = null |
|
|
|
|
|
|
|
|
const decryptCache = new LRUCache(500) |
|
|
|
|
|
|
|
|
function generateEncryptionKey() { |
|
|
if (!_encryptionKeyCache) { |
|
|
_encryptionKeyCache = crypto.scryptSync(config.security.encryptionKey, ENCRYPTION_SALT, 32) |
|
|
logger.info('🔑 Gemini encryption key derived and cached for performance optimization') |
|
|
} |
|
|
return _encryptionKeyCache |
|
|
} |
|
|
|
|
|
|
|
|
const GEMINI_ACCOUNT_KEY_PREFIX = 'gemini_account:' |
|
|
const SHARED_GEMINI_ACCOUNTS_KEY = 'shared_gemini_accounts' |
|
|
const ACCOUNT_SESSION_MAPPING_PREFIX = 'gemini_session_account_mapping:' |
|
|
|
|
|
|
|
|
function encrypt(text) { |
|
|
if (!text) { |
|
|
return '' |
|
|
} |
|
|
const key = generateEncryptionKey() |
|
|
const iv = crypto.randomBytes(IV_LENGTH) |
|
|
const cipher = crypto.createCipheriv(ALGORITHM, key, iv) |
|
|
let encrypted = cipher.update(text) |
|
|
encrypted = Buffer.concat([encrypted, cipher.final()]) |
|
|
return `${iv.toString('hex')}:${encrypted.toString('hex')}` |
|
|
} |
|
|
|
|
|
|
|
|
function decrypt(text) { |
|
|
if (!text) { |
|
|
return '' |
|
|
} |
|
|
|
|
|
|
|
|
const cacheKey = crypto.createHash('sha256').update(text).digest('hex') |
|
|
const cached = decryptCache.get(cacheKey) |
|
|
if (cached !== undefined) { |
|
|
return cached |
|
|
} |
|
|
|
|
|
try { |
|
|
const key = generateEncryptionKey() |
|
|
|
|
|
const ivHex = text.substring(0, 32) |
|
|
const encryptedHex = text.substring(33) |
|
|
|
|
|
const iv = Buffer.from(ivHex, 'hex') |
|
|
const encryptedText = Buffer.from(encryptedHex, 'hex') |
|
|
const decipher = crypto.createDecipheriv(ALGORITHM, key, iv) |
|
|
let decrypted = decipher.update(encryptedText) |
|
|
decrypted = Buffer.concat([decrypted, decipher.final()]) |
|
|
const result = decrypted.toString() |
|
|
|
|
|
|
|
|
decryptCache.set(cacheKey, result, 5 * 60 * 1000) |
|
|
|
|
|
|
|
|
if ((decryptCache.hits + decryptCache.misses) % 1000 === 0) { |
|
|
decryptCache.printStats() |
|
|
} |
|
|
|
|
|
return result |
|
|
} catch (error) { |
|
|
logger.error('Decryption error:', error) |
|
|
return '' |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
setInterval( |
|
|
() => { |
|
|
decryptCache.cleanup() |
|
|
logger.info('🧹 Gemini decrypt cache cleanup completed', decryptCache.getStats()) |
|
|
}, |
|
|
10 * 60 * 1000 |
|
|
) |
|
|
|
|
|
|
|
|
function createOAuth2Client(redirectUri = null, proxyConfig = null) { |
|
|
|
|
|
const uri = redirectUri || 'http://localhost:45462' |
|
|
|
|
|
|
|
|
const clientOptions = { |
|
|
clientId: OAUTH_CLIENT_ID, |
|
|
clientSecret: OAUTH_CLIENT_SECRET, |
|
|
redirectUri: uri |
|
|
} |
|
|
|
|
|
|
|
|
if (proxyConfig) { |
|
|
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig) |
|
|
if (proxyAgent) { |
|
|
|
|
|
clientOptions.transporterOptions = { |
|
|
agent: proxyAgent, |
|
|
httpsAgent: proxyAgent |
|
|
} |
|
|
logger.debug('Created OAuth2Client with proxy configuration') |
|
|
} |
|
|
} |
|
|
|
|
|
return new OAuth2Client(clientOptions) |
|
|
} |
|
|
|
|
|
|
|
|
async function generateAuthUrl(state = null, redirectUri = null, proxyConfig = null) { |
|
|
|
|
|
const finalRedirectUri = redirectUri || 'https://codeassist.google.com/authcode' |
|
|
const oAuth2Client = createOAuth2Client(finalRedirectUri, proxyConfig) |
|
|
|
|
|
if (proxyConfig) { |
|
|
logger.info( |
|
|
`🌐 Using proxy for Gemini auth URL generation: ${ProxyHelper.getProxyDescription(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🌐 No proxy configured for Gemini auth URL generation') |
|
|
} |
|
|
|
|
|
|
|
|
const codeVerifier = await oAuth2Client.generateCodeVerifierAsync() |
|
|
const stateValue = state || crypto.randomBytes(32).toString('hex') |
|
|
|
|
|
const authUrl = oAuth2Client.generateAuthUrl({ |
|
|
redirect_uri: finalRedirectUri, |
|
|
access_type: 'offline', |
|
|
scope: OAUTH_SCOPES, |
|
|
code_challenge_method: 'S256', |
|
|
code_challenge: codeVerifier.codeChallenge, |
|
|
state: stateValue, |
|
|
prompt: 'select_account' |
|
|
}) |
|
|
|
|
|
return { |
|
|
authUrl, |
|
|
state: stateValue, |
|
|
codeVerifier: codeVerifier.codeVerifier, |
|
|
redirectUri: finalRedirectUri |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function pollAuthorizationStatus(sessionId, maxAttempts = 60, interval = 2000) { |
|
|
let attempts = 0 |
|
|
const client = redisClient.getClientSafe() |
|
|
|
|
|
while (attempts < maxAttempts) { |
|
|
try { |
|
|
const sessionData = await client.get(`oauth_session:${sessionId}`) |
|
|
if (!sessionData) { |
|
|
throw new Error('OAuth session not found') |
|
|
} |
|
|
|
|
|
const session = JSON.parse(sessionData) |
|
|
if (session.code) { |
|
|
|
|
|
const tokens = await exchangeCodeForTokens(session.code) |
|
|
|
|
|
|
|
|
await client.del(`oauth_session:${sessionId}`) |
|
|
|
|
|
return { |
|
|
success: true, |
|
|
tokens |
|
|
} |
|
|
} |
|
|
|
|
|
if (session.error) { |
|
|
|
|
|
await client.del(`oauth_session:${sessionId}`) |
|
|
return { |
|
|
success: false, |
|
|
error: session.error |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
await new Promise((resolve) => setTimeout(resolve, interval)) |
|
|
attempts++ |
|
|
} catch (error) { |
|
|
logger.error('Error polling authorization status:', error) |
|
|
throw error |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
await client.del(`oauth_session:${sessionId}`) |
|
|
return { |
|
|
success: false, |
|
|
error: 'Authorization timeout' |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function exchangeCodeForTokens( |
|
|
code, |
|
|
redirectUri = null, |
|
|
codeVerifier = null, |
|
|
proxyConfig = null |
|
|
) { |
|
|
try { |
|
|
|
|
|
const oAuth2Client = createOAuth2Client(redirectUri, proxyConfig) |
|
|
|
|
|
if (proxyConfig) { |
|
|
logger.info( |
|
|
`🌐 Using proxy for Gemini token exchange: ${ProxyHelper.getProxyDescription(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🌐 No proxy configured for Gemini token exchange') |
|
|
} |
|
|
|
|
|
const tokenParams = { |
|
|
code, |
|
|
redirect_uri: redirectUri |
|
|
} |
|
|
|
|
|
|
|
|
if (codeVerifier) { |
|
|
tokenParams.codeVerifier = codeVerifier |
|
|
} |
|
|
|
|
|
const { tokens } = await oAuth2Client.getToken(tokenParams) |
|
|
|
|
|
|
|
|
return { |
|
|
access_token: tokens.access_token, |
|
|
refresh_token: tokens.refresh_token, |
|
|
scope: tokens.scope || OAUTH_SCOPES.join(' '), |
|
|
token_type: tokens.token_type || 'Bearer', |
|
|
expiry_date: tokens.expiry_date || Date.now() + tokens.expires_in * 1000 |
|
|
} |
|
|
} catch (error) { |
|
|
logger.error('Error exchanging code for tokens:', error) |
|
|
throw new Error('Failed to exchange authorization code') |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function refreshAccessToken(refreshToken, proxyConfig = null) { |
|
|
|
|
|
const oAuth2Client = createOAuth2Client(null, proxyConfig) |
|
|
|
|
|
try { |
|
|
|
|
|
oAuth2Client.setCredentials({ |
|
|
refresh_token: refreshToken |
|
|
}) |
|
|
|
|
|
if (proxyConfig) { |
|
|
logger.info( |
|
|
`🔄 Using proxy for Gemini token refresh: ${ProxyHelper.maskProxyInfo(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🔄 No proxy configured for Gemini token refresh') |
|
|
} |
|
|
|
|
|
|
|
|
const response = await oAuth2Client.refreshAccessToken() |
|
|
const { credentials } = response |
|
|
|
|
|
|
|
|
if (!credentials || !credentials.access_token) { |
|
|
throw new Error('No access token returned from refresh') |
|
|
} |
|
|
|
|
|
logger.info( |
|
|
`🔄 Successfully refreshed Gemini token. New expiry: ${new Date(credentials.expiry_date).toISOString()}` |
|
|
) |
|
|
|
|
|
return { |
|
|
access_token: credentials.access_token, |
|
|
refresh_token: credentials.refresh_token || refreshToken, |
|
|
scope: credentials.scope || OAUTH_SCOPES.join(' '), |
|
|
token_type: credentials.token_type || 'Bearer', |
|
|
expiry_date: credentials.expiry_date || Date.now() + 3600000 |
|
|
} |
|
|
} catch (error) { |
|
|
logger.error('Error refreshing access token:', { |
|
|
message: error.message, |
|
|
code: error.code, |
|
|
response: error.response?.data, |
|
|
hasProxy: !!proxyConfig, |
|
|
proxy: proxyConfig ? ProxyHelper.maskProxyInfo(proxyConfig) : 'No proxy' |
|
|
}) |
|
|
throw new Error(`Failed to refresh access token: ${error.message}`) |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function createAccount(accountData) { |
|
|
const id = uuidv4() |
|
|
const now = new Date().toISOString() |
|
|
|
|
|
|
|
|
let geminiOauth = null |
|
|
let accessToken = '' |
|
|
let refreshToken = '' |
|
|
let expiresAt = '' |
|
|
|
|
|
if (accountData.geminiOauth || accountData.accessToken) { |
|
|
|
|
|
if (accountData.geminiOauth) { |
|
|
geminiOauth = |
|
|
typeof accountData.geminiOauth === 'string' |
|
|
? accountData.geminiOauth |
|
|
: JSON.stringify(accountData.geminiOauth) |
|
|
|
|
|
const oauthData = |
|
|
typeof accountData.geminiOauth === 'string' |
|
|
? JSON.parse(accountData.geminiOauth) |
|
|
: accountData.geminiOauth |
|
|
|
|
|
accessToken = oauthData.access_token || '' |
|
|
refreshToken = oauthData.refresh_token || '' |
|
|
expiresAt = oauthData.expiry_date ? new Date(oauthData.expiry_date).toISOString() : '' |
|
|
} else { |
|
|
|
|
|
;({ accessToken } = accountData) |
|
|
refreshToken = accountData.refreshToken || '' |
|
|
|
|
|
|
|
|
geminiOauth = JSON.stringify({ |
|
|
access_token: accessToken, |
|
|
refresh_token: refreshToken, |
|
|
scope: accountData.scope || OAUTH_SCOPES.join(' '), |
|
|
token_type: accountData.tokenType || 'Bearer', |
|
|
expiry_date: accountData.expiryDate || Date.now() + 3600000 |
|
|
}) |
|
|
|
|
|
expiresAt = new Date(accountData.expiryDate || Date.now() + 3600000).toISOString() |
|
|
} |
|
|
} |
|
|
|
|
|
const account = { |
|
|
id, |
|
|
platform: 'gemini', |
|
|
name: accountData.name || 'Gemini Account', |
|
|
description: accountData.description || '', |
|
|
accountType: accountData.accountType || 'shared', |
|
|
isActive: 'true', |
|
|
status: 'active', |
|
|
|
|
|
|
|
|
schedulable: accountData.schedulable !== undefined ? String(accountData.schedulable) : 'true', |
|
|
priority: accountData.priority || 50, |
|
|
|
|
|
|
|
|
geminiOauth: geminiOauth ? encrypt(geminiOauth) : '', |
|
|
accessToken: accessToken ? encrypt(accessToken) : '', |
|
|
refreshToken: refreshToken ? encrypt(refreshToken) : '', |
|
|
expiresAt, |
|
|
|
|
|
scopes: accountData.geminiOauth ? accountData.scopes || OAUTH_SCOPES.join(' ') : '', |
|
|
|
|
|
|
|
|
subscriptionExpiresAt: accountData.subscriptionExpiresAt || null, |
|
|
|
|
|
|
|
|
proxy: accountData.proxy ? JSON.stringify(accountData.proxy) : '', |
|
|
|
|
|
|
|
|
projectId: accountData.projectId || '', |
|
|
|
|
|
|
|
|
tempProjectId: accountData.tempProjectId || '', |
|
|
|
|
|
|
|
|
supportedModels: accountData.supportedModels || [], |
|
|
|
|
|
|
|
|
createdAt: now, |
|
|
updatedAt: now, |
|
|
lastUsedAt: '', |
|
|
lastRefreshAt: '' |
|
|
} |
|
|
|
|
|
|
|
|
const client = redisClient.getClientSafe() |
|
|
await client.hset(`${GEMINI_ACCOUNT_KEY_PREFIX}${id}`, account) |
|
|
|
|
|
|
|
|
if (account.accountType === 'shared') { |
|
|
await client.sadd(SHARED_GEMINI_ACCOUNTS_KEY, id) |
|
|
} |
|
|
|
|
|
logger.info(`Created Gemini account: ${id}`) |
|
|
|
|
|
|
|
|
const returnAccount = { ...account } |
|
|
if (returnAccount.proxy) { |
|
|
try { |
|
|
returnAccount.proxy = JSON.parse(returnAccount.proxy) |
|
|
} catch (e) { |
|
|
returnAccount.proxy = null |
|
|
} |
|
|
} |
|
|
|
|
|
return returnAccount |
|
|
} |
|
|
|
|
|
|
|
|
async function getAccount(accountId) { |
|
|
const client = redisClient.getClientSafe() |
|
|
const accountData = await client.hgetall(`${GEMINI_ACCOUNT_KEY_PREFIX}${accountId}`) |
|
|
|
|
|
if (!accountData || Object.keys(accountData).length === 0) { |
|
|
return null |
|
|
} |
|
|
|
|
|
|
|
|
if (accountData.geminiOauth) { |
|
|
accountData.geminiOauth = decrypt(accountData.geminiOauth) |
|
|
} |
|
|
if (accountData.accessToken) { |
|
|
accountData.accessToken = decrypt(accountData.accessToken) |
|
|
} |
|
|
if (accountData.refreshToken) { |
|
|
accountData.refreshToken = decrypt(accountData.refreshToken) |
|
|
} |
|
|
|
|
|
|
|
|
if (accountData.proxy) { |
|
|
try { |
|
|
accountData.proxy = JSON.parse(accountData.proxy) |
|
|
} catch (e) { |
|
|
|
|
|
accountData.proxy = null |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
accountData.schedulable = accountData.schedulable !== 'false' |
|
|
|
|
|
return accountData |
|
|
} |
|
|
|
|
|
|
|
|
async function updateAccount(accountId, updates) { |
|
|
const existingAccount = await getAccount(accountId) |
|
|
if (!existingAccount) { |
|
|
throw new Error('Account not found') |
|
|
} |
|
|
|
|
|
const now = new Date().toISOString() |
|
|
updates.updatedAt = now |
|
|
|
|
|
|
|
|
|
|
|
const oldRefreshToken = existingAccount.refreshToken || '' |
|
|
let needUpdateExpiry = false |
|
|
|
|
|
|
|
|
if (updates.proxy !== undefined) { |
|
|
updates.proxy = updates.proxy ? JSON.stringify(updates.proxy) : '' |
|
|
} |
|
|
|
|
|
|
|
|
if (updates.schedulable !== undefined) { |
|
|
updates.schedulable = updates.schedulable.toString() |
|
|
} |
|
|
|
|
|
|
|
|
if (updates.geminiOauth) { |
|
|
updates.geminiOauth = encrypt( |
|
|
typeof updates.geminiOauth === 'string' |
|
|
? updates.geminiOauth |
|
|
: JSON.stringify(updates.geminiOauth) |
|
|
) |
|
|
} |
|
|
if (updates.accessToken) { |
|
|
updates.accessToken = encrypt(updates.accessToken) |
|
|
} |
|
|
if (updates.refreshToken) { |
|
|
updates.refreshToken = encrypt(updates.refreshToken) |
|
|
|
|
|
if (!oldRefreshToken && updates.refreshToken) { |
|
|
needUpdateExpiry = true |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
const client = redisClient.getClientSafe() |
|
|
if (updates.accountType && updates.accountType !== existingAccount.accountType) { |
|
|
if (updates.accountType === 'shared') { |
|
|
await client.sadd(SHARED_GEMINI_ACCOUNTS_KEY, accountId) |
|
|
} else { |
|
|
await client.srem(SHARED_GEMINI_ACCOUNTS_KEY, accountId) |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
if (needUpdateExpiry) { |
|
|
const newExpiry = new Date(Date.now() + 10 * 60 * 1000).toISOString() |
|
|
updates.expiresAt = newExpiry |
|
|
|
|
|
logger.info( |
|
|
`🔄 New refresh token added for Gemini account ${accountId}, setting token expiry to 10 minutes` |
|
|
) |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
if (updates.subscriptionExpiresAt !== undefined) { |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
if (updates.geminiOauth && !oldRefreshToken) { |
|
|
const oauthData = |
|
|
typeof updates.geminiOauth === 'string' |
|
|
? JSON.parse(decrypt(updates.geminiOauth)) |
|
|
: updates.geminiOauth |
|
|
|
|
|
if (oauthData.refresh_token) { |
|
|
|
|
|
const providedExpiry = oauthData.expiry_date || 0 |
|
|
const currentTime = Date.now() |
|
|
const oneHour = 60 * 60 * 1000 |
|
|
|
|
|
if (providedExpiry - currentTime > oneHour) { |
|
|
const newExpiry = new Date(currentTime + 10 * 60 * 1000).toISOString() |
|
|
updates.expiresAt = newExpiry |
|
|
logger.info( |
|
|
`🔄 Adjusted expiry time to 10 minutes for Gemini account ${accountId} with refresh token` |
|
|
) |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if (updates.isActive === 'false' && existingAccount.isActive !== 'false') { |
|
|
try { |
|
|
const webhookNotifier = require('../utils/webhookNotifier') |
|
|
await webhookNotifier.sendAccountAnomalyNotification({ |
|
|
accountId, |
|
|
accountName: updates.name || existingAccount.name || 'Unknown Account', |
|
|
platform: 'gemini', |
|
|
status: 'disabled', |
|
|
errorCode: 'GEMINI_MANUALLY_DISABLED', |
|
|
reason: 'Account manually disabled by administrator' |
|
|
}) |
|
|
} catch (webhookError) { |
|
|
logger.error('Failed to send webhook notification for manual account disable:', webhookError) |
|
|
} |
|
|
} |
|
|
|
|
|
await client.hset(`${GEMINI_ACCOUNT_KEY_PREFIX}${accountId}`, updates) |
|
|
|
|
|
logger.info(`Updated Gemini account: ${accountId}`) |
|
|
|
|
|
|
|
|
const updatedAccount = { ...existingAccount, ...updates } |
|
|
|
|
|
|
|
|
if (updatedAccount.proxy && typeof updatedAccount.proxy === 'string') { |
|
|
try { |
|
|
updatedAccount.proxy = JSON.parse(updatedAccount.proxy) |
|
|
} catch (e) { |
|
|
updatedAccount.proxy = null |
|
|
} |
|
|
} |
|
|
|
|
|
return updatedAccount |
|
|
} |
|
|
|
|
|
|
|
|
async function deleteAccount(accountId) { |
|
|
const account = await getAccount(accountId) |
|
|
if (!account) { |
|
|
throw new Error('Account not found') |
|
|
} |
|
|
|
|
|
|
|
|
const client = redisClient.getClientSafe() |
|
|
await client.del(`${GEMINI_ACCOUNT_KEY_PREFIX}${accountId}`) |
|
|
|
|
|
|
|
|
if (account.accountType === 'shared') { |
|
|
await client.srem(SHARED_GEMINI_ACCOUNTS_KEY, accountId) |
|
|
} |
|
|
|
|
|
|
|
|
const sessionMappings = await client.keys(`${ACCOUNT_SESSION_MAPPING_PREFIX}*`) |
|
|
for (const key of sessionMappings) { |
|
|
const mappedAccountId = await client.get(key) |
|
|
if (mappedAccountId === accountId) { |
|
|
await client.del(key) |
|
|
} |
|
|
} |
|
|
|
|
|
logger.info(`Deleted Gemini account: ${accountId}`) |
|
|
return true |
|
|
} |
|
|
|
|
|
|
|
|
async function getAllAccounts() { |
|
|
const client = redisClient.getClientSafe() |
|
|
const keys = await client.keys(`${GEMINI_ACCOUNT_KEY_PREFIX}*`) |
|
|
const accounts = [] |
|
|
|
|
|
for (const key of keys) { |
|
|
const accountData = await client.hgetall(key) |
|
|
if (accountData && Object.keys(accountData).length > 0) { |
|
|
|
|
|
const rateLimitInfo = await getAccountRateLimitInfo(accountData.id) |
|
|
|
|
|
|
|
|
if (accountData.proxy) { |
|
|
try { |
|
|
accountData.proxy = JSON.parse(accountData.proxy) |
|
|
} catch (e) { |
|
|
|
|
|
accountData.proxy = null |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
accountData.schedulable = accountData.schedulable !== 'false' |
|
|
|
|
|
const tokenExpiresAt = accountData.expiresAt || null |
|
|
const subscriptionExpiresAt = |
|
|
accountData.subscriptionExpiresAt && accountData.subscriptionExpiresAt !== '' |
|
|
? accountData.subscriptionExpiresAt |
|
|
: null |
|
|
|
|
|
|
|
|
accounts.push({ |
|
|
...accountData, |
|
|
geminiOauth: accountData.geminiOauth ? '[ENCRYPTED]' : '', |
|
|
accessToken: accountData.accessToken ? '[ENCRYPTED]' : '', |
|
|
refreshToken: accountData.refreshToken ? '[ENCRYPTED]' : '', |
|
|
|
|
|
|
|
|
|
|
|
tokenExpiresAt, |
|
|
subscriptionExpiresAt, |
|
|
expiresAt: subscriptionExpiresAt, |
|
|
|
|
|
|
|
|
|
|
|
scopes: |
|
|
accountData.scopes && accountData.scopes.trim() ? accountData.scopes.split(' ') : [], |
|
|
|
|
|
hasRefreshToken: !!accountData.refreshToken, |
|
|
|
|
|
rateLimitStatus: rateLimitInfo |
|
|
? { |
|
|
isRateLimited: rateLimitInfo.isRateLimited, |
|
|
rateLimitedAt: rateLimitInfo.rateLimitedAt, |
|
|
minutesRemaining: rateLimitInfo.minutesRemaining |
|
|
} |
|
|
: { |
|
|
isRateLimited: false, |
|
|
rateLimitedAt: null, |
|
|
minutesRemaining: 0 |
|
|
} |
|
|
}) |
|
|
} |
|
|
} |
|
|
|
|
|
return accounts |
|
|
} |
|
|
|
|
|
|
|
|
async function selectAvailableAccount(apiKeyId, sessionHash = null) { |
|
|
|
|
|
const client = redisClient.getClientSafe() |
|
|
if (sessionHash) { |
|
|
const mappedAccountId = await client.get(`${ACCOUNT_SESSION_MAPPING_PREFIX}${sessionHash}`) |
|
|
|
|
|
if (mappedAccountId) { |
|
|
const account = await getAccount(mappedAccountId) |
|
|
if (account && account.isActive === 'true' && !isTokenExpired(account)) { |
|
|
logger.debug(`Using sticky session account: ${mappedAccountId}`) |
|
|
return account |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
const apiKeyData = await client.hgetall(`api_key:${apiKeyId}`) |
|
|
|
|
|
|
|
|
if (apiKeyData.geminiAccountId) { |
|
|
const account = await getAccount(apiKeyData.geminiAccountId) |
|
|
if (account && account.isActive === 'true') { |
|
|
|
|
|
const isExpired = isTokenExpired(account) |
|
|
|
|
|
|
|
|
logTokenUsage(account.id, account.name, 'gemini', account.expiresAt, isExpired) |
|
|
|
|
|
if (isExpired) { |
|
|
await refreshAccountToken(account.id) |
|
|
return await getAccount(account.id) |
|
|
} |
|
|
|
|
|
|
|
|
if (sessionHash) { |
|
|
await client.setex( |
|
|
`${ACCOUNT_SESSION_MAPPING_PREFIX}${sessionHash}`, |
|
|
3600, |
|
|
account.id |
|
|
) |
|
|
} |
|
|
|
|
|
return account |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
const sharedAccountIds = await client.smembers(SHARED_GEMINI_ACCOUNTS_KEY) |
|
|
const availableAccounts = [] |
|
|
|
|
|
for (const accountId of sharedAccountIds) { |
|
|
const account = await getAccount(accountId) |
|
|
if ( |
|
|
account && |
|
|
account.isActive === 'true' && |
|
|
!isRateLimited(account) && |
|
|
!isSubscriptionExpired(account) |
|
|
) { |
|
|
availableAccounts.push(account) |
|
|
} else if (account && isSubscriptionExpired(account)) { |
|
|
logger.debug( |
|
|
`⏰ Skipping expired Gemini account: ${account.name}, expired at ${account.subscriptionExpiresAt}` |
|
|
) |
|
|
} |
|
|
} |
|
|
|
|
|
if (availableAccounts.length === 0) { |
|
|
throw new Error('No available Gemini accounts') |
|
|
} |
|
|
|
|
|
|
|
|
availableAccounts.sort((a, b) => { |
|
|
const aLastUsed = a.lastUsedAt ? new Date(a.lastUsedAt).getTime() : 0 |
|
|
const bLastUsed = b.lastUsedAt ? new Date(b.lastUsedAt).getTime() : 0 |
|
|
return aLastUsed - bLastUsed |
|
|
}) |
|
|
|
|
|
const selectedAccount = availableAccounts[0] |
|
|
|
|
|
|
|
|
const isExpired = isTokenExpired(selectedAccount) |
|
|
|
|
|
|
|
|
logTokenUsage( |
|
|
selectedAccount.id, |
|
|
selectedAccount.name, |
|
|
'gemini', |
|
|
selectedAccount.expiresAt, |
|
|
isExpired |
|
|
) |
|
|
|
|
|
if (isExpired) { |
|
|
await refreshAccountToken(selectedAccount.id) |
|
|
return await getAccount(selectedAccount.id) |
|
|
} |
|
|
|
|
|
|
|
|
if (sessionHash) { |
|
|
await client.setex(`${ACCOUNT_SESSION_MAPPING_PREFIX}${sessionHash}`, 3600, selectedAccount.id) |
|
|
} |
|
|
|
|
|
return selectedAccount |
|
|
} |
|
|
|
|
|
|
|
|
function isTokenExpired(account) { |
|
|
if (!account.expiresAt) { |
|
|
return true |
|
|
} |
|
|
|
|
|
const expiryTime = new Date(account.expiresAt).getTime() |
|
|
const now = Date.now() |
|
|
const buffer = 10 * 1000 |
|
|
|
|
|
return now >= expiryTime - buffer |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
function isSubscriptionExpired(account) { |
|
|
if (!account.subscriptionExpiresAt) { |
|
|
return false |
|
|
} |
|
|
const expiryDate = new Date(account.subscriptionExpiresAt) |
|
|
return expiryDate <= new Date() |
|
|
} |
|
|
|
|
|
|
|
|
function isRateLimited(account) { |
|
|
if (account.rateLimitStatus === 'limited' && account.rateLimitedAt) { |
|
|
const limitedAt = new Date(account.rateLimitedAt).getTime() |
|
|
const now = Date.now() |
|
|
const limitDuration = 60 * 60 * 1000 |
|
|
|
|
|
return now < limitedAt + limitDuration |
|
|
} |
|
|
return false |
|
|
} |
|
|
|
|
|
|
|
|
async function refreshAccountToken(accountId) { |
|
|
let lockAcquired = false |
|
|
let account = null |
|
|
|
|
|
try { |
|
|
account = await getAccount(accountId) |
|
|
if (!account) { |
|
|
throw new Error('Account not found') |
|
|
} |
|
|
|
|
|
if (!account.refreshToken) { |
|
|
throw new Error('No refresh token available') |
|
|
} |
|
|
|
|
|
|
|
|
lockAcquired = await tokenRefreshService.acquireRefreshLock(accountId, 'gemini') |
|
|
|
|
|
if (!lockAcquired) { |
|
|
|
|
|
logger.info( |
|
|
`🔒 Token refresh already in progress for Gemini account: ${account.name} (${accountId})` |
|
|
) |
|
|
logRefreshSkipped(accountId, account.name, 'gemini', 'already_locked') |
|
|
|
|
|
|
|
|
await new Promise((resolve) => setTimeout(resolve, 2000)) |
|
|
|
|
|
|
|
|
const updatedAccount = await getAccount(accountId) |
|
|
if (updatedAccount && updatedAccount.accessToken) { |
|
|
const accessToken = decrypt(updatedAccount.accessToken) |
|
|
return { |
|
|
access_token: accessToken, |
|
|
refresh_token: updatedAccount.refreshToken ? decrypt(updatedAccount.refreshToken) : '', |
|
|
expiry_date: updatedAccount.expiresAt ? new Date(updatedAccount.expiresAt).getTime() : 0, |
|
|
scope: updatedAccount.scope || OAUTH_SCOPES.join(' '), |
|
|
token_type: 'Bearer' |
|
|
} |
|
|
} |
|
|
|
|
|
throw new Error('Token refresh in progress by another process') |
|
|
} |
|
|
|
|
|
|
|
|
logRefreshStart(accountId, account.name, 'gemini', 'manual_refresh') |
|
|
logger.info(`🔄 Starting token refresh for Gemini account: ${account.name} (${accountId})`) |
|
|
|
|
|
|
|
|
|
|
|
const newTokens = await refreshAccessToken(account.refreshToken, account.proxy) |
|
|
|
|
|
|
|
|
const updates = { |
|
|
accessToken: newTokens.access_token, |
|
|
refreshToken: newTokens.refresh_token || account.refreshToken, |
|
|
expiresAt: new Date(newTokens.expiry_date).toISOString(), |
|
|
lastRefreshAt: new Date().toISOString(), |
|
|
geminiOauth: JSON.stringify(newTokens), |
|
|
status: 'active', |
|
|
errorMessage: '' |
|
|
} |
|
|
|
|
|
await updateAccount(accountId, updates) |
|
|
|
|
|
|
|
|
logRefreshSuccess(accountId, account.name, 'gemini', { |
|
|
accessToken: newTokens.access_token, |
|
|
refreshToken: newTokens.refresh_token, |
|
|
expiresAt: newTokens.expiry_date, |
|
|
scopes: newTokens.scope |
|
|
}) |
|
|
|
|
|
logger.info( |
|
|
`Refreshed token for Gemini account: ${accountId} - Access Token: ${maskToken(newTokens.access_token)}` |
|
|
) |
|
|
|
|
|
return newTokens |
|
|
} catch (error) { |
|
|
|
|
|
logRefreshError(accountId, account ? account.name : 'Unknown', 'gemini', error) |
|
|
|
|
|
logger.error(`Failed to refresh token for account ${accountId}:`, error) |
|
|
|
|
|
|
|
|
if (account) { |
|
|
try { |
|
|
await updateAccount(accountId, { |
|
|
status: 'error', |
|
|
errorMessage: error.message |
|
|
}) |
|
|
|
|
|
|
|
|
try { |
|
|
const webhookNotifier = require('../utils/webhookNotifier') |
|
|
await webhookNotifier.sendAccountAnomalyNotification({ |
|
|
accountId, |
|
|
accountName: account.name, |
|
|
platform: 'gemini', |
|
|
status: 'error', |
|
|
errorCode: 'GEMINI_ERROR', |
|
|
reason: `Token refresh failed: ${error.message}` |
|
|
}) |
|
|
} catch (webhookError) { |
|
|
logger.error('Failed to send webhook notification:', webhookError) |
|
|
} |
|
|
} catch (updateError) { |
|
|
logger.error('Failed to update account status after refresh error:', updateError) |
|
|
} |
|
|
} |
|
|
|
|
|
throw error |
|
|
} finally { |
|
|
|
|
|
if (lockAcquired) { |
|
|
await tokenRefreshService.releaseRefreshLock(accountId, 'gemini') |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function markAccountUsed(accountId) { |
|
|
await updateAccount(accountId, { |
|
|
lastUsedAt: new Date().toISOString() |
|
|
}) |
|
|
} |
|
|
|
|
|
|
|
|
async function setAccountRateLimited(accountId, isLimited = true) { |
|
|
const updates = isLimited |
|
|
? { |
|
|
rateLimitStatus: 'limited', |
|
|
rateLimitedAt: new Date().toISOString() |
|
|
} |
|
|
: { |
|
|
rateLimitStatus: '', |
|
|
rateLimitedAt: '' |
|
|
} |
|
|
|
|
|
await updateAccount(accountId, updates) |
|
|
} |
|
|
|
|
|
|
|
|
async function getAccountRateLimitInfo(accountId) { |
|
|
try { |
|
|
const account = await getAccount(accountId) |
|
|
if (!account) { |
|
|
return null |
|
|
} |
|
|
|
|
|
if (account.rateLimitStatus === 'limited' && account.rateLimitedAt) { |
|
|
const rateLimitedAt = new Date(account.rateLimitedAt) |
|
|
const now = new Date() |
|
|
const minutesSinceRateLimit = Math.floor((now - rateLimitedAt) / (1000 * 60)) |
|
|
|
|
|
|
|
|
const minutesRemaining = Math.max(0, 60 - minutesSinceRateLimit) |
|
|
const rateLimitEndAt = new Date(rateLimitedAt.getTime() + 60 * 60 * 1000).toISOString() |
|
|
|
|
|
return { |
|
|
isRateLimited: minutesRemaining > 0, |
|
|
rateLimitedAt: account.rateLimitedAt, |
|
|
minutesSinceRateLimit, |
|
|
minutesRemaining, |
|
|
rateLimitEndAt |
|
|
} |
|
|
} |
|
|
|
|
|
return { |
|
|
isRateLimited: false, |
|
|
rateLimitedAt: null, |
|
|
minutesSinceRateLimit: 0, |
|
|
minutesRemaining: 0, |
|
|
rateLimitEndAt: null |
|
|
} |
|
|
} catch (error) { |
|
|
logger.error(`❌ Failed to get rate limit info for Gemini account: ${accountId}`, error) |
|
|
return null |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function getOauthClient(accessToken, refreshToken, proxyConfig = null) { |
|
|
const client = createOAuth2Client(null, proxyConfig) |
|
|
|
|
|
const creds = { |
|
|
access_token: accessToken, |
|
|
refresh_token: refreshToken, |
|
|
scope: |
|
|
'https://www.googleapis.com/auth/cloud-platform https://www.googleapis.com/auth/userinfo.profile openid https://www.googleapis.com/auth/userinfo.email', |
|
|
token_type: 'Bearer', |
|
|
expiry_date: 1754269905646 |
|
|
} |
|
|
|
|
|
if (proxyConfig) { |
|
|
logger.info( |
|
|
`🌐 Using proxy for Gemini OAuth client: ${ProxyHelper.getProxyDescription(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🌐 No proxy configured for Gemini OAuth client') |
|
|
} |
|
|
|
|
|
|
|
|
client.setCredentials(creds) |
|
|
|
|
|
|
|
|
const { token } = await client.getAccessToken() |
|
|
|
|
|
if (!token) { |
|
|
return false |
|
|
} |
|
|
|
|
|
|
|
|
await client.getTokenInfo(token) |
|
|
|
|
|
logger.info('✅ OAuth客户端已创建') |
|
|
return client |
|
|
} |
|
|
|
|
|
|
|
|
async function loadCodeAssist(client, projectId = null, proxyConfig = null) { |
|
|
const axios = require('axios') |
|
|
const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com' |
|
|
const CODE_ASSIST_API_VERSION = 'v1internal' |
|
|
|
|
|
const { token } = await client.getAccessToken() |
|
|
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig) |
|
|
|
|
|
const tokenInfoConfig = { |
|
|
url: 'https://oauth2.googleapis.com/tokeninfo', |
|
|
method: 'POST', |
|
|
headers: { |
|
|
Authorization: `Bearer ${token}`, |
|
|
'Content-Type': 'application/x-www-form-urlencoded' |
|
|
}, |
|
|
data: new URLSearchParams({ access_token: token }).toString(), |
|
|
timeout: 15000 |
|
|
} |
|
|
|
|
|
if (proxyAgent) { |
|
|
tokenInfoConfig.httpAgent = proxyAgent |
|
|
tokenInfoConfig.httpsAgent = proxyAgent |
|
|
tokenInfoConfig.proxy = false |
|
|
} |
|
|
|
|
|
try { |
|
|
await axios(tokenInfoConfig) |
|
|
logger.info('📋 tokeninfo 接口验证成功') |
|
|
} catch (error) { |
|
|
logger.info('tokeninfo 接口获取失败', error) |
|
|
} |
|
|
|
|
|
const userInfoConfig = { |
|
|
url: 'https://www.googleapis.com/oauth2/v2/userinfo', |
|
|
method: 'GET', |
|
|
headers: { |
|
|
Authorization: `Bearer ${token}`, |
|
|
Accept: '*/*' |
|
|
}, |
|
|
timeout: 15000 |
|
|
} |
|
|
|
|
|
if (proxyAgent) { |
|
|
userInfoConfig.httpAgent = proxyAgent |
|
|
userInfoConfig.httpsAgent = proxyAgent |
|
|
userInfoConfig.proxy = false |
|
|
} |
|
|
|
|
|
try { |
|
|
await axios(userInfoConfig) |
|
|
logger.info('📋 userinfo 接口获取成功') |
|
|
} catch (error) { |
|
|
logger.info('userinfo 接口获取失败', error) |
|
|
} |
|
|
|
|
|
|
|
|
const clientMetadata = { |
|
|
ideType: 'IDE_UNSPECIFIED', |
|
|
platform: 'PLATFORM_UNSPECIFIED', |
|
|
pluginType: 'GEMINI' |
|
|
} |
|
|
|
|
|
|
|
|
if (projectId) { |
|
|
clientMetadata.duetProject = projectId |
|
|
} |
|
|
|
|
|
const request = { |
|
|
metadata: clientMetadata |
|
|
} |
|
|
|
|
|
|
|
|
if (projectId) { |
|
|
request.cloudaicompanionProject = projectId |
|
|
} |
|
|
|
|
|
const axiosConfig = { |
|
|
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:loadCodeAssist`, |
|
|
method: 'POST', |
|
|
headers: { |
|
|
Authorization: `Bearer ${token}`, |
|
|
'Content-Type': 'application/json' |
|
|
}, |
|
|
data: request, |
|
|
timeout: 30000 |
|
|
} |
|
|
|
|
|
|
|
|
if (proxyAgent) { |
|
|
axiosConfig.httpAgent = proxyAgent |
|
|
axiosConfig.httpsAgent = proxyAgent |
|
|
axiosConfig.proxy = false |
|
|
logger.info( |
|
|
`🌐 Using proxy for Gemini loadCodeAssist: ${ProxyHelper.getProxyDescription(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🌐 No proxy configured for Gemini loadCodeAssist') |
|
|
} |
|
|
|
|
|
const response = await axios(axiosConfig) |
|
|
|
|
|
logger.info('📋 loadCodeAssist API调用成功') |
|
|
return response.data |
|
|
} |
|
|
|
|
|
|
|
|
function getOnboardTier(loadRes) { |
|
|
|
|
|
const UserTierId = { |
|
|
LEGACY: 'LEGACY', |
|
|
FREE: 'FREE', |
|
|
PRO: 'PRO' |
|
|
} |
|
|
|
|
|
if (loadRes.currentTier) { |
|
|
return loadRes.currentTier |
|
|
} |
|
|
|
|
|
for (const tier of loadRes.allowedTiers || []) { |
|
|
if (tier.isDefault) { |
|
|
return tier |
|
|
} |
|
|
} |
|
|
|
|
|
return { |
|
|
name: '', |
|
|
description: '', |
|
|
id: UserTierId.LEGACY, |
|
|
userDefinedCloudaicompanionProject: true |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
async function onboardUser(client, tierId, projectId, clientMetadata, proxyConfig = null) { |
|
|
const axios = require('axios') |
|
|
const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com' |
|
|
const CODE_ASSIST_API_VERSION = 'v1internal' |
|
|
|
|
|
const { token } = await client.getAccessToken() |
|
|
|
|
|
const onboardReq = { |
|
|
tierId, |
|
|
metadata: clientMetadata |
|
|
} |
|
|
|
|
|
|
|
|
if (projectId) { |
|
|
onboardReq.cloudaicompanionProject = projectId |
|
|
} |
|
|
|
|
|
|
|
|
const baseAxiosConfig = { |
|
|
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:onboardUser`, |
|
|
method: 'POST', |
|
|
headers: { |
|
|
Authorization: `Bearer ${token}`, |
|
|
'Content-Type': 'application/json' |
|
|
}, |
|
|
data: onboardReq, |
|
|
timeout: 30000 |
|
|
} |
|
|
|
|
|
|
|
|
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig) |
|
|
if (proxyAgent) { |
|
|
baseAxiosConfig.httpAgent = proxyAgent |
|
|
baseAxiosConfig.httpsAgent = proxyAgent |
|
|
baseAxiosConfig.proxy = false |
|
|
logger.info( |
|
|
`🌐 Using proxy for Gemini onboardUser: ${ProxyHelper.getProxyDescription(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🌐 No proxy configured for Gemini onboardUser') |
|
|
} |
|
|
|
|
|
logger.info('📋 开始onboardUser API调用', { |
|
|
tierId, |
|
|
projectId, |
|
|
hasProjectId: !!projectId, |
|
|
isFreeTier: tierId === 'free-tier' || tierId === 'FREE' |
|
|
}) |
|
|
|
|
|
|
|
|
let lroRes = await axios(baseAxiosConfig) |
|
|
|
|
|
let attempts = 0 |
|
|
const maxAttempts = 12 |
|
|
|
|
|
while (!lroRes.data.done && attempts < maxAttempts) { |
|
|
logger.info(`⏳ 等待onboardUser完成... (${attempts + 1}/${maxAttempts})`) |
|
|
await new Promise((resolve) => setTimeout(resolve, 5000)) |
|
|
|
|
|
lroRes = await axios(baseAxiosConfig) |
|
|
attempts++ |
|
|
} |
|
|
|
|
|
if (!lroRes.data.done) { |
|
|
throw new Error('onboardUser操作超时') |
|
|
} |
|
|
|
|
|
logger.info('✅ onboardUser API调用完成') |
|
|
return lroRes.data |
|
|
} |
|
|
|
|
|
|
|
|
async function setupUser( |
|
|
client, |
|
|
initialProjectId = null, |
|
|
clientMetadata = null, |
|
|
proxyConfig = null |
|
|
) { |
|
|
logger.info('🚀 setupUser 开始', { initialProjectId, hasClientMetadata: !!clientMetadata }) |
|
|
|
|
|
let projectId = initialProjectId || process.env.GOOGLE_CLOUD_PROJECT || null |
|
|
logger.info('📋 初始项目ID', { projectId, fromEnv: !!process.env.GOOGLE_CLOUD_PROJECT }) |
|
|
|
|
|
|
|
|
if (!clientMetadata) { |
|
|
clientMetadata = { |
|
|
ideType: 'IDE_UNSPECIFIED', |
|
|
platform: 'PLATFORM_UNSPECIFIED', |
|
|
pluginType: 'GEMINI', |
|
|
duetProject: projectId |
|
|
} |
|
|
logger.info('🔧 使用默认 ClientMetadata') |
|
|
} |
|
|
|
|
|
|
|
|
logger.info('📞 调用 loadCodeAssist...') |
|
|
const loadRes = await loadCodeAssist(client, projectId, proxyConfig) |
|
|
logger.info('✅ loadCodeAssist 完成', { |
|
|
hasCloudaicompanionProject: !!loadRes.cloudaicompanionProject |
|
|
}) |
|
|
|
|
|
|
|
|
if (!projectId && loadRes.cloudaicompanionProject) { |
|
|
projectId = loadRes.cloudaicompanionProject |
|
|
logger.info('📋 从 loadCodeAssist 获取项目ID', { projectId }) |
|
|
} |
|
|
|
|
|
const tier = getOnboardTier(loadRes) |
|
|
logger.info('🎯 获取用户层级', { |
|
|
tierId: tier.id, |
|
|
userDefinedProject: tier.userDefinedCloudaicompanionProject |
|
|
}) |
|
|
|
|
|
if (tier.userDefinedCloudaiCompanionProject && !projectId) { |
|
|
throw new Error('此账号需要设置GOOGLE_CLOUD_PROJECT环境变量或提供projectId') |
|
|
} |
|
|
|
|
|
|
|
|
logger.info('📞 调用 onboardUser...', { tierId: tier.id, projectId }) |
|
|
const lroRes = await onboardUser(client, tier.id, projectId, clientMetadata, proxyConfig) |
|
|
logger.info('✅ onboardUser 完成', { hasDone: !!lroRes.done, hasResponse: !!lroRes.response }) |
|
|
|
|
|
const result = { |
|
|
projectId: lroRes.response?.cloudaicompanionProject?.id || projectId || '', |
|
|
userTier: tier.id, |
|
|
loadRes, |
|
|
onboardRes: lroRes.response || {} |
|
|
} |
|
|
|
|
|
logger.info('🎯 setupUser 完成', { resultProjectId: result.projectId, userTier: result.userTier }) |
|
|
return result |
|
|
} |
|
|
|
|
|
|
|
|
async function countTokens(client, contents, model = 'gemini-2.0-flash-exp', proxyConfig = null) { |
|
|
const axios = require('axios') |
|
|
const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com' |
|
|
const CODE_ASSIST_API_VERSION = 'v1internal' |
|
|
|
|
|
const { token } = await client.getAccessToken() |
|
|
|
|
|
|
|
|
const request = { |
|
|
request: { |
|
|
model: `models/${model}`, |
|
|
contents |
|
|
} |
|
|
} |
|
|
|
|
|
logger.info('📊 countTokens API调用开始', { model, contentsLength: contents.length }) |
|
|
|
|
|
const axiosConfig = { |
|
|
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:countTokens`, |
|
|
method: 'POST', |
|
|
headers: { |
|
|
Authorization: `Bearer ${token}`, |
|
|
'Content-Type': 'application/json' |
|
|
}, |
|
|
data: request, |
|
|
timeout: 30000 |
|
|
} |
|
|
|
|
|
|
|
|
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig) |
|
|
if (proxyAgent) { |
|
|
axiosConfig.httpAgent = proxyAgent |
|
|
axiosConfig.httpsAgent = proxyAgent |
|
|
axiosConfig.proxy = false |
|
|
logger.info( |
|
|
`🌐 Using proxy for Gemini countTokens: ${ProxyHelper.getProxyDescription(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🌐 No proxy configured for Gemini countTokens') |
|
|
} |
|
|
|
|
|
const response = await axios(axiosConfig) |
|
|
|
|
|
logger.info('✅ countTokens API调用成功', { totalTokens: response.data.totalTokens }) |
|
|
return response.data |
|
|
} |
|
|
|
|
|
|
|
|
async function generateContent( |
|
|
client, |
|
|
requestData, |
|
|
userPromptId, |
|
|
projectId = null, |
|
|
sessionId = null, |
|
|
proxyConfig = null |
|
|
) { |
|
|
const axios = require('axios') |
|
|
const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com' |
|
|
const CODE_ASSIST_API_VERSION = 'v1internal' |
|
|
|
|
|
const { token } = await client.getAccessToken() |
|
|
|
|
|
|
|
|
const request = { |
|
|
model: requestData.model, |
|
|
request: { |
|
|
...requestData.request, |
|
|
session_id: sessionId |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if (userPromptId) { |
|
|
request.user_prompt_id = userPromptId |
|
|
} |
|
|
|
|
|
|
|
|
if (projectId) { |
|
|
request.project = projectId |
|
|
} |
|
|
|
|
|
logger.info('🤖 generateContent API调用开始', { |
|
|
model: requestData.model, |
|
|
userPromptId, |
|
|
projectId, |
|
|
sessionId |
|
|
}) |
|
|
|
|
|
|
|
|
logger.info('📦 generateContent 请求详情', { |
|
|
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:generateContent`, |
|
|
requestBody: JSON.stringify(request, null, 2) |
|
|
}) |
|
|
|
|
|
const axiosConfig = { |
|
|
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:generateContent`, |
|
|
method: 'POST', |
|
|
headers: { |
|
|
Authorization: `Bearer ${token}`, |
|
|
'Content-Type': 'application/json' |
|
|
}, |
|
|
data: request, |
|
|
timeout: 60000 |
|
|
} |
|
|
|
|
|
|
|
|
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig) |
|
|
if (proxyAgent) { |
|
|
axiosConfig.httpAgent = proxyAgent |
|
|
axiosConfig.httpsAgent = proxyAgent |
|
|
axiosConfig.proxy = false |
|
|
logger.info( |
|
|
`🌐 Using proxy for Gemini generateContent: ${ProxyHelper.getProxyDescription(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🌐 No proxy configured for Gemini generateContent') |
|
|
} |
|
|
|
|
|
const response = await axios(axiosConfig) |
|
|
|
|
|
logger.info('✅ generateContent API调用成功') |
|
|
return response.data |
|
|
} |
|
|
|
|
|
|
|
|
async function generateContentStream( |
|
|
client, |
|
|
requestData, |
|
|
userPromptId, |
|
|
projectId = null, |
|
|
sessionId = null, |
|
|
signal = null, |
|
|
proxyConfig = null |
|
|
) { |
|
|
const axios = require('axios') |
|
|
const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com' |
|
|
const CODE_ASSIST_API_VERSION = 'v1internal' |
|
|
|
|
|
const { token } = await client.getAccessToken() |
|
|
|
|
|
|
|
|
const request = { |
|
|
model: requestData.model, |
|
|
request: { |
|
|
...requestData.request, |
|
|
session_id: sessionId |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if (userPromptId) { |
|
|
request.user_prompt_id = userPromptId |
|
|
} |
|
|
|
|
|
|
|
|
if (projectId) { |
|
|
request.project = projectId |
|
|
} |
|
|
|
|
|
logger.info('🌊 streamGenerateContent API调用开始', { |
|
|
model: requestData.model, |
|
|
userPromptId, |
|
|
projectId, |
|
|
sessionId |
|
|
}) |
|
|
|
|
|
const axiosConfig = { |
|
|
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:streamGenerateContent`, |
|
|
method: 'POST', |
|
|
params: { |
|
|
alt: 'sse' |
|
|
}, |
|
|
headers: { |
|
|
Authorization: `Bearer ${token}`, |
|
|
'Content-Type': 'application/json' |
|
|
}, |
|
|
data: request, |
|
|
responseType: 'stream', |
|
|
timeout: 60000 |
|
|
} |
|
|
|
|
|
|
|
|
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig) |
|
|
if (proxyAgent) { |
|
|
axiosConfig.httpAgent = proxyAgent |
|
|
axiosConfig.httpsAgent = proxyAgent |
|
|
axiosConfig.proxy = false |
|
|
logger.info( |
|
|
`🌐 Using proxy for Gemini streamGenerateContent: ${ProxyHelper.getProxyDescription(proxyConfig)}` |
|
|
) |
|
|
} else { |
|
|
logger.debug('🌐 No proxy configured for Gemini streamGenerateContent') |
|
|
} |
|
|
|
|
|
|
|
|
if (signal) { |
|
|
axiosConfig.signal = signal |
|
|
} |
|
|
|
|
|
const response = await axios(axiosConfig) |
|
|
|
|
|
logger.info('✅ streamGenerateContent API调用成功,开始流式传输') |
|
|
return response.data |
|
|
} |
|
|
|
|
|
|
|
|
async function updateTempProjectId(accountId, tempProjectId) { |
|
|
if (!tempProjectId) { |
|
|
return |
|
|
} |
|
|
|
|
|
try { |
|
|
const account = await getAccount(accountId) |
|
|
if (!account) { |
|
|
logger.warn(`Account ${accountId} not found when updating tempProjectId`) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
if (!account.projectId && tempProjectId !== account.tempProjectId) { |
|
|
await updateAccount(accountId, { tempProjectId }) |
|
|
logger.info(`Updated tempProjectId for account ${accountId}: ${tempProjectId}`) |
|
|
} |
|
|
} catch (error) { |
|
|
logger.error(`Failed to update tempProjectId for account ${accountId}:`, error) |
|
|
} |
|
|
} |
|
|
|
|
|
module.exports = { |
|
|
generateAuthUrl, |
|
|
pollAuthorizationStatus, |
|
|
exchangeCodeForTokens, |
|
|
refreshAccessToken, |
|
|
createAccount, |
|
|
getAccount, |
|
|
updateAccount, |
|
|
deleteAccount, |
|
|
getAllAccounts, |
|
|
selectAvailableAccount, |
|
|
refreshAccountToken, |
|
|
markAccountUsed, |
|
|
setAccountRateLimited, |
|
|
getAccountRateLimitInfo, |
|
|
isTokenExpired, |
|
|
getOauthClient, |
|
|
loadCodeAssist, |
|
|
getOnboardTier, |
|
|
onboardUser, |
|
|
setupUser, |
|
|
encrypt, |
|
|
decrypt, |
|
|
generateEncryptionKey, |
|
|
decryptCache, |
|
|
countTokens, |
|
|
generateContent, |
|
|
generateContentStream, |
|
|
updateTempProjectId, |
|
|
OAUTH_CLIENT_ID, |
|
|
OAUTH_SCOPES |
|
|
} |
|
|
|