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: 为 Claude Console 账户添加并发控制机制

Browse Source 
实现了完整的 Claude Console 账户并发任务数控制功能，防止单账户过载，提升服务稳定性。

  **核心功能**

  - 🔒 **原子性并发控制**: 基于 Redis Sorted Set 实现的抢占式并发槽位管理，防止竞态条件
  - 🔄 **自动租约刷新**: 流式请求每 5 分钟自动刷新租约，防止长连接租约过期
  - 🚨 **智能降级处理**: 并发满额时自动清理粘性会话并重试其他账户（最多 1 次）
  - 🎯 **专用错误码**: 引入 `CONSOLE_ACCOUNT_CONCURRENCY_FULL` 错误码，区分并发限制和其他错误
  - 📊 **批量性能优化**: 调度器使用 Promise.all 并行查询账户并发数，减少 Redis 往返

  **后端实现**

  1. **Redis 并发控制方法** (src/models/redis.js)
     - `incrConsoleAccountConcurrency()`: 增加并发计数（带租约）
     - `decrConsoleAccountConcurrency()`: 释放并发槽位
     - `refreshConsoleAccountConcurrencyLease()`: 刷新租约（流式请求）
     - `getConsoleAccountConcurrency()`: 查询当前并发数

  2. **账户服务增强** (src/services/claudeConsoleAccountService.js)
     - 添加 `maxConcurrentTasks` 字段（默认 0 表示无限制）
     - 获取账户时自动查询实时并发数 (`activeTaskCount`)
     - 支持更新并发限制配置

  3. **转发服务并发保护** (src/services/claudeConsoleRelayService.js)
     - 请求前原子性抢占槽位，超限则立即回滚并抛出专用错误
     - 流式请求启动定时器每 5 分钟刷新租约
     - `finally` 块确保槽位释放（即使发生异常）
     - 为每个请求分配唯一 `requestId` 用于并发追踪

  4. **统一调度器优化** (src/services/unifiedClaudeScheduler.js)
     - 获取可用账户时批量查询并发数（Promise.all 并行）
     - 预检查并发限制，避免选择已满的账户
     - 检查分组成员时也验证并发状态
     - 所有账户并发满额时抛出专用错误码

  5. **API 路由降级处理** (src/routes/api.js)
     - 捕获 `CONSOLE_ACCOUNT_CONCURRENCY_FULL` 错误
     - 自动清理粘性会话映射并重试（最多 1 次）
     - 重试失败返回 503 错误和友好提示
     - count_tokens 端点也支持并发满额重试

  6. **管理端点验证** (src/routes/admin.js)
     - 创建/更新账户时验证 `maxConcurrentTasks` 为非负整数
     - 支持前端传入并发限制配置

  **前端实现**

  1. **表单字段** (web/admin-spa/src/components/accounts/AccountForm.vue)
     - 添加"最大并发任务数"输入框（创建和编辑模式）
     - 支持占位符提示"0 表示不限制"
     - 表单数据自动映射到后端 API

  2. **实时监控** (web/admin-spa/src/views/AccountsView.vue)
     - 账户列表显示并发状态进度条和百分比
     - 颜色编码：绿色（<80%）、黄色（80%-100%）、红色（100%）
     - 显示"X / Y"格式的并发数（如"2 / 5"）
     - 未配置限制时显示"并发无限制"徽章
This commit is contained in:
sususu98
2025-10-21 13:43:57 +08:00
parent b61a3103e9
commit 1458d609ca
8 changed files with 706 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

292
src/routes/api.js
View File

@@ -56,6 +56,11 @@ async function handleMessagesRequest(req, res) {
})
}
// 🔄 并发满额重试标志：最多重试一次（使用req对象存储状态）
if (req._concurrencyRetryAttempted === undefined) {
req._concurrencyRetryAttempted = false
}
// 严格的输入验证
if (!req.body || typeof req.body !== 'object') {
return res.status(400).json({
@@ -676,9 +681,75 @@ async function handleMessagesRequest(req, res) {
logger.api(`✅ Request completed in ${duration}ms for key: ${req.apiKey.name}`)
return undefined
} catch (error) {
logger.error('❌ Claude relay error:', error.message, {
code: error.code,
stack: error.stack
let handledError = error
// 🔄 并发满额降级处理：捕获CONSOLE_ACCOUNT_CONCURRENCY_FULL错误
if (
handledError.code === 'CONSOLE_ACCOUNT_CONCURRENCY_FULL' &&
!req._concurrencyRetryAttempted
) {
req._concurrencyRetryAttempted = true
logger.warn(
`⚠️ Console account ${handledError.accountId} concurrency full, attempting fallback to other accounts...`
)
// 只有在响应头未发送时才能重试
if (!res.headersSent) {
try {
// 清理粘性会话映射（如果存在）
const sessionHash = sessionHelper.generateSessionHash(req.body)
await unifiedClaudeScheduler.clearSessionMapping(sessionHash)
logger.info('🔄 Session mapping cleared, retrying handleMessagesRequest...')
// 递归重试整个请求处理（会选择新账户）
return await handleMessagesRequest(req, res)
} catch (retryError) {
// 重试失败
if (retryError.code === 'CONSOLE_ACCOUNT_CONCURRENCY_FULL') {
logger.error('❌ All Console accounts reached concurrency limit after retry')
return res.status(503).json({
error: 'service_unavailable',
message:
'All available Claude Console accounts have reached their concurrency limit. Please try again later.'
})
}
// 其他错误继续向下处理
handledError = retryError
}
} else {
// 响应头已发送，无法重试
logger.error('❌ Cannot retry concurrency full error - response headers already sent')
if (!res.destroyed && !res.finished) {
res.end()
}
return undefined
}
}
// 🚫 第二次并发满额错误：已经重试过，直接返回503
if (
handledError.code === 'CONSOLE_ACCOUNT_CONCURRENCY_FULL' &&
req._concurrencyRetryAttempted
) {
logger.error('❌ All Console accounts reached concurrency limit (retry already attempted)')
if (!res.headersSent) {
return res.status(503).json({
error: 'service_unavailable',
message:
'All available Claude Console accounts have reached their concurrency limit. Please try again later.'
})
} else {
if (!res.destroyed && !res.finished) {
res.end()
}
return undefined
}
}
logger.error('❌ Claude relay error:', handledError.message, {
code: handledError.code,
stack: handledError.stack
})
// 确保在任何情况下都能返回有效的JSON响应
@@ -687,23 +758,29 @@ async function handleMessagesRequest(req, res) {
let statusCode = 500
let errorType = 'Relay service error'
if (error.message.includes('Connection reset') || error.message.includes('socket hang up')) {
if (
handledError.message.includes('Connection reset') ||
handledError.message.includes('socket hang up')
) {
statusCode = 502
errorType = 'Upstream connection error'
} else if (error.message.includes('Connection refused')) {
} else if (handledError.message.includes('Connection refused')) {
statusCode = 502
errorType = 'Upstream service unavailable'
} else if (error.message.includes('timeout')) {
} else if (handledError.message.includes('timeout')) {
statusCode = 504
errorType = 'Upstream timeout'
} else if (error.message.includes('resolve') || error.message.includes('ENOTFOUND')) {
} else if (
handledError.message.includes('resolve') ||
handledError.message.includes('ENOTFOUND')
) {
statusCode = 502
errorType = 'Upstream hostname resolution failed'
}
return res.status(statusCode).json({
error: errorType,
message: error.message || 'An unexpected error occurred',
message: handledError.message || 'An unexpected error occurred',
timestamp: new Date().toISOString()
})
} else {
@@ -860,84 +937,85 @@ router.get('/v1/organizations/:org_id/usage', authenticateApiKey, async (req, re
// 🔢 Token计数端点 - count_tokens beta API
router.post('/v1/messages/count_tokens', authenticateApiKey, async (req, res) => {
try {
// 检查权限
if (
req.apiKey.permissions &&
req.apiKey.permissions !== 'all' &&
req.apiKey.permissions !== 'claude'
) {
return res.status(403).json({
error: {
type: 'permission_error',
message: 'This API key does not have permission to access Claude'
}
})
}
// 检查权限
if (
req.apiKey.permissions &&
req.apiKey.permissions !== 'all' &&
req.apiKey.permissions !== 'claude'
) {
return res.status(403).json({
error: {
type: 'permission_error',
message: 'This API key does not have permission to access Claude'
}
})
}
logger.info(`🔢 Processing token count request for key: ${req.apiKey.name}`)
logger.info(`🔢 Processing token count request for key: ${req.apiKey.name}`)
// 生成会话哈希用于sticky会话
const sessionHash = sessionHelper.generateSessionHash(req.body)
const sessionHash = sessionHelper.generateSessionHash(req.body)
const requestedModel = req.body.model
const maxAttempts = 2
let attempt = 0
// 选择可用的Claude账户
const requestedModel = req.body.model
const processRequest = async () => {
const { accountId, accountType } = await unifiedClaudeScheduler.selectAccountForApiKey(
req.apiKey,
sessionHash,
requestedModel
)
let response
if (accountType === 'claude-official') {
// 使用官方Claude账号转发count_tokens请求
response = await claudeRelayService.relayRequest(
req.body,
req.apiKey,
req,
res,
req.headers,
{
skipUsageRecord: true, // 跳过usage记录，这只是计数请求
customPath: '/v1/messages/count_tokens' // 指定count_tokens路径
}
)
} else if (accountType === 'claude-console') {
// 使用Console Claude账号转发count_tokens请求
response = await claudeConsoleRelayService.relayRequest(
req.body,
req.apiKey,
req,
res,
req.headers,
accountId,
{
skipUsageRecord: true, // 跳过usage记录，这只是计数请求
customPath: '/v1/messages/count_tokens' // 指定count_tokens路径
}
)
} else if (accountType === 'ccr') {
// CCR不支持count_tokens
return res.status(501).json({
error: {
type: 'not_supported',
message: 'Token counting is not supported for CCR accounts'
}
})
} else {
// Bedrock不支持count_tokens
return res.status(501).json({
error: {
type: 'not_supported',
message: 'Token counting is not supported for Bedrock accounts'
if (accountType === 'ccr') {
throw Object.assign(new Error('Token counting is not supported for CCR accounts'), {
httpStatus: 501,
errorPayload: {
error: {
type: 'not_supported',
message: 'Token counting is not supported for CCR accounts'
}
}
})
}
// 直接返回响应，不记录token使用量
if (accountType === 'bedrock') {
throw Object.assign(new Error('Token counting is not supported for Bedrock accounts'), {
httpStatus: 501,
errorPayload: {
error: {
type: 'not_supported',
message: 'Token counting is not supported for Bedrock accounts'
}
}
})
}
const relayOptions = {
skipUsageRecord: true,
customPath: '/v1/messages/count_tokens'
}
const response =
accountType === 'claude-official'
? await claudeRelayService.relayRequest(
req.body,
req.apiKey,
req,
res,
req.headers,
relayOptions
)
: await claudeConsoleRelayService.relayRequest(
req.body,
req.apiKey,
req,
res,
req.headers,
accountId,
relayOptions
)
res.status(response.statusCode)
// 设置响应头
const skipHeaders = ['content-encoding', 'transfer-encoding', 'content-length']
Object.keys(response.headers).forEach((key) => {
if (!skipHeaders.includes(key.toLowerCase())) {
@@ -945,10 +1023,8 @@ router.post('/v1/messages/count_tokens', authenticateApiKey, async (req, res) =>
}
})
// 尝试解析并返回JSON响应
try {
const jsonData = JSON.parse(response.body)
// 对于非 2xx 响应，清理供应商特定信息
if (response.statusCode < 200 || response.statusCode >= 300) {
const sanitizedData = sanitizeUpstreamError(jsonData)
res.json(sanitizedData)
@@ -960,14 +1036,70 @@ router.post('/v1/messages/count_tokens', authenticateApiKey, async (req, res) =>
}
logger.info(`✅ Token count request completed for key: ${req.apiKey.name}`)
} catch (error) {
logger.error('❌ Token count error:', error)
res.status(500).json({
error: {
type: 'server_error',
message: 'Failed to count tokens'
}
while (attempt < maxAttempts) {
try {
await processRequest()
return
} catch (error) {
if (error.code === 'CONSOLE_ACCOUNT_CONCURRENCY_FULL') {
logger.warn(
`⚠️ Console account concurrency full during count_tokens (attempt ${attempt + 1}/${maxAttempts})`
)
if (attempt < maxAttempts - 1) {
try {
await unifiedClaudeScheduler.clearSessionMapping(sessionHash)
} catch (clearError) {
logger.error('❌ Failed to clear session mapping for count_tokens retry:', clearError)
if (!res.headersSent) {
return res.status(500).json({
error: {
type: 'server_error',
message: 'Failed to count tokens'
}
})
}
if (!res.destroyed && !res.finished) {
res.end()
}
return
}
attempt += 1
continue
}
if (!res.headersSent) {
return res.status(503).json({
error: 'service_unavailable',
message:
'All available Claude Console accounts have reached their concurrency limit. Please try again later.'
})
}
if (!res.destroyed && !res.finished) {
res.end()
}
return
}
})
if (error.httpStatus) {
return res.status(error.httpStatus).json(error.errorPayload)
}
logger.error('❌ Token count error:', error)
if (!res.headersSent) {
return res.status(500).json({
error: {
type: 'server_error',
message: 'Failed to count tokens'
}
})
}
if (!res.destroyed && !res.finished) {
res.end()
}
return
}
}
})