AstroResearch/src/agent/memory/mod.rs

// src/agent/memory/mod.rs
//
// 项目记忆管理器。
// 参考 Claude Code memdir 设计。
//
// 在 {library_dir}/memory/ 目录下维护：
//   - MEMORY.md — 索引文件（最多 200 行，25KB）
//   - {slug}.md — 每个记忆一个文件，YAML frontmatter + Markdown 内容
//
// 自动在 Agent 的 system prompt 中注入最近的记忆条目。
// 提供 save_memory 工具供 Agent 写入记忆。

pub mod age;
pub mod decay;
pub mod dedup;
pub mod extraction;
pub mod guardrails;
pub mod selection;
pub mod types;

use std::fs;
use std::path::PathBuf;
use tracing::{info, warn};

use self::types::{entry_from_frontmatter, parse_frontmatter, MemoryEntry, MemoryType};

/// 索引文件最大行数
const MAX_ENTRYPOINT_LINES: usize = 200;
/// 索引文件最大字节数
const MAX_ENTRYPOINT_BYTES: usize = 25_000;

/// 记忆管理器
pub struct MemoryManager {
    /// 记忆目录
    memory_dir: PathBuf,
    /// 已加载的记忆条目
    entries: Vec<MemoryEntry>,
    /// 自动提取追踪状态
    pub extraction_tracker: extraction::ExtractionTracker,
}

impl MemoryManager {
    /// 创建并加载记忆。
    /// `library_dir` 是项目配置中的 library 目录。
    pub fn new(library_dir: PathBuf) -> Self {
        let memory_dir = library_dir.join("memory");
        let mut manager = MemoryManager {
            memory_dir,
            entries: Vec::new(),
            extraction_tracker: extraction::ExtractionTracker::default(),
        };
        manager.reload();
        manager
    }

    /// 重新从磁盘加载所有记忆。
    pub fn reload(&mut self) {
        // 确保目录存在
        if let Err(e) = fs::create_dir_all(&self.memory_dir) {
            warn!("[Memory] 无法创建记忆目录 {:?}: {}", self.memory_dir, e);
            return;
        }

        self.entries.clear();

        // 扫描 .md 文件（排除 MEMORY.md 和目录）
        match fs::read_dir(&self.memory_dir) {
            Ok(dir_entries) => {
                for entry in dir_entries {
                    let entry = match entry {
                        Ok(e) => e,
                        Err(_) => continue,
                    };
                    let path = entry.path();
                    if path.is_dir() {
                        continue;
                    }
                    let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
                    if file_name == "MEMORY.md" || !file_name.ends_with(".md") {
                        continue;
                    }

                    let slug = file_name.strip_suffix(".md").unwrap_or(file_name);
                    let mtime = entry
                        .metadata()
                        .ok()
                        .and_then(|m| m.modified().ok())
                        .map(|t| {
                            t.duration_since(std::time::UNIX_EPOCH)
                                .unwrap_or_default()
                                .as_secs()
                        })
                        .unwrap_or(0);

                    match fs::read_to_string(&path) {
                        Ok(raw) => {
                            let (fields, content) = parse_frontmatter(&raw);
                            if let Some(mem_entry) =
                                entry_from_frontmatter(slug, &fields, &content, path.clone(), mtime)
                            {
                                self.entries.push(mem_entry);
                            }
                        }
                        Err(e) => {
                            warn!("[Memory] 无法读取 {:?}: {}", path, e);
                        }
                    }
                }
            }
            Err(e) => {
                warn!("[Memory] 无法扫描记忆目录: {}", e);
            }
        }

        // 按修改时间排序（最新在前）
        self.entries.sort_by(|a, b| b.mtime.cmp(&a.mtime));

        info!(
            "[Memory] 加载了 {} 条记忆从 {:?}",
            self.entries.len(),
            self.memory_dir
        );
    }

    /// 获取所有记忆条目
    pub fn entries(&self) -> &[MemoryEntry] {
        &self.entries
    }

    /// 获取记忆目录路径
    pub fn memory_dir(&self) -> &std::path::Path {
        &self.memory_dir
    }

    /// 标记主代理在本会话中已写入记忆（抑制自动提取）。
    pub fn mark_main_agent_wrote(&mut self) {
        self.extraction_tracker.main_agent_saved_this_session = true;
    }

    /// 语义匹配：委托给 selection 模块使用 LLM 结构化选择。
    ///
    /// 选择后应用指数时间衰减排序（更近的 active 记忆获得更高权重）。
    /// 失败时回退到 recency-based 选择（跳过已展示的条目）。
    pub async fn select_relevant_memories(
        &self,
        llm: &crate::clients::llm::LlmClient,
        context: &str,
        max_entries: usize,
    ) -> Vec<&MemoryEntry> {
        if self.entries.len() <= max_entries {
            return self.entries.iter().collect();
        }

        let sel_ctx = selection::SelectionContext {
            max_entries,
            ..Default::default()
        };

        let indices = selection::select_structured(&self.entries, llm, context, &sel_ctx).await;

        // 应用指数时间衰减排序（更近的 active 记忆在前，historical 在后）
        let sorted =
            selection::apply_decay_scoring(&indices, &self.entries, decay::DEFAULT_HALF_LIFE_DAYS);

        sorted.iter().filter_map(|&i| self.entries.get(i)).collect()
    }

    /// 从指定条目列表构建 system reminder（而非全部条目）
    pub fn build_system_reminder_from(&self, selected: &[&MemoryEntry]) -> Option<String> {
        if selected.is_empty() {
            return None;
        }

        let mut lines = vec![
            "<project-memory-context>".to_string(),
            String::new(),
            "[PROJECT MEMORY]".to_string(),
            String::new(),
        ];

        for entry in selected {
            let type_tag = match entry.memory_type {
                MemoryType::User => "[偏好]",
                MemoryType::Feedback => "[反馈]",
                MemoryType::Project => "[项目]",
                MemoryType::Reference => "[参考]",
            };
            // historical 记忆标记
            let status_tag = if !entry.status.is_active() {
                match entry.status.superseded_by() {
                    Some(new_slug) => format!(" [已更新→{}]", new_slug),
                    None => " [已更新]".to_string(),
                }
            } else {
                String::new()
            };

            let preview: String = entry
                .content
                .lines()
                .take(3)
                .collect::<Vec<_>>()
                .join("\n  ");
            lines.push(format!(
                "{} {}{}: {}\n  {}",
                type_tag, entry.name, status_tag, entry.description, preview
            ));
            // 注入时效警告（超过1天的记忆）
            let freshness = age::memory_freshness_note(entry.mtime);
            if !freshness.is_empty() {
                lines.push(freshness);
            }
        }

        lines.push(String::new());
        lines.push(
            "使用 save_memory 工具保存重要信息。记忆内容可能过时，请在使用前验证。".to_string(),
        );
        // 注入验证提醒（从记忆推荐前先核实）
        lines.push(String::new());
        lines.push(guardrails::build_verification_reminder());
        lines.push("</project-memory-context>".to_string());

        Some(lines.join("\n"))
    }

    /// 保存一条新的记忆。
    ///
    /// 如果 slug 已存在且内容为 Active，旧文件归档为 `{slug}_v1.md`
    /// 并将状态标记为 historical（永不删除旧记忆）。
    pub fn save_memory(
        &mut self,
        slug: &str,
        name: &str,
        description: &str,
        memory_type: MemoryType,
        content: &str,
    ) -> std::io::Result<()> {
        let file_path = self.memory_dir.join(format!("{}.md", slug));

        // 如果已有活跃版本，归档旧版本
        let old_path = self.memory_dir.join(format!("{}_v1.md", slug));
        if file_path.exists() {
            if let Ok(old_content) = fs::read_to_string(&file_path) {
                fs::write(&old_path, &old_content)?;
                info!("[Memory] 归档旧版本: {} → {}", slug, old_path.display());
            }
        }

        let frontmatter = format!(
            "---\nname: {}\ndescription: {}\ntype: {}\nstatus: active\n---\n",
            name,
            description,
            memory_type.as_str()
        );
        let full_content = format!("{}{}", frontmatter, content);

        fs::write(&file_path, &full_content)?;

        // 更新 MEMORY.md 索引
        self.update_index(slug, name, description, &memory_type)?;

        // 重新加载
        self.reload();

        info!("[Memory] 已保存记忆: {} ({})", name, slug);
        Ok(())
    }

    /// 更新 MEMORY.md 索引文件。
    fn update_index(
        &self,
        slug: &str,
        name: &str,
        description: &str,
        memory_type: &MemoryType,
    ) -> std::io::Result<()> {
        let index_path = self.memory_dir.join("MEMORY.md");
        let line = format!(
            "- [{}]({}.md) — {} (type: {})",
            name,
            slug,
            description,
            memory_type.as_str()
        );

        let mut content = if index_path.exists() {
            let existing = fs::read_to_string(&index_path).unwrap_or_default();
            // 检查是否已有此 slug 的条目
            let slug_marker = format!("]({}.md)", slug);
            let lines: Vec<&str> = existing.lines().collect();

            // 行数检查
            if lines.len() >= MAX_ENTRYPOINT_LINES {
                // 移除最旧的行（索引头部保持不变）
                warn!(
                    "[Memory] MEMORY.md 行数已满 ({}), 移除最旧条目",
                    lines.len()
                );
                let keep = MAX_ENTRYPOINT_LINES - 1;
                format!("{}\n{}", lines[..keep.min(lines.len())].join("\n"), line)
            } else {
                // 检查是否需要替换已存在的条目
                let has_entry = lines.iter().any(|l| l.contains(&slug_marker));
                if has_entry {
                    // 替换已存在的行
                    lines
                        .iter()
                        .map(|l| {
                            if l.contains(&slug_marker) {
                                line.as_str()
                            } else {
                                *l
                            }
                        })
                        .collect::<Vec<_>>()
                        .join("\n")
                } else {
                    format!("{}\n{}", existing, line)
                }
            }
        } else {
            format!("# Project Memory\n\n{}", line)
        };

        // 字节数检查（在大约 25KB 处截断）
        if content.len() > MAX_ENTRYPOINT_BYTES {
            let truncated: String = content
                .char_indices()
                .take_while(|(i, _)| *i < MAX_ENTRYPOINT_BYTES - 100)
                .map(|(_, c)| c)
                .collect();
            content = format!(
                "{}\n\n[MEMORY.md 已达到 {}KB 上限，旧条目已截断]",
                truncated,
                MAX_ENTRYPOINT_BYTES / 1024
            );
        }

        fs::write(&index_path, &content)?;
        Ok(())
    }

    /// 生成 system prompt 中注入的记忆段落。
    ///
    /// 包含最近的记忆条目（最多 10 条），并在前面注明可信度提醒。
    pub fn build_system_reminder(&self, max_entries: usize) -> Option<String> {
        if self.entries.is_empty() {
            return None;
        }

        let mut lines = vec![
            "<project-memory-context>".to_string(),
            "".to_string(),
            "[PROJECT MEMORY]".to_string(),
            "".to_string(),
        ];

        let count = max_entries.min(self.entries.len());
        for entry in self.entries.iter().take(count) {
            let type_tag = match entry.memory_type {
                MemoryType::User => "[偏好]",
                MemoryType::Feedback => "[反馈]",
                MemoryType::Project => "[项目]",
                MemoryType::Reference => "[参考]",
            };
            // historical 记忆标记
            let status_tag = if !entry.status.is_active() {
                match entry.status.superseded_by() {
                    Some(new_slug) => format!(" [已更新→{}]", new_slug),
                    None => " [已更新]".to_string(),
                }
            } else {
                String::new()
            };

            let preview: String = entry
                .content
                .lines()
                .take(3)
                .collect::<Vec<_>>()
                .join("\n  ");
            lines.push(format!(
                "{} {}{}: {}\n  {}",
                type_tag, entry.name, status_tag, entry.description, preview
            ));
            // 注入时效警告（超过1天的记忆）
            let freshness = age::memory_freshness_note(entry.mtime);
            if !freshness.is_empty() {
                lines.push(freshness);
            }
        }

        lines.push("".to_string());
        lines.push(
            "使用 save_memory 工具保存重要信息。记忆内容可能过时，请在使用前验证。".to_string(),
        );
        // 注入验证提醒（从记忆推荐前先核实）
        lines.push(String::new());
        lines.push(guardrails::build_verification_reminder());
        lines.push("</project-memory-context>".to_string());

        Some(lines.join("\n"))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_build_empty_reminder() {
        let manager = MemoryManager {
            memory_dir: PathBuf::from("/tmp/nonexistent"),
            entries: Vec::new(),
            extraction_tracker: extraction::ExtractionTracker::default(),
        };
        assert!(manager.build_system_reminder(10).is_none());
    }
}