📄 chat.ts  •  14784 bytes
/**
 * CmdCode V0.5 - AI 对话引擎
 * 流式调用 OpenAI 兼容 API,支持工具调用循环
 * 支持 429 错误自动密钥轮换
 */
import OpenAI from 'openai'
import { loadConfig, type Config } from './config.js'
import { ALL_TOOLS, executeTool, type ToolCall, type ToolResult } from './tools.js'
import { rotateChatApiKey, getChatKeyPoolStatus } from './apikeys.js'
import { t } from './i18n.js'
import { searchMemory } from './memory/memoryManager.js'
import { loadAppConfig } from './crypto-util.js'
import { SYSTEM_PROMPT_TEMPLATE } from './system-prompt.js'

/** 最大保留历史消息数(防止上下文超限,默认100,范围10-1000) */
let maxHistoryMessages = 100

export function setMaxHistoryMessages(n: number): void {
  if (n >= 10 && n <= 1000) {
    maxHistoryMessages = n
  }
}

export function getMaxHistoryMessages(): number {
  return maxHistoryMessages
}

/** 截断日志开关(默认开启,可关闭以减少噪声输出) */
export let showTruncationLog = false

export function setShowTruncationLog(show: boolean): void {
  showTruncationLog = show
}

export interface Message {
  role: 'system' | 'user' | 'assistant' | 'tool'
  content: string
  tool_calls?: ToolCall[]
  tool_call_id?: string
}

/** 获取用户自定义系统提示词(从 secrets.enc),无则返回空串 */
export function getSystemPrompt(): string {
  const appConfig = loadAppConfig()
  return appConfig?.systemPrompt || ''
}

/** 构建系统提示词(用户自定义 > 默认模板,工作目录动态) */
export function buildSystemPrompt(workspaceDir?: string): string {
  const custom = getSystemPrompt()
  if (custom) return custom
  const cwd = workspaceDir || process.cwd()
  return SYSTEM_PROMPT_TEMPLATE + `

Current user workspace: ${cwd}`
}

export class ChatEngine {
  private client: OpenAI
  private config: Config
  private messages: Message[]
  private maxTurns: number

  /** 安全过滤:将所有 API Key 替换为星号 */
  private maskKeys(text: string): string {
    return text
      .replace(/\bark-[a-z0-9\-]{20,60}\b/gi, '****')
      .replace(/\bsk-[a-zA-Z0-9\-]{20,80}\b/g, '****')
  }

  constructor(config?: Config, systemPrompt?: string, workspaceDir?: string) {
    this.config = config || loadConfig()
    // 首次连接超时30秒,避免启动时卡死;后续请求用配置超时
    this.client = new OpenAI({
      apiKey: this.config.apiKey,
      baseURL: this.config.baseUrl,
      timeout: Math.min(this.config.timeoutMs, 120_000), // 最大2分钟
      maxRetries: 1, // 减少重试,快速失败
    })
    this.maxTurns = 80 // SWE 任务复杂,最多 80 轮工具调用(原 40)
    this.messages = [
      { role: 'system', content: systemPrompt || buildSystemPrompt(workspaceDir) },
    ]
  }

  /** 从已有消息历史恢复(自动截断防止上下文超限) */
  static fromHistory(messages: Message[], config?: Config, workspaceDir?: string): ChatEngine {
    const engine = new ChatEngine(config, undefined, workspaceDir)
    // 保留系统提示,截断历史消息
    if (messages.length > 0 && messages[0].role === 'system') {
      const systemMsg = messages[0]
      const historyMsgs = messages.slice(1)
      // 保留最近 maxHistoryMessages 条历史
      const truncatedHistory = historyMsgs.slice(-maxHistoryMessages)
      engine.messages = [systemMsg, ...truncatedHistory]
      if (historyMsgs.length > maxHistoryMessages) {
        if (showTruncationLog) console.log(`  ${'\x1b[90m'}${t('history.truncated', {from: historyMsgs.length, to: truncatedHistory.length})}\x1b[0m`)
      }
    } else {
      // 没有系统提示,直接截断
      engine.messages = messages.slice(-maxHistoryMessages)
    }
    return engine
  }

  /** 获取当前消息历史 */
  getHistory(): Message[] {
    return [...this.messages]
  }

  /** 从历史对话中召回相关记忆,注入 System Prompt */
  private async injectMemoryContext(): Promise<void> {
    const nonSystemMsgs = this.messages.filter(m => m.role !== 'system')
    if (nonSystemMsgs.length < 10) return // 消息少时无需注入
    try {
      const lastUser = nonSystemMsgs.filter(m => m.role === 'user').pop()
      if (!lastUser) return
      const results = await searchMemory(lastUser.content, undefined, 5)
      if (results.length > 0) {
        const fragments = results.map(r => `[${r.source}] ${r.content.slice(0, 300)}`).join('\n---\n')
        const memoryInjection = `## 相关历史记忆(来自对话历史)\n${fragments}\n---\n`
        const sysMsg = this.messages.find(m => m.role === 'system')
        if (sysMsg && !sysMsg.content.includes('## 相关历史记忆')) {
          sysMsg.content = memoryInjection + sysMsg.content
        }
      }
    } catch {
      // 双路搜索失败不影响主流程
    }
  }

  /** 动态消息管理:注入记忆后仍保底截断 */
  private async manageContext(): Promise<void> {
    await this.injectMemoryContext()
    this.hardTruncate()
  }

  /**
   * 估算消息列表的token数量(粗略:中文4字=1token,英文1字符=1token)
   * MiniMax-M2.7 128K context,扣4K输出预留,输入上限约120K tokens
   */
  private estimateTokens(): number {
    let chars = 0
    for (const m of this.messages) {
      // role/content 开销 + 内容本身
      chars += m.content.length + 20
      if (m.role === 'tool') chars += 30 // tool_call_id 开销
    }
    // 1 token ≈ 1.5 字符(混合中英文)
    return Math.ceil(chars / 1.5)
  }

  /**
   * 智能上下文截断:同时受消息数量(maxHistoryMessages)和token上限(maxInputTokens)双重约束
   * - 消息数超限 → 保留最近 N 条
   * - token超限 → 从最旧的非system消息开始丢弃,直到token达标
   */
  private hardTruncate(): void {
    const maxTokens = 110_000 // 留10K给输出缓冲,实际更保守
    const systemMsg = this.messages.find(m => m.role === 'system')

    // 阶段1:消息数量截断(保留最近 maxHistoryMessages 条)
    let msgs = this.messages
    if (msgs.length > maxHistoryMessages) {
      msgs = msgs.slice(-maxHistoryMessages)
      if (showTruncationLog) {
        const kept = msgs.filter(m => m.role !== 'system').length
        console.log(`  ${'\x1b[90m'}[截断] 消息数 ${this.messages.length}→${msgs.length}(保留最近${kept}条对话)\x1b[0m`)
      }
    }

    // 阶段2:token数量截断(从最旧的非system消息开始丢弃)
    this.messages = msgs
    let tokens = this.estimateTokens()
    if (tokens <= maxTokens) return

    const nonSystem = msgs.filter(m => m.role !== 'system')
    let dropCount = 0
    for (let i = 0; i < nonSystem.length && tokens > maxTokens; i++) {
      const dropped = nonSystem[i]
      tokens -= Math.ceil(dropped.content.length / 1.5) + 30
      dropCount++
    }
    if (dropCount > 0) {
      const threshold = nonSystem[dropCount - 1]
      const idx = msgs.indexOf(threshold)
      this.messages = msgs.slice(idx + 1)
      if (showTruncationLog) {
        console.log(`  ${'\x1b[90m'}[截断] Token ${Math.ceil(tokens * 1.5)}>${maxTokens} → 丢弃最早${dropCount}条消息\x1b[0m`)
      }
    }
  }

  /** 单轮对话(含工具循环),流式输出到 stdout */
  async chat(userInput: string, onText?: (text: string) => void): Promise<string> {
    this.messages.push({ role: 'user', content: userInput })

    // 动态注入双路记忆 + 保底截断
    await this.manageContext()

    let finalText = ''
    let turns = 0
    // isRetry 防止 429 循环切换密钥后继续触发 429(密钥池耗尽时直接抛错)
    let isRetry = false

    while (turns < this.maxTurns) {
      turns++

      // 调用 API(流式)
      const openaiMessages = this.messages.map(m => {
        // 转换为 OpenAI 格式
        if (m.role === 'tool') {
          return { role: 'tool' as const, content: m.content, tool_call_id: m.tool_call_id! }
        }
        if (m.role === 'assistant' && m.tool_calls) {
          return {
            role: 'assistant' as const,
            content: m.content || null,
            tool_calls: m.tool_calls.map(tc => ({
              id: tc.id,
              type: 'function' as const,
              function: { name: tc.name, arguments: tc.arguments },
            })),
          }
        }
        return { role: m.role as any, content: m.content }
      })

      try {
        const stream = await this.client.chat.completions.create(
          {
            model: this.config.model,
            messages: openaiMessages as any,
            tools: ALL_TOOLS.map(t => ({
              type: 'function' as const,
              function: { name: t.name, description: t.description, parameters: t.parameters },
            })),
            stream: true,
          },
          { signal: AbortSignal.timeout(this.config.timeoutMs) },
        )

        // 收集流式响应
        let textContent = ''
        let toolCalls: Map<number, { id: string; name: string; arguments: string }> = new Map()

        for await (const chunk of stream) {
          const delta = chunk.choices[0]?.delta
          if (!delta) continue

          // 文本内容
          if (delta.content) {
            const safeDelta = this.maskKeys(delta.content)
            textContent += safeDelta
            if (onText) {
              onText(safeDelta)
            } else {
              process.stdout.write(safeDelta)
            }
          }

          // 工具调用
          if (delta.tool_calls) {
            for (const tc of delta.tool_calls) {
              const idx = tc.index
              if (!toolCalls.has(idx)) {
                toolCalls.set(idx, { id: tc.id || '', name: tc.function?.name || '', arguments: '' })
              }
              const existing = toolCalls.get(idx)!
              if (tc.id) existing.id = tc.id
              if (tc.function?.name) existing.name = tc.function.name
              if (tc.function?.arguments) existing.arguments += tc.function.arguments
            }
          }
        }

        // 换行
        if (textContent && !onText) process.stdout.write('\n')

        // 如果没有工具调用,对话结束
        if (toolCalls.size === 0) {
          finalText = textContent
          this.messages.push({ role: 'assistant', content: textContent })
          break
        }

        // 有工具调用,执行并继续循环
        const assistantToolCalls: ToolCall[] = []
        for (const [_, tc] of toolCalls) {
          assistantToolCalls.push({ id: tc.id, name: tc.name, arguments: tc.arguments })
        }

        this.messages.push({
          role: 'assistant',
          content: textContent || '',
          tool_calls: assistantToolCalls,
        })

        // 执行每个工具调用
        for (const tc of assistantToolCalls) {
          const icon = getToolIcon(tc.name)
          if (!onText) process.stdout.write(`\n${icon} ${tc.name}: ${summarizeArgs(tc)}\n`)
          
          const result = await executeTool(tc)
          
          if (!onText) {
            const preview = result.content.length > 200 
              ? this.maskKeys(result.content).slice(0, 200) + '...' 
              : this.maskKeys(result.content)
            process.stdout.write(`  → ${preview}\n`)
          }

          this.messages.push({
            role: 'tool',
            // P2 #2.5: 存储前过滤密钥,防止通过记忆系统泄露
            content: this.maskKeys(result.content),
            tool_call_id: tc.id,
          })
          
          // 工具消息后管理上下文
          await this.manageContext()
        }

        if (!onText) process.stdout.write('\n')
        finalText = textContent
      } catch (e: any) {
        // 检测 429 错误或配额用尽
        const is429 = e?.status === 429 || 
                      e?.error?.code === 429 || 
                      String(e).includes('429') ||
                      String(e).includes('rate limit') ||
                      String(e).includes('quota') ||
                      String(e).includes('exceeded')
        
        if (is429) {
          const poolStatus = getChatKeyPoolStatus()
          // 如果是本次请求内第二次触发 429(已切换过一次密钥),说明密钥池已耗尽
          if (isRetry || poolStatus.remaining <= 0) {
            console.log(`\n  \x1b[31m❌ ${t('ratelimit.exhausted')}\x1b[0m`)
            console.log(`  \x1b[33m${t('ratelimit.add_own')}\x1b[0m\n`)
            throw new Error('API key pool exhausted. Use /set to add your own API key.')
          }
          // 第一次 429:切换密钥后等待 15s(MiniMax m1 冷启动需要 8-15s)
          isRetry = true
          const nextKey = rotateChatApiKey()
          if (nextKey && nextKey.name !== this.config.keyName) {
            console.log(`\n  \x1b[33m⚠️ ${t('ratelimit.switching', {name: nextKey.name})}(等待15s)\x1b[0m\n`)
            await new Promise(r => setTimeout(r, 15000))
            const apiKeyStr = nextKey.apiKey.toString('utf-8')
            this.config.apiKey = apiKeyStr
            this.config.keyName = nextKey.name
            this.config.baseUrl = nextKey.baseUrl
            this.config.model = nextKey.model
            this.client = new OpenAI({
              apiKey: apiKeyStr,
              baseURL: nextKey.baseUrl,
              timeout: Math.min(this.config.timeoutMs, 120_000),
              maxRetries: 1,
            })
            continue // 重试当前请求
          }
          throw new Error('API key pool exhausted. Use /set to add your own API key.')
        }
        
        // 其他错误直接抛出
        throw e
      }
    }

    if (turns >= this.maxTurns) {
      const msg = `\n⚠️ Reached max turns (${this.maxTurns})`
      if (!onText) process.stdout.write(msg)
      finalText += msg
    }

    return finalText
  }
}

function getToolIcon(name: string): string {
  switch (name) {
    case 'file_read': return '📖'
    case 'file_write': return '✏️'
    case 'file_edit': return '🔧'
    case 'bash_run': return '⚡'
    case 'grep_search': return '🔍'
    case 'list_dir': return '📁'
    default: return '🔧'
  }
}

function summarizeArgs(tc: ToolCall): string {
  try {
    const args = JSON.parse(tc.arguments)
    switch (tc.name) {
      case 'file_read': return args.path
      case 'file_write': return args.path
      case 'file_edit': return args.path
      case 'bash_run': return args.command?.slice(0, 80)
      case 'grep_search': return `"${args.pattern}" in ${args.path || '.'}`
      case 'list_dir': return args.path || '.'
      default: return JSON.stringify(args).slice(0, 80)
    }
  } catch {
    return tc.arguments.slice(0, 80)
  }
}