mirror of
https://github.com/Wei-Shaw/claude-relay-service.git
synced 2026-01-22 16:40:25 +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
|
||||
let attempt = 0
|
||||
|
||||
// 引入 claudeConsoleAccountService 用于检查 count_tokens 可用性
|
||||
const claudeConsoleAccountService = require('../services/claudeConsoleAccountService')
|
||||
|
||||
const processRequest = async () => {
|
||||
const { accountId, accountType } = await unifiedClaudeScheduler.selectAccountForApiKey(
|
||||
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 = {
|
||||
skipUsageRecord: true,
|
||||
customPath: '/v1/messages/count_tokens'
|
||||
@@ -1028,6 +1042,23 @@ router.post('/v1/messages/count_tokens', authenticateApiKey, async (req, res) =>
|
||||
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)
|
||||
|
||||
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}`)
|
||||
return { fallbackResponse: false }
|
||||
}
|
||||
|
||||
while (attempt < maxAttempts) {
|
||||
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
|
||||
} catch (error) {
|
||||
if (error.code === 'CONSOLE_ACCOUNT_CONCURRENCY_FULL') {
|
||||
|
||||
@@ -1510,6 +1510,71 @@ class ClaudeConsoleAccountService {
|
||||
const expiryDate = new Date(account.subscriptionExpiresAt)
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user