feat: 添加API Key时间窗口限流功能并移除累计总量限制

- 新增时间窗口限流功能，支持按分钟设置时间窗口 - 支持在时间窗口内限制请求次数和Token使用量 - 移除原有的累计总量限制，只保留时间窗口限制 - Token统计包含所有4种类型：输入、输出、缓存创建、缓存读取 - 前端UI优化，明确显示限流参数的作用范围 - 限流触发时提供友好的错误提示和重置时间 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
2026-01-23 09:38:02 +00:00 · 2025-07-20 15:58:00 +08:00
parent 0aa986a0d8
commit 088ce266ba
7 changed files with 224 additions and 16 deletions
--- a/src/middleware/auth.js
+++ b/src/middleware/auth.js
@@ -103,6 +103,95 @@ const authenticateApiKey = async (req, res, next) => {
      };
    }

+    // 检查时间窗口限流
+    const rateLimitWindow = validation.keyData.rateLimitWindow || 0;
+    const rateLimitRequests = validation.keyData.rateLimitRequests || 0;
+    
+    if (rateLimitWindow > 0 && (rateLimitRequests > 0 || validation.keyData.tokenLimit > 0)) {
+      const windowStartKey = `rate_limit:window_start:${validation.keyData.id}`;
+      const requestCountKey = `rate_limit:requests:${validation.keyData.id}`;
+      const tokenCountKey = `rate_limit:tokens:${validation.keyData.id}`;
+      
+      const now = Date.now();
+      const windowDuration = rateLimitWindow * 60 * 1000; // 转换为毫秒
+      
+      // 获取窗口开始时间
+      let windowStart = await redis.getClient().get(windowStartKey);
+      
+      if (!windowStart) {
+        // 第一次请求，设置窗口开始时间
+        await redis.getClient().set(windowStartKey, now, 'PX', windowDuration);
+        await redis.getClient().set(requestCountKey, 0, 'PX', windowDuration);
+        await redis.getClient().set(tokenCountKey, 0, 'PX', windowDuration);
+        windowStart = now;
+      } else {
+        windowStart = parseInt(windowStart);
+        
+        // 检查窗口是否已过期
+        if (now - windowStart >= windowDuration) {
+          // 窗口已过期，重置
+          await redis.getClient().set(windowStartKey, now, 'PX', windowDuration);
+          await redis.getClient().set(requestCountKey, 0, 'PX', windowDuration);
+          await redis.getClient().set(tokenCountKey, 0, 'PX', windowDuration);
+          windowStart = now;
+        }
+      }
+      
+      // 获取当前计数
+      const currentRequests = parseInt(await redis.getClient().get(requestCountKey) || '0');
+      const currentTokens = parseInt(await redis.getClient().get(tokenCountKey) || '0');
+      
+      // 检查请求次数限制
+      if (rateLimitRequests > 0 && currentRequests >= rateLimitRequests) {
+        const resetTime = new Date(windowStart + windowDuration);
+        const remainingMinutes = Math.ceil((resetTime - now) / 60000);
+        
+        logger.security(`🚦 Rate limit exceeded (requests) for key: ${validation.keyData.id} (${validation.keyData.name}), requests: ${currentRequests}/${rateLimitRequests}`);
+        
+        return res.status(429).json({
+          error: 'Rate limit exceeded',
+          message: `已达到请求次数限制 (${rateLimitRequests} 次)，将在 ${remainingMinutes} 分钟后重置`,
+          currentRequests,
+          requestLimit: rateLimitRequests,
+          resetAt: resetTime.toISOString(),
+          remainingMinutes
+        });
+      }
+      
+      // 检查Token使用量限制
+      const tokenLimit = parseInt(validation.keyData.tokenLimit);
+      if (tokenLimit > 0 && currentTokens >= tokenLimit) {
+        const resetTime = new Date(windowStart + windowDuration);
+        const remainingMinutes = Math.ceil((resetTime - now) / 60000);
+        
+        logger.security(`🚦 Rate limit exceeded (tokens) for key: ${validation.keyData.id} (${validation.keyData.name}), tokens: ${currentTokens}/${tokenLimit}`);
+        
+        return res.status(429).json({
+          error: 'Rate limit exceeded',
+          message: `已达到 Token 使用限制 (${tokenLimit} tokens)，将在 ${remainingMinutes} 分钟后重置`,
+          currentTokens,
+          tokenLimit,
+          resetAt: resetTime.toISOString(),
+          remainingMinutes
+        });
+      }
+      
+      // 增加请求计数
+      await redis.getClient().incr(requestCountKey);
+      
+      // 存储限流信息到请求对象
+      req.rateLimitInfo = {
+        windowStart,
+        windowDuration,
+        requestCountKey,
+        tokenCountKey,
+        currentRequests: currentRequests + 1,
+        currentTokens,
+        rateLimitRequests,
+        tokenLimit
+      };
+    }
+    
    // 将验证信息添加到请求对象（只包含必要信息）
    req.apiKey = {
      id: validation.keyData.id,
@@ -110,6 +199,8 @@ const authenticateApiKey = async (req, res, next) => {
      tokenLimit: validation.keyData.tokenLimit,
      claudeAccountId: validation.keyData.claudeAccountId,
      concurrencyLimit: validation.keyData.concurrencyLimit,
+      rateLimitWindow: validation.keyData.rateLimitWindow,
+      rateLimitRequests: validation.keyData.rateLimitRequests,
      enableModelRestriction: validation.keyData.enableModelRestriction,
      restrictedModels: validation.keyData.restrictedModels
    };
--- a/src/routes/admin.js
+++ b/src/routes/admin.js
@@ -33,6 +33,8 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
      expiresAt,
      claudeAccountId,
      concurrencyLimit,
+      rateLimitWindow,
+      rateLimitRequests,
      enableModelRestriction,
      restrictedModels
    } = req.body;
@@ -58,6 +60,14 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
    if (concurrencyLimit !== undefined && concurrencyLimit !== null && concurrencyLimit !== '' && (!Number.isInteger(Number(concurrencyLimit)) || Number(concurrencyLimit) < 0)) {
      return res.status(400).json({ error: 'Concurrency limit must be a non-negative integer' });
    }
+    
+    if (rateLimitWindow !== undefined && rateLimitWindow !== null && rateLimitWindow !== '' && (!Number.isInteger(Number(rateLimitWindow)) || Number(rateLimitWindow) < 1)) {
+      return res.status(400).json({ error: 'Rate limit window must be a positive integer (minutes)' });
+    }
+    
+    if (rateLimitRequests !== undefined && rateLimitRequests !== null && rateLimitRequests !== '' && (!Number.isInteger(Number(rateLimitRequests)) || Number(rateLimitRequests) < 1)) {
+      return res.status(400).json({ error: 'Rate limit requests must be a positive integer' });
+    }

    // 验证模型限制字段
    if (enableModelRestriction !== undefined && typeof enableModelRestriction !== 'boolean') {
@@ -75,6 +85,8 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
      expiresAt,
      claudeAccountId,
      concurrencyLimit,
+      rateLimitWindow,
+      rateLimitRequests,
      enableModelRestriction,
      restrictedModels
    });
@@ -91,7 +103,7 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
 router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
  try {
    const { keyId } = req.params;
-    const { tokenLimit, concurrencyLimit, claudeAccountId, enableModelRestriction, restrictedModels } = req.body;
+    const { tokenLimit, concurrencyLimit, rateLimitWindow, rateLimitRequests, claudeAccountId, enableModelRestriction, restrictedModels } = req.body;

    // 只允许更新指定字段
    const updates = {};
@@ -109,6 +121,20 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
      }
      updates.concurrencyLimit = Number(concurrencyLimit);
    }
+    
+    if (rateLimitWindow !== undefined && rateLimitWindow !== null && rateLimitWindow !== '') {
+      if (!Number.isInteger(Number(rateLimitWindow)) || Number(rateLimitWindow) < 0) {
+        return res.status(400).json({ error: 'Rate limit window must be a non-negative integer (minutes)' });
+      }
+      updates.rateLimitWindow = Number(rateLimitWindow);
+    }
+    
+    if (rateLimitRequests !== undefined && rateLimitRequests !== null && rateLimitRequests !== '') {
+      if (!Number.isInteger(Number(rateLimitRequests)) || Number(rateLimitRequests) < 0) {
+        return res.status(400).json({ error: 'Rate limit requests must be a non-negative integer' });
+      }
+      updates.rateLimitRequests = Number(rateLimitRequests);
+    }

    if (claudeAccountId !== undefined) {
      // 空字符串表示解绑，null或空字符串都设置为空字符串
--- a/src/routes/api.js
+++ b/src/routes/api.js
@@ -3,6 +3,7 @@ const claudeRelayService = require('../services/claudeRelayService');
 const apiKeyService = require('../services/apiKeyService');
 const { authenticateApiKey } = require('../middleware/auth');
 const logger = require('../utils/logger');
+const redis = require('../models/redis');

 const router = express.Router();

@@ -66,6 +67,15 @@ router.post('/v1/messages', authenticateApiKey, async (req, res) => {
            logger.error('❌ Failed to record stream usage:', error);
          });
          
+          // 更新时间窗口内的token计数
+          if (req.rateLimitInfo) {
+            const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens;
+            redis.getClient().incrby(req.rateLimitInfo.tokenCountKey, totalTokens).catch(error => {
+              logger.error('❌ Failed to update rate limit token count:', error);
+            });
+            logger.api(`📊 Updated rate limit token count: +${totalTokens} tokens`);
+          }
+          
          usageDataCaptured = true;
          logger.api(`📊 Stream usage recorded (real) - Model: ${model}, Input: ${inputTokens}, Output: ${outputTokens}, Cache Create: ${cacheCreateTokens}, Cache Read: ${cacheReadTokens}, Total: ${inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens} tokens`);
        } else {
@@ -122,6 +132,13 @@ router.post('/v1/messages', authenticateApiKey, async (req, res) => {
          // 记录真实的token使用量（包含模型信息和所有4种token）
          await apiKeyService.recordUsage(req.apiKey.id, inputTokens, outputTokens, cacheCreateTokens, cacheReadTokens, model);
          
+          // 更新时间窗口内的token计数
+          if (req.rateLimitInfo) {
+            const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens;
+            await redis.getClient().incrby(req.rateLimitInfo.tokenCountKey, totalTokens);
+            logger.api(`📊 Updated rate limit token count: +${totalTokens} tokens`);
+          }
+          
          usageRecorded = true;
          logger.api(`📊 Non-stream usage recorded (real) - Model: ${model}, Input: ${inputTokens}, Output: ${outputTokens}, Cache Create: ${cacheCreateTokens}, Cache Read: ${cacheReadTokens}, Total: ${inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens} tokens`);
        } else {
--- a/src/services/apiKeyService.js
+++ b/src/services/apiKeyService.js
@@ -19,6 +19,8 @@ class ApiKeyService {
      claudeAccountId = null,
      isActive = true,
      concurrencyLimit = 0,
+      rateLimitWindow = null,
+      rateLimitRequests = null,
      enableModelRestriction = false,
      restrictedModels = []
    } = options;
@@ -35,6 +37,8 @@ class ApiKeyService {
      apiKey: hashedKey,
      tokenLimit: String(tokenLimit ?? 0),
      concurrencyLimit: String(concurrencyLimit ?? 0),
+      rateLimitWindow: String(rateLimitWindow ?? 0),
+      rateLimitRequests: String(rateLimitRequests ?? 0),
      isActive: String(isActive),
      claudeAccountId: claudeAccountId || '',
      enableModelRestriction: String(enableModelRestriction),
@@ -57,6 +61,8 @@ class ApiKeyService {
      description: keyData.description,
      tokenLimit: parseInt(keyData.tokenLimit),
      concurrencyLimit: parseInt(keyData.concurrencyLimit),
+      rateLimitWindow: parseInt(keyData.rateLimitWindow || 0),
+      rateLimitRequests: parseInt(keyData.rateLimitRequests || 0),
      isActive: keyData.isActive === 'true',
      claudeAccountId: keyData.claudeAccountId,
      enableModelRestriction: keyData.enableModelRestriction === 'true',
@@ -94,14 +100,8 @@ class ApiKeyService {
        return { valid: false, error: 'API key has expired' };
      }

-      // 检查使用限制
+      // 获取使用统计（供返回数据使用）
      const usage = await redis.getUsageStats(keyData.id);
-      const tokenLimit = parseInt(keyData.tokenLimit);
-      
-      if (tokenLimit > 0 && usage.total.tokens >= tokenLimit) {
-        return { valid: false, error: 'Token limit exceeded' };
-      }
-

      // 更新最后使用时间（优化：只在实际API调用时更新，而不是验证时）
      // 注意：lastUsedAt的更新已移至recordUsage方法中
@@ -124,6 +124,8 @@ class ApiKeyService {
          claudeAccountId: keyData.claudeAccountId,
          tokenLimit: parseInt(keyData.tokenLimit),
          concurrencyLimit: parseInt(keyData.concurrencyLimit || 0),
+          rateLimitWindow: parseInt(keyData.rateLimitWindow || 0),
+          rateLimitRequests: parseInt(keyData.rateLimitRequests || 0),
          enableModelRestriction: keyData.enableModelRestriction === 'true',
          restrictedModels: restrictedModels,
          usage
@@ -145,6 +147,8 @@ class ApiKeyService {
        key.usage = await redis.getUsageStats(key.id);
        key.tokenLimit = parseInt(key.tokenLimit);
        key.concurrencyLimit = parseInt(key.concurrencyLimit || 0);
+        key.rateLimitWindow = parseInt(key.rateLimitWindow || 0);
+        key.rateLimitRequests = parseInt(key.rateLimitRequests || 0);
        key.currentConcurrency = await redis.getConcurrency(key.id);
        key.isActive = key.isActive === 'true';
        key.enableModelRestriction = key.enableModelRestriction === 'true';
@@ -172,7 +176,7 @@ class ApiKeyService {
      }

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

      for (const [field, value] of Object.entries(updates)) {
--- a/src/services/claudeRelayService.js
+++ b/src/services/claudeRelayService.js
@@ -23,7 +23,7 @@ class ClaudeRelayService {
    
    try {
      // 调试日志：查看API Key数据
-      logger.info(`🔍 API Key data received:`, {
+      logger.info('🔍 API Key data received:', {
        apiKeyName: apiKeyData.name,
        enableModelRestriction: apiKeyData.enableModelRestriction,
        restrictedModels: apiKeyData.restrictedModels,
@@ -448,7 +448,7 @@ class ClaudeRelayService {
  async relayStreamRequestWithUsageCapture(requestBody, apiKeyData, responseStream, clientHeaders, usageCallback) {
    try {
      // 调试日志：查看API Key数据（流式请求）
-      logger.info(`🔍 [Stream] API Key data received:`, {
+      logger.info('🔍 [Stream] API Key data received:', {
        apiKeyName: apiKeyData.name,
        enableModelRestriction: apiKeyData.enableModelRestriction,
        restrictedModels: apiKeyData.restrictedModels,