Update src/lightweight-client-express.js

src/lightweight-client-express.js

@@ -7,11 +7,9 @@ import chalk from 'chalk';
 import {
   ChatMessage, ChatCompletionRequest, Choice, ChoiceDelta, ChatCompletionChunk
 } from './models.js';
-import {
-  initialize,
-  streamNotionResponse
-} from './lightweight-client.js';
 import { cookieManager } from './CookieManager.js';
+import { PassThrough } from 'stream';
+import fetch from 'node-fetch';
 
 // 获取当前文件的目录路径
 const __filename = fileURLToPath(import.meta.url);
@@ -54,7 +52,6 @@ async function validateProxy() {
   }
   
   try {
-    const { default: fetch } = await import('node-fetch');
     const { HttpsProxyAgent } = await import('https-proxy-agent');
     
     const testResponse = await fetch('https://httpbin.org/ip', {
@@ -78,6 +75,282 @@ async function validateProxy() {
   }
 }
 
+// 初始化函数
+async function initialize() {
+  logger.info('开始系统初始化...');
+  
+  try {
+    const envCookie = process.env.NOTION_COOKIE;
+    const envCookieFile = process.env.COOKIE_FILE;
+    
+    if (envCookie) {
+      logger.info('发现环境变量中的NOTION_COOKIE，正在初始化...');
+      const success = await cookieManager.initialize(envCookie);
+      if (success) {
+        logger.success('Cookie初始化成功');
+        return true;
+      } else {
+        logger.error('Cookie初始化失败');
+        return false;
+      }
+    } else if (envCookieFile) {
+      logger.info(`发现环境变量中的COOKIE_FILE: ${envCookieFile}，正在加载...`);
+      const success = await cookieManager.loadFromFile(envCookieFile);
+      if (success) {
+        logger.success('从文件加载Cookie成功');
+        return true;
+      } else {
+        logger.error('从文件加载Cookie失败');
+        return false;
+      }
+    } else {
+      logger.warning('未找到NOTION_COOKIE或COOKIE_FILE环境变量');
+      return false;
+    }
+  } catch (error) {
+    logger.error(`初始化过程出错: ${error.message}`);
+    return false;
+  }
+}
+
+// 流式响应函数
+async function streamNotionResponse(notionRequestBody) {
+  const stream = new PassThrough();
+  let streamClosed = false;
+  
+  const originalEnd = stream.end;
+  stream.end = function(...args) {
+    if (streamClosed) return;
+    streamClosed = true;
+    return originalEnd.apply(this, args);
+  };
+  
+  stream.write(':\n\n');
+  
+  const timeoutId = setTimeout(() => {
+    if (streamClosed) return;
+    
+    logger.warning(`请求超时，30秒内未收到响应`);
+    try {
+      const endChunk = new ChatCompletionChunk({
+        choices: [
+          new Choice({
+            delta: new ChoiceDelta({ content: "请求超时，未收到Notion响应。" }),
+            finish_reason: "timeout"
+          })
+        ]
+      });
+      stream.write(`data: ${JSON.stringify(endChunk)}\n\n`);
+      stream.write('data: [DONE]\n\n');
+      stream.end();
+    } catch (error) {
+      logger.error(`发送超时消息时出错: ${error}`);
+      if (!streamClosed) stream.end();
+    }
+  }, 30000);
+  
+  fetchNotionResponse(stream, notionRequestBody, timeoutId).catch((error) => {
+    if (streamClosed) return;
+    
+    logger.error(`流处理出错: ${error}`);
+    clearTimeout(timeoutId);
+    
+    try {
+      const errorChunk = new ChatCompletionChunk({
+        choices: [
+          new Choice({
+            delta: new ChoiceDelta({ content: `处理请求时出错: ${error.message}` }),
+            finish_reason: "error"
+          })
+        ]
+      });
+      stream.write(`data: ${JSON.stringify(errorChunk)}\n\n`);
+      stream.write('data: [DONE]\n\n');
+    } catch (e) {
+      logger.error(`发送错误消息时出错: ${e}`);
+    } finally {
+      if (!streamClosed) stream.end();
+    }
+  });
+  
+  return stream;
+}
+
+// 实际请求Notion API的函数
+async function fetchNotionResponse(chunkQueue, notionRequestBody, timeoutId) {
+  let responseReceived = false;
+  
+  const isStreamClosed = () => {
+    return chunkQueue.destroyed || (typeof chunkQueue.closed === 'boolean' && chunkQueue.closed);
+  };
+  
+  const safeWrite = (data) => {
+    if (!isStreamClosed()) {
+      try {
+        return chunkQueue.write(data);
+      } catch (error) {
+        logger.error(`流写入错误: ${error.message}`);
+        return false;
+      }
+    }
+    return false;
+  };
+  
+  try {
+    const cookieData = cookieManager.getNext();
+    if (!cookieData) {
+      throw new Error('没有可用的cookie');
+    }
+    
+    const headers = {
+      'Content-Type': 'application/json',
+      'accept': 'application/x-ndjson',
+      'accept-language': 'en-US,en;q=0.9',
+      'notion-audit-log-platform': 'web',
+      'notion-client-version': '23.13.0.3686',
+      'origin': 'https://www.notion.so',
+      'referer': 'https://www.notion.so/chat',
+      'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36',
+      'x-notion-active-user-header': cookieData.userId,
+      'x-notion-space-id': cookieData.spaceId,
+      'Cookie': cookieData.cookie
+    };
+    
+    const fetchOptions = {
+      method: 'POST',
+      headers: headers,
+      body: JSON.stringify(notionRequestBody),
+    };
+    
+    // 添加代理配置
+    if (PROXY_URL) {
+      const { HttpsProxyAgent } = await import('https-proxy-agent');
+      fetchOptions.agent = new HttpsProxyAgent(PROXY_URL);
+      logger.info(`使用固定代理连接Notion API`);
+    }
+    
+    const response = await fetch('https://www.notion.so/api/v3/runInferenceTranscript', fetchOptions);
+    
+    if (response.status === 401) {
+      logger.error(`收到401未授权错误，cookie可能已失效`);
+      cookieManager.markAsInvalid(cookieData.userId);
+      throw new Error('Cookie已失效');
+    }
+    
+    if (!response.ok) {
+      throw new Error(`HTTP error! status: ${response.status}`);
+    }
+    
+    if (!response.body) {
+      throw new Error("Response body is null");
+    }
+    
+    const reader = response.body;
+    let buffer = '';
+    
+    reader.on('data', (chunk) => {
+      if (isStreamClosed()) {
+        try {
+          reader.destroy();
+        } catch (error) {
+          logger.error(`销毁reader时出错: ${error.message}`);
+        }
+        return;
+      }
+      
+      try {
+        if (!responseReceived) {
+          responseReceived = true;
+          logger.info(`已连接Notion API`);
+          clearTimeout(timeoutId);
+        }
+        
+        const text = chunk.toString('utf8');
+        buffer += text;
+        
+        const lines = buffer.split('\n');
+        buffer = lines.pop() || '';
+        
+        for (const line of lines) {
+          if (!line.trim()) continue;
+          
+          try {
+            const jsonData = JSON.parse(line);
+            
+            if (jsonData?.type === "markdown-chat" && typeof jsonData?.value === "string") {
+              const content = jsonData.value;
+              if (!content) continue;
+              
+              const chunk = new ChatCompletionChunk({
+                choices: [
+                  new Choice({
+                    delta: new ChoiceDelta({ content }),
+                    finish_reason: null
+                  })
+                ]
+              });
+              
+              const dataStr = `data: ${JSON.stringify(chunk)}\n\n`;
+              if (!safeWrite(dataStr)) {
+                try {
+                  reader.destroy();
+                } catch (error) {
+                  logger.error(`写入失败后销毁reader时出错: ${error.message}`);
+                }
+                return;
+              }
+            }
+          } catch (parseError) {
+            logger.error(`解析JSON失败: ${parseError.message}`);
+          }
+        }
+      } catch (error) {
+        logger.error(`处理数据块时出错: ${error.message}`);
+      }
+    });
+    
+    reader.on('end', () => {
+      logger.info('Notion API响应流结束');
+      clearTimeout(timeoutId);
+      
+      try {
+        const endChunk = new ChatCompletionChunk({
+          choices: [
+            new Choice({
+              delta: new ChoiceDelta({ content: "" }),
+              finish_reason: "stop"
+            })
+          ]
+        });
+        safeWrite(`data: ${JSON.stringify(endChunk)}\n\n`);
+        safeWrite('data: [DONE]\n\n');
+      } catch (error) {
+        logger.error(`发送结束标记时出错: ${error.message}`);
+      }
+      
+      if (!isStreamClosed()) {
+        chunkQueue.end();
+      }
+    });
+    
+    reader.on('error', (error) => {
+      logger.error(`读取响应流时出错: ${error.message}`);
+      clearTimeout(timeoutId);
+      if (!isStreamClosed()) {
+        chunkQueue.end();
+      }
+    });
+    
+  } catch (error) {
+    logger.error(`请求Notion API失败: ${error.message}`);
+    clearTimeout(timeoutId);
+    if (!isStreamClosed()) {
+      chunkQueue.end();
+    }
+    throw error;
+  }
+}
+
 // 构建Notion请求的函数
 function buildNotionRequest(requestData) {
   const cookieData = cookieManager.getNext();
@@ -487,5 +760,3 @@ initialize().then(async (initResult) => {
     logger.warning(`服务已在 0.0.0.0:${PORT} 上启动（初始化失败模式）`);
   });
 });
-
-