mirror of
https://github.com/Wei-Shaw/claude-relay-service.git
synced 2026-01-22 16:43:35 +00:00
feat: 为console类型账号增加count_tokens端点判断
This commit is contained in:
@@ -972,6 +972,9 @@ router.post('/v1/messages/count_tokens', authenticateApiKey, async (req, res) =>
|
|||||||
const maxAttempts = 2
|
const maxAttempts = 2
|
||||||
let attempt = 0
|
let attempt = 0
|
||||||
|
|
||||||
|
// 引入 claudeConsoleAccountService 用于检查 count_tokens 可用性
|
||||||
|
const claudeConsoleAccountService = require('../services/claudeConsoleAccountService')
|
||||||
|
|
||||||
const processRequest = async () => {
|
const processRequest = async () => {
|
||||||
const { accountId, accountType } = await unifiedClaudeScheduler.selectAccountForApiKey(
|
const { accountId, accountType } = await unifiedClaudeScheduler.selectAccountForApiKey(
|
||||||
req.apiKey,
|
req.apiKey,
|
||||||
@@ -1003,6 +1006,17 @@ router.post('/v1/messages/count_tokens', authenticateApiKey, async (req, res) =>
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 🔍 claude-console 账户特殊处理:检查 count_tokens 端点是否可用
|
||||||
|
if (accountType === 'claude-console') {
|
||||||
|
const isUnavailable = await claudeConsoleAccountService.isCountTokensUnavailable(accountId)
|
||||||
|
if (isUnavailable) {
|
||||||
|
logger.info(
|
||||||
|
`⏭️ count_tokens unavailable for Claude Console account ${accountId}, returning fallback response`
|
||||||
|
)
|
||||||
|
return { fallbackResponse: true }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
const relayOptions = {
|
const relayOptions = {
|
||||||
skipUsageRecord: true,
|
skipUsageRecord: true,
|
||||||
customPath: '/v1/messages/count_tokens'
|
customPath: '/v1/messages/count_tokens'
|
||||||
@@ -1028,6 +1042,23 @@ router.post('/v1/messages/count_tokens', authenticateApiKey, async (req, res) =>
|
|||||||
relayOptions
|
relayOptions
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// 🔍 claude-console 账户:检测上游 404 响应并标记
|
||||||
|
if (accountType === 'claude-console' && response.statusCode === 404) {
|
||||||
|
logger.warn(
|
||||||
|
`⚠️ count_tokens endpoint returned 404 for Claude Console account ${accountId}, marking as unavailable`
|
||||||
|
)
|
||||||
|
// 标记失败不应影响 fallback 响应
|
||||||
|
try {
|
||||||
|
await claudeConsoleAccountService.markCountTokensUnavailable(accountId)
|
||||||
|
} catch (markError) {
|
||||||
|
logger.error(
|
||||||
|
`❌ Failed to mark count_tokens unavailable for account ${accountId}, but will still return fallback:`,
|
||||||
|
markError
|
||||||
|
)
|
||||||
|
}
|
||||||
|
return { fallbackResponse: true }
|
||||||
|
}
|
||||||
|
|
||||||
res.status(response.statusCode)
|
res.status(response.statusCode)
|
||||||
|
|
||||||
const skipHeaders = ['content-encoding', 'transfer-encoding', 'content-length']
|
const skipHeaders = ['content-encoding', 'transfer-encoding', 'content-length']
|
||||||
@@ -1050,11 +1081,21 @@ router.post('/v1/messages/count_tokens', authenticateApiKey, async (req, res) =>
|
|||||||
}
|
}
|
||||||
|
|
||||||
logger.info(`✅ Token count request completed for key: ${req.apiKey.name}`)
|
logger.info(`✅ Token count request completed for key: ${req.apiKey.name}`)
|
||||||
|
return { fallbackResponse: false }
|
||||||
}
|
}
|
||||||
|
|
||||||
while (attempt < maxAttempts) {
|
while (attempt < maxAttempts) {
|
||||||
try {
|
try {
|
||||||
await processRequest()
|
const result = await processRequest()
|
||||||
|
|
||||||
|
// 🔍 处理 fallback 响应(claude-console 账户 count_tokens 不可用)
|
||||||
|
if (result && result.fallbackResponse) {
|
||||||
|
if (!res.headersSent) {
|
||||||
|
return res.status(200).json({ input_tokens: 0 })
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
return
|
return
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
if (error.code === 'CONSOLE_ACCOUNT_CONCURRENCY_FULL') {
|
if (error.code === 'CONSOLE_ACCOUNT_CONCURRENCY_FULL') {
|
||||||
|
|||||||
@@ -1510,6 +1510,71 @@ class ClaudeConsoleAccountService {
|
|||||||
const expiryDate = new Date(account.subscriptionExpiresAt)
|
const expiryDate = new Date(account.subscriptionExpiresAt)
|
||||||
return expiryDate <= new Date()
|
return expiryDate <= new Date()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 🚫 标记账户的 count_tokens 端点不可用
|
||||||
|
async markCountTokensUnavailable(accountId) {
|
||||||
|
try {
|
||||||
|
const client = redis.getClientSafe()
|
||||||
|
const accountKey = `${this.ACCOUNT_KEY_PREFIX}${accountId}`
|
||||||
|
|
||||||
|
// 检查账户是否存在
|
||||||
|
const exists = await client.exists(accountKey)
|
||||||
|
if (!exists) {
|
||||||
|
logger.warn(
|
||||||
|
`⚠️ Cannot mark count_tokens unavailable for non-existent account: ${accountId}`
|
||||||
|
)
|
||||||
|
return { success: false, reason: 'Account not found' }
|
||||||
|
}
|
||||||
|
|
||||||
|
await client.hset(accountKey, {
|
||||||
|
countTokensUnavailable: 'true',
|
||||||
|
countTokensUnavailableAt: new Date().toISOString()
|
||||||
|
})
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
`🚫 Marked count_tokens endpoint as unavailable for Claude Console account: ${accountId}`
|
||||||
|
)
|
||||||
|
return { success: true }
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`❌ Failed to mark count_tokens unavailable for account ${accountId}:`, error)
|
||||||
|
throw error
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ✅ 移除账户的 count_tokens 不可用标记
|
||||||
|
async removeCountTokensUnavailable(accountId) {
|
||||||
|
try {
|
||||||
|
const client = redis.getClientSafe()
|
||||||
|
const accountKey = `${this.ACCOUNT_KEY_PREFIX}${accountId}`
|
||||||
|
|
||||||
|
await client.hdel(accountKey, 'countTokensUnavailable', 'countTokensUnavailableAt')
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
`✅ Removed count_tokens unavailable mark for Claude Console account: ${accountId}`
|
||||||
|
)
|
||||||
|
return { success: true }
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(
|
||||||
|
`❌ Failed to remove count_tokens unavailable mark for account ${accountId}:`,
|
||||||
|
error
|
||||||
|
)
|
||||||
|
throw error
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 🔍 检查账户的 count_tokens 端点是否不可用
|
||||||
|
async isCountTokensUnavailable(accountId) {
|
||||||
|
try {
|
||||||
|
const client = redis.getClientSafe()
|
||||||
|
const accountKey = `${this.ACCOUNT_KEY_PREFIX}${accountId}`
|
||||||
|
|
||||||
|
const value = await client.hget(accountKey, 'countTokensUnavailable')
|
||||||
|
return value === 'true'
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(`❌ Failed to check count_tokens availability for account ${accountId}:`, error)
|
||||||
|
return false // 出错时默认返回可用,避免误阻断
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
module.exports = new ClaudeConsoleAccountService()
|
module.exports = new ClaudeConsoleAccountService()
|
||||||
|
|||||||
Reference in New Issue
Block a user