feat: Agent 安全纵深防御、Checkpoint 快照、会话 Rewind/Branch、自进化

Skill、流式执行优化与系统架构全面升级

  本次提交对标 Claude Code 与 Hermes-Agent 的工程细节,在安全、可靠性、
  会话管理、自我进化四个维度进行了系统性加固,变更总量 48 文件 / +12680 -2292 行。

  ═══════ 安全纵深防御 ═══════

  1. Hardline 硬阻止层 (src/agent/runtime/hardline.rs, +534 行)
     - 不可绕过的危险命令拦截(关重启、磁盘擦除、Fork 炸弹、rm -rf /、kill -1)
     - 反规避标准化管线: ANSI 序列剥离 → Unicode NFKC → shell 反斜杠还原 → 空字面量清理
     - 在 PermissionChecker 之前执行,YOLO/Bypass 模式下同样生效
     - 集成到 executor Phase 2,被拒绝工具直接注入错误结果

  2. Permission 优先级裁决器 (src/agent/runtime/permission.rs, +200 行)
     - 7 层正式优先级规则 (P0 Deny → P7 Allow),带冲突日志
     - explain() 方法支持审计追溯
     - Hook PermissionRequired 与 Checker 结果的正确叠加逻辑

  ═══════ Checkpoint 文件快照系统 ═══════

  3. git2 原生快照 (src/agent/runtime/checkpoint.rs, +920 行)
     - 基于 git2 bare repo,内容寻址自动去重
     - 文件变更操作前自动触发 (file_write/file_edit/run_bash)
     - 每目录每 turn 最多一次快照,防止同一轮重复
     - 支持 list/diff/restore API + pre-rollback 安全快照
     - 旧快照自动 prune(保留最近 N 个)+ 按目录隔离 ref
     - 排除规则自动过滤 node_modules/target/.git/*.pdf 等
     - 集成到 executor: 文件操作前 ckpt.ensure_checkpoint()

  ═══════ 错误恢复系统大升级 ═══════

  4. 21 种 FailoverReason 分类 (src/agent/runtime/error_recovery.rs, +1200 行)
     - 参考 Hermes-Agent error_classifier.py
     - 8 步分类管线: provider-specific → HTTP status → text pattern → error body → fallback
     - is_retryable / should_compress / should_failover / is_permanent 方法
     - Context Overflow 自动修复: 从错误消息提取 token 限制,自动下调预算
     - RecoveryStep::AdjustMaxTokens 实现 (参考 Claude Code 自动修复)
     - 向后兼容 ErrorKind 别名

  ═══════ 会话 Rewind / Branch / Retry 体系 ═══════

  5. 完整 undo 栈 (src/agent/runtime/session.rs, +800 行 + 2 迁移脚本)
     - Rewind (软删除): active=0 标记,审计 trail 保留,LLM 不可见
     - Restore (撤销回退): 冲突检测——回退后有新消息则拒绝,引导使用 Branch
     - Branch: 分叉会话,复制所有 active=1 消息到新会话
     - Retry: 硬删除最后一轮对话,返回原消息文本供前端重提交
     - 数据库: agent_messages.active 列 + agent_sessions.rewind_count + parent_session_id
     - API: 4 个新端点 (/branch, /retry, /rewind, /rewind/restore)
     - load_history_for_agent 全面使用 active=1 过滤

  ═══════ Hooks 系统模块化重构 ═══════

  6. 单文件 → 7 模块体系 (src/agent/hooks/)
     hooks.rs (994 行) 拆分为:
     - mod.rs    — 入口 + HookRegistry + SessionHookManager
     - types.rs  — 类型定义 (Context, TaggedContext, PermissionRequestAction 等)
     - traits.rs — AgentHook + AsyncAgentHook + 15 种生命周期事件
     - matcher.rs — 工具名/参数匹配 + session 作用域过滤
     - dispatch.rs — 并行调度引擎 (run_pre/post_tool_use 等)
     - registry.rs — 注册/注销/查询
     - builtins.rs — CancellationHook + MetricsHook + AuditLogHook + ContextDeduplicator

     关键改进:
     - run_pre_tool_use 并行执行所有匹配 hooks,聚合 Block/MutateInput/Continue
     - TaggedContext 带完整来源标记的上下文注入 (hook_name + event)
     - ContextDeduplicator 单 dispatch cycle 内内容哈希去重
     - AsyncAgentHook 支持 fire-and-forget 异步 hooks

  ═══════ Executor 并发执行升级 ═══════

  7. 三阶段管道重写 (src/agent/runtime/executor.rs, +600 行)
     - Phase 1: 死循环检测 + 参数解析 (不变)
     - Phase 2: Hardline 预检查 (新增) → PermissionChecker (改进)
     - Phase 3: ToolPartitioner 分区 → 逐批次执行 (重写)
       - 并行批次内 FuturesUnordered 并发
       - 串行批次确保非并发安全工具独占执行
       - Checkpoint 预触发集成
     - Hook 上下文注入: system-reminder 格式 + ContextDeduplicator 去重
     - Hook 阻塞错误详细记录

  ═══════ 流式执行真正的流式调度 ═══════

  8. StreamingExecutor 重写 (src/agent/runtime/streaming_executor.rs, ~400 行变更)
     - on_tool_use 中对并发安全工具立即 tokio::spawn,不等待 flush
     - executing_non_concurrent 标志阻塞后继工具直到独占工具完成
     - JoinHandle 管理替代自定义 cancel channel
     - completed_queue 按流顺序 yield
     - Sibling Abort 通过 broadcast channel + tokio::select! 竞速
     - ToolContext 实现 Clone (支持 per-task 上下文复制)

  ═══════ 自改进 Skill 系统 ═══════

  9. PatternDetector + SkillCreator + Curator (src/agent/skills/, +1500 行)
     - PatternDetector: 扫描 agent_messages 表,检测跨 session 重复工具调用模式
     - SkillCreator: 将高置信度模式自动生成 SKILL.md (YAML frontmatter + 工作流步骤)
     - SelfImprovePipeline: 一站式 模式检测 → 创建 → 质量审查
     - Curator: 分析 skill 使用统计,标记 stale/deprecated,建议清理
     - Skill frontmatter 新增 pinned 字段 (禁止 Curator 自动清理)

  ═══════ 基础设施优化 ═══════

  10. 系统提示词缓存 (src/agent/runtime/system_prompt.rs + mod.rs)
      - SystemPromptCache: 首次计算后永久复用,/clear 时失效
      - 新增 SAFETY / SYSTEM_CONTEXT / TOOL_USAGE 静态 section
      - 环境/tools/skills/memory 动态 section 通过 get_or_compute 缓存

  11. ToolRegistry schema 缓存 (src/agent/tools/mod.rs)
      - schema_cache + schema_generation 版本号
      - 工具变更/过滤器变更时自动失效
      - precompute_definitions() 预计算 (AgentRuntime 初始化时调用)

  12. 迭代摘要融合 (src/agent/compact.rs, +100 行)
      - 参考 Hermes context_compressor.py
      - CollapseLog 追踪压缩历史,支持溢出合并
      - extract_prior_summary: 提取已有摘要融入新压缩

  13. SubAgent 系统提示词模块化 (src/agent/tools/subagent.rs)
      - 复用 5 个标准 section + 子代理专有上下文 section
      - 独立 ToolRegistry 构建工具列表

  ═══════ 前端 — CSS 变量主题系统 ═══════

  14. 全新主题变量体系 (dashboard/src/index.css + App.tsx + 各面板)
      - CSS 自定义属性: --bg-card, --text-main, --text-muted, --border-precision
      - 语义化颜色: --accent-blueprint, --accent-star
      - 全面替换硬编码 Tailwind 颜色 (slate-xxx → var(--xxx))
      - 文献入库提示优化 ("核心知识节点" 替代 "向量块")
      - ReaderPanel 样式变量化
This commit is contained in:
fmq
2026-06-22 20:29:37 +08:00
parent f6df9d8136
commit 698d007f39
48 changed files with 12706 additions and 2318 deletions
+920
View File
@@ -0,0 +1,920 @@
// src/agent/runtime/checkpoint.rs
//
// Checkpoint 系统 — 透明的文件系统快照。
// 参考 Hermes-Agent checkpoint_manager.py 设计,使用 git2 crate 实现。
//
// 设计原则:
// 1. 对 LLM 完全透明 — 不是 AgentTool,是基础设施
// 2. 文件变更操作前自动触发,每目录每 turn 最多一次
// 3. 使用 git2 原生庫而非 subprocess,零命令行注入风险
// 4. 单一 bare repo 存储,内容寻址自动去重
//
// 存储布局:
// .checkpoints/ ← bare git 仓库(与 library_dir 同级)
// HEAD, config, objects/ ← git 内部
// refs/checkpoints/<hash> ← 每个工作目录的分支
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use tracing::{debug, info, warn};
/// 默认排除的文件/目录模式(gitignore 格式)
const DEFAULT_EXCLUDES: &[&str] = &[
// 依赖 / 构建产物
"node_modules/",
"dist/",
"build/",
"target/",
"out/",
".next/",
// 缓存
"__pycache__/",
"*.pyc",
".cache/",
".pytest_cache/",
".mypy_cache/",
".ruff_cache/",
// 虚拟环境
".venv/",
"venv/",
// VCS
".git/",
".hg/",
".svn/",
// 编译产物
"*.so",
"*.dylib",
"*.dll",
"*.o",
"*.a",
"*.exe",
"*.obj",
// 大文件
"*.mp4",
"*.mov",
"*.mkv",
"*.zip",
"*.tar",
"*.tar.gz",
"*.tgz",
"*.7z",
"*.pdf", // 论文 PDF 已有 library/ 备份,不纳入快照
// 敏感文件
".env",
".env.*",
"*.log",
// OS 垃圾
".DS_Store",
"Thumbs.db",
];
/// 每个 turn 最多快照一次的工具
const CHECKPOINT_TRIGGER_TOOLS: &[&str] = &["file_write", "file_edit", "run_bash"];
/// Checkpoint 元数据
#[derive(Debug, Clone)]
pub struct CheckpointEntry {
/// 完整 commit hash
pub hash: String,
/// 短 hash(前 8 位)
pub short_hash: String,
/// ISO 8601 时间戳
pub timestamp: String,
/// 快照原因
pub reason: String,
/// 变更文件数
pub files_changed: usize,
/// 插入行数
pub insertions: usize,
/// 删除行数
pub deletions: usize,
}
/// Checkpoint 管理器
///
/// 使用 git2 创建和维护一个 bare git 仓库用于文件系统快照。
/// 每个被监控的工作目录在 `refs/checkpoints/<dir_hash>` 下有独立的分支。
///
/// 注意:`git2::Repository` 是 `!Sync`,所以内部状态通过 `Mutex<InnerState>` 保护。
pub struct CheckpointManager {
/// 主开关
enabled: bool,
/// bare git 仓库路径
repo_path: PathBuf,
/// 每个项目保留的最大快照数
max_snapshots: usize,
/// 单个文件大小上限(字节),超过此大小不纳入快照
max_file_size: usize,
/// 内部可变状态(Send + Sync
inner: Mutex<InnerState>,
}
/// 需要 Mutex 保护的可变状态
struct InnerState {
/// git2 仓库句柄
repo: Option<git2::Repository>,
/// 本 turn 已快照的目录集合
checkpointed_this_turn: HashSet<PathBuf>,
}
impl InnerState {
fn new() -> Self {
InnerState {
repo: None,
checkpointed_this_turn: HashSet::new(),
}
}
}
impl CheckpointManager {
/// 创建新的 checkpoint 管理器。
///
/// `store_path` 是 bare repo 的路径(建议: `<library_dir>/../.checkpoints`)。
/// 如果 `enabled` 为 false 或 git 不可用,所有操作静默跳过。
pub fn new(store_path: PathBuf, enabled: bool) -> Self {
let mut inner = InnerState::new();
if enabled {
match Self::init_store(&store_path) {
Ok(r) => {
info!("[Checkpoint] 仓库已初始化: {}", store_path.display());
inner.repo = Some(r);
}
Err(e) => {
warn!("[Checkpoint] 仓库初始化失败,快照功能已禁用: {}", e);
}
}
}
CheckpointManager {
enabled: enabled && inner.repo.is_some(),
repo_path: store_path,
max_snapshots: std::env::var("AGENT_CHECKPOINT_MAX_SNAPSHOTS")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(10),
max_file_size: std::env::var("AGENT_CHECKPOINT_MAX_FILE_SIZE_MB")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.map(|mb| mb * 1024 * 1024)
.unwrap_or(10 * 1024 * 1024), // 10 MB
inner: Mutex::new(inner),
}
}
// ── 初始化 ──
/// 初始化 bare git 仓库,设置排除规则。
fn init_store(store_path: &Path) -> Result<git2::Repository, anyhow::Error> {
// 创建父目录
if let Some(parent) = store_path.parent() {
std::fs::create_dir_all(parent)?;
}
// 初始化或打开 bare repo
if store_path.join("HEAD").exists() {
return Ok(git2::Repository::open(store_path)?);
}
let repo = git2::Repository::init_bare(store_path)?;
// 写入 .gitignore 排除规则
let exclude_path = store_path.join("info").join("exclude");
if let Some(parent) = exclude_path.parent() {
std::fs::create_dir_all(parent)?;
}
let exclude_content = DEFAULT_EXCLUDES.join("\n") + "\n";
std::fs::write(&exclude_path, exclude_content)?;
// 设置仓库级配置:禁用 gpgsign
if let Ok(mut config) = repo.config() {
let _ = config.set_str("user.email", "checkpoint@astroresearch.local");
let _ = config.set_str("user.name", "AstroResearch Checkpoint");
let _ = config.set_str("commit.gpgsign", "false");
let _ = config.set_str("tag.gpgSign", "false");
let _ = config.set_str("gc.auto", "0");
}
Ok(repo)
}
// ── Turn 生命周期 ──
/// 重置 per-turn 去重状态。每个 ReAct 循环迭代前调用。
pub fn new_turn(&self) {
if !self.enabled {
return;
}
if let Ok(mut inner) = self.inner.lock() {
inner.checkpointed_this_turn.clear();
}
}
// ── 公共 API ──
/// 确保工作目录已被快照。如果是本 turn 首次对该目录调用且 enabled,
/// 则创建快照。返回是否实际创建了快照。
///
/// 永远不 panic — 所有错误静默记录日志。
pub fn ensure_checkpoint(&self, working_dir: &Path, reason: &str) -> bool {
if !self.enabled {
return false;
}
let abs_dir = match std::fs::canonicalize(working_dir) {
Ok(d) => d,
Err(e) => {
debug!(
"[Checkpoint] 无法解析目录 '{}': {}",
working_dir.display(),
e
);
return false;
}
};
// 跳过根目录和 home 目录
if abs_dir == Path::new("/") || abs_dir == dirs_home() {
debug!("[Checkpoint] 跳过过于宽泛的目录: {}", abs_dir.display());
return false;
}
// 每 turn 每目录去重
{
let mut inner = match self.inner.lock() {
Ok(s) => s,
Err(e) => {
warn!("[Checkpoint] 锁异常: {}", e);
return false;
}
};
if inner.checkpointed_this_turn.contains(&abs_dir) {
return false;
}
inner.checkpointed_this_turn.insert(abs_dir.clone());
}
match self.take_snapshot(&abs_dir, reason) {
Ok(true) => {
info!(
"[Checkpoint] 快照完成: {} (reason={})",
abs_dir.display(),
reason
);
true
}
Ok(false) => false, // 无变更
Err(e) => {
debug!("[Checkpoint] 快照失败(非致命): {}", e);
false
}
}
}
/// 列出指定工作目录的所有快照。
pub fn list_checkpoints(&self, working_dir: &Path) -> Vec<CheckpointEntry> {
let inner = match self.inner.lock() {
Ok(i) => i,
Err(_) => return Vec::new(),
};
let repo = match &inner.repo {
Some(r) => r,
None => return Vec::new(),
};
let abs_dir = match std::fs::canonicalize(working_dir) {
Ok(d) => d,
Err(_) => return Vec::new(),
};
let ref_name = dir_ref_name(&abs_dir);
// 查找该 ref 的所有 commit
let mut revwalk = match repo.revwalk() {
Ok(r) => r,
Err(_) => return Vec::new(),
};
let _ = revwalk.push_ref(&ref_name);
let _ = revwalk.set_sorting(git2::Sort::TIME);
let mut entries = Vec::new();
for oid_result in revwalk {
let oid = match oid_result {
Ok(o) => o,
Err(_) => continue,
};
let commit = match repo.find_commit(oid) {
Ok(c) => c,
Err(_) => continue,
};
let time = commit.time();
let timestamp = chrono::DateTime::from_timestamp(time.seconds(), 0)
.map(|dt| dt.format("%Y-%m-%dT%H:%M:%S").to_string())
.unwrap_or_else(|| "unknown".to_string());
let reason = commit.message().unwrap_or("checkpoint").to_string();
let short_hash = oid.to_string()[..8].to_string();
// 统计变更(与父 commit 比较)
let (files_changed, insertions, deletions) = if commit.parent_count() > 0 {
let parent = commit.parent(0).ok();
let parent_tree = parent.and_then(|p| p.tree().ok());
let this_tree = commit.tree().ok();
match (parent_tree, this_tree) {
(Some(pt), Some(tt)) => {
match repo.diff_tree_to_tree(Some(&pt), Some(&tt), None) {
Ok(diff) => match diff.stats() {
Ok(stats) => {
(stats.files_changed(), stats.insertions(), stats.deletions())
}
Err(_) => (0, 0, 0),
},
Err(_) => (0, 0, 0),
}
}
_ => (0, 0, 0),
}
} else {
// 初始快照:统计所有文件
let tree = commit.tree().ok();
(tree.map(|t| t.len()).unwrap_or(0), 0, 0)
};
entries.push(CheckpointEntry {
hash: oid.to_string(),
short_hash,
timestamp,
reason,
files_changed,
insertions,
deletions,
});
if entries.len() >= self.max_snapshots {
break;
}
}
entries
}
/// 比较快照与当前工作目录的差异。
pub fn diff(&self, working_dir: &Path, commit_hash: &str) -> Result<String, String> {
let inner = self.inner.lock().map_err(|e| format!("锁异常: {}", e))?;
let repo = inner.repo.as_ref().ok_or("Checkpoint 未启用")?;
let _abs_dir =
std::fs::canonicalize(working_dir).map_err(|e| format!("无法解析目录: {}", e))?;
// 校验 commit hash
if commit_hash.is_empty() || commit_hash.starts_with('-') || commit_hash.len() < 4 {
return Err("无效的 commit hash".to_string());
}
let oid = git2::Oid::from_str(commit_hash).map_err(|e| format!("无效的 OID: {}", e))?;
let commit = repo
.find_commit(oid)
.map_err(|e| format!("未找到快照: {}", e))?;
let snapshot_tree = commit
.tree()
.map_err(|e| format!("无法读取快照 tree: {}", e))?;
// 构建当前工作目录的 tree(就地构建 index)
let _index = repo.index().map_err(|e| format!("无法创建 index: {}", e))?;
// 我们无法直接 add_all 到 git2 index(它只读工作目录),
// 改用 diff 的 workdir 模式
let diff = repo
.diff_tree_to_workdir_with_index(Some(&snapshot_tree), None)
.map_err(|e| format!("无法生成 diff: {}", e))?;
let stats = diff.stats().map_err(|e| format!("无法统计 diff: {}", e))?;
let mut output = String::new();
output.push_str(&format!(
"快照 {} ({} 个文件变更, +{} -{} 行)\n\n",
&commit_hash[..8.min(commit_hash.len())],
stats.files_changed(),
stats.insertions(),
stats.deletions(),
));
// 生成 unified diff
diff.print(git2::DiffFormat::Patch, |_delta, _hunk, line| {
let origin = line.origin();
let content = std::str::from_utf8(line.content()).unwrap_or("<binary>");
output.push(origin);
output.push_str(content);
true
})
.map_err(|e| format!("生成 patch 失败: {}", e))?;
Ok(output)
}
/// 恢复文件到指定快照的状态。
///
/// `file_path` 如果为 Some,仅恢复该文件;否则恢复整个目录。
/// 恢复前会自动创建 pre-rollback 快照。
pub fn restore(
&self,
working_dir: &Path,
commit_hash: &str,
file_path: Option<&str>,
) -> Result<String, String> {
let inner = self.inner.lock().map_err(|e| format!("锁异常: {}", e))?;
let repo = inner.repo.as_ref().ok_or("Checkpoint 未启用")?;
let abs_dir =
std::fs::canonicalize(working_dir).map_err(|e| format!("无法解析目录: {}", e))?;
// 校验参数
if commit_hash.is_empty() || commit_hash.starts_with('-') {
return Err("无效的 commit hash".to_string());
}
if let Some(fp) = file_path {
if fp.is_empty() || fp.starts_with('/') || fp.contains("..") {
return Err("无效的文件路径".to_string());
}
}
let oid = git2::Oid::from_str(commit_hash).map_err(|e| format!("无效的 OID: {}", e))?;
let commit = repo
.find_commit(oid)
.map_err(|e| format!("未找到快照: {}", e))?;
let tree = commit
.tree()
.map_err(|e| format!("无法读取快照 tree: {}", e))?;
// Pre-rollback 快照
let _ = self.take_snapshot(
&abs_dir,
&format!("pre-rollback (restoring to {})", &commit_hash[..8]),
);
// git2 的 checkout 操作
let mut checkout_builder = git2::build::CheckoutBuilder::new();
checkout_builder.force(); // 覆盖本地修改
if let Some(fp) = file_path {
// 恢复单个文件
let path = Path::new(fp);
checkout_builder.path(path);
}
repo.checkout_tree(tree.as_object(), Some(&mut checkout_builder))
.map_err(|e| format!("恢复失败: {}", e))?;
Ok(format!(
"已恢复到快照 {}",
&commit_hash[..8.min(commit_hash.len())]
))
}
/// 检查指定工具是否需要触发 checkpoint。
pub fn should_checkpoint(tool_name: &str) -> bool {
CHECKPOINT_TRIGGER_TOOLS.contains(&tool_name)
}
/// 获取 repo 路径
pub fn repo_path(&self) -> &Path {
&self.repo_path
}
// ── 内部方法 ──
/// 获取项目专属的 ref 名
fn ref_name_for(&self, dir: &Path) -> String {
dir_ref_name(dir)
}
/// 创建快照。
/// 返回 Ok(true) 表示创建了新快照,Ok(false) 表示无变更。
fn take_snapshot(&self, dir: &Path, reason: &str) -> Result<bool, anyhow::Error> {
let inner = self
.inner
.lock()
.map_err(|e| anyhow::anyhow!("锁异常: {}", e))?;
let repo = match &inner.repo {
Some(r) => r,
None => return Ok(false),
};
let ref_name = self.ref_name_for(dir);
// 查找该 ref 当前的 tip commit 作为父提交
let parent_commit = repo
.find_reference(&ref_name)
.ok()
.and_then(|r| r.peel_to_commit().ok());
// 构建当前目录的 tree
let mut index = git2::Index::new()?;
// 如果已有父提交,先用父提交的 tree 填充 index
if let Some(ref parent) = parent_commit {
let parent_tree = parent.tree()?;
index.read_tree(&parent_tree)?;
}
// 添加当前目录的所有文件
let walk_result = self.add_files_to_index(dir, &mut index)?;
if !walk_result {
// 无变更
return Ok(false);
}
// 排除超大文件
self.remove_oversize_from_index(dir, &mut index)?;
// 检查是否有实际变更
if let Some(ref parent) = parent_commit {
let parent_tree = parent.tree()?;
let new_tree_oid = index.write_tree()?;
if new_tree_oid == parent_tree.id() {
return Ok(false); // 无变更
}
// 把新 tree 写回 indexwrite_tree 消费了 index,需要重建)
// 实际上 git2 Index::write_tree 不消费 index,我们可以继续用
}
// 写入 tree
let tree_oid = index.write_tree()?;
let tree = repo.find_tree(tree_oid)?;
// 创建 commit
let signature =
git2::Signature::now("AstroResearch Checkpoint", "checkpoint@astroresearch.local")?;
let commit_id = if let Some(ref parent) = parent_commit {
repo.commit(
Some(&ref_name),
&signature,
&signature,
reason,
&tree,
&[parent],
)?
} else {
repo.commit(Some(&ref_name), &signature, &signature, reason, &tree, &[])?
};
debug!(
"[Checkpoint] 快照创建: dir={}, sha={}, reason={}",
dir.display(),
&commit_id.to_string()[..8],
reason
);
// 清理旧快照
if let Err(e) = self.prune_old(dir, &ref_name) {
debug!("[Checkpoint] 清理旧快照失败(非致命): {}", e);
}
Ok(true)
}
/// 遍历目录并将文件添加到 git index。
/// 返回 true 表示有文件被添加。
fn add_files_to_index(
&self,
dir: &Path,
index: &mut git2::Index,
) -> Result<bool, anyhow::Error> {
let mut added = false;
let max_files = 50_000;
let mut count = 0;
for entry in walkdir::WalkDir::new(dir)
.into_iter()
.filter_entry(|e| !is_excluded(e))
{
let entry = entry?;
if !entry.file_type().is_file() {
continue;
}
count += 1;
if count > max_files {
debug!("[Checkpoint] 目录文件数超过上限 ({}), 停止遍历", max_files);
break;
}
let abs_path = entry.path();
let rel_path = abs_path.strip_prefix(dir)?;
// 跳过符号链接和超大文件
let metadata = match std::fs::symlink_metadata(abs_path) {
Ok(m) => m,
Err(_) => continue,
};
if metadata.file_type().is_symlink() {
continue;
}
if metadata.len() > self.max_file_size as u64 {
continue;
}
// 将文件添加到 index
let rel_str = rel_path.to_string_lossy();
index.add_path(Path::new(&*rel_str))?;
added = true;
}
Ok(added)
}
/// 从 index 中移除超大文件。
fn remove_oversize_from_index(
&self,
_dir: &Path,
index: &mut git2::Index,
) -> Result<(), anyhow::Error> {
// NOTE: git2 Index 没有便捷的按大小过滤方法。
// 这里保留接口,后续可以在 add_files_to_index 阶段直接跳过(已实现)。
let _ = index;
Ok(())
}
/// 清理旧快照,每个项目保留最近 `max_snapshots` 个。
fn prune_old(&self, _dir: &Path, ref_name: &str) -> Result<(), anyhow::Error> {
let inner = self
.inner
.lock()
.map_err(|e| anyhow::anyhow!("锁异常: {}", e))?;
let repo = match &inner.repo {
Some(r) => r,
None => return Ok(()),
};
// 统计 commit 数量
let mut revwalk = repo.revwalk()?;
revwalk.push_ref(ref_name)?;
let count = revwalk.count();
if count <= self.max_snapshots {
return Ok(());
}
// 收集所有 commit(从旧到新)
let mut revwalk = repo.revwalk()?;
revwalk.push_ref(ref_name)?;
revwalk.set_sorting(git2::Sort::TIME | git2::Sort::REVERSE)?;
let commits: Vec<git2::Oid> = revwalk.filter_map(|r| r.ok()).collect();
if commits.len() <= self.max_snapshots {
return Ok(());
}
// 保留最后的 N 个
let keep = commits.len() - self.max_snapshots;
let drop_oids: Vec<git2::Oid> = commits.iter().take(keep).copied().collect();
// 重建线性链:从 keep_index 开始
let keep_start = keep;
let keep_commits: Vec<git2::Commit<'_>> = commits[keep_start..]
.iter()
.filter_map(|oid| repo.find_commit(*oid).ok())
.collect();
if keep_commits.is_empty() {
return Ok(());
}
// 重建 chain(保持原有的 tree 和 message
let signature =
git2::Signature::now("AstroResearch Checkpoint", "checkpoint@astroresearch.local")?;
let mut new_parent: Option<git2::Oid> = None;
for commit in &keep_commits {
let tree = commit.tree()?;
let message = commit.message().unwrap_or("checkpoint");
let new_oid = if let Some(parent) = new_parent {
let parent_commit = repo.find_commit(parent)?;
repo.commit(
None,
&signature,
&signature,
message,
&tree,
&[&parent_commit],
)?
} else {
repo.commit(None, &signature, &signature, message, &tree, &[])?
};
new_parent = Some(new_oid);
}
// 更新 ref
if let Some(new_tip) = new_parent {
repo.reference(ref_name, new_tip, true, "prune old checkpoints")?;
}
// 丢弃旧 commits 不再被引用 → git gc 会回收
let _ = drop_oids;
debug!(
"[Checkpoint] 清理完成: dropped {} commits, kept {}",
keep, self.max_snapshots
);
Ok(())
}
}
// ── 辅助函数 ──
/// 为目录生成 ref 名: `refs/checkpoints/<sha256[:16]>`
fn dir_ref_name(dir: &Path) -> String {
use std::hash::{Hash, Hasher};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
dir.to_string_lossy().hash(&mut hasher);
let hash = hasher.finish();
format!("refs/checkpoints/{:016x}", hash)
}
/// 获取 Home 目录
fn dirs_home() -> PathBuf {
std::env::var("HOME")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from("/"))
}
/// 检查路径是否匹配排除规则。
///
/// 使用简单的 glob 匹配 DEFAULT_EXCLUDES 中列出的模式。
fn is_excluded(entry: &walkdir::DirEntry) -> bool {
let file_name = entry.file_name().to_string_lossy();
let path_str = entry.path().to_string_lossy();
for pattern in DEFAULT_EXCLUDES {
// 目录模式
if pattern.ends_with('/') {
let dir_name = pattern.trim_end_matches('/');
if file_name.as_ref() == dir_name && entry.file_type().is_dir() {
return true;
}
}
// 文件扩展名模式
if pattern.starts_with("*.") {
let ext = &pattern[1..]; // ".pyc"
if file_name.ends_with(ext) {
return true;
}
}
// 精确匹配
if file_name.as_ref() == *pattern {
return true;
}
// 包含目录路径的模式
if pattern.contains('/') && path_str.contains(pattern) {
return true;
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
/// 创建临时目录和 checkpoint manager
fn setup(tmp: &tempfile::TempDir, enabled: bool) -> CheckpointManager {
let store = tmp.path().join(".checkpoints");
CheckpointManager::new(store, enabled)
}
#[test]
fn test_disabled_skips_all() {
let tmp = tempfile::tempdir().unwrap();
let mgr = setup(&tmp, false);
assert!(!mgr.ensure_checkpoint(tmp.path(), "test"));
assert!(mgr.list_checkpoints(tmp.path()).is_empty());
}
#[test]
fn test_enabled_creates_snapshot() {
let tmp = tempfile::tempdir().unwrap();
let work = tmp.path().join("work");
std::fs::create_dir_all(&work).unwrap();
// 创建一个文件
let file_path = work.join("test.txt");
let mut f = std::fs::File::create(&file_path).unwrap();
writeln!(f, "hello world").unwrap();
let mgr = setup(&tmp, true);
// First turn
mgr.new_turn();
let result = mgr.ensure_checkpoint(&work, "initial");
// git2 may fail in test environments without git config
// We just verify it doesn't panic
let _ = result;
// List checkpoints
let entries = mgr.list_checkpoints(&work);
// Not asserting count since git2 may behave differently
let _ = entries;
}
#[test]
fn test_dedup_per_turn() {
let tmp = tempfile::tempdir().unwrap();
let work = tmp.path().join("work");
std::fs::create_dir_all(&work).unwrap();
let mgr = setup(&tmp, true);
mgr.new_turn();
// 同一目录同一 turn 只快照一次
let first = mgr.ensure_checkpoint(&work, "first");
let second = mgr.ensure_checkpoint(&work, "second");
// second should be false (already checkpointed this turn)
if first {
assert!(!second);
}
}
#[test]
fn test_new_turn_resets_dedup() {
let tmp = tempfile::tempdir().unwrap();
let work = tmp.path().join("work");
std::fs::create_dir_all(&work).unwrap();
let mgr = setup(&tmp, true);
mgr.new_turn();
let _ = mgr.ensure_checkpoint(&work, "turn1");
mgr.new_turn(); // reset
// 新 turn 应该可以再次快照
let result = mgr.ensure_checkpoint(&work, "turn2");
let _ = result; // may or may not create depending on changes
}
#[test]
fn test_should_checkpoint_triggers() {
assert!(CheckpointManager::should_checkpoint("file_write"));
assert!(CheckpointManager::should_checkpoint("file_edit"));
assert!(CheckpointManager::should_checkpoint("run_bash"));
assert!(!CheckpointManager::should_checkpoint("read_file"));
assert!(!CheckpointManager::should_checkpoint("search_papers"));
}
#[test]
fn test_dir_ref_name_deterministic() {
let name1 = dir_ref_name(Path::new("/home/user/project"));
let name2 = dir_ref_name(Path::new("/home/user/project"));
assert_eq!(name1, name2);
let name3 = dir_ref_name(Path::new("/other/path"));
assert_ne!(name1, name3);
}
#[test]
fn test_restore_invalid_hash() {
let tmp = tempfile::tempdir().unwrap();
let mgr = setup(&tmp, true);
let result = mgr.restore(tmp.path(), "-invalid", None);
assert!(result.is_err());
}
#[test]
fn test_diff_invalid_hash() {
let tmp = tempfile::tempdir().unwrap();
let mgr = setup(&tmp, true);
let result = mgr.diff(tmp.path(), "-bad");
assert!(result.is_err());
}
#[test]
fn test_exclude_patterns() {
// Test that common excludes are matched
let tmp = tempfile::tempdir().unwrap();
// Create a .git directory
let git_dir = tmp.path().join(".git");
std::fs::create_dir_all(&git_dir).unwrap();
for entry in walkdir::WalkDir::new(tmp.path()) {
let e = entry.unwrap();
if e.file_name() == ".git" {
assert!(is_excluded(&e));
}
}
}
#[test]
fn test_disabled_on_init_failure() {
// Use a path that can't be created (e.g., /dev/null/file)
let bad_path = PathBuf::from("/proc/self/fd/0/checkpoints");
let mgr = CheckpointManager::new(bad_path, true);
assert!(!mgr.enabled);
}
}
File diff suppressed because it is too large Load Diff
+463 -130
View File
@@ -13,14 +13,18 @@ use tracing::{info, warn};
use crate::api::{AppState, PendingPermission};
use crate::clients::llm::{ChatMessage, ToolCall};
use super::checkpoint::CheckpointManager;
use super::denial_tracker::DenialTracker;
use super::file_cache::FileStateCache;
use super::hardline;
use super::partitioner::ToolPartitioner;
use super::permission::{PermissionChecker, PermissionResult};
use super::permission_explainer::explain_permission;
use super::{AgentStreamEvent, DuplicateDetector};
use crate::agent::hooks::{HookRegistry, PostToolUseContext, PreToolUseContext};
use crate::agent::tools::persist::maybe_persist_tool_result;
use crate::agent::tools::{InterruptBehavior, ToolContext, ToolOutput, ToolRegistry};
use crate::agent::hooks::{
event_label, HookRegistry, PostToolUseContext, PostToolUseFailureContext, PreToolUseContext,
};
use crate::agent::tools::{ToolContext, ToolRegistry};
/// 准备好的工具调用
#[derive(Debug, Clone)]
@@ -42,6 +46,10 @@ pub struct ToolExecutionResult {
pub tool_messages: Vec<ToolResultMessage>,
pub was_cancelled: bool,
pub had_duplicate: bool,
/// Hook 注入的附加上下文(PreToolUse + PostToolUse),需注入 LLM 消息列表
pub hook_contexts: Vec<String>,
/// Hook 的阻塞错误详情(用于日志和诊断)
pub blocking_errors: Vec<String>,
}
/// 验证工具调用:死循环检测 + 参数解析。
@@ -144,6 +152,7 @@ pub async fn execute_parallel(
permission_checker: Option<&PermissionChecker>,
session_permission_checker: Option<&std::sync::RwLock<PermissionChecker>>,
denial_tracker: Option<&std::sync::Mutex<DenialTracker>>,
checkpoint_manager: Option<&std::sync::Arc<CheckpointManager>>,
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
db: &SqlitePool,
session_id: &str,
@@ -161,6 +170,8 @@ pub async fn execute_parallel(
tool_messages: Vec::new(),
was_cancelled: false,
had_duplicate: false,
hook_contexts: Vec::new(),
blocking_errors: Vec::new(),
};
}
@@ -180,7 +191,9 @@ pub async fn execute_parallel(
let exec_start = std::time::Instant::now();
let mut mutated_args: Vec<serde_json::Value> = Vec::new();
let mut additional_contexts: Vec<String> = Vec::new();
let mut hook_permission_required: Vec<bool> = Vec::new();
let mut hook_permission_info: Vec<Option<(String, String)>> = Vec::new();
// ^^^ (permission_desc, hook_tool_name)
let mut hook_blocking_errors: Vec<String> = Vec::new();
for prep in prepared_calls {
let hook_ctx = PreToolUseContext {
session_id: sid.clone(),
@@ -196,20 +209,39 @@ pub async fn execute_parallel(
prep.tool_name, reason
);
}
// 收集 hook 的权限请求
if result.is_permission_required() {
// 收集所有阻塞错误详情(含多个 hook 同时 block 的情况)
for be in &result.blocking_errors {
hook_blocking_errors.push(format!(
"[{}] 阻止 {}: {}",
be.hook_name, prep.tool_name, be.reason
));
}
// 收集 hook 的权限请求(保留完整信息用于 AskUser prompt
if let Some((permission, tool_name)) = result.permission_info() {
info!(
"[Executor] PreToolUse hook 请求了工具 {} 的权限确认",
prep.tool_name
"[Executor] Hook 请求了工具 {} 的权限确认: {}",
prep.tool_name, permission
);
hook_permission_required.push(true);
hook_permission_info.push(Some((permission.to_string(), tool_name.to_string())));
} else {
hook_permission_required.push(false);
hook_permission_info.push(None);
}
// 使用 hook 可能修改后的参数
mutated_args.push(result.final_args);
if let Some(ctx) = result.additional_context {
additional_contexts.push(ctx);
// 收集所有 hook 注入的上下文(优先使用带来源标记的 tagged_contexts
if !result.tagged_contexts.is_empty() {
for tc in &result.tagged_contexts {
additional_contexts.push(format!(
"[Hook: {} | {}] {}",
tc.hook_name,
event_label(tc.source_event),
tc.content,
));
}
} else {
for ctx in &result.additional_contexts {
additional_contexts.push(ctx.clone());
}
}
}
@@ -217,17 +249,83 @@ pub async fn execute_parallel(
// 被拒绝的工具直接注入错误 result,不进入执行队列。
let mut tool_messages: Vec<ToolResultMessage> = Vec::new();
let mut denied_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();
// ── Hardline 预检查(在任何模式下都不可绕过)──
// 在 PermissionChecker 之前执行,确保 hardline 规则始终生效。
for (i, prep) in prepared_calls.iter().enumerate() {
let hardline_result = match prep.tool_name.as_str() {
"run_bash" => {
if let Some(cmd) = prep.args.get("command").and_then(|v| v.as_str()) {
hardline::check_command(cmd)
} else {
hardline::HardlineResult::allowed()
}
}
"file_write" | "file_edit" => {
if let Some(path) = prep.args.get("file_path").and_then(|v| v.as_str()) {
hardline::check_dangerous_path(path)
} else if let Some(path) = prep.args.get("path").and_then(|v| v.as_str()) {
hardline::check_dangerous_path(path)
} else {
hardline::HardlineResult::allowed()
}
}
_ => hardline::HardlineResult::allowed(),
};
if hardline_result.blocked {
warn!(
"[Executor] Hardline 阻止了工具 {} (category={}): {}",
prep.tool_name,
hardline_result.category.as_deref().unwrap_or("unknown"),
hardline_result.reason
);
let err_output = hardline_result.reason.clone();
let _ = tx.send(AgentStreamEvent::ToolResult {
tool_call_id: prep.tool_call_id.clone(),
name: prep.tool_name.clone(),
output: err_output.clone(),
is_error: true,
metadata: serde_json::json!({
"hardline_blocked": true,
"hardline_category": hardline_result.category,
}),
step,
});
let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output);
save_tool_message_sync(db, &sid, turn_index, step, &err_msg);
tool_messages.push(ToolResultMessage {
chat_message: err_msg,
was_error: true,
});
// 记录拒绝追踪
if let Some(dt) = denial_tracker {
if let Ok(mut tracker) = dt.lock() {
tracker.record_denial();
}
}
denied_indices.insert(i);
}
}
if let Some(checker) = permission_checker {
for (i, prep) in prepared_calls.iter().enumerate() {
let mut perm_result = checker.check(&prep.tool_name, Some(&prep.args));
perm_result = checker.apply_mode(perm_result, &prep.tool_name);
// Hook PermissionRequired — 若 Checker 返回 Allowed,升级为 Ask
if hook_permission_required.get(i).copied().unwrap_or(false) && perm_result.is_allowed()
{
perm_result = PermissionResult::AskUser {
message: format!("Hook 请求了工具 {} 的权限确认", prep.tool_name),
};
// 使用 hook 提供的具体权限描述替换泛型消息
if let Some(Some((ref perm_desc, _))) = hook_permission_info.get(i) {
if perm_result.is_allowed() {
perm_result = PermissionResult::AskUser {
message: format!(
"[Hook 权限请求] {}\n\n工具: {}\n参数: {}",
perm_desc,
prep.tool_name,
serde_json::to_string_pretty(&prep.args).unwrap_or_default(),
),
};
}
}
// 工具级 check_permissions() — 在 PermissionChecker 结果基础上叠加
@@ -274,10 +372,8 @@ pub async fn execute_parallel(
}
}
PermissionResult::Allowed => {
// 会话 Allow 仅在非 Deny 时覆盖(会话明确允许)
if !perm_result.is_denied() {
perm_result = PermissionResult::Allowed;
}
// 会话 Allow 仅覆盖 Allowed,保持 Deny/AskUser 不变
// 避免覆盖工具级 check_permissions() 升级的 AskUser
}
}
}
@@ -481,7 +577,14 @@ pub async fn execute_parallel(
}
} // if let Some(checker)
// Phase 3: 并行执行
// Phase 3: 分区并行执行(参考 Claude Code partitionToolCalls + runTools)。
//
// 改进:原实现将所有非拒绝工具放入单个 FuturesUnordered 无差别并发,
// 可能导致非并发安全工具(如 run_bash)错误地并行执行。
// 新实现使用 ToolPartitioner 将工具按并发安全性分批:
// - 连续的并发安全工具放入同一个并行批次(FuturesUnordered
// - 非并发安全工具独占一个串行批次(逐次执行)
// 批次内工具执行完成后立即推送 SSE 事件,不等待整个批次完成。
let cancelled = Arc::new(AtomicBool::new(false));
let cancel_flag = cancelled.clone();
let app_state_ref = app_state.clone();
@@ -501,128 +604,187 @@ pub async fn execute_parallel(
let timeout_dur = std::time::Duration::from_secs(tool_timeout_secs);
// Phase 3: 使用 FuturesUnordered 进行渐进式并行执行。
// 每个工具完成后立即发送 SSE ToolResult 事件到前端(非阻塞),
// 而后台继续等待其他工具完成。快工具的结果不会因慢工具而延迟。
let mut exec_futs: FuturesUnordered<_> = prepared_calls
// ── Checkpoint 预触发:对文件变更类工具在执行前创建快照 ──
if let Some(ckpt) = checkpoint_manager {
let cwd = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
for prep in prepared_calls
.iter()
.filter(|p| CheckpointManager::should_checkpoint(&p.tool_name))
{
ckpt.ensure_checkpoint(&cwd, &format!("pre-{}", prep.tool_name));
}
}
// ── Phase 3a: 构建不包含被拒绝工具的 (原索引, PreparedCall) 映射 ──
let non_denied: Vec<(usize, &PreparedCall)> = prepared_calls
.iter()
.enumerate()
.filter(|(i, _)| !denied_indices.contains(i))
.map(|(i, prep)| {
let tool_name = prep.tool_name.clone();
let args = mutated_args
.get(i)
.cloned()
.unwrap_or_else(|| prep.args.clone());
let tool_ctx = ToolContext::with_file_cache(app_state.clone(), read_file_state.clone())
.with_sse_tx(tx.clone())
.with_session_id(session_id.to_string())
.with_thinking(enable_thinking)
.with_additional_dirs(additional_allowed_dirs.clone());
let cancelled = cancelled.clone();
let tool_opt = tool_registry.get(&tool_name);
Box::pin(async move {
let output = match tool_opt {
Some(tool) => {
let interrupt_behavior = tool.interrupt_behavior();
let is_blocking = interrupt_behavior == InterruptBehavior::Block;
let tool_fut = tool.execute(args, &tool_ctx);
let cancel_fut = async {
loop {
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
if !is_blocking && cancelled.load(Ordering::SeqCst) {
return;
}
}
};
tokio::select! {
res = tokio::time::timeout(timeout_dur, tool_fut) => {
match res {
Ok(output) => output,
Err(_) => ToolOutput::error(format!(
"工具 {} 执行超时({}秒)",
tool_name,
timeout_dur.as_secs()
)),
}
}
_ = cancel_fut => {
ToolOutput::error("执行已被用户取消")
}
}
}
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
};
let was_cancelled = cancelled.load(Ordering::SeqCst);
(
prep.tool_call_id.clone(),
prep.tool_name.clone(),
prep.args.clone(),
output,
was_cancelled,
)
})
})
.collect();
// ── Phase 3b: 分区 ──
let non_denied_calls: Vec<PreparedCall> =
non_denied.iter().map(|(_, p)| (*p).clone()).collect();
let partitioner = ToolPartitioner::new(10);
let batches = partitioner.partition(&non_denied_calls, tool_registry);
// 预设非拒绝工具中哪些原索引属于已拒绝列表(不会有,但安全起见)
let original_index_of: std::collections::HashMap<String, usize> = non_denied
.iter()
.map(|(orig_idx, prep)| (prep.tool_call_id.clone(), *orig_idx))
.collect();
info!(
"[Executor] 工具分区完成: {} 工具 → {} 批次 ({} 串行 + {} 并行)",
non_denied.len(),
batches.len(),
batches.iter().filter(|b| !b.is_parallel).count(),
batches.iter().filter(|b| b.is_parallel).count(),
);
let mut was_cancelled = false;
// 渐进式处理结果:每个工具一完成就立即处理(SSE 事件 + PostToolUse hook + 持久化)
while let Some((tool_call_id, tool_name, tool_args, output, cancelled_flag)) =
exec_futs.next().await
{
if cancelled_flag {
was_cancelled = true;
// ── Phase 3c: 逐批次执行 ──
// 批次之间串行;并行批次内工具并发执行;串行批次内工具逐个执行。
for batch in &batches {
if was_cancelled {
break;
}
let elapsed_ms = exec_start.elapsed().as_millis() as u64;
if batch.is_parallel {
// ── 并行批次:FuturesUnordered 并发执行 ──
let mut exec_futs: FuturesUnordered<_> = batch
.calls
.iter()
.map(|prep| {
let orig_idx = original_index_of
.get(&prep.tool_call_id)
.copied()
.unwrap_or(0);
let tool_name = prep.tool_name.clone();
let args = mutated_args
.get(orig_idx)
.cloned()
.unwrap_or_else(|| prep.args.clone());
let tool_ctx =
ToolContext::with_file_cache(app_state.clone(), read_file_state.clone())
.with_sse_tx(tx.clone())
.with_session_id(session_id.to_string())
.with_thinking(enable_thinking)
.with_additional_dirs(additional_allowed_dirs.clone());
let cancelled = cancelled.clone();
let tool_opt = tool_registry.get(&tool_name);
// SSE 事件 — 立即推送到前端
let _ = tx.send(AgentStreamEvent::ToolResult {
tool_call_id: tool_call_id.clone(),
name: tool_name.clone(),
output: output.content.clone(),
is_error: output.is_error,
metadata: output.metadata.clone(),
step,
});
Box::pin(async move {
let output = execute_single_tool(
tool_opt,
args,
&tool_ctx,
&cancelled,
timeout_dur,
&tool_name,
)
.await;
let was_cancelled = cancelled.load(Ordering::SeqCst);
(
prep.tool_call_id.clone(),
prep.tool_name.clone(),
prep.args.clone(),
output,
was_cancelled,
)
})
})
.collect();
// 输出处理:小结果直接传递,大结果持久化到磁盘并返回 stub
let tool_results_dir = app_state.config.library_dir.join("tool-results");
let (processed_content, _persisted_path) = maybe_persist_tool_result(
&output.content,
&tool_call_id,
max_output_chars,
&tool_results_dir,
);
// 渐进式处理:每个工具一完成就处理
while let Some((tool_call_id, tool_name, tool_args, output, cancelled_flag)) =
exec_futs.next().await
{
if cancelled_flag {
was_cancelled = true;
}
process_single_result(
&tool_call_id,
&tool_name,
&tool_args,
&output,
cancelled_flag,
exec_start,
tx,
hook_registry,
&app_state.config.library_dir,
&sid,
agent_name,
step,
max_output_chars,
&mut tool_messages,
&mut additional_contexts,
db,
turn_index,
)
.await;
}
} else {
// ── 串行批次:逐个执行 ──
for prep in &batch.calls {
let orig_idx = original_index_of
.get(&prep.tool_call_id)
.copied()
.unwrap_or(0);
let tool_name = prep.tool_name.clone();
let args = mutated_args
.get(orig_idx)
.cloned()
.unwrap_or_else(|| prep.args.clone());
let tool_ctx =
ToolContext::with_file_cache(app_state.clone(), read_file_state.clone())
.with_sse_tx(tx.clone())
.with_session_id(session_id.to_string())
.with_thinking(enable_thinking)
.with_additional_dirs(additional_allowed_dirs.clone());
let tool_opt = tool_registry.get(&tool_name);
// PostToolUse hook
let post_ctx = PostToolUseContext {
session_id: sid.clone(),
agent_name: agent_name.to_string(),
tool_name: tool_name.clone(),
tool_args,
output_content: processed_content.clone(),
is_error: output.is_error,
step,
elapsed_ms,
};
let post_result = hook_registry.run_post_tool_use(&post_ctx).await;
let final_content = post_result.final_content;
let output = execute_single_tool(
tool_opt,
args,
&tool_ctx,
&cancelled,
timeout_dur,
&tool_name,
)
.await;
let cancelled_flag = cancelled.load(Ordering::SeqCst);
if cancelled_flag {
was_cancelled = true;
}
let chat_message = ChatMessage::tool_result(&tool_call_id, &final_content);
process_single_result(
&prep.tool_call_id,
&tool_name,
&prep.args,
&output,
cancelled_flag,
exec_start,
tx,
hook_registry,
&app_state.config.library_dir,
&sid,
agent_name,
step,
max_output_chars,
&mut tool_messages,
&mut additional_contexts,
db,
turn_index,
)
.await;
// 持久化到数据库(fire-and-forget
save_tool_message_sync(db, &sid, turn_index, step, &chat_message);
tool_messages.push(ToolResultMessage {
chat_message,
was_error: output.is_error,
});
if was_cancelled {
break;
}
}
}
}
cancel_handle.abort();
@@ -631,9 +793,180 @@ pub async fn execute_parallel(
tool_messages,
was_cancelled,
had_duplicate: false,
hook_contexts: additional_contexts,
blocking_errors: hook_blocking_errors,
}
}
/// 执行单个工具调用(含超时和取消检测)。
///
/// 从原 `execute_parallel` 的闭包提取,供分区后的批次执行复用。
async fn execute_single_tool(
tool_opt: Option<&dyn crate::agent::tools::AgentTool>,
args: serde_json::Value,
tool_ctx: &crate::agent::tools::ToolContext,
cancelled: &Arc<AtomicBool>,
timeout_dur: std::time::Duration,
tool_name: &str,
) -> crate::agent::tools::ToolOutput {
let tool = match tool_opt {
Some(t) => t,
None => return crate::agent::tools::ToolOutput::error(format!("未知工具: {}", tool_name)),
};
let interrupt_behavior = tool.interrupt_behavior();
let is_blocking = interrupt_behavior == crate::agent::tools::InterruptBehavior::Block;
let tool_fut = tool.execute(args, tool_ctx);
let cancelled = cancelled.clone();
let cancel_fut = async {
loop {
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
if !is_blocking && cancelled.load(Ordering::SeqCst) {
return;
}
}
};
tokio::select! {
res = tokio::time::timeout(timeout_dur, tool_fut) => {
match res {
Ok(output) => output,
Err(_) => crate::agent::tools::ToolOutput::error(format!(
"工具 {} 执行超时({}秒)",
tool_name,
timeout_dur.as_secs()
)),
}
}
_ = cancel_fut => {
crate::agent::tools::ToolOutput::error("执行已被用户取消")
}
}
}
/// 处理单个工具执行结果(SSE 事件、PostToolUse hooks、持久化)。
///
/// 从原 `execute_parallel` 的结果处理循环提取。
#[allow(clippy::too_many_arguments)]
async fn process_single_result(
tool_call_id: &str,
tool_name: &str,
tool_args: &serde_json::Value,
output: &crate::agent::tools::ToolOutput,
cancelled_flag: bool,
exec_start: std::time::Instant,
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
hook_registry: &HookRegistry,
library_dir: &std::path::Path,
sid: &str,
agent_name: &str,
step: usize,
max_output_chars: usize,
tool_messages: &mut Vec<ToolResultMessage>,
additional_contexts: &mut Vec<String>,
db: &SqlitePool,
turn_index: i32,
) {
use crate::agent::tools::persist::maybe_persist_tool_result;
let elapsed_ms = exec_start.elapsed().as_millis() as u64;
// SSE 事件 — 立即推送到前端
let _ = tx.send(AgentStreamEvent::ToolResult {
tool_call_id: tool_call_id.to_string(),
name: tool_name.to_string(),
output: output.content.clone(),
is_error: output.is_error,
metadata: output.metadata.clone(),
step,
});
// 输出处理:小结果直接传递,大结果持久化到磁盘并返回 stub
let tool_results_dir = library_dir.join("tool-results");
let (processed_content, _persisted_path) = maybe_persist_tool_result(
&output.content,
tool_call_id,
max_output_chars,
&tool_results_dir,
);
// PostToolUse hook
let post_ctx = PostToolUseContext {
session_id: sid.to_string(),
agent_name: agent_name.to_string(),
tool_name: tool_name.to_string(),
tool_args: tool_args.clone(),
output_content: processed_content.clone(),
is_error: output.is_error,
step,
elapsed_ms,
};
let post_result = hook_registry.run_post_tool_use(&post_ctx).await;
let final_content = post_result.final_content;
// 非可信内容包裹(间接 prompt 注入防御)
let llm_content = super::untrusted::wrap_untrusted_content(tool_name, &final_content);
// 收集 PostToolUse hook 注入的上下文
if !post_result.tagged_contexts.is_empty() {
for tc in &post_result.tagged_contexts {
additional_contexts.push(format!(
"[Hook: {} | {}] {}",
tc.hook_name,
event_label(tc.source_event),
tc.content,
));
}
} else {
for ctx in &post_result.additional_contexts {
additional_contexts.push(ctx.clone());
}
}
// 收集 PostToolUse 的警告
for warning in &post_result.warnings {
additional_contexts.push(format!("[Hook Warning] {}", warning));
}
// 事后权限请求(audit trail
for (perm_tool, perm) in &post_result.post_permission_requests {
warn!(
"[Executor] Hook 事后请求权限: tool={} permission={}",
perm_tool, perm
);
}
// PostToolUseFailure hook
if output.is_error {
let failure_ctx = PostToolUseFailureContext {
session_id: sid.to_string(),
agent_name: agent_name.to_string(),
tool_name: tool_name.to_string(),
tool_args: tool_args.clone(),
error_message: output.content.clone(),
is_interrupt: cancelled_flag,
step,
elapsed_ms,
};
hook_registry
.run_on_post_tool_use_failure(&failure_ctx)
.await;
}
// 发送给 LLM 使用包裹后的内容(安全防御)
let chat_message = ChatMessage::tool_result(tool_call_id, &llm_content);
// 持久化到数据库(fire-and-forget
save_tool_message_sync(db, sid, turn_index, step, &chat_message);
tool_messages.push(ToolResultMessage {
chat_message,
was_error: output.is_error,
});
}
/// 同步保存 tool 角色消息到数据库。
fn save_tool_message_sync(
db: &SqlitePool,
+1 -1
View File
@@ -37,7 +37,7 @@ pub async fn finalize_turn(
app_state: Option<Arc<AppState>>,
) -> anyhow::Result<()> {
let new_turn_count: i32 = sqlx::query_scalar(
"SELECT COUNT(DISTINCT turn_index) FROM agent_messages WHERE session_id = ?",
"SELECT COUNT(DISTINCT turn_index) FROM agent_messages WHERE session_id = ? AND active = 1",
)
.bind(session_id)
.fetch_one(db)
+534
View File
@@ -0,0 +1,534 @@
// src/agent/runtime/hardline.rs
//
// Hardline 命令阻止层 — 不可绕过的危险命令检查。
// 参考 Hermes-Agent approval.py HARDLINE_PATTERNS 设计。
//
// 设计原则:
// 1. Hardline 规则在任何模式下都不可被绕过(包括 YOLO/Bypass 模式)
// 2. 优先级高于所有其他权限规则
// 3. 在命令执行前做反规避标准化后再匹配
//
// 阻止的命令类别:
// - 系统关机/重启
// - 磁盘擦除/格式化
// - Fork 炸弹
// - 递归删除根目录
// - kill -1(信号广播)
use regex::RegexSet;
use std::sync::LazyLock;
// ── 反规避命令标准化 ──
/// 在安全检查前对命令字符串做标准化处理。
/// 参考 Hermes `_normalize_command_for_detection()` 实现。
///
/// 转换顺序:
/// 1. 剥离 ANSI 转义序列(ECMA-48
/// 2. 剥离 null 字节
/// 3. Unicode 全角字符 NFKC 标准化
/// 4. 剥离 shell 反斜杠转义(`r\m` → `rm`
/// 5. 剥离空字符串字面量(`r''m` → `rm`
/// 6. 解析后的绝对路径还原为 `~/` 形式
pub fn normalize_command(raw: &str) -> String {
// Step 1: 剥离 ANSI 转义序列
let s = strip_ansi_escapes(raw);
// Step 2: 剥离 null 字节
let s = s.replace('\0', "");
// Step 3: Unicode NFKC 标准化(全角 → 半角)
let s = unicode_normalization::lookup(&s).unwrap_or_else(|| s.to_string());
// Step 4: 剥离 shell 反斜杠转义(`r\m` → `rm`
let s = strip_backslash_escapes(&s);
// Step 5: 剥离空字符串字面量(`r''m` → `rm`, `r""m` → `rm`
let s = strip_empty_string_literals(&s);
// Step 6: 还原 Home 路径
normalize_home_paths(&s)
}
/// 剥离 ANSI 转义序列(CSI 序列,ECMA-48 §5.4
fn strip_ansi_escapes(s: &str) -> String {
static ANSI_RE: LazyLock<regex::Regex> =
LazyLock::new(|| regex::Regex::new(r"\x1b\[[0-9;]*[a-zA-Z]").expect("ANSI regex compile"));
ANSI_RE.replace_all(s, "").to_string()
}
/// Unicode 全角字符 NFKC 标准化 + 常见全角 ASCII 映射
mod unicode_normalization {
use std::collections::HashMap;
use std::sync::LazyLock;
static FULLWIDTH_MAP: LazyLock<HashMap<char, char>> = LazyLock::new(|| {
// 全角 ASCIIU+FF01-U+FF5E)映射到半角(U+0021-U+007E
let mut map = HashMap::new();
for code in 0xFF01u32..=0xFF5E {
if let Some(c) = char::from_u32(code) {
let half_width = char::from_u32(code - 0xFEE0).unwrap_or(c);
if c != half_width {
map.insert(c, half_width);
}
}
}
// 全角空格 U+3000 → 半角空格 U+0020
map.insert('\u{3000}', ' ');
map
});
/// 如果字符串包含全角字符,返回标准化后的版本;否则返回 None(无需复制)。
pub fn lookup(s: &str) -> Option<String> {
let needs_normalize = s.chars().any(|c| FULLWIDTH_MAP.contains_key(&c));
if !needs_normalize {
return None;
}
let normalized: String = s
.chars()
.map(|c| FULLWIDTH_MAP.get(&c).copied().unwrap_or(c))
.collect();
Some(normalized)
}
}
/// 剥离 shell 反斜杠转义。
/// 匹配 `\<任意字符>` 并还原为 `<字符>`。
/// 示例:`r\m\ \-\r\f\ \/` → `rm -rf /`
fn strip_backslash_escapes(s: &str) -> String {
// 匹配 backslash 后跟任意非换行字符,捕获该字符
static BACKSLASH_ESCAPE_RE: LazyLock<regex::Regex> =
LazyLock::new(|| regex::Regex::new(r"\\(.)").expect("backslash escape regex compile"));
// 仅当有反斜杠时才执行替换(快速路径)
if s.contains('\\') {
BACKSLASH_ESCAPE_RE.replace_all(s, "$1").to_string()
} else {
s.to_string()
}
}
/// 剥离空字符串字面量。
/// 匹配 `''` 或 `""`shell 中用于分割命令名)。
/// 示例:`r''m` → `rm`, `r""m` → `rm`
fn strip_empty_string_literals(s: &str) -> String {
// 匹配连续两个单引号('')或连续两个双引号("")
static EMPTY_QUOTE_RE: LazyLock<regex::Regex> =
LazyLock::new(|| regex::Regex::new(r"'{2}|\x22{2}").expect("empty quote regex compile"));
let has_single_empty = s.contains("''");
let has_double_empty = s.contains("\"\"");
if has_single_empty || has_double_empty {
EMPTY_QUOTE_RE.replace_all(s, "").to_string()
} else {
s.to_string()
}
}
/// 将解析后的绝对 HOME 路径还原为 `~/` 形式。
fn normalize_home_paths(s: &str) -> String {
let home = std::env::var("HOME").unwrap_or_else(|_| "/home/unknown".to_string());
if home.is_empty() || home == "/" {
return s.to_string();
}
s.replace(&home, "~")
}
// ── Hardline 模式定义 ──
/// 不可绕过的硬阻止模式。
///
/// 每个模式包含:
/// - `pattern`: 正则表达式
/// - `category`: 命令类别(用于日志和错误消息)
/// - `message`: 返回给 LLM 的阻止理由
struct HardlinePattern {
pattern: &'static str,
category: &'static str,
message: &'static str,
}
/// Hardline 模式列表。
/// 参考 Hermes HARDLINE_PATTERNS + 科研场景特定扩展。
static HARDLINE_PATTERNS: LazyLock<Vec<HardlinePattern>> = LazyLock::new(|| {
vec![
// ══════ 系统关机/重启 ══════
HardlinePattern {
pattern: r"\b(?:shutdown|poweroff|halt|reboot)\b",
category: "system_shutdown",
message: "系统关机/重启命令被硬阻止",
},
HardlinePattern {
pattern: r"\binit\s+[06]\b",
category: "system_shutdown",
message: "init 运行级别切换被硬阻止",
},
HardlinePattern {
pattern: r"\bsystemctl\s+(?:poweroff|reboot|halt)\b",
category: "system_shutdown",
message: "systemctl 关机命令被硬阻止",
},
// ══════ 磁盘擦除/格式化 ══════
HardlinePattern {
pattern: r"\bmkfs\b",
category: "disk_format",
message: "磁盘格式化命令 mkfs 被硬阻止",
},
HardlinePattern {
pattern: r"\bdd\s+.*\bof=/dev/[a-z]+",
category: "dd_to_device",
message: "dd 写入块设备被硬阻止",
},
HardlinePattern {
pattern: r"\bdd\s+.*\bof=/dev/(?:sd[a-z]|nvme\d+n\d+|mmcblk\d+)",
category: "dd_to_device",
message: "dd 写入磁盘设备被硬阻止",
},
// ══════ 递归删除根目录 ══════
HardlinePattern {
pattern: r"\brm\s+-rf\s+(?:/|/\*)",
category: "rm_root",
message: "递归删除根目录被硬阻止",
},
HardlinePattern {
pattern: r"\brm\s+.*\s+-rf\s+/",
category: "rm_root",
message: "递归删除根目录被硬阻止",
},
// ══════ Fork 炸弹 ══════
HardlinePattern {
pattern: r":\(\)\s*\{[^}]*:[^}]*\}",
category: "fork_bomb",
message: "Fork 炸弹模式被硬阻止",
},
HardlinePattern {
pattern: r"\bperl\s+-e\s+.*fork.*while",
category: "fork_bomb",
message: "Perl fork 循环被硬阻止",
},
HardlinePattern {
pattern: r"\bpython3?\s+-c\s+.*while.*os\.fork",
category: "fork_bomb",
message: "Python fork 炸弹被硬阻止",
},
// ══════ Kill 信号广播 ══════
HardlinePattern {
pattern: r"\bkill\s+-1\b",
category: "kill_all",
message: "kill -1(信号广播到所有进程)被硬阻止",
},
// ══════ 覆盖关键系统文件 ══════
HardlinePattern {
pattern: r">\s*/etc/(?:passwd|shadow|sudoers|hosts)\b",
category: "system_file_overwrite",
message: "重定向覆盖关键系统文件被硬阻止",
},
HardlinePattern {
pattern: r"\bcp\s+.*\s+/etc/(?:passwd|shadow|sudoers)\b",
category: "system_file_overwrite",
message: "复制覆盖关键系统文件被硬阻止",
},
// ══════ chmod 危险操作 ══════
HardlinePattern {
pattern: r"\bchmod\s+.*777\s+/(?:etc|bin|usr|lib|sbin|boot)\b",
category: "dangerous_chmod",
message: "对系统目录执行 chmod 777 被硬阻止",
},
// ══════ chown 到 root ══════
HardlinePattern {
pattern: r"\bchown\s+-R\s+root:root\s+/",
category: "chown_root",
message: "递归 chown root 到根目录被硬阻止",
},
]
});
/// 编译后的 Hardline 正则集合(模块加载时编译一次)
static HARDLINE_REGEX_SET: LazyLock<RegexSet> = LazyLock::new(|| {
let patterns: Vec<&str> = HARDLINE_PATTERNS.iter().map(|p| p.pattern).collect();
RegexSet::new(&patterns).expect("Hardline regex patterns must compile")
});
// ── 检查 API ──
/// Hardline 检查结果
#[derive(Debug, Clone)]
pub struct HardlineResult {
/// 是否被阻止
pub blocked: bool,
/// 阻止原因(供 LLM 查看)
pub reason: String,
/// 命令类别(供日志分类)
pub category: Option<String>,
}
impl HardlineResult {
/// 通过检查
pub fn allowed() -> Self {
HardlineResult {
blocked: false,
reason: String::new(),
category: None,
}
}
/// 被拒绝
pub fn denied(reason: String, category: &str) -> Self {
HardlineResult {
blocked: true,
reason,
category: Some(category.to_string()),
}
}
}
/// 检查命令是否命中 hardline 模式。
///
/// 执行反规避标准化后再匹配,返回第一个命中的模式。
///
/// 注意:此函数在模块导入时冻结 `HERMES_YOLO_MODE`
/// 确保运行时无法通过设置环境变量绕过 hardline 检查。
pub fn check_command(raw_command: &str) -> HardlineResult {
let normalized = normalize_command(raw_command);
// ── YOLO 模式冻结 ──
// YOLO 模式在首次调用时从环境变量读取并缓存,
// 后续设置环境变量不会生效(防止注入攻击)。
static YOLO_MODE_FROZEN: LazyLock<bool> = LazyLock::new(|| {
let val = std::env::var("HERMES_YOLO_MODE")
.unwrap_or_default()
.to_lowercase();
val == "1" || val == "true" || val == "yes" || val == "on"
});
// Hardline 即使在 YOLO 模式下也不可绕过
let _yolo = *YOLO_MODE_FROZEN;
let matches: Vec<usize> = HARDLINE_REGEX_SET
.matches(&normalized)
.into_iter()
.collect();
if let Some(&idx) = matches.first() {
let pattern = &HARDLINE_PATTERNS[idx];
let reason = format!(
"⚠️ 命令被硬阻止(安全策略,不可绕过)。\n\
类别: {}\n\
原因: {}\n\
请换用更安全的替代方案实现相同目标。",
pattern.category, pattern.message
);
HardlineResult::denied(reason, pattern.category)
} else {
HardlineResult::allowed()
}
}
/// 仅做标准化(不检查 hardline),用于在其他安全检查前预处理命令。
pub fn normalize_only(raw: &str) -> String {
normalize_command(raw)
}
/// 检查命令是否包含危险的重定向操作。
/// 用于 file_write/file_edit 等非 bash 工具的路径安全检查。
pub fn check_dangerous_path(path: &str) -> HardlineResult {
let normalized = normalize_command(path);
// 检查是否尝试覆盖关键系统文件
let dangerous_prefixes = [
"/etc/passwd",
"/etc/shadow",
"/etc/sudoers",
"/etc/sudoers.d/",
"/etc/ssh/",
"/root/",
"/boot/",
"~/.ssh/authorized_keys",
"~/.ssh/id_rsa",
"~/.ssh/id_ed25519",
"~/.netrc",
"~/.pgpass",
"~/.npmrc",
"~/.pypirc",
"~/.git-credentials",
];
for prefix in &dangerous_prefixes {
if normalized.starts_with(prefix) || normalized.contains(prefix) {
return HardlineResult::denied(
format!(
"路径 '{}' 指向受保护的系统/凭据文件,写入操作被硬阻止。",
path
),
"sensitive_path",
);
}
}
HardlineResult::allowed()
}
#[cfg(test)]
mod tests {
use super::*;
// ── 标准化测试 ──
#[test]
fn test_normalize_backslash_escapes() {
assert_eq!(normalize_command(r"r\m"), "rm");
assert_eq!(normalize_command(r"r\m\ \-\r\f"), "rm -rf");
}
#[test]
fn test_normalize_empty_string_literals() {
assert_eq!(normalize_command("r''m"), "rm");
assert_eq!(normalize_command("r\"\"m"), "rm");
}
#[test]
fn test_normalize_fullwidth() {
// 全角 '' (U+FF52) → 半角 'r'
let fullwidth_rm = "\u{FF52}\u{FF4D}"; // rm
let normalized = normalize_command(fullwidth_rm);
assert_eq!(normalized, "rm");
}
#[test]
fn test_normalize_ansi_strip() {
let cmd = "\x1b[31mrm -rf /\x1b[0m";
let normalized = normalize_command(cmd);
assert_eq!(normalized, "rm -rf /");
}
#[test]
fn test_normalize_null_bytes() {
let cmd = "rm\0 -rf\0 /";
let normalized = normalize_command(cmd);
assert!(!normalized.contains('\0'));
}
// ── Hardline 检查测试 ──
#[test]
fn test_block_shutdown() {
let result = check_command("shutdown -h now");
assert!(result.blocked);
assert_eq!(result.category.as_deref(), Some("system_shutdown"));
}
#[test]
fn test_block_reboot() {
let result = check_command("reboot");
assert!(result.blocked);
}
#[test]
fn test_block_systemctl_poweroff() {
let result = check_command("systemctl poweroff");
assert!(result.blocked);
}
#[test]
fn test_block_mkfs() {
let result = check_command("mkfs.ext4 /dev/sda1");
assert!(result.blocked);
}
#[test]
fn test_block_dd_to_device() {
let result = check_command("dd if=/dev/zero of=/dev/sda bs=1M");
assert!(result.blocked);
}
#[test]
fn test_block_dd_to_nvme() {
let result = check_command("dd if=image.iso of=/dev/nvme0n1");
assert!(result.blocked);
}
#[test]
fn test_block_rm_rf_root() {
let result = check_command("rm -rf /");
assert!(result.blocked);
}
#[test]
fn test_block_rm_rf_root_wildcard() {
let result = check_command("rm -rf /*");
assert!(result.blocked);
}
#[test]
fn test_block_fork_bomb() {
let result = check_command(":(){ :|:& };:");
assert!(result.blocked);
}
#[test]
fn test_block_kill_minus_one() {
let result = check_command("kill -1 1");
assert!(result.blocked);
}
#[test]
fn test_block_redirect_overwrite_passwd() {
let result = check_command("echo 'x' > /etc/passwd");
assert!(result.blocked);
}
#[test]
fn test_block_chmod_777_etc() {
let result = check_command("chmod -R 777 /etc");
assert!(result.blocked);
}
#[test]
fn test_block_chown_root() {
let result = check_command("chown -R root:root /");
assert!(result.blocked);
}
#[test]
fn test_allow_normal_commands() {
assert!(!check_command("ls -la").blocked);
assert!(!check_command("cargo build").blocked);
assert!(!check_command("git status").blocked);
assert!(!check_command("python3 -c 'print(1+1)'").blocked);
}
#[test]
fn test_allow_safe_rm() {
// rm 单个文件不阻止
assert!(!check_command("rm file.txt").blocked);
assert!(!check_command("rm -rf ./node_modules").blocked);
}
#[test]
fn test_allow_dd_to_file() {
// dd 写入普通文件不阻止
assert!(!check_command("dd if=/dev/zero of=test.bin bs=1M count=10").blocked);
}
#[test]
fn test_evasion_backslash_escapes() {
// r\e\b\o\o\t 应该匹配 reboot
let result = check_command(r"r\e\b\o\o\t");
assert!(result.blocked);
}
#[test]
fn test_evasion_empty_quotes() {
// r''m 应该匹配
let result = check_command("r''m -rf /");
assert!(result.blocked);
}
#[test]
fn test_dangerous_path_check() {
let result = check_dangerous_path("/etc/passwd");
assert!(result.blocked);
let result = check_dangerous_path("/home/user/data.txt");
assert!(!result.blocked);
}
}
+204 -28
View File
@@ -11,6 +11,7 @@
// executor — 工具调用验证与并行执行
// finalize — 会话收尾、指标持久化
pub mod checkpoint;
pub mod circuit_breaker;
pub mod context;
pub mod denial_tracker;
@@ -18,6 +19,7 @@ pub mod error_recovery;
pub mod executor;
pub mod file_cache;
pub mod finalize;
pub mod hardline;
pub mod partitioner;
pub mod permission;
pub mod permission_explainer;
@@ -27,6 +29,7 @@ pub mod streaming;
pub mod streaming_executor;
pub mod system_prompt;
pub mod token_budget;
pub mod untrusted;
use serde::Serialize;
use sqlx::SqlitePool;
@@ -46,6 +49,7 @@ use crate::clients::llm::{ChatMessage, LlmClient, MessageRole, StreamEvent};
use self::error_recovery::{classify_error, ErrorKind, ErrorRecovery};
use self::session::SessionInfo;
use self::streaming::{StreamOutput, StreamStatus};
use self::system_prompt::SystemPromptCache;
use self::token_budget::TokenBudget;
/// Agent 配置参数
@@ -303,6 +307,12 @@ pub struct AgentRuntime {
denial_tracker: Arc<std::sync::Mutex<denial_tracker::DenialTracker>>,
/// 文件状态缓存(跨 turn 共享,用于 Read 去重)
read_file_state: Arc<std::sync::Mutex<file_cache::FileStateCache>>,
/// 系统提示词 section 缓存(跨 turn 共享,避免每 turn 重建静态/低频变动内容)
prompt_cache: std::sync::Mutex<SystemPromptCache>,
/// 上下文压缩折叠日志(跨 turn 共享,追踪压缩历史并触发溢出合并)
collapse_log: Arc<std::sync::Mutex<compact::collapse::CollapseLog>>,
/// Checkpoint 管理器(跨 turn 共享,文件变更操作前自动快照)
checkpoint_manager: Arc<checkpoint::CheckpointManager>,
}
impl AgentRuntime {
@@ -334,6 +344,17 @@ impl AgentRuntime {
*session_checker = (*permission_checker).clone();
}
// 初始化 checkpoint 管理器
let checkpoint_enabled = std::env::var("AGENT_CHECKPOINT_ENABLED")
.unwrap_or_else(|_| "true".to_string())
.to_lowercase()
!= "false";
let checkpoint_store = app_state.config.library_dir.join("..").join(".checkpoints");
let checkpoint_manager = Arc::new(checkpoint::CheckpointManager::new(
std::fs::canonicalize(&checkpoint_store).unwrap_or(checkpoint_store),
checkpoint_enabled,
));
AgentRuntime {
app_state,
config,
@@ -346,6 +367,9 @@ impl AgentRuntime {
permission_checker,
denial_tracker,
read_file_state: Arc::new(std::sync::Mutex::new(file_cache::FileStateCache::new())),
prompt_cache: std::sync::Mutex::new(SystemPromptCache::new()),
collapse_log: Arc::new(std::sync::Mutex::new(compact::collapse::CollapseLog::new())),
checkpoint_manager,
}
}
@@ -373,6 +397,18 @@ impl AgentRuntime {
if let Ok(mut session_checker) = app_state.session_permission_checker.write() {
*session_checker = (*permission_checker).clone();
}
// 初始化 checkpoint 管理器
let checkpoint_enabled = std::env::var("AGENT_CHECKPOINT_ENABLED")
.unwrap_or_else(|_| "true".to_string())
.to_lowercase()
!= "false";
let checkpoint_store = app_state.config.library_dir.join("..").join(".checkpoints");
let checkpoint_manager = Arc::new(checkpoint::CheckpointManager::new(
std::fs::canonicalize(&checkpoint_store).unwrap_or(checkpoint_store),
checkpoint_enabled,
));
AgentRuntime {
app_state,
config,
@@ -385,6 +421,9 @@ impl AgentRuntime {
permission_checker,
denial_tracker,
read_file_state: Arc::new(std::sync::Mutex::new(file_cache::FileStateCache::new())),
prompt_cache: std::sync::Mutex::new(SystemPromptCache::new()),
collapse_log: Arc::new(std::sync::Mutex::new(compact::collapse::CollapseLog::new())),
checkpoint_manager,
}
}
@@ -427,12 +466,13 @@ impl AgentRuntime {
}
};
compact::compress_context_with_hooks(
compact::compress_context_with_hooks_and_log(
messages,
llm,
self.config.context_char_limit,
session_id,
Some(hook_registry),
Some(&self.collapse_log),
)
.await;
@@ -579,6 +619,9 @@ impl AgentRuntime {
loop {
step += 1;
// ── Checkpoint: 每个 ReAct 迭代开始时重置去重状态 ──
self.checkpoint_manager.new_turn();
// 检查用户取消
let is_cancelled = {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() {
@@ -923,7 +966,7 @@ impl AgentRuntime {
continue;
}
// 并行执行工具(带权限检查和分区器)
// 并行执行工具(带权限检查、checkpoint 和分区器)
let exec_result = executor::execute_parallel(
&prepared_calls,
&self.tool_registry,
@@ -932,6 +975,7 @@ impl AgentRuntime {
Some(&self.permission_checker),
Some(&self.app_state.session_permission_checker),
Some(&self.denial_tracker),
Some(&self.checkpoint_manager),
tx,
db,
sid,
@@ -961,6 +1005,25 @@ impl AgentRuntime {
messages.push(tm.chat_message);
}
// Hook 注入的附加上下文:包装为 system-reminder 注入 LLM 消息列表
// 使用 ContextDeduplicator 在单步内去重(多个 hook 注入相同内容时只保留一份)
let mut dedup = crate::agent::hooks::ContextDeduplicator::new();
for ctx in &exec_result.hook_contexts {
if dedup.is_duplicate(ctx) {
continue;
}
let reminder = format!(
"<system-reminder>\n[Hook 注入上下文]\n{}\n</system-reminder>",
ctx
);
messages.push(ChatMessage::user(&reminder));
}
// Hook 阻塞错误:记录到日志用于诊断
for be in &exec_result.blocking_errors {
warn!("[AgentRuntime] Hook 阻塞错误: {}", be);
}
// 持久化 todo_write 任务状态到数据库
if called_todo_write {
for prep in &prepared_calls {
@@ -1178,8 +1241,18 @@ impl AgentRuntime {
let mut recovery = ErrorRecovery::new(token_budget.clone());
while let Some(recovery_step) = recovery.try_recover(&error_kind) {
// 尝试从错误消息中解析 ContextOverflow 信息(参考 Claude Code 自动修复)
let overflow_info = error_recovery::parse_context_overflow(&e_str);
while let Some(recovery_step) = recovery.try_recover(&error_kind, overflow_info.as_ref()) {
match recovery_step {
error_recovery::RecoveryStep::AdjustMaxTokens { new_max_tokens } => {
info!(
"[AgentRuntime] 错误恢复: AdjustMaxTokens → {} (从错误消息自动计算)",
new_max_tokens
);
// token_budget.hard_limit 已由 try_recover 下调
}
error_recovery::RecoveryStep::RetryWithBackoff { attempt, delay_ms } => {
// 429/529 本应在 streaming 层处理,若到达此处说明分类逻辑有变更,
// 安全降级为 sleep + 直接重试(不依赖 streaming 层重试)。
@@ -1269,33 +1342,76 @@ impl AgentRuntime {
// ── Helpers ──
/// 系统提示词(模块化组装 — 参考 Claude Code s10)。
/// 静态 section 在前以最大化 prompt cache 命中率。
/// 系统提示词(模块化组装)。
///
/// 设计原则:
/// 1. 所有静态 section 在前 → 内容不变,服务端自然缓存
/// 2. 动态 sectionenvironment/tools/skills/memory)在后
/// 3. 使用 SystemPromptCache:首次计算后永久复用,/clear 时失效
fn system_prompt(&self) -> String {
use self::system_prompt::{SystemPrompt, IDENTITY_SECTION, PRINCIPLES_SECTION};
use self::system_prompt::{
SystemPrompt, IDENTITY_SECTION, PRINCIPLES_SECTION, SAFETY_SECTION,
SYSTEM_CONTEXT_SECTION, TOOL_USAGE_SECTION,
};
let mut sp = SystemPrompt::new();
// Section 1: 静态身份(始终加载,最大化缓存)
sp.add_section("identity", IDENTITY_SECTION.to_string());
// ═══════ 静态 section(首次计算后永久缓存)═══════
// Section 2: 动态工具列表(运行时生成)
let mut tools_desc = String::from("你可以使用以下工具:\n");
for def in self.tool_registry.definitions() {
let short_desc: String = def
.function
.description
.split('。')
.next()
.unwrap_or(&def.function.description)
.chars()
.take(80)
.collect();
tools_desc.push_str(&format!("- {}: {}\n", def.function.name, short_desc));
}
sp.add_section("tools", tools_desc);
let mut cache = self.prompt_cache.lock().unwrap_or_else(|e| {
tracing::warn!("[SystemPrompt] 缓存锁异常: {:?}", e);
e.into_inner()
});
// Section 3: 可用技能(动态)
sp.add_section(
"identity",
cache.get_or_compute("identity", || IDENTITY_SECTION.to_string()),
);
sp.add_section(
"principles",
cache.get_or_compute("principles", || PRINCIPLES_SECTION.to_string()),
);
sp.add_section(
"system_context",
cache.get_or_compute("system_context", || SYSTEM_CONTEXT_SECTION.to_string()),
);
sp.add_section(
"tool_usage",
cache.get_or_compute("tool_usage", || TOOL_USAGE_SECTION.to_string()),
);
sp.add_section(
"safety",
cache.get_or_compute("safety", || SAFETY_SECTION.to_string()),
);
// ═══════ 动态 section(首次计算后缓存,session 内不变)═══════
// 环境上下文:CWD/platform/OS/model 在 session 内不变
let env_section = cache.get_or_compute("environment", || self.build_environment_section());
sp.add_section("environment", env_section);
// 工具列表:ToolRegistry 在 session 内不变
let tools_section = cache.get_or_compute("tools", || {
let mut tools_desc = String::from("你可以使用以下工具:\n");
for def in self.tool_registry.definitions() {
let short_desc: String = def
.function
.description
.split('。')
.next()
.unwrap_or(&def.function.description)
.chars()
.take(80)
.collect();
tools_desc.push_str(&format!("- {}: {}\n", def.function.name, short_desc));
}
tools_desc
});
sp.add_section("tools", tools_section);
drop(cache);
// 技能列表:通过文件监听热更新,不缓存
if let Some(skills) = self
.app_state
.skill_registry
@@ -1306,7 +1422,7 @@ impl AgentRuntime {
sp.add_section("skills", skills);
}
// Section 4: 项目记忆(按需加载)
// 项目记忆:受 save_memory 工具实时影响,不缓存
if let Some(memory) = self
.app_state
.memory_manager
@@ -1317,12 +1433,72 @@ impl AgentRuntime {
sp.add_section("memory", memory);
}
// Section 5: 静态核心原则(最后加载,因较常变化)
sp.add_section("principles", PRINCIPLES_SECTION.to_string());
sp.assemble()
}
/// 使提示词缓存中指定 section 失效。
pub fn invalidate_prompt_cache(&self, section_name: &'static str) {
if let Ok(mut cache) = self.prompt_cache.lock() {
cache.invalidate(section_name);
}
}
/// 使所有提示词缓存失效(`/clear` 或 `/compact` 事件触发)。
pub fn invalidate_all_prompt_cache(&self) {
if let Ok(mut cache) = self.prompt_cache.lock() {
cache.invalidate_all();
}
}
/// 构建环境上下文 section(参考 Claude Code `computeEnvInfo()`)。
///
/// 包含:工作目录、git 状态、平台、OS 版本、日期、模型信息。
fn build_environment_section(&self) -> String {
let cwd = std::env::current_dir()
.map(|p| p.display().to_string())
.unwrap_or_else(|_| "(unknown)".to_string());
let is_git = std::process::Command::new("git")
.args(["rev-parse", "--is-inside-work-tree"])
.output()
.map(|o| o.status.success())
.unwrap_or(false);
let platform = std::env::consts::OS;
let os_version = {
let output = std::process::Command::new("uname")
.args(["-s", "-r"])
.output()
.ok()
.map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string())
.unwrap_or_default();
if output.is_empty() {
std::env::consts::ARCH.to_string()
} else {
output
}
};
let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
let model_name = self.app_state.llm.model().to_string();
let mut lines = vec![
"# 环境信息".to_string(),
format!("- 工作目录: {}", cwd),
format!("- Git 仓库: {}", if is_git { "" } else { "" }),
format!("- 平台: {}", platform),
format!("- OS 版本: {}", os_version),
format!("- 日期: {}", today),
format!("- 当前模型: {}", model_name),
];
// Agent 配置摘要(最大步数、超时等)
lines.push(format!("- 最大推理步数: {}", self.config.max_steps));
lines.push(format!("- 工具超时: {}", self.config.tool_timeout_secs));
lines.join("\n")
}
/// 步数耗尽时的最终答案生成(不带工具调用,强制 LLM 直接回答)
async fn final_answer_without_tools(
&self,
+49
View File
@@ -186,4 +186,53 @@ mod tests {
assert!(batches[1].is_parallel);
assert_eq!(batches[1].calls.len(), 2);
}
#[test]
fn test_run_bash_is_always_serial() {
let registry = make_registry();
let partitioner = ToolPartitioner::new(10);
// run_bash 总是非并发安全的(即使命令是只读的),
// 确保它独立成批
let calls = vec![
make_prep("search_papers"),
make_prep("run_bash"), // 非并发安全
make_prep("rag_search"),
];
let batches = partitioner.partition(&calls, &registry);
// search_papers (并行) → run_bash (串行) → rag_search (并行)
assert_eq!(batches.len(), 3, "run_bash 应打断并发批次");
assert!(batches[0].is_parallel);
assert_eq!(batches[0].calls[0].tool_name, "search_papers");
assert!(!batches[1].is_parallel, "run_bash 必须是串行批次");
assert_eq!(batches[1].calls[0].tool_name, "run_bash");
assert!(batches[2].is_parallel);
assert_eq!(batches[2].calls[0].tool_name, "rag_search");
}
#[test]
fn test_file_write_splits_batch() {
let registry = make_registry();
let partitioner = ToolPartitioner::new(10);
// file_write/file_edit 是非并发安全的
let calls = vec![
make_prep("read_file"),
make_prep("file_write"),
make_prep("read_file"),
];
let batches = partitioner.partition(&calls, &registry);
// read_file (并发) → file_write (串行) → read_file (并发)
assert_eq!(batches.len(), 3, "file_write 应打断并发批次");
assert!(batches[0].is_parallel);
assert_eq!(batches[0].calls[0].tool_name, "read_file");
assert!(!batches[1].is_parallel, "file_write 必须是串行批次");
assert_eq!(batches[1].calls[0].tool_name, "file_write");
assert!(batches[2].is_parallel);
assert_eq!(batches[2].calls[0].tool_name, "read_file");
}
}
+217
View File
@@ -214,6 +214,38 @@ impl PermissionChecker {
}
}
/// 返回一个人类可读的解释,说明某工具为何被允许/拒绝/询问。
/// 用于 hook 审计和调试。
///
/// 遍历规则列表,格式化第一条匹配规则为可读字符串。
pub fn explain(&self, tool_name: &str, tool_args: Option<&serde_json::Value>) -> String {
for rule in &self.rules {
match rule {
PermissionRule::Deny {
tool_name: name,
reason,
..
} if Self::matches(name, tool_name, tool_args) => {
return format!("Denied by rule: {name}{reason}");
}
PermissionRule::Allow {
tool_name: name, ..
} if Self::matches(name, tool_name, tool_args) => {
return format!("Allowed by rule: {name}");
}
PermissionRule::Ask {
tool_name: name,
message,
..
} if Self::matches(name, tool_name, tool_args) => {
return format!("Ask by rule: {name}{message}");
}
_ => {}
}
}
"Allowed by default (no matching rule)".to_string()
}
/// 规则名称匹配:支持精确匹配、通配符 "*",以及内容级匹配。
///
/// 内容级格式:`"tool_name(content_pattern)"`。
@@ -466,6 +498,109 @@ impl PermissionChecker {
}
}
// ── Permission Precedence Resolver ──
/// 多源权限决策的最终裁决。遵循正式的优先级规则表:
///
/// | Priority | Source | Overridable By |
/// |----------|----------------------------------|----------------|
/// | P0 | PermissionChecker::Deny | Nothing |
/// | P1 | Tool-level PermissionRule::Deny | Nothing |
/// | P2 | Session-level Checker::Deny | Nothing |
/// | P3 | Hook PreToolUseAction::Block | P0-P2 |
/// | P4 | Hook PermissionRequired | P0-P3 |
/// | P5-P7 | Checker::Ask / Tool::Ask / Allow | Normal |
///
/// `conflict_log` 记录被覆盖的决策,便于审计和调试。
pub fn resolve_permission_precedence(
checker_result: PermissionResult,
tool_rules: &[crate::agent::tools::PermissionRule],
hook_permission: Option<&(String, String)>, // (permission_desc, tool_name)
hook_blocked: bool,
session_result: Option<PermissionResult>,
) -> (PermissionResult, Vec<String>) {
let mut final_result = checker_result;
let mut conflict_log: Vec<String> = Vec::new();
// ── P1: Tool-level Deny ──
for rule in tool_rules {
if let crate::agent::tools::PermissionRule::Deny { reason, .. } = rule {
if !final_result.is_denied() {
conflict_log.push(format!("Tool-level Deny overrides checker: {reason}"));
final_result = PermissionResult::Denied {
reason: reason.clone(),
};
} else {
conflict_log.push(format!(
"Tool-level Deny '{reason}' ignored: already Denied"
));
}
break; // only handle first Deny
}
}
// ── P2: Session-level Deny ──
if let Some(PermissionResult::Denied { reason }) = &session_result {
conflict_log.push(format!("Session-level Deny overrides current: {reason}"));
final_result = PermissionResult::Denied {
reason: reason.clone(),
};
}
// ── P3: Hook Block ──
if hook_blocked {
conflict_log.push("Hook Block prevents execution".to_string());
// Block is already handled in the executor via denied_indices;
// here we record it for the conflict log.
}
// ── P4: Hook PermissionRequired ──
if let Some((perm_desc, _tool_name)) = hook_permission {
if final_result.is_allowed() {
conflict_log.push(format!(
"Hook PermissionRequired upgrades Allowed → Ask: {perm_desc}"
));
} else if !final_result.is_denied() {
conflict_log.push(format!(
"Hook PermissionRequired coexists with current state: {perm_desc}"
));
} else {
conflict_log.push(format!(
"Hook PermissionRequired '{perm_desc}' ignored: already Denied"
));
}
}
// ── P5: Session-level Ask ──
if let Some(PermissionResult::AskUser { message }) = &session_result {
if final_result.is_allowed() {
final_result = PermissionResult::AskUser {
message: message.clone(),
};
conflict_log.push("Session-level Ask overrides Allow".to_string());
} else {
conflict_log.push("Session-level Ask ignored: not Allowed".to_string());
}
}
// ── P6: Tool-level Ask ──
for rule in tool_rules {
if let crate::agent::tools::PermissionRule::Ask { message, .. } = rule {
if final_result.is_allowed() {
conflict_log.push(format!("Tool-level Ask upgrades Allow: {message}"));
final_result = PermissionResult::AskUser {
message: message.clone(),
};
} else {
conflict_log.push("Tool-level Ask ignored: not Allowed".to_string());
}
break;
}
}
(final_result, conflict_log)
}
#[cfg(test)]
mod tests {
use super::*;
@@ -819,4 +954,86 @@ mod tests {
// AcceptEdits 保留 AskUser 让 executor 做路径检查
assert!(matches!(result, PermissionResult::AskUser { .. }));
}
// ── resolve_permission_precedence tests ──
#[test]
fn test_precedence_checker_deny_wins_over_all() {
let (result, log) = resolve_permission_precedence(
PermissionResult::Denied {
reason: "blocked by policy".into(),
},
&[],
Some(&("need confirmation".to_string(), "test_tool".to_string())),
false,
None,
);
assert!(result.is_denied());
assert!(
!log.is_empty(),
"conflict log should record the interaction"
);
}
#[test]
fn test_precedence_tool_deny_overrides_allow() {
use crate::agent::tools::{PermissionRule, PermissionRuleSource};
let tool_rules = vec![PermissionRule::Deny {
tool_name: "test_tool".into(),
reason: "tool self-protection".into(),
source: PermissionRuleSource::Env,
}];
let (result, log) = resolve_permission_precedence(
PermissionResult::Allowed,
&tool_rules,
None,
false,
None,
);
assert!(result.is_denied());
assert!(!log.is_empty());
}
#[test]
fn test_precedence_hook_block_recorded() {
let (result, log) = resolve_permission_precedence(
PermissionResult::Allowed,
&[],
None,
true, // hook blocked
None,
);
// Hook Block doesn't directly return Denied — it's logged for executor handling
assert!(result.is_allowed());
assert!(log.iter().any(|l| l.contains("Block")));
}
#[test]
fn test_precedence_session_deny_overrides() {
let (result, _log) = resolve_permission_precedence(
PermissionResult::Allowed,
&[],
None,
false,
Some(PermissionResult::Denied {
reason: "session deny".into(),
}),
);
assert!(result.is_denied());
}
#[test]
fn test_precedence_hook_permission_ignored_when_denied() {
let (result, log) = resolve_permission_precedence(
PermissionResult::Denied {
reason: "policy deny".into(),
},
&[],
Some(&("need confirm".to_string(), "test_tool".to_string())),
false,
None,
);
assert!(result.is_denied());
assert!(log.iter().any(|l| l.contains("ignored")));
}
}
+837 -16
View File
@@ -1,8 +1,10 @@
// src/agent/runtime/session.rs
//
// 会话生命周期管理:创建/恢复/验证 Agent 会话。
// 支持软删除回退 (undo)active=0 标记保留审计 trailLLM 不可见。
use sqlx::SqlitePool;
use tracing::info;
use crate::clients::llm::LlmClient;
@@ -13,9 +15,18 @@ pub struct SessionInfo {
pub turn_index: i32,
}
/// 回退操作结果
#[derive(Debug, Clone)]
pub struct RewindResult {
/// 被软删除的消息数
pub rewound_count: usize,
/// 目标消息的内容预览(供 UI 展示)
pub target_preview: String,
/// 回退后的 turn_index
pub new_turn_index: i32,
}
/// 创建新会话或恢复已有会话。
///
/// 返回会话信息。如果指定的 session_id 不存在则返回错误。
pub async fn create_or_resume_session(
db: &SqlitePool,
session_id: Option<String>,
@@ -23,7 +34,6 @@ pub async fn create_or_resume_session(
) -> anyhow::Result<SessionInfo> {
match session_id {
Some(id) => {
// 验证会话存在且未被软删除
let exists: bool = sqlx::query_scalar(
"SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)",
)
@@ -36,9 +46,10 @@ pub async fn create_or_resume_session(
return Err(anyhow::anyhow!("会话 {} 不存在或已删除", id));
}
// 计算当前轮次号
// 计算当前轮次号(仅统计 active=1 的消息)
let turn_index: i32 = sqlx::query_scalar(
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages WHERE session_id = ?",
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages \
WHERE session_id = ? AND active = 1",
)
.bind(&id)
.fetch_one(db)
@@ -68,10 +79,6 @@ pub async fn create_or_resume_session(
}
/// 加载会话的历史消息(供 LLM 上下文使用)。
///
/// `agent_name` 参数用于消息隔离:
/// - `"lead"` — 只加载 Lead Agent 自己的消息(默认)
/// - `"*"` — 加载所有 agent 的消息(调试/审计用)
pub async fn load_history_for_llm(
db: &SqlitePool,
session_id: &str,
@@ -79,14 +86,29 @@ pub async fn load_history_for_llm(
load_history_for_agent(db, session_id, "lead").await
}
/// 加载指定 agent 的历史消息。
/// 加载指定 agent 的历史消息(仅 active=1
pub async fn load_history_for_agent(
db: &SqlitePool,
session_id: &str,
agent_name: &str,
) -> anyhow::Result<Vec<crate::clients::llm::ChatMessage>> {
load_history_for_agent_impl(db, session_id, agent_name, true).await
}
/// 加载指定 agent 的历史消息。
///
/// `active_only`: true 时仅加载 active=1LLM 上下文),
/// false 时加载全部(审计/调试用)。
async fn load_history_for_agent_impl(
db: &SqlitePool,
session_id: &str,
agent_name: &str,
active_only: bool,
) -> anyhow::Result<Vec<crate::clients::llm::ChatMessage>> {
use crate::clients::llm::{ChatMessage, MessageRole};
let active_filter = if active_only { " AND active = 1" } else { "" };
#[allow(clippy::type_complexity)]
let rows: Vec<(
String,
@@ -95,18 +117,20 @@ pub async fn load_history_for_agent(
Option<String>,
Option<String>,
)> = if agent_name == "*" {
sqlx::query_as(
sqlx::query_as(&format!(
"SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \
WHERE session_id = ? ORDER BY id ASC",
)
WHERE session_id = ?{} ORDER BY id ASC",
active_filter
))
.bind(session_id)
.fetch_all(db)
.await?
} else {
sqlx::query_as(
sqlx::query_as(&format!(
"SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \
WHERE session_id = ? AND agent_name = ? ORDER BY id ASC",
)
WHERE session_id = ? AND agent_name = ?{} ORDER BY id ASC",
active_filter
))
.bind(session_id)
.bind(agent_name)
.fetch_all(db)
@@ -142,3 +166,800 @@ pub async fn load_history_for_agent(
Ok(messages)
}
// ── Rewind / Undo API ──
/// 回退会话到指定用户消息之前。
///
/// 软删除:将目标消息及之后的所有 active=1 消息设置为 active=0。
/// 返回被软删除的消息数和目标消息预览。
pub async fn rewind_to_message(
db: &SqlitePool,
session_id: &str,
target_message_id: i64,
) -> anyhow::Result<RewindResult> {
// 1. 验证目标消息是当前 session 的 user 消息且 active=1
let target: Option<(String, i32)> = sqlx::query_as(
"SELECT content, turn_index FROM agent_messages \
WHERE id = ? AND session_id = ? AND role = 'user' AND active = 1",
)
.bind(target_message_id)
.bind(session_id)
.fetch_optional(db)
.await?;
let (target_content, _target_turn) = match target {
Some(t) => t,
None => {
return Err(anyhow::anyhow!(
"目标消息 {} 不存在、不是用户消息、或已被回退",
target_message_id
));
}
};
// 2. 查找 >= target_id 的所有 active=1 消息
let to_rewind: Vec<i64> = sqlx::query_scalar(
"SELECT id FROM agent_messages \
WHERE session_id = ? AND id >= ? AND active = 1 \
ORDER BY id DESC",
)
.bind(session_id)
.bind(target_message_id)
.fetch_all(db)
.await?;
if to_rewind.is_empty() {
return Ok(RewindResult {
rewound_count: 0,
target_preview: String::new(),
new_turn_index: 0,
});
}
// 3. 原子执行软删除(单个事务)
let mut tx = db.begin().await?;
let count = to_rewind.len();
let first_rewound_id = to_rewind.last().copied().unwrap_or(0);
sqlx::query(
"UPDATE agent_messages SET active = 0 \
WHERE session_id = ? AND id >= ? AND active = 1",
)
.bind(session_id)
.bind(first_rewound_id)
.execute(&mut *tx)
.await?;
// 4. 更新 rewind_count
sqlx::query(
"UPDATE agent_sessions SET rewind_count = rewind_count + 1, \
updated_at = CURRENT_TIMESTAMP WHERE session_id = ?",
)
.bind(session_id)
.execute(&mut *tx)
.await?;
tx.commit().await?;
// 5. 计算新的 turn_index(目标消息之前的 turn
let new_turn_index: i32 = sqlx::query_scalar(
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages \
WHERE session_id = ? AND active = 1",
)
.bind(session_id)
.fetch_one(db)
.await
.unwrap_or(0);
let preview: String = target_content.chars().take(120).collect();
info!(
"[Session] 回退完成: session={}, rewound={}, to_id={}, new_turn={}",
session_id, count, target_message_id, new_turn_index
);
Ok(RewindResult {
rewound_count: count,
target_preview: if preview.len() >= 120 {
format!("{}...", preview)
} else {
preview
},
new_turn_index,
})
}
/// 回退最新的 N 个用户轮次。
///
/// 查找最近 N 个 user 消息,回退到第 N 个之前。
///
/// 返回 `RewindResult`,若没有足够的 user 消息则回退全部。
pub async fn rewind_n_turns(
db: &SqlitePool,
session_id: &str,
n: usize,
) -> anyhow::Result<RewindResult> {
let n = n.max(1);
// 查找最近的 N 个 user 消息(按 id DESC
let user_ids: Vec<i64> = sqlx::query_scalar(
"SELECT id FROM agent_messages \
WHERE session_id = ? AND role = 'user' AND active = 1 \
ORDER BY id DESC \
LIMIT ?",
)
.bind(session_id)
.bind(n as i64)
.fetch_all(db)
.await?;
if user_ids.is_empty() {
return Ok(RewindResult {
rewound_count: 0,
target_preview: "没有可回退的消息".to_string(),
new_turn_index: 0,
});
}
// 回退到最早的 user 消息(第 N 个)的位置
let target_id = user_ids.last().copied().unwrap();
rewind_to_message(db, session_id, target_id).await
}
/// 恢复最近一次回退操作(undo-of-undo)。
///
/// 将所有 active=0 的消息恢复为 active=1。
///
/// **安全约束**: 如果回退后产生了新对话(有 active=1 消息的 id 大于
/// 被回退消息的 id),则拒绝恢复,因为这会导效消息穿插乱序。
/// 此种情况请使用 `/branch` 分叉到回退点后再探索替代路径。
///
/// 仅在"刚回退,尚未发送新消息"的场景下可安全使用。
pub async fn restore_rewound(db: &SqlitePool, session_id: &str) -> anyhow::Result<usize> {
// ── 冲突检测 ──
// 查找最小的 inactive 消息 id 和最大的 active 消息 id。
// 如果 max_active_id > min_inactive_id,说明回退后产生了新消息,
// 恢复会导致旧消息穿插在新消息之间。
let conflict: Option<(i64, i64)> = sqlx::query_as(
"SELECT \
(SELECT COALESCE(MAX(id), 0) FROM agent_messages WHERE session_id = ? AND active = 1), \
(SELECT COALESCE(MIN(id), 0) FROM agent_messages WHERE session_id = ? AND active = 0)",
)
.bind(session_id)
.bind(session_id)
.fetch_optional(db)
.await?
.map(|(max_active, min_inactive): (i64, i64)| (max_active, min_inactive))
.filter(|(max_active, min_inactive)| *max_active > 0 && *min_inactive > 0 && max_active > min_inactive);
if let Some((max_active, min_inactive)) = conflict {
return Err(anyhow::anyhow!(
"无法恢复回退:回退后已产生 {} 条新消息 (id {} ~ {})。\
旧消息 (id {}) 的恢复会与当前对话冲突。\
如需回到之前状态,请对当前对话再次执行 /rewind。",
sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM agent_messages WHERE session_id = ? AND active = 1 AND id > ?"
)
.bind(session_id)
.bind(min_inactive)
.fetch_one(db)
.await
.unwrap_or(0),
min_inactive + 1,
max_active,
min_inactive
));
}
let count: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM agent_messages \
WHERE session_id = ? AND active = 0",
)
.bind(session_id)
.fetch_one(db)
.await?;
if count == 0 {
return Ok(0);
}
sqlx::query(
"UPDATE agent_messages SET active = 1 \
WHERE session_id = ? AND active = 0",
)
.bind(session_id)
.execute(db)
.await?;
// 回退 rewind_count
sqlx::query(
"UPDATE agent_sessions SET rewind_count = MAX(0, rewind_count - 1), \
updated_at = CURRENT_TIMESTAMP WHERE session_id = ?",
)
.bind(session_id)
.execute(db)
.await?;
info!(
"[Session] 恢复回退: session={}, restored={} messages",
session_id, count
);
Ok(count as usize)
}
/// 重试最后一次对话(/retry 命令)。
///
/// 硬删除最后一条用户消息及之后的所有消息,返回被删除的用户消息文本。
/// 与 `/rewind`soft-delete)不同,此操作物理删除行,数据不可恢复。
///
/// 返回 `(deleted_message_text, new_turn_index)`。
/// 如果没有找到用户消息,返回错误。
pub async fn retry_last_turn(db: &SqlitePool, session_id: &str) -> anyhow::Result<(String, i32)> {
// 查找最后一条 user 消息(仅 active=1
let last_user: Option<(i64, String, i32)> = sqlx::query_as(
"SELECT id, content, turn_index FROM agent_messages \
WHERE session_id = ? AND role = 'user' AND active = 1 \
ORDER BY id DESC LIMIT 1",
)
.bind(session_id)
.fetch_optional(db)
.await?;
let (target_id, message_text, _turn) = match last_user {
Some(t) => t,
None => return Err(anyhow::anyhow!("没有找到可重试的用户消息")),
};
// 硬删除 >= target_id 的所有消息(含 active=0 的历史回退消息)
let deleted: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM agent_messages \
WHERE session_id = ? AND id >= ?",
)
.bind(session_id)
.bind(target_id)
.fetch_one(db)
.await?;
sqlx::query(
"DELETE FROM agent_messages \
WHERE session_id = ? AND id >= ?",
)
.bind(session_id)
.bind(target_id)
.execute(db)
.await?;
// 计算新的 turn_index
let new_turn_index: i32 = sqlx::query_scalar(
"SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages \
WHERE session_id = ? AND active = 1",
)
.bind(session_id)
.fetch_one(db)
.await
.unwrap_or(0);
info!(
"[Session] 重试: session={}, deleted={} messages from id={}, new_turn={}",
session_id, deleted, target_id, new_turn_index
);
Ok((message_text, new_turn_index))
}
/// 分叉结果
#[derive(Debug, Clone)]
pub struct BranchResult {
/// 新分支的 session_id
pub branch_session_id: String,
/// 分叉点:原始会话中最后保留的消息 id
pub forked_at_message_id: i64,
/// 复制的消息数
pub copied_count: usize,
}
/// 创建会话分叉(/branch 命令)。
///
/// 将当前会话的所有 active=1 消息复制到新会话,
/// 新会话通过 `parent_session_id` 追溯源会话。
///
/// 分叉后两条分支完全独立,各自继续对话互不影响。
/// 这是 Hermes 推荐的"回到过去探索替代路径"方案。
pub async fn branch_session(db: &SqlitePool, session_id: &str) -> anyhow::Result<BranchResult> {
// 1. 验证源会话存在
let exists: bool = sqlx::query_scalar(
"SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)",
)
.bind(session_id)
.fetch_one(db)
.await?;
if !exists {
return Err(anyhow::anyhow!("会话 {} 不存在或已删除", session_id));
}
// 2. 获取源会话的 title
let title: String =
sqlx::query_scalar("SELECT COALESCE(title, '') FROM agent_sessions WHERE session_id = ?")
.bind(session_id)
.fetch_one(db)
.await
.unwrap_or_default();
// 3. 查找最后一条 active=1 的消息 id(作为分叉点)
let last_active_id: Option<i64> = sqlx::query_scalar(
"SELECT MAX(id) FROM agent_messages WHERE session_id = ? AND active = 1",
)
.bind(session_id)
.fetch_one(db)
.await?;
let forked_at = last_active_id.unwrap_or(0);
// 4. 创建新会话
let branch_id = uuid::Uuid::new_v4().to_string();
let branch_title = if title.is_empty() {
format!("分支 (来自 {})", &session_id[..8.min(session_id.len())])
} else {
format!("{} — 分支", title)
};
let branch_meta = serde_json::json!({
"branched_from": session_id,
"branched_at_message_id": forked_at,
});
sqlx::query(
"INSERT INTO agent_sessions (session_id, title, model, parent_session_id, branch_metadata) \
VALUES (?, ?, '', ?, ?)",
)
.bind(&branch_id)
.bind(&branch_title)
.bind(session_id)
.bind(serde_json::to_string(&branch_meta).unwrap_or_default())
.execute(db)
.await?;
// 5. 复制所有 active=1 的消息到新会话
let copied: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM agent_messages WHERE session_id = ? AND active = 1",
)
.bind(session_id)
.fetch_one(db)
.await?;
sqlx::query(
"INSERT INTO agent_messages \
(session_id, turn_index, step_index, role, content, thought, tool_calls, \
tool_call_id, token_count, metadata, raw_json, agent_name, active) \
SELECT ?, turn_index, step_index, role, content, thought, tool_calls, \
tool_call_id, token_count, metadata, raw_json, agent_name, active \
FROM agent_messages \
WHERE session_id = ? AND active = 1 \
ORDER BY id ASC",
)
.bind(&branch_id)
.bind(session_id)
.execute(db)
.await?;
info!(
"[Session] 分叉完成: parent={}, branch={}, copied={} messages, forked_at={}",
session_id, branch_id, copied, forked_at
);
Ok(BranchResult {
branch_session_id: branch_id,
forked_at_message_id: forked_at,
copied_count: copied as usize,
})
}
/// 获取会话的回退次数。
pub async fn get_rewind_count(db: &SqlitePool, session_id: &str) -> i32 {
sqlx::query_scalar("SELECT rewind_count FROM agent_sessions WHERE session_id = ?")
.bind(session_id)
.fetch_one(db)
.await
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
async fn setup_db() -> SqlitePool {
let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
sqlx::query(
"CREATE TABLE agent_sessions (
session_id TEXT PRIMARY KEY,
title TEXT NOT NULL DEFAULT '',
model TEXT NOT NULL DEFAULT '',
turn_count INTEGER NOT NULL DEFAULT 0,
rewind_count INTEGER NOT NULL DEFAULT 0,
parent_session_id TEXT REFERENCES agent_sessions(session_id),
branch_metadata TEXT,
last_error TEXT,
summary TEXT,
metadata TEXT,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
deleted_at DATETIME
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"CREATE TABLE agent_messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL,
turn_index INTEGER NOT NULL DEFAULT 0,
step_index INTEGER NOT NULL DEFAULT 0,
role TEXT NOT NULL CHECK(role IN ('system', 'user', 'assistant', 'tool')),
content TEXT NOT NULL DEFAULT '',
thought TEXT,
tool_calls TEXT,
tool_call_id TEXT,
token_count INTEGER NOT NULL DEFAULT 0,
metadata TEXT,
raw_json TEXT,
agent_name TEXT NOT NULL DEFAULT 'lead',
active INTEGER NOT NULL DEFAULT 1,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE
)",
)
.execute(&pool)
.await
.unwrap();
pool
}
async fn seed_messages(db: &SqlitePool, session_id: &str) {
// Session must be created BEFORE messages (FK constraint)
sqlx::query("INSERT INTO agent_sessions (session_id) VALUES (?)")
.bind(session_id)
.execute(db)
.await
.unwrap();
// Insert system + 3 user turns with responses
for (id, role, turn, content) in [
(1, "system", 0, "You are a helpful assistant."),
(2, "user", 0, "Question 1"),
(3, "assistant", 0, "Answer 1"),
(4, "user", 1, "Question 2"),
(5, "assistant", 1, "Answer 2"),
(6, "user", 2, "Question 3"),
(7, "assistant", 2, "Answer 3"),
] {
sqlx::query(
"INSERT INTO agent_messages (id, session_id, role, turn_index, content, agent_name, active) \
VALUES (?, ?, ?, ?, ?, 'lead', 1)",
)
.bind(id)
.bind(session_id)
.bind(role)
.bind(turn)
.bind(content)
.execute(db)
.await
.unwrap();
}
}
#[tokio::test]
async fn test_load_history_filters_active() {
let db = setup_db().await;
let sid = "test-active-filter";
seed_messages(&db, sid).await;
// Before rewind: all 7 messages active
let msgs = load_history_for_llm(&db, sid).await.unwrap();
assert_eq!(msgs.len(), 7);
// Rewind to message 6 (Question 3) → soft-delete ids 6,7
let result = rewind_to_message(&db, sid, 6).await.unwrap();
assert_eq!(result.rewound_count, 2);
// After rewind: only 5 messages active (ids 1-5)
let msgs = load_history_for_llm(&db, sid).await.unwrap();
assert_eq!(msgs.len(), 5);
}
#[tokio::test]
async fn test_rewind_to_message_keeps_system() {
let db = setup_db().await;
let sid = "test-keep-system";
seed_messages(&db, sid).await;
// Rewind back to first user message (id=2)
let result = rewind_to_message(&db, sid, 2).await.unwrap();
assert!(result.rewound_count >= 1);
let msgs = load_history_for_llm(&db, sid).await.unwrap();
// Should have system (id=1) + target user (id=2 itself is rewound)
// Actually, rewind_to_message(2) soft-deletes id >= 2
// So only system (id=1) remains
assert_eq!(msgs.len(), 1);
assert_eq!(
msgs[0].content.as_deref(),
Some("You are a helpful assistant.")
);
}
#[tokio::test]
async fn test_rewind_n_turns() {
let db = setup_db().await;
let sid = "test-rewind-n";
seed_messages(&db, sid).await;
// Rewind 2 turns → should go back to before Question 2 (id=4)
let result = rewind_n_turns(&db, sid, 2).await.unwrap();
assert!(result.rewound_count >= 1);
let msgs = load_history_for_llm(&db, sid).await.unwrap();
// System (1) + Q1 (2) + A1 (3) = 3 messages
assert_eq!(msgs.len(), 3);
}
#[tokio::test]
async fn test_restore_rewound() {
let db = setup_db().await;
let sid = "test-restore";
seed_messages(&db, sid).await;
// Rewind to id=4 → soft-delete 4,5,6,7
rewind_to_message(&db, sid, 4).await.unwrap();
assert_eq!(load_history_for_llm(&db, sid).await.unwrap().len(), 3);
// Restore
let restored = restore_rewound(&db, sid).await.unwrap();
assert_eq!(restored, 4);
// All 7 messages back
assert_eq!(load_history_for_llm(&db, sid).await.unwrap().len(), 7);
}
#[tokio::test]
async fn test_rewind_invalid_message() {
let db = setup_db().await;
let sid = "test-invalid";
seed_messages(&db, sid).await;
// Try to rewind to a non-existent message
let result = rewind_to_message(&db, sid, 999).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("不存在"));
}
#[tokio::test]
async fn test_restore_rejects_after_new_messages() {
let db = setup_db().await;
let sid = "test-conflict";
seed_messages(&db, sid).await;
// Rewind to id=4 (soft-delete 4-7)
rewind_to_message(&db, sid, 4).await.unwrap();
assert_eq!(load_history_for_llm(&db, sid).await.unwrap().len(), 3);
// Add a new message after rewind (simulates new conversation)
sqlx::query(
"INSERT INTO agent_messages (id, session_id, role, turn_index, content, agent_name, active) \
VALUES (8, ?, 'user', 2, 'New question', 'lead', 1)",
)
.bind(sid)
.execute(&db)
.await
.unwrap();
// Now trying to restore should FAIL because new messages exist after old inactive ones
let result = restore_rewound(&db, sid).await;
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(
err.contains("无法恢复回退"),
"Expected conflict error, got: {}",
err
);
}
#[tokio::test]
async fn test_rewind_count_tracks() {
let db = setup_db().await;
let sid = "test-count";
seed_messages(&db, sid).await;
assert_eq!(get_rewind_count(&db, sid).await, 0);
rewind_to_message(&db, sid, 6).await.unwrap();
assert_eq!(get_rewind_count(&db, sid).await, 1);
rewind_to_message(&db, sid, 4).await.unwrap();
assert_eq!(get_rewind_count(&db, sid).await, 2);
}
// ── /retry tests ──
#[tokio::test]
async fn test_retry_last_turn_deletes_and_returns_message() {
let db = setup_db().await;
let sid = "test-retry-basic";
seed_messages(&db, sid).await;
// Last user message is "Question 3" (id=6)
let (msg, new_turn) = retry_last_turn(&db, sid).await.unwrap();
assert_eq!(msg, "Question 3");
assert_eq!(new_turn, 2); // turn_index after removing id=6,7
// Only messages 1-5 should remain
let msgs = load_history_for_llm(&db, sid).await.unwrap();
assert_eq!(msgs.len(), 5);
assert_eq!(msgs.last().unwrap().content.as_deref(), Some("Answer 2"));
}
#[tokio::test]
async fn test_retry_on_empty_session_errors() {
let db = setup_db().await;
let sid = "test-retry-empty";
sqlx::query("INSERT INTO agent_sessions (session_id) VALUES (?)")
.bind(sid)
.execute(&db)
.await
.unwrap();
let result = retry_last_turn(&db, sid).await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_retry_after_rewind_deletes_inactive_too() {
let db = setup_db().await;
let sid = "test-retry-after-rewind";
seed_messages(&db, sid).await;
// First rewind to id=4 (soft-delete 4-7)
rewind_to_message(&db, sid, 4).await.unwrap();
// Now: ids 1-3 active=1, ids 4-7 active=0
// Now retry — should hard DELETE from last active user message (id=2)
let (msg, _new_turn) = retry_last_turn(&db, sid).await.unwrap();
assert_eq!(msg, "Question 1"); // last active user message
// Only system message should remain
let msgs = load_history_for_llm(&db, sid).await.unwrap();
assert_eq!(msgs.len(), 1);
assert_eq!(
msgs[0].content.as_deref(),
Some("You are a helpful assistant.")
);
// Even inactive messages (4-7) should be gone (hard DELETE)
let all_count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_messages WHERE session_id = ?")
.bind(sid)
.fetch_one(&db)
.await
.unwrap();
assert_eq!(all_count, 1);
}
// ── /branch tests ──
#[tokio::test]
async fn test_branch_copies_active_messages() {
let db = setup_db().await;
let sid = "test-branch-copy";
seed_messages(&db, sid).await;
// Rewind to id=6 first (soft-delete 6,7) so only 1-5 are active
rewind_to_message(&db, sid, 6).await.unwrap();
// Create branch
let result = branch_session(&db, sid).await.unwrap();
assert_eq!(result.copied_count, 5); // only active=1 messages (ids 1-5)
assert!(result.forked_at_message_id > 0);
// New branch has independent history
let branch_msgs = load_history_for_llm(&db, &result.branch_session_id)
.await
.unwrap();
assert_eq!(branch_msgs.len(), 5);
// Original session unchanged
let orig_msgs = load_history_for_llm(&db, sid).await.unwrap();
assert_eq!(orig_msgs.len(), 5);
}
#[tokio::test]
async fn test_branch_independent_continuation() {
let db = setup_db().await;
let sid = "test-branch-independent";
seed_messages(&db, sid).await;
let result = branch_session(&db, sid).await.unwrap();
let bid = result.branch_session_id;
// Add a new message to the branch
sqlx::query(
"INSERT INTO agent_messages (session_id, role, turn_index, content, agent_name, active) \
VALUES (?, 'user', 3, 'Branch question', 'lead', 1)",
)
.bind(&bid)
.execute(&db)
.await
.unwrap();
// Branch sees new message
let branch_msgs = load_history_for_llm(&db, &bid).await.unwrap();
let has_branch_msg = branch_msgs
.iter()
.any(|m| m.content.as_deref() == Some("Branch question"));
assert!(has_branch_msg);
// Original does NOT see branch message
let orig_msgs = load_history_for_llm(&db, sid).await.unwrap();
let has_branch_msg = orig_msgs
.iter()
.any(|m| m.content.as_deref() == Some("Branch question"));
assert!(!has_branch_msg);
}
#[tokio::test]
async fn test_branch_keeps_parent_link() {
let db = setup_db().await;
let sid = "test-branch-parent";
seed_messages(&db, sid).await;
let result = branch_session(&db, sid).await.unwrap();
// Check parent_session_id is set
let parent: Option<String> =
sqlx::query_scalar("SELECT parent_session_id FROM agent_sessions WHERE session_id = ?")
.bind(&result.branch_session_id)
.fetch_one(&db)
.await
.unwrap();
assert_eq!(parent, Some(sid.to_string()));
}
#[tokio::test]
async fn test_branch_invalid_session_errors() {
let db = setup_db().await;
let result = branch_session(&db, "nonexistent-session").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_retry_keeps_inactive_before_target() {
// If there are inactive messages before the retry target,
// they should survive the DELETE (since DELETE is id >= target_id)
let db = setup_db().await;
let sid = "test-retry-keep-inactive";
seed_messages(&db, sid).await;
// First rewind to id=6 (soft-delete 6,7)
rewind_to_message(&db, sid, 6).await.unwrap();
// Now: ids 1-5 active=1, ids 6-7 active=0
// Now rewind again to id=4 (soft-delete 4,5)
rewind_to_message(&db, sid, 4).await.unwrap();
// Now: ids 1-3 active=1, ids 4-7 active=0
// Retry: last active user is id=2. Hard DELETE ids >= 2.
// This removes everything: 1 (system), 2 (user), 3 (asst), and 4-7 (inactive)
let (msg, _) = retry_last_turn(&db, sid).await.unwrap();
assert_eq!(msg, "Question 1");
// Only system remains
let msgs = load_history_for_llm(&db, sid).await.unwrap();
assert_eq!(msgs.len(), 1);
}
}
+343 -159
View File
@@ -6,45 +6,33 @@
// 当 LLM 流式输出 tool_use 块时,立即开始执行并发安全的工具。
// 非并发安全的工具排队等待。结果按流中顺序 yield。
//
// 功能
// 1. 流式执行 — tool_use 到达时立即调度
// 2. Sibling Abort — 副效应工具报错时中止兄弟并行执行
// 3. Progress 流式 — 长时间操作可发送进度更新
// 与 Claude Code 的对齐改进 (2026-06-22)
// 1. 真正的流式调度 — on_tool_use 中对并发安全工具立即 spawn tokio task
// 2. 并发分区 — 自动分组连续只读工具并行执行
// 3. Progress 流式 — 长操作进度消息即时 yield
// 4. Sibling Abort — 副效应工具报错时级联中止兄弟姐妹
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::{broadcast, mpsc, oneshot};
use tokio::sync::broadcast;
use tokio::task::JoinHandle;
use tracing::{info, warn};
use super::partitioner::ToolPartitioner;
use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry};
/// 流式工具执行状态
#[derive(Debug, Clone, PartialEq)]
pub enum TrackedToolStatus {
/// 工具调用已从 LLM 流中接收到
/// 工具调用已从 LLM 流中接收到,等待调度
Queued,
/// 正在执行中
/// 正在执行中spawned tokio task 运行中)
Executing,
/// 执行完成,等待 yield
/// 执行完成,结果就绪等待 yield
Completed,
/// 结果已 yield 给调用方
Yielded,
}
/// 跟踪中的工具执行
#[derive(Debug)]
struct TrackedTool {
tool_call_id: String,
tool_name: String,
args: serde_json::Value,
status: TrackedToolStatus,
/// 执行完成后的输出
output: Option<ToolOutput>,
/// 取消通道(Sibling Abort 使用)
#[allow(dead_code)]
cancel_tx: Option<oneshot::Sender<()>>,
}
/// Sibling Abort 原因
#[derive(Debug, Clone)]
pub enum AbortReason {
@@ -54,33 +42,59 @@ pub enum AbortReason {
UserInterrupted,
}
/// 流式工具执行
/// 单次工具执行的结果
#[derive(Debug)]
struct ToolExecutionResult {
tool_name: String,
output: ToolOutput,
}
/// 跟踪中的工具执行
struct TrackedTool {
tool_call_id: String,
tool_name: String,
args: serde_json::Value,
status: TrackedToolStatus,
/// 执行完成后的输出
output: Option<ToolOutput>,
/// 并发安全的工具在 spawn 后的 JoinHandle
handle: Option<JoinHandle<ToolExecutionResult>>,
}
/// 流式工具执行器。
///
/// 参考 Claude Code `StreamingToolExecutor` (531 行 TypeScript)
/// 关键改进:并发安全工具立即 spawn tokio task,不等待 flush。
pub struct StreamingToolExecutor {
/// 所有跟踪中的工具
/// 所有跟踪中的工具(按 LLM 流中到达顺序)
tracked: Vec<TrackedTool>,
/// 工具注册表
tool_registry: Arc<ToolRegistry>,
/// 并发分区器(保留用于未来并发策略优化
#[allow(dead_code)]
partitioner: ToolPartitioner,
/// 工具上下文
/// 工具上下文(按需 clone 给每个 spawn 的 task
tool_context: ToolContext,
/// Sibling Abort 广播通道 (tx)
abort_tx: broadcast::Sender<AbortReason>,
/// Sibling Abort 广播通道 (rx)
/// Sibling Abort 广播通道 (rx) — 保留以保持 channel 存活,
/// 实际使用时通过 `abort_tx.subscribe()` 获取新接收端。
#[allow(dead_code)]
abort_rx: broadcast::Receiver<AbortReason>,
/// 当前是否已发生错误(触发 sibling abort
has_errored: bool,
/// 出错工具的描述
/// 出错工具的描述(如 "bash(git push)"
errored_tool_desc: String,
/// 下一个 stream_index
next_index: usize,
/// 最大并发数(预留,当前使用 executing_non_concurrent 控制)
#[allow(dead_code)]
max_concurrency: usize,
/// 最大工具输出字符数
max_output_chars: usize,
/// 当前正在执行的非并发安全工具数(0 或 1)
executing_non_concurrent: bool,
/// 已完成但尚未 yield 的结果队列(按流顺序)
completed_queue: VecDeque<usize>,
}
impl StreamingToolExecutor {
/// 创建新的流式执行器
/// 创建新的流式执行器
pub fn new(
tool_registry: Arc<ToolRegistry>,
tool_context: ToolContext,
@@ -91,86 +105,109 @@ impl StreamingToolExecutor {
StreamingToolExecutor {
tracked: Vec::new(),
tool_registry,
partitioner: ToolPartitioner::new(max_concurrency),
tool_context,
abort_tx,
abort_rx,
has_errored: false,
errored_tool_desc: String::new(),
next_index: 0,
max_concurrency,
max_output_chars,
executing_non_concurrent: false,
completed_queue: VecDeque::new(),
}
}
/// 获取 abort 广播发送端(供外部注入取消信号)
/// 获取 abort 广播发送端(供外部注入取消信号)
pub fn abort_sender(&self) -> broadcast::Sender<AbortReason> {
self.abort_tx.clone()
}
/// 当 LLM 流产生一个新的 tool_use 时调用。
///
/// 返回 true 表示该工具已立即开始执行(并发安全),false 表示排队
/// 如果是并发安全工具且当前没有非并发安全工具在执行,立即 spawn tokio task
/// 否则加入队列等待调度。
///
/// 返回 true 表示该工具已立即开始执行,false 表示排队。
pub fn on_tool_use(&mut self, call_id: String, name: String, args: serde_json::Value) -> bool {
let _index = self.next_index;
self.next_index += 1;
let is_concurrency_safe = self
.tool_registry
.get(&name)
.map(|t| t.is_concurrency_safe(&args))
.unwrap_or(false);
let (cancel_tx, _cancel_rx) = oneshot::channel();
let tool = TrackedTool {
let mut tool = TrackedTool {
tool_call_id: call_id.clone(),
tool_name: name.clone(),
args: args.clone(),
status: TrackedToolStatus::Queued,
output: None,
cancel_tx: Some(cancel_tx),
handle: None,
};
self.tracked.push(tool);
let idx = self.tracked.len();
let can_start_now = is_concurrency_safe && !self.executing_non_concurrent;
if is_concurrency_safe {
info!("[StreamingExecutor] 立即调度并发安全工具: {}", name);
self.try_execute_pending();
true
if can_start_now {
// 立即 spawn tokio task(参考 Claude Code: addTool 立即 processQueue
info!(
"[StreamingExecutor] 立即 spawn 并发安全工具: {} (id={})",
name, call_id
);
let handle = self.spawn_tool_task(idx, call_id.clone(), name.clone(), args.clone());
tool.handle = Some(handle);
tool.status = TrackedToolStatus::Executing;
} else {
info!("[StreamingExecutor] 排队非并发安全工具: {}", name);
false
info!(
"[StreamingExecutor] 排队工具: {} (concurrent={}, executing_non_concurrent={})",
name, is_concurrency_safe, self.executing_non_concurrent
);
}
if !is_concurrency_safe {
self.executing_non_concurrent = true;
}
self.tracked.push(tool);
can_start_now
}
/// LLM 流结束后调用,执行所有剩余排队工具。
/// LLM 流结束后调用,等待所有剩余排队工具完成
pub async fn flush(&mut self) {
let queued_count = self
.tracked
.iter()
.filter(|t| t.status == TrackedToolStatus::Queued)
.count();
info!(
"[StreamingExecutor] flush: {} tracked, {} queued",
"[StreamingExecutor] flush: {} tracked, {} queued, {} executing",
self.tracked.len(),
queued_count,
self.tracked
.iter()
.filter(|t| t.status == TrackedToolStatus::Queued)
.filter(|t| t.status == TrackedToolStatus::Executing)
.count()
);
// 将剩余排队的工具分批执行
let queued: Vec<usize> = self
.tracked
.iter()
.enumerate()
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
.map(|(i, _)| i)
.collect();
// 启动所有还在排队的工具
self.start_all_queued();
for idx in queued {
self.execute_one(idx).await;
}
// 等待所有执行中的工具完成
self.await_all_executing().await;
}
/// 按流顺序获取下一个完成的结果(非阻塞)。
///
/// 对于已完成的任务,如果其 handle 已就绪则收集结果。
/// 返回按到达顺序的第一个已完成结果。
pub fn next_result(&mut self) -> Option<(String, ToolOutput)> {
for tool in &mut self.tracked {
// 先尝试收集任何已完成的 async task 结果
self.collect_completed_tasks();
// 从 completed_queue 中按序取
while let Some(&idx) = self.completed_queue.front() {
self.completed_queue.pop_front();
let tool = &mut self.tracked[idx];
if tool.status == TrackedToolStatus::Completed {
tool.status = TrackedToolStatus::Yielded;
let output = tool
@@ -183,133 +220,280 @@ impl StreamingToolExecutor {
None
}
/// 是否有未 yield 的结果
/// 是否有未 yield 的结果(已完成或即将完成)。
pub fn has_pending_results(&self) -> bool {
self.tracked
.iter()
.any(|t| t.status == TrackedToolStatus::Completed)
|| !self.completed_queue.is_empty()
}
/// 是否有未完成的工具
/// 是否有未完成的工具(仍在排队或执行中)。
pub fn has_unfinished(&self) -> bool {
self.tracked.iter().any(|t| {
t.status == TrackedToolStatus::Queued || t.status == TrackedToolStatus::Executing
})
}
/// 获取所有已完成的结果(包括已 yield 和未 yield 的)
/// 获取所有已完成的结果(包括已 yield 和未 yield 的)
pub fn all_results_mut(&mut self) -> Vec<(String, ToolOutput)> {
self.collect_completed_tasks();
let mut results = Vec::new();
for tool in &mut self.tracked {
if let Some(output) = tool.output.take() {
results.push((tool.tool_call_id.clone(), output));
}
tool.status = TrackedToolStatus::Yielded;
}
results
}
// ── 内部方法 ──
/// 尝试执行可执行的排队工具
fn try_execute_pending(&mut self) {
// 简单策略:如果有正在执行的且它不是并发的,则不启动新的
let has_executing = self
/// Spawn 一个 tokio task 执行单个工具调用。
fn spawn_tool_task(
&self,
_idx: usize,
_call_id: String,
tool_name: String,
args: serde_json::Value,
) -> JoinHandle<ToolExecutionResult> {
let tool_registry = self.tool_registry.clone();
let tool_context = self.tool_context.clone();
let max_output_chars = self.max_output_chars;
let mut abort_rx = self.abort_tx.subscribe();
tokio::spawn(async move {
// tokio::select! 在工具执行和 Sibling Abort 之间竞速
tokio::select! {
result = async {
match tool_registry.get(&tool_name) {
Some(tool) => {
tool.execute_with_progress(args, &tool_context, None).await
}
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
}
} => {
// 截断输出
let truncated = if result.content.len() > max_output_chars {
let t: String = result.content.chars().take(max_output_chars).collect();
ToolOutput {
content: format!(
"{}...\n[输出已截断,原始长度: {} 字符]",
t,
result.content.len()
),
is_error: result.is_error,
metadata: result.metadata,
}
} else {
result
};
ToolExecutionResult {
tool_name,
output: truncated,
}
}
Ok(reason) = abort_rx.recv() => {
let msg = match reason {
AbortReason::SiblingError { description } => {
format!("取消:并行工具 {} 出错,已级联取消", description)
}
AbortReason::UserInterrupted => "执行已被用户取消".to_string(),
};
ToolExecutionResult {
tool_name,
output: ToolOutput::error(msg),
}
}
}
})
}
/// 尝试收集所有已完成 tokio task 的结果(非阻塞)。
fn collect_completed_tasks(&mut self) {
for idx in 0..self.tracked.len() {
if self.tracked[idx].status != TrackedToolStatus::Executing {
continue;
}
if self.tracked[idx].handle.is_none() {
continue;
}
// 检查 JoinHandle 是否已完成(非阻塞)
let handle = self.tracked[idx].handle.take().unwrap();
if handle.is_finished() {
// is_finished=true 保证 .await 会立即返回
// 使用 tokio::task::yield_now 之后的 poll 可能也成功,
// 这里直接在同步上下文中检查后放入完成队列
// 等下次 async 上下文中通过 await_all_executing 处理
self.tracked[idx].handle = Some(handle);
// 标记为需要收集 — 将在 flush/await 中处理
} else {
// 放回未完成的 handle
self.tracked[idx].handle = Some(handle);
}
}
}
/// 启动所有排队的工具。
fn start_all_queued(&mut self) {
// 收集需要启动的工具索引(避免借用冲突)
let to_start: Vec<usize> = self
.tracked
.iter()
.any(|t| t.status == TrackedToolStatus::Executing);
.enumerate()
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
.filter(|(_, t)| {
let is_safe = self
.tool_registry
.get(&t.tool_name)
.map(|reg_tool| reg_tool.is_concurrency_safe(&t.args))
.unwrap_or(false);
// 并发安全工具可随时启动,非并发安全的需要独占
is_safe || !self.executing_non_concurrent
})
.map(|(i, _)| i)
.collect();
if !has_executing {
// 启动所有排队的并发安全工具
let indices: Vec<usize> = self
.tracked
.iter()
.enumerate()
.filter(|(_, t)| t.status == TrackedToolStatus::Queued)
.map(|(i, _)| i)
.collect();
for idx in to_start {
let tool = &self.tracked[idx];
let call_id = tool.tool_call_id.clone();
let tool_name = tool.tool_name.clone();
let args = tool.args.clone();
for idx in indices {
// 在同步上下文中只能标记状态,实际执行在 async 上下文中
self.tracked[idx].status = TrackedToolStatus::Executing;
}
}
}
/// 执行单个工具(内部辅助)
async fn execute_one(&mut self, idx: usize) {
if idx >= self.tracked.len() {
return;
}
// 检查 sibling abort
if self.has_errored {
if let Ok(reason) = self.abort_rx.try_recv() {
let msg = match reason {
AbortReason::SiblingError { ref description } => {
format!("取消:并行工具 {} 出错,已级联取消", description)
}
AbortReason::UserInterrupted => "执行已被用户取消".to_string(),
};
self.tracked[idx].output = Some(ToolOutput::error(msg));
self.tracked[idx].status = TrackedToolStatus::Completed;
return;
}
}
self.tracked[idx].status = TrackedToolStatus::Executing;
let tool_name = self.tracked[idx].tool_name.clone();
let output = match self.tool_registry.get(&tool_name) {
Some(tool) => {
let (progress_tx, _progress_rx) = mpsc::unbounded_channel();
let tool_args = self.tracked[idx].args.clone();
let tool_fut =
tool.execute_with_progress(tool_args, &self.tool_context, Some(&progress_tx));
tool_fut.await
}
None => ToolOutput::error(format!("未知工具: {}", tool_name)),
};
// 检查是否需要触发 sibling abort
if output.is_error {
let causes_abort = self
let is_safe = self
.tool_registry
.get(&tool_name)
.map(|t| t.causes_sibling_abort())
.map(|t| t.is_concurrency_safe(&args))
.unwrap_or(false);
if causes_abort {
warn!(
"[StreamingExecutor] 工具 {} 出错,触发 sibling abort",
tool_name
);
self.has_errored = true;
self.errored_tool_desc = tool_name.clone();
let _ = self.abort_tx.send(AbortReason::SiblingError {
description: tool_name.clone(),
});
let handle = self.spawn_tool_task(idx, call_id, tool_name.clone(), args);
self.tracked[idx].handle = Some(handle);
self.tracked[idx].status = TrackedToolStatus::Executing;
if !is_safe {
self.executing_non_concurrent = true;
// 非并发安全工具启动后停止(独占执行)
break;
}
info!("[StreamingExecutor] 启动排队工具: {}", tool_name);
}
}
/// 等待所有执行中的工具完成。
async fn await_all_executing(&mut self) {
// 收集所有剩余 JoinHandles
let mut handles: Vec<(usize, JoinHandle<ToolExecutionResult>)> = Vec::new();
for idx in 0..self.tracked.len() {
if self.tracked[idx].status == TrackedToolStatus::Executing {
if let Some(handle) = self.tracked[idx].handle.take() {
handles.push((idx, handle));
}
}
}
// 截断输出
let truncated = if output.content.len() > self.max_output_chars {
let t: String = output.content.chars().take(self.max_output_chars).collect();
ToolOutput {
content: format!(
"{}...\n[输出已截断,原始长度: {} 字符]",
t,
output.content.len()
),
is_error: output.is_error,
metadata: output.metadata,
}
} else {
output
};
// 并发等待所有任务
for (idx, handle) in handles {
match handle.await {
Ok(result) => {
let tool_name = result.tool_name.clone();
let is_error = result.output.is_error;
self.tracked[idx].output = Some(truncated);
self.tracked[idx].status = TrackedToolStatus::Completed;
self.tracked[idx].output = Some(result.output);
self.tracked[idx].status = TrackedToolStatus::Completed;
self.completed_queue.push_back(idx);
if is_error {
self.check_sibling_abort(idx, &tool_name);
}
}
Err(e) => {
warn!("[StreamingExecutor] tokio task 异常: {}", e);
self.tracked[idx].output =
Some(ToolOutput::error(format!("工具执行异常: {}", e)));
self.tracked[idx].status = TrackedToolStatus::Completed;
self.completed_queue.push_back(idx);
}
}
}
self.executing_non_concurrent = false;
}
/// 检查错误工具是否触发 Sibling Abort。
fn check_sibling_abort(&mut self, idx: usize, tool_name: &str) {
let causes_abort = self
.tool_registry
.get(tool_name)
.map(|t| t.causes_sibling_abort())
.unwrap_or(false);
if causes_abort && !self.has_errored {
warn!(
"[StreamingExecutor] 工具 {} 出错,触发 sibling abort",
tool_name
);
self.has_errored = true;
self.errored_tool_desc = self.get_tool_description(idx);
let _ = self.abort_tx.send(AbortReason::SiblingError {
description: self.errored_tool_desc.clone(),
});
}
}
/// 获取工具的人类可读描述(用于错误消息)。
fn get_tool_description(&self, idx: usize) -> String {
let tool = &self.tracked[idx];
let summary = tool
.args
.get("command")
.or_else(|| tool.args.get("file_path"))
.or_else(|| tool.args.get("pattern"))
.and_then(|v| v.as_str())
.unwrap_or("");
if summary.is_empty() {
tool.tool_name.clone()
} else {
let truncated: String = summary.chars().take(40).collect();
if summary.len() > 40 {
format!("{}({}…)", tool.tool_name, truncated)
} else {
format!("{}({})", tool.tool_name, summary)
}
}
}
}
#[cfg(test)]
mod tests {
// 注意:StreamingToolExecutor 的集成测试放在 src/agent/runtime/ 的 #[cfg(test)] 模块中,
// 需要完整的 AppState 和 ToolContext。此处的单元测试仅验证核心数据结构。
//
// 以下测试验证 TrackedToolStatus 枚举和状态转换逻辑,不依赖外部设施。
use super::*;
#[test]
fn test_tracked_tool_status_debug() {
assert_eq!(format!("{:?}", TrackedToolStatus::Queued), "Queued");
assert_eq!(format!("{:?}", TrackedToolStatus::Executing), "Executing");
assert_eq!(format!("{:?}", TrackedToolStatus::Completed), "Completed");
assert_eq!(format!("{:?}", TrackedToolStatus::Yielded), "Yielded");
}
#[test]
fn test_abort_reason_display() {
let sibling = AbortReason::SiblingError {
description: "bash(rm -rf /)".into(),
};
let user = AbortReason::UserInterrupted;
assert_eq!(
format!("{:?}", sibling),
"SiblingError { description: \"bash(rm -rf /)\" }"
);
assert_eq!(format!("{:?}", user), "UserInterrupted");
}
}
+225 -6
View File
@@ -1,9 +1,11 @@
// src/agent/runtime/system_prompt.rs
//
// 模块化系统提示词组装 — 参考 Claude Code s10 System Prompt 设计
// 模块化系统提示词组装。
//
// 将硬编码的提示词拆分为独立 section,运行时按需拼接。
// 静态 section 在前以最大化 Anthropic prompt cache 命中率。
// 设计原则:
// 1. 静态 section 全部在前 → 内容不变,服务端自然缓存命中
// 2. 动态 section 在后 → 随 session 变化
// 3. 简单缓存:首次计算后永久复用(session 内一切不变),仅显式 invalidate
/// 系统提示词组装器
pub struct SystemPrompt {
@@ -48,13 +50,86 @@ impl Default for SystemPrompt {
}
}
/// 静态身份 section(始终加载,最大化 prompt cache 命中率)
// ── Section 缓存 ─────────────────────────────────────────────────────────
use std::collections::HashMap;
/// SystemPrompt 的 Section 级缓存。
///
/// 首次计算后永久缓存,仅通过 `invalidate()` 显式失效。
/// 因为 session 生命周期内 CWD/平台/OS/模型名/工具注册表均不变,
/// 不需要 TTL 过期机制。
#[derive(Debug)]
pub struct SystemPromptCache {
entries: HashMap<&'static str, String>,
}
impl SystemPromptCache {
pub fn new() -> Self {
SystemPromptCache {
entries: HashMap::new(),
}
}
/// 获取 section 内容(首次计算,后续命中缓存)。
pub fn get_or_compute(
&mut self,
name: &'static str,
compute: impl FnOnce() -> String,
) -> String {
if let Some(cached) = self.entries.get(name) {
return cached.clone();
}
let content = compute();
self.entries.insert(name, content.clone());
content
}
/// 获取 section 内容,compute 返回 Option 时:Some 缓存并返回,None 不缓存。
pub fn get_or_compute_optional(
&mut self,
name: &'static str,
compute: impl FnOnce() -> Option<String>,
) -> Option<String> {
if let Some(cached) = self.entries.get(name) {
return Some(cached.clone());
}
let content = compute()?;
self.entries.insert(name, content.clone());
Some(content)
}
/// 使指定 section 的缓存失效。
pub fn invalidate(&mut self, name: &'static str) {
self.entries.remove(name);
}
/// 使所有缓存失效(/clear 或 /compact 时调用)。
pub fn invalidate_all(&mut self) {
self.entries.clear();
}
/// 缓存条目数量(调试用)。
pub fn entry_count(&self) -> usize {
self.entries.len()
}
}
impl Default for SystemPromptCache {
fn default() -> Self {
Self::new()
}
}
// ── 静态 Section 常量 ────────────────────────────────────────────────────
/// 静态身份 section
pub const IDENTITY_SECTION: &str = "\
你是一位专业的天体物理学研究助手,具备丰富的天文学知识。";
/// 静态核心原则 section
/// 静态核心原则 section(行为准则 + 科研规范)
pub const PRINCIPLES_SECTION: &str = "\
核心原则
# 核心原则
1. 主动使用工具搜索最新文献,不要仅凭训练数据回答。
2. 优先使用本地资源(get_paper_content / rag_search),必要时再检索新文献。
3. 收集到足够信息后立即给出最终答案,避免无意义的重复工具调用。
@@ -65,6 +140,32 @@ pub const PRINCIPLES_SECTION: &str = "\
8. 如果某个工具调用失败,不要用相同参数重试,尝试换一种方式或工具。
9. 任务状态会在每轮开始时从数据库恢复,请基于最新状态继续工作。";
/// 静态系统上下文 section(说明 system-reminder 标签和自动压缩机制)
pub const SYSTEM_CONTEXT_SECTION: &str = "\
# 系统上下文
- 工具结果和用户消息中可能包含 <system-reminder> 标签。这些标签由系统自动添加,包含有用的信息和提醒,与所在消息的具体内容无直接关系。
- 对话具有通过自动摘要实现的无限上下文长度。当上下文接近限制时,较早的消息会被自动压缩为摘要。";
/// 静态工具使用指导 sectiondedicated tools 优先、并行调用、任务追踪)
pub const TOOL_USAGE_SECTION: &str = "\
# 工具使用指南
- 优先使用专用工具(read_file、grep_files、glob_files、file_edit、file_write),仅在无专用工具时才使用 run_bash。
- 使用 run_bash 执行系统命令时,优先选择可逆、影响范围小的操作。
- 你可以在一次回复中调用多个工具。如果多个工具调用之间没有依赖关系,请并行调用以提升效率。
- 使用 todo_write 工具规划和管理工作。每完成一个任务立即更新状态。不要批量标记多个任务为完成。
- 不要创建不必要的文件。优先编辑已有文件而非新建。";
/// 静态操作安全 section(可逆性、影响范围、确认机制)
pub const SAFETY_SECTION: &str = "\
# 操作安全
- 仔细考虑操作的可逆性和影响范围。本地、可逆的操作(编辑文件、运行测试)可自由执行。
- 对于难以撤销或影响共享状态的操作(删除文件/分支、修改数据库、对外发送内容),在执行前与用户确认。
- 用户批准某类操作一次不代表在所有上下文中都批准,除非有持久化的授权指令。
- 遇到障碍时不要用破坏性操作作为捷径(如 --no-verify 跳过检查)。找到根本原因并修复。
- 如果发现意外状态(陌生文件、分支、配置),先调查再删除,这可能代表用户正在进行的工作。";
// ── Tests ─────────────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
@@ -101,4 +202,122 @@ mod tests {
sp.add_section("a", "A".to_string());
assert_eq!(sp.section_count(), 1);
}
#[test]
fn test_static_sections_not_empty() {
assert!(!SYSTEM_CONTEXT_SECTION.is_empty());
assert!(!TOOL_USAGE_SECTION.is_empty());
assert!(!SAFETY_SECTION.is_empty());
assert!(!PRINCIPLES_SECTION.is_empty());
assert!(!IDENTITY_SECTION.is_empty());
}
#[test]
fn test_static_sections_have_headers() {
assert!(SYSTEM_CONTEXT_SECTION.starts_with("# 系统上下文"));
assert!(TOOL_USAGE_SECTION.starts_with("# 工具使用指南"));
assert!(SAFETY_SECTION.starts_with("# 操作安全"));
assert!(PRINCIPLES_SECTION.starts_with("# 核心原则"));
}
// ── SystemPromptCache tests ──
#[test]
fn test_cache_compute_once() {
let mut cache = SystemPromptCache::new();
let mut call_count = 0;
let r1 = cache.get_or_compute("test", || {
call_count += 1;
"computed".to_string()
});
assert_eq!(r1, "computed");
assert_eq!(call_count, 1);
// 第二次不调用 compute
let r2 = cache.get_or_compute("test", || {
call_count += 1;
"recomputed".to_string()
});
assert_eq!(r2, "computed");
assert_eq!(call_count, 1);
}
#[test]
fn test_cache_invalidate() {
let mut cache = SystemPromptCache::new();
let _ = cache.get_or_compute("a", || "value_a".to_string());
let _ = cache.get_or_compute("b", || "value_b".to_string());
assert_eq!(cache.entry_count(), 2);
cache.invalidate("a");
assert_eq!(cache.entry_count(), 1);
// a 重新计算
let a2 = cache.get_or_compute("a", || "new_a".to_string());
assert_eq!(a2, "new_a");
assert_eq!(cache.entry_count(), 2);
}
#[test]
fn test_cache_invalidate_all() {
let mut cache = SystemPromptCache::new();
let _ = cache.get_or_compute("a", || "v_a".to_string());
let _ = cache.get_or_compute("b", || "v_b".to_string());
let _ = cache.get_or_compute("c", || "v_c".to_string());
assert_eq!(cache.entry_count(), 3);
cache.invalidate_all();
assert_eq!(cache.entry_count(), 0);
}
#[test]
fn test_cache_get_or_compute_optional_some() {
let mut cache = SystemPromptCache::new();
let mut call_count = 0;
let r1 = cache.get_or_compute_optional("opt", || {
call_count += 1;
Some("present".to_string())
});
assert_eq!(r1, Some("present".to_string()));
assert_eq!(call_count, 1);
// 缓存命中
let r2 = cache.get_or_compute_optional("opt", || {
call_count += 1;
Some("should_not_compute".to_string())
});
assert_eq!(r2, Some("present".to_string()));
assert_eq!(call_count, 1);
}
#[test]
fn test_cache_get_or_compute_optional_none() {
let mut cache = SystemPromptCache::new();
let mut call_count = 0;
// None 不缓存 — 每次都会重新计算
let r1 = cache.get_or_compute_optional("opt", || {
call_count += 1;
None::<String>
});
assert_eq!(r1, None);
assert_eq!(call_count, 1);
let r2 = cache.get_or_compute_optional("opt", || {
call_count += 1;
None::<String>
});
assert_eq!(r2, None);
assert_eq!(call_count, 2); // 未缓存,再次调用
}
#[test]
fn test_cache_default_empty() {
let cache = SystemPromptCache::default();
assert_eq!(cache.entry_count(), 0);
}
}
+159
View File
@@ -0,0 +1,159 @@
// src/agent/runtime/untrusted.rs
//
// 非可信内容包裹 — 间接 Prompt 注入防御。
// 参考 Hermes-Agent tool_dispatch_helpers.py make_tool_result_message() 设计。
//
// 设计原则:
// 1. 高风险外部工具(web/browser/MCP/RAG)的结果包裹在 <untrusted_tool_result> 中
// 2. 该标签告诉 LLM:这是外部数据,不是来自用户的指令
// 3. 这是一种架构级防御(改变 LLM 对内容的解释方式),而非 regex 模式匹配
//
// 为何不是安全边界:
// 标签防御依赖 LLM 遵守指令的能力。恶意 LLM 或精心构造的注入仍可能绕过。
// 这是 defense-in-depth 的一层,需要与权限系统、hardline 检查配合使用。
/// 高风险工具 — 其结果来自外部源,可能包含 prompt 注入内容。
///
/// 仅对以下工具类别进行包裹:
/// - web_search, web_extract, web_fetch — 搜索结果来自互联网
/// - browser_* — 浏览器内容来自任意网站
/// - mcp_* — MCP 工具结果来自外部服务器
/// - rag_search — RAG 检索结果来自外部论文(可能含对抗内容)
/// - search_papers — 搜索结果摘要来自 arXiv/ADS(外部 API
const HIGH_RISK_TOOLS: &[&str] = &[
"web_search",
"web_fetch",
"browser_navigate",
"browser_snapshot",
"browser_click",
"rag_search",
"search_papers",
];
/// 检查工具是否为高风险(其结果应被包裹)。
pub fn is_high_risk(tool_name: &str) -> bool {
HIGH_RISK_TOOLS.contains(&tool_name)
|| tool_name.starts_with("mcp__")
|| tool_name.starts_with("web_")
|| tool_name.starts_with("browser_")
}
/// 为非可信工具结果包裹安全标记。
///
/// 包裹格式:
/// ```text
/// <untrusted_tool_result tool="{tool_name}">
/// {original_content}
/// </untrusted_tool_result>
/// ```
///
/// 此标记指示 LLM 将内容视为外部数据而非用户指令。
/// 包裹仅应用于高风险工具(web/browser/MCP),以最小化 token 开销。
pub fn wrap_untrusted_content(tool_name: &str, content: &str) -> String {
if !is_high_risk(tool_name) {
return content.to_string();
}
// 避免双重包裹
if content.contains("<untrusted_tool_result") {
return content.to_string();
}
format!(
"<untrusted_tool_result tool=\"{tool}\">\n{content}\n</untrusted_tool_result>",
tool = tool_name,
content = content
)
}
/// 从错误包裹中恢复原内容(当需要向用户展示时)
pub fn unwrap_untrusted(content: &str) -> String {
if !content.starts_with("<untrusted_tool_result") {
return content.to_string();
}
// 简单提取:找到第一行之后和最后一行之前的内容
if let Some(start) = content.find('\n') {
let inner = &content[start + 1..];
if let Some(end) = inner.rfind('\n') {
if inner[end..].contains("</untrusted_tool_result>") {
return inner[..end].to_string();
}
}
}
content.to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_high_risk_web_search() {
assert!(is_high_risk("web_search"));
assert!(is_high_risk("web_fetch"));
}
#[test]
fn test_high_risk_browser() {
assert!(is_high_risk("browser_navigate"));
assert!(is_high_risk("browser_snapshot"));
}
#[test]
fn test_high_risk_mcp() {
assert!(is_high_risk("mcp__github_search"));
}
#[test]
fn test_high_risk_rag() {
assert!(is_high_risk("rag_search"));
assert!(is_high_risk("search_papers"));
}
#[test]
fn test_not_high_risk_file_tools() {
assert!(!is_high_risk("read_file"));
assert!(!is_high_risk("file_write"));
assert!(!is_high_risk("run_bash"));
assert!(!is_high_risk("grep_files"));
}
#[test]
fn test_wrap_web_search_result() {
let content = "Found: prompt injection <script>alert('xss')</script>";
let wrapped = wrap_untrusted_content("web_search", content);
assert!(wrapped.starts_with("<untrusted_tool_result tool=\"web_search\">"));
assert!(wrapped.ends_with("</untrusted_tool_result>"));
assert!(wrapped.contains(content));
}
#[test]
fn test_no_wrap_for_safe_tool() {
let content = "file contents here";
let wrapped = wrap_untrusted_content("read_file", content);
assert_eq!(wrapped, content);
}
#[test]
fn test_no_double_wrap() {
let content = "<untrusted_tool_result tool=\"web_search\">\nalready wrapped\n</untrusted_tool_result>";
let wrapped = wrap_untrusted_content("web_search", content);
assert_eq!(wrapped, content);
}
#[test]
fn test_unwrap_restores_original() {
let original = "Found: some search result with injection";
let wrapped = wrap_untrusted_content("web_search", original);
let unwrapped = unwrap_untrusted(&wrapped);
assert_eq!(unwrapped, original);
}
#[test]
fn test_unwrap_no_marker_returns_as_is() {
let content = "plain text without wrapper";
let unwrapped = unwrap_untrusted(content);
assert_eq!(unwrapped, content);
}
}