src/lightweight-client.js

import fetch from 'node-fetch';
import { JSDOM } from 'jsdom';
import dotenv from 'dotenv';
import { randomUUID } from 'crypto';
import { fileURLToPath } from 'url';
import { dirname, join } from 'path';
import { PassThrough } from 'stream';
import chalk from 'chalk';
import {
  NotionTranscriptConfigValue,
  NotionTranscriptContextValue, NotionTranscriptItem, NotionDebugOverrides,
  NotionRequestBody, ChoiceDelta, Choice, ChatCompletionChunk, NotionTranscriptItemByuser
} from './models.js';
import { proxyPool } from './ProxyPool.js';
import { proxyServer } from './ProxyServer.js';
import { cookieManager } from './CookieManager.js';
import { gotScraping } from 'got-scraping';

// 获取当前文件的目录路径
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

// 加载环境变量
dotenv.config({ path: join(dirname(__dirname), '.env') });

// 日志配置
const logger = {
  info: (message) => console.log(chalk.blue(`[info] ${message}`)),
  error: (message) => console.error(chalk.red(`[error] ${message}`)),
  warning: (message) => console.warn(chalk.yellow(`[warn] ${message}`)),
  success: (message) => console.log(chalk.green(`[success] ${message}`)),
};

// 配置
const NOTION_API_URL = "https://www.notion.so/api/v3/runInferenceTranscript";
// 这些变量将由cookieManager动态提供
let currentCookieData = null;
const USE_NATIVE_PROXY_POOL = process.env.USE_NATIVE_PROXY_POOL === 'true';
const ENABLE_PROXY_SERVER = process.env.ENABLE_PROXY_SERVER === 'true';
let proxy = null;

// 代理配置
const PROXY_URL = process.env.PROXY_URL || "";

// 标记是否成功初始化
let INITIALIZED_SUCCESSFULLY = false;

// 注册进程退出事件，确保代理服务器在程序退出时关闭
process.on('exit', () => {
  try {
    if (proxyServer) {
      proxyServer.stop();
    }
  } catch (error) {
    logger.error(`程序退出时关闭代理服务器出错: ${error.message}`);
  }
});

// 捕获意外退出信号
['SIGINT', 'SIGTERM', 'SIGQUIT'].forEach(signal => {
  process.on(signal, () => {
    logger.info(`收到${signal}信号，正在关闭代理服务器...`);
    try {
      if (proxyServer) {
        proxyServer.stop();
      }
    } catch (error) {
      logger.error(`关闭代理服务器出错: ${error.message}`);
    }
    process.exit(0);
  });
});

// 构建Notion请求
function buildNotionRequest(requestData) {
  // 确保我们有当前的cookie数据
  if (!currentCookieData) {
    currentCookieData = cookieManager.getNext();
    if (!currentCookieData) {
      throw new Error('没有可用的cookie');
    }
  }

  // 当前时间
  const now = new Date();
  // 格式化为ISO字符串，确保包含毫秒和时区
  const isoString = now.toISOString();
  
  // 生成随机名称，类似于Python版本
  const randomWords = ["Project", "Workspace", "Team", "Studio", "Lab", "Hub", "Zone", "Space"];
  const userName = `User${Math.floor(Math.random() * 900) + 100}`; // 生成100-999之间的随机数
  const spaceName = `${randomWords[Math.floor(Math.random() * randomWords.length)]} ${Math.floor(Math.random() * 99) + 1}`;
  
  // 创建transcript数组
  const transcript = [];
  
  // 添加配置项
  if(requestData.model === 'anthropic-sonnet-3.x-stable'){
    transcript.push(new NotionTranscriptItem({
      type: "config",
      value: new NotionTranscriptConfigValue({
      })
    }));
  }else{
      transcript.push(new NotionTranscriptItem({
    type: "config",
    value: new NotionTranscriptConfigValue({
      model: requestData.model
    })
  }));
  }

  
  // 添加上下文项
  transcript.push(new NotionTranscriptItem({
    type: "context",
    value: new NotionTranscriptContextValue({
      userId: currentCookieData.userId,
      spaceId: currentCookieData.spaceId,
      surface: "home_module",
      timezone: "America/Los_Angeles",
      userName: userName,
      spaceName: spaceName,
      spaceViewId: randomUUID(),
      currentDatetime: isoString
    })
  }));
  
  // 添加agent-integration项
  transcript.push(new NotionTranscriptItem({
    type: "agent-integration"
  }));
  
  // 添加消息
  for (const message of requestData.messages) {
    // 处理消息内容，确保格式一致
    let content = message.content;
    
    // 处理内容为数组的情况
    if (Array.isArray(content)) {
      let textContent = "";
      for (const part of content) {
        if (part && typeof part === 'object' && part.type === 'text') {
          if (typeof part.text === 'string') {
            textContent += part.text;
          }
        }
      }
      content = textContent || ""; // 使用提取的文本或空字符串
    } else if (typeof content !== 'string') {
      content = ""; // 如果不是字符串或数组，则默认为空字符串
    }
    
    if (message.role === "system") {
      // 系统消息作为用户消息添加
      transcript.push(new NotionTranscriptItemByuser({
        type: "user",
        value: [[content]],
        userId: currentCookieData.userId,
        createdAt: message.createdAt || isoString
      }));
    } else if (message.role === "user") {
      // 用户消息
      transcript.push(new NotionTranscriptItemByuser({
        type: "user",
        value: [[content]],
        userId: currentCookieData.userId,
        createdAt: message.createdAt || isoString
      }));
    } else if (message.role === "assistant") {
      // 助手消息
      transcript.push(new NotionTranscriptItem({
        type: "markdown-chat",
        value: content,
        traceId: message.traceId || randomUUID(),
        createdAt: message.createdAt || isoString
      }));
    }
  }
  
  // 创建请求体
  return new NotionRequestBody({
    spaceId: currentCookieData.spaceId,
    transcript: transcript,
    createThread: true,
    traceId: randomUUID(),
    debugOverrides: new NotionDebugOverrides({
      cachedInferences: {},
      annotationInferences: {},
      emitInferences: false
    }),
    generateTitle: false,
    saveAllThreadOperations: false
  });
}

// 流式处理Notion响应
async function streamNotionResponse(notionRequestBody) {
  // 确保我们有当前的cookie数据
  if (!currentCookieData) {
    currentCookieData = cookieManager.getNext();
    if (!currentCookieData) {
      throw new Error('没有可用的cookie');
    }
  }

  // 创建流
  const stream = new PassThrough();
  
  // 添加初始数据，确保连接建立
  stream.write(':\n\n');  // 发送一个空注释行，保持连接活跃
  
  // 设置HTTP头模板
  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': currentCookieData.userId,
    'x-notion-space-id': currentCookieData.spaceId
  };
  
  // 设置超时处理，确保流不会无限等待
  const timeoutId = setTimeout(() => {
    logger.warning(`请求超时，30秒内未收到响应`);
    try {
      if (!responseReceived) {
        stream.end();
      }
    } catch (e) {
      logger.error(`关闭超时流时出错: ${e.message}`);
    }
  }, 30000);

  let responseReceived = false;
  
  try {
    const notionCookie = cookieManager.getCookieString(currentCookieData.userId);

    // 配置 got-scraping
    const options = {
      method: 'POST',
      headers: headers,
      body: JSON.stringify(notionRequestBody),
      isStream: true,
      cookieJar: await cookieManager.getGotCookieJar(notionCookie),
      timeout: { request: 30000 },
      http2: true, // 使用HTTP/2，这对于模拟浏览器很重要
      headerGeneratorOptions: { // 伪造浏览器头
          browsers: ['chrome'],
          devices: ['desktop'],
          locales: ['en-US'],
          operatingSystems: ['windows'],
      },
    };

    // 如果配置了代理，则使用它
    if (PROXY_URL) {
      options.proxyUrl = PROXY_URL;
      logger.info(`使用代理: ${PROXY_URL}`);
    }

    const responseStream = gotScraping(NOTION_API_URL, options);

    // 监听流事件
    fetchNotionResponse(stream, responseStream, timeoutId);
  } catch (error) {
    logger.error(`创建 got-scraping 请求时出错: ${error.message}`);
    if (timeoutId) clearTimeout(timeoutId);
    stream.end();
  }
  
  return stream;
}

// 使用got-scraping的流处理函数
async function fetchNotionResponse(chunkQueue, responseStream, timeoutId) {
  let responseReceived = false;
  try {
    responseStream.on('data', (chunk) => {
      responseReceived = true;
      chunkQueue.write(chunk);
    });

    responseStream.on('end', () => {
      if (timeoutId) clearTimeout(timeoutId);
      if (!responseReceived) {
        logger.warning('流已结束但未收到任何数据');
      }
      chunkQueue.end();
    });

    responseStream.on('error', (error) => {
      logger.error(`got-scraping 响应流错误: ${error.message}`);
      if (timeoutId) clearTimeout(timeoutId);
      chunkQueue.end();
    });
  } catch (error) {
      logger.error(`处理 got-scraping 响应时出错: ${error.message}`);
      if (timeoutId) clearTimeout(timeoutId);
      chunkQueue.end();
  }
}

// 应用初始化
async function initialize() {
  logger.info(`初始化Notion配置...`);
  
  // 启动代理服务器 - 已禁用
  // try {
  //   await proxyServer.start();
  // } catch (error) {
  //   logger.error(`启动代理服务器失败: ${error.message}`);
  // }
  
  // 初始化cookie管理器
  let initResult = false;
  
  // 检查是否配置了cookie文件
  const cookieFilePath = process.env.COOKIE_FILE;
  if (cookieFilePath) {
    logger.info(`检测到COOKIE_FILE配置: ${cookieFilePath}`);
    initResult = await cookieManager.loadFromFile(cookieFilePath);
    
    if (!initResult) {
      logger.error(`从文件加载cookie失败，尝试使用环境变量中的NOTION_COOKIE`);
    }
  }
  
  // 如果文件加载失败或未配置文件，尝试从环境变量加载
  if (!initResult) {
    const cookiesString = process.env.NOTION_COOKIE;
    if (!cookiesString) {
      logger.error(`错误: 未设置NOTION_COOKIE环境变量或COOKIE_FILE路径，应用无法正常工作`);
      logger.error(`请在.env文件中设置有效的NOTION_COOKIE值或COOKIE_FILE路径`);
      INITIALIZED_SUCCESSFULLY = false;
      return;
    }
    
    logger.info(`正在从环境变量初始化cookie管理器...`);
    initResult = await cookieManager.initialize(cookiesString);
    
    if (!initResult) {
      logger.error(`初始化cookie管理器失败，应用无法正常工作`);
      INITIALIZED_SUCCESSFULLY = false;
      return;
    }
  }
  
  // 获取第一个可用的cookie数据
  currentCookieData = cookieManager.getNext();
  if (!currentCookieData) {
    logger.error(`没有可用的cookie，应用无法正常工作`);
    INITIALIZED_SUCCESSFULLY = false;
    return;
  }
  
  logger.success(`成功初始化cookie管理器，共有 ${cookieManager.getValidCount()} 个有效cookie`);
  logger.info(`当前使用的cookie对应的用户ID: ${currentCookieData.userId}`);
  logger.info(`当前使用的cookie对应的空间ID: ${currentCookieData.spaceId}`);
  
  // if (process.env.USE_NATIVE_PROXY_POOL === 'true') {
  //   logger.info(`正在初始化本地代理池...`);
  //   await proxyPool.initialize();
  // }
  
  INITIALIZED_SUCCESSFULLY = true;
}

// 导出函数
export {
  initialize,
  streamNotionResponse,
  buildNotionRequest,
  INITIALIZED_SUCCESSFULLY
};