AstroResearch/src/agent/runtime/file_cache.rs

// src/agent/runtime/file_cache.rs
//
// 文件状态缓存 — 参考 Claude Code FileStateCache 设计。
//
// 在 Read 工具调用前检查缓存：
//   1. 路径已缓存 → 读取磁盘 mtime → mtime 相同 + offset/limit 一致 → 返回 stub
//   2. mtime 不同或新文件 → 正常读取 → 写入缓存
//
// 压缩时：
//   - 压缩前：快照缓存到普通对象
//   - 压缩后：清空缓存，将最近 N 个文件作为上下文注入
//
// 缓存上限：100 个条目，25MB 内容总大小（LRU 自动淘汰）。

use lru::LruCache;
use std::num::NonZeroUsize;
use tracing::{info, warn};

/// 缓存条目最大数量
pub const MAX_ENTRIES: usize = 100;
/// 缓存内容总大小上限（25MB）
pub const MAX_CACHE_SIZE_BYTES: usize = 25 * 1024 * 1024;

/// 文件不变时的占位消息（参考 Claude Code FILE_UNCHANGED_STUB）
pub const FILE_UNCHANGED_STUB: &str =
    "File unchanged since last read. The content from the earlier read_file tool_result \
     in this conversation is still current — refer to that instead of re-reading.";

/// 压缩后恢复的最大文件数
pub const POST_COMPACT_MAX_FILES_TO_RESTORE: usize = 5;
/// 压缩后恢复的每文件最大 token 数（~字符数）
pub const POST_COMPACT_MAX_CHARS_PER_FILE: usize = 4_000;

/// 单个文件的缓存状态
#[derive(Debug, Clone)]
pub struct FileState {
    /// 上次读取的文件内容
    pub content: String,
    /// 文件修改时间（Unix 时间戳，秒级）
    pub timestamp: i64,
    /// 读取起始行（1-based）
    pub offset: usize,
    /// 行数限制
    pub limit: Option<usize>,
}

/// 文件状态快照（用于压缩前后传递，纯数据，不含 LRU 结构）
pub type FileStateSnapshot = Vec<(String, FileState)>;

/// 文件状态缓存。
///
/// 包装 `LruCache<String, FileState>` + 内容总大小追踪。
/// 通过 `Arc<Mutex<FileStateCache>>` 在工具调用间共享。
pub struct FileStateCache {
    cache: LruCache<String, FileState>,
    /// 当前缓存中所有内容的字节数总和（近似，使用 content.len()）
    current_size_bytes: usize,
    /// 最大字节数
    max_size_bytes: usize,
}

impl FileStateCache {
    /// 创建新的缓存实例
    pub fn new() -> Self {
        let max_entries = NonZeroUsize::new(MAX_ENTRIES).unwrap();
        FileStateCache {
            cache: LruCache::new(max_entries),
            current_size_bytes: 0,
            max_size_bytes: MAX_CACHE_SIZE_BYTES,
        }
    }

    /// 创建带自定义参数的新缓存
    pub fn with_limits(max_entries: usize, max_size_bytes: usize) -> Self {
        let max_entries =
            NonZeroUsize::new(max_entries).unwrap_or(NonZeroUsize::new(MAX_ENTRIES).unwrap());
        FileStateCache {
            cache: LruCache::new(max_entries),
            current_size_bytes: 0,
            max_size_bytes,
        }
    }

    /// 规范化路径 key（确保一致性）—— 仅做字符串规范化，不做磁盘 I/O
    fn normalize_key(path: &str) -> String {
        // 简单规范化：折叠重复的 / 和 \
        let mut result = String::with_capacity(path.len());
        let mut prev_slash = false;
        for ch in path.chars() {
            if ch == '/' || ch == '\\' {
                if !prev_slash {
                    result.push('/');
                    prev_slash = true;
                }
            } else {
                result.push(ch);
                prev_slash = false;
            }
        }
        // 去除尾随 /
        if result.ends_with('/') && result.len() > 1 {
            result.pop();
        }
        result
    }

    /// 获取缓存的条目，返回 Some(&FileState) 若存在
    pub fn get(&mut self, path: &str) -> Option<&FileState> {
        let key = Self::normalize_key(path);
        self.cache.get(&key)
    }

    /// 写入缓存条目，自动处理容量限制
    pub fn set(&mut self, path: &str, state: FileState) {
        let key = Self::normalize_key(path);
        let content_len = state.content.len();

        // 如果 key 已存在，先减去旧内容的 size
        if let Some(old) = self.cache.get(&key) {
            self.current_size_bytes = self.current_size_bytes.saturating_sub(old.content.len());
        }

        // 驱逐旧条目直到有足够空间
        while self.current_size_bytes + content_len > self.max_size_bytes && !self.cache.is_empty()
        {
            if let Some((_, evicted)) = self.cache.pop_lru() {
                self.current_size_bytes = self
                    .current_size_bytes
                    .saturating_sub(evicted.content.len());
            }
        }

        // 如果单个文件超过上限，仍存储但记录警告
        if content_len > self.max_size_bytes {
            warn!(
                "[FileCache] 单个文件内容 ({} bytes) 超过缓存上限 ({} bytes)",
                content_len, self.max_size_bytes
            );
        }

        self.current_size_bytes += content_len;
        self.cache.push(key, state);

        info!(
            "[FileCache] 缓存写入: path={}, size={} bytes, cache_entries={}, cache_size={}",
            path,
            content_len,
            self.cache.len(),
            self.current_size_bytes
        );
    }

    /// 检查 key 是否存在
    pub fn contains(&mut self, path: &str) -> bool {
        let key = Self::normalize_key(path);
        self.cache.contains(&key)
    }

    /// 删除缓存条目
    pub fn remove(&mut self, path: &str) -> bool {
        let key = Self::normalize_key(path);
        if let Some(removed) = self.cache.pop(&key) {
            self.current_size_bytes = self
                .current_size_bytes
                .saturating_sub(removed.content.len());
            true
        } else {
            false
        }
    }

    /// 清空缓存
    pub fn clear(&mut self) {
        self.cache.clear();
        self.current_size_bytes = 0;
        info!("[FileCache] 缓存已清空");
    }

    /// 缓存条目数
    pub fn len(&self) -> usize {
        self.cache.len()
    }

    /// 缓存是否为空
    pub fn is_empty(&self) -> bool {
        self.cache.len() == 0
    }

    /// 当前缓存内容总大小（近似字节数）
    pub fn current_size_bytes(&self) -> usize {
        self.current_size_bytes
    }

    // ── 压缩集成 ──

    /// 生成快照（纯数据，不含 LRU 结构）。
    /// 在压缩前调用，用于压缩后恢复文件上下文。
    pub fn to_snapshot(&mut self) -> FileStateSnapshot {
        // 按 timestamp 降序排序（最近读的在前）
        let mut entries: Vec<(String, FileState)> = self
            .cache
            .iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        entries.sort_by(|a, b| b.1.timestamp.cmp(&a.1.timestamp));
        entries
    }

    /// 从快照恢复指定数量的最近文件到缓存。
    /// 在压缩后调用。
    pub fn restore_from_snapshot(&mut self, snapshot: &FileStateSnapshot, max_files: usize) {
        for (path, state) in snapshot.iter().take(max_files) {
            // 不恢复过大的文件（已有提示说可能过时）
            if state.content.len() > POST_COMPACT_MAX_CHARS_PER_FILE {
                continue;
            }
            self.set(path, state.clone());
        }
        info!(
            "[FileCache] 从快照恢复了 {} 个文件 (快照大小: {})",
            self.len().min(max_files),
            snapshot.len()
        );
    }

    /// 从快照生成上 下文注入文本（用于压缩后注入到对话中）。
    /// 返回格式化的 markdown 块列表。
    pub fn build_restore_context(snapshot: &FileStateSnapshot, max_files: usize) -> Vec<String> {
        if snapshot.is_empty() {
            return Vec::new();
        }

        let mut contexts: Vec<String> = Vec::new();
        let mut used_chars = 0usize;
        let total_budget = POST_COMPACT_MAX_CHARS_PER_FILE * max_files;

        for (path, state) in snapshot.iter().take(max_files) {
            let preview: String = state
                .content
                .chars()
                .take(POST_COMPACT_MAX_CHARS_PER_FILE)
                .collect();
            let truncated = if state.content.len() > preview.len() {
                format!(
                    "{}…\n[内容已截断: {} 字符 → {} 字符]",
                    preview,
                    state.content.len(),
                    preview.len()
                )
            } else {
                preview
            };

            if used_chars + truncated.len() > total_budget {
                break;
            }

            let block = format!(
                "[压缩后恢复: {}]\n上次读取时间戳: {}\n内容:\n```\n{}\n```",
                path, state.timestamp, truncated
            );
            used_chars += block.len();
            contexts.push(block);
        }

        if !contexts.is_empty() {
            info!(
                "[FileCache] 生成压缩恢复上下文: {} 文件, {} chars",
                contexts.len(),
                used_chars
            );
        }

        contexts
    }
}

impl Default for FileStateCache {
    fn default() -> Self {
        Self::new()
    }
}

/// 获取文件的当前 mtime（Unix 时间戳，秒）。
/// 失败时返回 None。
pub fn get_file_mtime(path: &str) -> Option<i64> {
    match std::fs::metadata(path) {
        Ok(meta) => match meta.modified() {
            Ok(time) => match time.duration_since(std::time::UNIX_EPOCH) {
                Ok(d) => Some(d.as_secs() as i64),
                Err(_) => None,
            },
            Err(_) => None,
        },
        Err(_) => None,
    }
}

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

    #[test]
    fn test_new_cache_is_empty() {
        let cache = FileStateCache::new();
        assert!(cache.is_empty());
        assert_eq!(cache.len(), 0);
        assert_eq!(cache.current_size_bytes(), 0);
    }

    #[test]
    fn test_set_and_get() {
        let mut cache = FileStateCache::new();
        cache.set(
            "/tmp/test.txt",
            FileState {
                content: "hello world".to_string(),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );
        assert!(!cache.is_empty());

        let entry = cache.get("/tmp/test.txt");
        assert!(entry.is_some());
        assert_eq!(entry.unwrap().content, "hello world");
    }

    #[test]
    fn test_path_normalization() {
        let mut cache = FileStateCache::new();
        cache.set(
            "/tmp//test.txt",
            FileState {
                content: "test".to_string(),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );
        // 规范化后的路径应该能命中
        assert!(cache.get("/tmp/test.txt").is_some());
    }

    #[test]
    fn test_contains() {
        let mut cache = FileStateCache::new();
        assert!(!cache.contains("/tmp/test.txt"));

        cache.set(
            "/tmp/test.txt",
            FileState {
                content: "test".to_string(),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );
        assert!(cache.contains("/tmp/test.txt"));
    }

    #[test]
    fn test_remove() {
        let mut cache = FileStateCache::new();
        cache.set(
            "/tmp/test.txt",
            FileState {
                content: "test".to_string(),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );
        assert!(cache.remove("/tmp/test.txt"));
        assert!(cache.is_empty());
        assert!(!cache.remove("/tmp/test.txt"));
    }

    #[test]
    fn test_clear() {
        let mut cache = FileStateCache::new();
        cache.set(
            "/tmp/a.txt",
            FileState {
                content: "a".to_string(),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );
        cache.set(
            "/tmp/b.txt",
            FileState {
                content: "b".to_string(),
                timestamp: 2000,
                offset: 1,
                limit: None,
            },
        );
        cache.clear();
        assert!(cache.is_empty());
    }

    #[test]
    fn test_snapshot_sorted_by_timestamp_desc() {
        let mut cache = FileStateCache::new();
        cache.set(
            "/tmp/old.txt",
            FileState {
                content: "old".to_string(),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );
        cache.set(
            "/tmp/new.txt",
            FileState {
                content: "new".to_string(),
                timestamp: 3000,
                offset: 1,
                limit: None,
            },
        );
        cache.set(
            "/tmp/mid.txt",
            FileState {
                content: "mid".to_string(),
                timestamp: 2000,
                offset: 1,
                limit: None,
            },
        );

        let snapshot = cache.to_snapshot();
        assert_eq!(snapshot.len(), 3);
        // 按 timestamp 降序
        assert_eq!(snapshot[0].1.timestamp, 3000);
        assert_eq!(snapshot[1].1.timestamp, 2000);
        assert_eq!(snapshot[2].1.timestamp, 1000);
    }

    #[test]
    fn test_restore_from_snapshot() {
        let mut cache = FileStateCache::new();
        cache.set(
            "/tmp/a.txt",
            FileState {
                content: "a".to_string(),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );
        cache.set(
            "/tmp/b.txt",
            FileState {
                content: "b".to_string(),
                timestamp: 2000,
                offset: 1,
                limit: None,
            },
        );

        let snapshot = cache.to_snapshot();
        cache.clear();

        cache.restore_from_snapshot(&snapshot, 1);
        assert_eq!(cache.len(), 1);
        // 应该恢复 timestamp 最高的
        assert!(cache.get("/tmp/b.txt").is_some());
    }

    #[test]
    fn test_build_restore_context() {
        let mut cache = FileStateCache::new();
        cache.set(
            "/tmp/a.txt",
            FileState {
                content: "file a content here".to_string(),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );

        let snapshot = cache.to_snapshot();
        let contexts = FileStateCache::build_restore_context(&snapshot, 2);
        assert_eq!(contexts.len(), 1);
        assert!(contexts[0].contains("/tmp/a.txt"));
        assert!(contexts[0].contains("file a content here"));
    }

    #[test]
    fn test_empty_snapshot_builds_no_context() {
        let contexts = FileStateCache::build_restore_context(&Vec::new(), 5);
        assert!(contexts.is_empty());
    }

    #[test]
    fn test_lru_eviction_on_size() {
        // 创建一个小容量缓存（仅 200 bytes）
        let mut cache = FileStateCache::with_limits(100, 200);

        // 写入 3 个 100 字节内容 → 应触发 LRU 淘汰
        cache.set(
            "/tmp/1.txt",
            FileState {
                content: "x".repeat(100),
                timestamp: 1000,
                offset: 1,
                limit: None,
            },
        );
        cache.set(
            "/tmp/2.txt",
            FileState {
                content: "y".repeat(100),
                timestamp: 2000,
                offset: 1,
                limit: None,
            },
        );
        // 此时应该有 2 个条目（200 bytes）
        assert_eq!(cache.len(), 2);

        // 写入第 3 个 → 应淘汰最旧的（1.txt）
        cache.set(
            "/tmp/3.txt",
            FileState {
                content: "z".repeat(100),
                timestamp: 3000,
                offset: 1,
                limit: None,
            },
        );
        // 旧条目被淘汰
        assert!(!cache.contains("/tmp/1.txt"));
    }

    #[test]
    fn test_get_file_mtime() {
        // 创建临时文件
        let tmp = std::env::temp_dir().join("test_mtime.txt");
        std::fs::write(&tmp, "test").unwrap();

        let mtime = get_file_mtime(&tmp.to_string_lossy());
        assert!(mtime.is_some());
        assert!(mtime.unwrap() > 0);

        // 不存在的文件
        let mtime = get_file_mtime("/tmp/nonexistent_12345_xxx.txt");
        assert!(mtime.is_none());

        std::fs::remove_file(&tmp).ok();
    }

    #[test]
    fn test_build_restore_context_respects_limit() {
        let mut snapshot: FileStateSnapshot = Vec::new();
        // 按时间倒序排列（调用方已排序，最新的在前）
        for i in (0..5).rev() {
            snapshot.push((
                format!("/tmp/file_{}.txt", i),
                FileState {
                    content: format!("content_{}", i),
                    timestamp: 1000 + i as i64,
                    offset: 0,
                    limit: None,
                },
            ));
        }

        let context = FileStateCache::build_restore_context(&snapshot, 2);
        assert!(
            !context.is_empty(),
            "should produce context when snapshot has files"
        );
        // build_restore_context 返回 Vec<String>，包含最多 max_files 条
        assert!(
            context.len() <= 2,
            "should return at most 2 entries, got {}",
            context.len()
        );
        // 应包含时间戳最大的文件（排在最前面）
        assert!(
            context.iter().any(|s| s.contains("content_4")),
            "should include the newest file content"
        );
    }
}