mirror of
https://github.com/Wei-Shaw/claude-relay-service.git
synced 2026-01-23 21:17:30 +00:00
Merge branch 'antigravity'
This commit is contained in:
shaw
764
src/services/accountBalanceService.js
Normal file
764
src/services/accountBalanceService.js
Normal file
@@ -0,0 +1,764 @@
|
||||
const redis = require('../models/redis')
|
||||
const balanceScriptService = require('./balanceScriptService')
|
||||
const logger = require('../utils/logger')
|
||||
const CostCalculator = require('../utils/costCalculator')
|
||||
const { isBalanceScriptEnabled } = require('../utils/featureFlags')
|
||||
|
||||
class AccountBalanceService {
|
||||
constructor(options = {}) {
|
||||
this.redis = options.redis || redis
|
||||
this.logger = options.logger || logger
|
||||
|
||||
this.providers = new Map()
|
||||
|
||||
this.CACHE_TTL_SECONDS = 3600
|
||||
this.LOCAL_TTL_SECONDS = 300
|
||||
|
||||
this.LOW_BALANCE_THRESHOLD = 10
|
||||
this.HIGH_USAGE_THRESHOLD_PERCENT = 90
|
||||
this.DEFAULT_CONCURRENCY = 10
|
||||
}
|
||||
|
||||
getSupportedPlatforms() {
|
||||
return [
|
||||
'claude',
|
||||
'claude-console',
|
||||
'gemini',
|
||||
'gemini-api',
|
||||
'openai',
|
||||
'openai-responses',
|
||||
'azure_openai',
|
||||
'bedrock',
|
||||
'droid',
|
||||
'ccr'
|
||||
]
|
||||
}
|
||||
|
||||
normalizePlatform(platform) {
|
||||
if (!platform) {
|
||||
return null
|
||||
}
|
||||
|
||||
const value = String(platform).trim().toLowerCase()
|
||||
|
||||
// 兼容实施文档与历史命名
|
||||
if (value === 'claude-official') {
|
||||
return 'claude'
|
||||
}
|
||||
if (value === 'azure-openai') {
|
||||
return 'azure_openai'
|
||||
}
|
||||
|
||||
// 保持前端平台键一致
|
||||
return value
|
||||
}
|
||||
|
||||
registerProvider(platform, provider) {
|
||||
const normalized = this.normalizePlatform(platform)
|
||||
if (!normalized) {
|
||||
throw new Error('registerProvider: 缺少 platform')
|
||||
}
|
||||
if (!provider || typeof provider.queryBalance !== 'function') {
|
||||
throw new Error(`registerProvider: Provider 无效 (${normalized})`)
|
||||
}
|
||||
this.providers.set(normalized, provider)
|
||||
}
|
||||
|
||||
async getAccountBalance(accountId, platform, options = {}) {
|
||||
const normalizedPlatform = this.normalizePlatform(platform)
|
||||
const account = await this.getAccount(accountId, normalizedPlatform)
|
||||
if (!account) {
|
||||
return null
|
||||
}
|
||||
return await this._getAccountBalanceForAccount(account, normalizedPlatform, options)
|
||||
}
|
||||
|
||||
async refreshAccountBalance(accountId, platform) {
|
||||
const normalizedPlatform = this.normalizePlatform(platform)
|
||||
const account = await this.getAccount(accountId, normalizedPlatform)
|
||||
if (!account) {
|
||||
return null
|
||||
}
|
||||
|
||||
return await this._getAccountBalanceForAccount(account, normalizedPlatform, {
|
||||
queryApi: true,
|
||||
useCache: false
|
||||
})
|
||||
}
|
||||
|
||||
async getAllAccountsBalance(platform, options = {}) {
|
||||
const normalizedPlatform = this.normalizePlatform(platform)
|
||||
const accounts = await this.getAllAccountsByPlatform(normalizedPlatform)
|
||||
const queryApi = this._parseBoolean(options.queryApi) || false
|
||||
const useCache = options.useCache !== false
|
||||
|
||||
const results = await this._mapWithConcurrency(
|
||||
accounts,
|
||||
this.DEFAULT_CONCURRENCY,
|
||||
async (acc) => {
|
||||
try {
|
||||
const balance = await this._getAccountBalanceForAccount(acc, normalizedPlatform, {
|
||||
queryApi,
|
||||
useCache
|
||||
})
|
||||
return { ...balance, name: acc.name || '' }
|
||||
} catch (error) {
|
||||
this.logger.error(`批量获取余额失败: ${normalizedPlatform}:${acc?.id}`, error)
|
||||
return {
|
||||
success: true,
|
||||
data: {
|
||||
accountId: acc?.id,
|
||||
platform: normalizedPlatform,
|
||||
balance: null,
|
||||
quota: null,
|
||||
statistics: {},
|
||||
source: 'local',
|
||||
lastRefreshAt: new Date().toISOString(),
|
||||
cacheExpiresAt: null,
|
||||
status: 'error',
|
||||
error: error.message || '批量查询失败'
|
||||
},
|
||||
name: acc?.name || ''
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
async getBalanceSummary() {
|
||||
const platforms = this.getSupportedPlatforms()
|
||||
|
||||
const summary = {
|
||||
totalBalance: 0,
|
||||
totalCost: 0,
|
||||
lowBalanceCount: 0,
|
||||
platforms: {}
|
||||
}
|
||||
|
||||
for (const platform of platforms) {
|
||||
const accounts = await this.getAllAccountsByPlatform(platform)
|
||||
const platformData = {
|
||||
count: accounts.length,
|
||||
totalBalance: 0,
|
||||
totalCost: 0,
|
||||
lowBalanceCount: 0,
|
||||
accounts: []
|
||||
}
|
||||
|
||||
const balances = await this._mapWithConcurrency(
|
||||
accounts,
|
||||
this.DEFAULT_CONCURRENCY,
|
||||
async (acc) => {
|
||||
const balance = await this._getAccountBalanceForAccount(acc, platform, {
|
||||
queryApi: false,
|
||||
useCache: true
|
||||
})
|
||||
return { ...balance, name: acc.name || '' }
|
||||
}
|
||||
)
|
||||
|
||||
for (const item of balances) {
|
||||
platformData.accounts.push(item)
|
||||
|
||||
const amount = item?.data?.balance?.amount
|
||||
const percentage = item?.data?.quota?.percentage
|
||||
const totalCost = Number(item?.data?.statistics?.totalCost || 0)
|
||||
|
||||
const hasAmount = typeof amount === 'number' && Number.isFinite(amount)
|
||||
const isLowBalance = hasAmount && amount < this.LOW_BALANCE_THRESHOLD
|
||||
const isHighUsage =
|
||||
typeof percentage === 'number' &&
|
||||
Number.isFinite(percentage) &&
|
||||
percentage > this.HIGH_USAGE_THRESHOLD_PERCENT
|
||||
|
||||
if (hasAmount) {
|
||||
platformData.totalBalance += amount
|
||||
}
|
||||
|
||||
if (isLowBalance || isHighUsage) {
|
||||
platformData.lowBalanceCount += 1
|
||||
summary.lowBalanceCount += 1
|
||||
}
|
||||
|
||||
platformData.totalCost += totalCost
|
||||
}
|
||||
|
||||
summary.platforms[platform] = platformData
|
||||
summary.totalBalance += platformData.totalBalance
|
||||
summary.totalCost += platformData.totalCost
|
||||
}
|
||||
|
||||
return summary
|
||||
}
|
||||
|
||||
async clearCache(accountId, platform) {
|
||||
const normalizedPlatform = this.normalizePlatform(platform)
|
||||
if (!normalizedPlatform) {
|
||||
throw new Error('缺少 platform 参数')
|
||||
}
|
||||
|
||||
await this.redis.deleteAccountBalance(normalizedPlatform, accountId)
|
||||
this.logger.info(`余额缓存已清除: ${normalizedPlatform}:${accountId}`)
|
||||
}
|
||||
|
||||
async getAccount(accountId, platform) {
|
||||
if (!accountId || !platform) {
|
||||
return null
|
||||
}
|
||||
|
||||
const serviceMap = {
|
||||
claude: require('./claudeAccountService'),
|
||||
'claude-console': require('./claudeConsoleAccountService'),
|
||||
gemini: require('./geminiAccountService'),
|
||||
'gemini-api': require('./geminiApiAccountService'),
|
||||
openai: require('./openaiAccountService'),
|
||||
'openai-responses': require('./openaiResponsesAccountService'),
|
||||
azure_openai: require('./azureOpenaiAccountService'),
|
||||
bedrock: require('./bedrockAccountService'),
|
||||
droid: require('./droidAccountService'),
|
||||
ccr: require('./ccrAccountService')
|
||||
}
|
||||
|
||||
const service = serviceMap[platform]
|
||||
if (!service || typeof service.getAccount !== 'function') {
|
||||
return null
|
||||
}
|
||||
|
||||
return await service.getAccount(accountId)
|
||||
}
|
||||
|
||||
async getAllAccountsByPlatform(platform) {
|
||||
if (!platform) {
|
||||
return []
|
||||
}
|
||||
|
||||
const serviceMap = {
|
||||
claude: require('./claudeAccountService'),
|
||||
'claude-console': require('./claudeConsoleAccountService'),
|
||||
gemini: require('./geminiAccountService'),
|
||||
'gemini-api': require('./geminiApiAccountService'),
|
||||
openai: require('./openaiAccountService'),
|
||||
'openai-responses': require('./openaiResponsesAccountService'),
|
||||
azure_openai: require('./azureOpenaiAccountService'),
|
||||
bedrock: require('./bedrockAccountService'),
|
||||
droid: require('./droidAccountService'),
|
||||
ccr: require('./ccrAccountService')
|
||||
}
|
||||
|
||||
const service = serviceMap[platform]
|
||||
if (!service) {
|
||||
return []
|
||||
}
|
||||
|
||||
// Bedrock 特殊：返回 { success, data }
|
||||
if (platform === 'bedrock' && typeof service.getAllAccounts === 'function') {
|
||||
const result = await service.getAllAccounts()
|
||||
return result?.success ? result.data || [] : []
|
||||
}
|
||||
|
||||
if (platform === 'openai-responses') {
|
||||
return await service.getAllAccounts(true)
|
||||
}
|
||||
|
||||
if (typeof service.getAllAccounts !== 'function') {
|
||||
return []
|
||||
}
|
||||
|
||||
return await service.getAllAccounts()
|
||||
}
|
||||
|
||||
async _getAccountBalanceForAccount(account, platform, options = {}) {
|
||||
const queryMode = this._parseQueryMode(options.queryApi)
|
||||
const useCache = options.useCache !== false
|
||||
|
||||
const accountId = account?.id
|
||||
if (!accountId) {
|
||||
throw new Error('账户缺少 id')
|
||||
}
|
||||
|
||||
// 余额脚本配置状态（用于前端控制“刷新余额”按钮）
|
||||
let scriptConfig = null
|
||||
let scriptConfigured = false
|
||||
if (typeof this.redis?.getBalanceScriptConfig === 'function') {
|
||||
scriptConfig = await this.redis.getBalanceScriptConfig(platform, accountId)
|
||||
scriptConfigured = !!(
|
||||
scriptConfig &&
|
||||
scriptConfig.scriptBody &&
|
||||
String(scriptConfig.scriptBody).trim().length > 0
|
||||
)
|
||||
}
|
||||
const scriptEnabled = isBalanceScriptEnabled()
|
||||
const scriptMeta = { scriptEnabled, scriptConfigured }
|
||||
|
||||
const localBalance = await this._getBalanceFromLocal(accountId, platform)
|
||||
const localStatistics = localBalance.statistics || {}
|
||||
|
||||
const quotaFromLocal = this._buildQuotaFromLocal(account, localStatistics)
|
||||
|
||||
// 安全限制：queryApi=auto 仅用于 Antigravity（gemini + oauthProvider=antigravity）账户
|
||||
const effectiveQueryMode =
|
||||
queryMode === 'auto' && !(platform === 'gemini' && account?.oauthProvider === 'antigravity')
|
||||
? 'local'
|
||||
: queryMode
|
||||
|
||||
// local: 仅本地统计/缓存；auto: 优先缓存，无缓存则尝试远程 Provider（并缓存结果）
|
||||
if (effectiveQueryMode !== 'api') {
|
||||
if (useCache) {
|
||||
const cached = await this.redis.getAccountBalance(platform, accountId)
|
||||
if (cached && cached.status === 'success') {
|
||||
return this._buildResponse(
|
||||
{
|
||||
status: cached.status,
|
||||
errorMessage: cached.errorMessage,
|
||||
balance: quotaFromLocal.balance ?? cached.balance,
|
||||
currency: quotaFromLocal.currency || cached.currency || 'USD',
|
||||
quota: quotaFromLocal.quota || cached.quota || null,
|
||||
statistics: localStatistics,
|
||||
lastRefreshAt: cached.lastRefreshAt
|
||||
},
|
||||
accountId,
|
||||
platform,
|
||||
'cache',
|
||||
cached.ttlSeconds,
|
||||
scriptMeta
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if (effectiveQueryMode === 'local') {
|
||||
return this._buildResponse(
|
||||
{
|
||||
status: 'success',
|
||||
errorMessage: null,
|
||||
balance: quotaFromLocal.balance,
|
||||
currency: quotaFromLocal.currency || 'USD',
|
||||
quota: quotaFromLocal.quota,
|
||||
statistics: localStatistics,
|
||||
lastRefreshAt: localBalance.lastCalculated
|
||||
},
|
||||
accountId,
|
||||
platform,
|
||||
'local',
|
||||
null,
|
||||
scriptMeta
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// 强制查询：优先脚本（如启用且已配置），否则调用 Provider；失败自动降级到本地统计
|
||||
let providerResult
|
||||
|
||||
if (scriptEnabled && scriptConfigured) {
|
||||
providerResult = await this._getBalanceFromScript(scriptConfig, accountId, platform)
|
||||
} else {
|
||||
const provider = this.providers.get(platform)
|
||||
if (!provider) {
|
||||
return this._buildResponse(
|
||||
{
|
||||
status: 'error',
|
||||
errorMessage: `不支持的平台: ${platform}`,
|
||||
balance: quotaFromLocal.balance,
|
||||
currency: quotaFromLocal.currency || 'USD',
|
||||
quota: quotaFromLocal.quota,
|
||||
statistics: localStatistics,
|
||||
lastRefreshAt: new Date().toISOString()
|
||||
},
|
||||
accountId,
|
||||
platform,
|
||||
'local',
|
||||
null,
|
||||
scriptMeta
|
||||
)
|
||||
}
|
||||
providerResult = await this._getBalanceFromProvider(provider, account)
|
||||
}
|
||||
|
||||
const isRemoteSuccess =
|
||||
providerResult.status === 'success' && ['api', 'script'].includes(providerResult.queryMethod)
|
||||
|
||||
// 仅缓存“真实远程查询成功”的结果，避免把字段/本地降级结果当作 API 结果缓存 1h
|
||||
if (isRemoteSuccess) {
|
||||
await this.redis.setAccountBalance(
|
||||
platform,
|
||||
accountId,
|
||||
providerResult,
|
||||
this.CACHE_TTL_SECONDS
|
||||
)
|
||||
}
|
||||
|
||||
const source = isRemoteSuccess ? 'api' : 'local'
|
||||
|
||||
return this._buildResponse(
|
||||
{
|
||||
status: providerResult.status,
|
||||
errorMessage: providerResult.errorMessage,
|
||||
balance: quotaFromLocal.balance ?? providerResult.balance,
|
||||
currency: quotaFromLocal.currency || providerResult.currency || 'USD',
|
||||
quota: quotaFromLocal.quota || providerResult.quota || null,
|
||||
statistics: localStatistics,
|
||||
lastRefreshAt: providerResult.lastRefreshAt
|
||||
},
|
||||
accountId,
|
||||
platform,
|
||||
source,
|
||||
null,
|
||||
scriptMeta
|
||||
)
|
||||
}
|
||||
|
||||
async _getBalanceFromScript(scriptConfig, accountId, platform) {
|
||||
try {
|
||||
const result = await balanceScriptService.execute({
|
||||
scriptBody: scriptConfig.scriptBody,
|
||||
timeoutSeconds: scriptConfig.timeoutSeconds || 10,
|
||||
variables: {
|
||||
baseUrl: scriptConfig.baseUrl || '',
|
||||
apiKey: scriptConfig.apiKey || '',
|
||||
token: scriptConfig.token || '',
|
||||
accountId,
|
||||
platform,
|
||||
extra: scriptConfig.extra || ''
|
||||
}
|
||||
})
|
||||
|
||||
const mapped = result?.mapped || {}
|
||||
return {
|
||||
status: mapped.status || 'error',
|
||||
balance: typeof mapped.balance === 'number' ? mapped.balance : null,
|
||||
currency: mapped.currency || 'USD',
|
||||
quota: mapped.quota || null,
|
||||
queryMethod: 'api',
|
||||
rawData: mapped.rawData || result?.response?.data || null,
|
||||
lastRefreshAt: new Date().toISOString(),
|
||||
errorMessage: mapped.errorMessage || ''
|
||||
}
|
||||
} catch (error) {
|
||||
return {
|
||||
status: 'error',
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
quota: null,
|
||||
queryMethod: 'api',
|
||||
rawData: null,
|
||||
lastRefreshAt: new Date().toISOString(),
|
||||
errorMessage: error.message || '脚本执行失败'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async _getBalanceFromProvider(provider, account) {
|
||||
try {
|
||||
const result = await provider.queryBalance(account)
|
||||
return {
|
||||
status: 'success',
|
||||
balance: typeof result?.balance === 'number' ? result.balance : null,
|
||||
currency: result?.currency || 'USD',
|
||||
quota: result?.quota || null,
|
||||
queryMethod: result?.queryMethod || 'api',
|
||||
rawData: result?.rawData || null,
|
||||
lastRefreshAt: new Date().toISOString(),
|
||||
errorMessage: ''
|
||||
}
|
||||
} catch (error) {
|
||||
return {
|
||||
status: 'error',
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
quota: null,
|
||||
queryMethod: 'api',
|
||||
rawData: null,
|
||||
lastRefreshAt: new Date().toISOString(),
|
||||
errorMessage: error.message || '查询失败'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async _getBalanceFromLocal(accountId, platform) {
|
||||
const cached = await this.redis.getLocalBalance(platform, accountId)
|
||||
if (cached && cached.statistics) {
|
||||
return cached
|
||||
}
|
||||
|
||||
const statistics = await this._computeLocalStatistics(accountId)
|
||||
const localBalance = {
|
||||
status: 'success',
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
statistics,
|
||||
queryMethod: 'local',
|
||||
lastCalculated: new Date().toISOString()
|
||||
}
|
||||
|
||||
await this.redis.setLocalBalance(platform, accountId, localBalance, this.LOCAL_TTL_SECONDS)
|
||||
return localBalance
|
||||
}
|
||||
|
||||
async _computeLocalStatistics(accountId) {
|
||||
const safeNumber = (value) => {
|
||||
const num = Number(value)
|
||||
return Number.isFinite(num) ? num : 0
|
||||
}
|
||||
|
||||
try {
|
||||
const usageStats = await this.redis.getAccountUsageStats(accountId)
|
||||
const dailyCost = safeNumber(usageStats?.daily?.cost || 0)
|
||||
const monthlyCost = await this._computeMonthlyCost(accountId)
|
||||
const totalCost = await this._computeTotalCost(accountId)
|
||||
|
||||
return {
|
||||
totalCost,
|
||||
dailyCost,
|
||||
monthlyCost,
|
||||
totalRequests: safeNumber(usageStats?.total?.requests || 0),
|
||||
dailyRequests: safeNumber(usageStats?.daily?.requests || 0),
|
||||
monthlyRequests: safeNumber(usageStats?.monthly?.requests || 0)
|
||||
}
|
||||
} catch (error) {
|
||||
this.logger.debug(`本地统计计算失败: ${accountId}`, error)
|
||||
return {
|
||||
totalCost: 0,
|
||||
dailyCost: 0,
|
||||
monthlyCost: 0,
|
||||
totalRequests: 0,
|
||||
dailyRequests: 0,
|
||||
monthlyRequests: 0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async _computeMonthlyCost(accountId) {
|
||||
const tzDate = this.redis.getDateInTimezone(new Date())
|
||||
const currentMonth = `${tzDate.getUTCFullYear()}-${String(tzDate.getUTCMonth() + 1).padStart(
|
||||
2,
|
||||
'0'
|
||||
)}`
|
||||
|
||||
const pattern = `account_usage:model:monthly:${accountId}:*:${currentMonth}`
|
||||
return await this._sumModelCostsByKeysPattern(pattern)
|
||||
}
|
||||
|
||||
async _computeTotalCost(accountId) {
|
||||
const pattern = `account_usage:model:monthly:${accountId}:*:*`
|
||||
return await this._sumModelCostsByKeysPattern(pattern)
|
||||
}
|
||||
|
||||
async _sumModelCostsByKeysPattern(pattern) {
|
||||
try {
|
||||
const client = this.redis.getClientSafe()
|
||||
let totalCost = 0
|
||||
let cursor = '0'
|
||||
const scanCount = 200
|
||||
let iterations = 0
|
||||
const maxIterations = 2000
|
||||
|
||||
do {
|
||||
const [nextCursor, keys] = await client.scan(cursor, 'MATCH', pattern, 'COUNT', scanCount)
|
||||
cursor = nextCursor
|
||||
iterations += 1
|
||||
|
||||
if (!keys || keys.length === 0) {
|
||||
continue
|
||||
}
|
||||
|
||||
const pipeline = client.pipeline()
|
||||
keys.forEach((key) => pipeline.hgetall(key))
|
||||
const results = await pipeline.exec()
|
||||
|
||||
for (let i = 0; i < results.length; i += 1) {
|
||||
const [, data] = results[i] || []
|
||||
if (!data || Object.keys(data).length === 0) {
|
||||
continue
|
||||
}
|
||||
|
||||
const parts = String(keys[i]).split(':')
|
||||
const model = parts[4] || 'unknown'
|
||||
|
||||
const usage = {
|
||||
input_tokens: parseInt(data.inputTokens || 0),
|
||||
output_tokens: parseInt(data.outputTokens || 0),
|
||||
cache_creation_input_tokens: parseInt(data.cacheCreateTokens || 0),
|
||||
cache_read_input_tokens: parseInt(data.cacheReadTokens || 0)
|
||||
}
|
||||
|
||||
const costResult = CostCalculator.calculateCost(usage, model)
|
||||
totalCost += costResult.costs.total || 0
|
||||
}
|
||||
|
||||
if (iterations >= maxIterations) {
|
||||
this.logger.warn(`SCAN 次数超过上限，停止汇总：${pattern}`)
|
||||
break
|
||||
}
|
||||
} while (cursor !== '0')
|
||||
|
||||
return totalCost
|
||||
} catch (error) {
|
||||
this.logger.debug(`汇总模型费用失败: ${pattern}`, error)
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
_buildQuotaFromLocal(account, statistics) {
|
||||
if (!account || !Object.prototype.hasOwnProperty.call(account, 'dailyQuota')) {
|
||||
return { balance: null, currency: null, quota: null }
|
||||
}
|
||||
|
||||
const dailyQuota = Number(account.dailyQuota || 0)
|
||||
const used = Number(statistics?.dailyCost || 0)
|
||||
|
||||
const resetAt = this._computeNextResetAt(account.quotaResetTime || '00:00')
|
||||
|
||||
// 不限制
|
||||
if (!Number.isFinite(dailyQuota) || dailyQuota <= 0) {
|
||||
return {
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
quota: {
|
||||
daily: Infinity,
|
||||
used,
|
||||
remaining: Infinity,
|
||||
percentage: 0,
|
||||
unlimited: true,
|
||||
resetAt
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const remaining = Math.max(0, dailyQuota - used)
|
||||
const percentage = dailyQuota > 0 ? (used / dailyQuota) * 100 : 0
|
||||
|
||||
return {
|
||||
balance: remaining,
|
||||
currency: 'USD',
|
||||
quota: {
|
||||
daily: dailyQuota,
|
||||
used,
|
||||
remaining,
|
||||
resetAt,
|
||||
percentage: Math.round(percentage * 100) / 100
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_computeNextResetAt(resetTime) {
|
||||
const now = new Date()
|
||||
const tzNow = this.redis.getDateInTimezone(now)
|
||||
const offsetMs = tzNow.getTime() - now.getTime()
|
||||
|
||||
const [h, m] = String(resetTime || '00:00')
|
||||
.split(':')
|
||||
.map((n) => parseInt(n, 10))
|
||||
|
||||
const resetHour = Number.isFinite(h) ? h : 0
|
||||
const resetMinute = Number.isFinite(m) ? m : 0
|
||||
|
||||
const year = tzNow.getUTCFullYear()
|
||||
const month = tzNow.getUTCMonth()
|
||||
const day = tzNow.getUTCDate()
|
||||
|
||||
let resetAtMs = Date.UTC(year, month, day, resetHour, resetMinute, 0, 0) - offsetMs
|
||||
if (resetAtMs <= now.getTime()) {
|
||||
resetAtMs += 24 * 60 * 60 * 1000
|
||||
}
|
||||
|
||||
return new Date(resetAtMs).toISOString()
|
||||
}
|
||||
|
||||
_buildResponse(balanceData, accountId, platform, source, ttlSeconds = null, extraData = {}) {
|
||||
const now = new Date()
|
||||
|
||||
const amount = typeof balanceData.balance === 'number' ? balanceData.balance : null
|
||||
const currency = balanceData.currency || 'USD'
|
||||
|
||||
let cacheExpiresAt = null
|
||||
if (source === 'cache') {
|
||||
const ttl =
|
||||
typeof ttlSeconds === 'number' && ttlSeconds > 0 ? ttlSeconds : this.CACHE_TTL_SECONDS
|
||||
cacheExpiresAt = new Date(Date.now() + ttl * 1000).toISOString()
|
||||
}
|
||||
|
||||
return {
|
||||
success: true,
|
||||
data: {
|
||||
accountId,
|
||||
platform,
|
||||
balance:
|
||||
typeof amount === 'number'
|
||||
? {
|
||||
amount,
|
||||
currency,
|
||||
formattedAmount: this._formatCurrency(amount, currency)
|
||||
}
|
||||
: null,
|
||||
quota: balanceData.quota || null,
|
||||
statistics: balanceData.statistics || {},
|
||||
source,
|
||||
lastRefreshAt: balanceData.lastRefreshAt || now.toISOString(),
|
||||
cacheExpiresAt,
|
||||
status: balanceData.status || 'success',
|
||||
error: balanceData.errorMessage || null,
|
||||
...(extraData && typeof extraData === 'object' ? extraData : {})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_formatCurrency(amount, currency = 'USD') {
|
||||
try {
|
||||
if (typeof amount !== 'number' || !Number.isFinite(amount)) {
|
||||
return 'N/A'
|
||||
}
|
||||
return new Intl.NumberFormat('en-US', { style: 'currency', currency }).format(amount)
|
||||
} catch (error) {
|
||||
return `$${amount.toFixed(2)}`
|
||||
}
|
||||
}
|
||||
|
||||
_parseBoolean(value) {
|
||||
if (typeof value === 'boolean') {
|
||||
return value
|
||||
}
|
||||
if (typeof value !== 'string') {
|
||||
return null
|
||||
}
|
||||
const normalized = value.trim().toLowerCase()
|
||||
if (normalized === 'true' || normalized === '1' || normalized === 'yes') {
|
||||
return true
|
||||
}
|
||||
if (normalized === 'false' || normalized === '0' || normalized === 'no') {
|
||||
return false
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
_parseQueryMode(value) {
|
||||
if (value === 'auto') {
|
||||
return 'auto'
|
||||
}
|
||||
const parsed = this._parseBoolean(value)
|
||||
return parsed ? 'api' : 'local'
|
||||
}
|
||||
|
||||
async _mapWithConcurrency(items, limit, mapper) {
|
||||
const concurrency = Math.max(1, Number(limit) || 1)
|
||||
const list = Array.isArray(items) ? items : []
|
||||
|
||||
const results = new Array(list.length)
|
||||
let nextIndex = 0
|
||||
|
||||
const workers = new Array(Math.min(concurrency, list.length)).fill(null).map(async () => {
|
||||
while (nextIndex < list.length) {
|
||||
const currentIndex = nextIndex
|
||||
nextIndex += 1
|
||||
results[currentIndex] = await mapper(list[currentIndex], currentIndex)
|
||||
}
|
||||
})
|
||||
|
||||
await Promise.all(workers)
|
||||
return results
|
||||
}
|
||||
}
|
||||
|
||||
const accountBalanceService = new AccountBalanceService()
|
||||
module.exports = accountBalanceService
|
||||
module.exports.AccountBalanceService = AccountBalanceService
|
||||
3071
src/services/anthropicGeminiBridgeService.js
Normal file
3071
src/services/anthropicGeminiBridgeService.js
Normal file
File diff suppressed because it is too large
Load Diff
594
src/services/antigravityClient.js
Normal file
594
src/services/antigravityClient.js
Normal file
@@ -0,0 +1,594 @@
|
||||
const axios = require('axios')
|
||||
const https = require('https')
|
||||
const { v4: uuidv4 } = require('uuid')
|
||||
|
||||
const ProxyHelper = require('../utils/proxyHelper')
|
||||
const logger = require('../utils/logger')
|
||||
const {
|
||||
mapAntigravityUpstreamModel,
|
||||
normalizeAntigravityModelInput,
|
||||
getAntigravityModelMetadata
|
||||
} = require('../utils/antigravityModel')
|
||||
const { cleanJsonSchemaForGemini } = require('../utils/geminiSchemaCleaner')
|
||||
const { dumpAntigravityUpstreamRequest } = require('../utils/antigravityUpstreamDump')
|
||||
|
||||
const keepAliveAgent = new https.Agent({
|
||||
keepAlive: true,
|
||||
keepAliveMsecs: 30000,
|
||||
timeout: 120000,
|
||||
maxSockets: 100,
|
||||
maxFreeSockets: 10
|
||||
})
|
||||
|
||||
function getAntigravityApiUrl() {
|
||||
return process.env.ANTIGRAVITY_API_URL || 'https://daily-cloudcode-pa.sandbox.googleapis.com'
|
||||
}
|
||||
|
||||
function normalizeBaseUrl(url) {
|
||||
const str = String(url || '').trim()
|
||||
return str.endsWith('/') ? str.slice(0, -1) : str
|
||||
}
|
||||
|
||||
function getAntigravityApiUrlCandidates() {
|
||||
const configured = normalizeBaseUrl(getAntigravityApiUrl())
|
||||
const daily = 'https://daily-cloudcode-pa.sandbox.googleapis.com'
|
||||
const prod = 'https://cloudcode-pa.googleapis.com'
|
||||
|
||||
// 若显式配置了自定义 base url，则只使用该地址（不做 fallback，避免意外路由到别的环境）。
|
||||
if (process.env.ANTIGRAVITY_API_URL) {
|
||||
return [configured]
|
||||
}
|
||||
|
||||
// 默认行为：优先 daily（与旧逻辑一致），失败时再尝试 prod（对齐 CLIProxyAPI）。
|
||||
if (configured === normalizeBaseUrl(daily)) {
|
||||
return [configured, prod]
|
||||
}
|
||||
if (configured === normalizeBaseUrl(prod)) {
|
||||
return [configured, daily]
|
||||
}
|
||||
|
||||
return [configured, prod, daily].filter(Boolean)
|
||||
}
|
||||
|
||||
function getAntigravityHeaders(accessToken, baseUrl) {
|
||||
const resolvedBaseUrl = baseUrl || getAntigravityApiUrl()
|
||||
let host = 'daily-cloudcode-pa.sandbox.googleapis.com'
|
||||
try {
|
||||
host = new URL(resolvedBaseUrl).host || host
|
||||
} catch (e) {
|
||||
// ignore
|
||||
}
|
||||
|
||||
return {
|
||||
Host: host,
|
||||
'User-Agent': process.env.ANTIGRAVITY_USER_AGENT || 'antigravity/1.11.3 windows/amd64',
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
'Content-Type': 'application/json',
|
||||
'Accept-Encoding': 'gzip'
|
||||
}
|
||||
}
|
||||
|
||||
function generateAntigravityProjectId() {
|
||||
return `ag-${uuidv4().replace(/-/g, '').slice(0, 16)}`
|
||||
}
|
||||
|
||||
function generateAntigravitySessionId() {
|
||||
return `sess-${uuidv4()}`
|
||||
}
|
||||
|
||||
function resolveAntigravityProjectId(projectId, requestData) {
|
||||
const candidate = projectId || requestData?.project || requestData?.projectId || null
|
||||
return candidate || generateAntigravityProjectId()
|
||||
}
|
||||
|
||||
function resolveAntigravitySessionId(sessionId, requestData) {
|
||||
const candidate =
|
||||
sessionId || requestData?.request?.sessionId || requestData?.request?.session_id || null
|
||||
return candidate || generateAntigravitySessionId()
|
||||
}
|
||||
|
||||
function buildAntigravityEnvelope({ requestData, projectId, sessionId, userPromptId }) {
|
||||
const model = mapAntigravityUpstreamModel(requestData?.model)
|
||||
const resolvedProjectId = resolveAntigravityProjectId(projectId, requestData)
|
||||
const resolvedSessionId = resolveAntigravitySessionId(sessionId, requestData)
|
||||
const requestPayload = {
|
||||
...(requestData?.request || {})
|
||||
}
|
||||
|
||||
if (requestPayload.session_id !== undefined) {
|
||||
delete requestPayload.session_id
|
||||
}
|
||||
requestPayload.sessionId = resolvedSessionId
|
||||
|
||||
const envelope = {
|
||||
project: resolvedProjectId,
|
||||
requestId: `req-${uuidv4()}`,
|
||||
model,
|
||||
userAgent: 'antigravity',
|
||||
request: {
|
||||
...requestPayload
|
||||
}
|
||||
}
|
||||
|
||||
if (userPromptId) {
|
||||
envelope.user_prompt_id = userPromptId
|
||||
envelope.userPromptId = userPromptId
|
||||
}
|
||||
|
||||
normalizeAntigravityEnvelope(envelope)
|
||||
return { model, envelope }
|
||||
}
|
||||
|
||||
function normalizeAntigravityThinking(model, requestPayload) {
|
||||
if (!requestPayload || typeof requestPayload !== 'object') {
|
||||
return
|
||||
}
|
||||
|
||||
const { generationConfig } = requestPayload
|
||||
if (!generationConfig || typeof generationConfig !== 'object') {
|
||||
return
|
||||
}
|
||||
const { thinkingConfig } = generationConfig
|
||||
if (!thinkingConfig || typeof thinkingConfig !== 'object') {
|
||||
return
|
||||
}
|
||||
|
||||
const normalizedModel = normalizeAntigravityModelInput(model)
|
||||
if (thinkingConfig.thinkingLevel && !normalizedModel.startsWith('gemini-3-')) {
|
||||
delete thinkingConfig.thinkingLevel
|
||||
}
|
||||
|
||||
const metadata = getAntigravityModelMetadata(normalizedModel)
|
||||
if (metadata && !metadata.thinking) {
|
||||
delete generationConfig.thinkingConfig
|
||||
return
|
||||
}
|
||||
if (!metadata || !metadata.thinking) {
|
||||
return
|
||||
}
|
||||
|
||||
const budgetRaw = Number(thinkingConfig.thinkingBudget)
|
||||
if (!Number.isFinite(budgetRaw)) {
|
||||
return
|
||||
}
|
||||
let budget = Math.trunc(budgetRaw)
|
||||
|
||||
const minBudget = Number.isFinite(metadata.thinking.min) ? metadata.thinking.min : null
|
||||
const maxBudget = Number.isFinite(metadata.thinking.max) ? metadata.thinking.max : null
|
||||
|
||||
if (maxBudget !== null && budget > maxBudget) {
|
||||
budget = maxBudget
|
||||
}
|
||||
|
||||
let effectiveMax = Number.isFinite(generationConfig.maxOutputTokens)
|
||||
? generationConfig.maxOutputTokens
|
||||
: null
|
||||
let setDefaultMax = false
|
||||
if (!effectiveMax && metadata.maxCompletionTokens) {
|
||||
effectiveMax = metadata.maxCompletionTokens
|
||||
setDefaultMax = true
|
||||
}
|
||||
|
||||
if (effectiveMax && budget >= effectiveMax) {
|
||||
budget = Math.max(0, effectiveMax - 1)
|
||||
}
|
||||
|
||||
if (minBudget !== null && budget >= 0 && budget < minBudget) {
|
||||
delete generationConfig.thinkingConfig
|
||||
return
|
||||
}
|
||||
|
||||
thinkingConfig.thinkingBudget = budget
|
||||
if (setDefaultMax) {
|
||||
generationConfig.maxOutputTokens = effectiveMax
|
||||
}
|
||||
}
|
||||
|
||||
function normalizeAntigravityEnvelope(envelope) {
|
||||
if (!envelope || typeof envelope !== 'object') {
|
||||
return
|
||||
}
|
||||
const model = String(envelope.model || '')
|
||||
const requestPayload = envelope.request
|
||||
if (!requestPayload || typeof requestPayload !== 'object') {
|
||||
return
|
||||
}
|
||||
|
||||
if (requestPayload.safetySettings !== undefined) {
|
||||
delete requestPayload.safetySettings
|
||||
}
|
||||
|
||||
// 对齐 CLIProxyAPI：有 tools 时默认启用 VALIDATED（除非显式 NONE）
|
||||
if (Array.isArray(requestPayload.tools) && requestPayload.tools.length > 0) {
|
||||
const existing = requestPayload?.toolConfig?.functionCallingConfig || null
|
||||
if (existing?.mode !== 'NONE') {
|
||||
const nextCfg = { ...(existing || {}), mode: 'VALIDATED' }
|
||||
requestPayload.toolConfig = { functionCallingConfig: nextCfg }
|
||||
}
|
||||
}
|
||||
|
||||
// 对齐 CLIProxyAPI：非 Claude 模型移除 maxOutputTokens（Antigravity 环境不稳定）
|
||||
normalizeAntigravityThinking(model, requestPayload)
|
||||
if (!model.includes('claude')) {
|
||||
if (requestPayload.generationConfig && typeof requestPayload.generationConfig === 'object') {
|
||||
delete requestPayload.generationConfig.maxOutputTokens
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Claude 模型：parametersJsonSchema -> parameters + schema 清洗（避免 $schema / additionalProperties 等触发 400）
|
||||
if (!Array.isArray(requestPayload.tools)) {
|
||||
return
|
||||
}
|
||||
|
||||
for (const tool of requestPayload.tools) {
|
||||
if (!tool || typeof tool !== 'object') {
|
||||
continue
|
||||
}
|
||||
const decls = Array.isArray(tool.functionDeclarations)
|
||||
? tool.functionDeclarations
|
||||
: Array.isArray(tool.function_declarations)
|
||||
? tool.function_declarations
|
||||
: null
|
||||
|
||||
if (!decls) {
|
||||
continue
|
||||
}
|
||||
|
||||
for (const decl of decls) {
|
||||
if (!decl || typeof decl !== 'object') {
|
||||
continue
|
||||
}
|
||||
let schema =
|
||||
decl.parametersJsonSchema !== undefined ? decl.parametersJsonSchema : decl.parameters
|
||||
if (typeof schema === 'string' && schema) {
|
||||
try {
|
||||
schema = JSON.parse(schema)
|
||||
} catch (_) {
|
||||
schema = null
|
||||
}
|
||||
}
|
||||
|
||||
decl.parameters = cleanJsonSchemaForGemini(schema)
|
||||
delete decl.parametersJsonSchema
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function request({
|
||||
accessToken,
|
||||
proxyConfig = null,
|
||||
requestData,
|
||||
projectId = null,
|
||||
sessionId = null,
|
||||
userPromptId = null,
|
||||
stream = false,
|
||||
signal = null,
|
||||
params = null,
|
||||
timeoutMs = null
|
||||
}) {
|
||||
const { model, envelope } = buildAntigravityEnvelope({
|
||||
requestData,
|
||||
projectId,
|
||||
sessionId,
|
||||
userPromptId
|
||||
})
|
||||
|
||||
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig)
|
||||
let endpoints = getAntigravityApiUrlCandidates()
|
||||
|
||||
// Claude 模型在 sandbox(daily) 环境下对 tool_use/tool_result 的兼容性不稳定，优先走 prod。
|
||||
// 保持可配置优先：若用户显式设置了 ANTIGRAVITY_API_URL，则不改变顺序。
|
||||
if (!process.env.ANTIGRAVITY_API_URL && String(model).includes('claude')) {
|
||||
const prodHost = 'cloudcode-pa.googleapis.com'
|
||||
const dailyHost = 'daily-cloudcode-pa.sandbox.googleapis.com'
|
||||
const ordered = []
|
||||
for (const u of endpoints) {
|
||||
if (String(u).includes(prodHost)) {
|
||||
ordered.push(u)
|
||||
}
|
||||
}
|
||||
for (const u of endpoints) {
|
||||
if (!String(u).includes(prodHost)) {
|
||||
ordered.push(u)
|
||||
}
|
||||
}
|
||||
// 去重并保持 prod -> daily 的稳定顺序
|
||||
endpoints = Array.from(new Set(ordered)).sort((a, b) => {
|
||||
const av = String(a)
|
||||
const bv = String(b)
|
||||
const aScore = av.includes(prodHost) ? 0 : av.includes(dailyHost) ? 1 : 2
|
||||
const bScore = bv.includes(prodHost) ? 0 : bv.includes(dailyHost) ? 1 : 2
|
||||
return aScore - bScore
|
||||
})
|
||||
}
|
||||
|
||||
const isRetryable = (error) => {
|
||||
// 处理网络层面的连接重置或超时（常见于长请求被中间节点切断）
|
||||
if (error.code === 'ECONNRESET' || error.code === 'ETIMEDOUT') {
|
||||
return true
|
||||
}
|
||||
|
||||
const status = error?.response?.status
|
||||
if (status === 429) {
|
||||
return true
|
||||
}
|
||||
|
||||
// 400/404 的 “model unavailable / not found” 在不同环境间可能表现不同，允许 fallback。
|
||||
if (status === 400 || status === 404) {
|
||||
const data = error?.response?.data
|
||||
const safeToString = (value) => {
|
||||
if (typeof value === 'string') {
|
||||
return value
|
||||
}
|
||||
if (value === null || value === undefined) {
|
||||
return ''
|
||||
}
|
||||
// axios responseType=stream 时，data 可能是 stream（存在循环引用），不能 JSON.stringify
|
||||
if (typeof value === 'object' && typeof value.pipe === 'function') {
|
||||
return ''
|
||||
}
|
||||
if (Buffer.isBuffer(value)) {
|
||||
try {
|
||||
return value.toString('utf8')
|
||||
} catch (_) {
|
||||
return ''
|
||||
}
|
||||
}
|
||||
if (typeof value === 'object') {
|
||||
try {
|
||||
return JSON.stringify(value)
|
||||
} catch (_) {
|
||||
return ''
|
||||
}
|
||||
}
|
||||
return String(value)
|
||||
}
|
||||
|
||||
const text = safeToString(data)
|
||||
const msg = (text || '').toLowerCase()
|
||||
return (
|
||||
msg.includes('requested model is currently unavailable') ||
|
||||
msg.includes('tool_use') ||
|
||||
msg.includes('tool_result') ||
|
||||
msg.includes('requested entity was not found') ||
|
||||
msg.includes('not found')
|
||||
)
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
let lastError = null
|
||||
let retriedAfterDelay = false
|
||||
|
||||
const attemptRequest = async () => {
|
||||
for (let index = 0; index < endpoints.length; index += 1) {
|
||||
const baseUrl = endpoints[index]
|
||||
const url = `${baseUrl}/v1internal:${stream ? 'streamGenerateContent' : 'generateContent'}`
|
||||
|
||||
const axiosConfig = {
|
||||
url,
|
||||
method: 'POST',
|
||||
...(params ? { params } : {}),
|
||||
headers: getAntigravityHeaders(accessToken, baseUrl),
|
||||
data: envelope,
|
||||
timeout: stream ? 0 : timeoutMs || 600000,
|
||||
...(stream ? { responseType: 'stream' } : {})
|
||||
}
|
||||
|
||||
if (proxyAgent) {
|
||||
axiosConfig.httpsAgent = proxyAgent
|
||||
axiosConfig.proxy = false
|
||||
if (index === 0) {
|
||||
logger.info(
|
||||
`🌐 Using proxy for Antigravity ${stream ? 'streamGenerateContent' : 'generateContent'}: ${ProxyHelper.getProxyDescription(proxyConfig)}`
|
||||
)
|
||||
}
|
||||
} else {
|
||||
axiosConfig.httpsAgent = keepAliveAgent
|
||||
}
|
||||
|
||||
if (signal) {
|
||||
axiosConfig.signal = signal
|
||||
}
|
||||
|
||||
try {
|
||||
dumpAntigravityUpstreamRequest({
|
||||
requestId: envelope.requestId,
|
||||
model,
|
||||
stream,
|
||||
url,
|
||||
baseUrl,
|
||||
params: axiosConfig.params || null,
|
||||
headers: axiosConfig.headers,
|
||||
envelope
|
||||
}).catch(() => {})
|
||||
const response = await axios(axiosConfig)
|
||||
return { model, response }
|
||||
} catch (error) {
|
||||
lastError = error
|
||||
const status = error?.response?.status || null
|
||||
|
||||
const hasNext = index + 1 < endpoints.length
|
||||
if (hasNext && isRetryable(error)) {
|
||||
logger.warn('⚠️ Antigravity upstream error, retrying with fallback baseUrl', {
|
||||
status,
|
||||
from: baseUrl,
|
||||
to: endpoints[index + 1],
|
||||
model
|
||||
})
|
||||
continue
|
||||
}
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
throw lastError || new Error('Antigravity request failed')
|
||||
}
|
||||
|
||||
try {
|
||||
return await attemptRequest()
|
||||
} catch (error) {
|
||||
// 如果是 429 RESOURCE_EXHAUSTED 且尚未重试过，等待 2 秒后重试一次
|
||||
const status = error?.response?.status
|
||||
if (status === 429 && !retriedAfterDelay && !signal?.aborted) {
|
||||
const data = error?.response?.data
|
||||
|
||||
// 安全地将 data 转为字符串，避免 stream 对象导致循环引用崩溃
|
||||
const safeDataToString = (value) => {
|
||||
if (typeof value === 'string') {
|
||||
return value
|
||||
}
|
||||
if (value === null || value === undefined) {
|
||||
return ''
|
||||
}
|
||||
// stream 对象存在循环引用，不能 JSON.stringify
|
||||
if (typeof value === 'object' && typeof value.pipe === 'function') {
|
||||
return ''
|
||||
}
|
||||
if (Buffer.isBuffer(value)) {
|
||||
try {
|
||||
return value.toString('utf8')
|
||||
} catch (_) {
|
||||
return ''
|
||||
}
|
||||
}
|
||||
if (typeof value === 'object') {
|
||||
try {
|
||||
return JSON.stringify(value)
|
||||
} catch (_) {
|
||||
return ''
|
||||
}
|
||||
}
|
||||
return String(value)
|
||||
}
|
||||
|
||||
const msg = safeDataToString(data)
|
||||
if (
|
||||
msg.toLowerCase().includes('resource_exhausted') ||
|
||||
msg.toLowerCase().includes('no capacity')
|
||||
) {
|
||||
retriedAfterDelay = true
|
||||
logger.warn('⏳ Antigravity 429 RESOURCE_EXHAUSTED, waiting 2s before retry', { model })
|
||||
await new Promise((resolve) => setTimeout(resolve, 2000))
|
||||
return await attemptRequest()
|
||||
}
|
||||
}
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
async function fetchAvailableModels({ accessToken, proxyConfig = null, timeoutMs = 30000 }) {
|
||||
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig)
|
||||
const endpoints = getAntigravityApiUrlCandidates()
|
||||
|
||||
let lastError = null
|
||||
for (let index = 0; index < endpoints.length; index += 1) {
|
||||
const baseUrl = endpoints[index]
|
||||
const url = `${baseUrl}/v1internal:fetchAvailableModels`
|
||||
|
||||
const axiosConfig = {
|
||||
url,
|
||||
method: 'POST',
|
||||
headers: getAntigravityHeaders(accessToken, baseUrl),
|
||||
data: {},
|
||||
timeout: timeoutMs
|
||||
}
|
||||
|
||||
if (proxyAgent) {
|
||||
axiosConfig.httpsAgent = proxyAgent
|
||||
axiosConfig.proxy = false
|
||||
if (index === 0) {
|
||||
logger.info(
|
||||
`🌐 Using proxy for Antigravity fetchAvailableModels: ${ProxyHelper.getProxyDescription(proxyConfig)}`
|
||||
)
|
||||
}
|
||||
} else {
|
||||
axiosConfig.httpsAgent = keepAliveAgent
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await axios(axiosConfig)
|
||||
return response.data
|
||||
} catch (error) {
|
||||
lastError = error
|
||||
const status = error?.response?.status
|
||||
const hasNext = index + 1 < endpoints.length
|
||||
if (hasNext && (status === 429 || status === 404)) {
|
||||
continue
|
||||
}
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
throw lastError || new Error('Antigravity fetchAvailableModels failed')
|
||||
}
|
||||
|
||||
async function countTokens({
|
||||
accessToken,
|
||||
proxyConfig = null,
|
||||
contents,
|
||||
model,
|
||||
timeoutMs = 30000
|
||||
}) {
|
||||
const upstreamModel = mapAntigravityUpstreamModel(model)
|
||||
|
||||
const proxyAgent = ProxyHelper.createProxyAgent(proxyConfig)
|
||||
const endpoints = getAntigravityApiUrlCandidates()
|
||||
|
||||
let lastError = null
|
||||
for (let index = 0; index < endpoints.length; index += 1) {
|
||||
const baseUrl = endpoints[index]
|
||||
const url = `${baseUrl}/v1internal:countTokens`
|
||||
const axiosConfig = {
|
||||
url,
|
||||
method: 'POST',
|
||||
headers: getAntigravityHeaders(accessToken, baseUrl),
|
||||
data: {
|
||||
request: {
|
||||
model: `models/${upstreamModel}`,
|
||||
contents
|
||||
}
|
||||
},
|
||||
timeout: timeoutMs
|
||||
}
|
||||
|
||||
if (proxyAgent) {
|
||||
axiosConfig.httpsAgent = proxyAgent
|
||||
axiosConfig.proxy = false
|
||||
if (index === 0) {
|
||||
logger.info(
|
||||
`🌐 Using proxy for Antigravity countTokens: ${ProxyHelper.getProxyDescription(proxyConfig)}`
|
||||
)
|
||||
}
|
||||
} else {
|
||||
axiosConfig.httpsAgent = keepAliveAgent
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await axios(axiosConfig)
|
||||
return response.data
|
||||
} catch (error) {
|
||||
lastError = error
|
||||
const status = error?.response?.status
|
||||
const hasNext = index + 1 < endpoints.length
|
||||
if (hasNext && (status === 429 || status === 404)) {
|
||||
continue
|
||||
}
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
throw lastError || new Error('Antigravity countTokens failed')
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
getAntigravityApiUrl,
|
||||
getAntigravityApiUrlCandidates,
|
||||
getAntigravityHeaders,
|
||||
buildAntigravityEnvelope,
|
||||
request,
|
||||
fetchAvailableModels,
|
||||
countTokens
|
||||
}
|
||||
170
src/services/antigravityRelayService.js
Normal file
170
src/services/antigravityRelayService.js
Normal file
@@ -0,0 +1,170 @@
|
||||
const apiKeyService = require('./apiKeyService')
|
||||
const { convertMessagesToGemini, convertGeminiResponse } = require('./geminiRelayService')
|
||||
const { normalizeAntigravityModelInput } = require('../utils/antigravityModel')
|
||||
const antigravityClient = require('./antigravityClient')
|
||||
|
||||
function buildRequestData({ messages, model, temperature, maxTokens, sessionId }) {
|
||||
const requestedModel = normalizeAntigravityModelInput(model)
|
||||
const { contents, systemInstruction } = convertMessagesToGemini(messages)
|
||||
|
||||
const requestData = {
|
||||
model: requestedModel,
|
||||
request: {
|
||||
contents,
|
||||
generationConfig: {
|
||||
temperature,
|
||||
maxOutputTokens: maxTokens,
|
||||
candidateCount: 1,
|
||||
topP: 0.95,
|
||||
topK: 40
|
||||
},
|
||||
...(sessionId ? { sessionId } : {})
|
||||
}
|
||||
}
|
||||
|
||||
if (systemInstruction) {
|
||||
requestData.request.systemInstruction = { parts: [{ text: systemInstruction }] }
|
||||
}
|
||||
|
||||
return requestData
|
||||
}
|
||||
|
||||
async function* handleStreamResponse(response, model, apiKeyId, accountId) {
|
||||
let buffer = ''
|
||||
let totalUsage = {
|
||||
promptTokenCount: 0,
|
||||
candidatesTokenCount: 0,
|
||||
totalTokenCount: 0
|
||||
}
|
||||
let usageRecorded = false
|
||||
|
||||
try {
|
||||
for await (const chunk of response.data) {
|
||||
buffer += chunk.toString()
|
||||
|
||||
const lines = buffer.split('\n')
|
||||
buffer = lines.pop() || ''
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.trim()) {
|
||||
continue
|
||||
}
|
||||
|
||||
let jsonData = line
|
||||
if (line.startsWith('data: ')) {
|
||||
jsonData = line.substring(6).trim()
|
||||
}
|
||||
|
||||
if (!jsonData || jsonData === '[DONE]') {
|
||||
continue
|
||||
}
|
||||
|
||||
try {
|
||||
const data = JSON.parse(jsonData)
|
||||
const payload = data?.response || data
|
||||
|
||||
if (payload?.usageMetadata) {
|
||||
totalUsage = payload.usageMetadata
|
||||
}
|
||||
|
||||
const openaiChunk = convertGeminiResponse(payload, model, true)
|
||||
if (openaiChunk) {
|
||||
yield `data: ${JSON.stringify(openaiChunk)}\n\n`
|
||||
const finishReason = openaiChunk.choices?.[0]?.finish_reason
|
||||
if (finishReason === 'stop') {
|
||||
yield 'data: [DONE]\n\n'
|
||||
|
||||
if (apiKeyId && totalUsage.totalTokenCount > 0) {
|
||||
await apiKeyService.recordUsage(
|
||||
apiKeyId,
|
||||
totalUsage.promptTokenCount || 0,
|
||||
totalUsage.candidatesTokenCount || 0,
|
||||
0,
|
||||
0,
|
||||
model,
|
||||
accountId
|
||||
)
|
||||
usageRecorded = true
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
// ignore chunk parse errors
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
if (!usageRecorded && apiKeyId && totalUsage.totalTokenCount > 0) {
|
||||
await apiKeyService.recordUsage(
|
||||
apiKeyId,
|
||||
totalUsage.promptTokenCount || 0,
|
||||
totalUsage.candidatesTokenCount || 0,
|
||||
0,
|
||||
0,
|
||||
model,
|
||||
accountId
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function sendAntigravityRequest({
|
||||
messages,
|
||||
model,
|
||||
temperature = 0.7,
|
||||
maxTokens = 4096,
|
||||
stream = false,
|
||||
accessToken,
|
||||
proxy,
|
||||
apiKeyId,
|
||||
signal,
|
||||
projectId,
|
||||
accountId = null
|
||||
}) {
|
||||
const requestedModel = normalizeAntigravityModelInput(model)
|
||||
|
||||
const requestData = buildRequestData({
|
||||
messages,
|
||||
model: requestedModel,
|
||||
temperature,
|
||||
maxTokens,
|
||||
sessionId: apiKeyId
|
||||
})
|
||||
|
||||
const { response } = await antigravityClient.request({
|
||||
accessToken,
|
||||
proxyConfig: proxy,
|
||||
requestData,
|
||||
projectId,
|
||||
sessionId: apiKeyId,
|
||||
stream,
|
||||
signal,
|
||||
params: { alt: 'sse' }
|
||||
})
|
||||
|
||||
if (stream) {
|
||||
return handleStreamResponse(response, requestedModel, apiKeyId, accountId)
|
||||
}
|
||||
|
||||
const payload = response.data?.response || response.data
|
||||
const openaiResponse = convertGeminiResponse(payload, requestedModel, false)
|
||||
|
||||
if (apiKeyId && openaiResponse?.usage) {
|
||||
await apiKeyService.recordUsage(
|
||||
apiKeyId,
|
||||
openaiResponse.usage.prompt_tokens || 0,
|
||||
openaiResponse.usage.completion_tokens || 0,
|
||||
0,
|
||||
0,
|
||||
requestedModel,
|
||||
accountId
|
||||
)
|
||||
}
|
||||
|
||||
return openaiResponse
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
sendAntigravityRequest
|
||||
}
|
||||
@@ -37,6 +37,51 @@ const ACCOUNT_CATEGORY_MAP = {
|
||||
droid: 'droid'
|
||||
}
|
||||
|
||||
/**
|
||||
* 规范化权限数据，兼容旧格式（字符串）和新格式（数组）
|
||||
* @param {string|array} permissions - 权限数据
|
||||
* @returns {array} - 权限数组，空数组表示全部服务
|
||||
*/
|
||||
function normalizePermissions(permissions) {
|
||||
if (!permissions) {
|
||||
return [] // 空 = 全部服务
|
||||
}
|
||||
if (Array.isArray(permissions)) {
|
||||
return permissions
|
||||
}
|
||||
// 尝试解析 JSON 字符串（新格式存储）
|
||||
if (typeof permissions === 'string') {
|
||||
if (permissions.startsWith('[')) {
|
||||
try {
|
||||
const parsed = JSON.parse(permissions)
|
||||
if (Array.isArray(parsed)) {
|
||||
return parsed
|
||||
}
|
||||
} catch (e) {
|
||||
// 解析失败，继续处理为普通字符串
|
||||
}
|
||||
}
|
||||
// 旧格式 'all' 转为空数组
|
||||
if (permissions === 'all') {
|
||||
return []
|
||||
}
|
||||
// 旧单个字符串转为数组
|
||||
return [permissions]
|
||||
}
|
||||
return []
|
||||
}
|
||||
|
||||
/**
|
||||
* 检查是否有访问特定服务的权限
|
||||
* @param {string|array} permissions - 权限数据
|
||||
* @param {string} service - 服务名称（claude/gemini/openai/droid）
|
||||
* @returns {boolean} - 是否有权限
|
||||
*/
|
||||
function hasPermission(permissions, service) {
|
||||
const perms = normalizePermissions(permissions)
|
||||
return perms.length === 0 || perms.includes(service) // 空数组 = 全部服务
|
||||
}
|
||||
|
||||
function normalizeAccountTypeKey(type) {
|
||||
if (!type) {
|
||||
return null
|
||||
@@ -89,7 +134,7 @@ class ApiKeyService {
|
||||
azureOpenaiAccountId = null,
|
||||
bedrockAccountId = null, // 添加 Bedrock 账号ID支持
|
||||
droidAccountId = null,
|
||||
permissions = 'all', // 可选值：'claude'、'gemini'、'openai'、'droid' 或 'all'
|
||||
permissions = [], // 数组格式，空数组表示全部服务，如 ['claude', 'gemini']
|
||||
isActive = true,
|
||||
concurrencyLimit = 0,
|
||||
rateLimitWindow = null,
|
||||
@@ -132,7 +177,7 @@ class ApiKeyService {
|
||||
azureOpenaiAccountId: azureOpenaiAccountId || '',
|
||||
bedrockAccountId: bedrockAccountId || '', // 添加 Bedrock 账号ID
|
||||
droidAccountId: droidAccountId || '',
|
||||
permissions: permissions || 'all',
|
||||
permissions: JSON.stringify(normalizePermissions(permissions)),
|
||||
enableModelRestriction: String(enableModelRestriction),
|
||||
restrictedModels: JSON.stringify(restrictedModels || []),
|
||||
enableClientRestriction: String(enableClientRestriction || false),
|
||||
@@ -186,7 +231,7 @@ class ApiKeyService {
|
||||
azureOpenaiAccountId: keyData.azureOpenaiAccountId,
|
||||
bedrockAccountId: keyData.bedrockAccountId, // 添加 Bedrock 账号ID
|
||||
droidAccountId: keyData.droidAccountId,
|
||||
permissions: keyData.permissions,
|
||||
permissions: normalizePermissions(keyData.permissions),
|
||||
enableModelRestriction: keyData.enableModelRestriction === 'true',
|
||||
restrictedModels: JSON.parse(keyData.restrictedModels),
|
||||
enableClientRestriction: keyData.enableClientRestriction === 'true',
|
||||
@@ -338,7 +383,7 @@ class ApiKeyService {
|
||||
azureOpenaiAccountId: keyData.azureOpenaiAccountId,
|
||||
bedrockAccountId: keyData.bedrockAccountId, // 添加 Bedrock 账号ID
|
||||
droidAccountId: keyData.droidAccountId,
|
||||
permissions: keyData.permissions || 'all',
|
||||
permissions: normalizePermissions(keyData.permissions),
|
||||
tokenLimit: parseInt(keyData.tokenLimit),
|
||||
concurrencyLimit: parseInt(keyData.concurrencyLimit || 0),
|
||||
rateLimitWindow: parseInt(keyData.rateLimitWindow || 0),
|
||||
@@ -467,7 +512,7 @@ class ApiKeyService {
|
||||
azureOpenaiAccountId: keyData.azureOpenaiAccountId,
|
||||
bedrockAccountId: keyData.bedrockAccountId,
|
||||
droidAccountId: keyData.droidAccountId,
|
||||
permissions: keyData.permissions || 'all',
|
||||
permissions: normalizePermissions(keyData.permissions),
|
||||
tokenLimit: parseInt(keyData.tokenLimit),
|
||||
concurrencyLimit: parseInt(keyData.concurrencyLimit || 0),
|
||||
rateLimitWindow: parseInt(keyData.rateLimitWindow || 0),
|
||||
@@ -525,7 +570,7 @@ class ApiKeyService {
|
||||
key.isActive = key.isActive === 'true'
|
||||
key.enableModelRestriction = key.enableModelRestriction === 'true'
|
||||
key.enableClientRestriction = key.enableClientRestriction === 'true'
|
||||
key.permissions = key.permissions || 'all' // 兼容旧数据
|
||||
key.permissions = normalizePermissions(key.permissions)
|
||||
key.dailyCostLimit = parseFloat(key.dailyCostLimit || 0)
|
||||
key.totalCostLimit = parseFloat(key.totalCostLimit || 0)
|
||||
key.weeklyOpusCostLimit = parseFloat(key.weeklyOpusCostLimit || 0)
|
||||
@@ -1568,7 +1613,7 @@ class ApiKeyService {
|
||||
userId: keyData.userId,
|
||||
userUsername: keyData.userUsername,
|
||||
createdBy: keyData.createdBy,
|
||||
permissions: keyData.permissions,
|
||||
permissions: normalizePermissions(keyData.permissions),
|
||||
dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
|
||||
totalCostLimit: parseFloat(keyData.totalCostLimit || 0),
|
||||
// 所有平台账户绑定字段
|
||||
@@ -1820,4 +1865,8 @@ const apiKeyService = new ApiKeyService()
|
||||
// 为了方便其他服务调用，导出 recordUsage 方法
|
||||
apiKeyService.recordUsageMetrics = apiKeyService.recordUsage.bind(apiKeyService)
|
||||
|
||||
// 导出权限辅助函数供路由使用
|
||||
apiKeyService.hasPermission = hasPermission
|
||||
apiKeyService.normalizePermissions = normalizePermissions
|
||||
|
||||
module.exports = apiKeyService
|
||||
|
||||
133
src/services/balanceProviders/baseBalanceProvider.js
Normal file
133
src/services/balanceProviders/baseBalanceProvider.js
Normal file
@@ -0,0 +1,133 @@
|
||||
const axios = require('axios')
|
||||
const logger = require('../../utils/logger')
|
||||
const ProxyHelper = require('../../utils/proxyHelper')
|
||||
|
||||
/**
|
||||
* Provider 抽象基类
|
||||
* 各平台 Provider 需继承并实现 queryBalance(account)
|
||||
*/
|
||||
class BaseBalanceProvider {
|
||||
constructor(platform) {
|
||||
this.platform = platform
|
||||
this.logger = logger
|
||||
}
|
||||
|
||||
/**
|
||||
* 查询余额（抽象方法）
|
||||
* @param {object} account - 账户对象
|
||||
* @returns {Promise<object>}
|
||||
* 形如：
|
||||
* {
|
||||
* balance: number|null,
|
||||
* currency?: string,
|
||||
* quota?: { daily, used, remaining, resetAt, percentage, unlimited? },
|
||||
* queryMethod?: 'api'|'field'|'local',
|
||||
* rawData?: any
|
||||
* }
|
||||
*/
|
||||
async queryBalance(_account) {
|
||||
throw new Error('queryBalance 方法必须由子类实现')
|
||||
}
|
||||
|
||||
/**
|
||||
* 通用 HTTP 请求方法（支持代理）
|
||||
* @param {string} url
|
||||
* @param {object} options
|
||||
* @param {object} account
|
||||
*/
|
||||
async makeRequest(url, options = {}, account = {}) {
|
||||
const config = {
|
||||
url,
|
||||
method: options.method || 'GET',
|
||||
headers: options.headers || {},
|
||||
timeout: options.timeout || 15000,
|
||||
data: options.data,
|
||||
params: options.params,
|
||||
responseType: options.responseType
|
||||
}
|
||||
|
||||
const proxyConfig = account.proxyConfig || account.proxy
|
||||
if (proxyConfig) {
|
||||
const agent = ProxyHelper.createProxyAgent(proxyConfig)
|
||||
if (agent) {
|
||||
config.httpAgent = agent
|
||||
config.httpsAgent = agent
|
||||
config.proxy = false
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await axios(config)
|
||||
return {
|
||||
success: true,
|
||||
data: response.data,
|
||||
status: response.status,
|
||||
headers: response.headers
|
||||
}
|
||||
} catch (error) {
|
||||
const status = error.response?.status
|
||||
const message = error.response?.data?.message || error.message || '请求失败'
|
||||
this.logger.debug(`余额 Provider HTTP 请求失败: ${url} (${this.platform})`, {
|
||||
status,
|
||||
message
|
||||
})
|
||||
return { success: false, status, error: message }
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 从账户字段读取 dailyQuota / dailyUsage（通用降级方案）
|
||||
* 注意：部分平台 dailyUsage 字段可能不是实时值，最终以 AccountBalanceService 的本地统计为准
|
||||
*/
|
||||
readQuotaFromFields(account) {
|
||||
const dailyQuota = Number(account?.dailyQuota || 0)
|
||||
const dailyUsage = Number(account?.dailyUsage || 0)
|
||||
|
||||
// 无限制
|
||||
if (!Number.isFinite(dailyQuota) || dailyQuota <= 0) {
|
||||
return {
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
quota: {
|
||||
daily: Infinity,
|
||||
used: Number.isFinite(dailyUsage) ? dailyUsage : 0,
|
||||
remaining: Infinity,
|
||||
percentage: 0,
|
||||
unlimited: true
|
||||
},
|
||||
queryMethod: 'field'
|
||||
}
|
||||
}
|
||||
|
||||
const used = Number.isFinite(dailyUsage) ? dailyUsage : 0
|
||||
const remaining = Math.max(0, dailyQuota - used)
|
||||
const percentage = dailyQuota > 0 ? (used / dailyQuota) * 100 : 0
|
||||
|
||||
return {
|
||||
balance: remaining,
|
||||
currency: 'USD',
|
||||
quota: {
|
||||
daily: dailyQuota,
|
||||
used,
|
||||
remaining,
|
||||
percentage: Math.round(percentage * 100) / 100
|
||||
},
|
||||
queryMethod: 'field'
|
||||
}
|
||||
}
|
||||
|
||||
parseCurrency(data) {
|
||||
return data?.currency || data?.Currency || 'USD'
|
||||
}
|
||||
|
||||
async safeExecute(fn, fallbackValue = null) {
|
||||
try {
|
||||
return await fn()
|
||||
} catch (error) {
|
||||
this.logger.error(`余额 Provider 执行失败: ${this.platform}`, error)
|
||||
return fallbackValue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = BaseBalanceProvider
|
||||
30
src/services/balanceProviders/claudeBalanceProvider.js
Normal file
30
src/services/balanceProviders/claudeBalanceProvider.js
Normal file
@@ -0,0 +1,30 @@
|
||||
const BaseBalanceProvider = require('./baseBalanceProvider')
|
||||
const claudeAccountService = require('../claudeAccountService')
|
||||
|
||||
class ClaudeBalanceProvider extends BaseBalanceProvider {
|
||||
constructor() {
|
||||
super('claude')
|
||||
}
|
||||
|
||||
/**
|
||||
* Claude（OAuth）：优先尝试获取 OAuth usage（用于配额/使用信息），不强行提供余额金额
|
||||
*/
|
||||
async queryBalance(account) {
|
||||
this.logger.debug(`查询 Claude 余额（OAuth usage）: ${account?.id}`)
|
||||
|
||||
// 仅 OAuth 账户可用；失败时降级
|
||||
const usageData = await claudeAccountService.fetchOAuthUsage(account.id).catch(() => null)
|
||||
if (!usageData) {
|
||||
return { balance: null, currency: 'USD', queryMethod: 'local' }
|
||||
}
|
||||
|
||||
return {
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
queryMethod: 'api',
|
||||
rawData: usageData
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = ClaudeBalanceProvider
|
||||
@@ -0,0 +1,14 @@
|
||||
const BaseBalanceProvider = require('./baseBalanceProvider')
|
||||
|
||||
class ClaudeConsoleBalanceProvider extends BaseBalanceProvider {
|
||||
constructor() {
|
||||
super('claude-console')
|
||||
}
|
||||
|
||||
async queryBalance(account) {
|
||||
this.logger.debug(`查询 Claude Console 余额（字段）: ${account?.id}`)
|
||||
return this.readQuotaFromFields(account)
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = ClaudeConsoleBalanceProvider
|
||||
250
src/services/balanceProviders/geminiBalanceProvider.js
Normal file
250
src/services/balanceProviders/geminiBalanceProvider.js
Normal file
@@ -0,0 +1,250 @@
|
||||
const BaseBalanceProvider = require('./baseBalanceProvider')
|
||||
const antigravityClient = require('../antigravityClient')
|
||||
const geminiAccountService = require('../geminiAccountService')
|
||||
|
||||
const OAUTH_PROVIDER_ANTIGRAVITY = 'antigravity'
|
||||
|
||||
function clamp01(value) {
|
||||
if (typeof value !== 'number' || !Number.isFinite(value)) {
|
||||
return null
|
||||
}
|
||||
if (value < 0) {
|
||||
return 0
|
||||
}
|
||||
if (value > 1) {
|
||||
return 1
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
function round2(value) {
|
||||
if (typeof value !== 'number' || !Number.isFinite(value)) {
|
||||
return null
|
||||
}
|
||||
return Math.round(value * 100) / 100
|
||||
}
|
||||
|
||||
function normalizeQuotaCategory(displayName, modelId) {
|
||||
const name = String(displayName || '')
|
||||
const id = String(modelId || '')
|
||||
|
||||
if (name.includes('Gemini') && name.includes('Pro')) {
|
||||
return 'Gemini Pro'
|
||||
}
|
||||
if (name.includes('Gemini') && name.includes('Flash')) {
|
||||
return 'Gemini Flash'
|
||||
}
|
||||
if (name.includes('Gemini') && name.toLowerCase().includes('image')) {
|
||||
return 'Gemini Image'
|
||||
}
|
||||
|
||||
if (name.includes('Claude') || name.includes('GPT-OSS')) {
|
||||
return 'Claude'
|
||||
}
|
||||
|
||||
if (id.startsWith('gemini-3-pro-') || id.startsWith('gemini-2.5-pro')) {
|
||||
return 'Gemini Pro'
|
||||
}
|
||||
if (id.startsWith('gemini-3-flash') || id.startsWith('gemini-2.5-flash')) {
|
||||
return 'Gemini Flash'
|
||||
}
|
||||
if (id.includes('image')) {
|
||||
return 'Gemini Image'
|
||||
}
|
||||
if (id.includes('claude') || id.includes('gpt-oss')) {
|
||||
return 'Claude'
|
||||
}
|
||||
|
||||
return name || id || 'Unknown'
|
||||
}
|
||||
|
||||
function buildAntigravityQuota(modelsResponse) {
|
||||
const models = modelsResponse && typeof modelsResponse === 'object' ? modelsResponse.models : null
|
||||
|
||||
if (!models || typeof models !== 'object') {
|
||||
return null
|
||||
}
|
||||
|
||||
const parseRemainingFraction = (quotaInfo) => {
|
||||
if (!quotaInfo || typeof quotaInfo !== 'object') {
|
||||
return null
|
||||
}
|
||||
|
||||
const raw =
|
||||
quotaInfo.remainingFraction ??
|
||||
quotaInfo.remaining_fraction ??
|
||||
quotaInfo.remaining ??
|
||||
undefined
|
||||
|
||||
const num = typeof raw === 'number' ? raw : typeof raw === 'string' ? Number(raw) : NaN
|
||||
if (!Number.isFinite(num)) {
|
||||
return null
|
||||
}
|
||||
|
||||
return clamp01(num)
|
||||
}
|
||||
|
||||
const allowedCategories = new Set(['Gemini Pro', 'Claude', 'Gemini Flash', 'Gemini Image'])
|
||||
const fixedOrder = ['Gemini Pro', 'Claude', 'Gemini Flash', 'Gemini Image']
|
||||
|
||||
const categoryMap = new Map()
|
||||
|
||||
for (const [modelId, modelDataRaw] of Object.entries(models)) {
|
||||
if (!modelDataRaw || typeof modelDataRaw !== 'object') {
|
||||
continue
|
||||
}
|
||||
|
||||
const displayName = modelDataRaw.displayName || modelDataRaw.display_name || modelId
|
||||
const quotaInfo = modelDataRaw.quotaInfo || modelDataRaw.quota_info || null
|
||||
|
||||
const remainingFraction = parseRemainingFraction(quotaInfo)
|
||||
if (remainingFraction === null) {
|
||||
continue
|
||||
}
|
||||
|
||||
const remainingPercent = round2(remainingFraction * 100)
|
||||
const usedPercent = round2(100 - remainingPercent)
|
||||
const resetAt = quotaInfo?.resetTime || quotaInfo?.reset_time || null
|
||||
|
||||
const category = normalizeQuotaCategory(displayName, modelId)
|
||||
if (!allowedCategories.has(category)) {
|
||||
continue
|
||||
}
|
||||
const entry = {
|
||||
category,
|
||||
modelId,
|
||||
displayName: String(displayName || modelId || category),
|
||||
remainingPercent,
|
||||
usedPercent,
|
||||
resetAt: typeof resetAt === 'string' && resetAt.trim() ? resetAt : null
|
||||
}
|
||||
|
||||
const existing = categoryMap.get(category)
|
||||
if (!existing || entry.remainingPercent < existing.remainingPercent) {
|
||||
categoryMap.set(category, entry)
|
||||
}
|
||||
}
|
||||
|
||||
const buckets = fixedOrder.map((category) => {
|
||||
const existing = categoryMap.get(category) || null
|
||||
if (existing) {
|
||||
return existing
|
||||
}
|
||||
return {
|
||||
category,
|
||||
modelId: '',
|
||||
displayName: category,
|
||||
remainingPercent: null,
|
||||
usedPercent: null,
|
||||
resetAt: null
|
||||
}
|
||||
})
|
||||
|
||||
if (buckets.length === 0) {
|
||||
return null
|
||||
}
|
||||
|
||||
const critical = buckets
|
||||
.filter((item) => item.remainingPercent !== null)
|
||||
.reduce((min, item) => {
|
||||
if (!min) {
|
||||
return item
|
||||
}
|
||||
return (item.remainingPercent ?? 0) < (min.remainingPercent ?? 0) ? item : min
|
||||
}, null)
|
||||
|
||||
if (!critical) {
|
||||
return null
|
||||
}
|
||||
|
||||
return {
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
quota: {
|
||||
type: 'antigravity',
|
||||
total: 100,
|
||||
used: critical.usedPercent,
|
||||
remaining: critical.remainingPercent,
|
||||
percentage: critical.usedPercent,
|
||||
resetAt: critical.resetAt,
|
||||
buckets: buckets.map((item) => ({
|
||||
category: item.category,
|
||||
remaining: item.remainingPercent,
|
||||
used: item.usedPercent,
|
||||
percentage: item.usedPercent,
|
||||
resetAt: item.resetAt
|
||||
}))
|
||||
},
|
||||
queryMethod: 'api',
|
||||
rawData: {
|
||||
modelsCount: Object.keys(models).length,
|
||||
bucketCount: buckets.length
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
class GeminiBalanceProvider extends BaseBalanceProvider {
|
||||
constructor() {
|
||||
super('gemini')
|
||||
}
|
||||
|
||||
async queryBalance(account) {
|
||||
const oauthProvider = account?.oauthProvider
|
||||
if (oauthProvider !== OAUTH_PROVIDER_ANTIGRAVITY) {
|
||||
if (account && Object.prototype.hasOwnProperty.call(account, 'dailyQuota')) {
|
||||
return this.readQuotaFromFields(account)
|
||||
}
|
||||
return { balance: null, currency: 'USD', queryMethod: 'local' }
|
||||
}
|
||||
|
||||
const accessToken = String(account?.accessToken || '').trim()
|
||||
const refreshToken = String(account?.refreshToken || '').trim()
|
||||
const proxyConfig = account?.proxyConfig || account?.proxy || null
|
||||
|
||||
if (!accessToken) {
|
||||
throw new Error('Antigravity 账户缺少 accessToken')
|
||||
}
|
||||
|
||||
const fetch = async (token) =>
|
||||
await antigravityClient.fetchAvailableModels({
|
||||
accessToken: token,
|
||||
proxyConfig
|
||||
})
|
||||
|
||||
let data
|
||||
try {
|
||||
data = await fetch(accessToken)
|
||||
} catch (error) {
|
||||
const status = error?.response?.status
|
||||
if ((status === 401 || status === 403) && refreshToken) {
|
||||
const refreshed = await geminiAccountService.refreshAccessToken(
|
||||
refreshToken,
|
||||
proxyConfig,
|
||||
OAUTH_PROVIDER_ANTIGRAVITY
|
||||
)
|
||||
const nextToken = String(refreshed?.access_token || '').trim()
|
||||
if (!nextToken) {
|
||||
throw error
|
||||
}
|
||||
data = await fetch(nextToken)
|
||||
} else {
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
const mapped = buildAntigravityQuota(data)
|
||||
if (!mapped) {
|
||||
return {
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
quota: null,
|
||||
queryMethod: 'api',
|
||||
rawData: data || null
|
||||
}
|
||||
}
|
||||
|
||||
return mapped
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = GeminiBalanceProvider
|
||||
23
src/services/balanceProviders/genericBalanceProvider.js
Normal file
23
src/services/balanceProviders/genericBalanceProvider.js
Normal file
@@ -0,0 +1,23 @@
|
||||
const BaseBalanceProvider = require('./baseBalanceProvider')
|
||||
|
||||
class GenericBalanceProvider extends BaseBalanceProvider {
|
||||
constructor(platform) {
|
||||
super(platform)
|
||||
}
|
||||
|
||||
async queryBalance(account) {
|
||||
this.logger.debug(`${this.platform} 暂无专用余额 API，实现降级策略`)
|
||||
|
||||
if (account && Object.prototype.hasOwnProperty.call(account, 'dailyQuota')) {
|
||||
return this.readQuotaFromFields(account)
|
||||
}
|
||||
|
||||
return {
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
queryMethod: 'local'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = GenericBalanceProvider
|
||||
25
src/services/balanceProviders/index.js
Normal file
25
src/services/balanceProviders/index.js
Normal file
@@ -0,0 +1,25 @@
|
||||
const ClaudeBalanceProvider = require('./claudeBalanceProvider')
|
||||
const ClaudeConsoleBalanceProvider = require('./claudeConsoleBalanceProvider')
|
||||
const OpenAIResponsesBalanceProvider = require('./openaiResponsesBalanceProvider')
|
||||
const GenericBalanceProvider = require('./genericBalanceProvider')
|
||||
const GeminiBalanceProvider = require('./geminiBalanceProvider')
|
||||
|
||||
function registerAllProviders(balanceService) {
|
||||
// Claude
|
||||
balanceService.registerProvider('claude', new ClaudeBalanceProvider())
|
||||
balanceService.registerProvider('claude-console', new ClaudeConsoleBalanceProvider())
|
||||
|
||||
// OpenAI / Codex
|
||||
balanceService.registerProvider('openai-responses', new OpenAIResponsesBalanceProvider())
|
||||
balanceService.registerProvider('openai', new GenericBalanceProvider('openai'))
|
||||
balanceService.registerProvider('azure_openai', new GenericBalanceProvider('azure_openai'))
|
||||
|
||||
// 其他平台（降级）
|
||||
balanceService.registerProvider('gemini', new GeminiBalanceProvider())
|
||||
balanceService.registerProvider('gemini-api', new GenericBalanceProvider('gemini-api'))
|
||||
balanceService.registerProvider('bedrock', new GenericBalanceProvider('bedrock'))
|
||||
balanceService.registerProvider('droid', new GenericBalanceProvider('droid'))
|
||||
balanceService.registerProvider('ccr', new GenericBalanceProvider('ccr'))
|
||||
}
|
||||
|
||||
module.exports = { registerAllProviders }
|
||||
@@ -0,0 +1,54 @@
|
||||
const BaseBalanceProvider = require('./baseBalanceProvider')
|
||||
|
||||
class OpenAIResponsesBalanceProvider extends BaseBalanceProvider {
|
||||
constructor() {
|
||||
super('openai-responses')
|
||||
}
|
||||
|
||||
/**
|
||||
* OpenAI-Responses：
|
||||
* - 优先使用 dailyQuota 字段（如果配置了额度）
|
||||
* - 可选：尝试调用兼容 API（不同服务商实现不一，失败自动降级）
|
||||
*/
|
||||
async queryBalance(account) {
|
||||
this.logger.debug(`查询 OpenAI Responses 余额: ${account?.id}`)
|
||||
|
||||
// 配置了额度时直接返回（字段法）
|
||||
if (account?.dailyQuota && Number(account.dailyQuota) > 0) {
|
||||
return this.readQuotaFromFields(account)
|
||||
}
|
||||
|
||||
// 尝试调用 usage 接口（兼容性不保证）
|
||||
if (account?.apiKey && account?.baseApi) {
|
||||
const baseApi = String(account.baseApi).replace(/\/$/, '')
|
||||
const response = await this.makeRequest(
|
||||
`${baseApi}/v1/usage`,
|
||||
{
|
||||
method: 'GET',
|
||||
headers: {
|
||||
Authorization: `Bearer ${account.apiKey}`,
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
},
|
||||
account
|
||||
)
|
||||
|
||||
if (response.success) {
|
||||
return {
|
||||
balance: null,
|
||||
currency: this.parseCurrency(response.data),
|
||||
queryMethod: 'api',
|
||||
rawData: response.data
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
balance: null,
|
||||
currency: 'USD',
|
||||
queryMethod: 'local'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = OpenAIResponsesBalanceProvider
|
||||
210
src/services/balanceScriptService.js
Normal file
210
src/services/balanceScriptService.js
Normal file
@@ -0,0 +1,210 @@
|
||||
const vm = require('vm')
|
||||
const axios = require('axios')
|
||||
const { isBalanceScriptEnabled } = require('../utils/featureFlags')
|
||||
|
||||
/**
|
||||
* SSRF防护：检查URL是否访问内网或敏感地址
|
||||
* @param {string} url - 要检查的URL
|
||||
* @returns {boolean} - true表示URL安全
|
||||
*/
|
||||
function isUrlSafe(url) {
|
||||
try {
|
||||
const parsed = new URL(url)
|
||||
const hostname = parsed.hostname.toLowerCase()
|
||||
|
||||
// 禁止的协议
|
||||
if (!['http:', 'https:'].includes(parsed.protocol)) {
|
||||
return false
|
||||
}
|
||||
|
||||
// 禁止访问localhost和私有IP
|
||||
const privatePatterns = [
|
||||
/^localhost$/i,
|
||||
/^127\./,
|
||||
/^10\./,
|
||||
/^172\.(1[6-9]|2[0-9]|3[0-1])\./,
|
||||
/^192\.168\./,
|
||||
/^169\.254\./, // AWS metadata
|
||||
/^0\./, // 0.0.0.0
|
||||
/^::1$/,
|
||||
/^fc00:/i,
|
||||
/^fe80:/i,
|
||||
/\.local$/i,
|
||||
/\.internal$/i,
|
||||
/\.localhost$/i
|
||||
]
|
||||
|
||||
for (const pattern of privatePatterns) {
|
||||
if (pattern.test(hostname)) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
} catch {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 可配置脚本余额查询执行器
|
||||
* - 脚本格式：({ request: {...}, extractor: function(response){...} })
|
||||
* - 模板变量：{{baseUrl}}, {{apiKey}}, {{token}}, {{accountId}}, {{platform}}, {{extra}}
|
||||
*/
|
||||
class BalanceScriptService {
|
||||
/**
|
||||
* 执行脚本：返回标准余额结构 + 原始响应
|
||||
* @param {object} options
|
||||
* - scriptBody: string
|
||||
* - variables: Record<string,string>
|
||||
* - timeoutSeconds: number
|
||||
*/
|
||||
async execute(options = {}) {
|
||||
if (!isBalanceScriptEnabled()) {
|
||||
const error = new Error('余额脚本功能已禁用（可通过 BALANCE_SCRIPT_ENABLED=true 启用）')
|
||||
error.code = 'BALANCE_SCRIPT_DISABLED'
|
||||
throw error
|
||||
}
|
||||
|
||||
const scriptBody = options.scriptBody?.trim()
|
||||
if (!scriptBody) {
|
||||
throw new Error('脚本内容为空')
|
||||
}
|
||||
|
||||
const timeoutMs = Math.max(1, (options.timeoutSeconds || 10) * 1000)
|
||||
const sandbox = {
|
||||
console,
|
||||
Math,
|
||||
Date
|
||||
}
|
||||
|
||||
let scriptResult
|
||||
try {
|
||||
const wrapped = scriptBody.startsWith('(') ? scriptBody : `(${scriptBody})`
|
||||
const script = new vm.Script(wrapped)
|
||||
scriptResult = script.runInNewContext(sandbox, { timeout: timeoutMs })
|
||||
} catch (error) {
|
||||
throw new Error(`脚本解析失败: ${error.message}`)
|
||||
}
|
||||
|
||||
if (!scriptResult || typeof scriptResult !== 'object') {
|
||||
throw new Error('脚本返回格式无效（需返回 { request, extractor }）')
|
||||
}
|
||||
|
||||
const variables = options.variables || {}
|
||||
const request = this.applyTemplates(scriptResult.request || {}, variables)
|
||||
const { extractor } = scriptResult
|
||||
|
||||
if (!request?.url || typeof request.url !== 'string') {
|
||||
throw new Error('脚本 request.url 不能为空')
|
||||
}
|
||||
|
||||
// SSRF防护：验证URL安全性
|
||||
if (!isUrlSafe(request.url)) {
|
||||
throw new Error('脚本 request.url 不安全：禁止访问内网地址、localhost或使用非HTTP(S)协议')
|
||||
}
|
||||
|
||||
if (typeof extractor !== 'function') {
|
||||
throw new Error('脚本 extractor 必须是函数')
|
||||
}
|
||||
|
||||
const axiosConfig = {
|
||||
url: request.url,
|
||||
method: (request.method || 'GET').toUpperCase(),
|
||||
headers: request.headers || {},
|
||||
timeout: timeoutMs
|
||||
}
|
||||
|
||||
if (request.params) {
|
||||
axiosConfig.params = request.params
|
||||
}
|
||||
if (request.body || request.data) {
|
||||
axiosConfig.data = request.body || request.data
|
||||
}
|
||||
|
||||
let httpResponse
|
||||
try {
|
||||
httpResponse = await axios(axiosConfig)
|
||||
} catch (error) {
|
||||
const { response } = error || {}
|
||||
const { status, data } = response || {}
|
||||
throw new Error(
|
||||
`请求失败: ${status || ''} ${error.message}${data ? ` | ${JSON.stringify(data)}` : ''}`
|
||||
)
|
||||
}
|
||||
|
||||
const responseData = httpResponse?.data
|
||||
|
||||
let extracted = {}
|
||||
try {
|
||||
extracted = extractor(responseData) || {}
|
||||
} catch (error) {
|
||||
throw new Error(`extractor 执行失败: ${error.message}`)
|
||||
}
|
||||
|
||||
const mapped = this.mapExtractorResult(extracted, responseData)
|
||||
return {
|
||||
mapped,
|
||||
extracted,
|
||||
response: {
|
||||
status: httpResponse?.status,
|
||||
headers: httpResponse?.headers,
|
||||
data: responseData
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
applyTemplates(value, variables) {
|
||||
if (typeof value === 'string') {
|
||||
return value.replace(/{{(\w+)}}/g, (_, key) => {
|
||||
const trimmed = key.trim()
|
||||
return variables[trimmed] !== undefined ? String(variables[trimmed]) : ''
|
||||
})
|
||||
}
|
||||
if (Array.isArray(value)) {
|
||||
return value.map((item) => this.applyTemplates(item, variables))
|
||||
}
|
||||
if (value && typeof value === 'object') {
|
||||
const result = {}
|
||||
Object.keys(value).forEach((k) => {
|
||||
result[k] = this.applyTemplates(value[k], variables)
|
||||
})
|
||||
return result
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
mapExtractorResult(result = {}, responseData) {
|
||||
const isValid = result.isValid !== false
|
||||
const remaining = Number(result.remaining)
|
||||
const total = Number(result.total)
|
||||
const used = Number(result.used)
|
||||
const currency = result.unit || 'USD'
|
||||
|
||||
const quota =
|
||||
Number.isFinite(total) || Number.isFinite(used)
|
||||
? {
|
||||
total: Number.isFinite(total) ? total : null,
|
||||
used: Number.isFinite(used) ? used : null,
|
||||
remaining: Number.isFinite(remaining) ? remaining : null,
|
||||
percentage:
|
||||
Number.isFinite(total) && total > 0 && Number.isFinite(used)
|
||||
? (used / total) * 100
|
||||
: null
|
||||
}
|
||||
: null
|
||||
|
||||
return {
|
||||
status: isValid ? 'success' : 'error',
|
||||
errorMessage: isValid ? '' : result.invalidMessage || '套餐无效',
|
||||
balance: Number.isFinite(remaining) ? remaining : null,
|
||||
currency,
|
||||
quota,
|
||||
planName: result.planName || null,
|
||||
extra: result.extra || null,
|
||||
rawData: responseData || result.raw
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = new BalanceScriptService()
|
||||
@@ -16,11 +16,62 @@ const {
|
||||
} = require('../utils/tokenRefreshLogger')
|
||||
const tokenRefreshService = require('./tokenRefreshService')
|
||||
const LRUCache = require('../utils/lruCache')
|
||||
const antigravityClient = require('./antigravityClient')
|
||||
|
||||
// Gemini CLI OAuth 配置 - 这些是公开的 Gemini CLI 凭据
|
||||
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']
|
||||
// Gemini OAuth 配置 - 支持 Gemini CLI 与 Antigravity 两种 OAuth 应用
|
||||
const OAUTH_PROVIDER_GEMINI_CLI = 'gemini-cli'
|
||||
const OAUTH_PROVIDER_ANTIGRAVITY = 'antigravity'
|
||||
|
||||
const OAUTH_PROVIDERS = {
|
||||
[OAUTH_PROVIDER_GEMINI_CLI]: {
|
||||
// Gemini CLI OAuth 配置（公开）
|
||||
clientId:
|
||||
process.env.GEMINI_OAUTH_CLIENT_ID ||
|
||||
'681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com',
|
||||
clientSecret: process.env.GEMINI_OAUTH_CLIENT_SECRET || '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',
|
||||
scopes: ['https://www.googleapis.com/auth/cloud-platform']
|
||||
},
|
||||
[OAUTH_PROVIDER_ANTIGRAVITY]: {
|
||||
// Antigravity OAuth 配置（参考 gcli2api）
|
||||
clientId:
|
||||
process.env.ANTIGRAVITY_OAUTH_CLIENT_ID ||
|
||||
'1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com',
|
||||
clientSecret:
|
||||
process.env.ANTIGRAVITY_OAUTH_CLIENT_SECRET || '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',
|
||||
scopes: [
|
||||
'https://www.googleapis.com/auth/cloud-platform',
|
||||
'https://www.googleapis.com/auth/userinfo.email',
|
||||
'https://www.googleapis.com/auth/userinfo.profile',
|
||||
'https://www.googleapis.com/auth/cclog',
|
||||
'https://www.googleapis.com/auth/experimentsandconfigs'
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
if (!process.env.GEMINI_OAUTH_CLIENT_SECRET) {
|
||||
logger.warn(
|
||||
'⚠️ GEMINI_OAUTH_CLIENT_SECRET 未设置，使用内置默认值（建议在生产环境通过环境变量覆盖）'
|
||||
)
|
||||
}
|
||||
if (!process.env.ANTIGRAVITY_OAUTH_CLIENT_SECRET) {
|
||||
logger.warn(
|
||||
'⚠️ ANTIGRAVITY_OAUTH_CLIENT_SECRET 未设置，使用内置默认值（建议在生产环境通过环境变量覆盖）'
|
||||
)
|
||||
}
|
||||
|
||||
function normalizeOauthProvider(oauthProvider) {
|
||||
if (!oauthProvider) {
|
||||
return OAUTH_PROVIDER_GEMINI_CLI
|
||||
}
|
||||
return oauthProvider === OAUTH_PROVIDER_ANTIGRAVITY
|
||||
? OAUTH_PROVIDER_ANTIGRAVITY
|
||||
: OAUTH_PROVIDER_GEMINI_CLI
|
||||
}
|
||||
|
||||
function getOauthProviderConfig(oauthProvider) {
|
||||
const normalized = normalizeOauthProvider(oauthProvider)
|
||||
return OAUTH_PROVIDERS[normalized] || OAUTH_PROVIDERS[OAUTH_PROVIDER_GEMINI_CLI]
|
||||
}
|
||||
|
||||
// 🌐 TCP Keep-Alive Agent 配置
|
||||
// 解决长时间流式请求中 NAT/防火墙空闲超时导致的连接中断问题
|
||||
@@ -34,6 +85,117 @@ const keepAliveAgent = new https.Agent({
|
||||
|
||||
logger.info('🌐 Gemini HTTPS Agent initialized with TCP Keep-Alive support')
|
||||
|
||||
async function fetchAvailableModelsAntigravity(
|
||||
accessToken,
|
||||
proxyConfig = null,
|
||||
refreshToken = null
|
||||
) {
|
||||
try {
|
||||
let effectiveToken = accessToken
|
||||
if (refreshToken) {
|
||||
try {
|
||||
const client = await getOauthClient(
|
||||
accessToken,
|
||||
refreshToken,
|
||||
proxyConfig,
|
||||
OAUTH_PROVIDER_ANTIGRAVITY
|
||||
)
|
||||
if (client && client.getAccessToken) {
|
||||
const latest = await client.getAccessToken()
|
||||
if (latest?.token) {
|
||||
effectiveToken = latest.token
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn('Failed to refresh Antigravity access token for models list:', {
|
||||
message: error.message
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
const data = await antigravityClient.fetchAvailableModels({
|
||||
accessToken: effectiveToken,
|
||||
proxyConfig
|
||||
})
|
||||
const modelsDict = data?.models
|
||||
const created = Math.floor(Date.now() / 1000)
|
||||
|
||||
const models = []
|
||||
const seen = new Set()
|
||||
const {
|
||||
getAntigravityModelAlias,
|
||||
getAntigravityModelMetadata,
|
||||
normalizeAntigravityModelInput
|
||||
} = require('../utils/antigravityModel')
|
||||
|
||||
const pushModel = (modelId) => {
|
||||
if (!modelId || seen.has(modelId)) {
|
||||
return
|
||||
}
|
||||
seen.add(modelId)
|
||||
const metadata = getAntigravityModelMetadata(modelId)
|
||||
const entry = {
|
||||
id: modelId,
|
||||
object: 'model',
|
||||
created,
|
||||
owned_by: 'antigravity'
|
||||
}
|
||||
if (metadata?.name) {
|
||||
entry.name = metadata.name
|
||||
}
|
||||
if (metadata?.maxCompletionTokens) {
|
||||
entry.max_completion_tokens = metadata.maxCompletionTokens
|
||||
}
|
||||
if (metadata?.thinking) {
|
||||
entry.thinking = metadata.thinking
|
||||
}
|
||||
models.push(entry)
|
||||
}
|
||||
|
||||
if (modelsDict && typeof modelsDict === 'object') {
|
||||
for (const modelId of Object.keys(modelsDict)) {
|
||||
const normalized = normalizeAntigravityModelInput(modelId)
|
||||
const alias = getAntigravityModelAlias(normalized)
|
||||
if (!alias) {
|
||||
continue
|
||||
}
|
||||
pushModel(alias)
|
||||
|
||||
if (alias.endsWith('-thinking')) {
|
||||
pushModel(alias.replace(/-thinking$/, ''))
|
||||
}
|
||||
|
||||
if (alias.startsWith('gemini-claude-')) {
|
||||
pushModel(alias.replace(/^gemini-/, ''))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return models
|
||||
} catch (error) {
|
||||
logger.error('Failed to fetch Antigravity models:', error.response?.data || error.message)
|
||||
return [
|
||||
{
|
||||
id: 'gemini-2.5-flash',
|
||||
object: 'model',
|
||||
created: Math.floor(Date.now() / 1000),
|
||||
owned_by: 'antigravity'
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
async function countTokensAntigravity(client, contents, model, proxyConfig = null) {
|
||||
const { token } = await client.getAccessToken()
|
||||
const response = await antigravityClient.countTokens({
|
||||
accessToken: token,
|
||||
proxyConfig,
|
||||
contents,
|
||||
model
|
||||
})
|
||||
return response
|
||||
}
|
||||
|
||||
// 加密相关常量
|
||||
const ALGORITHM = 'aes-256-cbc'
|
||||
const ENCRYPTION_SALT = 'gemini-account-salt'
|
||||
@@ -124,14 +286,15 @@ setInterval(
|
||||
)
|
||||
|
||||
// 创建 OAuth2 客户端（支持代理配置）
|
||||
function createOAuth2Client(redirectUri = null, proxyConfig = null) {
|
||||
function createOAuth2Client(redirectUri = null, proxyConfig = null, oauthProvider = null) {
|
||||
// 如果没有提供 redirectUri，使用默认值
|
||||
const uri = redirectUri || 'http://localhost:45462'
|
||||
const oauthConfig = getOauthProviderConfig(oauthProvider)
|
||||
|
||||
// 准备客户端选项
|
||||
const clientOptions = {
|
||||
clientId: OAUTH_CLIENT_ID,
|
||||
clientSecret: OAUTH_CLIENT_SECRET,
|
||||
clientId: oauthConfig.clientId,
|
||||
clientSecret: oauthConfig.clientSecret,
|
||||
redirectUri: uri
|
||||
}
|
||||
|
||||
@@ -152,10 +315,17 @@ function createOAuth2Client(redirectUri = null, proxyConfig = null) {
|
||||
}
|
||||
|
||||
// 生成授权 URL (支持 PKCE 和代理)
|
||||
async function generateAuthUrl(state = null, redirectUri = null, proxyConfig = null) {
|
||||
async function generateAuthUrl(
|
||||
state = null,
|
||||
redirectUri = null,
|
||||
proxyConfig = null,
|
||||
oauthProvider = null
|
||||
) {
|
||||
// 使用新的 redirect URI
|
||||
const finalRedirectUri = redirectUri || 'https://codeassist.google.com/authcode'
|
||||
const oAuth2Client = createOAuth2Client(finalRedirectUri, proxyConfig)
|
||||
const normalizedProvider = normalizeOauthProvider(oauthProvider)
|
||||
const oauthConfig = getOauthProviderConfig(normalizedProvider)
|
||||
const oAuth2Client = createOAuth2Client(finalRedirectUri, proxyConfig, normalizedProvider)
|
||||
|
||||
if (proxyConfig) {
|
||||
logger.info(
|
||||
@@ -172,7 +342,7 @@ async function generateAuthUrl(state = null, redirectUri = null, proxyConfig = n
|
||||
const authUrl = oAuth2Client.generateAuthUrl({
|
||||
redirect_uri: finalRedirectUri,
|
||||
access_type: 'offline',
|
||||
scope: OAUTH_SCOPES,
|
||||
scope: oauthConfig.scopes,
|
||||
code_challenge_method: 'S256',
|
||||
code_challenge: codeVerifier.codeChallenge,
|
||||
state: stateValue,
|
||||
@@ -183,7 +353,8 @@ async function generateAuthUrl(state = null, redirectUri = null, proxyConfig = n
|
||||
authUrl,
|
||||
state: stateValue,
|
||||
codeVerifier: codeVerifier.codeVerifier,
|
||||
redirectUri: finalRedirectUri
|
||||
redirectUri: finalRedirectUri,
|
||||
oauthProvider: normalizedProvider
|
||||
}
|
||||
}
|
||||
|
||||
@@ -244,11 +415,14 @@ async function exchangeCodeForTokens(
|
||||
code,
|
||||
redirectUri = null,
|
||||
codeVerifier = null,
|
||||
proxyConfig = null
|
||||
proxyConfig = null,
|
||||
oauthProvider = null
|
||||
) {
|
||||
try {
|
||||
const normalizedProvider = normalizeOauthProvider(oauthProvider)
|
||||
const oauthConfig = getOauthProviderConfig(normalizedProvider)
|
||||
// 创建带代理配置的 OAuth2Client
|
||||
const oAuth2Client = createOAuth2Client(redirectUri, proxyConfig)
|
||||
const oAuth2Client = createOAuth2Client(redirectUri, proxyConfig, normalizedProvider)
|
||||
|
||||
if (proxyConfig) {
|
||||
logger.info(
|
||||
@@ -274,7 +448,7 @@ async function exchangeCodeForTokens(
|
||||
return {
|
||||
access_token: tokens.access_token,
|
||||
refresh_token: tokens.refresh_token,
|
||||
scope: tokens.scope || OAUTH_SCOPES.join(' '),
|
||||
scope: tokens.scope || oauthConfig.scopes.join(' '),
|
||||
token_type: tokens.token_type || 'Bearer',
|
||||
expiry_date: tokens.expiry_date || Date.now() + tokens.expires_in * 1000
|
||||
}
|
||||
@@ -285,9 +459,11 @@ async function exchangeCodeForTokens(
|
||||
}
|
||||
|
||||
// 刷新访问令牌
|
||||
async function refreshAccessToken(refreshToken, proxyConfig = null) {
|
||||
async function refreshAccessToken(refreshToken, proxyConfig = null, oauthProvider = null) {
|
||||
const normalizedProvider = normalizeOauthProvider(oauthProvider)
|
||||
const oauthConfig = getOauthProviderConfig(normalizedProvider)
|
||||
// 创建带代理配置的 OAuth2Client
|
||||
const oAuth2Client = createOAuth2Client(null, proxyConfig)
|
||||
const oAuth2Client = createOAuth2Client(null, proxyConfig, normalizedProvider)
|
||||
|
||||
try {
|
||||
// 设置 refresh_token
|
||||
@@ -319,7 +495,7 @@ async function refreshAccessToken(refreshToken, proxyConfig = null) {
|
||||
return {
|
||||
access_token: credentials.access_token,
|
||||
refresh_token: credentials.refresh_token || refreshToken, // 保留原 refresh_token 如果没有返回新的
|
||||
scope: credentials.scope || OAUTH_SCOPES.join(' '),
|
||||
scope: credentials.scope || oauthConfig.scopes.join(' '),
|
||||
token_type: credentials.token_type || 'Bearer',
|
||||
expiry_date: credentials.expiry_date || Date.now() + 3600000 // 默认1小时过期
|
||||
}
|
||||
@@ -339,6 +515,8 @@ async function refreshAccessToken(refreshToken, proxyConfig = null) {
|
||||
async function createAccount(accountData) {
|
||||
const id = uuidv4()
|
||||
const now = new Date().toISOString()
|
||||
const oauthProvider = normalizeOauthProvider(accountData.oauthProvider)
|
||||
const oauthConfig = getOauthProviderConfig(oauthProvider)
|
||||
|
||||
// 处理凭证数据
|
||||
let geminiOauth = null
|
||||
@@ -371,7 +549,7 @@ async function createAccount(accountData) {
|
||||
geminiOauth = JSON.stringify({
|
||||
access_token: accessToken,
|
||||
refresh_token: refreshToken,
|
||||
scope: accountData.scope || OAUTH_SCOPES.join(' '),
|
||||
scope: accountData.scope || oauthConfig.scopes.join(' '),
|
||||
token_type: accountData.tokenType || 'Bearer',
|
||||
expiry_date: accountData.expiryDate || Date.now() + 3600000 // 默认1小时
|
||||
})
|
||||
@@ -399,7 +577,8 @@ async function createAccount(accountData) {
|
||||
refreshToken: refreshToken ? encrypt(refreshToken) : '',
|
||||
expiresAt, // OAuth Token 过期时间（技术字段，自动刷新）
|
||||
// 只有OAuth方式才有scopes，手动添加的没有
|
||||
scopes: accountData.geminiOauth ? accountData.scopes || OAUTH_SCOPES.join(' ') : '',
|
||||
scopes: accountData.geminiOauth ? accountData.scopes || oauthConfig.scopes.join(' ') : '',
|
||||
oauthProvider,
|
||||
|
||||
// ✅ 新增：账户订阅到期时间（业务字段，手动管理）
|
||||
subscriptionExpiresAt: accountData.subscriptionExpiresAt || null,
|
||||
@@ -508,6 +687,10 @@ async function updateAccount(accountId, updates) {
|
||||
updates.schedulable = updates.schedulable.toString()
|
||||
}
|
||||
|
||||
if (updates.oauthProvider !== undefined) {
|
||||
updates.oauthProvider = normalizeOauthProvider(updates.oauthProvider)
|
||||
}
|
||||
|
||||
// 加密敏感字段
|
||||
if (updates.geminiOauth) {
|
||||
updates.geminiOauth = encrypt(
|
||||
@@ -885,12 +1068,13 @@ async function refreshAccountToken(accountId) {
|
||||
// 重新获取账户数据（可能已被其他进程刷新）
|
||||
const updatedAccount = await getAccount(accountId)
|
||||
if (updatedAccount && updatedAccount.accessToken) {
|
||||
const oauthConfig = getOauthProviderConfig(updatedAccount.oauthProvider)
|
||||
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(' '),
|
||||
scope: updatedAccount.scopes || oauthConfig.scopes.join(' '),
|
||||
token_type: 'Bearer'
|
||||
}
|
||||
}
|
||||
@@ -904,7 +1088,11 @@ async function refreshAccountToken(accountId) {
|
||||
|
||||
// account.refreshToken 已经是解密后的值（从 getAccount 返回）
|
||||
// 传入账户的代理配置
|
||||
const newTokens = await refreshAccessToken(account.refreshToken, account.proxy)
|
||||
const newTokens = await refreshAccessToken(
|
||||
account.refreshToken,
|
||||
account.proxy,
|
||||
account.oauthProvider
|
||||
)
|
||||
|
||||
// 更新账户信息
|
||||
const updates = {
|
||||
@@ -1036,14 +1224,15 @@ async function getAccountRateLimitInfo(accountId) {
|
||||
}
|
||||
|
||||
// 获取配置的OAuth客户端 - 参考GeminiCliSimulator的getOauthClient方法（支持代理）
|
||||
async function getOauthClient(accessToken, refreshToken, proxyConfig = null) {
|
||||
const client = createOAuth2Client(null, proxyConfig)
|
||||
async function getOauthClient(accessToken, refreshToken, proxyConfig = null, oauthProvider = null) {
|
||||
const normalizedProvider = normalizeOauthProvider(oauthProvider)
|
||||
const oauthConfig = getOauthProviderConfig(normalizedProvider)
|
||||
const client = createOAuth2Client(null, proxyConfig, normalizedProvider)
|
||||
|
||||
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',
|
||||
scope: oauthConfig.scopes.join(' '),
|
||||
token_type: 'Bearer',
|
||||
expiry_date: 1754269905646
|
||||
}
|
||||
@@ -1509,6 +1698,43 @@ async function generateContent(
|
||||
return response.data
|
||||
}
|
||||
|
||||
// 调用 Antigravity 上游生成内容（非流式）
|
||||
async function generateContentAntigravity(
|
||||
client,
|
||||
requestData,
|
||||
userPromptId,
|
||||
projectId = null,
|
||||
sessionId = null,
|
||||
proxyConfig = null
|
||||
) {
|
||||
const { token } = await client.getAccessToken()
|
||||
const { model } = antigravityClient.buildAntigravityEnvelope({
|
||||
requestData,
|
||||
projectId,
|
||||
sessionId,
|
||||
userPromptId
|
||||
})
|
||||
|
||||
logger.info('🪐 Antigravity generateContent API调用开始', {
|
||||
model,
|
||||
userPromptId,
|
||||
projectId,
|
||||
sessionId
|
||||
})
|
||||
|
||||
const { response } = await antigravityClient.request({
|
||||
accessToken: token,
|
||||
proxyConfig,
|
||||
requestData,
|
||||
projectId,
|
||||
sessionId,
|
||||
userPromptId,
|
||||
stream: false
|
||||
})
|
||||
logger.info('✅ Antigravity generateContent API调用成功')
|
||||
return response.data
|
||||
}
|
||||
|
||||
// 调用 Code Assist API 生成内容（流式）
|
||||
async function generateContentStream(
|
||||
client,
|
||||
@@ -1593,6 +1819,46 @@ async function generateContentStream(
|
||||
return response.data // 返回流对象
|
||||
}
|
||||
|
||||
// 调用 Antigravity 上游生成内容（流式）
|
||||
async function generateContentStreamAntigravity(
|
||||
client,
|
||||
requestData,
|
||||
userPromptId,
|
||||
projectId = null,
|
||||
sessionId = null,
|
||||
signal = null,
|
||||
proxyConfig = null
|
||||
) {
|
||||
const { token } = await client.getAccessToken()
|
||||
const { model } = antigravityClient.buildAntigravityEnvelope({
|
||||
requestData,
|
||||
projectId,
|
||||
sessionId,
|
||||
userPromptId
|
||||
})
|
||||
|
||||
logger.info('🌊 Antigravity streamGenerateContent API调用开始', {
|
||||
model,
|
||||
userPromptId,
|
||||
projectId,
|
||||
sessionId
|
||||
})
|
||||
|
||||
const { response } = await antigravityClient.request({
|
||||
accessToken: token,
|
||||
proxyConfig,
|
||||
requestData,
|
||||
projectId,
|
||||
sessionId,
|
||||
userPromptId,
|
||||
stream: true,
|
||||
signal,
|
||||
params: { alt: 'sse' }
|
||||
})
|
||||
logger.info('✅ Antigravity streamGenerateContent API调用成功，开始流式传输')
|
||||
return response.data
|
||||
}
|
||||
|
||||
// 更新账户的临时项目 ID
|
||||
async function updateTempProjectId(accountId, tempProjectId) {
|
||||
if (!tempProjectId) {
|
||||
@@ -1687,10 +1953,12 @@ module.exports = {
|
||||
generateEncryptionKey,
|
||||
decryptCache, // 暴露缓存对象以便测试和监控
|
||||
countTokens,
|
||||
countTokensAntigravity,
|
||||
generateContent,
|
||||
generateContentStream,
|
||||
generateContentAntigravity,
|
||||
generateContentStreamAntigravity,
|
||||
fetchAvailableModelsAntigravity,
|
||||
updateTempProjectId,
|
||||
resetAccountStatus,
|
||||
OAUTH_CLIENT_ID,
|
||||
OAUTH_SCOPES
|
||||
resetAccountStatus
|
||||
}
|
||||
|
||||
@@ -4,11 +4,35 @@ const accountGroupService = require('./accountGroupService')
|
||||
const redis = require('../models/redis')
|
||||
const logger = require('../utils/logger')
|
||||
|
||||
const OAUTH_PROVIDER_GEMINI_CLI = 'gemini-cli'
|
||||
const OAUTH_PROVIDER_ANTIGRAVITY = 'antigravity'
|
||||
const KNOWN_OAUTH_PROVIDERS = [OAUTH_PROVIDER_GEMINI_CLI, OAUTH_PROVIDER_ANTIGRAVITY]
|
||||
|
||||
function normalizeOauthProvider(oauthProvider) {
|
||||
if (!oauthProvider) {
|
||||
return OAUTH_PROVIDER_GEMINI_CLI
|
||||
}
|
||||
return oauthProvider === OAUTH_PROVIDER_ANTIGRAVITY
|
||||
? OAUTH_PROVIDER_ANTIGRAVITY
|
||||
: OAUTH_PROVIDER_GEMINI_CLI
|
||||
}
|
||||
|
||||
class UnifiedGeminiScheduler {
|
||||
constructor() {
|
||||
this.SESSION_MAPPING_PREFIX = 'unified_gemini_session_mapping:'
|
||||
}
|
||||
|
||||
_getSessionMappingKey(sessionHash, oauthProvider = null) {
|
||||
if (!sessionHash) {
|
||||
return null
|
||||
}
|
||||
if (!oauthProvider) {
|
||||
return `${this.SESSION_MAPPING_PREFIX}${sessionHash}`
|
||||
}
|
||||
const normalized = normalizeOauthProvider(oauthProvider)
|
||||
return `${this.SESSION_MAPPING_PREFIX}${normalized}:${sessionHash}`
|
||||
}
|
||||
|
||||
// 🔧 辅助方法：检查账户是否可调度（兼容字符串和布尔值）
|
||||
_isSchedulable(schedulable) {
|
||||
// 如果是 undefined 或 null，默认为可调度
|
||||
@@ -32,7 +56,8 @@ class UnifiedGeminiScheduler {
|
||||
requestedModel = null,
|
||||
options = {}
|
||||
) {
|
||||
const { allowApiAccounts = false } = options
|
||||
const { allowApiAccounts = false, oauthProvider = null } = options
|
||||
const normalizedOauthProvider = oauthProvider ? normalizeOauthProvider(oauthProvider) : null
|
||||
|
||||
try {
|
||||
// 如果API Key绑定了专属账户或分组，优先使用
|
||||
@@ -83,14 +108,23 @@ class UnifiedGeminiScheduler {
|
||||
this._isActive(boundAccount.isActive) &&
|
||||
boundAccount.status !== 'error'
|
||||
) {
|
||||
logger.info(
|
||||
`🎯 Using bound dedicated Gemini account: ${boundAccount.name} (${apiKeyData.geminiAccountId}) for API key ${apiKeyData.name}`
|
||||
)
|
||||
// 更新账户的最后使用时间
|
||||
await geminiAccountService.markAccountUsed(apiKeyData.geminiAccountId)
|
||||
return {
|
||||
accountId: apiKeyData.geminiAccountId,
|
||||
accountType: 'gemini'
|
||||
if (
|
||||
normalizedOauthProvider &&
|
||||
normalizeOauthProvider(boundAccount.oauthProvider) !== normalizedOauthProvider
|
||||
) {
|
||||
logger.warn(
|
||||
`⚠️ Bound Gemini OAuth account ${boundAccount.name} oauthProvider=${normalizeOauthProvider(boundAccount.oauthProvider)} does not match requested oauthProvider=${normalizedOauthProvider}, falling back to pool`
|
||||
)
|
||||
} else {
|
||||
logger.info(
|
||||
`🎯 Using bound dedicated Gemini account: ${boundAccount.name} (${apiKeyData.geminiAccountId}) for API key ${apiKeyData.name}`
|
||||
)
|
||||
// 更新账户的最后使用时间
|
||||
await geminiAccountService.markAccountUsed(apiKeyData.geminiAccountId)
|
||||
return {
|
||||
accountId: apiKeyData.geminiAccountId,
|
||||
accountType: 'gemini'
|
||||
}
|
||||
}
|
||||
} else {
|
||||
logger.warn(
|
||||
@@ -102,7 +136,7 @@ class UnifiedGeminiScheduler {
|
||||
|
||||
// 如果有会话哈希，检查是否有已映射的账户
|
||||
if (sessionHash) {
|
||||
const mappedAccount = await this._getSessionMapping(sessionHash)
|
||||
const mappedAccount = await this._getSessionMapping(sessionHash, normalizedOauthProvider)
|
||||
if (mappedAccount) {
|
||||
// 验证映射的账户是否仍然可用
|
||||
const isAvailable = await this._isAccountAvailable(
|
||||
@@ -111,7 +145,7 @@ class UnifiedGeminiScheduler {
|
||||
)
|
||||
if (isAvailable) {
|
||||
// 🚀 智能会话续期（续期 unified 映射键，按配置）
|
||||
await this._extendSessionMappingTTL(sessionHash)
|
||||
await this._extendSessionMappingTTL(sessionHash, normalizedOauthProvider)
|
||||
logger.info(
|
||||
`🎯 Using sticky session account: ${mappedAccount.accountId} (${mappedAccount.accountType}) for session ${sessionHash}`
|
||||
)
|
||||
@@ -132,11 +166,10 @@ class UnifiedGeminiScheduler {
|
||||
}
|
||||
|
||||
// 获取所有可用账户
|
||||
const availableAccounts = await this._getAllAvailableAccounts(
|
||||
apiKeyData,
|
||||
requestedModel,
|
||||
allowApiAccounts
|
||||
)
|
||||
const availableAccounts = await this._getAllAvailableAccounts(apiKeyData, requestedModel, {
|
||||
allowApiAccounts,
|
||||
oauthProvider: normalizedOauthProvider
|
||||
})
|
||||
|
||||
if (availableAccounts.length === 0) {
|
||||
// 提供更详细的错误信息
|
||||
@@ -160,7 +193,8 @@ class UnifiedGeminiScheduler {
|
||||
await this._setSessionMapping(
|
||||
sessionHash,
|
||||
selectedAccount.accountId,
|
||||
selectedAccount.accountType
|
||||
selectedAccount.accountType,
|
||||
normalizedOauthProvider
|
||||
)
|
||||
logger.info(
|
||||
`🎯 Created new sticky session mapping: ${selectedAccount.name} (${selectedAccount.accountId}, ${selectedAccount.accountType}) for session ${sessionHash}`
|
||||
@@ -189,7 +223,18 @@ class UnifiedGeminiScheduler {
|
||||
}
|
||||
|
||||
// 📋 获取所有可用账户
|
||||
async _getAllAvailableAccounts(apiKeyData, requestedModel = null, allowApiAccounts = false) {
|
||||
async _getAllAvailableAccounts(
|
||||
apiKeyData,
|
||||
requestedModel = null,
|
||||
allowApiAccountsOrOptions = false
|
||||
) {
|
||||
const options =
|
||||
allowApiAccountsOrOptions && typeof allowApiAccountsOrOptions === 'object'
|
||||
? allowApiAccountsOrOptions
|
||||
: { allowApiAccounts: allowApiAccountsOrOptions }
|
||||
const { allowApiAccounts = false, oauthProvider = null } = options
|
||||
const normalizedOauthProvider = oauthProvider ? normalizeOauthProvider(oauthProvider) : null
|
||||
|
||||
const availableAccounts = []
|
||||
|
||||
// 如果API Key绑定了专属账户，优先返回
|
||||
@@ -254,6 +299,12 @@ class UnifiedGeminiScheduler {
|
||||
this._isActive(boundAccount.isActive) &&
|
||||
boundAccount.status !== 'error'
|
||||
) {
|
||||
if (
|
||||
normalizedOauthProvider &&
|
||||
normalizeOauthProvider(boundAccount.oauthProvider) !== normalizedOauthProvider
|
||||
) {
|
||||
return availableAccounts
|
||||
}
|
||||
const isRateLimited = await this.isAccountRateLimited(boundAccount.id)
|
||||
if (!isRateLimited) {
|
||||
// 检查模型支持
|
||||
@@ -303,6 +354,12 @@ class UnifiedGeminiScheduler {
|
||||
(account.accountType === 'shared' || !account.accountType) && // 兼容旧数据
|
||||
this._isSchedulable(account.schedulable)
|
||||
) {
|
||||
if (
|
||||
normalizedOauthProvider &&
|
||||
normalizeOauthProvider(account.oauthProvider) !== normalizedOauthProvider
|
||||
) {
|
||||
continue
|
||||
}
|
||||
// 检查是否可调度
|
||||
|
||||
// 检查token是否过期
|
||||
@@ -437,9 +494,10 @@ class UnifiedGeminiScheduler {
|
||||
}
|
||||
|
||||
// 🔗 获取会话映射
|
||||
async _getSessionMapping(sessionHash) {
|
||||
async _getSessionMapping(sessionHash, oauthProvider = null) {
|
||||
const client = redis.getClientSafe()
|
||||
const mappingData = await client.get(`${this.SESSION_MAPPING_PREFIX}${sessionHash}`)
|
||||
const key = this._getSessionMappingKey(sessionHash, oauthProvider)
|
||||
const mappingData = key ? await client.get(key) : null
|
||||
|
||||
if (mappingData) {
|
||||
try {
|
||||
@@ -454,27 +512,42 @@ class UnifiedGeminiScheduler {
|
||||
}
|
||||
|
||||
// 💾 设置会话映射
|
||||
async _setSessionMapping(sessionHash, accountId, accountType) {
|
||||
async _setSessionMapping(sessionHash, accountId, accountType, oauthProvider = null) {
|
||||
const client = redis.getClientSafe()
|
||||
const mappingData = JSON.stringify({ accountId, accountType })
|
||||
// 依据配置设置TTL（小时）
|
||||
const appConfig = require('../../config/config')
|
||||
const ttlHours = appConfig.session?.stickyTtlHours || 1
|
||||
const ttlSeconds = Math.max(1, Math.floor(ttlHours * 60 * 60))
|
||||
await client.setex(`${this.SESSION_MAPPING_PREFIX}${sessionHash}`, ttlSeconds, mappingData)
|
||||
const key = this._getSessionMappingKey(sessionHash, oauthProvider)
|
||||
if (!key) {
|
||||
return
|
||||
}
|
||||
await client.setex(key, ttlSeconds, mappingData)
|
||||
}
|
||||
|
||||
// 🗑️ 删除会话映射
|
||||
async _deleteSessionMapping(sessionHash) {
|
||||
const client = redis.getClientSafe()
|
||||
await client.del(`${this.SESSION_MAPPING_PREFIX}${sessionHash}`)
|
||||
if (!sessionHash) {
|
||||
return
|
||||
}
|
||||
|
||||
const keys = [this._getSessionMappingKey(sessionHash)]
|
||||
for (const provider of KNOWN_OAUTH_PROVIDERS) {
|
||||
keys.push(this._getSessionMappingKey(sessionHash, provider))
|
||||
}
|
||||
await client.del(keys.filter(Boolean))
|
||||
}
|
||||
|
||||
// 🔁 续期统一调度会话映射TTL（针对 unified_gemini_session_mapping:* 键），遵循会话配置
|
||||
async _extendSessionMappingTTL(sessionHash) {
|
||||
async _extendSessionMappingTTL(sessionHash, oauthProvider = null) {
|
||||
try {
|
||||
const client = redis.getClientSafe()
|
||||
const key = `${this.SESSION_MAPPING_PREFIX}${sessionHash}`
|
||||
const key = this._getSessionMappingKey(sessionHash, oauthProvider)
|
||||
if (!key) {
|
||||
return false
|
||||
}
|
||||
const remainingTTL = await client.ttl(key)
|
||||
|
||||
if (remainingTTL === -2) {
|
||||
|
||||
Reference in New Issue
Block a user