Merge remote-tracking branch 'origin/main' into dev

2026-01-23 09:38:02 +00:00 · 2025-07-28 09:30:33 +08:00
parent 2a6bb2ab01 7c91e9424f
commit 4913be3b21
45 changed files with 9066 additions and 553 deletions
--- a/src/services/apiKeyService.js
+++ b/src/services/apiKeyService.js
@@ -24,7 +24,10 @@ class ApiKeyService {
      rateLimitWindow = null,
      rateLimitRequests = null,
      enableModelRestriction = false,
-      restrictedModels = []
+      restrictedModels = [],
+      enableClientRestriction = false,
+      allowedClients = [],
+      dailyCostLimit = 0
    } = options;

    // 生成简单的API Key (64字符十六进制)
@@ -47,6 +50,9 @@ class ApiKeyService {
      permissions: permissions || 'all',
      enableModelRestriction: String(enableModelRestriction),
      restrictedModels: JSON.stringify(restrictedModels || []),
+      enableClientRestriction: String(enableClientRestriction || false),
+      allowedClients: JSON.stringify(allowedClients || []),
+      dailyCostLimit: String(dailyCostLimit || 0),
      createdAt: new Date().toISOString(),
      lastUsedAt: '',
      expiresAt: expiresAt || '',
@@ -73,6 +79,9 @@ class ApiKeyService {
      permissions: keyData.permissions,
      enableModelRestriction: keyData.enableModelRestriction === 'true',
      restrictedModels: JSON.parse(keyData.restrictedModels),
+      enableClientRestriction: keyData.enableClientRestriction === 'true',
+      allowedClients: JSON.parse(keyData.allowedClients || '[]'),
+      dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
      createdAt: keyData.createdAt,
      expiresAt: keyData.expiresAt,
      createdBy: keyData.createdBy
@@ -108,6 +117,9 @@ class ApiKeyService {

      // 获取使用统计（供返回数据使用）
      const usage = await redis.getUsageStats(keyData.id);
+      
+      // 获取当日费用统计
+      const dailyCost = await redis.getDailyCost(keyData.id);

      // 更新最后使用时间（优化：只在实际API调用时更新，而不是验证时）
      // 注意：lastUsedAt的更新已移至recordUsage方法中
@@ -122,11 +134,22 @@ class ApiKeyService {
        restrictedModels = [];
      }

+      // 解析允许的客户端
+      let allowedClients = [];
+      try {
+        allowedClients = keyData.allowedClients ? JSON.parse(keyData.allowedClients) : [];
+      } catch (e) {
+        allowedClients = [];
+      }
+
      return {
        valid: true,
        keyData: {
          id: keyData.id,
          name: keyData.name,
+          description: keyData.description,
+          createdAt: keyData.createdAt,
+          expiresAt: keyData.expiresAt,
          claudeAccountId: keyData.claudeAccountId,
          geminiAccountId: keyData.geminiAccountId,
          permissions: keyData.permissions || 'all',
@@ -136,6 +159,10 @@ class ApiKeyService {
          rateLimitRequests: parseInt(keyData.rateLimitRequests || 0),
          enableModelRestriction: keyData.enableModelRestriction === 'true',
          restrictedModels: restrictedModels,
+          enableClientRestriction: keyData.enableClientRestriction === 'true',
+          allowedClients: allowedClients,
+          dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
+          dailyCost: dailyCost || 0,
          usage
        }
      };
@@ -160,12 +187,20 @@ class ApiKeyService {
        key.currentConcurrency = await redis.getConcurrency(key.id);
        key.isActive = key.isActive === 'true';
        key.enableModelRestriction = key.enableModelRestriction === 'true';
+        key.enableClientRestriction = key.enableClientRestriction === 'true';
        key.permissions = key.permissions || 'all'; // 兼容旧数据
+        key.dailyCostLimit = parseFloat(key.dailyCostLimit || 0);
+        key.dailyCost = await redis.getDailyCost(key.id) || 0;
        try {
          key.restrictedModels = key.restrictedModels ? JSON.parse(key.restrictedModels) : [];
        } catch (e) {
          key.restrictedModels = [];
        }
+        try {
+          key.allowedClients = key.allowedClients ? JSON.parse(key.allowedClients) : [];
+        } catch (e) {
+          key.allowedClients = [];
+        }
        delete key.apiKey; // 不返回哈希后的key
      }

@@ -185,15 +220,15 @@ class ApiKeyService {
      }

      // 允许更新的字段
-      const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'geminiAccountId', 'permissions', 'expiresAt', 'enableModelRestriction', 'restrictedModels'];
+      const allowedUpdates = ['name', 'description', 'tokenLimit', 'concurrencyLimit', 'rateLimitWindow', 'rateLimitRequests', 'isActive', 'claudeAccountId', 'geminiAccountId', 'permissions', 'expiresAt', 'enableModelRestriction', 'restrictedModels', 'enableClientRestriction', 'allowedClients', 'dailyCostLimit'];
      const updatedData = { ...keyData };

      for (const [field, value] of Object.entries(updates)) {
        if (allowedUpdates.includes(field)) {
-          if (field === 'restrictedModels') {
-            // 特殊处理 restrictedModels 数组
+          if (field === 'restrictedModels' || field === 'allowedClients') {
+            // 特殊处理数组字段
            updatedData[field] = JSON.stringify(value || []);
-          } else if (field === 'enableModelRestriction') {
+          } else if (field === 'enableModelRestriction' || field === 'enableClientRestriction') {
            // 布尔值转字符串
            updatedData[field] = String(value);
          } else {
@@ -234,18 +269,45 @@ class ApiKeyService {
    }
  }

-  // 📊 记录使用情况（支持缓存token）
-  async recordUsage(keyId, inputTokens = 0, outputTokens = 0, cacheCreateTokens = 0, cacheReadTokens = 0, model = 'unknown') {
+  // 📊 记录使用情况（支持缓存token和账户级别统计）
+  async recordUsage(keyId, inputTokens = 0, outputTokens = 0, cacheCreateTokens = 0, cacheReadTokens = 0, model = 'unknown', accountId = null) {
    try {
      const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens;
+      
+      // 计算费用
+      const CostCalculator = require('../utils/costCalculator');
+      const costInfo = CostCalculator.calculateCost({
+        input_tokens: inputTokens,
+        output_tokens: outputTokens,
+        cache_creation_input_tokens: cacheCreateTokens,
+        cache_read_input_tokens: cacheReadTokens
+      }, model);
+      
+      // 记录API Key级别的使用统计
      await redis.incrementTokenUsage(keyId, totalTokens, inputTokens, outputTokens, cacheCreateTokens, cacheReadTokens, model);
      
-      // 更新最后使用时间（性能优化：只在实际使用时更新）
+      // 记录费用统计
+      if (costInfo.costs.total > 0) {
+        await redis.incrementDailyCost(keyId, costInfo.costs.total);
+        logger.database(`💰 Recorded cost for ${keyId}: $${costInfo.costs.total.toFixed(6)}, model: ${model}`);
+      } else {
+        logger.debug(`💰 No cost recorded for ${keyId} - zero cost for model: ${model}`);
+      }
+      
+      // 获取API Key数据以确定关联的账户
      const keyData = await redis.getApiKey(keyId);
      if (keyData && Object.keys(keyData).length > 0) {
+        // 更新最后使用时间
        keyData.lastUsedAt = new Date().toISOString();
-        // 使用记录时不需要重新建立哈希映射
        await redis.setApiKey(keyId, keyData);
+        
+        // 记录账户级别的使用统计（只统计实际处理请求的账户）
+        if (accountId) {
+          await redis.incrementAccountUsage(accountId, totalTokens, inputTokens, outputTokens, cacheCreateTokens, cacheReadTokens, model);
+          logger.database(`📊 Recorded account usage: ${accountId} - ${totalTokens} tokens (API Key: ${keyId})`);
+        } else {
+          logger.debug('⚠️ No accountId provided for usage recording, skipping account-level statistics');
+        }
      }
      
      const logParts = [`Model: ${model}`, `Input: ${inputTokens}`, `Output: ${outputTokens}`];
@@ -274,6 +336,16 @@ class ApiKeyService {
    return await redis.getUsageStats(keyId);
  }

+  // 📊 获取账户使用统计
+  async getAccountUsageStats(accountId) {
+    return await redis.getAccountUsageStats(accountId);
+  }
+
+  // 📈 获取所有账户使用统计
+  async getAllAccountsUsageStats() {
+    return await redis.getAllAccountsUsageStats();
+  }
+

  // 🧹 清理过期的API Keys
  async cleanupExpiredKeys() {
@@ -283,14 +355,17 @@ class ApiKeyService {
      let cleanedCount = 0;

      for (const key of apiKeys) {
-        if (key.expiresAt && new Date(key.expiresAt) < now) {
-          await redis.deleteApiKey(key.id);
+        // 检查是否已过期且仍处于激活状态
+        if (key.expiresAt && new Date(key.expiresAt) < now && key.isActive === 'true') {
+          // 将过期的 API Key 标记为禁用状态，而不是直接删除
+          await this.updateApiKey(key.id, { isActive: false });
+          logger.info(`🔒 API Key ${key.id} (${key.name}) has expired and been disabled`);
          cleanedCount++;
        }
      }

      if (cleanedCount > 0) {
-        logger.success(`🧹 Cleaned up ${cleanedCount} expired API keys`);
+        logger.success(`🧹 Disabled ${cleanedCount} expired API keys`);
      }

      return cleanedCount;
--- a/src/services/claudeAccountService.js
+++ b/src/services/claudeAccountService.js
@@ -444,11 +444,11 @@ class ClaudeAccountService {
      }

      // 如果没有映射或映射无效，选择新账户
-      // 优先选择最近刷新过token的账户
+      // 优先选择最久未使用的账户（负载均衡）
      const sortedAccounts = activeAccounts.sort((a, b) => {
-        const aLastRefresh = new Date(a.lastRefreshAt || 0).getTime();
-        const bLastRefresh = new Date(b.lastRefreshAt || 0).getTime();
-        return bLastRefresh - aLastRefresh;
+        const aLastUsed = new Date(a.lastUsedAt || 0).getTime();
+        const bLastUsed = new Date(b.lastUsedAt || 0).getTime();
+        return aLastUsed - bLastUsed; // 最久未使用的优先
      });

      const selectedAccountId = sortedAccounts[0].id;
@@ -544,11 +544,11 @@ class ClaudeAccountService {
          return aRateLimitedAt - bRateLimitedAt; // 最早限流的优先
        });
      } else {
-        // 非限流账户按最近刷新时间排序
+        // 非限流账户按最后使用时间排序（最久未使用的优先）
        candidateAccounts = candidateAccounts.sort((a, b) => {
-          const aLastRefresh = new Date(a.lastRefreshAt || 0).getTime();
-          const bLastRefresh = new Date(b.lastRefreshAt || 0).getTime();
-          return bLastRefresh - aLastRefresh;
+          const aLastUsed = new Date(a.lastUsedAt || 0).getTime();
+          const bLastUsed = new Date(b.lastUsedAt || 0).getTime();
+          return aLastUsed - bLastUsed; // 最久未使用的优先
        });
      }

--- a/src/services/claudeRelayService.js
+++ b/src/services/claudeRelayService.js
@@ -181,6 +181,8 @@ class ClaudeRelayService {
      
      logger.info(`✅ API request completed - Key: ${apiKeyData.name}, Account: ${accountId}, Model: ${requestBody.model}, Input: ~${Math.round(inputTokens)} tokens, Output: ~${Math.round(outputTokens)} tokens`);
      
+      // 在响应中添加accountId，以便调用方记录账户级别统计
+      response.accountId = accountId;
      return response;
    } catch (error) {
      logger.error(`❌ Claude relay request failed for key: ${apiKeyData.name || apiKeyData.id}:`, error.message);
@@ -619,7 +621,10 @@ class ClaudeRelayService {
      const proxyAgent = await this._getProxyAgent(accountId);
      
      // 发送流式请求并捕获usage数据
-      return await this._makeClaudeStreamRequestWithUsageCapture(processedBody, accessToken, proxyAgent, clientHeaders, responseStream, usageCallback, accountId, sessionHash, streamTransformer, options);
+      return await this._makeClaudeStreamRequestWithUsageCapture(processedBody, accessToken, proxyAgent, clientHeaders, responseStream, (usageData) => {
+        // 在usageCallback中添加accountId
+        usageCallback({ ...usageData, accountId });
+      }, accountId, sessionHash, streamTransformer, options);
    } catch (error) {
      logger.error('❌ Claude stream relay with usage capture failed:', error);
      throw error;
--- a/src/services/costInitService.js
+++ b/src/services/costInitService.js
@@ -0,0 +1,182 @@
+const redis = require('../models/redis');
+const apiKeyService = require('./apiKeyService');
+const CostCalculator = require('../utils/costCalculator');
+const logger = require('../utils/logger');
+
+class CostInitService {
+  /**
+   * 初始化所有API Key的费用数据
+   * 扫描历史使用记录并计算费用
+   */
+  async initializeAllCosts() {
+    try {
+      logger.info('💰 Starting cost initialization for all API Keys...');
+      
+      const apiKeys = await apiKeyService.getAllApiKeys();
+      const client = redis.getClientSafe();
+      
+      let processedCount = 0;
+      let errorCount = 0;
+      
+      for (const apiKey of apiKeys) {
+        try {
+          await this.initializeApiKeyCosts(apiKey.id, client);
+          processedCount++;
+          
+          if (processedCount % 10 === 0) {
+            logger.info(`💰 Processed ${processedCount} API Keys...`);
+          }
+        } catch (error) {
+          errorCount++;
+          logger.error(`❌ Failed to initialize costs for API Key ${apiKey.id}:`, error);
+        }
+      }
+      
+      logger.success(`💰 Cost initialization completed! Processed: ${processedCount}, Errors: ${errorCount}`);
+      return { processed: processedCount, errors: errorCount };
+    } catch (error) {
+      logger.error('❌ Failed to initialize costs:', error);
+      throw error;
+    }
+  }
+  
+  /**
+   * 初始化单个API Key的费用数据
+   */
+  async initializeApiKeyCosts(apiKeyId, client) {
+    // 获取所有时间的模型使用统计
+    const modelKeys = await client.keys(`usage:${apiKeyId}:model:*:*:*`);
+    
+    // 按日期分组统计
+    const dailyCosts = new Map(); // date -> cost
+    const monthlyCosts = new Map(); // month -> cost
+    const hourlyCosts = new Map(); // hour -> cost
+    
+    for (const key of modelKeys) {
+      // 解析key格式: usage:{keyId}:model:{period}:{model}:{date}
+      const match = key.match(/usage:(.+):model:(daily|monthly|hourly):(.+):(\d{4}-\d{2}(?:-\d{2})?(?::\d{2})?)$/);
+      if (!match) continue;
+      
+      const [, , period, model, dateStr] = match;
+      
+      // 获取使用数据
+      const data = await client.hgetall(key);
+      if (!data || Object.keys(data).length === 0) continue;
+      
+      // 计算费用
+      const usage = {
+        input_tokens: parseInt(data.totalInputTokens) || parseInt(data.inputTokens) || 0,
+        output_tokens: parseInt(data.totalOutputTokens) || parseInt(data.outputTokens) || 0,
+        cache_creation_input_tokens: parseInt(data.totalCacheCreateTokens) || parseInt(data.cacheCreateTokens) || 0,
+        cache_read_input_tokens: parseInt(data.totalCacheReadTokens) || parseInt(data.cacheReadTokens) || 0
+      };
+      
+      const costResult = CostCalculator.calculateCost(usage, model);
+      const cost = costResult.costs.total;
+      
+      // 根据period分组累加费用
+      if (period === 'daily') {
+        const currentCost = dailyCosts.get(dateStr) || 0;
+        dailyCosts.set(dateStr, currentCost + cost);
+      } else if (period === 'monthly') {
+        const currentCost = monthlyCosts.get(dateStr) || 0;
+        monthlyCosts.set(dateStr, currentCost + cost);
+      } else if (period === 'hourly') {
+        const currentCost = hourlyCosts.get(dateStr) || 0;
+        hourlyCosts.set(dateStr, currentCost + cost);
+      }
+    }
+    
+    // 将计算出的费用写入Redis
+    const promises = [];
+    
+    // 写入每日费用
+    for (const [date, cost] of dailyCosts) {
+      const key = `usage:cost:daily:${apiKeyId}:${date}`;
+      promises.push(
+        client.set(key, cost.toString()),
+        client.expire(key, 86400 * 30) // 30天过期
+      );
+    }
+    
+    // 写入每月费用
+    for (const [month, cost] of monthlyCosts) {
+      const key = `usage:cost:monthly:${apiKeyId}:${month}`;
+      promises.push(
+        client.set(key, cost.toString()),
+        client.expire(key, 86400 * 90) // 90天过期
+      );
+    }
+    
+    // 写入每小时费用
+    for (const [hour, cost] of hourlyCosts) {
+      const key = `usage:cost:hourly:${apiKeyId}:${hour}`;
+      promises.push(
+        client.set(key, cost.toString()),
+        client.expire(key, 86400 * 7) // 7天过期
+      );
+    }
+    
+    // 计算总费用
+    let totalCost = 0;
+    for (const cost of dailyCosts.values()) {
+      totalCost += cost;
+    }
+    
+    // 写入总费用
+    if (totalCost > 0) {
+      const totalKey = `usage:cost:total:${apiKeyId}`;
+      promises.push(client.set(totalKey, totalCost.toString()));
+    }
+    
+    await Promise.all(promises);
+    
+    logger.debug(`💰 Initialized costs for API Key ${apiKeyId}: Daily entries: ${dailyCosts.size}, Total cost: $${totalCost.toFixed(2)}`);
+  }
+  
+  /**
+   * 检查是否需要初始化费用数据
+   */
+  async needsInitialization() {
+    try {
+      const client = redis.getClientSafe();
+      
+      // 检查是否有任何费用数据
+      const costKeys = await client.keys('usage:cost:*');
+      
+      // 如果没有费用数据，需要初始化
+      if (costKeys.length === 0) {
+        logger.info('💰 No cost data found, initialization needed');
+        return true;
+      }
+      
+      // 检查是否有使用数据但没有对应的费用数据
+      const sampleKeys = await client.keys('usage:*:model:daily:*:*');
+      if (sampleKeys.length > 10) {
+        // 抽样检查
+        const sampleSize = Math.min(10, sampleKeys.length);
+        for (let i = 0; i < sampleSize; i++) {
+          const usageKey = sampleKeys[Math.floor(Math.random() * sampleKeys.length)];
+          const match = usageKey.match(/usage:(.+):model:daily:(.+):(\d{4}-\d{2}-\d{2})$/);
+          if (match) {
+            const [, keyId, , date] = match;
+            const costKey = `usage:cost:daily:${keyId}:${date}`;
+            const hasCost = await client.exists(costKey);
+            if (!hasCost) {
+              logger.info(`💰 Found usage without cost data for key ${keyId} on ${date}, initialization needed`);
+              return true;
+            }
+          }
+        }
+      }
+      
+      logger.info('💰 Cost data appears to be up to date');
+      return false;
+    } catch (error) {
+      logger.error('❌ Failed to check initialization status:', error);
+      return false;
+    }
+  }
+}
+
+module.exports = new CostInitService();