Update index.js

index.js

@@ -18,7 +18,7 @@ const CONFIG = {
     RedisToken: process.env.RedisToken
   },
   RETRY: {
-    MAX_ATTEMPTS: 1,
+    MAX_ATTEMPTS: 2,
     DELAY_BASE: 1000
   },
   SERVER: {
@@ -95,35 +95,37 @@ class TokenManager {
   }
 
   async updateCacheTokens() {
-    CONFIG.DEFAULT_HEADERS["anti-csrftoken-a2z"] = Tokens[currentIndex]["anti_csrftoken_a2z"];
-    CONFIG.DEFAULT_HEADERS.Cookie = `idToken=${Tokens[currentIndex]["idToken"]}; pr_refresh_token=${Tokens[currentIndex]["pr_refresh_token"]};aws-waf-token=${Tokens[currentIndex]["aws_waf_token"]}`;
-    CONFIG.DEFAULT_HEADERS.referer = Tokens[currentIndex]["refreshUrl"];
+    CONFIG.DEFAULT_HEADERS["anti-csrftoken-a2z"] = Tokens[currentIndex].anti_csrftoken_a2z;
+    CONFIG.DEFAULT_HEADERS.Cookie = `idToken=${Tokens[currentIndex].idToken}; pr_refresh_token=${Tokens[currentIndex].pr_refresh_token};aws-waf-token=${Tokens[currentIndex].aws_waf_token}`;
+    CONFIG.DEFAULT_HEADERS.referer = Tokens[currentIndex].refreshUrl;
   }
 
   async updateTokens(response, isWaf = false) {
     if (isWaf) {
       console.log("更新waf");
-      if (await Utils.extractWaf()) {
-        Tokens[currentIndex]["aws_waf_token"] = await Utils.extractWaf();
+      var wafToken = await Utils.extractWaf();
+      if (wafToken) {
+        Tokens[currentIndex].aws_waf_token = wafToken;
         return;
       } else {
         console.log("提取aws-waf-token失败");
         currentIndex = (currentIndex + 1) % Tokens.length;
       }
-    }
-    const newCsrfToken = response.headers.get('anti-csrftoken-a2z');
-    const cookies = response.headers.get('set-cookie');
-    if (newCsrfToken && cookies) {
-      console.log("更新缓存");
-      Tokens[currentIndex]["anti_csrftoken_a2z"] = newCsrfToken;
-      const idTokenMatch = cookies.match(/idToken=([^;]+)/);
-      if (idTokenMatch && idTokenMatch[1]) {
-        Tokens[currentIndex]["idToken"] = idTokenMatch[1];
+    }else{
+      const newCsrfToken = response.headers.get('anti-csrftoken-a2z');
+      const cookies = response.headers.get('set-cookie');
+      if (newCsrfToken && cookies) {
+        console.log("更新缓存");
+        Tokens[currentIndex].anti_csrftoken_a2z = newCsrfToken;
+        const idTokenMatch = cookies.match(/idToken=([^;]+)/);
+        if (idTokenMatch && idTokenMatch[1]) {
+          Tokens[currentIndex].idToken = idTokenMatch[1];
+        }
+        await this.updateRedisTokens();//最后更新redis数据库缓存
+        console.log("更新缓存完毕");
       }
-      await this.updateRedisTokens();//最后更新redis数据库缓存
-      console.log("更新缓存完毕");
+      currentIndex = (currentIndex + 1) % Tokens.length;
     }
-    currentIndex = (currentIndex + 1) % Tokens.length;
   }
 }
 
@@ -143,10 +145,10 @@ class Utils {
     try {
       const page = await browser.newPage();
       await page.setExtraHTTPHeaders({
-        cookie: `pr_refresh_token=${Tokens[currentIndex]["pr_refresh_token"]};aws-waf-token=${Tokens[currentIndex]["aws_waf_token"]};idToken=${Tokens[currentIndex]["idToken"]}`
+        cookie: `pr_refresh_token=${Tokens[currentIndex].pr_refresh_token};idToken=${Tokens[currentIndex].idToken}`
       });
 
-      await page.goto(Tokens[currentIndex]["refreshUrl"], { waitUntil: 'networkidle0' });
+      await page.goto(Tokens[currentIndex].refreshUrl, { waitUntil: 'networkidle0' });
 
       // 直接从页面 cookies 中提取 aws-waf-token
       const awsWafToken = (await page.cookies()).find(
@@ -154,7 +156,6 @@ class Utils {
       )?.value;
 
       if (awsWafToken) {
-        Tokens[currentIndex]["aws_waf_token"] = awsWafToken;
         console.log("成功提取 aws-waf-token");
         await browser.close();
         return awsWafToken;
@@ -264,8 +265,6 @@ async function initializeService() {
 
 await initializeService();
 
-
-
 class ApiClient {
   constructor(modelId) {
     console.log(modelId);
@@ -497,8 +496,6 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
       return res.status(401).json({ error: "Unauthorized" });
     }
     await tokenManager.updateCacheTokens();
-
-
     const apiClient = new ApiClient(req.body.model);
     const requestPayload = await apiClient.transformMessages(req.body);
     let retryCount = 0;
@@ -514,23 +511,30 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
           },
           body: JSON.stringify(requestPayload)
         });
-        if (response.status == 200) {
-          console.log("请求成功");
-          break; // 如果请求成功，跳出重试循环
-        }
-        if (response.status != 200) {
-          reqStatus = response.status;
-          if (response.status == 202) {
-            console.log("请求失败，已达到速率，重新获取waf然后更新token库");
+        reqStatus = response.status;
+        switch (reqStatus) {
+          case 200:
+            console.log("请求成功");
+            // 异步更新token
+            tokenManager.updateTokens(response)
+            // 处理响应
+            if (req.body.stream) {
+              await ResponseHandler.handleStreamResponse(response, req.body.model, res);
+            } else {
+              await ResponseHandler.handleNormalResponse(response, req.body.model, res);
+            }
+            return;
+          case 202:
+            console.log("请求受限，更新WAF");
             await tokenManager.updateTokens(response, true);
-          }
+            break;
+          default:
+            throw new Error(`请求失败! status: ${response.status}`);
         }
         retryCount++;
         if (retryCount >= CONFIG.RETRY.MAX_ATTEMPTS) {
-          reqStatus = response.status;
           throw new Error(`上游服务请求失败! status: ${response.status}`);
         }
-        // 等待一段时间后重试
         await new Promise(resolve => setTimeout(resolve, CONFIG.RETRY.DELAY_BASE * retryCount));
 
       } catch (error) {
@@ -538,17 +542,9 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
         if (retryCount >= CONFIG.RETRY.MAX_ATTEMPTS) {
           throw error;
         }
-        // 等待一段时间后重试
         await new Promise(resolve => setTimeout(resolve, CONFIG.RETRY.DELAY_BASE * retryCount));
       }
     }
-    if (req.body.stream) {
-      await ResponseHandler.handleStreamResponse(response, req.body.model, res);
-    } else {
-      await ResponseHandler.handleNormalResponse(response, req.body.model, res);
-    }
-    //请求完毕后再尝试更新redis数据库缓存，避免阻塞
-    await tokenManager.updateTokens(response);
   } catch (error) {
     res.status(reqStatus).json({
       error: {