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:
@@ -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 写回 index(write_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);
|
||||
}
|
||||
}
|
||||
+1204
-133
File diff suppressed because it is too large
Load Diff
+463
-130
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(|| {
|
||||
// 全角 ASCII(U+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() {
|
||||
// 全角 'r' (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
@@ -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. 动态 section(environment/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,
|
||||
|
||||
@@ -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, ®istry);
|
||||
|
||||
// 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, ®istry);
|
||||
|
||||
// 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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -1,8 +1,10 @@
|
||||
// src/agent/runtime/session.rs
|
||||
//
|
||||
// 会话生命周期管理:创建/恢复/验证 Agent 会话。
|
||||
// 支持软删除回退 (undo):active=0 标记保留审计 trail,LLM 不可见。
|
||||
|
||||
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=1(LLM 上下文),
|
||||
/// 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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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> 标签。这些标签由系统自动添加,包含有用的信息和提醒,与所在消息的具体内容无直接关系。
|
||||
- 对话具有通过自动摘要实现的无限上下文长度。当上下文接近限制时,较早的消息会被自动压缩为摘要。";
|
||||
|
||||
/// 静态工具使用指导 section(dedicated 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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user