//! 工具命名空间存储 — 大工具结果不进主消息队列 //! //! 当工具执行结果超过阈值时,存入 `NamespaceStore`,主队列只保留轻量引用路径。 //! 运行时为内存 HashMap(LRU 淘汰),持久化时由调用方展开引用写 DB。 //! //! 存储格式: "namespace://tool_name/args_hash" //! 例: namespace://read_file/a1b2c3d4 use std::collections::HashMap; use std::time::Instant; /// 引用路径前缀 pub const NAMESPACE_REF_PREFIX: &str = "namespace://"; /// namespace 条目被 LRU 淘汰(或跨会话残留引用)后,展开点 read_only()=None 时的提示文案。 /// /// 旧实现 None 分支保留 "namespace://" 字面 URI,LLM 收到无意义串且无告警。 /// 现统一替换为此文案:既告知用户结果已淘汰,又提示可重新调用工具取回。 pub const EVICTED_PLACEHOLDER: &str = "[该工具结果已超出内存上限被淘汰,如需请重新调用对应工具]"; /// namespace 字节阈值(> 2048 bytes 触发) pub const NAMESPACE_BYTE_THRESHOLD: usize = 2048; /// namespace 行数阈值(> 50 行 触发) pub const NAMESPACE_LINE_THRESHOLD: usize = 50; /// 始终进 namespace 的工具名(其输出几乎总是大结果) pub const ALWAYS_LARGE_TOOLS: &[&str] = &["read_file", "list_directory", "grep", "diff_files"]; /// 判定工具结果是否应进入 namespace pub fn should_use_namespace(content: &str, tool_name: &str) -> bool { if ALWAYS_LARGE_TOOLS.contains(&tool_name) { return true; } content.len() > NAMESPACE_BYTE_THRESHOLD || content.lines().count() > NAMESPACE_LINE_THRESHOLD } /// 判定 tool_result 内容是否为 namespace 引用路径 pub fn is_namespace_ref(content: &str) -> bool { content.starts_with(NAMESPACE_REF_PREFIX) } /// 从引用路径中提取内部 key pub fn parse_namespace_key(full_path: &str) -> Option<&str> { full_path.strip_prefix(NAMESPACE_REF_PREFIX) .and_then(|s| s.split('/').nth(1)) } /// namespace 条目 #[derive(Debug, Clone)] pub struct NamespaceEntry { pub key: String, pub tool_name: String, pub content: String, pub created_at: Instant, pub access_count: u64, } /// namespace 存储(内存 HashMap + LRU 淘汰) #[derive(Debug, Clone)] pub struct NamespaceStore { entries: HashMap, max_bytes: usize, current_bytes: usize, } impl NamespaceStore { /// 创建 namespace 存储 pub fn new(max_bytes: usize) -> Self { Self { entries: HashMap::new(), max_bytes, current_bytes: 0, } } /// 默认大小:10MB pub fn default() -> Self { Self::new(10 * 1024 * 1024) } /// 存入 namespace。返回引用路径字符串。 pub fn store(&mut self, tool_name: &str, content: &str) -> String { let key = simple_hash(content); let path = format!("{}{}/{}", NAMESPACE_REF_PREFIX, tool_name, key); // 已存在 → 刷新 access_count,不重复占用空间 if self.entries.contains_key(&key) { if let Some(entry) = self.entries.get_mut(&key) { entry.access_count += 1; } return path; } // 检查容量,超限淘汰 let new_bytes = content.len(); self.evict_if_needed(new_bytes); let entry = NamespaceEntry { key: key.clone(), tool_name: tool_name.to_string(), content: content.to_string(), created_at: Instant::now(), access_count: 1, }; self.current_bytes += new_bytes; self.entries.insert(key, entry); path } /// 读取完整内容 pub fn read(&mut self, full_path: &str) -> Option<&str> { let key = parse_namespace_key(full_path)?; let entry = self.entries.get_mut(key)?; entry.access_count += 1; Some(entry.content.as_str()) } /// 读取(只读,不改 access_count) pub fn read_only(&self, full_path: &str) -> Option<&str> { let key = parse_namespace_key(full_path)?; self.entries.get(key).map(|e| e.content.as_str()) } /// 检查 key 是否存在 pub fn contains(&self, full_path: &str) -> bool { parse_namespace_key(full_path) .and_then(|k| self.entries.get(k)) .is_some() } /// 条目数 pub fn len(&self) -> usize { self.entries.len() } pub fn is_empty(&self) -> bool { self.entries.is_empty() } /// LRU 淘汰:删除最旧条目直到腾出所需空间 fn evict_if_needed(&mut self, needed: usize) { while self.current_bytes + needed > self.max_bytes && !self.entries.is_empty() { // 找最旧的条目 let oldest_key = self.entries.iter() .min_by_key(|(_, e)| e.created_at) .map(|(k, _)| k.clone()); if let Some(key) = oldest_key { if let Some(entry) = self.entries.remove(&key) { self.current_bytes = self.current_bytes.saturating_sub(entry.content.len()); } } } } /// 清空 pub fn clear(&mut self) { self.entries.clear(); self.current_bytes = 0; } } /// 内容 → 引用 key 段(确定性,不含 `/` 等会破坏 `parse_namespace_key` 的字符)。 /// /// 历史 bug:旧实现用 `"{len}_{前16字符}"`,前 16 字符原样进 key;内容前缀含 `/`(目录列表、 /// 带路径的 grep 结果、文件头是路径)时,`namespace://tool/{key}` 被 `parse_namespace_key` /// 的 `split('/').nth(1)` 从首个 `/` 截断 → read/read_only 拿回 None → 原文彻底丢失。 /// 现改用 SipHash u64 → 16 位 hex:确定性 + 极低碰撞 + 不含特殊字符。 fn simple_hash(content: &str) -> String { use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); content.hash(&mut hasher); format!("{:016x}", hasher.finish()) } #[cfg(test)] mod tests { use super::*; #[test] fn should_use_namespace_small_content() { assert!(!should_use_namespace("short", "custom_tool")); } #[test] fn should_use_namespace_large_bytes() { let large = "x".repeat(3000); assert!(should_use_namespace(&large, "custom_tool")); } #[test] fn should_use_namespace_always_large_tool() { assert!(should_use_namespace("short", "read_file")); assert!(should_use_namespace("short", "list_directory")); assert!(should_use_namespace("short", "diff_files")); } #[test] fn should_use_namespace_many_lines() { let many = (0..60).map(|i| format!("line {}", i)).collect::>().join("\n"); assert!(should_use_namespace(&many, "custom_tool")); } #[test] fn is_namespace_ref_detects_prefix() { assert!(is_namespace_ref("namespace://read_file/abc")); assert!(!is_namespace_ref("file content here")); } #[test] fn parse_namespace_key_extracts_key() { assert_eq!(parse_namespace_key("namespace://read_file/abc_123"), Some("abc_123")); assert_eq!(parse_namespace_key("namespace://grep/xyz"), Some("xyz")); assert_eq!(parse_namespace_key("no_prefix"), None); } #[test] fn store_and_read() { let mut ns = NamespaceStore::new(100_000); let path = ns.store("read_file", "hello world content"); assert!(is_namespace_ref(&path)); assert!(path.contains("read_file")); assert_eq!(ns.read(&path), Some("hello world content")); } #[test] fn store_duplicate_key_does_not_double_count() { let mut ns = NamespaceStore::new(100_000); let content = "test content for duplicate check"; let p1 = ns.store("tool_a", content); let p2 = ns.store("tool_b", content); assert_eq!(ns.len(), 1, "相同内容应去重"); // read(&mut self) 返回 Option<&str> 借用 ns,两次调用须各自转 owned 避免双重可变借用 let r1 = ns.read(&p1).map(str::to_owned); let r2 = ns.read(&p2).map(str::to_owned); assert_eq!(r1, r2); } #[test] fn store_content_with_slash_in_prefix_reads_back() { // 回归:旧 simple_hash 把内容前 16 字符原样进 key,前缀含 '/' 时(目录列表/带路径 grep // 结果/文件头是路径)namespace://tool/{key} 被 parse_namespace_key 的 split('/') 截断 // → read None,原文丢失。hex hash 修复后 store→read 正确取回。 let mut ns = NamespaceStore::new(100_000); let content = "foo/bar/baz/qux\nline2\nline3"; let path = ns.store("list_directory", content); assert!(is_namespace_ref(&path)); assert_eq!(ns.read(&path).map(str::to_owned), Some(content.to_string())); } #[test] fn eviction_oldest_removed_when_over_limit() { let mut ns = NamespaceStore::new(100); // 极小容量 let small = "a"; // 1 byte ns.store("t1", small); assert_eq!(ns.len(), 1); // 填充到超限 let big = "x".repeat(200); ns.store("t2", &big); // t1 可能被淘汰 if ns.current_bytes > 100 { assert!(ns.len() < 2, "超限时应已淘汰旧条目"); } } #[test] fn clear_resets_all() { let mut ns = NamespaceStore::new(100_000); ns.store("t1", "content"); assert!(!ns.is_empty()); ns.clear(); assert!(ns.is_empty()); assert_eq!(ns.current_bytes, 0); } }