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

feat: 给key增加总用量限制

Browse Source
This commit is contained in:
itzhan
2025-09-19 21:57:24 +08:00
parent 4d78471891
commit ec28b66e7f
11 changed files with 604 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
src/middleware/auth.js
View File

@@ -308,6 +308,29 @@ const authenticateApiKey = async (req, res, next) => {
}
}
// 检查总额度限制（基于累计费用）
const totalUsageLimit = Number(validation.keyData.totalUsageLimit || 0)
if (totalUsageLimit > 0) {
const totalCost = Number(validation.keyData.totalCost || 0)
if (totalCost >= totalUsageLimit) {
logger.security(
`📉 Total usage limit exceeded for key: ${validation.keyData.id} (${validation.keyData.name}), cost: $${totalCost.toFixed(2)}/$${totalUsageLimit.toFixed(2)}`
)
return res.status(429).json({
error: 'Total usage limit exceeded',
message: `已达到总额度限制 ($${totalUsageLimit.toFixed(2)})`,
currentCost: totalCost,
costLimit: totalUsageLimit
})
}
logger.api(
`📉 Total usage for key: ${validation.keyData.id} (${validation.keyData.name}), cost: $${totalCost.toFixed(2)}/$${totalUsageLimit.toFixed(2)}`
)
}
// 检查每日费用限制
const dailyCostLimit = validation.keyData.dailyCostLimit || 0
if (dailyCostLimit > 0) {
@@ -333,6 +356,29 @@ const authenticateApiKey = async (req, res, next) => {
)
}
// 检查总费用限制
const totalCostLimit = validation.keyData.totalCostLimit || 0
if (totalCostLimit > 0) {
const totalCost = validation.keyData.totalCost || 0
if (totalCost >= totalCostLimit) {
logger.security(
`💰 Total cost limit exceeded for key: ${validation.keyData.id} (${validation.keyData.name}), cost: $${totalCost.toFixed(2)}/$${totalCostLimit}`
)
return res.status(429).json({
error: 'Total cost limit exceeded',
message: `已达到总费用限制 ($${totalCostLimit})`,
currentCost: totalCost,
costLimit: totalCostLimit
})
}
logger.api(
`💰 Total cost usage for key: ${validation.keyData.id} (${validation.keyData.name}), current: $${totalCost.toFixed(2)}/$${totalCostLimit}`
)
}
// 检查 Opus 周费用限制（仅对 Opus 模型生效）
const weeklyOpusCostLimit = validation.keyData.weeklyOpusCostLimit || 0
if (weeklyOpusCostLimit > 0) {
@@ -394,6 +440,9 @@ const authenticateApiKey = async (req, res, next) => {
allowedClients: validation.keyData.allowedClients,
dailyCostLimit: validation.keyData.dailyCostLimit,
dailyCost: validation.keyData.dailyCost,
totalUsageLimit: validation.keyData.totalUsageLimit,
totalCostLimit: validation.keyData.totalCostLimit,
totalCost: validation.keyData.totalCost,
usage: validation.keyData.usage
}
req.usage = validation.keyData.usage

42
src/models/redis.js
View File

@@ -636,6 +636,48 @@ class RedisClient {
}
}
async addUsageRecord(keyId, record, maxRecords = 200) {
const listKey = `usage:records:${keyId}`
const client = this.getClientSafe()
try {
await client
.multi()
.lpush(listKey, JSON.stringify(record))
.ltrim(listKey, 0, Math.max(0, maxRecords - 1))
.expire(listKey, 86400 * 90) // 默认保留90天
.exec()
} catch (error) {
logger.error(`❌ Failed to append usage record for key ${keyId}:`, error)
}
}
async getUsageRecords(keyId, limit = 50) {
const listKey = `usage:records:${keyId}`
const client = this.getClient()
if (!client) {
return []
}
try {
const rawRecords = await client.lrange(listKey, 0, Math.max(0, limit - 1))
return rawRecords
.map((entry) => {
try {
return JSON.parse(entry)
} catch (error) {
logger.warn('⚠️ Failed to parse usage record entry:', error)
return null
}
})
.filter(Boolean)
} catch (error) {
logger.error(`❌ Failed to load usage records for key ${keyId}:`, error)
return []
}
}
// 💰 获取当日费用
async getDailyCost(keyId) {
const today = getDateStringInTimezone()

61
src/routes/admin.js
View File

@@ -533,6 +533,8 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
enableClientRestriction,
allowedClients,
dailyCostLimit,
totalUsageLimit,
totalCostLimit,
weeklyOpusCostLimit,
tags,
activationDays, // 新增：激活后有效天数
@@ -615,6 +617,35 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
return res.status(400).json({ error: 'All tags must be non-empty strings' })
}
if (
totalUsageLimit !== undefined &&
totalUsageLimit !== null &&
totalUsageLimit !== ''
) {
const usageLimit = Number(totalUsageLimit)
if (Number.isNaN(usageLimit) || usageLimit < 0) {
return res.status(400).json({ error: 'Total usage limit must be a non-negative number' })
}
}
if (
totalCostLimit !== undefined &&
totalCostLimit !== null &&
totalCostLimit !== '' &&
(Number.isNaN(Number(totalCostLimit)) || Number(totalCostLimit) < 0)
) {
return res.status(400).json({ error: 'Total cost limit must be a non-negative number' })
}
if (
totalCostLimit !== undefined &&
totalCostLimit !== null &&
totalCostLimit !== '' &&
(Number.isNaN(Number(totalCostLimit)) || Number(totalCostLimit) < 0)
) {
return res.status(400).json({ error: 'Total cost limit must be a non-negative number' })
}
// 验证激活相关字段
if (expirationMode && !['fixed', 'activation'].includes(expirationMode)) {
return res
@@ -660,6 +691,8 @@ router.post('/api-keys', authenticateAdmin, async (req, res) => {
enableClientRestriction,
allowedClients,
dailyCostLimit,
totalUsageLimit,
totalCostLimit,
weeklyOpusCostLimit,
tags,
activationDays,
@@ -699,6 +732,8 @@ router.post('/api-keys/batch', authenticateAdmin, async (req, res) => {
enableClientRestriction,
allowedClients,
dailyCostLimit,
totalUsageLimit,
totalCostLimit,
weeklyOpusCostLimit,
tags,
activationDays,
@@ -748,6 +783,8 @@ router.post('/api-keys/batch', authenticateAdmin, async (req, res) => {
enableClientRestriction,
allowedClients,
dailyCostLimit,
totalUsageLimit,
totalCostLimit,
weeklyOpusCostLimit,
tags,
activationDays,
@@ -865,6 +902,12 @@ router.put('/api-keys/batch', authenticateAdmin, async (req, res) => {
if (updates.dailyCostLimit !== undefined) {
finalUpdates.dailyCostLimit = updates.dailyCostLimit
}
if (updates.totalUsageLimit !== undefined) {
finalUpdates.totalUsageLimit = updates.totalUsageLimit
}
if (updates.totalCostLimit !== undefined) {
finalUpdates.totalCostLimit = updates.totalCostLimit
}
if (updates.weeklyOpusCostLimit !== undefined) {
finalUpdates.weeklyOpusCostLimit = updates.weeklyOpusCostLimit
}
@@ -993,6 +1036,8 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
allowedClients,
expiresAt,
dailyCostLimit,
totalUsageLimit,
totalCostLimit,
weeklyOpusCostLimit,
tags,
ownerId // 新增：所有者ID字段
@@ -1142,6 +1187,22 @@ router.put('/api-keys/:keyId', authenticateAdmin, async (req, res) => {
updates.dailyCostLimit = costLimit
}
if (totalCostLimit !== undefined && totalCostLimit !== null && totalCostLimit !== '') {
const costLimit = Number(totalCostLimit)
if (isNaN(costLimit) || costLimit < 0) {
return res.status(400).json({ error: 'Total cost limit must be a non-negative number' })
}
updates.totalCostLimit = costLimit
}
if (totalUsageLimit !== undefined && totalUsageLimit !== null && totalUsageLimit !== '') {
const usageLimit = Number(totalUsageLimit)
if (Number.isNaN(usageLimit) || usageLimit < 0) {
return res.status(400).json({ error: 'Total usage limit must be a non-negative number' })
}
updates.totalUsageLimit = usageLimit
}
// 处理 Opus 周费用限制
if (
weeklyOpusCostLimit !== undefined &&

7
src/routes/apiStats.js
View File

@@ -114,6 +114,7 @@ router.post('/api/user-stats', async (req, res) => {
// 获取当日费用统计
const dailyCost = await redis.getDailyCost(keyId)
const costStats = await redis.getCostStats(keyId)
// 处理数据格式，与 validateApiKey 返回的格式保持一致
// 解析限制模型数据
@@ -140,7 +141,10 @@ router.post('/api/user-stats', async (req, res) => {
rateLimitWindow: parseInt(keyData.rateLimitWindow) || 0,
rateLimitRequests: parseInt(keyData.rateLimitRequests) || 0,
dailyCostLimit: parseFloat(keyData.dailyCostLimit) || 0,
totalUsageLimit: parseFloat(keyData.totalUsageLimit) || 0,
totalCostLimit: parseFloat(keyData.totalCostLimit) || 0,
dailyCost: dailyCost || 0,
totalCost: costStats.total || 0,
enableModelRestriction: keyData.enableModelRestriction === 'true',
restrictedModels,
enableClientRestriction: keyData.enableClientRestriction === 'true',
@@ -372,11 +376,14 @@ router.post('/api/user-stats', async (req, res) => {
rateLimitRequests: fullKeyData.rateLimitRequests || 0,
rateLimitCost: parseFloat(fullKeyData.rateLimitCost) || 0, // 新增：费用限制
dailyCostLimit: fullKeyData.dailyCostLimit || 0,
totalUsageLimit: fullKeyData.totalUsageLimit || 0,
totalCostLimit: fullKeyData.totalCostLimit || 0,
// 当前使用量
currentWindowRequests,
currentWindowTokens,
currentWindowCost, // 新增：当前窗口费用
currentDailyCost,
currentTotalCost: totalCost,
// 时间窗口信息
windowStartTime,
windowEndTime,

43
src/routes/userRoutes.js
View File

@@ -258,6 +258,9 @@ router.get('/api-keys', authenticateUser, async (req, res) => {
usage: flatUsage,
dailyCost: key.dailyCost,
dailyCostLimit: key.dailyCostLimit,
totalUsageLimit: key.totalUsageLimit,
totalCost: key.totalCost,
totalCostLimit: key.totalCostLimit,
// 不返回实际的key值，只返回前缀和后几位
keyPreview: key.key
? `${key.key.substring(0, 8)}...${key.key.substring(key.key.length - 4)}`
@@ -287,7 +290,15 @@ router.get('/api-keys', authenticateUser, async (req, res) => {
// 🔑 创建新的API Key
router.post('/api-keys', authenticateUser, async (req, res) => {
try {
const { name, description, tokenLimit, expiresAt, dailyCostLimit } = req.body
const {
name,
description,
tokenLimit,
expiresAt,
dailyCostLimit,
totalUsageLimit,
totalCostLimit
} = req.body
if (!name || !name.trim()) {
return res.status(400).json({
@@ -296,6 +307,32 @@ router.post('/api-keys', authenticateUser, async (req, res) => {
})
}
if (
totalCostLimit !== undefined &&
totalCostLimit !== null &&
totalCostLimit !== '' &&
(Number.isNaN(Number(totalCostLimit)) || Number(totalCostLimit) < 0)
) {
return res.status(400).json({
error: 'Invalid total cost limit',
message: 'Total cost limit must be a non-negative number'
})
}
if (
totalUsageLimit !== undefined &&
totalUsageLimit !== null &&
totalUsageLimit !== ''
) {
const usageLimit = Number(totalUsageLimit)
if (Number.isNaN(usageLimit) || usageLimit < 0) {
return res.status(400).json({
error: 'Invalid total usage limit',
message: 'Total usage limit must be a non-negative number'
})
}
}
// 检查用户API Key数量限制
const userApiKeys = await apiKeyService.getUserApiKeys(req.user.id)
if (userApiKeys.length >= config.userManagement.maxApiKeysPerUser) {
@@ -314,6 +351,8 @@ router.post('/api-keys', authenticateUser, async (req, res) => {
tokenLimit: tokenLimit || null,
expiresAt: expiresAt || null,
dailyCostLimit: dailyCostLimit || null,
totalUsageLimit: totalUsageLimit || null,
totalCostLimit: totalCostLimit || null,
createdBy: 'user',
// 设置服务权限为全部服务，确保前端显示“服务权限”为“全部服务”且具备完整访问权限
permissions: 'all'
@@ -337,6 +376,8 @@ router.post('/api-keys', authenticateUser, async (req, res) => {
tokenLimit: newApiKey.tokenLimit,
expiresAt: newApiKey.expiresAt,
dailyCostLimit: newApiKey.dailyCostLimit,
totalUsageLimit: newApiKey.totalUsageLimit,
totalCostLimit: newApiKey.totalCostLimit,
createdAt: newApiKey.createdAt
}
})

99
src/services/apiKeyService.js
View File

@@ -33,6 +33,8 @@ class ApiKeyService {
enableClientRestriction = false,
allowedClients = [],
dailyCostLimit = 0,
totalUsageLimit = 0,
totalCostLimit = 0,
weeklyOpusCostLimit = 0,
tags = [],
activationDays = 0, // 新增：激活后有效天数（0表示不使用此功能）
@@ -68,6 +70,8 @@ class ApiKeyService {
enableClientRestriction: String(enableClientRestriction || false),
allowedClients: JSON.stringify(allowedClients || []),
dailyCostLimit: String(dailyCostLimit || 0),
totalUsageLimit: String(totalUsageLimit || 0),
totalCostLimit: String(totalCostLimit || 0),
weeklyOpusCostLimit: String(weeklyOpusCostLimit || 0),
tags: JSON.stringify(tags || []),
activationDays: String(activationDays || 0), // 新增：激活后有效天数
@@ -111,6 +115,8 @@ class ApiKeyService {
enableClientRestriction: keyData.enableClientRestriction === 'true',
allowedClients: JSON.parse(keyData.allowedClients || '[]'),
dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
totalUsageLimit: parseFloat(keyData.totalUsageLimit || 0),
totalCostLimit: parseFloat(keyData.totalCostLimit || 0),
weeklyOpusCostLimit: parseFloat(keyData.weeklyOpusCostLimit || 0),
tags: JSON.parse(keyData.tags || '[]'),
activationDays: parseInt(keyData.activationDays || 0),
@@ -188,8 +194,12 @@ class ApiKeyService {
// 获取使用统计（供返回数据使用）
const usage = await redis.getUsageStats(keyData.id)
// 获取当日费用统计
const dailyCost = await redis.getDailyCost(keyData.id)
// 获取费用统计
const [dailyCost, costStats] = await Promise.all([
redis.getDailyCost(keyData.id),
redis.getCostStats(keyData.id)
])
const totalCost = costStats?.total || 0
// 更新最后使用时间（优化：只在实际API调用时更新，而不是验证时）
// 注意：lastUsedAt的更新已移至recordUsage方法中
@@ -245,8 +255,11 @@ class ApiKeyService {
enableClientRestriction: keyData.enableClientRestriction === 'true',
allowedClients,
dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
totalUsageLimit: parseFloat(keyData.totalUsageLimit || 0),
totalCostLimit: parseFloat(keyData.totalCostLimit || 0),
weeklyOpusCostLimit: parseFloat(keyData.weeklyOpusCostLimit || 0),
dailyCost: dailyCost || 0,
totalCost,
weeklyOpusCost: (await redis.getWeeklyOpusCost(keyData.id)) || 0,
tags,
usage
@@ -306,7 +319,10 @@ class ApiKeyService {
}
// 获取当日费用
const dailyCost = (await redis.getDailyCost(keyData.id)) || 0
const [dailyCost, costStats] = await Promise.all([
redis.getDailyCost(keyData.id),
redis.getCostStats(keyData.id)
])
// 获取使用统计
const usage = await redis.getUsageStats(keyData.id)
@@ -365,8 +381,11 @@ class ApiKeyService {
enableClientRestriction: keyData.enableClientRestriction === 'true',
allowedClients,
dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
totalUsageLimit: parseFloat(keyData.totalUsageLimit || 0),
totalCostLimit: parseFloat(keyData.totalCostLimit || 0),
weeklyOpusCostLimit: parseFloat(keyData.weeklyOpusCostLimit || 0),
dailyCost: dailyCost || 0,
totalCost: costStats?.total || 0,
weeklyOpusCost: (await redis.getWeeklyOpusCost(keyData.id)) || 0,
tags,
usage
@@ -405,12 +424,17 @@ class ApiKeyService {
key.rateLimitWindow = parseInt(key.rateLimitWindow || 0)
key.rateLimitRequests = parseInt(key.rateLimitRequests || 0)
key.rateLimitCost = parseFloat(key.rateLimitCost || 0) // 新增：速率限制费用字段
key.totalUsageLimit = parseFloat(key.totalUsageLimit || 0)
if (Number.isNaN(key.totalUsageLimit)) {
key.totalUsageLimit = 0
}
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.totalCostLimit = parseFloat(key.totalCostLimit || 0)
key.weeklyOpusCostLimit = parseFloat(key.weeklyOpusCostLimit || 0)
key.dailyCost = (await redis.getDailyCost(key.id)) || 0
key.weeklyOpusCost = (await redis.getWeeklyOpusCost(key.id)) || 0
@@ -532,6 +556,8 @@ class ApiKeyService {
'enableClientRestriction',
'allowedClients',
'dailyCostLimit',
'totalUsageLimit',
'totalCostLimit',
'weeklyOpusCostLimit',
'tags',
'userId', // 新增：用户ID（所有者变更）
@@ -827,6 +853,21 @@ class ApiKeyService {
}
}
// 记录单次请求的使用详情
const usageCost = costInfo && costInfo.costs ? costInfo.costs.total || 0 : 0
await redis.addUsageRecord(keyId, {
timestamp: new Date().toISOString(),
model,
accountId: accountId || null,
inputTokens,
outputTokens,
cacheCreateTokens,
cacheReadTokens,
totalTokens,
cost: Number(usageCost.toFixed(6)),
costBreakdown: costInfo && costInfo.costs ? costInfo.costs : undefined
})
const logParts = [`Model: ${model}`, `Input: ${inputTokens}`, `Output: ${outputTokens}`]
if (cacheCreateTokens > 0) {
logParts.push(`Cache Create: ${cacheCreateTokens}`)
@@ -973,6 +1014,32 @@ class ApiKeyService {
}
}
const usageRecord = {
timestamp: new Date().toISOString(),
model,
accountId: accountId || null,
accountType: accountType || null,
inputTokens,
outputTokens,
cacheCreateTokens,
cacheReadTokens,
ephemeral5mTokens,
ephemeral1hTokens,
totalTokens,
cost: Number((costInfo.totalCost || 0).toFixed(6)),
costBreakdown: {
input: costInfo.inputCost || 0,
output: costInfo.outputCost || 0,
cacheCreate: costInfo.cacheCreateCost || 0,
cacheRead: costInfo.cacheReadCost || 0,
ephemeral5m: costInfo.ephemeral5mCost || 0,
ephemeral1h: costInfo.ephemeral1hCost || 0
},
isLongContext: costInfo.isLongContextRequest || false
}
await redis.addUsageRecord(keyId, usageRecord)
const logParts = [`Model: ${model}`, `Input: ${inputTokens}`, `Output: ${outputTokens}`]
if (cacheCreateTokens > 0) {
logParts.push(`Cache Create: ${cacheCreateTokens}`)
@@ -1014,8 +1081,24 @@ class ApiKeyService {
}
// 📈 获取使用统计
async getUsageStats(keyId) {
return await redis.getUsageStats(keyId)
async getUsageStats(keyId, options = {}) {
const usageStats = await redis.getUsageStats(keyId)
// options 可能是字符串（兼容旧接口），仅当为对象时才解析
const optionObject =
options && typeof options === 'object' && !Array.isArray(options) ? options : {}
if (optionObject.includeRecords === false) {
return usageStats
}
const recordLimit = optionObject.recordLimit || 20
const recentRecords = await redis.getUsageRecords(keyId, recordLimit)
return {
...usageStats,
recentRecords
}
}
// 📊 获取账户使用统计
@@ -1067,6 +1150,8 @@ class ApiKeyService {
dailyCost,
totalCost: costStats.total,
dailyCostLimit: parseFloat(key.dailyCostLimit || 0),
totalUsageLimit: parseFloat(key.totalUsageLimit || 0),
totalCostLimit: parseFloat(key.totalCostLimit || 0),
userId: key.userId,
userUsername: key.userUsername,
createdBy: key.createdBy,
@@ -1112,7 +1197,9 @@ class ApiKeyService {
userUsername: keyData.userUsername,
createdBy: keyData.createdBy,
permissions: keyData.permissions,
dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0)
dailyCostLimit: parseFloat(keyData.dailyCostLimit || 0),
totalUsageLimit: parseFloat(keyData.totalUsageLimit || 0),
totalCostLimit: parseFloat(keyData.totalCostLimit || 0)
}
} catch (error) {
logger.error('❌ Failed to get API key by ID:', error)