From 02018e10f38a9be9e7ac28f7bae0af124d3a6f0d Mon Sep 17 00:00:00 2001
From: shaw <shaw-wei@foxmail.com>
Date: Mon, 1 Dec 2025 10:14:12 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E4=B8=BAconsole=E7=B1=BB=E5=9E=8B?=
 =?UTF-8?q?=E8=B4=A6=E5=8F=B7=E5=A2=9E=E5=8A=A0count=5Ftokens=E7=AB=AF?=
 =?UTF-8?q?=E7=82=B9=E5=88=A4=E6=96=AD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/routes/api.js                           | 43 +++++++++++++-
 src/services/claudeConsoleAccountService.js | 65 +++++++++++++++++++++
 2 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/src/routes/api.js b/src/routes/api.js
index e5e7e22a..ce31440d 100644
--- a/src/routes/api.js
+++ b/src/routes/api.js
@@ -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') {
diff --git a/src/services/claudeConsoleAccountService.js b/src/services/claudeConsoleAccountService.js
index 9121ee55..3dee9c90 100644
--- a/src/services/claudeConsoleAccountService.js
+++ b/src/services/claudeConsoleAccountService.js
@@ -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()