281 lines
9.4 KiB
Rust
281 lines
9.4 KiB
Rust
//! 工具命名空间存储 — 大工具结果不进主消息队列
|
||
//!
|
||
//! 当工具执行结果超过阈值时，存入 `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<String, NamespaceEntry>,
|
||
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::<Vec<_>>().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);
|
||
}
|
||
}
|