diff --git a/src/routes/geminiRoutes.js b/src/routes/geminiRoutes.js index ba10dd01..0b1c737b 100644 --- a/src/routes/geminiRoutes.js +++ b/src/routes/geminiRoutes.js @@ -7,6 +7,7 @@ const { sendGeminiRequest, getAvailableModels } = require('../services/geminiRel const crypto = require('crypto'); const sessionHelper = require('../utils/sessionHelper'); const unifiedGeminiScheduler = require('../services/unifiedGeminiScheduler'); +const apiKeyService = require('../services/apiKeyService'); // const { OAuth2Client } = require('google-auth-library'); // OAuth2Client is not used in this file // 生成会话哈希 @@ -195,7 +196,13 @@ router.get('/models', authenticateApiKey, async (req, res) => { } // 选择账户获取模型列表 - const account = await geminiAccountService.selectAvailableAccount(apiKeyData.id); + let account = null; + try { + const accountSelection = await unifiedGeminiScheduler.selectAccountForApiKey(apiKeyData, null, null); + account = await geminiAccountService.getAccount(accountSelection.accountId); + } catch (error) { + logger.warn('Failed to select Gemini account for models endpoint:', error); + } if (!account) { // 返回默认模型列表 @@ -470,6 +477,25 @@ async function handleGenerateContent(req, res) { req.apiKey?.id // 使用 API Key ID 作为 session ID ); + // 记录使用统计 + if (response?.response?.usageMetadata) { + try { + const usage = response.response.usageMetadata; + await apiKeyService.recordUsage( + req.apiKey.id, + usage.promptTokenCount || 0, + usage.candidatesTokenCount || 0, + 0, // cacheCreateTokens + 0, // cacheReadTokens + model, + account.id + ); + logger.info(`📊 Recorded Gemini usage - Input: ${usage.promptTokenCount}, Output: ${usage.candidatesTokenCount}, Total: ${usage.totalTokenCount}`); + } catch (error) { + logger.error('Failed to record Gemini usage:', error); + } + } + res.json(response); } catch (error) { console.log(321, error.response); @@ -565,11 +591,73 @@ async function handleStreamGenerateContent(req, res) { res.setHeader('Connection', 'keep-alive'); res.setHeader('X-Accel-Buffering', 'no'); - // 直接管道转发流式响应，不进行额外处理 - streamResponse.pipe(res, { end: false }); + // 处理流式响应并捕获usage数据 + let buffer = ''; + let totalUsage = { + promptTokenCount: 0, + candidatesTokenCount: 0, + totalTokenCount: 0 + }; + let usageReported = false; - streamResponse.on('end', () => { + streamResponse.on('data', (chunk) => { + try { + const chunkStr = chunk.toString(); + + // 直接转发数据到客户端 + if (!res.destroyed) { + res.write(chunkStr); + } + + // 同时解析数据以捕获usage信息 + buffer += chunkStr; + const lines = buffer.split('\n'); + buffer = lines.pop() || ''; + + for (const line of lines) { + if (line.startsWith('data: ') && line.length > 6) { + try { + const jsonStr = line.slice(6); + if (jsonStr && jsonStr !== '[DONE]') { + const data = JSON.parse(jsonStr); + + // 从响应中提取usage数据 + if (data.response?.usageMetadata) { + totalUsage = data.response.usageMetadata; + logger.debug('📊 Captured Gemini usage data:', totalUsage); + } + } + } catch (e) { + // 忽略解析错误 + } + } + } + } catch (error) { + logger.error('Error processing stream chunk:', error); + } + }); + + streamResponse.on('end', async () => { logger.info('Stream completed successfully'); + + // 记录使用统计 + if (!usageReported && totalUsage.totalTokenCount > 0) { + try { + await apiKeyService.recordUsage( + req.apiKey.id, + totalUsage.promptTokenCount || 0, + totalUsage.candidatesTokenCount || 0, + 0, // cacheCreateTokens + 0, // cacheReadTokens + model, + account.id + ); + logger.info(`📊 Recorded Gemini stream usage - Input: ${totalUsage.promptTokenCount}, Output: ${totalUsage.candidatesTokenCount}, Total: ${totalUsage.totalTokenCount}`); + } catch (error) { + logger.error('Failed to record Gemini usage:', error); + } + } + res.end(); }); diff --git a/src/routes/openaiClaudeRoutes.js b/src/routes/openaiClaudeRoutes.js index 247fc127..6b36f1fb 100644 --- a/src/routes/openaiClaudeRoutes.js +++ b/src/routes/openaiClaudeRoutes.js @@ -12,7 +12,7 @@ const { authenticateApiKey } = require('../middleware/auth'); const claudeRelayService = require('../services/claudeRelayService'); const openaiToClaude = require('../services/openaiToClaude'); const apiKeyService = require('../services/apiKeyService'); -const claudeAccountService = require('../services/claudeAccountService'); +const unifiedClaudeScheduler = require('../services/unifiedClaudeScheduler'); const claudeCodeHeadersService = require('../services/claudeCodeHeadersService'); const sessionHelper = require('../utils/sessionHelper'); @@ -206,7 +206,8 @@ async function handleChatCompletion(req, res, apiKeyData) { const sessionHash = sessionHelper.generateSessionHash(claudeRequest); // 选择可用的Claude账户 - const accountId = await claudeAccountService.selectAccountForApiKey(apiKeyData, sessionHash); + const accountSelection = await unifiedClaudeScheduler.selectAccountForApiKey(apiKeyData, sessionHash, claudeRequest.model); + const accountId = accountSelection.accountId; // 获取该账号存储的 Claude Code headers const claudeCodeHeaders = await claudeCodeHeadersService.getAccountHeaders(accountId); diff --git a/src/routes/openaiGeminiRoutes.js b/src/routes/openaiGeminiRoutes.js index 8b1af00b..9a9714e2 100644 --- a/src/routes/openaiGeminiRoutes.js +++ b/src/routes/openaiGeminiRoutes.js @@ -3,6 +3,7 @@ const router = express.Router(); const logger = require('../utils/logger'); const { authenticateApiKey } = require('../middleware/auth'); const geminiAccountService = require('../services/geminiAccountService'); +const unifiedGeminiScheduler = require('../services/unifiedGeminiScheduler'); const { getAvailableModels } = require('../services/geminiRelayService'); const crypto = require('crypto'); @@ -167,6 +168,8 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => { const startTime = Date.now(); let abortController = null; let account = null; // Declare account outside try block for error handling + let accountSelection = null; // Declare accountSelection for error handling + let sessionHash = null; // Declare sessionHash for error handling try { const apiKeyData = req.apiKey; @@ -263,13 +266,16 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => { } // 生成会话哈希用于粘性会话 - const sessionHash = generateSessionHash(req); + sessionHash = generateSessionHash(req); // 选择可用的 Gemini 账户 - account = await geminiAccountService.selectAvailableAccount( - apiKeyData.id, - sessionHash - ); + try { + accountSelection = await unifiedGeminiScheduler.selectAccountForApiKey(apiKeyData, sessionHash, model); + account = await geminiAccountService.getAccount(accountSelection.accountId); + } catch (error) { + logger.error('Failed to select Gemini account:', error); + account = null; + } if (!account) { return res.status(503).json({ @@ -339,6 +345,14 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => { }; res.write(`data: ${JSON.stringify(initialChunk)}\n\n`); + // 用于收集usage数据 + let totalUsage = { + promptTokenCount: 0, + candidatesTokenCount: 0, + totalTokenCount: 0 + }; + let usageReported = false; + streamResponse.on('data', (chunk) => { try { const chunkStr = chunk.toString(); @@ -365,6 +379,12 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => { try { const data = JSON.parse(jsonData); + // 捕获usage数据 + if (data.response?.usageMetadata) { + totalUsage = data.response.usageMetadata; + logger.debug('📊 Captured Gemini usage data:', totalUsage); + } + // 转换为 OpenAI 流式格式 if (data.response?.candidates && data.response.candidates.length > 0) { const candidate = data.response.candidates[0]; @@ -430,8 +450,28 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => { } }); - streamResponse.on('end', () => { + streamResponse.on('end', async () => { logger.info('Stream completed successfully'); + + // 记录使用统计 + if (!usageReported && totalUsage.totalTokenCount > 0) { + try { + const apiKeyService = require('../services/apiKeyService'); + await apiKeyService.recordUsage( + apiKeyData.id, + totalUsage.promptTokenCount || 0, + totalUsage.candidatesTokenCount || 0, + 0, // cacheCreateTokens + 0, // cacheReadTokens + model, + account.id + ); + logger.info(`📊 Recorded Gemini stream usage - Input: ${totalUsage.promptTokenCount}, Output: ${totalUsage.candidatesTokenCount}, Total: ${totalUsage.totalTokenCount}`); + } catch (error) { + logger.error('Failed to record Gemini usage:', error); + } + } + if (!res.headersSent) { res.write('data: [DONE]\n\n'); } @@ -473,6 +513,26 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => { // 转换为 OpenAI 格式并返回 const openaiResponse = convertGeminiResponseToOpenAI(response, model, false); + + // 记录使用统计 + if (openaiResponse.usage) { + try { + const apiKeyService = require('../services/apiKeyService'); + await apiKeyService.recordUsage( + apiKeyData.id, + openaiResponse.usage.prompt_tokens || 0, + openaiResponse.usage.completion_tokens || 0, + 0, // cacheCreateTokens + 0, // cacheReadTokens + model, + account.id + ); + logger.info(`📊 Recorded Gemini usage - Input: ${openaiResponse.usage.prompt_tokens}, Output: ${openaiResponse.usage.completion_tokens}, Total: ${openaiResponse.usage.total_tokens}`); + } catch (error) { + logger.error('Failed to record Gemini usage:', error); + } + } + res.json(openaiResponse); } @@ -484,8 +544,8 @@ router.post('/v1/chat/completions', authenticateApiKey, async (req, res) => { // 处理速率限制 if (error.status === 429) { - if (req.apiKey && account) { - await geminiAccountService.setAccountRateLimited(account.id, true); + if (req.apiKey && account && accountSelection) { + await unifiedGeminiScheduler.markAccountRateLimited(account.id, 'gemini', sessionHash); } } @@ -525,7 +585,13 @@ router.get('/v1/models', authenticateApiKey, async (req, res) => { } // 选择账户获取模型列表 - const account = await geminiAccountService.selectAvailableAccount(apiKeyData.id); + let account = null; + try { + const accountSelection = await unifiedGeminiScheduler.selectAccountForApiKey(apiKeyData, null, null); + account = await geminiAccountService.getAccount(accountSelection.accountId); + } catch (error) { + logger.warn('Failed to select Gemini account for models endpoint:', error); + } let models = []; diff --git a/src/services/claudeRelayService.js b/src/services/claudeRelayService.js index 0a8ac6b9..09ddd013 100644 --- a/src/services/claudeRelayService.js +++ b/src/services/claudeRelayService.js @@ -5,6 +5,7 @@ const path = require('path'); const { SocksProxyAgent } = require('socks-proxy-agent'); const { HttpsProxyAgent } = require('https-proxy-agent'); const claudeAccountService = require('./claudeAccountService'); +const unifiedClaudeScheduler = require('./unifiedClaudeScheduler'); const sessionHelper = require('../utils/sessionHelper'); const logger = require('../utils/logger'); const config = require('../../config/config'); @@ -91,9 +92,11 @@ class ClaudeRelayService { const sessionHash = sessionHelper.generateSessionHash(requestBody); // 选择可用的Claude账户（支持专属绑定和sticky会话） - const accountId = await claudeAccountService.selectAccountForApiKey(apiKeyData, sessionHash); + const accountSelection = await unifiedClaudeScheduler.selectAccountForApiKey(apiKeyData, sessionHash, requestBody.model); + const accountId = accountSelection.accountId; + const accountType = accountSelection.accountType; - logger.info(`📤 Processing API request for key: ${apiKeyData.name || apiKeyData.id}, account: ${accountId}${sessionHash ? `, session: ${sessionHash}` : ''}`); + logger.info(`📤 Processing API request for key: ${apiKeyData.name || apiKeyData.id}, account: ${accountId} (${accountType})${sessionHash ? `, session: ${sessionHash}` : ''}`); // 获取有效的访问token const accessToken = await claudeAccountService.getValidAccessToken(accountId); @@ -172,13 +175,13 @@ class ClaudeRelayService { if (isRateLimited) { logger.warn(`🚫 Rate limit detected for account ${accountId}, status: ${response.statusCode}`); // 标记账号为限流状态并删除粘性会话映射，传递准确的重置时间戳 - await claudeAccountService.markAccountRateLimited(accountId, sessionHash, rateLimitResetTimestamp); + await unifiedClaudeScheduler.markAccountRateLimited(accountId, accountType, sessionHash, rateLimitResetTimestamp); } } else if (response.statusCode === 200 || response.statusCode === 201) { // 如果请求成功，检查并移除限流状态 - const isRateLimited = await claudeAccountService.isAccountRateLimited(accountId); + const isRateLimited = await unifiedClaudeScheduler.isAccountRateLimited(accountId, accountType); if (isRateLimited) { - await claudeAccountService.removeAccountRateLimit(accountId); + await unifiedClaudeScheduler.removeAccountRateLimit(accountId, accountType); } // 只有真实的 Claude Code 请求才更新 headers @@ -621,9 +624,11 @@ class ClaudeRelayService { const sessionHash = sessionHelper.generateSessionHash(requestBody); // 选择可用的Claude账户（支持专属绑定和sticky会话） - const accountId = await claudeAccountService.selectAccountForApiKey(apiKeyData, sessionHash); + const accountSelection = await unifiedClaudeScheduler.selectAccountForApiKey(apiKeyData, sessionHash, requestBody.model); + const accountId = accountSelection.accountId; + const accountType = accountSelection.accountType; - logger.info(`📡 Processing streaming API request with usage capture for key: ${apiKeyData.name || apiKeyData.id}, account: ${accountId}${sessionHash ? `, session: ${sessionHash}` : ''}`); + logger.info(`📡 Processing streaming API request with usage capture for key: ${apiKeyData.name || apiKeyData.id}, account: ${accountId} (${accountType})${sessionHash ? `, session: ${sessionHash}` : ''}`); // 获取有效的访问token const accessToken = await claudeAccountService.getValidAccessToken(accountId); @@ -638,7 +643,7 @@ class ClaudeRelayService { return await this._makeClaudeStreamRequestWithUsageCapture(processedBody, accessToken, proxyAgent, clientHeaders, responseStream, (usageData) => { // 在usageCallback中添加accountId usageCallback({ ...usageData, accountId }); - }, accountId, sessionHash, streamTransformer, options); + }, accountId, accountType, sessionHash, streamTransformer, options); } catch (error) { logger.error('❌ Claude stream relay with usage capture failed:', error); throw error; @@ -646,7 +651,7 @@ class ClaudeRelayService { } // 🌊 发送流式请求到Claude API（带usage数据捕获） - async _makeClaudeStreamRequestWithUsageCapture(body, accessToken, proxyAgent, clientHeaders, responseStream, usageCallback, accountId, sessionHash, streamTransformer = null, requestOptions = {}) { + async _makeClaudeStreamRequestWithUsageCapture(body, accessToken, proxyAgent, clientHeaders, responseStream, usageCallback, accountId, accountType, sessionHash, streamTransformer = null, requestOptions = {}) { // 获取过滤后的客户端 headers const filteredHeaders = this._filterClientHeaders(clientHeaders); @@ -854,12 +859,12 @@ class ClaudeRelayService { } // 标记账号为限流状态并删除粘性会话映射 - await claudeAccountService.markAccountRateLimited(accountId, sessionHash, rateLimitResetTimestamp); + await unifiedClaudeScheduler.markAccountRateLimited(accountId, accountType, sessionHash, rateLimitResetTimestamp); } else if (res.statusCode === 200) { // 如果请求成功，检查并移除限流状态 - const isRateLimited = await claudeAccountService.isAccountRateLimited(accountId); + const isRateLimited = await unifiedClaudeScheduler.isAccountRateLimited(accountId, accountType); if (isRateLimited) { - await claudeAccountService.removeAccountRateLimit(accountId); + await unifiedClaudeScheduler.removeAccountRateLimit(accountId, accountType); } // 只有真实的 Claude Code 请求才更新 headers（流式请求） diff --git a/src/services/unifiedClaudeScheduler.js b/src/services/unifiedClaudeScheduler.js index c2b07b7e..64c10eab 100644 --- a/src/services/unifiedClaudeScheduler.js +++ b/src/services/unifiedClaudeScheduler.js @@ -9,6 +9,16 @@ class UnifiedClaudeScheduler { this.SESSION_MAPPING_PREFIX = 'unified_claude_session_mapping:'; } + // 🔧 辅助方法：检查账户是否可调度（兼容字符串和布尔值） + _isSchedulable(schedulable) { + // 如果是 undefined 或 null，默认为可调度 + if (schedulable === undefined || schedulable === null) { + return true; + } + // 明确设置为 false（布尔值）或 'false'（字符串）时不可调度 + return schedulable !== false && schedulable !== 'false'; + } + // 🎯 统一调度Claude账号（官方和Console） async selectAccountForApiKey(apiKeyData, sessionHash = null, requestedModel = null) { try { @@ -152,7 +162,7 @@ class UnifiedClaudeScheduler { account.status !== 'error' && account.status !== 'blocked' && (account.accountType === 'shared' || !account.accountType) && // 兼容旧数据 - account.schedulable !== 'false') { // 检查是否可调度 + this._isSchedulable(account.schedulable)) { // 检查是否可调度 // 检查是否被限流 const isRateLimited = await claudeAccountService.isAccountRateLimited(account.id); @@ -179,7 +189,7 @@ class UnifiedClaudeScheduler { if (account.isActive === true && account.status === 'active' && account.accountType === 'shared' && - account.schedulable !== false) { // 检查是否可调度 + this._isSchedulable(account.schedulable)) { // 检查是否可调度 // 检查模型支持（如果有请求的模型） if (requestedModel && account.supportedModels) { @@ -246,7 +256,7 @@ class UnifiedClaudeScheduler { return false; } // 检查是否可调度 - if (account.schedulable === 'false') { + if (!this._isSchedulable(account.schedulable)) { logger.info(`🚫 Account ${accountId} is not schedulable`); return false; } @@ -257,7 +267,7 @@ class UnifiedClaudeScheduler { return false; } // 检查是否可调度 - if (account.schedulable === false) { + if (!this._isSchedulable(account.schedulable)) { logger.info(`🚫 Claude Console account ${accountId} is not schedulable`); return false; } @@ -444,7 +454,7 @@ class UnifiedClaudeScheduler { ? account.status !== 'error' && account.status !== 'blocked' : account.status === 'active'; - if (isActive && status && account.schedulable !== false) { + if (isActive && status && this._isSchedulable(account.schedulable)) { // 检查模型支持（Console账户） if (accountType === 'claude-console' && requestedModel && account.supportedModels && account.supportedModels.length > 0) { if (!account.supportedModels.includes(requestedModel)) { diff --git a/src/services/unifiedGeminiScheduler.js b/src/services/unifiedGeminiScheduler.js index f4f056b6..92dc4868 100644 --- a/src/services/unifiedGeminiScheduler.js +++ b/src/services/unifiedGeminiScheduler.js @@ -8,6 +8,16 @@ class UnifiedGeminiScheduler { this.SESSION_MAPPING_PREFIX = 'unified_gemini_session_mapping:'; } + // 🔧 辅助方法：检查账户是否可调度（兼容字符串和布尔值） + _isSchedulable(schedulable) { + // 如果是 undefined 或 null，默认为可调度 + if (schedulable === undefined || schedulable === null) { + return true; + } + // 明确设置为 false（布尔值）或 'false'（字符串）时不可调度 + return schedulable !== false && schedulable !== 'false'; + } + // 🎯 统一调度Gemini账号 async selectAccountForApiKey(apiKeyData, sessionHash = null, requestedModel = null) { try { @@ -128,7 +138,7 @@ class UnifiedGeminiScheduler { if (account.isActive === 'true' && account.status !== 'error' && (account.accountType === 'shared' || !account.accountType) && // 兼容旧数据 - account.schedulable !== 'false') { // 检查是否可调度 + this._isSchedulable(account.schedulable)) { // 检查是否可调度 // 检查token是否过期 const isExpired = geminiAccountService.isTokenExpired(account); @@ -192,7 +202,7 @@ class UnifiedGeminiScheduler { return false; } // 检查是否可调度 - if (account.schedulable === 'false') { + if (!this._isSchedulable(account.schedulable)) { logger.info(`🚫 Gemini account ${accountId} is not schedulable`); return false; } @@ -347,7 +357,7 @@ class UnifiedGeminiScheduler { // 检查账户是否可用 if (account.isActive === 'true' && account.status !== 'error' && - account.schedulable !== 'false') { + this._isSchedulable(account.schedulable)) { // 检查token是否过期 const isExpired = geminiAccountService.isTokenExpired(account); diff --git a/web/admin-spa/src/components/apikeys/UsageDetailModal.vue b/web/admin-spa/src/components/apikeys/UsageDetailModal.vue index f9c4b64d..44037bb2 100644 --- a/web/admin-spa/src/components/apikeys/UsageDetailModal.vue +++ b/web/admin-spa/src/components/apikeys/UsageDetailModal.vue @@ -1,28 +1,37 @@ \ No newline at end of file