github/claude-relay-service
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/claude-relay-service.git synced 2026-01-23 00:53:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: 实现基于费用的速率限制功能

Browse Source 
- 新增 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:
shaw
2025-08-31 17:27:37 +08:00
parent a54622e3d7
commit e84c6a5555
21 changed files with 1662 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
src/routes/admin.js
View File

@@ -397,11 +397,13 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
concurrencyLimit,
rateLimitWindow,
rateLimitRequests,
rateLimitCost,
enableModelRestriction,
restrictedModels,
enableClientRestriction,
allowedClients,
dailyCostLimit,
weeklyOpusCostLimit,
tags
} = req.body
@@ -494,11 +496,13 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
concurrencyLimit,
rateLimitWindow,
rateLimitRequests,
rateLimitCost,
enableModelRestriction,
restrictedModels,
enableClientRestriction,
allowedClients,
dailyCostLimit,
weeklyOpusCostLimit,
tags
})
@@ -532,6 +536,7 @@ router.post('/api-keys/batch', authenticateAdmin, async (req, res) => {
enableClientRestriction,
allowedClients,
dailyCostLimit,
weeklyOpusCostLimit,
tags
} = req.body
@@ -575,6 +580,7 @@ router.post('/api-keys/batch', authenticateAdmin, async (req, res) => {
enableClientRestriction,
allowedClients,
dailyCostLimit,
weeklyOpusCostLimit,
tags
})
@@ -685,6 +691,9 @@ router.put('/api-keys/batch', authenticateAdmin, async (req, res) => {
if (updates.dailyCostLimit !== undefined) {
finalUpdates.dailyCostLimit = updates.dailyCostLimit
}
if (updates.weeklyOpusCostLimit !== undefined) {
finalUpdates.weeklyOpusCostLimit = updates.weeklyOpusCostLimit
}
if (updates.permissions !== undefined) {
finalUpdates.permissions = updates.permissions
}
@@ -795,6 +804,7 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
concurrencyLimit,
rateLimitWindow,
rateLimitRequests,
rateLimitCost,
isActive,
claudeAccountId,
claudeConsoleAccountId,
@@ -808,6 +818,7 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
allowedClients,
expiresAt,
dailyCostLimit,
weeklyOpusCostLimit,
tags
} = req.body
@@ -844,6 +855,14 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
updates.rateLimitRequests = Number(rateLimitRequests)
}
if (rateLimitCost !== undefined && rateLimitCost !== null && rateLimitCost !== '') {
const cost = Number(rateLimitCost)
if (isNaN(cost) || cost < 0) {
return res.status(400).json({ error: 'Rate limit cost must be a non-negative number' })
}
updates.rateLimitCost = cost
}
if (claudeAccountId !== undefined) {
// 空字符串表示解绑，null或空字符串都设置为空字符串
updates.claudeAccountId = claudeAccountId || ''
@@ -935,6 +954,22 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
updates.dailyCostLimit = costLimit
}
// 处理 Opus 周费用限制
if (
weeklyOpusCostLimit !== undefined &&
weeklyOpusCostLimit !== null &&
weeklyOpusCostLimit !== ''
) {
const costLimit = Number(weeklyOpusCostLimit)
// 明确验证非负数（0 表示禁用，负数无意义）
if (isNaN(costLimit) || costLimit < 0) {
return res
.status(400)
.json({ error: 'Weekly Opus cost limit must be a non-negative number' })
}
updates.weeklyOpusCostLimit = costLimit
}
// 处理标签
if (tags !== undefined) {
if (!Array.isArray(tags)) {
@@ -1468,13 +1503,53 @@ router.get('/claude-accounts', authenticateAdmin, async (req, res) => {
const usageStats = await redis.getAccountUsageStats(account.id)
const groupInfos = await accountGroupService.getAccountGroup(account.id)
// 获取会话窗口使用统计（仅对有活跃窗口的账户）
let sessionWindowUsage = null
if (account.sessionWindow && account.sessionWindow.hasActiveWindow) {
const windowUsage = await redis.getAccountSessionWindowUsage(
account.id,
account.sessionWindow.windowStart,
account.sessionWindow.windowEnd
)
// 计算会话窗口的总费用
let totalCost = 0
const modelCosts = {}
for (const [modelName, usage] of Object.entries(windowUsage.modelUsage)) {
const usageData = {
input_tokens: usage.inputTokens,
output_tokens: usage.outputTokens,
cache_creation_input_tokens: usage.cacheCreateTokens,
cache_read_input_tokens: usage.cacheReadTokens
}
const costResult = CostCalculator.calculateCost(usageData, modelName)
modelCosts[modelName] = {
...usage,
cost: costResult.costs.total
}
totalCost += costResult.costs.total
}
sessionWindowUsage = {
totalTokens: windowUsage.totalAllTokens,
totalRequests: windowUsage.totalRequests,
totalCost,
modelUsage: modelCosts
}
}
return {
...account,
// 转换schedulable为布尔值
schedulable: account.schedulable === 'true' || account.schedulable === true,
groupInfos,
usage: {
daily: usageStats.daily,
total: usageStats.total,
averages: usageStats.averages
averages: usageStats.averages,
sessionWindow: sessionWindowUsage
}
}
} catch (statsError) {
@@ -1488,7 +1563,8 @@ router.get('/claude-accounts', authenticateAdmin, async (req, res) => {
usage: {
daily: { tokens: 0, requests: 0, allTokens: 0 },
total: { tokens: 0, requests: 0, allTokens: 0 },
averages: { rpm: 0, tpm: 0 }
averages: { rpm: 0, tpm: 0 },
sessionWindow: null
}
}
} catch (groupError) {
@@ -1502,7 +1578,8 @@ router.get('/claude-accounts', authenticateAdmin, async (req, res) => {
usage: {
daily: { tokens: 0, requests: 0, allTokens: 0 },
total: { tokens: 0, requests: 0, allTokens: 0 },
averages: { rpm: 0, tpm: 0 }
averages: { rpm: 0, tpm: 0 },
sessionWindow: null
}
}
}
@@ -1531,7 +1608,8 @@ router.post('/claude-accounts', authenticateAdmin, async (req, res) => {
accountType,
platform = 'claude',
priority,
groupId
groupId,
autoStopOnWarning
} = req.body
if (!name) {
@@ -1568,7 +1646,8 @@ router.post('/claude-accounts', authenticateAdmin, async (req, res) => {
proxy,
accountType: accountType || 'shared', // 默认为共享类型
platform,
priority: priority || 50 // 默认优先级为50
priority: priority || 50, // 默认优先级为50
autoStopOnWarning: autoStopOnWarning === true // 默认为false
})
// 如果是分组类型，将账户添加到分组
@@ -1826,6 +1905,8 @@ router.get('/claude-console-accounts', authenticateAdmin, async (req, res) => {
return {
...account,
// 转换schedulable为布尔值
schedulable: account.schedulable === 'true' || account.schedulable === true,
groupInfos,
usage: {
daily: usageStats.daily,
@@ -1842,6 +1923,8 @@ router.get('/claude-console-accounts', authenticateAdmin, async (req, res) => {
const groupInfos = await accountGroupService.getAccountGroup(account.id)
return {
...account,
// 转换schedulable为布尔值
schedulable: account.schedulable === 'true' || account.schedulable === true,
groupInfos,
usage: {
daily: { tokens: 0, requests: 0, allTokens: 0 },

84
src/routes/api.js
View File

@@ -5,6 +5,7 @@ const bedrockRelayService = require('../services/bedrockRelayService')
const bedrockAccountService = require('../services/bedrockAccountService')
const unifiedClaudeScheduler = require('../services/unifiedClaudeScheduler')
const apiKeyService = require('../services/apiKeyService')
const pricingService = require('../services/pricingService')
const { authenticateApiKey } = require('../middleware/auth')
const logger = require('../utils/logger')
const redis = require('../models/redis')
@@ -131,14 +132,16 @@ async function handleMessagesRequest(req, res) {
}
apiKeyService
.recordUsageWithDetails(req.apiKey.id, usageObject, model, usageAccountId)
.recordUsageWithDetails(req.apiKey.id, usageObject, model, usageAccountId, 'claude')
.catch((error) => {
logger.error('❌ Failed to record stream usage:', error)
})
// 更新时间窗口内的token计数
// 更新时间窗口内的token计数和费用
if (req.rateLimitInfo) {
const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens
// 更新Token计数（向后兼容）
redis
.getClient()
.incrby(req.rateLimitInfo.tokenCountKey, totalTokens)
@@ -146,6 +149,22 @@ async function handleMessagesRequest(req, res) {
logger.error('❌ Failed to update rate limit token count:', error)
})
logger.api(`📊 Updated rate limit token count: +${totalTokens} tokens`)
// 计算并更新费用计数（新功能）
if (req.rateLimitInfo.costCountKey) {
const costInfo = pricingService.calculateCost(usageData, model)
if (costInfo.totalCost > 0) {
redis
.getClient()
.incrbyfloat(req.rateLimitInfo.costCountKey, costInfo.totalCost)
.catch((error) => {
logger.error('❌ Failed to update rate limit cost count:', error)
})
logger.api(
`💰 Updated rate limit cost count: +$${costInfo.totalCost.toFixed(6)}`
)
}
}
}
usageDataCaptured = true
@@ -216,14 +235,22 @@ async function handleMessagesRequest(req, res) {
}
apiKeyService
.recordUsageWithDetails(req.apiKey.id, usageObject, model, usageAccountId)
.recordUsageWithDetails(
req.apiKey.id,
usageObject,
model,
usageAccountId,
'claude-console'
)
.catch((error) => {
logger.error('❌ Failed to record stream usage:', error)
})
// 更新时间窗口内的token计数
// 更新时间窗口内的token计数和费用
if (req.rateLimitInfo) {
const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens
// 更新Token计数（向后兼容）
redis
.getClient()
.incrby(req.rateLimitInfo.tokenCountKey, totalTokens)
@@ -231,6 +258,22 @@ async function handleMessagesRequest(req, res) {
logger.error('❌ Failed to update rate limit token count:', error)
})
logger.api(`📊 Updated rate limit token count: +${totalTokens} tokens`)
// 计算并更新费用计数（新功能）
if (req.rateLimitInfo.costCountKey) {
const costInfo = pricingService.calculateCost(usageData, model)
if (costInfo.totalCost > 0) {
redis
.getClient()
.incrbyfloat(req.rateLimitInfo.costCountKey, costInfo.totalCost)
.catch((error) => {
logger.error('❌ Failed to update rate limit cost count:', error)
})
logger.api(
`💰 Updated rate limit cost count: +$${costInfo.totalCost.toFixed(6)}`
)
}
}
}
usageDataCaptured = true
@@ -271,9 +314,11 @@ async function handleMessagesRequest(req, res) {
logger.error('❌ Failed to record Bedrock stream usage:', error)
})
// 更新时间窗口内的token计数
// 更新时间窗口内的token计数和费用
if (req.rateLimitInfo) {
const totalTokens = inputTokens + outputTokens
// 更新Token计数（向后兼容）
redis
.getClient()
.incrby(req.rateLimitInfo.tokenCountKey, totalTokens)
@@ -281,6 +326,20 @@ async function handleMessagesRequest(req, res) {
logger.error('❌ Failed to update rate limit token count:', error)
})
logger.api(`📊 Updated rate limit token count: +${totalTokens} tokens`)
// 计算并更新费用计数（新功能）
if (req.rateLimitInfo.costCountKey) {
const costInfo = pricingService.calculateCost(result.usage, result.model)
if (costInfo.totalCost > 0) {
redis
.getClient()
.incrbyfloat(req.rateLimitInfo.costCountKey, costInfo.totalCost)
.catch((error) => {
logger.error('❌ Failed to update rate limit cost count:', error)
})
logger.api(`💰 Updated rate limit cost count: +$${costInfo.totalCost.toFixed(6)}`)
}
}
}
usageDataCaptured = true
@@ -438,11 +497,24 @@ async function handleMessagesRequest(req, res) {
responseAccountId
)
// 更新时间窗口内的token计数
// 更新时间窗口内的token计数和费用
if (req.rateLimitInfo) {
const totalTokens = inputTokens + outputTokens + cacheCreateTokens + cacheReadTokens
// 更新Token计数（向后兼容）
await redis.getClient().incrby(req.rateLimitInfo.tokenCountKey, totalTokens)
logger.api(`📊 Updated rate limit token count: +${totalTokens} tokens`)
// 计算并更新费用计数（新功能）
if (req.rateLimitInfo.costCountKey) {
const costInfo = pricingService.calculateCost(jsonData.usage, model)
if (costInfo.totalCost > 0) {
await redis
.getClient()
.incrbyfloat(req.rateLimitInfo.costCountKey, costInfo.totalCost)
logger.api(`💰 Updated rate limit cost count: +$${costInfo.totalCost.toFixed(6)}`)
}
}
}
usageRecorded = true

8
src/routes/apiStats.js
View File

@@ -278,21 +278,24 @@ router.post('/api/user-stats', async (req, res) => {
// 获取当前使用量
let currentWindowRequests = 0
let currentWindowTokens = 0
let currentWindowCost = 0 // 新增：当前窗口费用
let currentDailyCost = 0
let windowStartTime = null
let windowEndTime = null
let windowRemainingSeconds = null
try {
// 获取当前时间窗口的请求次数和Token使用量
// 获取当前时间窗口的请求次数、Token使用量和费用
if (fullKeyData.rateLimitWindow > 0) {
const client = redis.getClientSafe()
const requestCountKey = `rate_limit:requests:${keyId}`
const tokenCountKey = `rate_limit:tokens:${keyId}`
const costCountKey = `rate_limit:cost:${keyId}` // 新增：费用计数key
const windowStartKey = `rate_limit:window_start:${keyId}`
currentWindowRequests = parseInt((await client.get(requestCountKey)) || '0')
currentWindowTokens = parseInt((await client.get(tokenCountKey)) || '0')
currentWindowCost = parseFloat((await client.get(costCountKey)) || '0') // 新增：获取当前窗口费用
// 获取窗口开始时间和计算剩余时间
const windowStart = await client.get(windowStartKey)
@@ -313,6 +316,7 @@ router.post('/api/user-stats', async (req, res) => {
// 重置计数为0，因为窗口已过期
currentWindowRequests = 0
currentWindowTokens = 0
currentWindowCost = 0 // 新增：重置窗口费用
}
}
}
@@ -356,10 +360,12 @@ router.post('/api/user-stats', async (req, res) => {
concurrencyLimit: fullKeyData.concurrencyLimit || 0,
rateLimitWindow: fullKeyData.rateLimitWindow || 0,
rateLimitRequests: fullKeyData.rateLimitRequests || 0,
rateLimitCost: parseFloat(fullKeyData.rateLimitCost) || 0, // 新增：费用限制
dailyCostLimit: fullKeyData.dailyCostLimit || 0,
// 当前使用量
currentWindowRequests,
currentWindowTokens,
currentWindowCost, // 新增：当前窗口费用
currentDailyCost,
// 时间窗口信息
windowStartTime,