diff --git a/config/config.example.js b/config/config.example.js index 9cf26002..e5e0c340 100644 --- a/config/config.example.js +++ b/config/config.example.js @@ -205,6 +205,14 @@ const config = { hotReload: process.env.HOT_RELOAD === 'true' }, + // 💰 账户余额相关配置 + accountBalance: { + // 是否允许执行自定义余额脚本（安全开关） + // 说明：脚本能力可发起任意 HTTP 请求并在服务端执行 extractor 逻辑，建议仅在受控环境开启 + // 默认保持开启；如需禁用请显式设置：BALANCE_SCRIPT_ENABLED=false + enableBalanceScript: process.env.BALANCE_SCRIPT_ENABLED !== 'false' + }, + // 📬 用户消息队列配置 // 优化说明：锁在请求发送成功后立即释放（而非请求完成后），因为 Claude API 限流基于请求发送时刻计算 userMessageQueue: { diff --git a/src/app.js b/src/app.js index db15df2e..46391325 100644 --- a/src/app.js +++ b/src/app.js @@ -52,6 +52,16 @@ class Application { await redis.connect() logger.success('✅ Redis connected successfully') + // 💳 初始化账户余额查询服务（Provider 注册） + try { + const accountBalanceService = require('./services/accountBalanceService') + const { registerAllProviders } = require('./services/balanceProviders') + registerAllProviders(accountBalanceService) + logger.info('✅ 账户余额查询服务已初始化') + } catch (error) { + logger.warn('⚠️ 账户余额查询服务初始化失败:', error.message) + } + // 💰 初始化价格服务 logger.info('🔄 Initializing pricing service...') await pricingService.initialize() diff --git a/src/models/redis.js b/src/models/redis.js index 8b41cabc..e69ba727 100644 --- a/src/models/redis.js +++ b/src/models/redis.js @@ -1521,6 +1521,123 @@ class RedisClient { return await this.client.del(key) } + // 💰 账户余额缓存（API 查询结果） + async setAccountBalance(platform, accountId, balanceData, ttl = 3600) { + const key = `account_balance:${platform}:${accountId}` + + const payload = { + balance: + balanceData && balanceData.balance !== null && balanceData.balance !== undefined + ? String(balanceData.balance) + : '', + currency: balanceData?.currency || 'USD', + lastRefreshAt: balanceData?.lastRefreshAt || new Date().toISOString(), + queryMethod: balanceData?.queryMethod || 'api', + status: balanceData?.status || 'success', + errorMessage: balanceData?.errorMessage || balanceData?.error || '', + rawData: balanceData?.rawData ? JSON.stringify(balanceData.rawData) : '', + quota: balanceData?.quota ? JSON.stringify(balanceData.quota) : '' + } + + await this.client.hset(key, payload) + await this.client.expire(key, ttl) + } + + async getAccountBalance(platform, accountId) { + const key = `account_balance:${platform}:${accountId}` + const [data, ttlSeconds] = await Promise.all([this.client.hgetall(key), this.client.ttl(key)]) + + if (!data || Object.keys(data).length === 0) { + return null + } + + let rawData = null + if (data.rawData) { + try { + rawData = JSON.parse(data.rawData) + } catch (error) { + rawData = null + } + } + + let quota = null + if (data.quota) { + try { + quota = JSON.parse(data.quota) + } catch (error) { + quota = null + } + } + + return { + balance: data.balance ? parseFloat(data.balance) : null, + currency: data.currency || 'USD', + lastRefreshAt: data.lastRefreshAt || null, + queryMethod: data.queryMethod || null, + status: data.status || null, + errorMessage: data.errorMessage || '', + rawData, + quota, + ttlSeconds: Number.isFinite(ttlSeconds) ? ttlSeconds : null + } + } + + // 📊 账户余额缓存（本地统计） + async setLocalBalance(platform, accountId, statisticsData, ttl = 300) { + const key = `account_balance_local:${platform}:${accountId}` + + await this.client.hset(key, { + estimatedBalance: JSON.stringify(statisticsData || {}), + lastCalculated: new Date().toISOString() + }) + await this.client.expire(key, ttl) + } + + async getLocalBalance(platform, accountId) { + const key = `account_balance_local:${platform}:${accountId}` + const data = await this.client.hgetall(key) + + if (!data || !data.estimatedBalance) { + return null + } + + try { + return JSON.parse(data.estimatedBalance) + } catch (error) { + return null + } + } + + async deleteAccountBalance(platform, accountId) { + const key = `account_balance:${platform}:${accountId}` + const localKey = `account_balance_local:${platform}:${accountId}` + await this.client.del(key, localKey) + } + + // 🧩 账户余额脚本配置 + async setBalanceScriptConfig(platform, accountId, scriptConfig) { + const key = `account_balance_script:${platform}:${accountId}` + await this.client.set(key, JSON.stringify(scriptConfig || {})) + } + + async getBalanceScriptConfig(platform, accountId) { + const key = `account_balance_script:${platform}:${accountId}` + const raw = await this.client.get(key) + if (!raw) { + return null + } + try { + return JSON.parse(raw) + } catch (error) { + return null + } + } + + async deleteBalanceScriptConfig(platform, accountId) { + const key = `account_balance_script:${platform}:${accountId}` + return await this.client.del(key) + } + // 📈 系统统计 async getSystemStats() { const keys = await Promise.all([ diff --git a/src/routes/admin/accountBalance.js b/src/routes/admin/accountBalance.js new file mode 100644 index 00000000..7f1d18db --- /dev/null +++ b/src/routes/admin/accountBalance.js @@ -0,0 +1,214 @@ +const express = require('express') +const { authenticateAdmin } = require('../../middleware/auth') +const logger = require('../../utils/logger') +const accountBalanceService = require('../../services/accountBalanceService') +const balanceScriptService = require('../../services/balanceScriptService') +const { isBalanceScriptEnabled } = require('../../utils/featureFlags') + +const router = express.Router() + +const ensureValidPlatform = (rawPlatform) => { + const normalized = accountBalanceService.normalizePlatform(rawPlatform) + if (!normalized) { + return { ok: false, status: 400, error: '缺少 platform 参数' } + } + + const supported = accountBalanceService.getSupportedPlatforms() + if (!supported.includes(normalized)) { + return { ok: false, status: 400, error: `不支持的平台: ${normalized}` } + } + + return { ok: true, platform: normalized } +} + +// 1) 获取账户余额（默认本地统计优先，可选触发 Provider） +// GET /admin/accounts/:accountId/balance?platform=xxx&queryApi=false +router.get('/accounts/:accountId/balance', authenticateAdmin, async (req, res) => { + try { + const { accountId } = req.params + const { platform, queryApi } = req.query + + const valid = ensureValidPlatform(platform) + if (!valid.ok) { + return res.status(valid.status).json({ success: false, error: valid.error }) + } + + const balance = await accountBalanceService.getAccountBalance(accountId, valid.platform, { + queryApi + }) + + if (!balance) { + return res.status(404).json({ success: false, error: 'Account not found' }) + } + + return res.json(balance) + } catch (error) { + logger.error('获取账户余额失败', error) + return res.status(500).json({ success: false, error: error.message }) + } +}) + +// 2) 强制刷新账户余额（强制触发查询：优先脚本；Provider 仅为降级） +// POST /admin/accounts/:accountId/balance/refresh +// Body: { platform: 'xxx' } +router.post('/accounts/:accountId/balance/refresh', authenticateAdmin, async (req, res) => { + try { + const { accountId } = req.params + const { platform } = req.body || {} + + const valid = ensureValidPlatform(platform) + if (!valid.ok) { + return res.status(valid.status).json({ success: false, error: valid.error }) + } + + logger.info(`手动刷新余额: ${valid.platform}:${accountId}`) + + const balance = await accountBalanceService.refreshAccountBalance(accountId, valid.platform) + if (!balance) { + return res.status(404).json({ success: false, error: 'Account not found' }) + } + + return res.json(balance) + } catch (error) { + logger.error('刷新账户余额失败', error) + return res.status(500).json({ success: false, error: error.message }) + } +}) + +// 3) 批量获取平台所有账户余额 +// GET /admin/accounts/balance/platform/:platform?queryApi=false +router.get('/accounts/balance/platform/:platform', authenticateAdmin, async (req, res) => { + try { + const { platform } = req.params + const { queryApi } = req.query + + const valid = ensureValidPlatform(platform) + if (!valid.ok) { + return res.status(valid.status).json({ success: false, error: valid.error }) + } + + const balances = await accountBalanceService.getAllAccountsBalance(valid.platform, { queryApi }) + + return res.json({ success: true, data: balances }) + } catch (error) { + logger.error('批量获取余额失败', error) + return res.status(500).json({ success: false, error: error.message }) + } +}) + +// 4) 获取余额汇总（Dashboard 用） +// GET /admin/accounts/balance/summary +router.get('/accounts/balance/summary', authenticateAdmin, async (req, res) => { + try { + const summary = await accountBalanceService.getBalanceSummary() + return res.json({ success: true, data: summary }) + } catch (error) { + logger.error('获取余额汇总失败', error) + return res.status(500).json({ success: false, error: error.message }) + } +}) + +// 5) 清除缓存 +// DELETE /admin/accounts/:accountId/balance/cache?platform=xxx +router.delete('/accounts/:accountId/balance/cache', authenticateAdmin, async (req, res) => { + try { + const { accountId } = req.params + const { platform } = req.query + + const valid = ensureValidPlatform(platform) + if (!valid.ok) { + return res.status(valid.status).json({ success: false, error: valid.error }) + } + + await accountBalanceService.clearCache(accountId, valid.platform) + + return res.json({ success: true, message: '缓存已清除' }) + } catch (error) { + logger.error('清除缓存失败', error) + return res.status(500).json({ success: false, error: error.message }) + } +}) + +// 6) 获取/保存/测试余额脚本配置（单账户） +router.get('/accounts/:accountId/balance/script', authenticateAdmin, async (req, res) => { + try { + const { accountId } = req.params + const { platform } = req.query + + const valid = ensureValidPlatform(platform) + if (!valid.ok) { + return res.status(valid.status).json({ success: false, error: valid.error }) + } + + const config = await accountBalanceService.redis.getBalanceScriptConfig( + valid.platform, + accountId + ) + return res.json({ success: true, data: config || null }) + } catch (error) { + logger.error('获取余额脚本配置失败', error) + return res.status(500).json({ success: false, error: error.message }) + } +}) + +router.put('/accounts/:accountId/balance/script', authenticateAdmin, async (req, res) => { + try { + const { accountId } = req.params + const { platform } = req.query + const valid = ensureValidPlatform(platform) + if (!valid.ok) { + return res.status(valid.status).json({ success: false, error: valid.error }) + } + + const payload = req.body || {} + await accountBalanceService.redis.setBalanceScriptConfig(valid.platform, accountId, payload) + return res.json({ success: true, data: payload }) + } catch (error) { + logger.error('保存余额脚本配置失败', error) + return res.status(500).json({ success: false, error: error.message }) + } +}) + +router.post('/accounts/:accountId/balance/script/test', authenticateAdmin, async (req, res) => { + try { + const { accountId } = req.params + const { platform } = req.query + const valid = ensureValidPlatform(platform) + if (!valid.ok) { + return res.status(valid.status).json({ success: false, error: valid.error }) + } + + if (!isBalanceScriptEnabled()) { + return res.status(403).json({ + success: false, + error: '余额脚本功能已禁用（可通过 BALANCE_SCRIPT_ENABLED=true 启用）' + }) + } + + const payload = req.body || {} + const { scriptBody } = payload + if (!scriptBody) { + return res.status(400).json({ success: false, error: '脚本内容不能为空' }) + } + + const result = await balanceScriptService.execute({ + scriptBody, + timeoutSeconds: payload.timeoutSeconds || 10, + variables: { + baseUrl: payload.baseUrl || '', + apiKey: payload.apiKey || '', + token: payload.token || '', + accountId, + platform: valid.platform, + extra: payload.extra || '' + } + }) + + return res.json({ success: true, data: result }) + } catch (error) { + logger.error('测试余额脚本失败', error) + return res.status(400).json({ success: false, error: error.message }) + } +}) + +module.exports = router diff --git a/src/routes/admin/balanceScripts.js b/src/routes/admin/balanceScripts.js new file mode 100644 index 00000000..ef7ffa01 --- /dev/null +++ b/src/routes/admin/balanceScripts.js @@ -0,0 +1,41 @@ +const express = require('express') +const { authenticateAdmin } = require('../../middleware/auth') +const balanceScriptService = require('../../services/balanceScriptService') +const router = express.Router() + +// 获取全部脚本配置列表 +router.get('/balance-scripts', authenticateAdmin, (req, res) => { + const items = balanceScriptService.listConfigs() + return res.json({ success: true, data: items }) +}) + +// 获取单个脚本配置 +router.get('/balance-scripts/:name', authenticateAdmin, (req, res) => { + const { name } = req.params + const config = balanceScriptService.getConfig(name || 'default') + return res.json({ success: true, data: config }) +}) + +// 保存脚本配置 +router.put('/balance-scripts/:name', authenticateAdmin, (req, res) => { + try { + const { name } = req.params + const saved = balanceScriptService.saveConfig(name || 'default', req.body || {}) + return res.json({ success: true, data: saved }) + } catch (error) { + return res.status(400).json({ success: false, error: error.message }) + } +}) + +// 测试脚本（不落库） +router.post('/balance-scripts/:name/test', authenticateAdmin, async (req, res) => { + try { + const { name } = req.params + const result = await balanceScriptService.testScript(name || 'default', req.body || {}) + return res.json({ success: true, data: result }) + } catch (error) { + return res.status(400).json({ success: false, error: error.message }) + } +}) + +module.exports = router diff --git a/src/routes/admin/index.js b/src/routes/admin/index.js index 0b8cbecd..7fe901c7 100644 --- a/src/routes/admin/index.js +++ b/src/routes/admin/index.js @@ -21,6 +21,7 @@ const openaiResponsesAccountsRoutes = require('./openaiResponsesAccounts') const droidAccountsRoutes = require('./droidAccounts') const dashboardRoutes = require('./dashboard') const usageStatsRoutes = require('./usageStats') +const accountBalanceRoutes = require('./accountBalance') const systemRoutes = require('./system') const concurrencyRoutes = require('./concurrency') const claudeRelayConfigRoutes = require('./claudeRelayConfig') @@ -37,6 +38,7 @@ router.use('/', openaiResponsesAccountsRoutes) router.use('/', droidAccountsRoutes) router.use('/', dashboardRoutes) router.use('/', usageStatsRoutes) +router.use('/', accountBalanceRoutes) router.use('/', systemRoutes) router.use('/', concurrencyRoutes) router.use('/', claudeRelayConfigRoutes) diff --git a/src/services/accountBalanceService.js b/src/services/accountBalanceService.js new file mode 100644 index 00000000..3265c4b8 --- /dev/null +++ b/src/services/accountBalanceService.js @@ -0,0 +1,748 @@ +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 queryApi = this._parseBoolean(options.queryApi) || false + 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) + + // 非强制查询：优先读缓存 + if (!queryApi) { + 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 + ) + } + } + + 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 + } + + 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 diff --git a/src/services/balanceProviders/baseBalanceProvider.js b/src/services/balanceProviders/baseBalanceProvider.js new file mode 100644 index 00000000..ececd2e5 --- /dev/null +++ b/src/services/balanceProviders/baseBalanceProvider.js @@ -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} + * 形如： + * { + * 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 diff --git a/src/services/balanceProviders/claudeBalanceProvider.js b/src/services/balanceProviders/claudeBalanceProvider.js new file mode 100644 index 00000000..89783028 --- /dev/null +++ b/src/services/balanceProviders/claudeBalanceProvider.js @@ -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 diff --git a/src/services/balanceProviders/claudeConsoleBalanceProvider.js b/src/services/balanceProviders/claudeConsoleBalanceProvider.js new file mode 100644 index 00000000..f5441047 --- /dev/null +++ b/src/services/balanceProviders/claudeConsoleBalanceProvider.js @@ -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 diff --git a/src/services/balanceProviders/genericBalanceProvider.js b/src/services/balanceProviders/genericBalanceProvider.js new file mode 100644 index 00000000..6b3efe2b --- /dev/null +++ b/src/services/balanceProviders/genericBalanceProvider.js @@ -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 diff --git a/src/services/balanceProviders/index.js b/src/services/balanceProviders/index.js new file mode 100644 index 00000000..d55fda5b --- /dev/null +++ b/src/services/balanceProviders/index.js @@ -0,0 +1,24 @@ +const ClaudeBalanceProvider = require('./claudeBalanceProvider') +const ClaudeConsoleBalanceProvider = require('./claudeConsoleBalanceProvider') +const OpenAIResponsesBalanceProvider = require('./openaiResponsesBalanceProvider') +const GenericBalanceProvider = require('./genericBalanceProvider') + +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 GenericBalanceProvider('gemini')) + 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 } diff --git a/src/services/balanceProviders/openaiResponsesBalanceProvider.js b/src/services/balanceProviders/openaiResponsesBalanceProvider.js new file mode 100644 index 00000000..9ff8433e --- /dev/null +++ b/src/services/balanceProviders/openaiResponsesBalanceProvider.js @@ -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 diff --git a/src/services/balanceScriptService.js b/src/services/balanceScriptService.js new file mode 100644 index 00000000..5bf06801 --- /dev/null +++ b/src/services/balanceScriptService.js @@ -0,0 +1,161 @@ +const vm = require('vm') +const axios = require('axios') +const { isBalanceScriptEnabled } = require('../utils/featureFlags') + +/** + * 可配置脚本余额查询执行器 + * - 脚本格式：({ request: {...}, extractor: function(response){...} }) + * - 模板变量：{{baseUrl}}, {{apiKey}}, {{token}}, {{accountId}}, {{platform}}, {{extra}} + */ +class BalanceScriptService { + /** + * 执行脚本：返回标准余额结构 + 原始响应 + * @param {object} options + * - scriptBody: string + * - variables: Record + * - 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 不能为空') + } + + 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() diff --git a/src/utils/featureFlags.js b/src/utils/featureFlags.js new file mode 100644 index 00000000..35802d55 --- /dev/null +++ b/src/utils/featureFlags.js @@ -0,0 +1,44 @@ +let config = {} +try { + // config/config.js 可能在某些环境不存在（例如仅拷贝了 config.example.js） + // 为保证可运行，这里做容错处理 + // eslint-disable-next-line global-require + config = require('../../config/config') +} catch (error) { + config = {} +} + +const parseBooleanEnv = (value) => { + if (typeof value === 'boolean') { + return value + } + if (typeof value !== 'string') { + return false + } + const normalized = value.trim().toLowerCase() + return normalized === 'true' || normalized === '1' || normalized === 'yes' || normalized === 'on' +} + +/** + * 是否允许执行“余额脚本”（安全开关） + * 默认开启，便于保持现有行为；如需禁用请显式设置 BALANCE_SCRIPT_ENABLED=false（环境变量优先） + */ +const isBalanceScriptEnabled = () => { + if ( + process.env.BALANCE_SCRIPT_ENABLED !== undefined && + process.env.BALANCE_SCRIPT_ENABLED !== '' + ) { + return parseBooleanEnv(process.env.BALANCE_SCRIPT_ENABLED) + } + + const fromConfig = + config?.accountBalance?.enableBalanceScript ?? + config?.features?.balanceScriptEnabled ?? + config?.security?.enableBalanceScript + + return typeof fromConfig === 'boolean' ? fromConfig : true +} + +module.exports = { + isBalanceScriptEnabled +} diff --git a/tests/accountBalanceService.test.js b/tests/accountBalanceService.test.js new file mode 100644 index 00000000..c2a9c3a8 --- /dev/null +++ b/tests/accountBalanceService.test.js @@ -0,0 +1,218 @@ +// Mock logger，避免测试输出污染控制台 +jest.mock('../src/utils/logger', () => ({ + debug: jest.fn(), + info: jest.fn(), + warn: jest.fn(), + error: jest.fn() +})) + +const accountBalanceServiceModule = require('../src/services/accountBalanceService') + +const { AccountBalanceService } = accountBalanceServiceModule + +describe('AccountBalanceService', () => { + const originalBalanceScriptEnabled = process.env.BALANCE_SCRIPT_ENABLED + + afterEach(() => { + if (originalBalanceScriptEnabled === undefined) { + delete process.env.BALANCE_SCRIPT_ENABLED + } else { + process.env.BALANCE_SCRIPT_ENABLED = originalBalanceScriptEnabled + } + }) + + const mockLogger = { + debug: jest.fn(), + info: jest.fn(), + warn: jest.fn(), + error: jest.fn() + } + + const buildMockRedis = () => ({ + getLocalBalance: jest.fn().mockResolvedValue(null), + setLocalBalance: jest.fn().mockResolvedValue(undefined), + getAccountBalance: jest.fn().mockResolvedValue(null), + setAccountBalance: jest.fn().mockResolvedValue(undefined), + deleteAccountBalance: jest.fn().mockResolvedValue(undefined), + getBalanceScriptConfig: jest.fn().mockResolvedValue(null), + getAccountUsageStats: jest.fn().mockResolvedValue({ + total: { requests: 10 }, + daily: { requests: 2, cost: 20 }, + monthly: { requests: 5 } + }), + getDateInTimezone: (date) => new Date(date.getTime() + 8 * 3600 * 1000) + }) + + it('should normalize platform aliases', () => { + const service = new AccountBalanceService({ redis: buildMockRedis(), logger: mockLogger }) + expect(service.normalizePlatform('claude-official')).toBe('claude') + expect(service.normalizePlatform('azure-openai')).toBe('azure_openai') + expect(service.normalizePlatform('gemini-api')).toBe('gemini-api') + }) + + it('should build local quota/balance from dailyQuota and local dailyCost', async () => { + const mockRedis = buildMockRedis() + const service = new AccountBalanceService({ redis: mockRedis, logger: mockLogger }) + + service._computeMonthlyCost = jest.fn().mockResolvedValue(30) + service._computeTotalCost = jest.fn().mockResolvedValue(123.45) + + const account = { id: 'acct-1', name: 'A', dailyQuota: '100', quotaResetTime: '00:00' } + const result = await service._getAccountBalanceForAccount(account, 'claude-console', { + queryApi: false, + useCache: true + }) + + expect(result.success).toBe(true) + expect(result.data.source).toBe('local') + expect(result.data.balance.amount).toBeCloseTo(80, 6) + expect(result.data.quota.percentage).toBeCloseTo(20, 6) + expect(result.data.statistics.totalCost).toBeCloseTo(123.45, 6) + expect(mockRedis.setLocalBalance).toHaveBeenCalled() + }) + + it('should use cached balance when account has no dailyQuota', async () => { + const mockRedis = buildMockRedis() + mockRedis.getAccountBalance.mockResolvedValue({ + status: 'success', + balance: 12.34, + currency: 'USD', + quota: null, + errorMessage: '', + lastRefreshAt: '2025-01-01T00:00:00Z', + ttlSeconds: 120 + }) + + const service = new AccountBalanceService({ redis: mockRedis, logger: mockLogger }) + service._computeMonthlyCost = jest.fn().mockResolvedValue(0) + service._computeTotalCost = jest.fn().mockResolvedValue(0) + + const account = { id: 'acct-2', name: 'B' } + const result = await service._getAccountBalanceForAccount(account, 'openai', { + queryApi: false, + useCache: true + }) + + expect(result.data.source).toBe('cache') + expect(result.data.balance.amount).toBeCloseTo(12.34, 6) + expect(result.data.lastRefreshAt).toBe('2025-01-01T00:00:00Z') + }) + + it('should not cache provider errors and fallback to local when queryApi=true', async () => { + const mockRedis = buildMockRedis() + const service = new AccountBalanceService({ redis: mockRedis, logger: mockLogger }) + + service._computeMonthlyCost = jest.fn().mockResolvedValue(0) + service._computeTotalCost = jest.fn().mockResolvedValue(0) + + service.registerProvider('openai', { + queryBalance: () => { + throw new Error('boom') + } + }) + + const account = { id: 'acct-3', name: 'C' } + const result = await service._getAccountBalanceForAccount(account, 'openai', { + queryApi: true, + useCache: false + }) + + expect(mockRedis.setAccountBalance).not.toHaveBeenCalled() + expect(result.data.source).toBe('local') + expect(result.data.status).toBe('error') + expect(result.data.error).toBe('boom') + }) + + it('should ignore script config when balance script is disabled', async () => { + process.env.BALANCE_SCRIPT_ENABLED = 'false' + + const mockRedis = buildMockRedis() + mockRedis.getBalanceScriptConfig.mockResolvedValue({ + scriptBody: '({ request: { url: "http://example.com" }, extractor: function(){ return {} } })' + }) + + const service = new AccountBalanceService({ redis: mockRedis, logger: mockLogger }) + service._computeMonthlyCost = jest.fn().mockResolvedValue(0) + service._computeTotalCost = jest.fn().mockResolvedValue(0) + + const provider = { queryBalance: jest.fn().mockResolvedValue({ balance: 1, currency: 'USD' }) } + service.registerProvider('openai', provider) + + const scriptSpy = jest.spyOn(service, '_getBalanceFromScript') + + const account = { id: 'acct-script-off', name: 'S' } + const result = await service._getAccountBalanceForAccount(account, 'openai', { + queryApi: true, + useCache: false + }) + + expect(provider.queryBalance).toHaveBeenCalled() + expect(scriptSpy).not.toHaveBeenCalled() + expect(result.data.source).toBe('api') + }) + + it('should prefer script when configured and enabled', async () => { + process.env.BALANCE_SCRIPT_ENABLED = 'true' + + const mockRedis = buildMockRedis() + mockRedis.getBalanceScriptConfig.mockResolvedValue({ + scriptBody: '({ request: { url: "http://example.com" }, extractor: function(){ return {} } })' + }) + + const service = new AccountBalanceService({ redis: mockRedis, logger: mockLogger }) + service._computeMonthlyCost = jest.fn().mockResolvedValue(0) + service._computeTotalCost = jest.fn().mockResolvedValue(0) + + const provider = { queryBalance: jest.fn().mockResolvedValue({ balance: 2, currency: 'USD' }) } + service.registerProvider('openai', provider) + + jest.spyOn(service, '_getBalanceFromScript').mockResolvedValue({ + status: 'success', + balance: 3, + currency: 'USD', + quota: null, + queryMethod: 'script', + rawData: { ok: true }, + lastRefreshAt: '2025-01-01T00:00:00Z', + errorMessage: '' + }) + + const account = { id: 'acct-script-on', name: 'T' } + const result = await service._getAccountBalanceForAccount(account, 'openai', { + queryApi: true, + useCache: false + }) + + expect(provider.queryBalance).not.toHaveBeenCalled() + expect(result.data.source).toBe('api') + expect(result.data.balance.amount).toBeCloseTo(3, 6) + expect(result.data.lastRefreshAt).toBe('2025-01-01T00:00:00Z') + }) + + it('should count low balance once per account in summary', async () => { + const mockRedis = buildMockRedis() + const service = new AccountBalanceService({ redis: mockRedis, logger: mockLogger }) + + service.getSupportedPlatforms = () => ['claude-console'] + service.getAllAccountsByPlatform = async () => [{ id: 'acct-4', name: 'D' }] + service._getAccountBalanceForAccount = async () => ({ + success: true, + data: { + accountId: 'acct-4', + platform: 'claude-console', + balance: { amount: 5, currency: 'USD', formattedAmount: '$5.00' }, + quota: { percentage: 95 }, + statistics: { totalCost: 1 }, + source: 'local', + lastRefreshAt: '2025-01-01T00:00:00Z', + cacheExpiresAt: null, + status: 'success', + error: null + } + }) + + const summary = await service.getBalanceSummary() + expect(summary.lowBalanceCount).toBe(1) + expect(summary.platforms['claude-console'].lowBalanceCount).toBe(1) + }) +}) diff --git a/web/admin-spa/src/components/accounts/AccountBalanceScriptModal.vue b/web/admin-spa/src/components/accounts/AccountBalanceScriptModal.vue new file mode 100644 index 00000000..17f2be00 --- /dev/null +++ b/web/admin-spa/src/components/accounts/AccountBalanceScriptModal.vue @@ -0,0 +1,302 @@ + + + + + diff --git a/web/admin-spa/src/components/accounts/BalanceDisplay.vue b/web/admin-spa/src/components/accounts/BalanceDisplay.vue new file mode 100644 index 00000000..dc8985b2 --- /dev/null +++ b/web/admin-spa/src/components/accounts/BalanceDisplay.vue @@ -0,0 +1,281 @@ + + + diff --git a/web/admin-spa/src/views/AccountsView.vue b/web/admin-spa/src/views/AccountsView.vue index 65e0bcfe..2f95c589 100644 --- a/web/admin-spa/src/views/AccountsView.vue +++ b/web/admin-spa/src/views/AccountsView.vue @@ -141,6 +141,28 @@ + +
+ + + +
+ + +
@@ -1425,6 +1469,26 @@
+ +
+

余额/配额

+ +
+ +
+
+
@@ -1906,6 +1970,13 @@ @saved="handleScheduledTestSaved" /> + + { showToast('定时测试配置已保存', 'success') } +// 余额脚本配置 +const showBalanceScriptModal = ref(false) +const selectedAccountForScript = ref(null) + +const openBalanceScriptModal = (account) => { + selectedAccountForScript.value = account + showBalanceScriptModal.value = true +} + +const closeBalanceScriptModal = () => { + showBalanceScriptModal.value = false + selectedAccountForScript.value = null +} + +const handleBalanceScriptSaved = async () => { + showToast('余额脚本已保存', 'success') + const account = selectedAccountForScript.value + closeBalanceScriptModal() + + if (!account?.id || !account?.platform) { + return + } + + // 重新拉取一次余额信息，用于刷新 scriptConfigured 状态（启用“刷新余额”按钮） + try { + const res = await apiClient.get(`/admin/accounts/${account.id}/balance`, { + params: { platform: account.platform, queryApi: false } + }) + if (res?.success && res.data) { + handleBalanceRefreshed(account.id, res.data) + } + } catch (error) { + console.debug('Failed to reload balance after saving script:', error) + } +} + // 计算排序后的账户列表 const sortedAccounts = computed(() => { let sourceAccounts = accounts.value @@ -2768,6 +2878,104 @@ const paginatedAccounts = computed(() => { return sortedAccounts.value.slice(start, end) }) +const canRefreshVisibleBalances = computed(() => { + const targets = paginatedAccounts.value + if (!Array.isArray(targets) || targets.length === 0) { + return false + } + + return targets.some((account) => { + const info = account?.balanceInfo + return info?.scriptEnabled !== false && !!info?.scriptConfigured + }) +}) + +const refreshBalanceTooltip = computed(() => { + if (accountsLoading.value) return '正在加载账户...' + if (refreshingBalances.value) return '刷新中...' + if (!canRefreshVisibleBalances.value) return '当前页未配置余额脚本，无法刷新' + return '刷新当前页余额（仅对已配置余额脚本的账户生效）' +}) + +// 余额刷新成功回调 +const handleBalanceRefreshed = (accountId, balanceInfo) => { + accounts.value = accounts.value.map((account) => { + if (account.id !== accountId) return account + return { ...account, balanceInfo } + }) +} + +// 余额请求错误回调（仅提示，不中断页面） +const handleBalanceError = (_accountId, error) => { + const message = error?.message || '余额查询失败' + showToast(message, 'error') +} + +// 批量刷新当前页余额（触发查询） +const refreshVisibleBalances = async () => { + if (refreshingBalances.value) return + + const targets = paginatedAccounts.value + if (!targets || targets.length === 0) { + return + } + + const eligibleTargets = targets.filter((account) => { + const info = account?.balanceInfo + return info?.scriptEnabled !== false && !!info?.scriptConfigured + }) + + if (eligibleTargets.length === 0) { + showToast('当前页没有配置余额脚本的账户', 'warning') + return + } + + const skippedCount = targets.length - eligibleTargets.length + + refreshingBalances.value = true + try { + const results = await Promise.all( + eligibleTargets.map(async (account) => { + try { + const response = await apiClient.post(`/admin/accounts/${account.id}/balance/refresh`, { + platform: account.platform + }) + return { id: account.id, success: !!response?.success, data: response?.data || null } + } catch (error) { + return { id: account.id, success: false, error: error?.message || '刷新失败' } + } + }) + ) + + const updatedMap = results.reduce((map, item) => { + if (item.success && item.data) { + map[item.id] = item.data + } + return map + }, {}) + + const successCount = results.filter((r) => r.success).length + const failCount = results.length - successCount + + const skippedText = skippedCount > 0 ? `，跳过 ${skippedCount} 个未配置脚本` : '' + if (Object.keys(updatedMap).length > 0) { + accounts.value = accounts.value.map((account) => { + const balanceInfo = updatedMap[account.id] + if (!balanceInfo) return account + return { ...account, balanceInfo } + }) + } + + if (failCount === 0) { + showToast(`成功刷新 ${successCount} 个账户余额${skippedText}`, 'success') + } else { + showToast(`刷新完成：${successCount} 成功，${failCount} 失败${skippedText}`, 'warning') + } + } finally { + refreshingBalances.value = false + } +} + const updateSelectAllState = () => { const currentIds = paginatedAccounts.value.map((account) => account.id) const selectedInCurrentPage = currentIds.filter((id) => @@ -2818,6 +3026,54 @@ const cleanupSelectedAccounts = () => { updateSelectAllState() } +// 异步加载余额缓存（按平台批量拉取，避免逐行请求） +const loadBalanceCacheForAccounts = async () => { + const current = accounts.value + if (!Array.isArray(current) || current.length === 0) { + return + } + + const platforms = Array.from(new Set(current.map((acc) => acc.platform).filter(Boolean))) + if (platforms.length === 0) { + return + } + + const responses = await Promise.all( + platforms.map(async (platform) => { + try { + const res = await apiClient.get(`/admin/accounts/balance/platform/${platform}`, { + params: { queryApi: false } + }) + return { platform, success: !!res?.success, data: res?.data || [] } + } catch (error) { + console.debug(`Failed to load balance cache for ${platform}:`, error) + return { platform, success: false, data: [] } + } + }) + ) + + const balanceMap = responses.reduce((map, item) => { + if (!item.success) return map + const list = Array.isArray(item.data) ? item.data : [] + list.forEach((entry) => { + const accountId = entry?.data?.accountId + if (accountId) { + map[accountId] = entry.data + } + }) + return map + }, {}) + + if (Object.keys(balanceMap).length === 0) { + return + } + + accounts.value = accounts.value.map((account) => ({ + ...account, + balanceInfo: balanceMap[account.id] || account.balanceInfo || null + })) +} + // 加载账户列表 const loadAccounts = async (forceReload = false) => { accountsLoading.value = true @@ -3010,6 +3266,11 @@ const loadAccounts = async (forceReload = false) => { console.debug('Claude usage loading failed:', err) }) } + + // 异步加载余额缓存（按平台批量） + loadBalanceCacheForAccounts().catch((err) => { + console.debug('Balance cache loading failed:', err) + }) } catch (error) { showToast('加载账户失败', 'error') } finally { diff --git a/web/admin-spa/src/views/BalanceScriptsView.vue b/web/admin-spa/src/views/BalanceScriptsView.vue new file mode 100644 index 00000000..1e4334da --- /dev/null +++ b/web/admin-spa/src/views/BalanceScriptsView.vue @@ -0,0 +1,312 @@ + + + + + diff --git a/web/admin-spa/src/views/DashboardView.vue b/web/admin-spa/src/views/DashboardView.vue index 84d60d27..0e04eee3 100644 --- a/web/admin-spa/src/views/DashboardView.vue +++ b/web/admin-spa/src/views/DashboardView.vue @@ -196,6 +196,105 @@
+ +
+
+
+
+

+ 账户余额/配额 +

+

+ {{ formatCurrencyUsd(balanceSummary.totalBalance || 0) }} +

+

+ 低余额: {{ balanceSummary.lowBalanceCount || 0 }} | 总成本: + {{ formatCurrencyUsd(balanceSummary.totalCost || 0) }} +

+
+
+ +
+
+ +
+

+ 更新时间: {{ formatLastUpdate(balanceSummaryUpdatedAt) }} +

+ +
+
+ +
+
+

低余额账户

+ + {{ lowBalanceAccounts.length }} 个 + +
+ +
+ 正在加载... +
+
+ 全部正常 +
+
+
+
+
+ {{ account.name || account.accountId }} +
+ + {{ getBalancePlatformLabel(account.platform) }} + +
+
+ 余额: {{ account.balance.formattedAmount }} + 今日成本: {{ formatCurrencyUsd(account.statistics?.dailyCost || 0) }} +
+
+
+ 配额使用 + + {{ account.quota.percentage.toFixed(1) }}% + +
+
+
+
+
+
+
+
+
+
{ + const map = { + claude: 'Claude', + 'claude-console': 'Claude Console', + gemini: 'Gemini', + 'gemini-api': 'Gemini API', + openai: 'OpenAI', + 'openai-responses': 'OpenAI Responses', + azure_openai: 'Azure OpenAI', + bedrock: 'Bedrock', + droid: 'Droid', + ccr: 'CCR' + } + return map[platform] || platform +} + +const lowBalanceAccounts = computed(() => { + const result = [] + const platforms = balanceSummary.value?.platforms || {} + + Object.entries(platforms).forEach(([platform, data]) => { + const list = Array.isArray(data?.accounts) ? data.accounts : [] + list.forEach((entry) => { + const accountData = entry?.data + if (!accountData) return + + const amount = accountData.balance?.amount + const percentage = accountData.quota?.percentage + + const isLowBalance = typeof amount === 'number' && amount < 10 + const isHighUsage = typeof percentage === 'number' && percentage > 90 + + if (isLowBalance || isHighUsage) { + result.push({ + ...accountData, + name: entry?.name || accountData.accountId, + platform: accountData.platform || platform + }) + } + }) + }) + + return result +}) + +const formatCurrencyUsd = (amount) => { + const value = Number(amount) + if (!Number.isFinite(value)) return '$0.00' + if (value >= 1) return `$${value.toFixed(2)}` + if (value >= 0.01) return `$${value.toFixed(3)}` + return `$${value.toFixed(6)}` +} + +const formatLastUpdate = (isoString) => { + if (!isoString) return '未知' + const date = new Date(isoString) + if (Number.isNaN(date.getTime())) return '未知' + return date.toLocaleTimeString('zh-CN', { hour: '2-digit', minute: '2-digit' }) +} + +const loadBalanceSummary = async () => { + loadingBalanceSummary.value = true + try { + const response = await apiClient.get('/admin/accounts/balance/summary') + if (response?.success) { + balanceSummary.value = response.data || { + totalBalance: 0, + totalCost: 0, + lowBalanceCount: 0, + platforms: {} + } + balanceSummaryUpdatedAt.value = new Date().toISOString() + } + } catch (error) { + console.debug('加载余额汇总失败:', error) + showToast('加载余额汇总失败', 'error') + } finally { + loadingBalanceSummary.value = false + } +} + // 自动刷新相关 const autoRefreshEnabled = ref(false) const autoRefreshInterval = ref(30) // 秒 @@ -1488,7 +1680,7 @@ async function refreshAllData() { isRefreshing.value = true try { - await Promise.all([loadDashboardData(), refreshChartsData()]) + await Promise.all([loadDashboardData(), refreshChartsData(), loadBalanceSummary()]) } finally { isRefreshing.value = false }