Update index.js

index.js

@@ -3,22 +3,29 @@ import fetch from 'node-fetch';
 import FormData from 'form-data';
 import dotenv from 'dotenv';
 import cors from 'cors';
-import puppeteer from 'puppeteer';
+import puppeteer from 'puppeteer-extra'
+import StealthPlugin from 'puppeteer-extra-plugin-stealth'
 import { v4 as uuidv4 } from 'uuid';
+import Logger from './logger.js';
 
 dotenv.config();
+
 // 配置常量
 const CONFIG = {
     MODELS: {
-        'grok-latest': 'grok-latest',
-        'grok-latest-image': 'grok-latest',
-        'grok-latest-search': 'grok-latest'
+        'grok-2': 'grok-latest',
+        'grok-2-imageGen': 'grok-latest',
+        'grok-2-search': 'grok-latest',
+        "grok-3": "grok-3",
+        "grok-3-imageGen": "grok-3",
+        "grok-3-deepsearch": "grok-3",
+        "grok-3-reasoning": "grok-3"
     },
     API: {
         BASE_URL: "https://grok.com",
         API_KEY: process.env.API_KEY || "sk-123456",
-        SSO_TOKEN: null,//登录时才有的认证cookie,这里暂时用不到，之后可能需要
         SIGNATURE_COOKIE: null,
+        TEMP_COOKIE: null,
         PICGO_KEY: process.env.PICGO_KEY || null //想要生图的话需要填入这个PICGO图床的key
     },
     SERVER: {
@@ -26,13 +33,16 @@ const CONFIG = {
         BODY_LIMIT: '5mb'
     },
     RETRY: {
-        MAX_ATTEMPTS: 1,//重试次数
-        DELAY_BASE: 1000 // 基础延迟时间（毫秒）
+        MAX_ATTEMPTS: 2//重试次数
     },
+    IS_THINKING: false,
+    IS_IMG_GEN: false,
+    IS_IMG_GEN2: false,
+    SSO_INDEX: 0,//sso的索引
     ISSHOW_SEARCH_RESULTS: process.env.ISSHOW_SEARCH_RESULTS === 'true',//是否显示搜索结果,默认关闭
     CHROME_PATH: process.env.CHROME_PATH || "/usr/bin/chromium"//chrome路径
 };
-
+puppeteer.use(StealthPlugin())
 // 请求头配置
 const DEFAULT_HEADERS = {
     'accept': '*/*',
@@ -52,9 +62,133 @@ const DEFAULT_HEADERS = {
     'baggage': 'sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c'
 };
 
+
+async function initialization() {
+    const ssoArray = process.env.SSO.split(',');
+    const ssorwArray = process.env.SSO_RW.split(',');
+    ssoArray.forEach((sso, index) => {
+        tokenManager.addToken(`sso-rw=${ssorwArray[index]};sso=${sso}`);
+    });
+    console.log(JSON.stringify(tokenManager.getActiveTokens(), null, 2));
+    await Utils.get_signature()
+    Logger.info("初始化完成", 'Server');
+}
+
+
+class AuthTokenManager {
+    constructor() {
+        this.activeTokens = [];
+        this.expiredTokens = new Map();
+        this.tokenModelFrequency = new Map();
+        this.modelRateLimit = {
+            "grok-3": { RequestFrequency: 20 },
+            "grok-3-deepsearch": { RequestFrequency: 5 },
+            "grok-3-reasoning": { RequestFrequency: 5 }
+        };
+    }
+
+    addToken(token) {
+        if (!this.activeTokens.includes(token)) {
+            this.activeTokens.push(token);
+            this.tokenModelFrequency.set(token, {
+                "grok-3": 0,
+                "grok-3-deepsearch": 0,
+                "grok-3-reasoning": 0
+            });
+        }
+    }
+
+    getTokenByIndex(index, model) {
+        if (!this.modelRateLimit[model]) return;
+        if(this.activeTokens.length === 0){
+            return null;
+        }
+        const token = this.activeTokens[index];
+        this.recordModelRequest(token, model);
+        return token;
+    }
+
+    recordModelRequest(token, model) {
+        if (!this.modelRateLimit[model]) return;
+        const tokenFrequency = this.tokenModelFrequency.get(token);
+        if (tokenFrequency && tokenFrequency[model] !== undefined) {
+            tokenFrequency[model]++;
+        }
+        this.checkAndRemoveTokenIfLimitReached(token);
+    }
+    setModelLimit(index, model) {
+        if (!this.modelRateLimit[model]) return;
+        const tokenFrequency = this.tokenModelFrequency.get(this.activeTokens[index]);
+        tokenFrequency[model] = 9999;
+    }
+    isTokenModelLimitReached(index, model) {
+        if (!this.modelRateLimit[model]) return;
+        const token = this.activeTokens[index];
+        const tokenFrequency = this.tokenModelFrequency.get(token);
+
+        if (!tokenFrequency) {
+            return false;
+        }
+        return tokenFrequency[model] >= this.modelRateLimit[model].RequestFrequency;
+    }
+    checkAndRemoveTokenIfLimitReached(token) {
+        const tokenFrequency = this.tokenModelFrequency.get(token);
+        if (!tokenFrequency) return;
+
+        const isLimitReached = Object.keys(tokenFrequency).every(model =>
+            tokenFrequency[model] >= this.modelRateLimit[model].RequestFrequency
+        );
+
+        if (isLimitReached) {
+            const tokenIndex = this.activeTokens.indexOf(token);
+            if (tokenIndex !== -1) {
+                this.removeTokenByIndex(tokenIndex);
+            }
+        }
+    }
+
+    removeTokenByIndex(index) {
+        if (!this.isRecoveryProcess) {
+            this.startTokenRecoveryProcess();
+        }
+        const token = this.activeTokens[index];
+        this.expiredTokens.set(token, Date.now());
+        this.activeTokens.splice(index, 1);
+        this.tokenModelFrequency.delete(token);
+        Logger.info(`令牌${token}已达到上限，已移除`, 'TokenManager');
+    }
+
+    startTokenRecoveryProcess() {
+        setInterval(() => {
+            const now = Date.now();
+            for (const [token, expiredTime] of this.expiredTokens.entries()) {
+                if (now - expiredTime >= 2 * 60 * 60 * 1000) {
+                    this.tokenModelUsage.set(token, {
+                        "grok-3": 0,
+                        "grok-3-imageGen": 0,
+                        "grok-3-deepsearch": 0,
+                        "grok-3-reasoning": 0
+                    });
+                    this.activeTokens.push(token);
+                    this.expiredTokens.delete(token);
+                    Logger.info(`令牌${token}已恢复，已添加到可用令牌列表`, 'TokenManager');
+                }
+            }
+        }, 2 * 60 * 60 * 1000);
+    }
+
+    getTokenCount() {
+        return this.activeTokens.length || 0;
+    }
+
+    getActiveTokens() {
+        return [...this.activeTokens];
+    }
+}
+
 class Utils {
     static async extractGrokHeaders() {
-        console.log("开始提取头信息");
+        Logger.info("开始提取头信息", 'Server');
         try {
             // 启动浏览器
             const browser = await puppeteer.launch({
@@ -69,8 +203,10 @@ class Utils {
             });
 
             const page = await browser.newPage();
-            await page.goto('https://grok.com/', { waitUntil: 'networkidle0' });
-
+            await page.goto('https://grok.com/', { waitUntil: 'domcontentloaded' });
+            await page.evaluate(() => {
+                return new Promise(resolve => setTimeout(resolve, 5000))
+            })
             // 获取所有 Cookies
             const cookies = await page.cookies();
             const targetHeaders = ['x-anonuserid', 'x-challenge', 'x-signature'];
@@ -85,80 +221,32 @@ class Utils {
             // 关闭浏览器
             await browser.close();
             // 打印并返回提取的头信息
-            console.log('提取的头信息:', JSON.stringify(extractedHeaders, null, 2));
+            Logger.info('提取的头信息:', JSON.stringify(extractedHeaders, null, 2), 'Server');
             return extractedHeaders;
 
         } catch (error) {
-            console.error('获取头信息出错:', error);
+            Logger.error('获取头信息出错:', error, 'Server');
             return null;
         }
     }
     static async get_signature() {
-        if (CONFIG.API.SIGNATURE_COOKIE) {
-            return CONFIG.API.SIGNATURE_COOKIE;
+        if (CONFIG.API.TEMP_COOKIE) {
+            return CONFIG.API.TEMP_COOKIE;
         }
-        console.log("刷新认证信息");
+        Logger.info("刷新认证信息", 'Server');
         let retryCount = 0;
-        while (retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
+        while (!CONFIG.API.TEMP_COOKIE || retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
             let headers = await Utils.extractGrokHeaders();
             if (headers) {
-                console.log("获取认证信息成功");
-                CONFIG.API.SIGNATURE_COOKIE = { cookie: `x-anonuserid=${headers["x-anonuserid"]}; x-challenge=${headers["x-challenge"]}; x-signature=${headers["x-signature"]}` };
-                return CONFIG.API.SIGNATURE_COOKIE;
+                Logger.info("获取认证信息成功", 'Server');
+                CONFIG.API.TEMP_COOKIE = { cookie: `x-anonuserid=${headers["x-anonuserid"]}; x-challenge=${headers["x-challenge"]}; x-signature=${headers["x-signature"]}` };
+                return;
             }
             retryCount++;
             if (retryCount >= CONFIG.RETRY.MAX_ATTEMPTS) {
                 throw new Error(`获取认证信息失败!`);
             }
-            await new Promise(resolve => setTimeout(resolve, CONFIG.RETRY.DELAY_BASE * retryCount));
-        }
-    }
-    static async handleError(error, res) {
-        // 如果是500错误且提供了原始请求函数，尝试重新获取签名并重试
-        if (String(error).includes("status: 500")) {
-            try {
-                await Utils.get_signature();
-                if (CONFIG.API.SIGNATURE_COOKIE) {
-                    return res.status(500).json({
-                        error: {
-                            message: `${retryError.message}已重新刷新认证信息，请重新对话`,
-                            type: 'server_error',
-                            param: null,
-                            code: error.code || null
-                        }
-                    });
-                } else {
-                    return res.status(500).json({
-                        error: {
-                            message: "认证信息获取失败",
-                            type: 'server_error',
-                            param: null,
-                            code: null
-                        }
-                    });
-                }
-            } catch (retryError) {
-                console.error('重试失败:', retryError);
-                return res.status(500).json({
-                    error: {
-                        message: `${retryError.message},认证信息获取失败`,
-                        type: 'server_error',
-                        param: null,
-                        code: retryError.code || null
-                    }
-                });
-            }
         }
-
-        // 其他错误直接返回
-        res.status(500).json({
-            error: {
-                message: error.message,
-                type: 'server_error',
-                param: null,
-                code: error.code || null
-            }
-        });
     }
     static async organizeSearchResults(searchResults) {
         // 确保传入的是有效的搜索结果对象
@@ -177,6 +265,11 @@ class Utils {
         });
         return formattedResults.join('\n\n');
     }
+    static createAuthHeaders(model) {
+        return {
+            'cookie': `${tokenManager.getTokenByIndex(CONFIG.SSO_INDEX, model)}`
+        };
+    }
 }
 
 
@@ -228,7 +321,7 @@ class GrokApiClient {
                     content: imageBuffer
                 }
             };
-            console.log("发送图片请求");
+            Logger.info("发送图片请求", 'Server');
             // 发送请求
             const response = await fetch(url, {
                 method: 'POST',
@@ -240,29 +333,26 @@ class GrokApiClient {
             });
 
             if (!response.ok) {
-                console.error(`上传图片失败,状态码:${response.status},原因:${response.error}`);
+                Logger.error(`上传图片失败,状态码:${response.status},原因:${response.error}`, 'Server');
                 return '';
             }
 
             const result = await response.json();
-            console.log('上传图片成功:', result);
+            Logger.info('上传图片成功:', result, 'Server');
             return result.fileMetadataId;
 
         } catch (error) {
-            console.error('上传图片失败:', error);
+            Logger.error(error, 'Server');
             return '';
         }
     }
 
     async prepareChatRequest(request) {
-
-        if (request.model === 'grok-latest-image' && !CONFIG.API.PICGO_KEY) {
+        if ((request.model === 'grok-2-imageGen' || request.model === 'grok-3-imageGen') && !CONFIG.API.PICGO_KEY) {
             throw new Error(`该模型需要配置PICGO图床密钥!`);
         }
         var todoMessages = request.messages;
-        if (request.model === 'grok-latest-image' || request.model === 'grok-latest-search') {
-            todoMessages = Array.isArray(todoMessages) ? [todoMessages[todoMessages.length - 1]] : [todoMessages];;
-        }
+
         let fileAttachments = [];
         let messages = '';
         let lastRole = null;
@@ -283,52 +373,43 @@ class GrokApiClient {
         for (const current of todoMessages) {
             const role = current.role === 'assistant' ? 'assistant' : 'user';
             let textContent = '';
-            // 处理消息内容
+
             if (Array.isArray(current.content)) {
-                // 处理数组内的所有内容
                 for (const item of current.content) {
                     if (item.type === 'image_url') {
-                        // 如果是图片且是最后一条消息，则处理图片
                         if (current === todoMessages[todoMessages.length - 1]) {
                             const processedImage = await processImageUrl(item);
                             if (processedImage) fileAttachments.push(processedImage);
                         }
-                        textContent += (textContent ? '\n' : '') + "[图片]";//图片占位符
+                        textContent += (textContent ? '\n' : '') + "[图片]";
                     } else if (item.type === 'text') {
                         textContent += (textContent ? '\n' : '') + item.text;
                     }
                 }
             } else if (typeof current.content === 'object' && current.content !== null) {
-                // 处理单个对象内容
                 if (current.content.type === 'image_url') {
-                    // 如果是图片且是最后一条消息，则处理图片
                     if (current === todoMessages[todoMessages.length - 1]) {
                         const processedImage = await processImageUrl(current.content);
                         if (processedImage) fileAttachments.push(processedImage);
                     }
-                    textContent += (textContent ? '\n' : '') + "[图片]";//图片占位符
+                    textContent += (textContent ? '\n' : '') + "[图片]";
                 } else if (current.content.type === 'text') {
                     textContent = current.content.text;
                 }
             } else {
-                // 处理普通文本内容
                 textContent = this.processMessageContent(current.content);
             }
-            // 添加文本内容��消息字符串
             if (textContent) {
                 if (role === lastRole) {
-                    // 如果角色相同，合并消息内容
                     lastContent += '\n' + textContent;
                     messages = messages.substring(0, messages.lastIndexOf(`${role.toUpperCase()}: `)) +
                         `${role.toUpperCase()}: ${lastContent}\n`;
                 } else {
-                    // 如果角色不同，添加新的消息
                     messages += `${role.toUpperCase()}: ${textContent}\n`;
                     lastContent = textContent;
                     lastRole = role;
                 }
             } else if (current === todoMessages[todoMessages.length - 1] && fileAttachments.length > 0) {
-                // 如果是最后一条消息且有图片附件，添加空消息占位
                 messages += `${role.toUpperCase()}: [图片]\n`;
             }
         }
@@ -339,27 +420,35 @@ class GrokApiClient {
 
         messages = messages.trim();
 
-        if (request.model === 'grok-latest-search') {
+        if (request.model === 'grok-2-search') {
             search = true;
         }
         return {
-            message: messages,
             modelName: this.modelId,
-            disableSearch: false,
+            message: messages,
+            fileAttachments: fileAttachments,
             imageAttachments: [],
+            disableSearch: false,
+            enableImageGeneration: true,
             returnImageBytes: false,
             returnRawGrokInXaiRequest: false,
-            fileAttachments: fileAttachments,
             enableImageStreaming: false,
             imageGenerationCount: 1,
+            forceConcise: false,
             toolOverrides: {
-                imageGen: request.model === 'grok-latest-image',
+                imageGen: request.model === 'grok-2-imageGen' || request.model === 'grok-3-imageGen',
                 webSearch: search,
                 xSearch: search,
                 xMediaSearch: search,
                 trendsSearch: search,
                 xPostAnalyze: search
-            }
+            },
+            enableSideBySide: true,
+            isPreset: false,
+            sendFinalMetadata: true,
+            customInstructions: "",
+            deepsearchPreset: request.model === 'grok-3-deepsearch' ? "default" : "",
+            isReasoning: request.model === 'grok-3-reasoning'
         };
     }
 }
@@ -400,247 +489,146 @@ class MessageProcessor {
         };
     }
 }
-
-// 中间件配置
-const app = express();
-app.use(express.json({ limit: CONFIG.SERVER.BODY_LIMIT }));
-app.use(express.urlencoded({ extended: true, limit: CONFIG.SERVER.BODY_LIMIT }));
-app.use(cors({
-    origin: '*',
-    methods: ['GET', 'POST', 'OPTIONS'],
-    allowedHeaders: ['Content-Type', 'Authorization']
-}));
-// API路由
-app.get('/hf/v1/models', (req, res) => {
-    res.json({
-        object: "list",
-        data: Object.keys(CONFIG.MODELS).map((model, index) => ({
-            id: model,
-            object: "model",
-            created: Math.floor(Date.now() / 1000),
-            owned_by: "xai",
-        }))
-    });
-});
-
-app.post('/hf/v1/chat/completions', async (req, res) => {
-    const authToken = req.headers.authorization?.replace('Bearer ', '');
-    if (authToken !== CONFIG.API.API_KEY) {
-        return res.status(401).json({ error: 'Unauthorized' });
-    }
-    const makeRequest = async () => {
-        if (!CONFIG.API.SIGNATURE_COOKIE) {
-            await Utils.get_signature();
-        }
-        const grokClient = new GrokApiClient(req.body.model);
-        const requestPayload = await grokClient.prepareChatRequest(req.body);
-        //创建新对话        
-        const newMessageReq = await fetch(`${CONFIG.API.BASE_URL}/api/rpc`, {
-            method: 'POST',
-            headers: {
-                ...DEFAULT_HEADERS,
-                ...CONFIG.API.SIGNATURE_COOKIE
-            },
-            body: JSON.stringify({
-                rpc: "createConversation",
-                req: {
-                    temporary: false
-                }
-            })
-        });
-        if (!newMessageReq.ok) {
-            throw new Error(`上游服务请求失败! status: ${newMessageReq.status}`);
-        }
-
-        // 获取响应文本
-        const responseText = await newMessageReq.json();
-        const conversationId = responseText.conversationId;
-        console.log("会话ID:conversationId", conversationId);
-        if (!conversationId) {
-            throw new Error(`创建会话失败! status: ${newMessageReq.status}`);
-        }
-        //发送对话
-        const response = await fetch(`${CONFIG.API.BASE_URL}/api/conversations/${conversationId}/responses`, {
-            method: 'POST',
-            headers: {
-                "accept": "text/event-stream",
-                "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
-                "content-type": "text/plain;charset=UTF-8",
-                "Connection": "keep-alive",
-                ...CONFIG.API.SIGNATURE_COOKIE
-            },
-            body: JSON.stringify(requestPayload)
-        });
-
-        if (!response.ok) {
-            throw new Error(`上游服务请求失败! status: ${response.status}`);
-        }
-        if (req.body.stream) {
-            await handleStreamResponse(response, req.body.model, res);
-        } else {
-            await handleNormalResponse(response, req.body.model, res);
+async function processModelResponse(linejosn, model) {
+    let result = { token: '', imageUrl: null }
+    if (CONFIG.IS_IMG_GEN) {
+        if (linejosn?.cachedImageGenerationResponse && !CONFIG.IS_IMG_GEN2) {
+            result.imageUrl = linejosn.cachedImageGenerationResponse.imageUrl;
         }
+        return result;
     }
-    try {
-        await makeRequest();
-    } catch (error) {
-        CONFIG.API.SIGNATURE_COOKIE = null;
-        await Utils.handleError(error, res);
+    var token = linejosn?.token
+    if (!token) return result;
+    //非生图模型的处理
+    switch (model) {
+        case 'grok-2':
+            result.token = token;
+            return result;
+        case 'grok-2-search':
+            if (linejosn?.webSearchResults && CONFIG.ISSHOW_SEARCH_RESULTS) {
+                result.token = `\r\n<think>${await Utils.organizeSearchResults(linejosn.webSearchResults)}</think>\r\n`;
+            } else {
+                result.token = token;
+            }
+            return result;
+        case 'grok-3':
+            result.token = token;
+            return result;
+        case 'grok-3-deepsearch':
+            if (linejosn.messageTag === "final") {
+                result.token = token;
+            }
+            return result;
+        case 'grok-3-reasoning':
+            if (linejosn?.isThinking && !CONFIG.IS_THINKING) {
+                result.token = "<think>" + token;
+                CONFIG.IS_THINKING = true;
+            } else if (CONFIG.IS_THINKING && !linejosn.isThinking) {
+                result.token = "</think>" + token;
+                CONFIG.IS_THINKING = false;
+            } else {
+                result.token = token;
+            }
+            return result;
     }
-});
-
-async function handleStreamResponse(response, model, res) {
-    res.setHeader('Content-Type', 'text/event-stream');
-    res.setHeader('Cache-Control', 'no-cache');
-    res.setHeader('Connection', 'keep-alive');
+}
 
+async function handleResponse(response, model, res, isStream) {
     try {
         const stream = response.body;
         let buffer = '';
-
-        stream.on('data', async (chunk) => {
-            buffer += chunk.toString();
-            const lines = buffer.split('\n');
-            buffer = lines.pop() || '';
-
-            for (const line of lines) {
-                if (!line.trim()) continue;
-                const trimmedLine = line.trim();
-                if (trimmedLine.startsWith('data: ')) {
-                    const data = trimmedLine.substring(6);
-                    if(data === "[DONE]"){
-                        console.log("流结束");
-                        res.write('data: [DONE]\n\n');
-                        return res.end();
-                    };
-                    try {
-                        if(!data.trim())continue;
-                        const linejosn = JSON.parse(data);
-                        if (linejosn?.error?.name === "RateLimitError") {
-                            var responseData = MessageProcessor.createChatResponse(`${linejosn.error.name},请重新对话`, model, true);
-                            CONFIG.API.SIGNATURE_COOKIE = null;
-                            console.log("认证信息已删除");
-                            res.write(`data: ${JSON.stringify(responseData)}\n\n`);
-                            res.write('data: [DONE]\n\n');
-                            return res.end();
-                        }
-                        switch (model) {
-                            case 'grok-latest-image':
-                                if (linejosn.response === "modelResponse" && linejosn?.modelResponse?.generatedImageUrls) {
-                                    const dataImage = await handleImageResponse(linejosn.modelResponse.generatedImageUrls);
-                                    const responseData = MessageProcessor.createChatResponse(dataImage, model, true);
-                                    res.write(`data: ${JSON.stringify(responseData)}\n\n`);   
-                                }
-                                break;
-                            case 'grok-latest':
-                                if (linejosn.response === "token") {
-                                    const token = linejosn.token;
-                                    if (token && token.length > 0) {
-                                        const responseData = MessageProcessor.createChatResponse(token, model, true);
-                                        res.write(`data: ${JSON.stringify(responseData)}\n\n`);
+        let fullResponse = '';
+        const dataPromises = [];
+
+        return new Promise((resolve, reject) => {
+            stream.on('data', async (chunk) => {
+                buffer += chunk.toString();
+                const lines = buffer.split('\n');
+                buffer = lines.pop() || '';
+
+                for (const line of lines) {
+                    if (!line.trim()) continue;
+                    const trimmedLine = line.trim();
+                    if (trimmedLine.startsWith('data: ')) {
+                        const data = trimmedLine.substring(6);
+                        try {
+                            if (!data.trim()) continue;
+                            if(data === "[DONE]")continue;
+                            const linejosn = JSON.parse(data);
+                            if (linejosn?.error) {
+                                reject(new Error(linejosn));
+                                return;
+                            }
+                            if (linejosn?.doImgGen || linejosn?.imageAttachmentInfo) {
+                                CONFIG.IS_IMG_GEN = true;
+                            }
+                            const processPromise = (async () => {
+                                const result = await processModelResponse(linejosn, model);
+                                if (result.token) {
+                                    if (isStream) {
+                                        res.write(`data: ${JSON.stringify(MessageProcessor.createChatResponse(result.token, model, true))}\n\n`);
+                                    } else {
+                                        fullResponse += result.token;
                                     }
                                 }
-                                break;
-                            case 'grok-latest-search':
-                                if (linejosn.response === "token") {
-                                    const token = linejosn.token;
-                                    if (token && token.length > 0) {
-                                        const responseData = MessageProcessor.createChatResponse(token, model, true);
-                                        res.write(`data: ${JSON.stringify(responseData)}\n\n`);
+                                if (result.imageUrl) {
+                                    CONFIG.IS_IMG_GEN2 = true;
+                                    const dataImage = await handleImageResponse(result.imageUrl);
+                                    if (isStream) {
+                                        res.write(`data: ${JSON.stringify(MessageProcessor.createChatResponse(dataImage, model, true))}\n\n`);
+                                    } else {
+                                        res.json(MessageProcessor.createChatResponse(dataImage, model));
                                     }
                                 }
-                                if (linejosn.response === "webSearchResults" && CONFIG.ISSHOW_SEARCH_RESULTS) {
-                                    const searchResults = await Utils.organizeSearchResults(linejosn.webSearchResults);
-                                    const responseData = MessageProcessor.createChatResponse(`<think>\r\n${searchResults}\r\n</think>\r\n`, model, true);
-                                    res.write(`data: ${JSON.stringify(responseData)}\n\n`);
-                                }
-                                break;
+                            })();
+                            dataPromises.push(processPromise);
+                        } catch (error) {
+                            console.log(error);
+                            continue;
                         }
-                    } catch (error) {
-                        console.error('JSON解析错误:', error);
                     }
                 }
-            }
-        });
-        stream.on('error', (error) => {
-            console.error('流处理错误:', error);
-            res.write('data: [DONE]\n\n');
-            res.end();
-        });
-
-    } catch (error) {
-        console.error('处理响应错误:', error);
-        res.write('data: [DONE]\n\n');
-        res.end();
-    }
-}
-
-async function handleNormalResponse(response, model, res) {
-    let fullResponse = '';
-    let imageUrl = '';
+            });
 
-    try {
-        const responseText = await response.text();
-        const lines = responseText.split('\n');
-
-        for (const line of lines) {
-            if (!line.trim()) continue;
-            const data = line.slice(6);
-            if (data === '[DONE]') continue;
-            let linejosn = JSON.parse(data);
-            if (linejosn?.error?.name === "RateLimitError") {
-                fullResponse = `${linejosn.error.name},请重新对话`;
-                CONFIG.API.SIGNATURE_COOKIE = null;
-                return res.json(MessageProcessor.createChatResponse(fullResponse, model));
-            }
-            switch (model) {
-                case 'grok-latest-image':
-                    if (linejosn.response === "modelResponse" && linejosn?.modelResponse?.generatedImageUrls) {
-                        imageUrl = linejosn.modelResponse.generatedImageUrls;
-                    }
-                    break;
-                case 'grok-latest':
-                    if (linejosn.response === "token") {
-                        const token = linejosn.token;
-                        if (token && token.length > 0) {
-                            fullResponse += token;
-                        }
-                    }
-                    break;
-                case 'grok-latest-search':
-                    if (linejosn.response === "token") {
-                        const token = linejosn.token;
-                        if (token && token.length > 0) {
-                            fullResponse += token;
-                        }
-                    }
-                    if (linejosn.response === "webSearchResults" && CONFIG.ISSHOW_SEARCH_RESULTS) {
-                        fullResponse += `\r\n<think>${await Utils.organizeSearchResults(linejosn.webSearchResults)}</think>\r\n`;
+            stream.on('end', async () => {
+                try {
+                    await Promise.all(dataPromises);
+                    if (isStream) {
+                        res.write('data: [DONE]\n\n');
+                        res.end();
+                    } else {
+                        if (!CONFIG.IS_IMG_GEN2) {
+                            res.json(MessageProcessor.createChatResponse(fullResponse, model));
+                        }                        
                     }
-                    break;
-            }
-        }
-        if (imageUrl) {
-            const dataImage = await handleImageResponse(imageUrl);
-            const responseData = MessageProcessor.createChatResponse(dataImage, model);
-            res.json(responseData);
-        } else {
-            const responseData = MessageProcessor.createChatResponse(fullResponse, model);
-            res.json(responseData);
-        }
+                    CONFIG.IS_IMG_GEN = false;
+                    CONFIG.IS_IMG_GEN2 = false;
+                    resolve();
+                } catch (error) {
+                    reject(error);
+                }
+            });
+            stream.on('error', (error) => {
+                Logger.error(error, 'Server');
+                reject(error);
+            });
+        });
     } catch (error) {
-        Utils.handleError(error, res);
+        Logger.error(error, 'Server');
+        CONFIG.IS_IMG_GEN = false;
+        CONFIG.IS_IMG_GEN2 = false;
+        throw error;
     }
+
 }
+
 async function handleImageResponse(imageUrl) {
-    //对服务器发送图片请求
-    const MAX_RETRIES = 3;
+    const MAX_RETRIES = 2;
     let retryCount = 0;
     let imageBase64Response;
-    
+
     while (retryCount < MAX_RETRIES) {
         try {
-            //发送图片请求获取图片
             imageBase64Response = await fetch(`https://assets.grok.com/${imageUrl}`, {
                 method: 'GET',
                 headers: {
@@ -649,16 +637,11 @@ async function handleImageResponse(imageUrl) {
                 }
             });
 
-            if (imageBase64Response.ok) {
-                break; // 如果请求成功，跳出重试循环
-            }
-
+            if (imageBase64Response.ok) break; 
             retryCount++;
             if (retryCount === MAX_RETRIES) {
                 throw new Error(`上游服务请求失败! status: ${imageBase64Response.status}`);
             }
-
-            // 等待一段时间后重试（可以使用指数退避）
             await new Promise(resolve => setTimeout(resolve, CONFIG.API.RETRY_TIME * retryCount));
 
         } catch (error) {
@@ -666,7 +649,6 @@ async function handleImageResponse(imageUrl) {
             if (retryCount === MAX_RETRIES) {
                 throw error;
             }
-            // 等待一段时间后重试
             await new Promise(resolve => setTimeout(resolve, CONFIG.API.RETRY_TIME * retryCount));
         }
     }
@@ -692,17 +674,163 @@ async function handleImageResponse(imageUrl) {
     if (!responseURL.ok) {
         return "生图失败，请查看图床密钥是否设置正确"
     } else {
+        console.log("生图成功");
         const result = await responseURL.json();
         return `![image](${result.image.url})`
     }
 }
 
-// 404处理
+const tokenManager = new AuthTokenManager();
+await initialization();
+
+// 中间件配置
+const app = express();
+app.use(Logger.requestLogger);
+app.use(express.json({ limit: CONFIG.SERVER.BODY_LIMIT }));
+app.use(express.urlencoded({ extended: true, limit: CONFIG.SERVER.BODY_LIMIT }));
+app.use(cors({
+    origin: '*',
+    methods: ['GET', 'POST', 'OPTIONS'],
+    allowedHeaders: ['Content-Type', 'Authorization']
+}));
+
+app.get('/hf/v1/models', (req, res) => {
+    res.json({
+        object: "list",
+        data: Object.keys(CONFIG.MODELS).map((model, index) => ({
+            id: model,
+            object: "model",
+            created: Math.floor(Date.now() / 1000),
+            owned_by: "grok",
+        }))
+    });
+});
+
+
+app.post('/hf/v1/chat/completions', async (req, res) => {
+    try {
+        const authToken = req.headers.authorization?.replace('Bearer ', '');
+        if (authToken !== CONFIG.API.API_KEY) {
+            return res.status(401).json({ error: 'Unauthorized' });
+        }
+        let isTempCookie = req.body.model.includes("grok-2");
+        let retryCount = 0;
+        
+        while (retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
+            retryCount++;
+            const grokClient = new GrokApiClient(req.body.model);
+            const requestPayload = await grokClient.prepareChatRequest(req.body);
+
+            if (!CONFIG.API.TEMP_COOKIE) {
+                await Utils.get_signature();
+            }
+
+            if (isTempCookie) {
+                CONFIG.API.SIGNATURE_COOKIE = CONFIG.API.TEMP_COOKIE;
+                Logger.info(`已切换为临时令牌`, 'Server');
+            } else {
+                CONFIG.API.SIGNATURE_COOKIE = Utils.createAuthHeaders(req.body.model);
+            }
+            Logger.info(`当前令牌索引: ${CONFIG.SSO_INDEX}`, 'Server');
+            const newMessageReq = await fetch(`${CONFIG.API.BASE_URL}/api/rpc`, {
+                method: 'POST',
+                headers: {
+                    ...DEFAULT_HEADERS,
+                    ...CONFIG.API.SIGNATURE_COOKIE
+                },
+                body: JSON.stringify({
+                    rpc: "createConversation",
+                    req: {
+                        temporary: false
+                    }
+                })
+            });
+
+            const responseText = await newMessageReq.json();
+            const conversationId = responseText.conversationId;
+
+            const response = await fetch(`${CONFIG.API.BASE_URL}/api/conversations/${conversationId}/responses`, {
+                method: 'POST',
+                headers: {
+                    "accept": "text/event-stream",
+                    "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
+                    "content-type": "text/plain;charset=UTF-8",
+                    "Connection": "keep-alive",
+                    ...CONFIG.API.SIGNATURE_COOKIE
+                },
+                body: JSON.stringify(requestPayload)
+            });
+
+            if (response.ok) {
+                Logger.info(`请求成功`, 'Server');
+                CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
+                Logger.info(`当前剩余可用令牌数: ${tokenManager.getTokenCount()}`, 'Server');
+                try {
+                    await handleResponse(response, req.body.model, res, req.body.stream);
+                    return; 
+                } catch (error) {
+                    if(isTempCookie){
+                        await Utils.get_signature();
+                    }else{
+                        tokenManager.setModelLimit(CONFIG.SSO_INDEX, req.body.model);
+                        for (let i = 1; i <= tokenManager.getTokenCount(); i++) {
+                            CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
+                            if (!tokenManager.isTokenModelLimitReached(CONFIG.SSO_INDEX, req.body.model)) {
+                                break;
+                            } else if (i >= tokenManager.getTokenCount()) {
+                                throw new Error(`${req.body.model} 次数已达上限，请切换其他模型或者重新对话`);
+                            }
+                        }
+                    }
+                }
+            } else {
+                if (response.status === 429) {
+                    if (isTempCookie) {
+                        await Utils.get_signature();
+                    } else {
+                        tokenManager.setModelLimit(CONFIG.SSO_INDEX, req.body.model);
+                        for (let i = 1; i <= tokenManager.getTokenCount(); i++) {
+                            CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
+                            if (!tokenManager.isTokenModelLimitReached(CONFIG.SSO_INDEX, req.body.model)) {
+                                break;
+                            } else if (i >= tokenManager.getTokenCount()) {
+                                throw new Error(`${req.body.model} 次数已达上限，请切换其他模型或者重新对话`);
+                            }
+                        }
+                    }
+                } else {
+                    // 非429错误直接抛出
+                    if (isTempCookie) {
+                        await Utils.get_signature();
+                    } else {
+                        Logger.error(`令牌异常错误状态!status: ${response.status}， 已移除当前令牌${CONFIG.SSO_INDEX.cookie}`, 'Server');
+                        tokenManager.removeTokenByIndex(CONFIG.SSO_INDEX);
+                        Logger.info(`当前剩余可用令牌数: ${tokenManager.getTokenCount()}`, 'Server');
+                        CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
+                    }
+                }
+            }
+        }
+        throw new Error('当前模型所有令牌都已耗尽');
+    } catch (error) {
+        Logger.error(error, 'ChatAPI');
+        res.status(500).json({
+            error: {
+                message: error.message,
+                type: 'server_error',
+                param: null,
+                code: error.code || null
+            }
+        });
+    }
+});
+
+
 app.use((req, res) => {
     res.status(200).send('api运行正常');
 });
 
-// 启动服务器
+
 app.listen(CONFIG.SERVER.PORT, () => {
-    console.log(`服务器已启动，监听端口: ${CONFIG.SERVER.PORT}`);
+    Logger.info(`服务器已启动，监听端口: ${CONFIG.SERVER.PORT}`, 'Server');
 });