mirror of
https://github.com/Wei-Shaw/claude-relay-service.git
synced 2026-01-23 00:53:33 +00:00
feat: 实现基于费用的速率限制功能
- 新增 rateLimitCost 字段,支持按费用进行速率限制 - 新增 weeklyOpusCostLimit 字段,支持 Opus 模型周费用限制 - 优化速率限制逻辑,支持费用、请求数、token多维度控制 - 更新前端界面,添加费用限制配置选项 - 增强账户管理功能,支持费用统计和限制 - 改进 Redis 数据模型,支持费用计数器 - 优化价格计算服务,支持更精确的成本核算 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
const apiKeyService = require('../services/apiKeyService')
|
||||
const logger = require('../utils/logger')
|
||||
const redis = require('../models/redis')
|
||||
const { RateLimiterRedis } = require('rate-limiter-flexible')
|
||||
// const { RateLimiterRedis } = require('rate-limiter-flexible') // 暂时未使用
|
||||
const config = require('../../config/config')
|
||||
|
||||
// 🔑 API Key验证中间件(优化版)
|
||||
@@ -182,11 +182,18 @@ const authenticateApiKey = async (req, res, next) => {
|
||||
// 检查时间窗口限流
|
||||
const rateLimitWindow = validation.keyData.rateLimitWindow || 0
|
||||
const rateLimitRequests = validation.keyData.rateLimitRequests || 0
|
||||
const rateLimitCost = validation.keyData.rateLimitCost || 0 // 新增:费用限制
|
||||
|
||||
if (rateLimitWindow > 0 && (rateLimitRequests > 0 || validation.keyData.tokenLimit > 0)) {
|
||||
// 兼容性检查:如果tokenLimit仍有值,使用tokenLimit;否则使用rateLimitCost
|
||||
const hasRateLimits =
|
||||
rateLimitWindow > 0 &&
|
||||
(rateLimitRequests > 0 || validation.keyData.tokenLimit > 0 || rateLimitCost > 0)
|
||||
|
||||
if (hasRateLimits) {
|
||||
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 costCountKey = `rate_limit:cost:${validation.keyData.id}` // 新增:费用计数器
|
||||
|
||||
const now = Date.now()
|
||||
const windowDuration = rateLimitWindow * 60 * 1000 // 转换为毫秒
|
||||
@@ -199,6 +206,7 @@ const authenticateApiKey = async (req, res, next) => {
|
||||
await redis.getClient().set(windowStartKey, now, 'PX', windowDuration)
|
||||
await redis.getClient().set(requestCountKey, 0, 'PX', windowDuration)
|
||||
await redis.getClient().set(tokenCountKey, 0, 'PX', windowDuration)
|
||||
await redis.getClient().set(costCountKey, 0, 'PX', windowDuration) // 新增:重置费用
|
||||
windowStart = now
|
||||
} else {
|
||||
windowStart = parseInt(windowStart)
|
||||
@@ -209,6 +217,7 @@ const authenticateApiKey = async (req, res, next) => {
|
||||
await redis.getClient().set(windowStartKey, now, 'PX', windowDuration)
|
||||
await redis.getClient().set(requestCountKey, 0, 'PX', windowDuration)
|
||||
await redis.getClient().set(tokenCountKey, 0, 'PX', windowDuration)
|
||||
await redis.getClient().set(costCountKey, 0, 'PX', windowDuration) // 新增:重置费用
|
||||
windowStart = now
|
||||
}
|
||||
}
|
||||
@@ -216,6 +225,7 @@ const authenticateApiKey = async (req, res, next) => {
|
||||
// 获取当前计数
|
||||
const currentRequests = parseInt((await redis.getClient().get(requestCountKey)) || '0')
|
||||
const currentTokens = parseInt((await redis.getClient().get(tokenCountKey)) || '0')
|
||||
const currentCost = parseFloat((await redis.getClient().get(costCountKey)) || '0') // 新增:当前费用
|
||||
|
||||
// 检查请求次数限制
|
||||
if (rateLimitRequests > 0 && currentRequests >= rateLimitRequests) {
|
||||
@@ -236,24 +246,46 @@ const authenticateApiKey = async (req, res, next) => {
|
||||
})
|
||||
}
|
||||
|
||||
// 检查Token使用量限制
|
||||
// 兼容性检查:优先使用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)
|
||||
if (tokenLimit > 0) {
|
||||
// 使用Token限制(向后兼容)
|
||||
if (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}`
|
||||
)
|
||||
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
|
||||
})
|
||||
return res.status(429).json({
|
||||
error: 'Rate limit exceeded',
|
||||
message: `已达到 Token 使用限制 (${tokenLimit} tokens),将在 ${remainingMinutes} 分钟后重置`,
|
||||
currentTokens,
|
||||
tokenLimit,
|
||||
resetAt: resetTime.toISOString(),
|
||||
remainingMinutes
|
||||
})
|
||||
}
|
||||
} else if (rateLimitCost > 0) {
|
||||
// 使用费用限制(新功能)
|
||||
if (currentCost >= rateLimitCost) {
|
||||
const resetTime = new Date(windowStart + windowDuration)
|
||||
const remainingMinutes = Math.ceil((resetTime - now) / 60000)
|
||||
|
||||
logger.security(
|
||||
`💰 Rate limit exceeded (cost) for key: ${validation.keyData.id} (${validation.keyData.name}), cost: $${currentCost.toFixed(2)}/$${rateLimitCost}`
|
||||
)
|
||||
|
||||
return res.status(429).json({
|
||||
error: 'Rate limit exceeded',
|
||||
message: `已达到费用限制 ($${rateLimitCost}),将在 ${remainingMinutes} 分钟后重置`,
|
||||
currentCost,
|
||||
costLimit: rateLimitCost,
|
||||
resetAt: resetTime.toISOString(),
|
||||
remainingMinutes
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 增加请求计数
|
||||
@@ -265,10 +297,13 @@ const authenticateApiKey = async (req, res, next) => {
|
||||
windowDuration,
|
||||
requestCountKey,
|
||||
tokenCountKey,
|
||||
costCountKey, // 新增:费用计数器
|
||||
currentRequests: currentRequests + 1,
|
||||
currentTokens,
|
||||
currentCost, // 新增:当前费用
|
||||
rateLimitRequests,
|
||||
tokenLimit
|
||||
tokenLimit,
|
||||
rateLimitCost // 新增:费用限制
|
||||
}
|
||||
}
|
||||
|
||||
@@ -297,6 +332,46 @@ const authenticateApiKey = async (req, res, next) => {
|
||||
)
|
||||
}
|
||||
|
||||
// 检查 Opus 周费用限制(仅对 Opus 模型生效)
|
||||
const weeklyOpusCostLimit = validation.keyData.weeklyOpusCostLimit || 0
|
||||
if (weeklyOpusCostLimit > 0) {
|
||||
// 从请求中获取模型信息
|
||||
const requestBody = req.body || {}
|
||||
const model = requestBody.model || ''
|
||||
|
||||
// 判断是否为 Opus 模型
|
||||
if (model && model.toLowerCase().includes('claude-opus')) {
|
||||
const weeklyOpusCost = validation.keyData.weeklyOpusCost || 0
|
||||
|
||||
if (weeklyOpusCost >= weeklyOpusCostLimit) {
|
||||
logger.security(
|
||||
`💰 Weekly Opus cost limit exceeded for key: ${validation.keyData.id} (${validation.keyData.name}), cost: $${weeklyOpusCost.toFixed(2)}/$${weeklyOpusCostLimit}`
|
||||
)
|
||||
|
||||
// 计算下周一的重置时间
|
||||
const now = new Date()
|
||||
const dayOfWeek = now.getDay()
|
||||
const daysUntilMonday = dayOfWeek === 0 ? 1 : (8 - dayOfWeek) % 7 || 7
|
||||
const resetDate = new Date(now)
|
||||
resetDate.setDate(now.getDate() + daysUntilMonday)
|
||||
resetDate.setHours(0, 0, 0, 0)
|
||||
|
||||
return res.status(429).json({
|
||||
error: 'Weekly Opus cost limit exceeded',
|
||||
message: `已达到 Opus 模型周费用限制 ($${weeklyOpusCostLimit})`,
|
||||
currentCost: weeklyOpusCost,
|
||||
costLimit: weeklyOpusCostLimit,
|
||||
resetAt: resetDate.toISOString() // 下周一重置
|
||||
})
|
||||
}
|
||||
|
||||
// 记录当前 Opus 费用使用情况
|
||||
logger.api(
|
||||
`💰 Opus weekly cost usage for key: ${validation.keyData.id} (${validation.keyData.name}), current: $${weeklyOpusCost.toFixed(2)}/$${weeklyOpusCostLimit}`
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// 将验证信息添加到请求对象(只包含必要信息)
|
||||
req.apiKey = {
|
||||
id: validation.keyData.id,
|
||||
@@ -311,6 +386,7 @@ const authenticateApiKey = async (req, res, next) => {
|
||||
concurrencyLimit: validation.keyData.concurrencyLimit,
|
||||
rateLimitWindow: validation.keyData.rateLimitWindow,
|
||||
rateLimitRequests: validation.keyData.rateLimitRequests,
|
||||
rateLimitCost: validation.keyData.rateLimitCost, // 新增:费用限制
|
||||
enableModelRestriction: validation.keyData.enableModelRestriction,
|
||||
restrictedModels: validation.keyData.restrictedModels,
|
||||
enableClientRestriction: validation.keyData.enableClientRestriction,
|
||||
@@ -713,35 +789,41 @@ const errorHandler = (error, req, res, _next) => {
|
||||
}
|
||||
|
||||
// 🌐 全局速率限制中间件(延迟初始化)
|
||||
let rateLimiter = null
|
||||
// const rateLimiter = null // 暂时未使用
|
||||
|
||||
const getRateLimiter = () => {
|
||||
if (!rateLimiter) {
|
||||
try {
|
||||
const client = redis.getClient()
|
||||
if (!client) {
|
||||
logger.warn('⚠️ Redis client not available for rate limiter')
|
||||
return null
|
||||
}
|
||||
// 暂时注释掉未使用的函数
|
||||
// const getRateLimiter = () => {
|
||||
// if (!rateLimiter) {
|
||||
// try {
|
||||
// const client = redis.getClient()
|
||||
// if (!client) {
|
||||
// logger.warn('⚠️ Redis client not available for rate limiter')
|
||||
// return null
|
||||
// }
|
||||
//
|
||||
// rateLimiter = new RateLimiterRedis({
|
||||
// storeClient: client,
|
||||
// keyPrefix: 'global_rate_limit',
|
||||
// points: 1000, // 请求数量
|
||||
// duration: 900, // 15分钟 (900秒)
|
||||
// blockDuration: 900 // 阻塞时间15分钟
|
||||
// })
|
||||
//
|
||||
// logger.info('✅ Rate limiter initialized successfully')
|
||||
// } catch (error) {
|
||||
// logger.warn('⚠️ Rate limiter initialization failed, using fallback', { error: error.message })
|
||||
// return null
|
||||
// }
|
||||
// }
|
||||
// return rateLimiter
|
||||
// }
|
||||
|
||||
rateLimiter = new RateLimiterRedis({
|
||||
storeClient: client,
|
||||
keyPrefix: 'global_rate_limit',
|
||||
points: 1000, // 请求数量
|
||||
duration: 900, // 15分钟 (900秒)
|
||||
blockDuration: 900 // 阻塞时间15分钟
|
||||
})
|
||||
const globalRateLimit = async (req, res, next) =>
|
||||
// 已禁用全局IP限流 - 直接跳过所有请求
|
||||
next()
|
||||
|
||||
logger.info('✅ Rate limiter initialized successfully')
|
||||
} catch (error) {
|
||||
logger.warn('⚠️ Rate limiter initialization failed, using fallback', { error: error.message })
|
||||
return null
|
||||
}
|
||||
}
|
||||
return rateLimiter
|
||||
}
|
||||
|
||||
const globalRateLimit = async (req, res, next) => {
|
||||
// 以下代码已被禁用
|
||||
/*
|
||||
// 跳过健康检查和内部请求
|
||||
if (req.path === '/health' || req.path === '/api/health') {
|
||||
return next()
|
||||
@@ -777,7 +859,7 @@ const globalRateLimit = async (req, res, next) => {
|
||||
retryAfter: Math.round(msBeforeNext / 1000)
|
||||
})
|
||||
}
|
||||
}
|
||||
*/
|
||||
|
||||
// 📊 请求大小限制中间件
|
||||
const requestSizeLimit = (req, res, next) => {
|
||||
|
||||
Reference in New Issue
Block a user