feat: 为console类型账号增加count_tokens端点判断

This commit is contained in:
shaw
2025-12-01 10:14:12 +08:00
parent e17cd1d61b
commit 02018e10f3
2 changed files with 107 additions and 1 deletions

View File

@@ -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') {

View File

@@ -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()