宜博 / lemonboss

/**
 * MemoryService — Agent 自动记忆读写
 *
 * Phase 1 (MVP):
 *   - 对话结束 → 自动生成 session 摘要 → short-term/
 *   - 对话结束 → 关键事实追加 → memory.md
 *   - 新对话开始 → 加载最近 N 个 session 摘要注入 prompt
 *
 * Phase 2 (TODO):
 *   - 海马体关键词匹配
 *   - 小脑方案固化
 */

const fs = require('fs');
const path = require('path');

const DATA_DIR = path.resolve(__dirname, '../../data/agents');
const MAX_SHORT_TERM = 10;            // 最多保留最近 N 个 session 摘要
const MAX_MEMORY_MD_SIZE = 3000;      // memory.md 最大字符数
const INJECT_SESSION_COUNT = 5;       // prompt 注入最近 N 个 session

// ════════════════════════════════════════
//  1. 对话结束 → 自动保存 Session 摘要
// ════════════════════════════════════════

/**
 * 从 messages.log 和 meta 中提取信息，生成 session 摘要，写入 short-term/
 *
 * @param {string} agentId
 * @param {string} convId
 * @param {object} meta - { status, stats, createdAt, updatedAt, title }
 * @param {string} lastOutput - Agent 最后的输出文本
 */
function saveSessionSummary(agentId, convId, meta, lastOutput) {
  try {
    const agentDir = path.join(DATA_DIR, agentId);
    const shortTermDir = path.join(agentDir, 'memory', 'short-term');

    if (!fs.existsSync(shortTermDir)) {
      fs.mkdirSync(shortTermDir, { recursive: true });
    }

    // 提取任务描述 (从 title 或 messages.log 的第一条)
    const task = extractTask(meta.title || '', agentId, convId);

    // 提取结果摘要 (Agent 最后输出的前 500 字)
    const resultSummary = extractResultSummary(lastOutput);

    // 提取提到的文件路径
    const files = extractMentionedFiles(lastOutput);

    // 计算持续时间
    const createdAt = new Date(meta.createdAt);
    const updatedAt = new Date(meta.updatedAt || Date.now());
    const durationSec = Math.round((updatedAt.getTime() - createdAt.getTime()) / 1000);
    const durationStr = durationSec > 60
      ? `${Math.round(durationSec / 60)}分钟`
      : `${durationSec}秒`;

    // 构造摘要
    const timestamp = new Date().toISOString().replace('T', ' ').substring(0, 19);
    const summary = [
      `# Session 摘要`,
      ``,
      `- **时间**: ${timestamp}`,
      `- **对话ID**: ${convId}`,
      `- **状态**: ${meta.status || 'completed'}`,
      `- **工具调用**: ${meta.stats?.toolCallCount || 0} 次`,
      `- **耗时**: ${durationStr}`,
      ``,
      `## 任务`,
      task || '(无明确任务描述)',
      ``,
      `## 结果`,
      resultSummary || '(无输出)',
    ];

    if (files.length > 0) {
      summary.push('', '## 涉及文件', ...files.map(f => `- ${f}`));
    }

    summary.push('');

    // 写入文件
    const filename = `session-${Date.now()}.md`;
    const filepath = path.join(shortTermDir, filename);
    fs.writeFileSync(filepath, summary.join('\n'), 'utf-8');

    // 清理过旧的 session 文件（保留最近 MAX_SHORT_TERM 个）
    cleanupShortTerm(shortTermDir);

    // 追加关键事实到 memory.md
    appendToMemoryMd(agentId, task, resultSummary, meta);

    console.log(`[Memory] ✅ Session 摘要已保存: ${agentId}/${filename}`);
    return filepath;
  } catch (err) {
    console.error(`[Memory] ⚠️ Session 摘要保存失败:`, err.message);
    return null;
  }
}

// ════════════════════════════════════════
//  2. 新对话开始 → 加载近期记忆
// ════════════════════════════════════════

/**
 * 加载最近 N 个 session 摘要，用于注入 system prompt
 *
 * @param {string} agentId
 * @returns {string} 格式化的近期记忆文本
 */
function loadRecentMemory(agentId) {
  try {
    const shortTermDir = path.join(DATA_DIR, agentId, 'memory', 'short-term');
    if (!fs.existsSync(shortTermDir)) return '';

    // 按文件名倒序（时间戳降序）取最近 N 个
    const files = fs.readdirSync(shortTermDir)
      .filter(f => f.startsWith('session-') && f.endsWith('.md'))
      .sort()
      .reverse()
      .slice(0, INJECT_SESSION_COUNT);

    if (files.length === 0) return '';

    const sessions = files.map(f => {
      const content = fs.readFileSync(path.join(shortTermDir, f), 'utf-8');
      return compactSession(content);
    });

    return [
      `# 近期记忆（最近 ${sessions.length} 次对话）`,
      '',
      ...sessions,
    ].join('\n');
  } catch (err) {
    console.error(`[Memory] ⚠️ 加载近期记忆失败:`, err.message);
    return '';
  }
}

// ════════════════════════════════════════
//  3. memory.md 自动追加
// ════════════════════════════════════════

/**
 * 追加关键事实到 memory.md
 * 只在有实质操作（toolCallCount > 0）时追加
 */
function appendToMemoryMd(agentId, task, result, meta) {
  if (!meta.stats?.toolCallCount || meta.stats.toolCallCount === 0) return;

  const memoryPath = path.join(DATA_DIR, agentId, 'memory', 'memory.md');
  let content = '';
  try {
    content = fs.readFileSync(memoryPath, 'utf-8');
  } catch { /* 文件不存在 */ }

  const timestamp = new Date().toISOString().substring(0, 10);
  const taskOneLine = (task || '').split('\n')[0].substring(0, 80);
  const resultOneLine = (result || '').split('\n')[0].substring(0, 100);

  // 构造新条目
  const entry = `- [${timestamp}] ${taskOneLine} → ${resultOneLine}`;

  // 追加到"关键事实"段
  if (content.includes('## 关键事实')) {
    content = content.replace(
      /## 关键事实\n([\s\S]*?)(?=\n## |$)/,
      (match, existing) => {
        const trimmed = existing.trim();
        if (trimmed === '（尚无记录）') {
          return `## 关键事实\n${entry}\n`;
        }
        return `## 关键事实\n${trimmed}\n${entry}\n`;
      }
    );
  } else {
    // 没有"关键事实"段，追加到末尾
    content += `\n\n## 关键事实\n${entry}\n`;
  }

  // 大小控制：超过限制时截断最早的条目
  if (content.length > MAX_MEMORY_MD_SIZE) {
    content = truncateMemory(content);
  }

  fs.writeFileSync(memoryPath, content, 'utf-8');
}

// ════════════════════════════════════════
//  内部工具函数
// ════════════════════════════════════════

/** 从 title 提取任务描述 */
function extractTask(title, agentId, convId) {
  // 去掉 [TASK from xxx] 前缀
  let task = title.replace(/\[TASK from \w+\]\s*/g, '').trim();
  // 去掉任务编号行
  task = task.replace(/任务编号:\s*T-\w+\s*/g, '').trim();
  task = task.replace(/优先级:\s*P\d+\s*/g, '').trim();

  // 如果 title 太短，尝试从 messages.log 读取
  if (task.length < 10) {
    try {
      const logPath = path.join(DATA_DIR, agentId, 'conversations', convId, 'messages.log');
      const log = fs.readFileSync(logPath, 'utf-8');
      const userMsg = log.match(/\[USER\]\n([\s\S]*?)(?:\n\[AGENT\]|$)/);
      if (userMsg) task = userMsg[1].trim().substring(0, 300);
    } catch { /* ignore */ }
  }

  return task.substring(0, 300);
}

/** 从 Agent 输出提取结果摘要 */
function extractResultSummary(output) {
  if (!output) return '';
  // 优先取 [RESULT: DONE]...[/RESULT] 块
  const resultBlock = output.match(/\[RESULT:\s*DONE\]([\s\S]*?)\[\/RESULT\]/);
  if (resultBlock) return resultBlock[1].trim().substring(0, 500);

  // 否则取最后 500 字
  const trimmed = output.trim();
  if (trimmed.length > 500) {
    return '...' + trimmed.substring(trimmed.length - 500);
  }
  return trimmed;
}

/** 提取输出中提到的文件路径 */
function extractMentionedFiles(output) {
  if (!output) return [];
  // 匹配常见文件路径模式
  const patterns = [
    /(?:写入|创建|保存|输出)[到至：:]\s*[`"']?([^\s`"']+\.\w+)/g,
    /(?:data\/\w+\/[^\s]+)/g,
    /(?:workspace\/[^\s]+)/g,
  ];

  const files = new Set();
  for (const pat of patterns) {
    let m;
    while ((m = pat.exec(output)) !== null) {
      const f = m[1] || m[0];
      if (f.length > 5 && f.length < 200) files.add(f);
    }
  }
  return [...files].slice(0, 10);  // 最多 10 个
}

/** 压缩 session 内容为紧凑格式（用于 prompt 注入） */
function compactSession(content) {
  // 提取关键字段
  const time = content.match(/时间\**[：:]\s*(.+)/)?.[1] || '';
  const status = content.match(/状态\**[：:]\s*(.+)/)?.[1] || '';
  const tools = content.match(/工具调用\**[：:]\s*(.+)/)?.[1] || '';
  const duration = content.match(/耗时\**[：:]\s*(.+)/)?.[1] || '';

  // 提取任务和结果段
  const taskMatch = content.match(/## 任务\n([\s\S]*?)(?=\n## |$)/);
  const resultMatch = content.match(/## 结果\n([\s\S]*?)(?=\n## |$)/);
  const filesMatch = content.match(/## 涉及文件\n([\s\S]*?)(?=\n## |$)/);

  const task = taskMatch?.[1]?.trim().substring(0, 100) || '';
  const result = resultMatch?.[1]?.trim().substring(0, 200) || '';
  const filesList = filesMatch?.[1]?.trim() || '';

  const parts = [`### ${time} [${status}] (${tools}, ${duration})`];
  if (task) parts.push(`- 任务: ${task}`);
  if (result) parts.push(`- 结果: ${result}`);
  if (filesList) parts.push(`- 文件: ${filesList.replace(/\n/g, ', ').replace(/- /g, '')}`);
  parts.push('');

  return parts.join('\n');
}

/** 清理过旧的 short-term 文件 */
function cleanupShortTerm(dir) {
  try {
    const files = fs.readdirSync(dir)
      .filter(f => f.startsWith('session-') && f.endsWith('.md'))
      .sort();

    // 保留最近 MAX_SHORT_TERM 个
    while (files.length > MAX_SHORT_TERM) {
      const oldest = files.shift();
      fs.unlinkSync(path.join(dir, oldest));
    }
  } catch { /* ignore */ }
}

/** memory.md 超大时截断最早条目 */
function truncateMemory(content) {
  const lines = content.split('\n');
  // 找到 ## 关键事实 段，删除最早的条目
  const factIdx = lines.findIndex(l => l.startsWith('## 关键事实'));
  if (factIdx === -1) return content;

  // 找到条目行（以 - [ 开头）
  const entries = [];
  for (let i = factIdx + 1; i < lines.length; i++) {
    if (lines[i].startsWith('## ')) break;
    if (lines[i].startsWith('- [')) entries.push(i);
  }

  // 删除最早的 1/3 条目
  const removeCount = Math.max(1, Math.floor(entries.length / 3));
  const removeIndices = new Set(entries.slice(0, removeCount));
  const newLines = lines.filter((_, i) => !removeIndices.has(i));

  return newLines.join('\n');
}

module.exports = {
  saveSessionSummary,
  loadRecentMemory,
  appendToMemoryMd,
};