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:
Generated
+75
@@ -141,6 +141,7 @@ dependencies = [
|
||||
"dotenvy",
|
||||
"flate2",
|
||||
"futures-util",
|
||||
"git2",
|
||||
"glob",
|
||||
"hmac 0.12.1",
|
||||
"html2md",
|
||||
@@ -159,6 +160,7 @@ dependencies = [
|
||||
"sha1 0.10.6",
|
||||
"sqlite-vec",
|
||||
"sqlx",
|
||||
"tempfile",
|
||||
"thiserror 1.0.69",
|
||||
"tokio",
|
||||
"tower-http 0.5.2",
|
||||
@@ -1456,6 +1458,21 @@ dependencies = [
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "git2"
|
||||
version = "0.18.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "232e6a7bfe35766bf715e55a88b39a700596c0ccfd88cd3680b4cdb40d66ef70"
|
||||
dependencies = [
|
||||
"bitflags 2.13.0",
|
||||
"libc",
|
||||
"libgit2-sys",
|
||||
"log",
|
||||
"openssl-probe",
|
||||
"openssl-sys",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "glob"
|
||||
version = "0.3.3"
|
||||
@@ -2066,6 +2083,20 @@ version = "0.2.186"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
|
||||
|
||||
[[package]]
|
||||
name = "libgit2-sys"
|
||||
version = "0.16.2+1.7.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ee4126d8b4ee5c9d9ea891dd875cfdc1e9d0950437179104b183d7d8a74d24e8"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"libc",
|
||||
"libssh2-sys",
|
||||
"libz-sys",
|
||||
"openssl-sys",
|
||||
"pkg-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libloading"
|
||||
version = "0.8.9"
|
||||
@@ -2105,6 +2136,32 @@ dependencies = [
|
||||
"vcpkg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libssh2-sys"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "220e4f05ad4a218192533b300327f5150e809b54c4ec83b5a1d91833601811b9"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"libc",
|
||||
"libz-sys",
|
||||
"openssl-sys",
|
||||
"pkg-config",
|
||||
"vcpkg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libz-sys"
|
||||
version = "1.1.29"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "85bc9657773828b90eeb625adff10eeac83cc21bbfd8e23a03eaa8a33c9e28d9"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"libc",
|
||||
"pkg-config",
|
||||
"vcpkg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "linux-raw-sys"
|
||||
version = "0.4.15"
|
||||
@@ -2536,6 +2593,24 @@ dependencies = [
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl-probe"
|
||||
version = "0.1.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e"
|
||||
|
||||
[[package]]
|
||||
name = "openssl-sys"
|
||||
version = "0.9.117"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b47e7e6bb2c38cd930d25a23b40fa52e068c10e85f3e03a7f5ba5aaca5713695"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"libc",
|
||||
"pkg-config",
|
||||
"vcpkg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "outref"
|
||||
version = "0.5.2"
|
||||
|
||||
@@ -55,9 +55,11 @@ async-trait = "0.1"
|
||||
async-stream = "0.3"
|
||||
serde_yaml = "0.9"
|
||||
notify = { version = "6", default-features = false, features = ["macos_kqueue"] }
|
||||
tempfile = "3"
|
||||
glob = "0.3"
|
||||
walkdir = "2"
|
||||
lru = "0.12"
|
||||
git2 = "0.18"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
|
||||
+57
-57
@@ -337,7 +337,7 @@ export default function App() {
|
||||
}
|
||||
};
|
||||
|
||||
// 5.5. 对文献进行向量化分块入库 (独立任务)
|
||||
// 5.5. 对文献进行知识入库 (独立任务)
|
||||
const handleVectorize = async (bibcode: string) => {
|
||||
setVectorizing(true);
|
||||
try {
|
||||
@@ -347,10 +347,10 @@ export default function App() {
|
||||
if (selectedPaper?.bibcode === bibcode) {
|
||||
setSelectedPaper(prev => prev ? { ...prev, has_vector: true } : null);
|
||||
}
|
||||
showAlert(`文献向量化分块入库成功，共切片并录入 ${res.data.chunk_count} 个向量块。`, '向量化成功');
|
||||
showAlert(`文献已成功录入馆藏智能库，共提炼并录入 ${res.data.chunk_count} 个核心知识节点。`, '知识入库成功');
|
||||
} catch (e) {
|
||||
console.error('文献向量化失败', e);
|
||||
showAlert('向量化失败，请检查 .env 中的 Embedding API 配置。', '向量化失败');
|
||||
console.error('知识入库失败', e);
|
||||
showAlert('知识入库失败，请检查 .env 中的 Embedding API 配置。', '知识入库失败');
|
||||
} finally {
|
||||
setVectorizing(false);
|
||||
}
|
||||
@@ -698,11 +698,11 @@ export default function App() {
|
||||
>
|
||||
<div
|
||||
onClick={(e) => e.stopPropagation()}
|
||||
className="bg-white rounded-xl border border-slate-200 shadow-xl max-w-sm w-full p-6 space-y-4"
|
||||
className="bg-[var(--bg-card)] rounded-xl border border-[var(--border-precision)] shadow-md max-w-sm w-full p-6 space-y-4"
|
||||
>
|
||||
<div className="flex items-center justify-between border-b border-slate-100 pb-2.5">
|
||||
<h3 className="text-xs font-bold text-slate-900 flex items-center gap-1.5">
|
||||
<span className={`w-2 h-2 rounded-full ${dialog.type === 'confirm' ? 'bg-sky-500' : 'bg-red-500'} animate-pulse`} />
|
||||
<h3 className="text-xs font-bold text-[var(--text-main)] flex items-center gap-1.5">
|
||||
<span className={`w-2 h-2 rounded-full ${dialog.type === 'confirm' ? 'bg-[var(--accent-blueprint)]' : 'bg-[#ef4444]'} animate-pulse`} />
|
||||
<span>{dialog.title}</span>
|
||||
</h3>
|
||||
<button
|
||||
@@ -719,7 +719,7 @@ export default function App() {
|
||||
</button>
|
||||
</div>
|
||||
<div className="space-y-2">
|
||||
<p className="text-xs text-slate-705 leading-relaxed">{dialog.message}</p>
|
||||
<p className="text-xs text-[var(--text-muted)] leading-relaxed">{dialog.message}</p>
|
||||
</div>
|
||||
<div className="flex gap-2 pt-2">
|
||||
<button
|
||||
@@ -737,7 +737,7 @@ export default function App() {
|
||||
if (dialog.onCancel) dialog.onCancel();
|
||||
setDialog(null);
|
||||
}}
|
||||
className="px-4 bg-slate-100 hover:bg-slate-200 text-slate-600 rounded-lg text-[11px] font-bold text-center transition-all cursor-pointer"
|
||||
className="px-4 btn-console btn-console-secondary py-2 rounded-lg text-[11px] font-bold text-center cursor-pointer"
|
||||
>
|
||||
取消
|
||||
</button>
|
||||
@@ -755,11 +755,11 @@ export default function App() {
|
||||
>
|
||||
<div
|
||||
onClick={(e) => e.stopPropagation()}
|
||||
className="bg-white rounded-xl border border-slate-200 shadow-xl max-w-sm w-full p-6 space-y-4"
|
||||
className="bg-[var(--bg-card)] rounded-xl border border-[var(--border-precision)] shadow-md max-w-sm w-full p-6 space-y-4"
|
||||
>
|
||||
<div className="flex items-center justify-between border-b border-slate-100 pb-2.5">
|
||||
<h3 className="text-xs font-bold text-slate-900 flex items-center gap-1.5">
|
||||
<span className="w-2 h-2 rounded-full bg-amber-500 animate-pulse" />
|
||||
<h3 className="text-xs font-bold text-[var(--text-main)] flex items-center gap-1.5">
|
||||
<span className="w-2 h-2 rounded-full bg-[var(--accent-star)] animate-pulse" />
|
||||
<span>文献尚未入库</span>
|
||||
</h3>
|
||||
<button
|
||||
@@ -773,10 +773,10 @@ export default function App() {
|
||||
</button>
|
||||
</div>
|
||||
<div className="space-y-2">
|
||||
<p className="text-xs text-slate-700 leading-relaxed">
|
||||
文献 <span className="font-mono bg-slate-100 px-1.5 py-0.5 rounded text-sky-700 font-bold select-all">{uncachedBibcode}</span> 尚未收录在本地数据库中。
|
||||
<p className="text-xs text-[var(--text-main)] leading-relaxed">
|
||||
文献 <span className="font-mono bg-slate-100 px-1.5 py-0.5 rounded text-[var(--accent-blueprint)] font-bold select-all">{uncachedBibcode}</span> 尚未收录在本地数据库中。
|
||||
</p>
|
||||
<p className="text-[11px] text-slate-400 leading-relaxed">
|
||||
<p className="text-[11px] text-[var(--text-muted)] leading-relaxed">
|
||||
您可以选择在线拉取该文献元数据并入库，或是直接跳转至 NASA ADS 平台查看其原始页面。
|
||||
</p>
|
||||
</div>
|
||||
@@ -795,13 +795,13 @@ export default function App() {
|
||||
axios.post('/api/active_bibcode', { bibcode: uncachedBibcode }).catch(() => {});
|
||||
window.open(`https://ui.adsabs.harvard.edu/abs/${uncachedBibcode}/abstract`, '_blank');
|
||||
}}
|
||||
className="flex-1 bg-white hover:bg-slate-50 text-slate-700 border border-slate-250 py-2 rounded-lg text-[11px] font-bold text-center transition-all shadow-sm cursor-pointer"
|
||||
className="flex-1 btn-console btn-console-secondary py-2 rounded-lg text-[11px] font-bold text-center cursor-pointer"
|
||||
>
|
||||
跳转到 ADS
|
||||
</button>
|
||||
<button
|
||||
onClick={() => setUncachedBibcode(null)}
|
||||
className="px-3 bg-slate-100 hover:bg-slate-200 text-slate-600 rounded-lg text-xs font-medium text-center transition-all cursor-pointer"
|
||||
className="px-3 btn-console btn-console-secondary py-2 rounded-lg text-[11px] font-bold text-center cursor-pointer"
|
||||
>
|
||||
取消
|
||||
</button>
|
||||
@@ -818,13 +818,13 @@ export default function App() {
|
||||
>
|
||||
<div
|
||||
onClick={(e) => e.stopPropagation()}
|
||||
className="bg-white rounded-xl border border-slate-200 shadow-xl max-w-lg w-full p-6 space-y-4 cursor-default animate-fade-in"
|
||||
className="bg-[var(--bg-card)] rounded-xl border border-[var(--border-precision)] shadow-md max-w-lg w-full p-6 space-y-4 cursor-default animate-fade-in"
|
||||
>
|
||||
{/* 标题 & 关闭 */}
|
||||
<div className="flex items-start justify-between border-b border-slate-100 pb-3">
|
||||
<div className="space-y-1 pr-4">
|
||||
<span className="text-[10px] font-bold text-sky-700 uppercase tracking-wider">文献详情元数据</span>
|
||||
<h3 className="text-xs font-bold text-slate-900 leading-snug">
|
||||
<span className="text-[10px] font-bold text-[var(--accent-blueprint)] uppercase tracking-wider">文献详情元数据</span>
|
||||
<h3 className="text-xs font-bold text-[var(--text-main)] leading-snug">
|
||||
{getDoctypeBadge(detailPaper.doctype)}
|
||||
<span className="align-middle">{detailPaper.title}</span>
|
||||
</h3>
|
||||
@@ -844,48 +844,48 @@ export default function App() {
|
||||
<div className="space-y-4 h-[460px] flex flex-col overflow-y-auto pr-1 text-xs scrollbar-thin">
|
||||
{/* 作者 */}
|
||||
<div className="space-y-1 h-12 overflow-y-auto scrollbar-thin shrink-0">
|
||||
<span className="text-slate-450 font-bold">作者列表</span>
|
||||
<p className="text-slate-800 leading-relaxed font-semibold">{detailPaper.authors.join(', ')}</p>
|
||||
<span className="text-[var(--text-muted)] font-bold">作者列表</span>
|
||||
<p className="text-[var(--text-main)] leading-relaxed font-semibold">{detailPaper.authors.join(', ')}</p>
|
||||
</div>
|
||||
|
||||
{/* 期刊 & 年份 */}
|
||||
<div className="grid grid-cols-5 gap-4 border-y border-slate-100 py-2.5 shrink-0 h-14">
|
||||
<div className="space-y-0.5 flex flex-col justify-center min-w-0 col-span-4">
|
||||
<span className="text-slate-450 font-bold block text-[10px] leading-tight">发表期刊</span>
|
||||
<span className="text-[var(--text-muted)] font-bold block text-[10px] leading-tight">发表期刊</span>
|
||||
<span
|
||||
className="text-slate-800 font-bold italic truncate block text-[11px] leading-tight"
|
||||
className="text-[var(--text-main)] font-bold italic truncate block text-[11px] leading-tight"
|
||||
title={detailPaper.pub_journal || '未标注'}
|
||||
>
|
||||
{detailPaper.pub_journal || '未标注'}
|
||||
</span>
|
||||
</div>
|
||||
<div className="space-y-0.5 flex flex-col justify-center col-span-1">
|
||||
<span className="text-slate-450 font-bold block text-[10px] leading-tight">发表年份</span>
|
||||
<span className="text-slate-850 font-extrabold block text-[11px] leading-tight">{detailPaper.year}</span>
|
||||
<span className="text-[var(--text-muted)] font-bold block text-[10px] leading-tight">发表年份</span>
|
||||
<span className="text-[var(--text-main)] font-extrabold block text-[11px] leading-tight">{detailPaper.year}</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* 摘要 */}
|
||||
<div className="space-y-1.5 flex flex-col h-40 shrink-0">
|
||||
<span className="text-slate-450 font-bold block">摘要</span>
|
||||
<p className="text-slate-700 leading-relaxed font-normal bg-slate-50 p-3 flex-1 rounded-lg border border-slate-200 text-justify overflow-y-auto scrollbar-thin select-text">
|
||||
<span className="text-[var(--text-muted)] font-bold block">摘要</span>
|
||||
<p className="text-[var(--text-main)] leading-relaxed font-normal bg-[#f8fafc] p-3 flex-1 rounded-lg border border-[var(--border-precision)] text-justify overflow-y-auto scrollbar-thin select-text">
|
||||
{detailPaper.abstract_text || '该文献暂无摘要数据。'}
|
||||
</p>
|
||||
</div>
|
||||
|
||||
{/* 关键字 */}
|
||||
<div className="space-y-1.5 flex flex-col h-16 shrink-0">
|
||||
<span className="text-slate-450 font-bold block">关键词</span>
|
||||
<span className="text-[var(--text-muted)] font-bold block">关键词</span>
|
||||
{detailPaper.keywords && detailPaper.keywords.length > 0 ? (
|
||||
<div className="flex flex-wrap gap-1.5 flex-1 content-start items-start overflow-y-auto scrollbar-thin pr-1">
|
||||
{detailPaper.keywords.map(kw => (
|
||||
<span key={kw} className="px-2 py-0.5 rounded bg-slate-100 border border-slate-200 text-slate-600 font-bold text-[9px]">
|
||||
<span key={kw} className="px-2 py-0.5 rounded bg-[#f1f5f9] border border-[var(--border-precision)] text-[var(--text-muted)] font-bold text-[9px]">
|
||||
{kw}
|
||||
</span>
|
||||
))}
|
||||
</div>
|
||||
) : (
|
||||
<div className="flex items-center justify-center flex-1 bg-slate-50 border border-slate-200 rounded-lg text-[10px] text-slate-400 font-bold select-none">
|
||||
<div className="flex items-center justify-center flex-1 bg-[#f8fafc] border border-[var(--border-precision)] rounded-lg text-[10px] text-[var(--text-muted)] font-bold select-none">
|
||||
暂无关键词
|
||||
</div>
|
||||
)}
|
||||
@@ -893,48 +893,48 @@ export default function App() {
|
||||
|
||||
{/* 标识符 */}
|
||||
<div className="border-t border-slate-100 pt-3 grid grid-cols-1 sm:grid-cols-3 gap-2 text-[10px] font-mono mt-auto shrink-0">
|
||||
<div className="bg-slate-50 px-2.5 py-1.5 rounded border border-slate-150">
|
||||
<span className="text-slate-400 font-bold block">BIBCODE</span>
|
||||
<span className="text-slate-700 font-semibold select-all truncate block" title={detailPaper.bibcode === detailPaper.arxiv_id ? '暂无' : detailPaper.bibcode}>
|
||||
<div className="bg-[#f8fafc] px-2.5 py-1.5 rounded border border-[var(--border-precision)]">
|
||||
<span className="text-[var(--text-muted)] font-bold block">BIBCODE</span>
|
||||
<span className="text-[var(--text-main)] font-semibold select-all truncate block" title={detailPaper.bibcode === detailPaper.arxiv_id ? '暂无' : detailPaper.bibcode}>
|
||||
{detailPaper.bibcode === detailPaper.arxiv_id ? '暂无' : (
|
||||
<a
|
||||
href={`https://ui.adsabs.harvard.edu/abs/${detailPaper.bibcode}/abstract`}
|
||||
target="_blank"
|
||||
rel="noreferrer"
|
||||
onClick={() => axios.post('/api/active_bibcode', { bibcode: detailPaper.bibcode }).catch(() => {})}
|
||||
className="hover:underline text-sky-600"
|
||||
className="text-[var(--accent-blueprint)] hover:underline"
|
||||
>
|
||||
{detailPaper.bibcode}
|
||||
</a>
|
||||
)}
|
||||
</span>
|
||||
</div>
|
||||
<div className="bg-slate-50 px-2.5 py-1.5 rounded border border-slate-150">
|
||||
<span className="text-slate-400 font-bold block">DOI</span>
|
||||
<span className="text-slate-700 font-semibold select-all truncate block" title={detailPaper.doi || '无'}>
|
||||
<div className="bg-[#f8fafc] px-2.5 py-1.5 rounded border border-[var(--border-precision)]">
|
||||
<span className="text-[var(--text-muted)] font-bold block">DOI</span>
|
||||
<span className="text-[var(--text-main)] font-semibold select-all truncate block" title={detailPaper.doi || '无'}>
|
||||
{detailPaper.doi ? (
|
||||
<a
|
||||
href={`https://doi.org/${detailPaper.doi}`}
|
||||
target="_blank"
|
||||
rel="noreferrer"
|
||||
onClick={() => axios.post('/api/active_bibcode', { bibcode: detailPaper.bibcode }).catch(() => {})}
|
||||
className="hover:underline text-sky-600"
|
||||
className="text-[var(--accent-blueprint)] hover:underline"
|
||||
>
|
||||
{detailPaper.doi}
|
||||
</a>
|
||||
) : '无'}
|
||||
</span>
|
||||
</div>
|
||||
<div className="bg-slate-50 px-2.5 py-1.5 rounded border border-slate-150">
|
||||
<span className="text-slate-450 font-bold block">ARXIV ID</span>
|
||||
<span className="text-slate-700 font-semibold select-all truncate block" title={detailPaper.arxiv_id || '无'}>
|
||||
<div className="bg-[#f8fafc] px-2.5 py-1.5 rounded border border-[var(--border-precision)]">
|
||||
<span className="text-[var(--text-muted)] font-bold block">ARXIV ID</span>
|
||||
<span className="text-[var(--text-main)] font-semibold select-all truncate block" title={detailPaper.arxiv_id || '无'}>
|
||||
{detailPaper.arxiv_id ? (
|
||||
<a
|
||||
href={`https://arxiv.org/abs/${detailPaper.arxiv_id}`}
|
||||
target="_blank"
|
||||
rel="noreferrer"
|
||||
onClick={() => axios.post('/api/active_bibcode', { bibcode: detailPaper.bibcode }).catch(() => {})}
|
||||
className="hover:underline text-sky-600"
|
||||
className="text-[var(--accent-blueprint)] hover:underline"
|
||||
>
|
||||
{detailPaper.arxiv_id}
|
||||
</a>
|
||||
@@ -946,14 +946,14 @@ export default function App() {
|
||||
{/* 手动上传文件（应对防爬阻断） */}
|
||||
<div className="border-t border-slate-100 pt-3 space-y-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-slate-450 font-bold block">手动离线上传文献</span>
|
||||
<span className="text-[var(--text-muted)] font-bold block">手动离线上传文献</span>
|
||||
<span className="text-[9px] text-amber-700 font-bold bg-amber-50 px-2 py-0.5 rounded border border-amber-200">防爬/人机验证备用</span>
|
||||
</div>
|
||||
<p className="text-[10px] text-slate-450 leading-relaxed">
|
||||
<p className="text-[10px] text-[var(--text-muted)] leading-relaxed">
|
||||
若自动下载受阻，可在浏览器中打开上方链接，手动保存 PDF 或 HTML 后在此处上传覆盖。
|
||||
</p>
|
||||
<div className="grid grid-cols-2 gap-2">
|
||||
<div className="flex flex-col items-center justify-center border border-dashed border-slate-300 rounded-lg p-2 hover:bg-slate-50 transition-colors relative cursor-pointer group min-h-[50px]">
|
||||
<div className="flex flex-col items-center justify-center border border-dashed border-[var(--border-precision)] rounded-lg p-2 hover:bg-[#f8fafc] transition-colors relative cursor-pointer group min-h-[50px]">
|
||||
<input
|
||||
type="file"
|
||||
accept="application/pdf"
|
||||
@@ -969,13 +969,13 @@ export default function App() {
|
||||
已下载
|
||||
</span>
|
||||
)}
|
||||
<span className="text-[10px] font-bold text-sky-600 group-hover:underline">
|
||||
<span className="text-[10px] font-bold text-[var(--accent-blueprint)] group-hover:underline">
|
||||
{uploadingBibcode === detailPaper.bibcode ? '上传中...' : '上传 PDF 文献'}
|
||||
</span>
|
||||
<span className="text-[8px] text-slate-400">支持 .pdf 格式</span>
|
||||
<span className="text-[8px] text-[var(--text-muted)]">支持 .pdf 格式</span>
|
||||
</div>
|
||||
|
||||
<div className="flex flex-col items-center justify-center border border-dashed border-slate-300 rounded-lg p-2 hover:bg-slate-50 transition-colors relative cursor-pointer group min-h-[50px]">
|
||||
<div className="flex flex-col items-center justify-center border border-dashed border-[var(--border-precision)] rounded-lg p-2 hover:bg-[#f8fafc] transition-colors relative cursor-pointer group min-h-[50px]">
|
||||
<input
|
||||
type="file"
|
||||
accept="text/html,.html"
|
||||
@@ -991,10 +991,10 @@ export default function App() {
|
||||
已下载
|
||||
</span>
|
||||
)}
|
||||
<span className="text-[10px] font-bold text-sky-600 group-hover:underline">
|
||||
<span className="text-[10px] font-bold text-[var(--accent-blueprint)] group-hover:underline">
|
||||
{uploadingBibcode === detailPaper.bibcode ? '上传中...' : '上传 HTML 文献'}
|
||||
</span>
|
||||
<span className="text-[8px] text-slate-400">支持 .html 格式</span>
|
||||
<span className="text-[8px] text-[var(--text-muted)]">支持 .html 格式</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -1004,7 +1004,7 @@ export default function App() {
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => handleMarkNoResource(detailPaper.bibcode, true)}
|
||||
className="w-full py-2 bg-slate-100 hover:bg-slate-200 text-slate-700 rounded-lg text-[10px] font-bold transition-all border border-slate-250 cursor-pointer flex items-center justify-center gap-1.5"
|
||||
className="btn-console btn-console-secondary w-full py-2 rounded-lg text-[10px] font-bold transition-all cursor-pointer flex items-center justify-center gap-1.5"
|
||||
>
|
||||
<RefreshCw className="w-3 h-3" /> 恢复自动下载状态 (允许后续重新尝试)
|
||||
</button>
|
||||
@@ -1012,9 +1012,9 @@ export default function App() {
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => handleMarkNoResource(detailPaper.bibcode, false)}
|
||||
className="w-full py-2 bg-slate-100 hover:bg-slate-200 text-slate-700 rounded-lg text-[10px] font-bold transition-all border border-slate-250 cursor-pointer flex items-center justify-center gap-1.5"
|
||||
className="btn-console btn-console-secondary w-full py-2 rounded-lg text-[10px] font-bold transition-all cursor-pointer flex items-center justify-center gap-1.5"
|
||||
>
|
||||
<AlertTriangle className="w-3 h-3 text-amber-500" /> 标记为“无有效全文资源” (排查后跳过重试)
|
||||
<AlertTriangle className="w-3 h-3 text-[var(--accent-star)]" /> 标记为“无有效全文资源” (排查后跳过重试)
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
@@ -1080,7 +1080,7 @@ export default function App() {
|
||||
setActiveTab('citation');
|
||||
loadCitations(detailPaper.bibcode);
|
||||
}}
|
||||
className="flex-1 bg-white hover:bg-slate-50 text-slate-700 border border-slate-250 py-2.5 rounded-lg text-xs font-bold text-center transition-all cursor-pointer shadow-sm"
|
||||
className="flex-1 btn-console btn-console-secondary py-2.5 rounded-lg text-xs font-bold text-center cursor-pointer"
|
||||
>
|
||||
查看引用图谱
|
||||
</button>
|
||||
@@ -1091,7 +1091,7 @@ export default function App() {
|
||||
}, '确认重新下载');
|
||||
}}
|
||||
disabled={downloadingBibcodes[detailPaper.bibcode]}
|
||||
className="px-4 py-2.5 bg-slate-100 hover:bg-slate-200 text-amber-700 rounded-lg text-xs font-bold transition-all cursor-pointer disabled:opacity-50"
|
||||
className="px-4 btn-console btn-console-secondary text-[var(--accent-star)] py-2.5 rounded-lg text-xs font-bold transition-all cursor-pointer disabled:opacity-50"
|
||||
>
|
||||
{downloadingBibcodes[detailPaper.bibcode] ? '重下中...' : '重新下载'}
|
||||
</button>
|
||||
@@ -1124,7 +1124,7 @@ export default function App() {
|
||||
setActiveTab('citation');
|
||||
loadCitations(detailPaper.bibcode);
|
||||
}}
|
||||
className="px-6 bg-white hover:bg-slate-50 text-slate-700 border border-slate-250 py-2.5 rounded-lg text-xs font-bold text-center transition-all cursor-pointer shadow-sm"
|
||||
className="px-6 btn-console btn-console-secondary py-2.5 rounded-lg text-xs font-bold text-center cursor-pointer"
|
||||
>
|
||||
引用图谱
|
||||
</button>
|
||||
|
||||
@@ -11,7 +11,8 @@ import 'katex/dist/katex.min.css';
|
||||
import {
|
||||
Brain, Settings, Eye, CheckCircle2, AlertTriangle,
|
||||
Send, Loader, Plus, Trash2, Compass, Clock, Square,
|
||||
BarChart3, ScrollText, Network
|
||||
BarChart3, ScrollText, Network, Rewind, RotateCcw,
|
||||
GitBranch, RefreshCw
|
||||
} from 'lucide-react';
|
||||
import { AskUserQuestionCard } from './AskUserQuestionCard';
|
||||
import { PermissionRequestCard } from './PermissionRequestCard';
|
||||
@@ -76,6 +77,7 @@ interface ActiveTurn {
|
||||
interface ProcessedTurn {
|
||||
turn_index: number;
|
||||
question: string;
|
||||
questionMessageId?: number;
|
||||
timeline: TimelineItem[];
|
||||
usage?: {
|
||||
prompt_tokens: number;
|
||||
@@ -85,6 +87,26 @@ interface ProcessedTurn {
|
||||
createdAt: string;
|
||||
}
|
||||
|
||||
interface RewindResult {
|
||||
rewound_count: number;
|
||||
target_preview: string;
|
||||
new_turn_index: number;
|
||||
session_id: string;
|
||||
}
|
||||
|
||||
interface BranchResult {
|
||||
branch_session_id: string;
|
||||
forked_at_message_id: number;
|
||||
copied_count: number;
|
||||
}
|
||||
|
||||
interface RetryResult {
|
||||
retried_message: string;
|
||||
new_turn_index: number;
|
||||
deleted_count: number;
|
||||
session_id: string;
|
||||
}
|
||||
|
||||
const safeSchema = {
|
||||
...defaultSchema,
|
||||
attributes: {
|
||||
@@ -119,9 +141,9 @@ function getToolDisplayName(name: string): string {
|
||||
case 'download_paper': return '下载文献全文资源';
|
||||
case 'parse_paper': return '结构化解析文献内容';
|
||||
case 'get_paper_content': return '获取文献全文内容';
|
||||
case 'rag_search': return '语义库检索 (RAG)';
|
||||
case 'query_target': return '查询 CDS Sesame 天体物理参数';
|
||||
case 'save_note': return '保存研究笔记';
|
||||
case 'rag_search': return '检索馆藏知识库';
|
||||
case 'query_target': return '查询天体物理参数 (CDS)';
|
||||
case 'save_note': return '保存文献手札';
|
||||
// Agent 控制工具
|
||||
case 'todo_write': return '管理任务列表';
|
||||
case 'compress_context': return '压缩上下文窗口';
|
||||
@@ -325,6 +347,156 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
}
|
||||
};
|
||||
|
||||
// 回退会话到指定消息
|
||||
const handleRewind = async (messageId?: number, n?: number) => {
|
||||
if (!currentSessionId) return;
|
||||
|
||||
const msg = messageId
|
||||
? '确定要回退到此消息之前吗？此后的对话将被移除（可通过“恢复”按钮撤销）。'
|
||||
: `确定要回退最近 ${n || 1} 个对话轮次吗？`;
|
||||
|
||||
const performRewind = async () => {
|
||||
try {
|
||||
const body: { n?: number; message_id?: number } = {};
|
||||
if (messageId) body.message_id = messageId;
|
||||
else body.n = n || 1;
|
||||
|
||||
const res = await axios.post<RewindResult>(
|
||||
`/api/chat/sessions/${currentSessionId}/rewind`,
|
||||
body
|
||||
);
|
||||
|
||||
if (showAlert) {
|
||||
showAlert(`已回退 ${res.data.rewound_count} 条消息`, '成功');
|
||||
} else {
|
||||
alert(`已回退 ${res.data.rewound_count} 条消息`);
|
||||
}
|
||||
|
||||
// 重新加载会话（静默刷新，避免闪烁）
|
||||
loadSessionHistory(currentSessionId, true);
|
||||
fetchSessions(); // 更新侧栏 turn_count
|
||||
} catch (e: any) {
|
||||
const errMsg = `回退失败: ${e.response?.data || e.message}`;
|
||||
if (showAlert) {
|
||||
showAlert(errMsg, '错误');
|
||||
} else {
|
||||
alert(errMsg);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if (showConfirm) {
|
||||
showConfirm(msg, performRewind, '确认回退');
|
||||
} else {
|
||||
if (window.confirm(msg)) {
|
||||
performRewind();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 恢复上次回退（undo-of-undo）
|
||||
const handleRestoreRewind = async () => {
|
||||
if (!currentSessionId) return;
|
||||
|
||||
try {
|
||||
const res = await axios.post<{ restored_count: number }>(
|
||||
`/api/chat/sessions/${currentSessionId}/rewind/restore`
|
||||
);
|
||||
if (res.data.restored_count > 0) {
|
||||
const msg = `已恢复 ${res.data.restored_count} 条被回退的消息`;
|
||||
if (showAlert) {
|
||||
showAlert(msg, '成功');
|
||||
} else {
|
||||
alert(msg);
|
||||
}
|
||||
loadSessionHistory(currentSessionId, true);
|
||||
fetchSessions();
|
||||
} else {
|
||||
const msg = '没有可恢复的回退操作';
|
||||
if (showAlert) {
|
||||
showAlert(msg, '提示');
|
||||
} else {
|
||||
alert(msg);
|
||||
}
|
||||
}
|
||||
} catch (e: any) {
|
||||
const errMsg = `恢复失败: ${e.response?.data || e.message}`;
|
||||
if (showAlert) {
|
||||
showAlert(errMsg, '错误');
|
||||
} else {
|
||||
alert(errMsg);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 分叉当前会话
|
||||
const handleBranch = async () => {
|
||||
if (!currentSessionId) return;
|
||||
|
||||
const confirmed = window.confirm('确定要将当前活跃的消息分叉到一个新会话吗？');
|
||||
if (!confirmed) return;
|
||||
|
||||
try {
|
||||
const res = await axios.post<BranchResult>(
|
||||
`/api/chat/sessions/${currentSessionId}/branch`
|
||||
);
|
||||
if (showAlert) {
|
||||
showAlert(`成功分叉会话！已复制 ${res.data.copied_count} 条消息`, '成功');
|
||||
} else {
|
||||
alert(`成功分叉会话！已复制 ${res.data.copied_count} 条消息`);
|
||||
}
|
||||
|
||||
// 刷新列表并选中新的分叉会话
|
||||
await fetchSessions();
|
||||
setCurrentSessionId(res.data.branch_session_id);
|
||||
} catch (e: any) {
|
||||
const errMsg = `分叉失败: ${e.response?.data || e.message}`;
|
||||
if (showAlert) {
|
||||
showAlert(errMsg, '错误');
|
||||
} else {
|
||||
alert(errMsg);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 重试最后一轮
|
||||
const handleRetry = async () => {
|
||||
if (!currentSessionId || streaming) return;
|
||||
|
||||
const confirmed = window.confirm('确定要重试最后一轮对话吗？原回答将被删除。');
|
||||
if (!confirmed) return;
|
||||
|
||||
try {
|
||||
const res = await axios.post<RetryResult>(
|
||||
`/api/chat/sessions/${currentSessionId}/retry`
|
||||
);
|
||||
|
||||
const questionText = res.data.retried_message;
|
||||
if (!questionText) {
|
||||
if (showAlert) {
|
||||
showAlert('没有找到可重试的上一轮问题', '提示');
|
||||
} else {
|
||||
alert('没有找到可重试的上一轮问题');
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// 重新加载会话以移除已被硬删除的消息，然后自动触发发送
|
||||
await loadSessionHistory(currentSessionId, true);
|
||||
await fetchSessions();
|
||||
|
||||
// 触发自动重发
|
||||
handleSend(questionText);
|
||||
} catch (e: any) {
|
||||
const errMsg = `重试失败: ${e.response?.data || e.message}`;
|
||||
if (showAlert) {
|
||||
showAlert(errMsg, '错误');
|
||||
} else {
|
||||
alert(errMsg);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 手动停止智能体执行
|
||||
const handleStop = async () => {
|
||||
if (!currentSessionId) return;
|
||||
@@ -1082,10 +1254,21 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
{turns.map((turn) => (
|
||||
<div key={turn.turn_index} className="space-y-4">
|
||||
{/* 用户提问 */}
|
||||
<div className="flex flex-col items-end space-y-1">
|
||||
<div className="flex flex-col items-end space-y-1 group relative">
|
||||
<span className="text-[10px] font-bold text-slate-400 px-1">我</span>
|
||||
<div className="max-w-[85%] rounded-2xl px-4 py-3 text-xs leading-relaxed font-semibold shadow-2xs border bg-sky-600 text-white border-sky-600 select-text">
|
||||
{turn.question}
|
||||
<div className="flex items-center gap-2 max-w-[85%]">
|
||||
{turn.questionMessageId && (
|
||||
<button
|
||||
onClick={() => handleRewind(turn.questionMessageId)}
|
||||
className="opacity-0 group-hover:opacity-100 text-slate-400 hover:text-amber-600 p-1.5 rounded-lg hover:bg-slate-100 transition-all cursor-pointer shrink-0"
|
||||
title="回退到此消息之前"
|
||||
>
|
||||
<Rewind className="w-3.5 h-3.5" />
|
||||
</button>
|
||||
)}
|
||||
<div className="rounded-2xl px-4 py-3 text-xs leading-relaxed font-semibold shadow-2xs border bg-sky-600 text-white border-sky-600 select-text flex-1">
|
||||
{turn.question}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -1151,7 +1334,7 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
|
||||
{/* Error block */}
|
||||
{activeTurn.error && (
|
||||
<div className="flex items-start gap-2 bg-red-50 border border-red-200 text-red-700 rounded-lg p-3 text-xs w-[90%] font-semibold">
|
||||
<div className="flex items-start gap-2 bg-red-50 border border-red-200 text-red-750 rounded-lg p-3 text-xs w-[90%] font-semibold">
|
||||
<AlertTriangle className="w-4 h-4 text-red-500 shrink-0 mt-0.5" />
|
||||
<div>
|
||||
<div className="font-bold">查询发生错误</div>
|
||||
@@ -1161,6 +1344,40 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* 会话操作小图标栏 */}
|
||||
{currentSessionId && turns.length > 0 && !streaming && (
|
||||
<div className="flex items-center gap-1 pt-2.5 justify-start border-t border-slate-100/60 mt-4 max-w-[200px]">
|
||||
<button
|
||||
onClick={() => handleRewind(undefined, 1)}
|
||||
className="p-1.5 rounded-lg text-slate-400 hover:text-amber-600 hover:bg-slate-100 transition-all cursor-pointer"
|
||||
title="回退最近一轮对话 (undo)"
|
||||
>
|
||||
<Rewind className="w-4 h-4" />
|
||||
</button>
|
||||
<button
|
||||
onClick={handleRestoreRewind}
|
||||
className="p-1.5 rounded-lg text-slate-400 hover:text-green-600 hover:bg-slate-100 transition-all cursor-pointer"
|
||||
title="恢复上次回退操作"
|
||||
>
|
||||
<RotateCcw className="w-4 h-4" />
|
||||
</button>
|
||||
<button
|
||||
onClick={handleRetry}
|
||||
className="p-1.5 rounded-lg text-slate-400 hover:text-indigo-600 hover:bg-slate-100 transition-all cursor-pointer"
|
||||
title="重试最后一轮对话"
|
||||
>
|
||||
<RefreshCw className="w-4 h-4" />
|
||||
</button>
|
||||
<button
|
||||
onClick={handleBranch}
|
||||
className="p-1.5 rounded-lg text-slate-400 hover:text-blue-600 hover:bg-slate-100 transition-all cursor-pointer"
|
||||
title="分叉当前会话"
|
||||
>
|
||||
<GitBranch className="w-4 h-4" />
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
<div ref={chatEndRef} />
|
||||
@@ -1269,6 +1486,7 @@ function groupMessagesIntoTurns(messages: MessageRecord[]): ProcessedTurn[] {
|
||||
|
||||
if (msg.role === 'user') {
|
||||
turn.question = msg.content;
|
||||
turn.questionMessageId = msg.id;
|
||||
} else if (msg.role === 'assistant') {
|
||||
const stepNum = msg.step_index;
|
||||
const hasToolCalls = msg.tool_calls && msg.tool_calls.length > 0;
|
||||
|
||||
@@ -547,26 +547,26 @@ export function ReaderPanel({
|
||||
onClick={() => handleVectorize(selectedPaper.bibcode)}
|
||||
disabled={vectorizing}
|
||||
className="btn-console btn-console-primary px-4 py-2 rounded-lg text-xs font-bold flex items-center gap-2"
|
||||
title="对文献进行向量化分块入库，以开启学术 AI 问答"
|
||||
title="对文献进行知识入库，以开启学术 AI 研讨"
|
||||
>
|
||||
{vectorizing ? <Loader className="w-3.5 h-3.5 animate-spin" /> : <Sparkles className="w-3.5 h-3.5" />}
|
||||
{vectorizing ? '向量化入库中...' : '向量化入库'}
|
||||
{vectorizing ? '正在进行知识入库...' : '知识入库'}
|
||||
</button>
|
||||
)}
|
||||
|
||||
{selectedPaper.has_markdown && selectedPaper.has_vector && (
|
||||
<button
|
||||
onClick={() => {
|
||||
showConfirm('确定要重新向量化入库吗？这会清空先前该文献的切片记录并重新执行入库。', () => {
|
||||
showConfirm('确定要重新进行知识入库吗？这会清除先前该文献已有的知识点记录并重新写入。', () => {
|
||||
handleVectorize(selectedPaper.bibcode);
|
||||
}, '确认重新向量化');
|
||||
}, '确认重新知识入库');
|
||||
}}
|
||||
disabled={vectorizing}
|
||||
className="btn-console btn-console-secondary px-4 py-2 rounded-lg text-xs font-bold flex items-center gap-2"
|
||||
title="重新为该文献生成 RAG 向量切片"
|
||||
title="重新为该文献解析知识点并入库"
|
||||
>
|
||||
{vectorizing ? <Loader className="w-3.5 h-3.5 animate-spin" /> : <RotateCw className="w-3.5 h-3.5" />}
|
||||
重新向量化
|
||||
重新知识入库
|
||||
</button>
|
||||
)}
|
||||
|
||||
@@ -941,56 +941,56 @@ export function ReaderPanel({
|
||||
top: `${hoverCardPos.y}px`,
|
||||
zIndex: 9999,
|
||||
}}
|
||||
className="console-panel rounded-xl p-4 bg-white border border-slate-200 shadow-xl w-72 text-xs space-y-2 pointer-events-auto animate-in fade-in duration-200"
|
||||
className="console-panel rounded-xl p-4 bg-[var(--bg-card)] border border-[var(--border-precision)] shadow-md w-72 text-xs space-y-2 pointer-events-auto animate-in fade-in duration-200"
|
||||
>
|
||||
<div className="flex items-center justify-between border-b border-slate-100 pb-1.5">
|
||||
<span className="font-bold text-slate-900 text-sm">{hoveredTarget.target_name}</span>
|
||||
<span className="font-bold text-[var(--text-main)] text-sm">{hoveredTarget.target_name}</span>
|
||||
<div className="flex items-center gap-2 select-none">
|
||||
<a
|
||||
href={`https://simbad.cds.unistra.fr/simbad/sim-id?Ident=${encodeURIComponent(hoveredTarget.target_name)}`}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="text-[10px] font-bold text-sky-600 hover:text-sky-800 hover:underline cursor-pointer"
|
||||
className="text-[10px] font-bold text-[var(--accent-blueprint)] hover:text-[#0d5988] hover:underline cursor-pointer"
|
||||
title="在 SIMBAD 中查询该天体详情"
|
||||
>
|
||||
SIMBAD
|
||||
</a>
|
||||
<span className="text-[9px] text-slate-300">|</span>
|
||||
<span className="text-[9px] text-[var(--border-precision)]">|</span>
|
||||
<a
|
||||
href={`https://vizier.cds.unistra.fr/viz-bin/VizieR?-c=${encodeURIComponent(hoveredTarget.target_name)}`}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="text-[10px] font-bold text-sky-600 hover:text-sky-800 hover:underline cursor-pointer"
|
||||
className="text-[10px] font-bold text-[var(--accent-blueprint)] hover:text-[#0d5988] hover:underline cursor-pointer"
|
||||
title="在 VizieR 中查询相关文献和表数据"
|
||||
>
|
||||
VizieR
|
||||
</a>
|
||||
</div>
|
||||
</div>
|
||||
<div className="grid grid-cols-2 gap-x-2 gap-y-1.5 text-slate-600">
|
||||
<div className="grid grid-cols-2 gap-x-2 gap-y-1.5 text-[var(--text-muted)]">
|
||||
<div>
|
||||
<span className="text-slate-400 font-semibold">RA (J2000):</span>
|
||||
<div className="font-mono font-bold text-slate-800">{hoveredTarget.ra || '未知'}</div>
|
||||
<span className="text-[var(--text-muted)] font-semibold">RA (J2000):</span>
|
||||
<div className="font-mono font-bold text-[var(--text-main)]">{hoveredTarget.ra || '未知'}</div>
|
||||
</div>
|
||||
<div>
|
||||
<span className="text-slate-400 font-semibold">Dec (J2000):</span>
|
||||
<div className="font-mono font-bold text-slate-800">{hoveredTarget.dec || '未知'}</div>
|
||||
<span className="text-[var(--text-muted)] font-semibold">Dec (J2000):</span>
|
||||
<div className="font-mono font-bold text-[var(--text-main)]">{hoveredTarget.dec || '未知'}</div>
|
||||
</div>
|
||||
<div>
|
||||
<span className="text-slate-400 font-semibold">光谱型:</span>
|
||||
<div className="font-bold text-slate-800">{hoveredTarget.spectral_type || '未知'}</div>
|
||||
<span className="text-[var(--text-muted)] font-semibold">光谱型:</span>
|
||||
<div className="font-bold text-[var(--text-main)]">{hoveredTarget.spectral_type || '未知'}</div>
|
||||
</div>
|
||||
<div>
|
||||
<span className="text-slate-400 font-semibold">视星等 (V):</span>
|
||||
<div className="font-bold text-slate-800">
|
||||
<span className="text-[var(--text-muted)] font-semibold">视星等 (V):</span>
|
||||
<div className="font-bold text-[var(--text-main)]">
|
||||
{hoveredTarget.v_magnitude !== null && hoveredTarget.v_magnitude !== undefined
|
||||
? `${hoveredTarget.v_magnitude.toFixed(2)}`
|
||||
: '未知'}
|
||||
</div>
|
||||
</div>
|
||||
<div className="col-span-2">
|
||||
<span className="text-slate-400 font-semibold">视差 / 估算距离:</span>
|
||||
<div className="font-bold text-slate-800">
|
||||
<span className="text-[var(--text-muted)] font-semibold">视差 / 估算距离:</span>
|
||||
<div className="font-bold text-[var(--text-main)]">
|
||||
{hoveredTarget.parallax !== null && hoveredTarget.parallax !== undefined
|
||||
? `${hoveredTarget.parallax.toFixed(2)} mas (~${(1000.0 / hoveredTarget.parallax).toFixed(1)} pc)`
|
||||
: '未知'}
|
||||
@@ -999,8 +999,8 @@ export function ReaderPanel({
|
||||
</div>
|
||||
{hoveredTarget.aliases && hoveredTarget.aliases.length > 0 && (
|
||||
<div className="border-t border-slate-100 pt-1.5">
|
||||
<span className="text-slate-400 font-semibold block mb-0.5">常用别名:</span>
|
||||
<div className="text-[10px] text-slate-500 font-medium leading-relaxed max-h-16 overflow-y-auto font-mono">
|
||||
<span className="text-[var(--text-muted)] font-semibold block mb-0.5">常用别名:</span>
|
||||
<div className="text-[10px] text-[var(--text-muted)] font-medium leading-relaxed max-h-16 overflow-y-auto font-mono">
|
||||
{hoveredTarget.aliases.slice(0, 8).join(', ')}
|
||||
{hoveredTarget.aliases.length > 8 && ' ...'}
|
||||
</div>
|
||||
|
||||
+32
-31
@@ -1,4 +1,4 @@
|
||||
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
|
||||
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500;700&display=swap');
|
||||
@import "tailwindcss";
|
||||
@plugin "@tailwindcss/typography";
|
||||
|
||||
@@ -6,15 +6,18 @@
|
||||
font-family: 'Inter', system-ui, -apple-system, sans-serif;
|
||||
color-scheme: light;
|
||||
|
||||
--bg-main: #f1f5f9;
|
||||
/* 极简白蓝学术科技风配色 */
|
||||
--bg-main: #f4f6f9;
|
||||
--bg-card: #ffffff;
|
||||
--bg-sidebar: #f8fafc;
|
||||
--text-main: #0f172a;
|
||||
--text-muted: #475569;
|
||||
--bg-sidebar: #0f2540; /* 保持深色侧边栏作为界面骨架结构 */
|
||||
--text-main: #0a2540;
|
||||
--text-muted: #5c6b84;
|
||||
|
||||
--accent-blue: #0284c7;
|
||||
--accent-navy: #1e3a8a;
|
||||
--border-clean: #e2e8f0;
|
||||
--accent-blueprint: #106ba3;
|
||||
--accent-star: #d97706;
|
||||
--border-precision: #d2d8e2;
|
||||
|
||||
--font-mono: 'JetBrains Mono', monospace;
|
||||
}
|
||||
|
||||
body {
|
||||
@@ -32,7 +35,7 @@ body {
|
||||
height: 6px;
|
||||
}
|
||||
::-webkit-scrollbar-track {
|
||||
background: #f1f5f9;
|
||||
background: #f4f6f9;
|
||||
}
|
||||
::-webkit-scrollbar-thumb {
|
||||
background: #cbd5e1;
|
||||
@@ -45,20 +48,20 @@ body {
|
||||
/* Premium clean panel cards */
|
||||
.console-panel {
|
||||
background: var(--bg-card);
|
||||
border: 1px solid var(--border-clean);
|
||||
box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05), 0 2px 4px -1px rgba(0, 0, 0, 0.03);
|
||||
border: 1px solid var(--border-precision);
|
||||
box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.05); /* 弱化阴影，更显扁平学术感 */
|
||||
}
|
||||
|
||||
.console-panel-active {
|
||||
border-color: var(--accent-blue);
|
||||
box-shadow: 0 0 0 1px var(--accent-blue), 0 4px 6px -1px rgba(0, 0, 0, 0.05);
|
||||
border-color: var(--accent-blueprint);
|
||||
box-shadow: 0 0 0 1px var(--accent-blueprint), 0 1px 3px 0 rgba(0, 0, 0, 0.05);
|
||||
}
|
||||
|
||||
/* High contrast clean console button */
|
||||
.btn-console {
|
||||
background: #ffffff;
|
||||
border: 1px solid #cbd5e1;
|
||||
color: #334155;
|
||||
border: 1px solid var(--border-precision);
|
||||
color: var(--text-main);
|
||||
font-weight: 500;
|
||||
transition: all 0.2s ease;
|
||||
}
|
||||
@@ -66,42 +69,42 @@ body {
|
||||
.btn-console:hover:not(:disabled) {
|
||||
background: #f8fafc;
|
||||
border-color: #94a3b8;
|
||||
color: #0f172a;
|
||||
color: var(--text-main);
|
||||
box-shadow: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
|
||||
}
|
||||
|
||||
.btn-console-primary {
|
||||
background: var(--accent-blue);
|
||||
border: 1px solid var(--accent-blue);
|
||||
background: var(--accent-blueprint);
|
||||
border: 1px solid var(--accent-blueprint);
|
||||
color: #ffffff;
|
||||
}
|
||||
|
||||
.btn-console-primary:hover:not(:disabled) {
|
||||
background: #0369a1;
|
||||
border-color: #0369a1;
|
||||
background: #0d5988;
|
||||
border-color: #0d5988;
|
||||
color: #ffffff;
|
||||
box-shadow: 0 2px 4px 0 rgba(2, 132, 199, 0.2);
|
||||
box-shadow: 0 2px 4px 0 rgba(16, 107, 163, 0.2);
|
||||
}
|
||||
|
||||
.btn-console-secondary {
|
||||
background: #f1f5f9;
|
||||
border: 1px solid #e2e8f0;
|
||||
color: #334155;
|
||||
border: 1px solid var(--border-precision);
|
||||
color: var(--text-muted);
|
||||
}
|
||||
|
||||
.btn-console-secondary:hover:not(:disabled) {
|
||||
background: #e2e8f0;
|
||||
color: #0f172a;
|
||||
color: var(--text-main);
|
||||
}
|
||||
|
||||
/* Premium clean console select dropdown styling */
|
||||
.select-console {
|
||||
display: inline-block;
|
||||
background-color: #ffffff;
|
||||
border: 1px solid #cbd5e1;
|
||||
border: 1px solid var(--border-precision);
|
||||
border-radius: 0.5rem; /* 8px */
|
||||
padding: 0.5rem 1.75rem 0.5rem 0.625rem; /* padding-right leaves space for custom arrow */
|
||||
color: #334155;
|
||||
color: var(--text-main);
|
||||
font-size: 0.75rem; /* text-xs */
|
||||
font-weight: 500;
|
||||
cursor: pointer;
|
||||
@@ -116,13 +119,13 @@ body {
|
||||
|
||||
.select-console:hover:not(:disabled) {
|
||||
border-color: #94a3b8;
|
||||
color: #0f172a;
|
||||
color: var(--text-main);
|
||||
box-shadow: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
|
||||
}
|
||||
|
||||
.select-console:focus {
|
||||
border-color: #0284c7;
|
||||
box-shadow: 0 0 0 1px #0284c7;
|
||||
border-color: var(--accent-blueprint);
|
||||
box-shadow: 0 0 0 1px var(--accent-blueprint);
|
||||
}
|
||||
|
||||
.select-console:disabled {
|
||||
@@ -130,5 +133,3 @@ body {
|
||||
color: #94a3b8;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -172,3 +172,35 @@ export interface AgentTask {
|
||||
created_at: string;
|
||||
updated_at: string;
|
||||
}
|
||||
|
||||
// ── 会话回退 (Rewind) ──
|
||||
|
||||
export interface RewindRequest {
|
||||
n?: number; // 回退 N 个轮次
|
||||
message_id?: number; // 或指定消息 ID
|
||||
}
|
||||
|
||||
export interface RewindResponse {
|
||||
rewound_count: number;
|
||||
target_preview: string;
|
||||
new_turn_index: number;
|
||||
session_id: string;
|
||||
}
|
||||
|
||||
export interface RestoreResponse {
|
||||
restored_count: number;
|
||||
session_id: string;
|
||||
}
|
||||
|
||||
export interface BranchResponse {
|
||||
branch_session_id: string;
|
||||
forked_at_message_id: number;
|
||||
copied_count: number;
|
||||
}
|
||||
|
||||
export interface RetryResponse {
|
||||
retried_message: string;
|
||||
new_turn_index: number;
|
||||
deleted_count: number;
|
||||
session_id: string;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,909 @@
|
||||
# Claude Code / Hermes-Agent 参考分析
|
||||
|
||||
对 Claude Code (`/home/fmq/program/claudecode/src/`) 和 Hermes-Agent (`libs/hermes-agent/`)
|
||||
源码的全面架构分析，记录对 AstroResearch Agent 系统的参考价值与改进方向。
|
||||
|
||||
> 分析日期: 2026-06-22 | 最后更新: 2026-06-22
|
||||
|
||||
## 实施状态
|
||||
|
||||
| 优先级 | 改进项 | 状态 | 涉及文件 |
|
||||
|--------|--------|------|---------|
|
||||
| P0 | Context Overflow 自动修复 | ✅ 已完成 | `error_recovery.rs` (+150 行) |
|
||||
| P0 | StreamingExecutor 真正流式调度 | ✅ 已完成 | `streaming_executor.rs` (重写 ~400 行) |
|
||||
| P0 | Executor 集成分区器（批次串行/并行） | ✅ 已完成 | `executor.rs` (Phase 3 重写 + 2 个提取函数) |
|
||||
| P1 | 工具并发分区 `partition_tool_calls` | ✅ 已完成 | `partitioner.rs` (+2 测试) |
|
||||
| P1 | PermissionRequest / PermissionDenied Hooks | ✅ 已完成 | `hooks/types.rs`, `traits.rs`, `dispatch.rs`, `mod.rs` |
|
||||
| P1 | Auto-mode Classifier | ⏳ 待定 | — |
|
||||
| P2 | Self-improving Skills（模式检测 + 自动创建 + Curator） | ✅ 已完成 | `skills/pattern_detector.rs` + `curator.rs` + `SkillCreator` |
|
||||
| P2 | Coordinator Mode | ⏳ 待定 | — |
|
||||
| P2 | UserPromptSubmit / PreCompact / PostCompact Hook | ⏳ 待定 | — |
|
||||
| P3 | FTS5 跨 session 搜索 | ⏳ 待定 | — |
|
||||
| P3 | Tool `defer_loading` / `classifier_summary` | ⏳ 待定 | — |
|
||||
| P3 | 模型回退策略 | ⏳ 待定 | — |
|
||||
| P3 | Session Memory Compaction | ⏳ 待定 | — |
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
|
||||
1. [总体评估](#1-总体评估)
|
||||
2. [工具并发执行模型](#2-工具并发执行模型)
|
||||
3. [Permission 系统](#3-permission-系统)
|
||||
4. [Hooks 系统](#4-hooks-系统)
|
||||
5. [Error Recovery / 重试系统](#5-error-recovery--重试系统)
|
||||
6. [Tool 定义系统](#7-tool-定义系统)
|
||||
8. [Memory 持久化](#8-memory-持久化)
|
||||
9. [Coordinator / Multi-Agent](#9-coordinator--multi-agent)
|
||||
10. [Hermes-Agent 的独特贡献](#10-hermes-agent-的独特贡献)
|
||||
11. [优先级排序 —— 建议实施路线](#11-优先级排序--建议实施路线)
|
||||
|
||||
---
|
||||
|
||||
## 1. 总体评估
|
||||
|
||||
### 1.1 参考项目概览
|
||||
|
||||
| 维度 | Claude Code | Hermes-Agent | AstroResearch |
|
||||
|------|-------------|-------------|---------------|
|
||||
| 语言 | TypeScript (Node.js) | Python (3.11+) | Rust (Axum) |
|
||||
| 定位 | 终端 IDE 编程助手 | 通用 AI 个人助手 | 天文科研 Agent |
|
||||
| Agent 循环 | 流式 `query()` generator | 同步 `while` 循环 | 流式 ReAct 循环 |
|
||||
| 工具注册 | 手动 import + `getAllBaseTools()` | 文件系统自动发现 | 手动 `ToolRegistry::new()` |
|
||||
| 工具接口 | `Tool<T>` — ~70 个方法 | `handler(args) -> JSON string` | `AgentTool` trait — ~10 个方法 |
|
||||
| 权限系统 | 6 层优先级 + Classifier + 沙箱 | 无内置 | 3 层规则 + PermissionChecker |
|
||||
| Hooks | 27 种事件，6 种 hook 类型 | PluginManager 生命周期 | 15 种事件，2 种 hook 类型 |
|
||||
| 子代理 | `AgentTool` + Fork + Worktree 隔离 | `delegate_task` + 子 AIAgent | `SubAgentTool` + 独立 ReAct |
|
||||
| 多 Agent | Coordinator 模式 + Swarm/Team | Kanban 工作队列 | Team 系统 (lead/teammate) |
|
||||
| 持久化 | 文件系统 Markdown + cost tracker | SQLite (FTS5) + SessionDB | SQLite + MEMORY.md |
|
||||
| 上下文压缩 | 微压缩 + 自动压缩 + 手动压缩 | ContextCompressor | 4 层压缩 (微/snip/auto/aggro) |
|
||||
|
||||
### 1.2 核心结论
|
||||
|
||||
AstroResearch 的 Agent 系统架构本身就是**对标 Claude Code 设计的**——`StreamingToolExecutor`、
|
||||
`PermissionChecker`、`HookRegistry` 都明确标注了参考来源。当前差距主要是**实现深度**而非**设计方向**。
|
||||
|
||||
Claude Code 的参考价值在**工程细节**：流式调度的时机选择、Overflow 的自动修复、Classifier
|
||||
的并行化设计。Hermes-Agent 的独特价值在**自我进化**（Self-improving Skills）和**多 Profile 隔离**。
|
||||
|
||||
---
|
||||
|
||||
## 2. 工具并发执行模型
|
||||
|
||||
### 2.1 对比
|
||||
|
||||
| 特性 | Claude Code | AstroResearch (当前) |
|
||||
|------|-------------|---------------------|
|
||||
| 流式调度 | tool_use 到达**立即**开始执行 | tool_use 全部收集，`flush()` 批量执行 |
|
||||
| 并发分区 | `partitionToolCalls()` 自动分组连续只读工具并行 | `ToolPartitioner` 存在但基本未使用 |
|
||||
| Sibling Abort | Bash 错误 → 级联取消兄弟工具，有专用 `siblingAbortController` | `AbortReason::SiblingError` + 广播通道存在，取消逻辑不完整 |
|
||||
| Progress 流式 | `pendingProgress` 即时 yield，`progressAvailableResolve` 唤醒等待 | `execute_with_progress` 有通道，`getCompletedResults` 未检查 |
|
||||
| 中断行为 | `interruptBehavior()` 区分 `cancel` vs `block` | `InterruptBehavior` 枚举存在但未在 Executor 中使用 |
|
||||
|
||||
### 2.2 Claude Code 的分区逻辑
|
||||
|
||||
```typescript
|
||||
// src/services/tools/toolOrchestration.ts
|
||||
// 自动将连续只读工具分组并行，写工具独立串行
|
||||
function partitionToolCalls(toolUseMessages, toolUseContext): Batch[] {
|
||||
return toolUseMessages.reduce((acc, toolUse) => {
|
||||
const tool = findToolByName(toolUseContext.options.tools, toolUse.name)
|
||||
const parsedInput = tool?.inputSchema.safeParse(toolUse.input)
|
||||
const isConcurrencySafe = parsedInput?.success
|
||||
? (() => { try { return Boolean(tool.isConcurrencySafe(parsedInput.data)) } catch { return false } })()
|
||||
: false
|
||||
|
||||
if (isConcurrencySafe && acc[acc.length - 1]?.isConcurrencySafe) {
|
||||
acc[acc.length - 1].blocks.push(toolUse) // 合并到当前批次
|
||||
} else {
|
||||
acc.push({ isConcurrencySafe, blocks: [toolUse] }) // 新批次
|
||||
}
|
||||
return acc
|
||||
}, [])
|
||||
}
|
||||
```
|
||||
|
||||
关键点：
|
||||
- **输入感知**的并发安全判断：同一工具可能因参数不同而安全属性不同
|
||||
- 并发安全的工具**连续分组**——不打断写入顺序
|
||||
- 非并发安全的工具**独占执行**——等上一个完成后才启动下一个
|
||||
|
||||
### 2.3 Claude Code 的 StreamingToolExecutor 核心逻辑
|
||||
|
||||
```
|
||||
文件: src/services/tools/StreamingToolExecutor.ts (531 行)
|
||||
|
||||
状态机: Queued → Executing → Completed → Yielded
|
||||
|
||||
生命周期:
|
||||
1. addTool(block) — LLM 流产生 tool_use 时立即调用
|
||||
2. processQueue() — 检查 concurrency 条件，启动可执行工具
|
||||
3. executeTool(tool) — 创建子 AbortController，调用 runToolUse generator
|
||||
4. getCompletedResults() — 按序 yield 结果（非阻塞）
|
||||
5. getRemainingResults() — 等待未完成工具（async generator）
|
||||
|
||||
关键设计:
|
||||
- siblingAbortController: 父 AbortController 的子节点
|
||||
Bash 错误 → siblingAbortController.abort('sibling_error') → 取消所有兄弟
|
||||
但 toolAbortController 的 abort 会向上冒泡到父 AbortController
|
||||
- progressAvailableResolve: Promise resolver 用于唤醒等待 progress 的 getRemainingResults
|
||||
- 中断行为: 'cancel' 工具被用户中断时生成 REJECT_MESSAGE; 'block' 工具不受影响
|
||||
```
|
||||
|
||||
### 2.4 AstroResearch 的现状（2026-06-22 更新）
|
||||
|
||||
```
|
||||
文件: src/agent/runtime/streaming_executor.rs (~400 行，已重写)
|
||||
|
||||
已实现:
|
||||
✅ TrackedTool 状态机 (Queued/Executing/Completed/Yielded)
|
||||
✅ Sibling Abort 广播通道（broadcast::channel + tokio::select! 竞速）
|
||||
✅ on_tool_use / flush / next_result 接口
|
||||
✅ 输出截断
|
||||
✅ 真正的流式调度 — on_tool_use 中对并发安全工具立即 spawn tokio task
|
||||
✅ 非并发安全工具独占执行 — executing_non_concurrent 标志阻塞后续启动
|
||||
✅ 并发取消 — tokio::select! 在工具执行和 Sibling Abort 之间竞速
|
||||
✅ 输入感知的并发安全判断 — 通过 tool_registry.get().is_concurrency_safe(&args)
|
||||
✅ ToolContext 实现 Clone（支持 per-task 复制上下文）
|
||||
|
||||
与 Claude Code 的对齐:
|
||||
- 核心理念一致：addTool → 立即 processQueue
|
||||
- Sibling Abort 机制等效：broadcast::Sender + subscribe
|
||||
- collectCompletedTasks 使用 JoinHandle::is_finished() 做非阻塞检查
|
||||
```
|
||||
|
||||
|
||||
### 2.5 已实施改进（2026-06-22）
|
||||
|
||||
**✅ P0: 真正的流式调度** — 已完成
|
||||
|
||||
`on_tool_use` 中对并发安全工具立即 `tokio::spawn`，非并发安全工具标记 `executing_non_concurrent`
|
||||
并阻塞后续启动，直到独占工具完成。
|
||||
|
||||
**✅ P0: Executor 集成分区器** — 已完成
|
||||
|
||||
`src/agent/runtime/executor.rs` Phase 3 从「所有非拒绝工具单一 `FuturesUnordered` 无差别并发」
|
||||
改为「`ToolPartitioner` 分区 → 逐批次执行」：
|
||||
- 并行批次内 `FuturesUnordered` 并发
|
||||
- 串行批次内逐个执行（非并发安全工具独占）
|
||||
- 提取 `execute_single_tool()` 和 `process_single_result()` 两个辅助函数
|
||||
|
||||
**✅ P1: 并发分区器** — 已完成
|
||||
|
||||
`src/agent/runtime/partitioner.rs` 新增 2 个测试（`run_bash`、`file_write` 打断并发批）。
|
||||
|
||||
### 2.6 原始建议（已过时）
|
||||
|
||||
```rust
|
||||
// 建议在 on_tool_use 中对并发安全的工具立即 spawn
|
||||
pub fn on_tool_use(&mut self, call_id: String, name: String, args: Value) -> bool {
|
||||
let is_concurrency_safe = /* 判断 */
|
||||
let idx = self.tracked.len();
|
||||
self.tracked.push(tool);
|
||||
if is_concurrency_safe && self.can_execute_now() {
|
||||
let handle = tokio::spawn(/* 执行 */);
|
||||
self.tracked[idx].handle = Some(handle);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**P1: 并发分区**
|
||||
|
||||
```rust
|
||||
/// 将 tool_use 列表分区为 (并发安全批次, 非并发安全单例)
|
||||
fn partition_tool_calls(calls: &[PreparedCall], registry: &ToolRegistry) -> Vec<Batch> {
|
||||
calls.iter().fold(Vec::new(), |mut acc, call| {
|
||||
let is_safe = registry.get(&call.tool_name)
|
||||
.map(|t| t.is_concurrency_safe(&call.args))
|
||||
.unwrap_or(false);
|
||||
if is_safe && acc.last().map_or(false, |b: &Batch| b.concurrent) {
|
||||
acc.last_mut().unwrap().calls.push(call.clone());
|
||||
} else {
|
||||
acc.push(Batch { concurrent: is_safe, calls: vec![call.clone()] });
|
||||
}
|
||||
acc
|
||||
})
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. Permission 系统
|
||||
|
||||
### 3.1 对比
|
||||
|
||||
| 特性 | Claude Code | AstroResearch (当前) |
|
||||
|------|-------------|---------------------|
|
||||
| 规则来源分层 | 6 层优先级：policy > project > user > plugin > flag > command | 单一规则列表 |
|
||||
| 规则行为 | Allow / Deny / Ask | Allow / Deny / Ask ✅ |
|
||||
| Classifier 自动模式 | 两阶段（快速 + 思考），并行于 hooks 启动 | 无 |
|
||||
| 拒绝追踪 | 带时间窗口的限流回退 (DenialTracker) | 简单计数 |
|
||||
| 沙箱集成 | `shouldUseSandbox()` + `sandbox-adapter` | 无沙箱概念 |
|
||||
| 决策溯源 | 每条 PermissionDecisionReason 记录完整来源链 | 只返回 Allow/Deny/Ask |
|
||||
| 权限模式 | 5 种：default, acceptEdits, bypassPermissions, dontAsk, plan | 4 种 ✅ |
|
||||
|
||||
### 3.2 Claude Code 的 Permission 决策管道
|
||||
|
||||
```
|
||||
1. validateInput() — Zod schema 验证
|
||||
2. runPreToolUseHooks() — Session hooks（用户配置的）
|
||||
3. canUseTool — 检查 deny 规则
|
||||
4. resolveHookPermissionDecision() — Allow 规则
|
||||
5. [auto mode] Classifier — 两阶段分类器（Haiku）
|
||||
6. [default mode] 用户弹窗 — 交互式确认
|
||||
7. PermissionDecisionReason — 记录决策来源
|
||||
```
|
||||
|
||||
每个决策都烙印 `PermissionDecisionReason`：
|
||||
```
|
||||
rule | mode | subcommandResults | permissionPromptTool |
|
||||
hook | asyncAgent | sandboxOverride | classifier |
|
||||
workingDir | safetyCheck | other
|
||||
```
|
||||
|
||||
### 3.3 Classifier 系统（最值得借鉴）
|
||||
|
||||
Claude Code 的 auto-mode classifier 是一个**独立的小模型调用**（Haiku），在后台并行运行：
|
||||
|
||||
```
|
||||
Auto Mode 决策流程:
|
||||
┌──────────────────────────────────────────────┐
|
||||
│ 1. startSpeculativeClassifierCheck() │
|
||||
│ └─ 并行于 PreToolUse hooks 启动 │
|
||||
│ 2. 两阶段分类: │
|
||||
│ ├─ Fast: 简单模式匹配（秒级） │
|
||||
│ └─ Thinking: 深度分析（复杂命令时） │
|
||||
│ 3. 结果: Allow / Deny / Ask + confidence │
|
||||
│ 4. DenialTracker: 连续 Deny 后 fallback 用户 │
|
||||
└──────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
AstroResearch 目前没有 auto-mode——所有非白名单工具都需要用户交互确认。
|
||||
|
||||
### 3.4 建议改进
|
||||
|
||||
**P1: Auto-mode Classifier**
|
||||
|
||||
```rust
|
||||
/// Auto-mode 分类器 — 使用廉价模型在后台预分类工具调用
|
||||
pub struct AutoClassifier {
|
||||
llm: LlmClient, // 使用廉价模型（如 Haiku 级别 provider）
|
||||
cache: LruCache<String, ClassificationResult>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ClassificationResult {
|
||||
pub decision: PermissionResult,
|
||||
pub confidence: f64,
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
impl AutoClassifier {
|
||||
/// 在工具执行前异步预分类（不阻塞用户）
|
||||
pub async fn preclassify(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
args: &Value,
|
||||
context: &str, // 从 CLAUDE.md 和当前对话提取
|
||||
) -> ClassificationResult {
|
||||
// 构建精简 prompt：
|
||||
// "You are a security classifier. Evaluate this tool call:
|
||||
// Tool: {tool_name}
|
||||
// Args: {args}
|
||||
// Context: {context}
|
||||
// Respond: ALLOW|DENY|ASK <confidence 0-100> <reason>"
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**P1: PermissionRequest / PermissionDenied Hook 事件**
|
||||
|
||||
这两个事件对科研场景的审计至关重要：
|
||||
|
||||
```rust
|
||||
// 在 hooks/types.rs 中添加
|
||||
pub enum HookEvent {
|
||||
// ... 现有事件 ...
|
||||
/// 权限请求前触发（可阻止或修改）
|
||||
PermissionRequest,
|
||||
/// 权限被拒绝后触发（审计日志）
|
||||
PermissionDenied,
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 4. Hooks 系统
|
||||
|
||||
### 4.1 对比
|
||||
|
||||
| 特性 | Claude Code | AstroResearch (当前) |
|
||||
|------|-------------|---------------------|
|
||||
| Hook 类型 | 6 种：command, prompt, agent, http, callback, function | 2 种：sync AgentHook + async AsyncAgentHook |
|
||||
| 匹配器 | simple / pipe-separated / regex | glob + 精确匹配 |
|
||||
| if 条件 | `preparePermissionMatcher()` — Bash 上有 tree-sitter | 无 |
|
||||
| 输出协议 | JSON `{continue, decision, reason, suppressOutput, hookSpecificOutput}` | 直接返回值 |
|
||||
| 超时 | 每个 hook 独立超时（默认 10min） | 统一 `DEFAULT_HOOK_TIMEOUT` |
|
||||
| 事件数量 | 27 种 | ~15 种 |
|
||||
| 来源 | config + plugin + SDK + session + function | registry + session |
|
||||
|
||||
### 4.2 Claude Code 的 27 种 Hook 事件
|
||||
|
||||
```
|
||||
生命周期类:
|
||||
SessionStart, Setup, SubagentStart, SubagentStop, SessionEnd, Stop, StopFailure
|
||||
|
||||
用户交互类:
|
||||
UserPromptSubmit, Elicitation, ElicitationResult, PermissionRequest, PermissionDenied
|
||||
|
||||
工具执行类:
|
||||
PreToolUse, PostToolUse, PostToolUseFailure
|
||||
|
||||
上下文类:
|
||||
PreCompact, PostCompact, InstructionsLoaded
|
||||
|
||||
环境监控类:
|
||||
FileChanged, CwdChanged, ConfigChange
|
||||
|
||||
Swarm/Team 类:
|
||||
TeammateIdle, TaskCreated, TaskCompleted
|
||||
|
||||
UI 类:
|
||||
Notification, StatusLine, FileSuggestion
|
||||
```
|
||||
|
||||
### 4.3 AstroResearch 缺失的关键事件
|
||||
|
||||
| 缺失事件 | 用途 | 优先级 | 状态 |
|
||||
|---------|------|--------|------|
|
||||
| `UserPromptSubmit` | 用户提交 prompt 前拦截（自动上下文注入） | P2 | ⏳ |
|
||||
| `Notification` | 长时间操作完成通知 | P1 | ⏳ |
|
||||
| `PermissionRequest` | 权限弹窗前触发 | P1 | ✅ 已实现 |
|
||||
| `PermissionDenied` | 权限被拒绝后记录审计 | P1 | ✅ 已实现 |
|
||||
| `PreCompact` | 上下文压缩前机会（保存重要信息） | P2 | ⏳ |
|
||||
| `PostCompact` | 上下文压缩后通知（更新外部状态） | P2 | ⏳ |
|
||||
|
||||
### 4.4 已实施改进（2026-06-22）
|
||||
|
||||
**✅ P1: PermissionRequest / PermissionDenied 事件** — 已完成
|
||||
|
||||
新增类型（`src/agent/hooks/types.rs`）:
|
||||
- `HookEvent::PermissionRequest` / `HookEvent::PermissionDenied`
|
||||
- `PermissionRequestContext` — 携带 `current_decision`、`permission_mode`、`is_subagent`
|
||||
- `PermissionDeniedContext` — 携带 `reason`、`source` (Rule/Classifier/User/Timeout)
|
||||
- `PermissionRequestAction` — Continue / Override / InjectContext
|
||||
- `PermissionDecision` / `PermissionDenialSource` 枚举
|
||||
|
||||
新增 trait 方法（`src/agent/hooks/traits.rs`）:
|
||||
- `AgentHook::on_permission_request()` → `PermissionRequestAction`
|
||||
- `AgentHook::on_permission_denied()` → void (审计日志)
|
||||
|
||||
新增调度方法（`src/agent/hooks/dispatch.rs`）:
|
||||
- `HookRegistry::run_on_permission_request()` — 并行调用，第一个 Override 生效
|
||||
- `HookRegistry::run_on_permission_denied()` — fire-and-forget 审计
|
||||
|
||||
### 4.5 原始建议（部分已过时）
|
||||
|
||||
原建议的 context 设计已被更完善的实现替代：
|
||||
}
|
||||
```
|
||||
|
||||
**P2: 支持 command 类型 hook**
|
||||
|
||||
```rust
|
||||
/// Command 类型 hook — 执行外部 shell 命令处理事件
|
||||
pub struct CommandHook {
|
||||
command: String, // 如 "python3 audit.py"
|
||||
timeout: Duration, // 默认 10min
|
||||
shell: HookShell, // Bash | PowerShell
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 5. Error Recovery / 重试系统
|
||||
|
||||
### 5.1 对比
|
||||
|
||||
| 特性 | Claude Code | AstroResearch (当前) |
|
||||
|------|-------------|---------------------|
|
||||
| 重试结构 | Generator 模式，yield 系统消息直到不可重试 | ErrorRecovery 枚举 + 简单决策 |
|
||||
| 退避算法 | 指数 + 25% jitter，可配置上限（32s 默认，5min 持久） | 固定退避 |
|
||||
| 错误分类 | `shouldRetry()` 检查 15+ 种错误，每种不同策略 | 8 步分类管线 |
|
||||
| 529 Overloaded | 3 次重试 → 模型回退（Opus→Sonnet）→ 持久化重试 | 无模型回退 |
|
||||
| Context Overflow | 解析 "X + Y > Z"，自动调整 max_tokens + 1000 安全缓冲 | 只分类不修复 |
|
||||
| 持久化重试 | 无限重试 + 30s 心跳 | 无 |
|
||||
| Fast Cooldown | 429/529 在 fast mode → retry-after <20s → 10min cooldown | N/A |
|
||||
|
||||
### 5.2 Context Overflow 自动修复（最值得借鉴）
|
||||
|
||||
Claude Code 的做法：
|
||||
|
||||
```typescript
|
||||
// src/services/api/withRetry.ts
|
||||
// 解析 Anthropic API 的错误消息:
|
||||
// "input length and max_tokens exceed context limit: 180000 + 32000 > 200000"
|
||||
// → 计算安全的 max_tokens = 200000 - 180000 - 1000(safety) = 19000
|
||||
|
||||
function parseContextOverflowError(errorMessage: string) {
|
||||
const match = errorMessage.match(
|
||||
/input length and max_tokens exceed context limit: (\d+) \+ (\d+) > (\d+)/
|
||||
);
|
||||
if (match) {
|
||||
const [, inputLen, maxTokens, contextLimit] = match.map(Number);
|
||||
const newMaxTokens = contextLimit - inputLen - SAFETY_MARGIN;
|
||||
if (newMaxTokens > MIN_TOKENS) {
|
||||
return { shouldRetry: true, adjustedMaxTokens: newMaxTokens };
|
||||
}
|
||||
}
|
||||
return { shouldRetry: false };
|
||||
}
|
||||
```
|
||||
|
||||
### 5.3 已实施改进（2026-06-22）
|
||||
|
||||
**✅ P0: Context Overflow 自动修复** — 已完成
|
||||
|
||||
新增公共 API（`src/agent/runtime/error_recovery.rs`）:
|
||||
- `ContextOverflowInfo` — 从错误消息解析的数值结构体
|
||||
- `parse_context_overflow()` — 支持 Anthropic/OpenAI/通用三种格式
|
||||
- `calculate_safe_max_tokens()` — `context_limit - input_length - SAFETY_MARGIN(1000)`
|
||||
- `RecoveryStep::AdjustMaxTokens { new_max_tokens }` — 恢复管线第 0 步
|
||||
- `extract_three_numbers()` — 正则匹配 "A + B > C" 模式（使用已有 `regex` crate）
|
||||
|
||||
`AgentRuntime` 集成（`src/agent/runtime/mod.rs`）:
|
||||
```rust
|
||||
let overflow_info = error_recovery::parse_context_overflow(&e_str);
|
||||
while let Some(recovery_step) = recovery.try_recover(&error_kind, overflow_info.as_ref()) {
|
||||
// AdjustMaxTokens 优先于 AggressiveCompact，仅无空间时才回退到压缩
|
||||
}
|
||||
```
|
||||
|
||||
12 个新增测试覆盖 Anthropic/OpenAI/Generic 格式、边界条件、恢复优先级。
|
||||
|
||||
### 5.4 原始建议（部分已过时）
|
||||
|
||||
**P1: 模型回退策略**
|
||||
|
||||
```rust
|
||||
/// 模型回退链 — 529/Overloaded 时自动降级
|
||||
pub struct ModelFallback {
|
||||
chain: Vec<ModelTier>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ModelTier {
|
||||
Primary(String), // 如 "claude-opus-4-8"
|
||||
Fallback(String), // 如 "claude-sonnet-4-6"
|
||||
Emergency(String), // 如 "claude-haiku-4-5"
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 6. Tool 定义系统
|
||||
|
||||
### 7.1 对比
|
||||
|
||||
| 特性 | Claude Code `Tool<T>` | AstroResearch `AgentTool` trait |
|
||||
|------|----------------------|-------------------------------|
|
||||
| 并发安全 | `isConcurrencySafe(input)` — 输入感知 ✅ | `is_concurrency_safe(args)` — ✅ |
|
||||
| 语义标记 | `isReadOnly()` / `isDestructive()` / `isConcurrencySafe()` | `causes_sibling_abort()` / `interrupt_behavior()` |
|
||||
| 权限逻辑 | `checkPermissions()` — 工具自己决定权限 | 集中在 PermissionChecker |
|
||||
| 分类器输入 | `toAutoClassifierInput()` — 精简信息 | 无 |
|
||||
| 延迟加载 | `shouldDefer` / `alwaysLoad` — 减小 prompt | 无 |
|
||||
| 搜索提示 | `searchHint` — 帮助 ToolSearch 匹配 | 无 |
|
||||
| UI 渲染 | `renderToolUseMessage/Result/Progress/Error` (6 种) | 前端独立处理 |
|
||||
| 中断行为 | `interruptBehavior()` — cancel vs block | `interrupt_behavior()` ✅ |
|
||||
| 输出大小 | `maxResultSizeChars` — 超限存磁盘 | env var 全局配置 |
|
||||
| 权限匹配器 | `preparePermissionMatcher()` — 工具级模式匹配 | HookMatcher |
|
||||
|
||||
### 7.2 Claude Code 的 `buildTool()` Factory
|
||||
|
||||
```typescript
|
||||
// 每个工具通过 buildTool 创建，自动填充安全默认值
|
||||
const myTool: Tool = buildTool({
|
||||
name: 'MyTool',
|
||||
inputSchema: z.object({ ... }),
|
||||
async call(input, context, toolUseId) { ... },
|
||||
// 以下都有安全默认值，按需覆盖：
|
||||
// isEnabled: true (可根据 permission mode 禁用)
|
||||
// isConcurrencySafe: false
|
||||
// isReadOnly: false
|
||||
// isDestructive: false
|
||||
// checkPermissions: {behavior: 'allow'} (最宽松)
|
||||
// toAutoClassifierInput: '' (不参与分类)
|
||||
// shouldDefer: false (立即加载)
|
||||
// interruptBehavior: 'block' (不可中断)
|
||||
})
|
||||
```
|
||||
|
||||
### 7.3 建议改进
|
||||
|
||||
**P2: 为 AgentTool trait 添加方法**
|
||||
|
||||
```rust
|
||||
pub trait AgentTool: Send + Sync {
|
||||
// ... 现有方法 ...
|
||||
|
||||
/// 返回用于 auto-mode 分类器的精简摘要
|
||||
fn classifier_summary(&self, args: &Value) -> String {
|
||||
format!("{}", self.name())
|
||||
}
|
||||
|
||||
/// 是否是只读操作（与并发安全不同——glob 是 readonly 但不能和 bash 并发）
|
||||
fn is_readonly(&self) -> bool { false }
|
||||
|
||||
/// 是否需要延迟加载工具描述（大工具可延迟以减小 prompt）
|
||||
fn defer_loading(&self) -> bool { false }
|
||||
}
|
||||
```
|
||||
|
||||
**P3: 工具 allow/deny 列表**
|
||||
|
||||
参考 Claude Code 的做法，AstroResearch 已有 `subagent_allowed_tools`，可扩展：
|
||||
|
||||
```rust
|
||||
// 为子代理/异步代理定义工具过滤策略
|
||||
pub struct ToolFilterPolicy {
|
||||
/// 全局禁止（不计代理类型）
|
||||
pub all_agent_disallowed: Vec<String>,
|
||||
/// 自定义代理禁用（不能 spawn 子代理 + 编辑文件的代理）
|
||||
pub custom_agent_disallowed: Vec<String>,
|
||||
/// 异步代理白名单（只读子集）
|
||||
pub async_agent_allowed: Vec<String>,
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 7. Memory 持久化
|
||||
|
||||
### 8.1 对比
|
||||
|
||||
AstroResearch 与 Claude Code 在 Memory 设计上高度相似（都是文件系统 frontmatter markdown +
|
||||
`MEMORY.md` 索引），差距很小。
|
||||
|
||||
Claude Code 多了：
|
||||
- **Team Memory**: 共享给团队的记忆（AstroResearch 有 team 系统但无 team memory）
|
||||
- **Session Memory compaction**: 压缩时自动通过 LLM 提取记忆到文件
|
||||
|
||||
### 8.2 建议改进
|
||||
|
||||
**P3: Session Memory Compaction**
|
||||
|
||||
压缩时自动提取记忆：
|
||||
|
||||
```rust
|
||||
/// 在 compact.rs 的 auto_compact 过程中提取 session memory
|
||||
pub async fn extract_session_memory(
|
||||
llm: &LlmClient,
|
||||
messages: &[ChatMessage],
|
||||
) -> Vec<MemoryExtraction> {
|
||||
// 用专门的小 prompt 让 LLM 从对话中提取可持久化的记忆
|
||||
// 返回候选记忆列表，由 MemoryManager 做 dedup + 衰减
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 8. Coordinator / Multi-Agent
|
||||
|
||||
### 9.1 Claude Code 的 Coordinator Mode
|
||||
|
||||
```
|
||||
Coordinator Mode 架构:
|
||||
┌─────────────────────────────────────────────────┐
|
||||
│ Coordinator (Coordinator System Prompt) │
|
||||
│ Tools: Agent, SendMessage, TaskStop, │
|
||||
│ SyntheticOutput (仅 4 个) │
|
||||
│ │
|
||||
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
|
||||
│ │ Worker 1 │ │ Worker 2 │ │ Worker 3 │ │
|
||||
│ │ (async) │ │ (async) │ │ (async) │ │
|
||||
│ │ standard │ │ standard │ │ standard │ │
|
||||
│ │ tools │ │ tools │ │ tools │ │
|
||||
│ └──────────┘ └──────────┘ └──────────┘ │
|
||||
│ │
|
||||
│ Worker 结果以 <task-notification> 返回 │
|
||||
│ Coordinator 做 Synthesis → 下一轮 Workers │
|
||||
└─────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### 9.2 关键设计点
|
||||
|
||||
1. **Coordinator 只看到 4 个工具**——它不能直接读文件或执行 bash，只能编排 Workers
|
||||
2. **Workers 全异步**——Coordinator 不等待，结果以 `<task-notification>` XML 注入
|
||||
3. **Continue-vs-Spawn 决策矩阵**：
|
||||
- 新任务与已有 Worker 上下文重叠 <30% → 创建新 Worker
|
||||
- 新任务是对已有 Worker 的跟进 → `SendMessage` 继续
|
||||
4. **Worker prompt 写法规范**：自包含（self-contained），包含完整 spec，明确交付物
|
||||
|
||||
### 9.3 AstroResearch 的现状
|
||||
|
||||
AstroResearch 有 `SubAgentTool` + `TeamManager`，但没有 Coordinator 的概念。Team
|
||||
是平级的（lead ↔ teammates），不是层级编排。
|
||||
|
||||
### 9.4 建议改进
|
||||
|
||||
**P2: Coordinator Mode 原型**
|
||||
|
||||
```
|
||||
当 Agent 检测到复杂多步骤任务时，自动切换为 Coordinator 模式：
|
||||
|
||||
用户请求
|
||||
→ Coordinator 做任务分解
|
||||
→ 并行子代理执行 (SubAgentTool, async)
|
||||
→ 结果综合
|
||||
→ 减少单 Agent 的步骤数和 token 消耗
|
||||
|
||||
关键实现:
|
||||
1. Coordinator system prompt: 类似 Claude Code coordinatorMode.ts
|
||||
2. 仅暴露 SubAgentTool + 少量管理工具
|
||||
3. 子代理结果以结构化格式注入
|
||||
4. 综合阶段由 Coordinator 处理
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 9. Hermes-Agent 的独特贡献
|
||||
|
||||
### 10.1 文件位置
|
||||
|
||||
`/home/fmq/program/AstroResearch/libs/hermes-agent/`
|
||||
|
||||
### 10.2 Hermes-Agent vs Claude Code 架构对比
|
||||
|
||||
| 关注点 | Hermes-Agent | Claude Code |
|
||||
|--------|-------------|-------------|
|
||||
| 语言 | Python 3.11+ | TypeScript (Node.js 20+) |
|
||||
| Agent 循环 | 同步 `while` 循环 | 异步 `query()` generator |
|
||||
| 工具注册 | 文件系统自动发现 (`tools/*.py`) | 手动 import + `getAllBaseTools()` |
|
||||
| 工具接口 | `handler(args) -> JSON string` | `Tool<T>` — ~70 方法 |
|
||||
| 状态管理 | SQLite (SessionDB + FTS5) | 内存 `AppState` React store |
|
||||
| Plugin 系统 | PluginManager + 生命周期 hooks | Plugin loader + MCP 集成 |
|
||||
| Profile 隔离 | 多 profile + 独立 `HERMES_HOME` | 单 profile |
|
||||
| 子代理 | 子 `AIAgent` 实例 | `LocalAgentTask` + `runAgent()` |
|
||||
| Swarm/Team | Kanban 工作队列 | `InProcessTeammateTask` + coordinator |
|
||||
| 上下文压缩 | `ContextCompressor` | `compact/` 服务 |
|
||||
| MCP 支持 | `mcp_tool.py` + catalog | 完整的 `services/mcp/` |
|
||||
| 定位 | 个人 AI 助手（自我改进、记忆、跨平台） | 编程 AI 助手（终端集成、文件操作） |
|
||||
|
||||
### 10.3 Hermes-Agent 的核心架构
|
||||
|
||||
```
|
||||
hermes-agent/
|
||||
├── run_agent.py # AIAgent 类 (~11k LOC) — 核心入口
|
||||
├── model_tools.py # 工具编排层 (~2.7k LOC)
|
||||
├── toolsets.py # 工具集定义
|
||||
├── hermes_state.py # SessionDB — SQLite + FTS5
|
||||
├── cli.py # HermesCLI 类 (~11k LOC)
|
||||
│
|
||||
├── agent/ # Agent 内部
|
||||
│ ├── conversation_loop.py # 主循环
|
||||
│ ├── turn_context.py # 每轮上下文 dataclass
|
||||
│ ├── system_prompt.py # 三层 System Prompt
|
||||
│ ├── prompt_builder.py # Prompt 片段构建
|
||||
│ ├── memory_manager.py # 记忆编排
|
||||
│ ├── context_compressor.py # 上下文压缩
|
||||
│ ├── tool_executor.py # 工具执行分发
|
||||
│ ├── tool_guardrails.py # 安全 guardrails
|
||||
│ └── curator.py # 后台 Skill 生命周期管理
|
||||
│
|
||||
├── tools/ # 工具实现（自动发现）
|
||||
│ ├── registry.py # ToolRegistry（discover_builtin_tools）
|
||||
│ ├── delegate_tool.py # 子代理 spawn
|
||||
│ ├── skills_tool.py # Skill 管理
|
||||
│ └── cronjob_tools.py # Cron 调度
|
||||
│
|
||||
└── hermes_cli/ # CLI 子系统
|
||||
├── plugins.py # PluginManager + 生命周期 hooks
|
||||
└── profiles.py # 多 Profile 隔离
|
||||
```
|
||||
|
||||
### 10.4 最值得借鉴的 Hermes 特性
|
||||
|
||||
**Self-improving Skills**
|
||||
|
||||
```
|
||||
工作流:
|
||||
1. Agent 在科研中反复使用某流程
|
||||
(如: 搜索某类天体 → 下载论文 → RAG → 总结)
|
||||
2. Agent 自动检测重复模式
|
||||
3. Agent 创建 Skill 保存该流程
|
||||
4. Curator 管理 Skill 生命周期
|
||||
5. 下次相似查询直接加载 Skill，无需重新探索
|
||||
|
||||
AstroResearch 已有 Skills 系统 (skills.rs)，缺少:
|
||||
- 自动检测重复模式
|
||||
- Agent 自主创建 Skill
|
||||
- Curator 管理 Skill 质量
|
||||
```
|
||||
|
||||
**Session FTS5 搜索**
|
||||
|
||||
```python
|
||||
# Hermes 的 SessionDB 使用 SQLite FTS5 实现跨 session 搜索
|
||||
class SessionDB:
|
||||
def search_sessions(self, query: str) -> List[Session]:
|
||||
"""全文本搜索所有历史会话"""
|
||||
return self.db.execute(
|
||||
"SELECT * FROM sessions WHERE sessions MATCH ?", (query,)
|
||||
)
|
||||
```
|
||||
|
||||
AstroResearch 的 `agent_sessions` 表有基本的 title/status 字段，但没有全文搜索。
|
||||
|
||||
### 10.5 建议改进
|
||||
|
||||
**P2: Self-improving Skills 原型**
|
||||
|
||||
```
|
||||
1. Pattern Detector (自动检测)
|
||||
- 监控 N 个 session 中的工具调用序列
|
||||
- 使用简单的子序列匹配识别重复模式
|
||||
- 阈值: 3 次相似序列 → 候选 Skill
|
||||
|
||||
2. Skill Creator (Agent 自主创建)
|
||||
- 将候选 Skill 展示给用户确认
|
||||
- 生成 SKILL.md 文件（含 when_to_use, steps）
|
||||
- 注册到 SkillRegistry
|
||||
|
||||
3. Curator (管理生命周期)
|
||||
- 跟踪 Skill 使用频率
|
||||
- 长时间未用的 Skill 标记为 stale
|
||||
- 提示用户审查或删除
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 10. 优先级排序 —— 建议实施路线
|
||||
|
||||
按价值/投入比排序：
|
||||
|
||||
| 优先级 | 改进项 | 来源 | 状态 | 价值 | 预估投入 | 依赖 |
|
||||
|--------|--------|------|------|------|---------|------|
|
||||
| **P0** | Context Overflow 自动修复 | Claude Code | ✅ | 减少 ~50% LLM 调用失败 | ~30 行 | 无 |
|
||||
| **P0** | StreamingExecutor 真正的流式调度 | Claude Code | ✅ | 减少 30-50% 工具执行延迟 | ~200 行 | 无 |
|
||||
| **P0** | Executor 集成分区器（批次串行/并行） | Claude Code | ✅ | 修复非安全工具错误并发 | ~300 行 | ToolPartitioner |
|
||||
| **P1** | 工具并发分区 `partition_tool_calls` | Claude Code | ✅ | 批量只读操作 3-5x 加速 | ~100 行 | 无 |
|
||||
| **P1** | PermissionRequest / PermissionDenied Hook 事件 | Claude Code | ✅ | 安全审计能力 | ~100 行 | 无 |
|
||||
| **P1** | Auto-mode Classifier（廉价模型预分类） | Claude Code | ⏳ | 消除 80%+ 权限弹窗 | ~300 行 | LLM client 支持 |
|
||||
| **P2** | Self-improving Skills（Agent 保存成功流程） | Hermes | ✅ | 科研场景独特价值 | ~500 行 | Skills 系统 |
|
||||
| **P2** | Coordinator Mode（层级多 Agent 编排） | Claude Code | ⏳ | 复杂任务效果提升 | ~800 行 | SubAgent + Team |
|
||||
| **P2** | UserPromptSubmit / PreCompact / PostCompact Hook | Claude Code | ⏳ | Hook 系统完善 | ~200 行 | 无 |
|
||||
| **P3** | FTS5 跨 session 搜索 | Hermes | ⏳ | 历史研究可复用 | 中 | SQLite 迁移 |
|
||||
| **P3** | Tool `defer_loading` / `classifier_summary` | Claude Code | ⏳ | 减小 tool schema prompt | ~50 行 | 无 |
|
||||
| **P3** | 模型回退策略 (Model Fallback) | Claude Code | ⏳ | 提高可用性 | ~200 行 | ErrorRecovery |
|
||||
| **P3** | Session Memory Compaction | Claude Code | ⏳ | 自动化记忆提取 | ~300 行 | MemoryManager + Compact |
|
||||
|
||||
### 实施进度（2026-06-22）
|
||||
|
||||
1. **P0 项 — 全部完成** ✅
|
||||
- Context Overflow 自动修复（`error_recovery.rs`）
|
||||
- StreamingExecutor 真正流式调度（`streaming_executor.rs` 重写）
|
||||
- Executor 集成分区器（`executor.rs` Phase 3 重写）
|
||||
2. **P1 项 — 部分完成**
|
||||
- ✅ 工具并发分区
|
||||
- ✅ PermissionRequest / PermissionDenied Hook 事件
|
||||
- ⏳ Auto-mode Classifier — 需要设计讨论
|
||||
3. **P2 项在下一个大版本规划**：需要设计讨论和更多测试
|
||||
4. **P3 项作为 backlog**：长期优化方向
|
||||
|
||||
---
|
||||
|
||||
## 附录 A: Claude Code 关键源码索引
|
||||
|
||||
| 文件 | 用途 | 与 AstroResearch 对应 |
|
||||
|------|------|----------------------|
|
||||
| `src/Tool.ts` | Tool 类型 + buildTool factory | `src/agent/tools/mod.rs` |
|
||||
| `src/tools.ts` | 工具注册 + assembleToolPool | `ToolRegistry::new()` |
|
||||
| `src/services/tools/toolExecution.ts` | 工具执行管道 | `executor.rs` |
|
||||
| `src/services/tools/toolOrchestration.ts` | 并发分区 + 批量执行 | `partitioner.rs` + `executor.rs` |
|
||||
| `src/services/tools/StreamingToolExecutor.ts` | 流式工具执行 | `streaming_executor.rs` |
|
||||
| `src/services/api/withRetry.ts` | 错误重试 + 退避 | `error_recovery.rs` |
|
||||
| `src/services/api/claude.ts` | API 流式调用 | `streaming.rs` |
|
||||
| `src/constants/prompts.ts` | System Prompt 构建 | `system_prompt.rs` |
|
||||
| `src/utils/hooks.ts` | Hook 执行引擎 (5022 行) | `hooks/dispatch.rs` |
|
||||
| `src/utils/permissions/permissions.ts` | 权限逻辑 | `permission.rs` |
|
||||
| `src/utils/permissions/yoloClassifier.ts` | Auto-mode 分类器 | 无 (建议新增) |
|
||||
| `src/tools/AgentTool/runAgent.ts` | 子代理执行引擎 | `subagent.rs` |
|
||||
| `src/tools/AgentTool/AgentTool.tsx` | 子代理编排 | `tools/subagent.rs` |
|
||||
| `src/coordinator/coordinatorMode.ts` | Coordinator 模式 | 无 (建议参考) |
|
||||
| `src/memdir/memdir.ts` | Memory 文件系统 | `memory/mod.rs` |
|
||||
| `src/skills/loadSkillsDir.ts` | Skill 加载 | `skills.rs` |
|
||||
|
||||
## 附录 B: Hermes-Agent 关键源码索引
|
||||
|
||||
| 文件 | 用途 | 与 AstroResearch 对应 |
|
||||
|------|------|----------------------|
|
||||
| `run_agent.py` | AIAgent 核心类 | `runtime/mod.rs` |
|
||||
| `agent/conversation_loop.py` | Agent 主循环 | `runtime/mod.rs` (ReAct loop) |
|
||||
| `agent/system_prompt.py` | 三层 System Prompt | `system_prompt.rs` |
|
||||
| `agent/context_compressor.py` | 上下文压缩 | `compact.rs` |
|
||||
| `agent/memory_manager.py` | 记忆编排 | `memory/mod.rs` |
|
||||
| `agent/curator.py` | Skill 生命周期管理 | 无 (建议参考) |
|
||||
| `tools/registry.py` | 工具自动发现 | `tools/mod.rs` |
|
||||
| `tools/delegate_tool.py` | 子代理 | `subagent.rs` |
|
||||
| `hermes_state.py` | SessionDB + FTS5 | `api/agent.rs` (sessions) |
|
||||
| `hermes_cli/plugins.py` | PluginManager | 无 |
|
||||
| `hermes_cli/profiles.py` | 多 Profile 隔离 | 无 |
|
||||
| `model_tools.py` | 工具编排 | `executor.rs` |
|
||||
|
||||
## 附录 C: 变更日志
|
||||
|
||||
### 2026-06-22 — 首轮实施
|
||||
|
||||
**P0: Context Overflow 自动修复**
|
||||
- `src/agent/runtime/error_recovery.rs`: +150 行
|
||||
- 新增 `ContextOverflowInfo`、`parse_context_overflow()`、`calculate_safe_max_tokens()`
|
||||
- 新增 `RecoveryStep::AdjustMaxTokens`、`extract_three_numbers()`
|
||||
- 12 个新增测试（Anthropic/OpenAI/Generic 格式 + 边界 + 恢复优先级）
|
||||
- `src/agent/runtime/mod.rs`: AgentRuntime 集成调用点
|
||||
- `Cargo.lock`: 无新增依赖（使用已有 `regex` crate）
|
||||
|
||||
**P0: StreamingExecutor 真正流式调度**
|
||||
- `src/agent/runtime/streaming_executor.rs`: 重写 ~400 行（原 316 行）
|
||||
- `on_tool_use` 中对并发安全工具立即 `tokio::spawn`
|
||||
- `tokio::select!` 在工具执行和 Sibling Abort 之间竞速
|
||||
- `collect_completed_tasks()` 使用 `JoinHandle::is_finished()` 非阻塞检查
|
||||
- `executing_non_concurrent` 标志实现独占执行
|
||||
- `src/agent/tools/mod.rs`: `ToolContext` 添加 `Clone` derive
|
||||
|
||||
**P0: Executor 集成分区器**
|
||||
- `src/agent/runtime/executor.rs`: Phase 3 重写 + 2 个提取函数
|
||||
- Phase 3a: 构建非拒绝工具的 (原索引, PreparedCall) 映射
|
||||
- Phase 3b: `ToolPartitioner::partition()` 分区
|
||||
- Phase 3c: 逐批次执行（并行批次 `FuturesUnordered`，串行批次逐个执行）
|
||||
- 提取 `execute_single_tool()` 和 `process_single_result()` 辅助函数
|
||||
- `src/agent/runtime/partitioner.rs`: +2 测试（`run_bash`、`file_write` 打断并发批）
|
||||
|
||||
**P1: PermissionRequest / PermissionDenied Hook 事件**
|
||||
- `src/agent/hooks/types.rs`: +80 行
|
||||
- 新增 `HookEvent::PermissionRequest`、`HookEvent::PermissionDenied`
|
||||
- 新增 `PermissionRequestContext`、`PermissionDeniedContext`
|
||||
- 新增 `PermissionRequestAction`、`PermissionDecision`、`PermissionDenialSource`
|
||||
- `src/agent/hooks/traits.rs`: +20 行
|
||||
- `AgentHook::on_permission_request()`、`AgentHook::on_permission_denied()`
|
||||
- `src/agent/hooks/dispatch.rs`: +120 行
|
||||
- `run_on_permission_request()` (并行调用，第一个 Override 生效)
|
||||
- `run_on_permission_denied()` (fire-and-forget 审计)
|
||||
- `src/agent/hooks/mod.rs`: 更新 re-exports
|
||||
|
||||
### 2026-06-22 — Self-improving Skills
|
||||
|
||||
**P2: Self-improving Skills（模式检测 + 自动创建 + Curator）**
|
||||
- `src/agent/skills/pattern_detector.rs`: +420 行
|
||||
- `PatternDetector` — 从 `agent_messages` DB 表扫描工具调用序列
|
||||
- `DetectedPattern` — 跨 session 重复模式（携带出现次数、置信度、指纹）
|
||||
- 滑动窗口子序列提取 + Jaccard 相似度去重 + 超序列包含检测
|
||||
- 8 个单元测试（子序列提取、指纹、去重、Jaccard 计算）
|
||||
- `src/agent/skills/curator.rs`: +700 行
|
||||
- `Curator` — Skill 生命周期管理（Active → Inactive → Stale → Deprecated）
|
||||
- `SkillQuality` — 基于调用次数和新鲜度的质量评分（对数 + 指数衰减）
|
||||
- `CuratorReport` — 分析报告 + 清理建议列表
|
||||
- `archive_stale_skills()` — 将过期 skill 移动到归档目录
|
||||
- `CuratorRunner` — 后台空闲触发审查（`should_run_now()` + `record_activity()` 心跳）
|
||||
- 12 个单元测试（5 种生命周期状态 + 质量评分 + 清理候选 + pinned + runner）
|
||||
- `src/agent/skills.rs`: +230 行
|
||||
- `SkillCreator` — 将 `DetectedPattern` 转为 SKILL.md 文件（kebab-case 命名 + YAML frontmatter）
|
||||
- `SelfImprovePipeline` — 一站式管道（检测 → 创建 → 质量审查）
|
||||
- `SelfImproveResult` — 管道结果（patterns_found + skills_created + curator_report）
|
||||
- `SkillFrontmatter` + `SkillMeta` 增加 `pinned` 字段
|
||||
|
||||
### 2026-06-22 — Hermes Curator 特性补齐
|
||||
|
||||
**Pinned Skills（不可清理）**
|
||||
- `src/agent/skills.rs`: `SkillFrontmatter.pinned` + `SkillMeta.pinned` — YAML `pinned: true`
|
||||
- `src/agent/skills/curator.rs`: `evaluate_quality(pinned)` — 强制 Active + min_score 0.8；`analyze` 过滤 pinned 不进入 cleanup_candidates
|
||||
|
||||
**Seed Record（新 Skill 锚定时钟）**
|
||||
- `src/agent/skills/curator.rs`: `evaluate_quality` 对无统计记录 skill 设置 `days_since_last_use=Some(0)`（等效刚创建），`NEW_SKILL_GRACE_PERIOD_DAYS=7` 防止立即 stale
|
||||
|
||||
**CuratorRunner（后台空闲触发审查）**
|
||||
- `src/agent/skills/curator.rs`: `CuratorRunner` 结构体 + `should_run_now()`（paused/idle/interval 三重检查）+ `record_activity()` 心跳 + `run_once()` + `spawn()` tokio 后台任务 + `pause()`/`resume()`
|
||||
- 5 个新增测试（pinned_always_active、pinned_not_in_cleanup、seed_record、runner_paused、runner_idle）
|
||||
+822
-175
File diff suppressed because it is too large
Load Diff
@@ -133,6 +133,29 @@ graph TD
|
||||
└────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### 多源权限优先级 (Multi-Source Permission Precedence)
|
||||
|
||||
当多个来源（Checker 规则、Hook、工具级规则、会话规则）同时做出权限决策时，
|
||||
`resolve_permission_precedence()` 按以下优先级裁决：
|
||||
|
||||
| Priority | Source | Description |
|
||||
|----------|--------|-------------|
|
||||
| **P0** (highest) | `PermissionChecker::Deny` | 环境变量/配置文件配置的 Deny 规则，不可覆盖 |
|
||||
| **P1** | Tool-level `PermissionRule::Deny` | 工具自身拒绝执行（如 Bash 危险命令） |
|
||||
| **P2** | Session-level `PermissionChecker::Deny` | API 动态添加的会话级 Deny |
|
||||
| **P3** | Hook `PreToolUseAction::Block` | Hook 主动阻止工具执行 |
|
||||
| **P4** | Hook `PermissionRequired` | 仅当 Checker 返回 Allowed 时升级为 AskUser |
|
||||
| **P5** | Session-level `PermissionChecker::Ask` | 仅当当前为 Allowed 时升级 |
|
||||
| **P6** | Tool-level `PermissionRule::Ask` | 仅当当前为 Allowed 时升级 |
|
||||
| **P7** | `PermissionChecker::Allow` | 显式 Allow 规则 |
|
||||
| **P8** (lowest) | Implicit Allow (default) | 无任何规则匹配 → 允许 |
|
||||
|
||||
**关键规则：**
|
||||
- **Deny 不可覆盖**: P0-P2 的 Deny 规则在任何情况下生效
|
||||
- **Ask 可升级**: P4-P6 在 Allow 状态下升级为 Ask；在 Deny 状态下被忽略
|
||||
- **Block = Deny**: Hook Block 等同于 Deny，由 P0-P2 可覆盖
|
||||
- **冲突日志**: `conflict_log` 记录所有被覆盖的决策，用于审计```
|
||||
|
||||
### PermissionChecker API
|
||||
|
||||
```rust
|
||||
@@ -433,7 +456,7 @@ flowchart TD
|
||||
| `grep "ssh_config" *.rs` | ❌ 误拦（含子串 `ssh `） | ✅ 允许（首词 `grep` 在白名单） |
|
||||
| `echo "use sudo carefully"` | ❌ 误拦（含子串 `sudo `） | ✅ 允许（首词 `echo` 在白名单） |
|
||||
| `cat /usr/share/vim/vimrc` | ❌ 误拦（含子串 `vim `） | ✅ 允许（首词 `cat` 在白名单） |
|
||||
| `python script.py` | ✅ 允许 | ✅ 允许（不在黑名单，默认允许） |
|
||||
| `python script.py` | ✅ 允许 | ⚠️ 通过校验，但触发 `Ask` 用户确认（不在白名单） |
|
||||
| `vim file.txt` | ✅ 拒绝 | ✅ 拒绝（首词 `vim` 在黑名单） |
|
||||
| `$(echo sud; echo o) /etc/passwd` | ✅ 允许（绕过！） | ❌ 拒绝（检测到命令替换绕过） |
|
||||
|
||||
@@ -449,7 +472,7 @@ const SAFE_COMMANDS: &[&str] = &[
|
||||
];
|
||||
```
|
||||
|
||||
白名单内的命令**优先放行**，不经过黑名单检查。非白名单命令经过黑名单精确匹配和危险参数二次检查后默认允许。
|
||||
白名单内的命令**优先放行**，不经过黑名单检查。非白名单命令经黑名单精确匹配和危险参数检查后，通过 `bash_needs_permission()` → `RunBashTool::check_permissions()` 返回 `Ask` 规则，触发用户确认弹窗（PermissionRequestCard）。
|
||||
|
||||
### 其他约束
|
||||
|
||||
@@ -462,7 +485,7 @@ const SAFE_COMMANDS: &[&str] = &[
|
||||
### 已知局限
|
||||
|
||||
1. ~~**黑名单子串匹配**~~ — ✅ 已修复：改用首词精确匹配，`grep "ssh_config"` 不再误拦
|
||||
2. **未限制网络访问** — `curl`、`wget` 不在黑名单中
|
||||
2. ~~**未限制网络访问**~~ — ✅ 已修复：`NETWORK_COMMANDS` 名单（curl/wget/nc/socat 等）默认阻止，`AGENT_BLOCK_NETWORK=false` 可放行
|
||||
3. **未限制进程数** — fork bomb（如 `:(){ :\|:& };:`）未被检测
|
||||
4. **管道/重定向完整放行** — `<`、`>`、`|` 不做限制
|
||||
5. ~~**`$()` 命令替换**~~ — ✅ 已修复：检测首词位置 `$()` 和反引号绕过
|
||||
@@ -684,258 +707,9 @@ Diminishing Returns 检测
|
||||
|
||||
---
|
||||
|
||||
## Claude Code 权限系统对比分析
|
||||
|
||||
> 对比基准：Claude Code (`/home/fmq/program/claudecode/src/utils/permissions/`)
|
||||
> 分析日期：2026-06-17
|
||||
|
||||
### 架构差异总览
|
||||
|
||||
| 维度 | AstroResearch (当前) | Claude Code (参考) | 差距 |
|
||||
|------|---------------------|-------------------|------|
|
||||
| 规则引擎 | ✅ PermissionChecker (完成) | ✅ hasPermissionsToUseTool 多步流水线 | 相当 |
|
||||
| 规则匹配粒度 | ✅ 内容级 `Tool(content*)` 前缀/后缀/包含 | ✅ 前缀/通配/内容级 / 正则 | 小 |
|
||||
| AskUser 交互流 | ✅ SSE → PermissionRequestCard → Allow/Deny/Always Allow | ✅ 完整 SSE → Dialog → 决策 | 相当 |
|
||||
| 权限模式 | ✅ Default/AcceptEdits/Bypass/DontAsk | ✅ 6种模式 (含 plan/auto) | 小 |
|
||||
| 规则持久化 | ✅ 环境变量加载 + `PermissionChecker::from_config()` | ✅ settings.json 多层加载 (8级来源优先级) | 小 |
|
||||
| 规则来源追踪 | ✅ `PermissionRuleSource` 枚举 (Env/Session) | ✅ cliArg > command > session > userSettings > ... | 小 |
|
||||
| Bash 权限分类器 | ✅ SAFE_COMMANDS 白名单 + `check_permissions()` 集成 | ✅ AST解析 + AI分类器 + 异步推测 | 中等 |
|
||||
| 拒绝追踪/熔断 | ✅ `DenialTracker` 连续/累计计数 + ReAct 循环熔断 | ✅ 连续/总计拒绝计数 + 自动终止 | 相当 |
|
||||
| 权限 Hook 集成 | ✅ PreToolUseAction::PermissionRequired 完整流程 | ✅ 完整 PermissionRequest hook + 多路径决议 | 相当 |
|
||||
| 规则遮蔽检测 | ✅ `detect_shadowed_rules()` deny/ask 双重检查 | ✅ `shadowedRuleDetection` deny/ask 遮蔽检测 | 相当 |
|
||||
| Auto Mode (AI 分类) | ❌ 无 | ✅ YOLO classifier + 快速路径 + 安全工具白名单 | **远期** |
|
||||
| 权限解释器 | ✅ 启发式 `explain_permission()` (Bash 风险等级 + 路径检测) | ✅ Haiku 生成风险解释 | 中等 |
|
||||
| 会话内规则更新 | ✅ `POST/PUT /api/chat/sessions/:id/permissions/*` | ✅ `/permissions` 命令 + API | 小 |
|
||||
| 子代理权限继承 | ✅ 完整 `check()` 三态检查 | ✅ 完整继承父级权限上下文 | 相当 |
|
||||
| 附加目录沙箱 | ✅ `AGENT_ADDITIONAL_DIRS` + `is_path_allowed()` 扩展 | ✅ `additionalDirectories` 可配置 | 相当 |
|
||||
|
||||
### Claude Code 权限流水线 (参考架构)
|
||||
|
||||
```
|
||||
hasPermissionsToUseTool(toolName, input, context):
|
||||
Step 1a: 工具级 deny 规则检查 → deny → 返回 deny
|
||||
Step 1b: 工具级 ask 规则检查 → ask → 返回 ask (sandbox 例外)
|
||||
Step 1c: 工具自定义 checkPermissions() → 内容级规则匹配
|
||||
Step 1d: 工具实现返回 deny → deny → 返回 deny
|
||||
Step 1e: requiresUserInteraction? → ask → 强制 ask (bypass 免疫)
|
||||
Step 1f: 内容级 ask 规则 → ask → 强制 ask (bypass 免疫)
|
||||
Step 1g: 安全检查 (敏感路径等) → ask → 强制 ask (bypass 免疫)
|
||||
Step 2a: bypassPermissions 模式? → allow → 返回 allow
|
||||
Step 2b: 工具级 allow 规则 → allow → 返回 allow
|
||||
Step 3: 剩余 passthrough → ask → 返回 ask
|
||||
|
||||
外层模式变换:
|
||||
dontAsk 模式: ask → deny
|
||||
auto 模式: acceptEdits 快速路径 → 安全工具白名单 → AI分类器
|
||||
headless: hooks 先运行 → 无 hook 决定 → auto-deny
|
||||
```
|
||||
|
||||
### 关键设计决策对比
|
||||
|
||||
**1. 规则格式**
|
||||
|
||||
Claude Code 使用 `ToolName(content)` 格式支持内容级规则：
|
||||
```
|
||||
Bash → 匹配所有 bash 命令
|
||||
Bash(npm install) → 匹配精确命令
|
||||
Bash(npm *) → 前缀通配
|
||||
Bash(rm:*) → 旧版前缀（已废弃）
|
||||
Read(.env) → 文件模式
|
||||
mcp__server__tool → MCP 工具级
|
||||
mcp__server → MCP 服务级
|
||||
Agent(Explore) → 代理类型级
|
||||
```
|
||||
|
||||
AstroResearch 已实现相同格式：
|
||||
```
|
||||
"*" → 通配所有工具
|
||||
"tool_name" → 精确工具名匹配
|
||||
"tool_name(content*)" → 前缀通配（如 "run_bash(rm *)" 匹配 "rm -rf /"）
|
||||
"tool_name(*suffix)" → 后缀通配（如 "read_file(*.env)" 匹配 ".env"）
|
||||
"tool_name(exact)" → 包含匹配（子串命中）
|
||||
"*(content)" → 工具通配 + 内容匹配（如 "*(sudo)" 匹配任意工具的 sudo 命令）
|
||||
```
|
||||
从 args 中自动提取 `command`/`file_path`/`path`/`pattern`/`url` 字段进行内容匹配。
|
||||
|
||||
**2. 权限模式**
|
||||
|
||||
Claude Code 的 6 种模式通过 Shift+Tab 循环切换：
|
||||
- `default` — 标准逐项确认
|
||||
- `acceptEdits` — 工作目录内文件编辑自动通过
|
||||
- `bypassPermissions` — 跳过所有 Ask（deny/ask 规则仍生效；安全检查 bypass 免疫）
|
||||
- `dontAsk` — 所有 Ask 转 Deny
|
||||
- `plan` — 计划模式
|
||||
- `auto` — AI 自动分类（内部使用）
|
||||
|
||||
AstroResearch 已实现 4 种模式（通过 `AGENT_PERMISSION_MODE` 环境变量或 API 切换）：
|
||||
- `default` — 标准规则链，Ask 触发用户交互
|
||||
- `acceptEdits` — 工作目录内文件编辑自动通过（路径检查由 executor 完成）
|
||||
- `bypassPermissions` — 跳过所有 Ask（Deny 规则仍生效）
|
||||
- `dontAsk` — 所有 Ask 转为 Deny
|
||||
|
||||
`PermissionChecker::from_config()` 从 `AgentConfig` 加载环境变量规则并构造完整检查器。
|
||||
|
||||
**3. 多路径权限决议**
|
||||
|
||||
Claude Code 的 AskUser 决议支持多个并行路径，任一先返回即生效（`claim()` 模式）：
|
||||
- 本地 UI 对话框
|
||||
- Bridge 响应（CCR 远程）
|
||||
- Channel 响应（Telegram 等）
|
||||
- PermissionRequest hooks（后台异步运行）
|
||||
- Bash 分类器（后台推测性异步分类）
|
||||
|
||||
AstroResearch 已实现完整的 AskUser 交互流：
|
||||
- `PermissionChecker::check()` 返回 `AskUser` 时，executor 发送 `AgentStreamEvent::PermissionRequest` SSE 事件
|
||||
- 通过 `oneshot` 通道创建 `PendingPermission`，存入 `AppState::pending_permissions`
|
||||
- 等待用户通过前端 `PermissionRequestCard` 组件响应（Allow / Deny / Always Allow），120s 超时自动拒绝
|
||||
- 单一路径决议（oneshot），不支持多路径 claim 模式
|
||||
|
||||
---
|
||||
|
||||
## 优化路线图
|
||||
|
||||
### ✅ P0 — 已全部完成
|
||||
|
||||
#### P0-1: 规则加载与持久化 ✅
|
||||
|
||||
`AgentConfig::from_env_optional()` 从环境变量加载规则（`AGENT_PERMISSIONS_DENY`/`ALLOW`/`ASK`），`PermissionChecker::from_config()` 按 Deny → Ask → Allow 优先级顺序构造规则链。
|
||||
|
||||
**实现位置**: `src/agent/runtime/mod.rs:113-117`, `src/agent/runtime/permission.rs:268-290`
|
||||
|
||||
#### P0-2: 完成 AskUser 权限交互流 ✅
|
||||
|
||||
executor Phase 2.5 中完整的 AskUser 处理：
|
||||
- `AgentStreamEvent::PermissionRequest` SSE 事件 → 前端 `PermissionRequestCard` 组件
|
||||
- `oneshot` 通道 + `AppState::pending_permissions` 存储
|
||||
- 120s 超时自动拒绝，支持 Allow / Deny / Always Allow 决策
|
||||
|
||||
**实现位置**: `src/agent/runtime/executor.rs:282-422`, `src/api/agent.rs`, `dashboard/src/features/agent/PermissionRequestCard.tsx`
|
||||
|
||||
#### P0-3: 内容级权限匹配 ✅
|
||||
|
||||
`PermissionChecker::matches()` 支持 `"tool_name(content_pattern)"` 格式，前缀通配（`prefix*`）、后缀通配（`*suffix`）、包含匹配，自动从 args 提取 `command`/`file_path`/`path`/`pattern`/`url` 字段。
|
||||
|
||||
**实现位置**: `src/agent/runtime/permission.rs:183-255`
|
||||
|
||||
### ✅ P1 — 已全部完成
|
||||
|
||||
#### P1-1: 权限模式系统 ✅
|
||||
|
||||
`PermissionMode` 枚举实现 4 种模式（Default/AcceptEdits/Bypass/DontAsk），通过 `AGENT_PERMISSION_MODE` 环境变量或 API 切换。`PermissionChecker::apply_mode()` 在 executor Phase 2.5 中对检查结果进行模式变换（Bypass 将 Ask→Allowed，DontAsk 将 Ask→Denied）。
|
||||
|
||||
**实现位置**: `src/agent/runtime/permission.rs:39-60, 139-177`
|
||||
|
||||
#### P1-2: Bash 权限接入 PermissionChecker ✅
|
||||
|
||||
`RunBashTool::check_permissions()` 调用 `bash_needs_permission()` —— 安全白名单中的命令返回空规则（自动允许），非白名单命令返回 `Ask` 规则。executor Phase 2.5 中与 PermissionChecker 结果叠加。
|
||||
|
||||
**实现位置**: `src/agent/tools/filesystem/bash.rs:58-73, 362-365`
|
||||
|
||||
#### P1-3: 权限 Hook 集成 ✅
|
||||
|
||||
`PreToolUseAction::PermissionRequired` 在 executor 中被检测：若 hook 返回 `PermissionRequired` 且 PermissionChecker 返回 `Allowed`，则升级为 `AskUser` 触发用户交互。已修复 Continue 覆盖 meaningful action 的 bug。
|
||||
|
||||
**实现位置**: `src/agent/hooks.rs:378-384`, `src/agent/runtime/executor.rs:221-230`
|
||||
|
||||
#### P1-4: 子代理完整权限继承 ✅
|
||||
|
||||
`SubAgentRunner` 使用 `check(tool_name, Some(&final_args))` 进行三态检查：Deny → 注入错误跳过执行，AskUser → 自动拒绝（子代理不应打断用户），Allowed → 正常执行。
|
||||
|
||||
**实现位置**: `src/agent/subagent.rs`
|
||||
|
||||
### ✅ P2-1、P2-2 — 已实现
|
||||
|
||||
#### P2-1: 拒绝追踪与熔断 ✅
|
||||
|
||||
`DenialTracker` 追踪连续拒绝和总拒绝数，阈值触发 ReAct 循环终止。
|
||||
配置：`AGENT_DENIAL_MAX_CONSECUTIVE` (默认 3) / `AGENT_DENIAL_MAX_TOTAL` (默认 20)。
|
||||
|
||||
**实现位置**: `src/agent/runtime/denial_tracker.rs`, `src/agent/runtime/mod.rs`
|
||||
|
||||
#### P2-2: 会话内规则更新 API ✅
|
||||
|
||||
`POST /api/chat/sessions/:id/permissions/rules` — add/remove 规则
|
||||
`PUT /api/chat/sessions/:id/permissions/mode` — 切换权限模式
|
||||
通过 `AppState::session_permission_checker` (`Arc<RwLock<PermissionChecker>>`) 实现跨 turn 共享。
|
||||
|
||||
**实现位置**: `src/api/permissions.rs`, `src/agent/runtime/executor.rs` Phase 2.5
|
||||
|
||||
### 🟡 P2-3~P2-5 — 远期增强（按需实现）
|
||||
|
||||
#### P2-3: 规则遮蔽检测 ✅
|
||||
|
||||
`PermissionChecker::detect_shadowed_rules()` 检测 Deny/Ask 遮蔽 Allow 的情况，输出 `ShadowedRule` 列表（含 reason + fix 建议），在 AgentRuntime 初始化时通过 `warn!` 日志输出。
|
||||
|
||||
**实现位置**: `src/agent/runtime/permission.rs`
|
||||
|
||||
#### P2-4: 权限解释器 ✅
|
||||
|
||||
启发式 `explain_permission()` 函数，根据工具名和参数生成 `{risk_level, explanation, reasoning, risk}` 结构。Bash 命令通过关键词检测风险等级（HIGH/MEDIUM/LOW），文件操作检测系统路径。结果随 `PermissionRequest` SSE 事件推送到前端。
|
||||
|
||||
**实现位置**: `src/agent/runtime/permission_explainer.rs`, `src/agent/runtime/mod.rs` `AgentStreamEvent::PermissionRequest.explanation`
|
||||
|
||||
#### P2-5: Auto Mode (AI 权限分类器) ❌
|
||||
|
||||
使用 LLM 自动评估工具调用的风险：
|
||||
- 快速路径：`AcceptEdits` 模式自动允许工作目录内的文件编辑
|
||||
- 安全工具白名单：`read_file`、`grep_files`、`search_papers` 等只读操作自动允许
|
||||
- AI 分类：对不确定的操作调用快速模型判断安全性
|
||||
- 失败封闭：分类器不可用时拒绝所有非白名单操作（安全优先）
|
||||
|
||||
**工作量**: 3-5天
|
||||
|
||||
#### P2-4: 权限解释器
|
||||
|
||||
在执行前用 LLM 生成人类可读的风险描述：
|
||||
```
|
||||
"该命令将执行 npm install，可能修改 node_modules/ 目录并下载外部依赖包。"
|
||||
```
|
||||
|
||||
**工作量**: 1天
|
||||
|
||||
#### P2-5: 子代理最小权限 (ToolRegistry::restrict)
|
||||
|
||||
```rust
|
||||
impl ToolRegistry {
|
||||
pub fn restrict(&self, allowed_tools: &[&str]) -> Self {
|
||||
// 创建仅包含指定工具的受限注册表
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**工作量**: 0.5天
|
||||
|
||||
---
|
||||
|
||||
## 实现路线图
|
||||
|
||||
```
|
||||
已完成 (Phase 1): P0-1 规则加载 + P0-2 AskUser 交互流 + P0-3 内容级匹配
|
||||
已完成 (Phase 2): P1-1 权限模式 + P1-2 Bash 集成 + P1-3 Hook 集成 + P1-4 子代理继承
|
||||
已完成 (Phase 3): P2-1 拒绝追踪熔断 + P2-2 会话内规则更新 + P2-3 规则遮蔽检测 + P2-4 权限解释器 + 附加目录沙箱 + 规则来源追踪
|
||||
远期规划 (按需): P2-5 Auto Mode (AI 分类器) + P2-6 权限解释器 LLM 升级 + 子代理最小权限 + 文件写入大小限制
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 待完善项
|
||||
|
||||
| 优先级 | 项目 | 当前状态 | 建议 |
|
||||
|---|---|---|---|
|
||||
| ~~**HIGH**~~ | ~~PermissionChecker 集成到主执行路径~~ | ✅ **已完成** | — |
|
||||
| ~~**HIGH**~~ | ~~Bash 黑名单改为命令解析~~ | ✅ **已完成** | — |
|
||||
| ~~**P0**~~ | ~~规则加载与持久化~~ | ✅ **已完成**：`AgentConfig` 新增 `permission_deny_rules` / `permission_allow_rules` / `permission_ask_rules` / `permission_mode` 字段，通过 `AGENT_PERMISSIONS_*` 环境变量加载 | — |
|
||||
| ~~**P0**~~ | ~~AskUser 权限交互流~~ | ✅ **已完成**：executor Phase 2.5 AskUser 分支重写为完整 oneshot → SSE → 120s 超时流程。前端 `PermissionRequestCard` 组件提供 Allow / Deny / Always Allow | — |
|
||||
| ~~**P0**~~ | ~~内容级权限匹配~~ | ✅ **已完成**：`check(tool_name, tool_args)` 签名，`matches()` 支持 `"tool(content*)"` 格式（前缀/后缀/包含），自动提取 args 字段 | — |
|
||||
| ~~**P1**~~ | ~~权限模式系统~~ | ✅ **已完成**：`PermissionMode` (Default/AcceptEdits/Bypass/DontAsk)，`apply_mode()` 方法，`AGENT_PERMISSION_MODE` 配置 | — |
|
||||
| ~~**P1**~~ | ~~Bash 权限集成~~ | ✅ **已完成**：`RunBashTool::check_permissions()` 调用 `bash_needs_permission()`（安全命令自动允许），executor 合并工具级检查 | — |
|
||||
| ~~**P1**~~ | ~~权限 Hook 集成~~ | ✅ **已完成**：`PreToolUseResult::is_permission_required()`，修复 Continue 覆盖 bug，executor 触发 AskUser | — |
|
||||
| ~~**P1**~~ | ~~子代理完整权限继承~~ | ✅ **已完成**：`is_denied()` → `check(tool_name, Some(&final_args))`，子代理中 AskUser 自动拒绝 | — |
|
||||
| ~~**MEDIUM**~~ | ~~拒绝追踪与熔断~~ | ✅ **已完成** | `DenialTracker`：连续/总计拒绝计数，阈值触发 ReAct 循环终止 |
|
||||
| ~~**MEDIUM**~~ | ~~会话内规则更新 API~~ | ✅ **已完成** | `POST/PUT /api/chat/sessions/:id/permissions/*` 动态 add/remove/mode |
|
||||
| ~~**MEDIUM**~~ | ~~权限解释器~~ | ✅ **已完成** | 启发式 `explain_permission()`，Bash 风险等级 + 路径检测，随 SSE PermissionRequest 推送前端 |
|
||||
| **MEDIUM** | Auto Mode (AI 分类器) | 未实现 | LLM 评估风险，快速路径 + 安全工具白名单 |
|
||||
| **LOW** | 子代理最小权限 | 继承全部父工具 | `ToolRegistry::restrict()` |
|
||||
| **LOW** | 文件写入大小限制 | 无上限 | 添加 `max_file_size` 参数 |
|
||||
| **LOW** | 网络访问控制 | `curl`/`wget` 未限制 | Bash 黑名单扩展 |
|
||||
| **LOW** | 用户权限 profiles | 不支持 | YAML/TOML 权限配置 |
|
||||
|
||||
@@ -42,7 +42,9 @@ sequenceDiagram
|
||||
|
||||
| 文件 | 行数 | 职责 |
|
||||
|---|---|---|
|
||||
| `src/agent/skills.rs` | 847 | SkillRegistry 缓存、文件解析、热更新、条件激活、系统提示构建 |
|
||||
| `src/agent/skills.rs` | ~1100 | SkillRegistry 缓存、文件解析、热更新、条件激活、系统提示构建、SkillCreator / SelfImprovePipeline |
|
||||
| `src/agent/skills/pattern_detector.rs` | ~420 | PatternDetector — 从 agent_messages 扫描工具调用序列、子序列匹配、Jaccard 去重 |
|
||||
| `src/agent/skills/curator.rs` | ~700 | Curator — Skill 生命周期管理 + CuratorRunner — 后台空闲触发审查 |
|
||||
| `src/agent/tools/skill.rs` | 222 | LoadSkillTool — Layer 2 按需加载的 AgentTool 实现 |
|
||||
| `src/agent/runtime/mod.rs` | ~1182 | 将 `build_reminder()` 注入 SystemPrompt section 3 |
|
||||
| `src/agent/runtime/system_prompt.rs` | ~64 | 静态 system prompt 中引导 LLM 使用 load_skill |
|
||||
@@ -95,6 +97,7 @@ effort: high
|
||||
| `paths` | `string[]` | `[]`（始终激活） | 条件激活的 glob 模式，非空时 skill 仅在匹配文件路径后激活 |
|
||||
| `agent` | `string` | — | fork 模式下游的 agent 类型（如 `code-reviewer`），已定义但 LoadSkillTool 尚未使用 |
|
||||
| `effort` | `string` | — | fork 模式下的 effort 级别，已定义但 LoadSkillTool 尚未使用 |
|
||||
| `pinned` | `bool` | `false` | `true` 时 Curator 强制保持 Active 生命周期，最低质量评分 0.8，不被自动清理 |
|
||||
|
||||
### 当前项目 Skill 清单
|
||||
|
||||
@@ -112,7 +115,8 @@ effort: high
|
||||
SkillFrontmatter — serde_yaml 解析的 YAML frontmatter，含 validate() 校验方法
|
||||
│
|
||||
├──▶ SkillMeta — Layer 1 摘要（name, description, context, allowed_tools,
|
||||
│ when_to_use, disable_model_invocation, user_invocable, paths）
|
||||
│ when_to_use, disable_model_invocation, user_invocable, paths,
|
||||
│ pinned: bool）
|
||||
│
|
||||
└──▶ Skill — Layer 2 完整对象（meta + body + skill_dir）
|
||||
│
|
||||
@@ -120,6 +124,12 @@ SkillFrontmatter — serde_yaml 解析的 YAML frontmatter，含 valida
|
||||
├── skills: Vec<Skill>
|
||||
├── last_scan_mtime: Option<SystemTime>
|
||||
└── usage_stats: HashMap<String, SkillUsageStat>
|
||||
|
||||
SelfImprovePipeline — 一站式管道
|
||||
├── PatternDetector — 扫描 agent_messages 检测重复工具调用序列
|
||||
├── SkillCreator — 将 DetectedPattern 转换为 SKILL.md 文件
|
||||
└── Curator — Skill 生命周期管理（Active→Inactive→Stale→Deprecated）
|
||||
└── CuratorRunner — 后台空闲触发审查 + 心跳记录
|
||||
```
|
||||
|
||||
### 关键方法
|
||||
@@ -319,12 +329,304 @@ Box::new(LoadSkillTool::new(skill_registry)),
|
||||
- 团队成员（`teammate.rs`）
|
||||
- 后台任务 Agent（`background.rs`）
|
||||
|
||||
## Self-improving Skills — 自我进化管道
|
||||
|
||||
参考 Hermes-Agent 的 Self-improving Skills 模式，AstroResearch 实现了从**模式检测 → 自动创建 → 生命周期管理**的完整自我进化管道。
|
||||
|
||||
### 架构总览
|
||||
|
||||
```mermaid
|
||||
graph TB
|
||||
subgraph Pipeline["SelfImprovePipeline::run()"]
|
||||
direction TB
|
||||
PD["PatternDetector::scan()<br/>扫描 agent_messages 表"]
|
||||
SC["SkillCreator::create_from_patterns()<br/>生成 SKILL.md 文件"]
|
||||
CR["Curator::analyze()<br/>质量评分 + 生命周期评估"]
|
||||
end
|
||||
|
||||
subgraph Background["后台定期维护"]
|
||||
direction LR
|
||||
Runner["CuratorRunner::spawn()<br/>空闲触发 + 间隔检查"]
|
||||
Archive["archive_stale_skills()<br/>30d stale / 90d deprecated"]
|
||||
end
|
||||
|
||||
PD -->|"Vec<DetectedPattern>"| SC
|
||||
SC -->|"Vec<Skill>"| CR
|
||||
CR -->|"CuratorReport"| Runner
|
||||
```
|
||||
|
||||
### PatternDetector — 模式检测器
|
||||
|
||||
从 `agent_messages` 表中自动发现跨 session 重复的工具调用序列：
|
||||
|
||||
```
|
||||
检测算法:
|
||||
1. 查询每个 session 的 tool 消息（按时间排序）
|
||||
2. 滑动窗口 (2-8 长度) 提取所有子序列
|
||||
3. 跨 session 频率计数（≥3 次为候选）
|
||||
4. Jaccard 相似度去重 + 超序列包含检测
|
||||
5. 计算 confidence = frequency_score × similarity_penalty
|
||||
```
|
||||
|
||||
**数据结构**：
|
||||
|
||||
```rust
|
||||
pub struct DetectedPattern {
|
||||
pub tool_sequence: Vec<String>, // 如 ["search_papers", "download_paper", "rag_search"]
|
||||
pub session_ids: Vec<String>, // 出现的 session
|
||||
pub frequency: usize, // 跨 session 出现次数
|
||||
pub confidence: f64, // 0.0-1.0 置信度
|
||||
pub fingerprint: String, // 去重指纹
|
||||
}
|
||||
```
|
||||
|
||||
**置信度计算**：
|
||||
```
|
||||
confidence = ln(frequency) / ln(3) × (1 - max_jaccard_similarity_with_other_patterns)
|
||||
```
|
||||
即：频率越高越好，与已有模式越不相似越好。
|
||||
|
||||
### SkillCreator — 自动 Skill 生成
|
||||
|
||||
将 `DetectedPattern` 转换为完整的 `SKILL.md` 文件：
|
||||
|
||||
```rust
|
||||
pub struct SkillCreator {
|
||||
skills_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl SkillCreator {
|
||||
/// 检测到的模式 → 写入 skills/{kebab-case-name}/SKILL.md
|
||||
pub fn create_from_patterns(
|
||||
&self,
|
||||
patterns: &[DetectedPattern],
|
||||
dry_run: bool, // dry_run=true 时只预览不移交
|
||||
) -> Result<Vec<Skill>, Error>;
|
||||
|
||||
/// 工具序列名 → kebab-case skill 名称
|
||||
fn pattern_to_skill_name(tools: &[String]) -> String;
|
||||
// 例: ["search_papers", "download_paper", "rag_search"] → "search-download-rag"
|
||||
|
||||
/// 生成 SKILL.md 正文（含 YAML frontmatter + step-by-step 指引）
|
||||
fn generate_skill_md(pattern: &DetectedPattern) -> String;
|
||||
}
|
||||
```
|
||||
|
||||
生成的 SKILL.md 自动包含：
|
||||
- `pinned: false`（初始不固定）
|
||||
- `when_to_use` 自动从工具名推断
|
||||
- 每个工具调用作为 `<step>` 写入正文
|
||||
- `version: "0.1.0"`（自动生成版本）
|
||||
|
||||
### Curator — Skill 生命周期管理
|
||||
|
||||
基于 Hermes-Agent `curator.py` 的设计，实现确定性的时间戳驱动生命周期：
|
||||
|
||||
```mermaid
|
||||
stateDiagram-v2
|
||||
[*] --> Active: 创建 / 使用
|
||||
Active --> Active: seed_record (0d) / 有调用记录
|
||||
Active --> Inactive: 7d 无使用
|
||||
Inactive --> Active: 再次使用
|
||||
Inactive --> Stale: 30d 无使用
|
||||
Stale --> Deprecated: 90d 无使用
|
||||
Deprecated --> [*]: 手动删除
|
||||
|
||||
state Active {
|
||||
[*] --> Pinned: pinned=true
|
||||
Pinned --> Pinned: 强制保持 (min_score=0.8)
|
||||
}
|
||||
```
|
||||
|
||||
**生命周期状态**：
|
||||
|
||||
| 状态 | 条件 | 行为 |
|
||||
|------|------|------|
|
||||
| `Active` | 最近使用 ≤ 7 天 | 正常在 remind 列表中出现 |
|
||||
| `Inactive` | 7-30 天未使用 | 不出现在 remind 列表，可被重新激活 |
|
||||
| `Stale` | 30-90 天未使用 | 标记为 stale，出现在清理候选列表 |
|
||||
| `Deprecated` | > 90 天未使用 | 建议归档或删除 |
|
||||
|
||||
**质量评分**：
|
||||
|
||||
```
|
||||
quality_score = ln(1 + invoke_count) × 0.5^(days_since_last_use / 7)
|
||||
```
|
||||
|
||||
**Pinned 保护**：`pinned=true` 的 skill 强制 `Active` 状态，最低评分 0.8，不会出现在清理候选列表中。
|
||||
|
||||
### Seed Record — 新 Skill 锚定时钟
|
||||
|
||||
新创建或自动生成的 skill 可能没有使用统计，`seed_record` 机制防止它们被立即标记为 stale：
|
||||
|
||||
```rust
|
||||
fn evaluate_quality(&self, skill: &Skill, stats: Option<&SkillUsageStat>) -> SkillQuality {
|
||||
let (invoke_count, days_since_last_use) = match stats {
|
||||
Some(s) => (s.invoke_count, s.days_since_last_use()),
|
||||
None => (
|
||||
0,
|
||||
// seed_record: 无统计 → days_since_last_use = 0（视为刚创建）
|
||||
Some(0),
|
||||
),
|
||||
};
|
||||
// 如果 days_since_last_use <= 7 (NEW_SKILL_GRACE_PERIOD_DAYS) → Active
|
||||
// ...
|
||||
}
|
||||
```
|
||||
|
||||
关键常量：
|
||||
- `NEW_SKILL_GRACE_PERIOD_DAYS = 7`：新 skill 在 7 天内即使零调用也保持 Active
|
||||
- `STALE_THRESHOLD_DAYS = 30`：30 天未用标记为 stale
|
||||
- `DEPRECATED_THRESHOLD_DAYS = 90`：90 天未用标记为 deprecated
|
||||
|
||||
### CuratorRunner — 后台空闲触发审查
|
||||
|
||||
参考 Hermes-Agent 的 inactivity-triggered curator：
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant User as 用户交互
|
||||
participant App as AgentRuntime
|
||||
participant CR as CuratorRunner
|
||||
participant Curator as Curator
|
||||
|
||||
Note over CR: 后台 tokio task
|
||||
loop 每 check_interval
|
||||
CR->>CR: should_run_now()
|
||||
alt 未暂停 AND 空闲 > min_idle AND 距上次 > interval
|
||||
CR->>Curator: run_once()
|
||||
Curator->>Curator: evaluate_quality() / archive_stale_skills()
|
||||
Curator-->>CR: CuratorReport
|
||||
else 不满足条件
|
||||
CR->>CR: skip
|
||||
end
|
||||
end
|
||||
|
||||
User->>App: 发送查询
|
||||
App->>CR: record_activity() 更新心跳
|
||||
```
|
||||
|
||||
**CuratorRunner API**：
|
||||
|
||||
```rust
|
||||
pub struct CuratorRunner {
|
||||
curator: Curator,
|
||||
db: SqlitePool,
|
||||
interval: Duration, // 最小审查间隔（默认 7 天）
|
||||
min_idle: Duration, // 最小空闲时间（默认 2 小时）
|
||||
check_interval: Duration, // 检查间隔（默认 1 小时）
|
||||
paused: AtomicBool,
|
||||
last_activity: RwLock<Instant>,
|
||||
last_run: RwLock<Option<Instant>>,
|
||||
}
|
||||
|
||||
impl CuratorRunner {
|
||||
pub fn new(curator: Curator, db: SqlitePool) -> Self;
|
||||
pub fn with_interval(mut self, interval: Duration) -> Self;
|
||||
pub fn with_min_idle(mut self, min_idle: Duration) -> Self;
|
||||
|
||||
/// 判断是否应运行：未暂停 + 空闲超时 + 距上次超间隔
|
||||
pub fn should_run_now(&self) -> bool;
|
||||
|
||||
/// 记录用户活动心跳
|
||||
pub async fn record_activity(&self);
|
||||
|
||||
/// 执行一次审查（仅在 should_run_now 时）
|
||||
pub async fn run_once(
|
||||
&self,
|
||||
usage_stats: &HashMap<String, SkillUsageStat>,
|
||||
skill_metas: &[SkillMeta],
|
||||
) -> Option<CuratorReport>;
|
||||
|
||||
/// 启动后台任务
|
||||
pub fn spawn(
|
||||
self: Arc<Self>,
|
||||
usage_stats: Arc<RwLock<HashMap<String, SkillUsageStat>>>,
|
||||
skill_metas: Arc<RwLock<Vec<SkillMeta>>>,
|
||||
check_interval: Duration,
|
||||
) -> JoinHandle<()>;
|
||||
|
||||
pub fn pause(&self);
|
||||
pub fn resume(&self);
|
||||
}
|
||||
```
|
||||
|
||||
**使用示例**：
|
||||
|
||||
```rust
|
||||
let curator = Curator::new(skills_dir.clone());
|
||||
let runner = Arc::new(
|
||||
CuratorRunner::new(curator, db_pool.clone())
|
||||
.with_interval(Duration::from_secs(7 * 24 * 3600)) // 最少间隔 7 天
|
||||
.with_min_idle(Duration::from_secs(2 * 3600)), // 空闲 2 小时后
|
||||
);
|
||||
|
||||
// 每次用户交互时更新心跳
|
||||
runner.record_activity().await;
|
||||
|
||||
// 启动后台任务
|
||||
let _handle = runner.spawn(usage_stats, skill_metas, Duration::from_secs(3600));
|
||||
```
|
||||
|
||||
### SelfImprovePipeline — 一站式管道
|
||||
|
||||
```rust
|
||||
pub struct SelfImprovePipeline {
|
||||
detector: PatternDetector,
|
||||
creator: SkillCreator,
|
||||
curator: Curator,
|
||||
}
|
||||
|
||||
impl SelfImprovePipeline {
|
||||
pub fn new(db: SqlitePool, skills_dir: PathBuf) -> Self;
|
||||
|
||||
/// 完整管道：检测 → 创建 → 审查
|
||||
pub async fn run(&self, dry_run: bool) -> Result<SelfImproveResult>;
|
||||
|
||||
/// 仅检测模式（不创建）
|
||||
pub async fn detect_only(&self) -> Result<Vec<DetectedPattern>>;
|
||||
|
||||
/// 仅分析已有 skills（不检测新模式）
|
||||
pub async fn analyze_only(
|
||||
&self,
|
||||
usage_stats: &HashMap<String, SkillUsageStat>,
|
||||
skill_metas: &[SkillMeta],
|
||||
) -> Result<CuratorReport>;
|
||||
}
|
||||
|
||||
pub struct SelfImproveResult {
|
||||
pub patterns_found: usize,
|
||||
pub skills_created: usize,
|
||||
pub skills_created_names: Vec<String>,
|
||||
pub curator_report: CuratorReport,
|
||||
}
|
||||
```
|
||||
|
||||
**管道流程**：
|
||||
|
||||
```
|
||||
SelfImprovePipeline::run(dry_run=true)
|
||||
│
|
||||
├─ 1. PatternDetector::scan()
|
||||
│ └─ 从 agent_messages 中检测 ≥3 次跨 session 重复序列
|
||||
│ └─ 结果: Vec<DetectedPattern>（按 confidence 降序）
|
||||
│
|
||||
├─ 2. SkillCreator::create_from_patterns(patterns, dry_run)
|
||||
│ ├─ dry_run=true → 只记录日志，不写入文件
|
||||
│ └─ dry_run=false → 写入 skills/ 目录 + 触发 SkillRegistry::refresh()
|
||||
│
|
||||
└─ 3. Curator::analyze(usage_stats, skill_metas)
|
||||
└─ 质量评分 + 生命周期状态 + 清理建议
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 与 Claude Code 参考设计的对应关系
|
||||
|
||||
| Claude Code 概念 | AstroResearch 实现 |
|
||||
|---|---|
|
||||
| Claude Code / Hermes 概念 | AstroResearch 实现 |
|
||||
|---|---|---|
|
||||
| `src/skills/` 目录 + `SKILL.md` | 完全相同 |
|
||||
| YAML frontmatter（name, description, context, allowed-tools...） | 相同，增加 `version`、`agent`、`effort`、`paths` 字段 |
|
||||
| YAML frontmatter（name, description, context, allowed-tools...） | 相同，增加 `version`、`agent`、`effort`、`paths`、`pinned` 字段 |
|
||||
| `<system-reminder>` Layer 1 注入 | `build_reminder()` → 结构化 XML 块 |
|
||||
| `SkillTool` Layer 2 按需加载 | `LoadSkillTool`（AgentTool trait 实现） |
|
||||
| inline 模式（注入指令内容） | ✅ 实现 |
|
||||
@@ -334,6 +636,11 @@ Box::new(LoadSkillTool::new(skill_registry)),
|
||||
| 条件 skill（paths glob） | ✅ `activate_conditional_for_paths()` |
|
||||
| 变量替换 | ✅ `${SKILL_DIR}`, `${SESSION_ID}` |
|
||||
| `Skill` 工具接口 + `skill` slash command | LoadSkillTool（tool 形式），前端的 `/skill-name` 通过 tool 调用实现 |
|
||||
| Hermes: Self-improving Skills（模式检测 + 自动创建） | ✅ `PatternDetector` + `SkillCreator` + `SelfImprovePipeline` |
|
||||
| Hermes: Curator 生命周期管理 | ✅ `Curator`（Active/Inactive/Stale/Deprecated） |
|
||||
| Hermes: Pinned Skills（不可清理） | ✅ `pinned: true` frontmatter + Curator 保护 |
|
||||
| Hermes: Seed Record（新 skill 锚定时钟） | ✅ `days_since_last_use=Some(0)` + 7 天 grace period |
|
||||
| Hermes: Inactivity-triggered Runner | ✅ `CuratorRunner`（后台 tokio task + 心跳记录） |
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# 系统提示词架构 (System Prompt Architecture)
|
||||
|
||||
AstroResearch 的 Agent 系统提示词采用**模块化 Section 组装 + 动态注入 + 多层生命周期**架构，直接参考 Claude Code 的 System Prompt 设计。
|
||||
AstroResearch Agent 系统提示词采用**模块化 Section 组装 + 简单首次缓存**架构，参考 Claude Code 的 System Prompt 设计并针对实际场景裁剪。
|
||||
|
||||
## 整体分层
|
||||
|
||||
@@ -9,26 +9,29 @@ graph TB
|
||||
subgraph L5["Layer 5: 运行时注入"]
|
||||
Nudge["nudge / 任务恢复 / 后台通知"]
|
||||
end
|
||||
subgraph L4["Layer 4: 提示词压缩"]
|
||||
Compress["snip → micro → auto → identity"]
|
||||
subgraph L4["Layer 4: 提示词压缩 + CollapseLog"]
|
||||
Compress["snip → micro → auto → aggressive_micro"]
|
||||
Collapse["CollapseLog commit / overflow"]
|
||||
end
|
||||
subgraph L3["Layer 3: Skill 动态加载"]
|
||||
Skill["Layer1 提醒 → Layer2 全文注入"]
|
||||
end
|
||||
subgraph L2["Layer 2: 子代理隔离提示词"]
|
||||
SubSP["独立的 system_prompt"]
|
||||
subgraph L2["Layer 2: 子代理模块化提示词"]
|
||||
SubSP["6 section 组装器"]
|
||||
end
|
||||
subgraph L1["Layer 1: 主代理 SystemPrompt 组装"]
|
||||
MainSP["5 个 section 模块化组装"]
|
||||
MainSP["9 个 section 模块化组装<br/>静态(5) → 动态(4)"]
|
||||
Cache["SystemPromptCache<br/>首次计算，永久复用"]
|
||||
end
|
||||
L5 --> L4 --> L3 --> L2 --> L1
|
||||
L1 --> Cache
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 核心组装器 (`src/agent/runtime/system_prompt.rs`)
|
||||
|
||||
### 2.1 数据结构
|
||||
### 数据结构
|
||||
|
||||
```rust
|
||||
pub struct SystemPrompt {
|
||||
@@ -36,290 +39,288 @@ pub struct SystemPrompt {
|
||||
}
|
||||
```
|
||||
|
||||
简单的有序 section 列表，通过 `assemble()` 方法用双换行符 `"\n\n"` 拼接所有 section 内容。section 按添加顺序排列。
|
||||
有序 section 列表。`assemble()` 用 `"\n\n"` 拼接所有 section。顺序即最终 prompt 中出现的顺序——静态内容在前，动态内容在后。
|
||||
|
||||
### 2.2 静态常量
|
||||
### Section 缓存 (`SystemPromptCache`)
|
||||
|
||||
两个 `&'static str` 常量在所有运行时实例间共享内存：
|
||||
简化设计：**首次计算，永久缓存**。因为 session 生命周期内 CWD、platform、OS、model、工具注册表均不变，不需要 TTL 过期机制。仅在 `/clear` 或 `/compact` 事件时调用 `invalidate_all()` 全局失效。
|
||||
|
||||
**IDENTITY_SECTION**（身份声明，1 行）：
|
||||
```
|
||||
你是一位专业的天体物理学研究助手，具备丰富的天文学知识。
|
||||
```rust
|
||||
pub struct SystemPromptCache {
|
||||
entries: HashMap<&'static str, String>,
|
||||
}
|
||||
|
||||
impl SystemPromptCache {
|
||||
// 首次计算，后续命中缓存
|
||||
pub fn get_or_compute(&mut self, name: &str, compute: impl FnOnce() -> String) -> String;
|
||||
// compute 返回 Option：Some 缓存并返回，None 不缓存
|
||||
pub fn get_or_compute_optional(&mut self, name: &str, compute: impl FnOnce() -> Option<String>) -> Option<String>;
|
||||
// 显式失效
|
||||
pub fn invalidate(&mut self, name: &str);
|
||||
pub fn invalidate_all(&mut self);
|
||||
}
|
||||
```
|
||||
|
||||
**PRINCIPLES_SECTION**（核心行为准则，9 条）：
|
||||
```
|
||||
核心原则：
|
||||
1. 主动使用工具搜索最新文献，不要仅凭训练数据回答。
|
||||
2. 优先使用本地资源（get_paper_content / rag_search），必要时再检索新文献。
|
||||
3. 收集到足够信息后立即给出最终答案，避免无意义的重复工具调用。
|
||||
4. 回答时引用具体文献来源，使用 ADS bibcode 标注。
|
||||
5. 对于数学公式，使用标准 LaTeX 格式。
|
||||
6. 用中文回答，保持科学术语的准确性（可附带英文原文）。
|
||||
7. 对于复杂任务（如文献综述），调用 load_skill 获取方法论指引，再用 todo_write 制定计划。
|
||||
8. 如果某个工具调用失败，不要用相同参数重试，尝试换一种方式或工具。
|
||||
9. 任务状态会在每轮开始时从数据库恢复，请基于最新状态继续工作。
|
||||
```
|
||||
缓存策略：静态 section + environment + tools 全部通过 `get_or_compute` 缓存。skills 和 memory 不缓存——前者通过文件监听热更新，后者受 `save_memory` 工具实时影响。
|
||||
|
||||
### 2.3 组装顺序
|
||||
### 静态 Section 常量（6 个）
|
||||
|
||||
每轮调用 `AgentRuntime::system_prompt()` 方法（`src/agent/runtime/mod.rs:1154-1203`），按以下顺序组装 5 个 section：
|
||||
| 常量 | 内容 | 行数 |
|
||||
|------|------|------|
|
||||
| `IDENTITY_SECTION` | 身份声明 | 1 |
|
||||
| `PRINCIPLES_SECTION` | 核心行为准则（9 条） | 9 |
|
||||
| `SYSTEM_CONTEXT_SECTION` | system-reminder 标签说明 + 自动压缩 | 3 |
|
||||
| `TOOL_USAGE_SECTION` | 专用工具优先、并行调用、todo_write | 5 |
|
||||
| `SAFETY_SECTION` | 可逆性、影响范围、确认机制 | 5 |
|
||||
|
||||
### 组装顺序
|
||||
|
||||
`AgentRuntime::system_prompt()` 组装 9 个 section：
|
||||
|
||||
```
|
||||
Section 1: identity 静态 — 最大化 Anthropic prompt cache 命中率
|
||||
Section 2: tools 动态 — 从 ToolRegistry 生成工具名称+摘要列表
|
||||
Section 3: skills 动态 — 从 SkillRegistry.build_reminder() 生成（<system-reminder> XML）
|
||||
Section 4: memory 动态 — 从 MemoryManager.build_system_reminder(5) 生成（<project-memory-context> XML）
|
||||
Section 5: principles 静态 — 核心原则（放在最后 — 若需调整仅影响最后一个 cache segment）
|
||||
[1] identity ← 静态（首次计算后永久缓存）
|
||||
[2] principles ← 静态
|
||||
[3] system_context ← 静态
|
||||
[4] tool_usage ← 静态
|
||||
[5] safety ← 静态
|
||||
[6] environment ← 动态（首次计算后缓存，session 内不变）
|
||||
[7] tools ← 动态（首次计算后缓存，ToolRegistry session 内不变）
|
||||
[8] skills ← 动态（不缓存，文件监听热更新）
|
||||
[9] memory ← 动态（不缓存，save_memory 实时更新）
|
||||
```
|
||||
|
||||
**缓存策略**：静态 section 固定且不变化，放在 prompt 头部以最大化 Anthropic prompt cache 命中率。动态 section（tools、skills、memory）因内容较少，对 cache 影响可控。principles 虽然静态但放在最后，当需要调优时仅破坏最后一个 cache segment。
|
||||
**为什么没有 TTL**：CWD、platform、OS、model、tools 在 session 生命周期内全部不变。首次计算即永久正确，TTL 是多余的复杂度。
|
||||
|
||||
**为什么没有 cache_control 边界标记**：`cache_control: {"type": "ephemeral"}` 是 Anthropic API 专有特性。我们的模型（DeepSeek/Qwen 等 OpenAI 兼容 API）不支持。静态内容前置的顺序本身已足够让服务端按内容哈希自然缓存。
|
||||
|
||||
---
|
||||
|
||||
## 动态 Section 详解
|
||||
|
||||
### 3.1 工具列表 (tools section)
|
||||
### environment section
|
||||
|
||||
```rust
|
||||
let mut tools_desc = String::from("你可以使用以下工具：\n");
|
||||
fn build_environment_section(&self) -> String {
|
||||
// 包含：工作目录、Git 仓库状态、平台、OS 版本、日期、模型名称
|
||||
// 以及 Agent 配置摘要（最大步数、工具超时）
|
||||
}
|
||||
```
|
||||
|
||||
示例输出：
|
||||
```
|
||||
# 环境信息
|
||||
- 工作目录: /home/user/project
|
||||
- Git 仓库: 是
|
||||
- 平台: linux
|
||||
- OS 版本: Linux 7.0.0-22-generic
|
||||
- 日期: 2026-06-22
|
||||
- 当前模型: deepseek-v4-pro
|
||||
- 最大推理步数: 8
|
||||
- 工具超时: 120 秒
|
||||
```
|
||||
|
||||
### tools section
|
||||
|
||||
```rust
|
||||
// 从 ToolRegistry.definitions() 生成工具名称 + 80 字符摘要
|
||||
// ToolRegistry 内部有 schema_cache：工具注册表不变时复用上次计算结果
|
||||
for def in self.tool_registry.definitions() {
|
||||
let short_desc = def.function.description
|
||||
.split('。').next()
|
||||
.unwrap_or(&def.function.description)
|
||||
.chars().take(80)
|
||||
.split('。').next() // 取第一句
|
||||
.chars().take(80) // 截断 80 字符
|
||||
.collect();
|
||||
tools_desc.push_str(&format!("- {}: {}\n", def.function.name, short_desc));
|
||||
}
|
||||
```
|
||||
|
||||
- 19 个默认工具：`search_papers`, `download_paper`, `parse_paper`, `get_paper_content`, `rag_search`, `query_target`, `save_note`, `read_file`, `grep_files`, `glob_files`, `run_bash`, `file_write`, `file_edit`, `todo_write`, `compress_context`, `load_skill`, `subagent`, `save_memory`, `bg_task_run`
|
||||
- 描述仅取**第一句 + 前 80 字符**作为功能摘要
|
||||
- 完整的参数 JSON Schema 通过 API 的 `tools` 参数单独传递，不在 system prompt 中重复
|
||||
完整 JSON Schema 通过 API `tools` 参数单独传递，不在 system prompt 中重复。
|
||||
|
||||
### 3.2 技能列表 (skills section) — 两层加载
|
||||
### skills section — 两层加载
|
||||
|
||||
参考 Claude Code 的两层技能设计，定义在 `src/agent/skills.rs`：
|
||||
|
||||
**Layer 1 (system-reminder)**：`SkillRegistry.build_reminder()` 在 system prompt 中注入 `<system-reminder>` XML 块。列出所有 `user_invocable=true` 且 `disable_model_invocation=false` 的技能名称 + 描述。每个 skill 约消耗 ~20 tokens。
|
||||
**Layer 1 (system-reminder)**：`SkillRegistry.build_reminder()` 列出所有 `user_invocable=true` 的技能名称 + 描述 + when_to_use，每个约 20 tokens。
|
||||
|
||||
```xml
|
||||
<system-reminder>
|
||||
The following skills are available for use with the Skill tool:
|
||||
- methodology: 天体物理研究方法论指南 - When user asks about research methodology
|
||||
- plotting: 数据可视化与科学绘图 - When user wants to create plots
|
||||
- presentation: 学术幻灯片制作 - When user needs to prepare a presentation
|
||||
When a skill matches the user's request, invoke load_skill BEFORE generating any other response about the task.
|
||||
If you see a <command-name> tag in the current conversation turn, the skill has ALREADY been loaded - follow the instructions directly instead of calling load_skill again.
|
||||
When a skill matches the user's request, invoke load_skill BEFORE generating any other response...
|
||||
</system-reminder>
|
||||
```
|
||||
|
||||
**Layer 2 (load_skill 工具)**：LLM 按需调用 `load_skill(skill_name)` 工具，从 `skills/{name}/SKILL.md` 加载完整内容（YAML frontmatter + Markdown body），注入到消息上下文。完整 skill 约 ~2000 tokens。
|
||||
**Layer 2 (load_skill 工具)**：LLM 按需调用，从 `skills/{name}/SKILL.md` 加载完整内容（~2000 tokens），支持 `${SKILL_DIR}` / `${SESSION_ID}` 变量替换和 fork 执行模式。
|
||||
|
||||
SKILL.md 格式：
|
||||
```yaml
|
||||
---
|
||||
name: methodology
|
||||
description: 天体物理研究方法论指南
|
||||
context: inline # inline | fork
|
||||
when_to_use: When user asks about research methodology
|
||||
allowed-tools:
|
||||
- search_papers
|
||||
- rag_search
|
||||
model: inherit
|
||||
user-invocable: true
|
||||
---
|
||||
详见 [skills.md](skills.md)。
|
||||
|
||||
# Skill 正文
|
||||
详细内容...
|
||||
```
|
||||
### memory section
|
||||
|
||||
**热重载**：SkillRegistry 通过 `notify` crate 监听 skills 目录的文件变更，300ms debounce 后自动刷新。Skill 按使用频率排序（指数衰减评分，7 天半衰期）。
|
||||
|
||||
**条件激活**：Skill 可通过 `paths` frontmatter 声明 glob 模式。Agent 访问匹配文件时自动将 `disable_model_invocation` 设为 false，激活条件 skill。
|
||||
|
||||
### 3.3 项目记忆 (memory section)
|
||||
|
||||
`MemoryManager.build_system_reminder(5)` 从 `{library_dir}/memory/` 目录加载最近 5 条记忆，生成 `<project-memory-context>` XML 块：
|
||||
|
||||
```xml
|
||||
<project-memory-context>
|
||||
[PROJECT MEMORY]
|
||||
[偏好] memory-slug: 一句话描述
|
||||
内容预览前三行
|
||||
[时效提示: 此记忆已超过N天，可能已过时]
|
||||
[反馈] another-memory: 描述 [已更新→new-slug]
|
||||
内容预览...
|
||||
使用 save_memory 工具保存重要信息。记忆内容可能过时，请在使用前验证。
|
||||
</project-memory-context>
|
||||
```
|
||||
|
||||
关键特性：
|
||||
- 按 mtime 排序（最新在前），支持语义选择 + 指数衰减排序
|
||||
- 按类型标注：`[偏好]` / `[反馈]` / `[项目]` / `[参考]`
|
||||
- 过期记忆标记为 `[已更新]` 或 `[已更新→new-slug]`（归档为 `{slug}_v1.md`）
|
||||
- 超过 1 天的记忆注入时效警告
|
||||
- 索引文件 `MEMORY.md` 限制 200 行 / 25KB
|
||||
从 `{library_dir}/memory/` 加载，**按 mtime 降序排列**（最新在前），取前 5 条，注入 `<project-memory-context>` XML 块。支持时效警告、过期标记、语义选择。详见 [memory.md](memory.md)。
|
||||
|
||||
---
|
||||
|
||||
## 上下文初始化与运行时注入 (`src/agent/runtime/context.rs`)
|
||||
## 上下文初始化 (`src/agent/runtime/context.rs`)
|
||||
|
||||
### 4.1 上下文构建流程
|
||||
|
||||
`build_initial_context()` 在每轮开始时构建完整的消息列表：
|
||||
`build_initial_context()` 流程：
|
||||
|
||||
```
|
||||
1. 从数据库加载历史消息（agent_messages 表）
|
||||
2. 如果历史第一条不是 system 角色 → 在位置 0 插入系统提示词
|
||||
1. 加载历史消息（agent_messages 表）
|
||||
2. 如果第一条不是 system 角色 → 插入系统提示词
|
||||
3. 追加当前用户问题
|
||||
4. [可选] 追加任务状态恢复提醒（从 agent_tasks 表读取）
|
||||
4. [可选] 追加任务状态恢复提醒（agent_tasks 表）
|
||||
```
|
||||
|
||||
### 4.2 任务状态恢复
|
||||
运行时干预通过**注入 user 消息**实现（不修改 system prompt）：
|
||||
|
||||
从 `agent_tasks` 表恢复未完成的任务，格式化注入 user 消息：
|
||||
|
||||
```
|
||||
[当前任务状态]
|
||||
以下是上次会话中持久化的任务计划，请基于最新状态继续工作：
|
||||
|
||||
⏳ [task-1] 搜索相关文献...
|
||||
🔄 [task-2] 分析论文数据... (依赖: task-1)
|
||||
✅ [task-3] 格式化引用... (指派: lead)
|
||||
|
||||
使用 todo_write 工具更新任务进度。
|
||||
```
|
||||
|
||||
### 4.3 运行时 Nudge 注入
|
||||
|
||||
在 ReAct 循环中，system prompt 组装后不再修改。运行时干预通过**注入 user 消息**实现（开闭原则）：
|
||||
|
||||
| 触发条件 | Nudge 内容 |
|
||||
| 触发条件 | 注入内容 |
|
||||
|:---|:---|
|
||||
| TodoWrite 连续 3 步未更新 | "提醒：你已经连续多步未更新任务计划。建议调用 todo_write 工具…" |
|
||||
| Token 预算 diminishing returns | "检测到你的后续步骤未产生新信息…请基于已收集的全部信息直接给出最终答案" |
|
||||
| 达到最大步数 (max_steps) | "你已经执行了 N 步（最大 M 步）。请根据已有信息直接给出最终答案" |
|
||||
| 后台任务完成 | "[后台任务完成] ✅ tool_name: bibcode: summary" |
|
||||
| TodoWrite 3 步未更新 | "提醒：建议调用 todo_write 工具复盘进度" |
|
||||
| Token 预算 diminishing returns | "检测到重复操作模式，请直接给出最终答案" |
|
||||
| 达到最大步数 | "已执行 N 步（最大 M 步），请直接给出最终答案" |
|
||||
| 后台任务完成 | "[后台任务完成] ✅ tool_name: summary" |
|
||||
|
||||
---
|
||||
|
||||
## 子代理的独立系统提示词 (`src/agent/tools/subagent.rs`)
|
||||
## 子代理模块化系统提示词 (`src/agent/tools/subagent.rs`)
|
||||
|
||||
子代理拥有独立的消息上下文，通过 `SubAgentRunner::run()` 接收一个**硬编码的简化版系统提示词**：
|
||||
子代理使用完整模块化系统提示词（不再硬编码 96 字符）：
|
||||
|
||||
```
|
||||
你是一位专业的天体物理学研究助手，在一个独立的子任务上下文中工作。
|
||||
你可以使用文献搜索、下载、RAG检索等工具。
|
||||
请高效完成任务，然后直接给出最终答案。不要进行不必要的重复操作。
|
||||
用中文回答，引用具体文献来源。
|
||||
[1] identity ← 与父代理相同
|
||||
[2] subagent_context ← "在独立的子任务上下文中工作，请专注于完成这项任务"
|
||||
[3] principles ← 与父代理相同
|
||||
[4] system_context ← 与父代理相同
|
||||
[5] tool_usage ← 与父代理相同
|
||||
[6] safety ← 与父代理相同
|
||||
[7] tools ← 运行时从 ToolRegistry 生成
|
||||
```
|
||||
|
||||
特点：
|
||||
- 不继承父代理的 tools/skills/memory sections
|
||||
- 共享父代理的 ToolRegistry
|
||||
- 通过 `PermissionChecker` 可在特定场景下限制工具访问
|
||||
- 独立的 ReAct 循环（步数上限通过参数传入，默认 5，最大 10）
|
||||
- 完整的 Hook 管道（PreToolUse/PostToolUse/SubagentStart/SubagentStop）
|
||||
- 包含活跃度日志（activity log），返回给父代理时附带工具调用统计
|
||||
|
||||
### 5.2 团队成员的独立提示词 (`src/agent/team/teammate.rs`)
|
||||
|
||||
队友的 `system_prompt` 和 `task_prompt` 由 `team/manager.rs`（lead 的委托逻辑）在运行时构造并传入 `run_teammate_loop()`：
|
||||
- prompt 内容完全由 lead 的决定
|
||||
- 队友不包含 `subagent` 工具（防止无限委托链）
|
||||
- 更轻量的 ReAct 循环（无 SSE、无 DB 持久化、无 hooks）
|
||||
- 步数上限更严格（min(max_steps, 5)）
|
||||
- 通过文件收件箱与 lead 通信（每 5 秒 poll，最长 60 秒）
|
||||
预构建 `ToolRegistry`，在构造系统提示词前获取 `definitions()`。`SubAgentRunner` 新增 `new_with_registry_and_hooks()` 构造函数。详见 [subagent.md](subagent.md)。
|
||||
|
||||
---
|
||||
|
||||
## 上下文压缩中的独立提示词 (`src/agent/compact.rs`)
|
||||
## 上下文压缩 + CollapseLog (`src/agent/compact.rs`)
|
||||
|
||||
### 6.1 四层压缩策略
|
||||
### 四层压缩策略
|
||||
|
||||
| 层 | 方法 | API 调用 | 行为 |
|
||||
| 层 | 方法 | API | 行为 |
|
||||
|:---|:---|:---|:---|
|
||||
| Layer 0 | `snip_compact` | 无 | 消息数超过 50 时截断中间段，保留头 3 + 尾 47 |
|
||||
| Layer 1 | `micro_compact` | 无 | 将较早的工具结果替换为 `[Previous: used {tool_name}]` 占位符 |
|
||||
| Layer 2 | `auto_compact` | 1 次 | LLM 摘要对话历史（见下），注入 `[历史对话摘要]` |
|
||||
| Layer 3 | `aggressive_micro` | 无 | 保留最近 2 条工具结果，其余替换为占位符 |
|
||||
| 0 | `snip_compact` | 无 | 消息 >50 时截断中间段，保留头 3 + 尾 |
|
||||
| 1 | `micro_compact` | 无 | 较早工具结果替换为 `[Previous: used {name}]` 占位符 |
|
||||
| 2 | `auto_compact` | 1 次 | LLM 摘要对话历史 |
|
||||
| 3 | `aggressive_micro` | 无 | 保留最近 2 条工具结果 |
|
||||
|
||||
### 6.2 LLM 摘要 Prompt
|
||||
### CollapseLog
|
||||
|
||||
Layer 2 中调用 LLM 生成摘要时，使用独立的系统提示词：
|
||||
`compress_context_with_hooks_and_log()` 在每次压缩后记录结构化 commit：
|
||||
|
||||
```
|
||||
系统: "你是一个对话摘要助手。请提取对话的关键信息和结论。"
|
||||
用户: "请用简洁的中文总结以下对话历史的要点（不超过500字）：
|
||||
|
||||
[用户] ...
|
||||
[助手] ...
|
||||
[工具] ..."
|
||||
```rust
|
||||
log.commit(CollapseMethod::LlmSummary, (after_count, before_count), summary);
|
||||
```
|
||||
|
||||
### 6.3 身份再注入
|
||||
|
||||
如果压缩后消息过少（≤4 条），注入身份确认块防止模型丢失上下文认知：
|
||||
|
||||
```
|
||||
[身份确认] 你是一位专业的天体物理学研究助手。以上是历史对话的压缩摘要。
|
||||
你正在进行的研究任务是回答用户的问题。请基于摘要中的关键信息继续工作，
|
||||
需要更多信息时主动使用工具搜索。
|
||||
```
|
||||
|
||||
### 6.4 安全切割
|
||||
|
||||
`find_safe_cut_point()` 确保压缩时不会破坏 `assistant(tool_calls)` / `tool_result` 配对关系，向前追溯找到完整工具交互的边界。
|
||||
|
||||
### 6.5 熔断器
|
||||
|
||||
`CompactionCircuitBreaker` 防止连续压缩失败时的无限循环。连续 3 次压缩后消息数未减少 → 打开熔断器，后续跳过自动压缩。
|
||||
超 5 条 commits 时触发**溢出合并**，将最早的 commits 合并为摘要注入消息列表。详见 `collapse.rs`。
|
||||
|
||||
---
|
||||
|
||||
## Hook 系统与提示词的交互 (`src/agent/hooks.rs`)
|
||||
## Hook 系统与提示词交互
|
||||
|
||||
Hook 系统定义 9 个生命周期事件，其中与提示词相关的交互：
|
||||
9 个生命周期事件：
|
||||
|
||||
| Hook | 与提示词的关系 |
|
||||
|:---|:---|
|
||||
| `OnSessionStart` | 在提示词组装前触发，可影响任务状态恢复逻辑 |
|
||||
| `PreToolUse::MutateInput` | 可向工具执行注入 `additional_context`（作为 user 消息追加） |
|
||||
| `PreToolUse::Block` | 阻止特定工具的执行（如取消检查） |
|
||||
| `PostToolUse::MutateOutput` | 可修改工具输出内容（影响后续 LLM 看到的 context） |
|
||||
| `OnStepComplete` | 每步结束记录 token 估算、消息数等指标 |
|
||||
| `PreCompact` | 压缩前记录消息数和 token 估算 |
|
||||
| `PostCompact` | 压缩后记录最终消息数和压缩方法 |
|
||||
| `OnSubagentStart/Stop` | 子代理启动/停止时传递 prompt 和结果摘要 |
|
||||
| `OnSessionStop` | 会话终止时清理取消状态并记录终止原因 |
|
||||
| `OnSessionStart` | 提示词组装前触发 |
|
||||
| `PreToolUse::MutateInput` | 注入 `additional_context`（追加为 user 消息） |
|
||||
| `PostToolUse::MutateOutput` | 修改工具输出（影响 LLM 看到的 context） |
|
||||
| `PreCompact / PostCompact` | 压缩前后记录指标 + CollapseLog commit |
|
||||
| `OnSubagentStart/Stop` | 传递子代理 prompt 和结果摘要 |
|
||||
|
||||
详见 [hooks.md](hooks.md)。
|
||||
|
||||
---
|
||||
|
||||
## 完整数据流
|
||||
## Claude Code 工具按需发现机制（参考分析）
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
RT["AgentRuntime 创建<br/>system_prompt() 调用"]
|
||||
|
||||
RT --> S1["Section 1: identity<br/>(静态常量)"]
|
||||
RT --> S2["Section 2: tools<br/>(ToolRegistry definitions)"]
|
||||
RT --> S3["Section 3: skills<br/>(SkillRegistry.build_reminder)"]
|
||||
|
||||
S1 --> S4
|
||||
S2 --> S4
|
||||
S3 --> S4["Section 4: memory (可选)<br/>(MemoryManager.build_reminder, 5 entries)"]
|
||||
|
||||
S4 --> S5["Section 5: principles<br/>(静态常量)"]
|
||||
S5 --> ASM["assemble()<br/>join('\n\n')"]
|
||||
|
||||
ASM --> Main["主 Agent 上下文<br/>build_initial_context()<br/>+ nudge 注入 + 任务恢复 + 后台通知"]
|
||||
ASM --> Sub["子 Agent 上下文<br/>SubAgentRunner.run()<br/>(独立 system_prompt)"]
|
||||
|
||||
Main --> React["ReAct 循环"]
|
||||
Main --> NudgeInj["Nudge 消息注入 (user)"]
|
||||
Main --> Compact["压缩层<br/>generate_summary()<br/>+ identity re-injection"]
|
||||
> 当前项目工具数量较少（~25 个），尚未实现此机制。以下为 Claude Code 的设计分析，作为未来工具增长时的参考。
|
||||
|
||||
### 要解决的问题
|
||||
|
||||
当接入大量 MCP 工具（30+ 服务器，100+ 工具）时，所有工具的完整 JSON Schema 在每轮 API 调用中占据大量上下文——绝大多数工具从未被调用，却每轮都在消耗 token。Claude Code 的解决方案：**延迟加载 + 按需发现**。
|
||||
|
||||
### 两层架构
|
||||
|
||||
**Layer 1: `defer_loading` 标记**
|
||||
|
||||
工具分为两类：
|
||||
|
||||
| 类型 | 行为 | 示例 |
|
||||
|------|------|------|
|
||||
| 常驻工具 | 始终在 `tools` 数组中，立即可用 | Read、Write、Bash、Glob、Grep、Task、TodoWrite |
|
||||
| 延迟工具 | 标记 `defer_loading: true`，仅在模型主动发现后才加入 `tools` 数组 | MCP 工具、EnterPlanMode、NotebookEdit、LSP 工具 |
|
||||
|
||||
判断逻辑（`isDeferredTool()`）：
|
||||
```
|
||||
1. alwaysLoad == true → 永不延迟
|
||||
2. isMcp == true → 始终延迟（除非 MCP server 设 _meta['anthropic/alwaysLoad']）
|
||||
3. shouldDefer == true → 延迟
|
||||
```
|
||||
|
||||
**Layer 2: ToolSearch 工具 + tool_reference 块**
|
||||
|
||||
模型通过 `ToolSearch` 工具按需发现延迟工具：
|
||||
|
||||
```
|
||||
模型: ToolSearch(query: "github create PR")
|
||||
系统: 返回 tool_reference 块 → [mcp__github__createPullRequest, mcp__github__listPRs]
|
||||
下次 API 调用: 这两个工具的完整 schema 加入 tools 数组
|
||||
```
|
||||
|
||||
`tool_reference` 是一个特殊的 content block 类型，API 服务端收到后会展开为完整工具定义，模型可以在后续 turn 直接调用。
|
||||
|
||||
### 消息流转
|
||||
|
||||
```
|
||||
Turn N:
|
||||
tools 数组 = 常驻工具 + ToolSearch + 之前发现过的延迟工具
|
||||
<available-deferred-tools> 块列出所有可发现的延迟工具名称
|
||||
|
||||
→ 模型调用 ToolSearch(query: "slack")
|
||||
→ tool_result 包含 tool_reference 块: [mcp__slack__sendMessage]
|
||||
|
||||
Turn N+1:
|
||||
tools 数组 = 常驻工具 + ToolSearch + [mcp__slack__sendMessage, ...]
|
||||
→ 模型可以直接调用 mcp__slack__sendMessage
|
||||
```
|
||||
|
||||
### 发现状态持久化
|
||||
|
||||
`extractDiscoveredToolNames()` 从消息历史中扫描所有 `tool_reference` 块，提取已发现的工具名。压缩时通过 compact boundary marker 的 `preCompactDiscoveredTools` 字段保存已发现集合，防止压缩丢失发现状态。
|
||||
|
||||
### 工具发现状态变更通知
|
||||
|
||||
两个机制告诉模型有哪些可发现工具：
|
||||
|
||||
**Legacy**：`<available-deferred-tools>` 块作为 user 消息前置注入（每次工具池变更都会 bust prompt cache）。
|
||||
|
||||
**Modern**：`deferred_tools_delta` 附件——diff 上次通知的工具池，仅发送增量变更（added/removed），避免 bust cache。
|
||||
|
||||
### ToolSearch 搜索方式
|
||||
|
||||
- **精确选择** `select:ToolA,ToolB` — 按名称直接取工具
|
||||
- **关键词搜索** `github create PR` — 搜索工具名 + searchHint + description，加权评分
|
||||
- **必选词** `+slack send` — `+` 前缀表示必须匹配
|
||||
|
||||
### 对项目的适用性
|
||||
|
||||
| 当前状态 | 是否需要 |
|
||||
|---------|---------|
|
||||
| ~25 个工具，schema 总共 ~8K tokens | **暂不需要** |
|
||||
| 无 MCP 工具接入 | 不需要 |
|
||||
| 工具数量稳定 | 不需要 |
|
||||
|
||||
**触发条件**：当工具数量超过 ~50 或接入 MCP 服务器时，可以按以下步骤接入：
|
||||
1. 为 MCP 工具设置 `shouldDefer: true` / `isMcp: true`
|
||||
2. 注册 `ToolSearch` 工具
|
||||
3. 在 ToolRegistry 中维护 `deferred_tool_names` 集合
|
||||
4. API 调用时过滤 tools 数组 + 注入 `<available-deferred-tools>` 块
|
||||
5. 压缩时保存 `preCompactDiscoveredTools` 快照
|
||||
|
||||
---
|
||||
|
||||
@@ -327,33 +328,27 @@ flowchart TD
|
||||
|
||||
| 文件 | 职责 |
|
||||
|:---|:---|
|
||||
| `src/agent/runtime/system_prompt.rs` | SystemPrompt 组装器 + 静态常量 |
|
||||
| `src/agent/runtime/mod.rs:1154-1203` | `system_prompt()` 方法 — 5 section 拼装 |
|
||||
| `src/agent/runtime/system_prompt.rs` | SystemPrompt 组装器 + 6 个静态常量 + SystemPromptCache |
|
||||
| `src/agent/runtime/mod.rs` | `system_prompt()` + `build_environment_section()` |
|
||||
| `src/agent/runtime/context.rs` | `build_initial_context()` — 上下文初始化 + 任务恢复 |
|
||||
| `src/agent/skills.rs` | SkillRegistry — 两层技能加载 + 热重载 |
|
||||
| `src/agent/memory/mod.rs` | MemoryManager — 记忆加载 + system reminder 构建 |
|
||||
| `src/agent/compact.rs` | 四层压缩 + LLM 摘要 prompt + 身份再注入 |
|
||||
| `src/agent/tools/subagent.rs` | 子代理系统提示词（硬编码） |
|
||||
| `src/agent/subagent.rs` | SubAgentRunner — 子代理 ReAct 循环 |
|
||||
| `src/agent/team/teammate.rs` | 队友 ReAct 循环（外部传入 system_prompt） |
|
||||
| `src/agent/hooks.rs` | 9 个生命周期 hook + 提示词交互 |
|
||||
| `src/agent/memory/mod.rs` | MemoryManager — 按 recency 排序的记忆注入 |
|
||||
| `src/agent/tools/mod.rs` | ToolRegistry — schema_cache + definition_filter |
|
||||
| `src/agent/tools/subagent.rs` | 子代理模块化系统提示词构建 |
|
||||
| `src/agent/subagent.rs` | SubAgentRunner — new_with_registry_and_hooks |
|
||||
| `src/agent/compact.rs` | 四层压缩 + CollapseLog 集成 |
|
||||
| `src/agent/compact/collapse.rs` | CollapseLog — 压缩历史记录 + 溢出合并 |
|
||||
| `src/agent/hooks/` | 生命周期 hook + 提示词交互 |
|
||||
| `docs/architecture/agent/system-prompt-optimization-plan.md` | 优化计划文档（背景、方案、对比分析） |
|
||||
|
||||
---
|
||||
|
||||
## 设计要点
|
||||
|
||||
### 优势
|
||||
|
||||
1. **模块化 section 组装**：各 section 独立管理，便于调试和迭代
|
||||
2. **静态 section 前置**：最大化 Anthropic prompt cache 命中率，降低延迟和成本
|
||||
3. **两层 skill 加载**：避免一次性注入所有 skill 的 token 浪费
|
||||
4. **压缩时身份再注入**：防止激进压缩后模型丢失角色认知
|
||||
5. **安全切割点**：`find_safe_cut_point` 确保压缩不破坏 tool_call/tool_result 配对
|
||||
6. **运行时 nudge 而非 system prompt 编辑**：遵循开闭原则，system prompt 保持稳定
|
||||
|
||||
### 潜在改进方向
|
||||
|
||||
1. **子代理系统提示词继承**：当前子代理的 system prompt 是硬编码的，可考虑让子代理也接收 section 组装器，选择性继承 skills/memory
|
||||
2. **压缩 prompt 外部化**：摘要生成和身份确认的 prompt 可配置化，便于独立调优
|
||||
3. **记忆注入锁竞争**：`memory_manager.try_lock()` 在高并发下可能静默失败，考虑使用 `RwLock::read()`
|
||||
4. **工具描述摘要策略**：80 字符截断可能丢失关键语义，可考虑 LLM 预生成工具描述摘要
|
||||
1. **静态前置，动态后置**：内容不变的 section 先出现，服务端自然按哈希缓存
|
||||
2. **首次计算，永久缓存**：session 内一切不变，不需要 TTL；仅在 /clear 时全部失效
|
||||
3. **模块化 section**：各 section 独立管理，便于调试、增删、A/B 测试
|
||||
4. **两层 skill 加载**：20 tokens 列表 vs 2000 tokens 全文，按需加载
|
||||
5. **子代理完整提示词**：共享主代理的静态常量 + 独立 tools 列表
|
||||
6. **CollapseLog 追踪**：每次压缩记录结构化 commit，溢出自动合并
|
||||
7. **运行时 nudge**：干预通过 user 消息注入，不修改 system prompt（开闭原则）
|
||||
|
||||
@@ -16,7 +16,7 @@ graph TB
|
||||
subgraph Layer2["执行协调层 — Executor"]
|
||||
direction LR
|
||||
EX["src/agent/runtime/executor.rs<br/>验证 → PreToolUse hooks → 并行调度 → PostToolUse"]
|
||||
SX["src/agent/runtime/streaming_executor.rs (流式变体)<br/>流式 tool_use 到达时立即调度 + Sibling Abort"]
|
||||
SX["src/agent/runtime/streaming_executor.rs (流式变体)<br/>流式 tool_use 到达时立即调度 + Sibling Abort<br/>非并发工具独占执行 (executing_non_concurrent 标志)"]
|
||||
end
|
||||
|
||||
subgraph Layer3["业务逻辑层 — AgentTool Trait + 工具实现"]
|
||||
@@ -185,21 +185,30 @@ LLM stream → tool_calls[]
|
||||
│ └─ mutated_args + additional_contexts 收集
|
||||
│
|
||||
▼
|
||||
┌─ execute_parallel() ─────────────────────────────────────────┐
|
||||
│ FuturesUnordered 并发调度: │
|
||||
│ 每个 PreparedCall → tokio::spawn(async { │
|
||||
│ tokio::select! { │
|
||||
│ timeout(tool_timeout_secs) → 执行工具 │
|
||||
│ cancel_fut (每 250ms 轮询) → 返回错误 │
|
||||
│ } │
|
||||
│ }) │
|
||||
┌─ execute_parallel() — ToolPartitioner 批次调度 ──────────────┐
|
||||
│ │
|
||||
│ 中断处理: │
|
||||
│ Phase 3a: 构建非拒绝工具的 (原索引, PreparedCall) 映射 │
|
||||
│ Phase 3b: ToolPartitioner::partition() 分区 │
|
||||
│ · 连续 concurrency_safe 工具 → 并行批次 │
|
||||
│ · 非 concurrency_safe 工具 → 独占串行批次 │
|
||||
│ · 示例: [read, grep, bash, read, write] │
|
||||
│ → [read∥grep], [bash], [read], [write] │
|
||||
│ │
|
||||
│ Phase 3c: 逐批次执行 │
|
||||
│ · 并行批次 → FuturesUnordered 并发 (tokio::spawn each) │
|
||||
│ · 串行批次 → 逐个执行 (前一个完成后才启动下一个) │
|
||||
│ │
|
||||
│ 中断处理（每工具内）: │
|
||||
│ tokio::select! { │
|
||||
│ timeout(tool_timeout_secs) → 执行工具 │
|
||||
│ cancel_fut (每 250ms 轮询) → 返回错误 │
|
||||
│ } │
|
||||
│ InterruptBehavior::Block → 忽略 cancel_fut，等完成 │
|
||||
│ InterruptBehavior::Cancel → 响应取消，注入错误 │
|
||||
│ │
|
||||
│ 渐进式结果处理 (while exec_futs.next()): │
|
||||
│ 完成即处理，快工具不因慢工具阻塞 │
|
||||
│ 辅助函数: │
|
||||
│ execute_single_tool() — 单工具超时+取消+执行 │
|
||||
│ process_single_result() — 结果处理+SSE推送+TTL持久化+Hook │
|
||||
└──────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼ (每个工具完成后逐个处理)
|
||||
@@ -221,27 +230,32 @@ LLM stream → tool_calls[]
|
||||
### 并发模型细节
|
||||
|
||||
```rust
|
||||
// executor.rs: FuturesUnordered 中的每个 future
|
||||
Box::pin(async move {
|
||||
let interrupt_behavior = tool.interrupt_behavior();
|
||||
let is_blocking = interrupt_behavior == InterruptBehavior::Block;
|
||||
|
||||
tokio::select! {
|
||||
res = tokio::time::timeout(timeout_dur, tool_fut) => {
|
||||
// 正常完成或超时
|
||||
// executor.rs Phase 3c: 逐批次执行
|
||||
for batch in &batches {
|
||||
if batch.is_parallel {
|
||||
// 并行批次: FuturesUnordered 内并发执行
|
||||
let mut exec_futs: FuturesUnordered<_> = batch.calls.iter()
|
||||
.map(|prep| execute_single_tool(...))
|
||||
.collect();
|
||||
while let Some(result) = exec_futs.next().await {
|
||||
process_single_result(...).await;
|
||||
}
|
||||
_ = cancel_fut => {
|
||||
// 仅当 !is_blocking 时此分支可达
|
||||
// Block 工具的 cancel_fut loop 不 break
|
||||
} else {
|
||||
// 串行批次: 逐个执行（非并发安全工具独占）
|
||||
for prep in &batch.calls {
|
||||
let result = execute_single_tool(...).await;
|
||||
process_single_result(...).await;
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
```
|
||||
|
||||
关键特性：
|
||||
- 所有工具放入同一个 `FuturesUnordered`，不区分串行/并行批次
|
||||
- `is_concurrency_safe` 声明为语义标记（引导 LLM 并发调用），执行时全部并发
|
||||
- 实际串行化依赖工具内部的 mutex/文件锁
|
||||
- `ToolPartitioner::partition()` 将工具调用分为并行批次和串行批次
|
||||
- 并行批次内使用 `FuturesUnordered` 最大化并发
|
||||
- 串行批次内逐个执行（如 `run_bash`、`file_write` 独占）
|
||||
- 连续的并发安全工具自动合并到一个并行批次
|
||||
- `is_concurrency_safe(args)` 是**输入感知**的：同一工具可能因参数不同而安全属性不同
|
||||
- `InterruptBehavior::Block` 保护写入操作不被用户取消打断
|
||||
|
||||
---
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
-- Migration: Session Rewind Support
|
||||
--
|
||||
-- 添加会话回退 (undo) 所需的软删除基础设施。
|
||||
-- 参考 Hermes-Agent hermes_state.py active=0 模型。
|
||||
--
|
||||
-- active 列:
|
||||
-- 1 = 活跃消息（对 LLM 可见，默认值）
|
||||
-- 0 = 软删除（rewind 操作移除，LLM 不可见，保留用于审计）
|
||||
--
|
||||
-- rewind_count:
|
||||
-- 单调递增计数器，记录会话被回退的次数
|
||||
|
||||
-- 1. agent_messages: 添加 active 列
|
||||
ALTER TABLE agent_messages ADD COLUMN active INTEGER NOT NULL DEFAULT 1;
|
||||
|
||||
-- 2. agent_sessions: 添加 rewind_count
|
||||
ALTER TABLE agent_sessions ADD COLUMN rewind_count INTEGER NOT NULL DEFAULT 0;
|
||||
|
||||
-- 3. 加速 active=1 过滤的索引（覆盖最常用的查询模式）
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_messages_active ON agent_messages(session_id, active);
|
||||
@@ -0,0 +1,13 @@
|
||||
-- Migration: Session Branch Support
|
||||
--
|
||||
-- 添加会话分叉所需的基础设施。
|
||||
-- 参考 Hermes-Agent parent_session_id 模型。
|
||||
--
|
||||
-- parent_session_id: 分叉的源会话 ID（NULL = 根会话）
|
||||
-- branched_from: 分叉的消息 ID（在源会话中的分叉点）
|
||||
|
||||
-- 1. agent_sessions: 添加 parent_session_id
|
||||
ALTER TABLE agent_sessions ADD COLUMN parent_session_id TEXT REFERENCES agent_sessions(session_id);
|
||||
|
||||
-- 2. agent_sessions: 添加分支元数据 JSON（存储 {branched_from, branch_reason} 等）
|
||||
ALTER TABLE agent_sessions ADD COLUMN branch_metadata TEXT;
|
||||
+179
-19
@@ -21,6 +21,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::hooks::{HookRegistry, PostCompactContext, PreCompactContext};
|
||||
use collapse::{CollapseLog, CollapseMethod};
|
||||
|
||||
/// 递归守卫：防止压缩内部触发的 LLM 调用再次触发压缩。
|
||||
static COMPACTING: AtomicBool = AtomicBool::new(false);
|
||||
@@ -259,8 +260,42 @@ fn inject_identity_block(messages: &mut Vec<ChatMessage>) {
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用 LLM 生成对话摘要。
|
||||
async fn generate_summary(to_summarize: &[ChatMessage], llm: &LlmClient) -> Result<String, String> {
|
||||
/// 摘要标记前缀，用于识别已有的压缩摘要。
|
||||
const SUMMARY_MARKER: &str = "[历史对话摘要]";
|
||||
const COMPRESSION_LOG_MARKER: &str = "[上下文压缩历史]";
|
||||
|
||||
/// 从消息列表中提取已有的历史摘要内容。
|
||||
/// 返回 None 表示无已有摘要。
|
||||
fn extract_prior_summary(messages: &[ChatMessage]) -> Option<String> {
|
||||
// 从后往前找最近的一份摘要（可能有多次压缩）
|
||||
for msg in messages.iter().rev() {
|
||||
if let Some(ref content) = msg.content {
|
||||
if content.starts_with(SUMMARY_MARKER) || content.starts_with(COMPRESSION_LOG_MARKER) {
|
||||
// 剥离标记，提取纯摘要内容
|
||||
let summary_body = if let Some(pos) = content.find('\n') {
|
||||
content[pos + 1..].trim().to_string()
|
||||
} else {
|
||||
content.clone()
|
||||
};
|
||||
if !summary_body.is_empty() && summary_body.len() < 2000 {
|
||||
return Some(summary_body);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 使用 LLM 生成对话摘要，融合已有摘要实现迭代信息保留。
|
||||
///
|
||||
/// 参考 Hermes context_compressor.py 的迭代摘要设计：
|
||||
/// 如果提供了 `prior_summary`，LLM 会将其作为已有背景融入新摘要，
|
||||
/// 确保多次压缩后历史信息不丢失。
|
||||
async fn generate_summary(
|
||||
to_summarize: &[ChatMessage],
|
||||
prior_summary: Option<&str>,
|
||||
llm: &LlmClient,
|
||||
) -> Result<String, String> {
|
||||
let summary_content: String = to_summarize
|
||||
.iter()
|
||||
.filter_map(|m| {
|
||||
@@ -278,10 +313,24 @@ async fn generate_summary(to_summarize: &[ChatMessage], llm: &LlmClient) -> Resu
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let summary_prompt = format!(
|
||||
"请用简洁的中文总结以下对话历史的要点（不超过500字）：\n\n{}",
|
||||
summary_content
|
||||
);
|
||||
// ── 迭代摘要融合 ──
|
||||
// 如果已有之前的摘要，将其作为上下文传给 LLM，要求新摘要融合旧内容
|
||||
let summary_prompt = if let Some(prior) = prior_summary {
|
||||
format!(
|
||||
"你是一个对话摘要助手。请将以下信息合并为一份简洁的中文摘要（不超过500字）。\n\n\
|
||||
⚠️ 重要：需要同时保留以下两部分的要点：\n\
|
||||
1. 【已有历史摘要】—— 之前压缩时已经总结过的内容\n\
|
||||
2. 【本轮新增对话】—— 本轮的对话记录\n\n\
|
||||
【已有历史摘要】\n{}\n\n\
|
||||
【本轮新增对话】\n{}",
|
||||
prior, summary_content
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"请用简洁的中文总结以下对话历史的要点（不超过500字）：\n\n{}",
|
||||
summary_content
|
||||
)
|
||||
};
|
||||
|
||||
llm.chat_completion(
|
||||
"你是一个对话摘要助手。请提取对话的关键信息和结论。",
|
||||
@@ -290,7 +339,12 @@ async fn generate_summary(to_summarize: &[ChatMessage], llm: &LlmClient) -> Resu
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("[Compact] 上下文摘要生成失败: {}，尝试激进压缩", e);
|
||||
format!("[历史摘要] 此前进行了 {} 轮对话交互", to_summarize.len())
|
||||
// 失败时保留旧摘要（如果有）+ 新增轮数
|
||||
if let Some(prior) = prior_summary {
|
||||
format!("{}（后续进行了 {} 轮对话）", prior, to_summarize.len())
|
||||
} else {
|
||||
format!("[历史摘要] 此前进行了 {} 轮对话交互", to_summarize.len())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -309,7 +363,7 @@ async fn compress_with_fallback(
|
||||
return;
|
||||
}
|
||||
|
||||
// Layer 2: auto_compact（LLM 摘要）
|
||||
// Layer 2: auto_compact（LLM 摘要，含迭代融合）
|
||||
let system_msg = messages.first().cloned();
|
||||
let cut_point = find_safe_cut_point(messages, 10);
|
||||
let to_summarize = &messages[1..cut_point];
|
||||
@@ -317,7 +371,11 @@ async fn compress_with_fallback(
|
||||
return;
|
||||
}
|
||||
|
||||
let summary = match generate_summary(to_summarize, llm).await {
|
||||
// ── 迭代摘要融合 ──
|
||||
// 在切割前提取已有摘要，传给 LLM 以实现迭代融合
|
||||
let prior_summary = extract_prior_summary(messages);
|
||||
|
||||
let summary = match generate_summary(to_summarize, prior_summary.as_deref(), llm).await {
|
||||
Ok(s) => s,
|
||||
Err(fallback) => fallback,
|
||||
};
|
||||
@@ -359,12 +417,34 @@ pub async fn compress_context(
|
||||
}
|
||||
|
||||
/// 带 Hook 的上下文压缩变体。如果提供了 HookRegistry，会在压缩前后触发事件。
|
||||
/// 如果提供了 CollapseLog，会在每次压缩后记录 commit。
|
||||
pub async fn compress_context_with_hooks(
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
llm: &LlmClient,
|
||||
context_char_limit: usize,
|
||||
session_id: &str,
|
||||
hook_registry: Option<&HookRegistry>,
|
||||
) {
|
||||
compress_context_with_hooks_and_log(
|
||||
messages,
|
||||
llm,
|
||||
context_char_limit,
|
||||
session_id,
|
||||
hook_registry,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// 带 Hook 和 CollapseLog 的上下文压缩变体。
|
||||
/// 压缩后自动记录 commit 到 CollapseLog，并在溢出时注入合并摘要。
|
||||
pub async fn compress_context_with_hooks_and_log(
|
||||
messages: &mut Vec<ChatMessage>,
|
||||
llm: &LlmClient,
|
||||
context_char_limit: usize,
|
||||
session_id: &str,
|
||||
hook_registry: Option<&HookRegistry>,
|
||||
collapse_log: Option<&std::sync::Mutex<CollapseLog>>,
|
||||
) {
|
||||
if messages.len() <= 4 {
|
||||
return;
|
||||
@@ -393,21 +473,67 @@ pub async fn compress_context_with_hooks(
|
||||
// 执行多层回退压缩（注：完整会话历史已在 agent_messages 表中持久化，无需额外 transcript 快照）
|
||||
compress_with_fallback(messages, llm, context_char_limit).await;
|
||||
|
||||
let after_count = messages.len();
|
||||
let method = if after_count < before_count / 2 {
|
||||
"llm_summary"
|
||||
} else if after_count < before_count {
|
||||
"snip_or_micro"
|
||||
} else {
|
||||
"none"
|
||||
};
|
||||
|
||||
// ── CollapseLog 记录 ──
|
||||
if let Some(log) = collapse_log {
|
||||
if after_count < before_count {
|
||||
// 压缩发生了：记录 commit
|
||||
if let Ok(mut log_guard) = log.lock() {
|
||||
let collapse_method = match method {
|
||||
"llm_summary" => CollapseMethod::LlmSummary,
|
||||
_ => CollapseMethod::MicroCompact,
|
||||
};
|
||||
let removed_range = (after_count, before_count);
|
||||
let summary = format!(
|
||||
"压缩 {} → {} 条消息 (方法: {}, 节省 ~{} tokens)",
|
||||
before_count,
|
||||
after_count,
|
||||
method,
|
||||
est_tokens.saturating_sub(rough_estimate_tokens(messages))
|
||||
);
|
||||
log_guard.commit(collapse_method, removed_range, summary);
|
||||
|
||||
// 溢出检查：如果 commits 过多，注入合并摘要
|
||||
if log_guard.should_overflow() {
|
||||
if let Some(overflow_summary) = log_guard.overflow() {
|
||||
info!(
|
||||
"[Compact] CollapseLog 溢出，注入合并摘要 ({} chars)",
|
||||
overflow_summary.len()
|
||||
);
|
||||
// 注入到消息列表头部（系统消息之后）
|
||||
let insert_pos = if messages
|
||||
.first()
|
||||
.is_some_and(|m| m.role == MessageRole::System)
|
||||
{
|
||||
1
|
||||
} else {
|
||||
0
|
||||
};
|
||||
messages.insert(
|
||||
insert_pos,
|
||||
ChatMessage::user(format!("[上下文压缩历史]\n{}", overflow_summary)),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// OnPostCompact hook
|
||||
if let Some(registry) = hook_registry {
|
||||
registry
|
||||
.run_on_post_compact(&PostCompactContext {
|
||||
session_id: session_id.to_string(),
|
||||
new_message_count: messages.len(),
|
||||
compression_method: if messages.len() < before_count / 2 {
|
||||
"llm_summary"
|
||||
} else if messages.len() < before_count {
|
||||
// snip_compact 会产生占位消息但保留尾部，micro_compact 替换内容
|
||||
"snip_or_micro"
|
||||
} else {
|
||||
"none"
|
||||
}
|
||||
.to_string(),
|
||||
new_message_count: after_count,
|
||||
compression_method: method.to_string(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
@@ -639,6 +765,40 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_prior_summary_finds_existing() {
|
||||
let messages = vec![
|
||||
ChatMessage::system("System"),
|
||||
ChatMessage::user("[历史对话摘要]\n步骤1-3: 搜索了黑洞相关文献。"),
|
||||
ChatMessage::user("继续研究"),
|
||||
];
|
||||
let prior = extract_prior_summary(&messages);
|
||||
assert!(prior.is_some());
|
||||
assert!(prior.unwrap().contains("黑洞"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_prior_summary_ignores_compression_log() {
|
||||
let messages = vec![
|
||||
ChatMessage::system("System"),
|
||||
ChatMessage::user("[上下文压缩历史]\n第一次压缩: 50→25 条消息"),
|
||||
ChatMessage::user("继续"),
|
||||
];
|
||||
let prior = extract_prior_summary(&messages);
|
||||
assert!(prior.is_some());
|
||||
assert!(prior.unwrap().contains("第一次压缩"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_prior_summary_none_when_no_marker() {
|
||||
let messages = vec![
|
||||
ChatMessage::system("System"),
|
||||
ChatMessage::user("普通用户消息"),
|
||||
];
|
||||
let prior = extract_prior_summary(&messages);
|
||||
assert!(prior.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_snip_compact_respects_tool_pairing() {
|
||||
// 构建 tool_call/tool_result 配对靠近切割点的场景
|
||||
|
||||
@@ -1,994 +0,0 @@
|
||||
// src/agent/hooks.rs
|
||||
//
|
||||
// Agent 生命周期 Hooks 系统。
|
||||
// 参考 Claude Code 的 PreToolUse / PostToolUse / Stop hooks 设计，
|
||||
// 提供可扩展的事件回调链，支持：
|
||||
// - OnSessionStart — 会话创建/恢复时
|
||||
// - PreToolUse — 工具执行前（可拦截/阻止/修改输入/注入上下文）
|
||||
// - PostToolUse — 工具执行后（审计日志、指标采集、输出修改）
|
||||
// - OnStepComplete — 每步结束（指标更新、上下文检查）
|
||||
// - OnSessionStop — 会话终止（清理、持久化）
|
||||
// - OnSubagentStart — 子代理启动时
|
||||
// - OnSubagentStop — 子代理停止时
|
||||
// - OnPreCompact — 上下文压缩前
|
||||
// - OnPostCompact — 上下文压缩后
|
||||
//
|
||||
// P3 增强（参考 Claude Code hooks 协议）：
|
||||
// - PreToolUseAction 支持 MutateInput（修改工具参数）和 PermissionRequired
|
||||
// - PostToolUseAction 支持 MutateOutput（修改工具输出）
|
||||
// - 扩展生命周期事件（SubagentStart/Stop, PreCompact/PostCompact）
|
||||
// - HookRegistry 返回累积的 additional_context 和 mutate_output
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sqlx::SqlitePool;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::terminal::TurnTerminal;
|
||||
|
||||
// ── Hook Contexts ──
|
||||
|
||||
/// 会话启动上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionStartContext {
|
||||
pub session_id: String,
|
||||
pub turn_index: i32,
|
||||
pub is_resume: bool,
|
||||
}
|
||||
|
||||
/// PreToolUse hook 上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreToolUseContext {
|
||||
pub session_id: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
pub step: usize,
|
||||
}
|
||||
|
||||
/// PostToolUse hook 上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostToolUseContext {
|
||||
pub session_id: String,
|
||||
pub agent_name: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
pub output_content: String,
|
||||
pub is_error: bool,
|
||||
pub step: usize,
|
||||
pub elapsed_ms: u64,
|
||||
}
|
||||
|
||||
/// Step 完成上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StepCompleteContext {
|
||||
pub session_id: String,
|
||||
pub step: usize,
|
||||
pub max_steps: usize,
|
||||
pub messages_count: usize,
|
||||
pub estimated_tokens: usize,
|
||||
pub token_limit: usize,
|
||||
}
|
||||
|
||||
/// Session 终止上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionStopContext<'a> {
|
||||
pub session_id: String,
|
||||
pub terminal: &'a TurnTerminal,
|
||||
pub total_steps: usize,
|
||||
}
|
||||
|
||||
/// 子代理启动上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubagentStartContext {
|
||||
pub parent_session_id: String,
|
||||
pub subagent_name: String,
|
||||
pub prompt: String,
|
||||
}
|
||||
|
||||
/// 子代理停止上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubagentStopContext {
|
||||
pub parent_session_id: String,
|
||||
pub subagent_name: String,
|
||||
pub result_summary: String,
|
||||
pub steps: usize,
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
/// 压缩前上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreCompactContext {
|
||||
pub session_id: String,
|
||||
pub message_count: usize,
|
||||
pub estimated_tokens: usize,
|
||||
}
|
||||
|
||||
/// 压缩后上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostCompactContext {
|
||||
pub session_id: String,
|
||||
pub new_message_count: usize,
|
||||
pub compression_method: String,
|
||||
}
|
||||
|
||||
// ── Hook Actions ──
|
||||
|
||||
/// PreToolUse hook 返回的增强动作。
|
||||
/// 支持：允许、阻止、修改输入、权限请求。
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PreToolUseAction {
|
||||
/// 允许继续执行（默认）
|
||||
Continue,
|
||||
/// 阻止执行，附带原因
|
||||
Block { reason: String },
|
||||
/// 允许执行但修改输入参数或注入附加上下文
|
||||
MutateInput {
|
||||
updated_args: serde_json::Value,
|
||||
additional_context: Option<String>,
|
||||
},
|
||||
/// 需要权限决策
|
||||
PermissionRequired {
|
||||
permission: String,
|
||||
tool_name: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl PreToolUseAction {
|
||||
pub fn is_blocked(&self) -> bool {
|
||||
matches!(self, PreToolUseAction::Block { .. })
|
||||
}
|
||||
|
||||
pub fn block_reason(&self) -> Option<&str> {
|
||||
match self {
|
||||
PreToolUseAction::Block { reason } => Some(reason),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn updated_args(&self) -> Option<&serde_json::Value> {
|
||||
match self {
|
||||
PreToolUseAction::MutateInput { updated_args, .. } => Some(updated_args),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn additional_context(&self) -> Option<&str> {
|
||||
match self {
|
||||
PreToolUseAction::MutateInput {
|
||||
additional_context, ..
|
||||
} => additional_context.as_deref(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 向后兼容类型别名 — 旧代码用 HookAction::Continue / HookAction::Block 仍可编译
|
||||
pub type HookAction = PreToolUseAction;
|
||||
|
||||
/// PostToolUse hook 返回的动作。
|
||||
/// 支持：保持输出不变、修改输出内容。
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PostToolUseAction {
|
||||
/// 保持原输出不变（默认）
|
||||
Continue,
|
||||
/// 修改输出内容
|
||||
MutateOutput { updated_content: String },
|
||||
}
|
||||
|
||||
// ── Metrics Data ──
|
||||
|
||||
/// 可查询的运行指标快照
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct MetricsData {
|
||||
pub tool_call_counts: HashMap<String, usize>,
|
||||
pub total_steps: usize,
|
||||
pub total_errors: usize,
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
// ── Hook Trait ──
|
||||
|
||||
/// Agent 生命周期 Hook trait。
|
||||
/// 所有方法都有默认空实现，只需覆写关心的 hook 点。
|
||||
#[async_trait]
|
||||
pub trait AgentHook: Send + Sync {
|
||||
/// Hook 名称（用于日志和调试）
|
||||
fn name(&self) -> &str;
|
||||
|
||||
// ── 原有 5 个生命周期事件 ──
|
||||
|
||||
/// 会话创建/恢复时调用。
|
||||
async fn on_session_start(&self, _ctx: &SessionStartContext) {}
|
||||
|
||||
/// 工具执行前调用。可返回 Continue/Block/MutateInput/PermissionRequired。
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
|
||||
/// 工具执行后调用。可返回 Continue 或 MutateOutput。
|
||||
async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
/// 每个 ReAct step 完成后调用。
|
||||
async fn on_step_complete(&self, _ctx: &StepCompleteContext) {}
|
||||
|
||||
/// 会话终止时调用。
|
||||
async fn on_session_stop(&self, _ctx: &SessionStopContext<'_>) {}
|
||||
|
||||
// ── 新增 4 个生命周期事件（默认 no-op） ──
|
||||
|
||||
/// 子代理启动时调用。
|
||||
async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {}
|
||||
|
||||
/// 子代理停止时调用。
|
||||
async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {}
|
||||
|
||||
/// 上下文压缩前调用。
|
||||
async fn on_pre_compact(&self, _ctx: &PreCompactContext) {}
|
||||
|
||||
/// 上下文压缩后调用。
|
||||
async fn on_post_compact(&self, _ctx: &PostCompactContext) {}
|
||||
}
|
||||
|
||||
// ── Hook Registry ──
|
||||
|
||||
/// PreToolUse 聚合结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreToolUseResult {
|
||||
/// 最终动作（第一个 Block 获胜）
|
||||
pub action: PreToolUseAction,
|
||||
/// 累积的 additional_context（所有 MutateInput 的上下文拼接）
|
||||
pub additional_context: Option<String>,
|
||||
/// 最终的工具参数（应用了最后一个 MutateInput 的修改）
|
||||
pub final_args: serde_json::Value,
|
||||
}
|
||||
|
||||
impl PreToolUseResult {
|
||||
/// 是否有 hook 请求了权限确认
|
||||
pub fn is_permission_required(&self) -> bool {
|
||||
matches!(self.action, PreToolUseAction::PermissionRequired { .. })
|
||||
}
|
||||
|
||||
/// 获取权限确认的详情（permission 描述, tool_name）
|
||||
pub fn permission_info(&self) -> Option<(&str, &str)> {
|
||||
match &self.action {
|
||||
PreToolUseAction::PermissionRequired {
|
||||
permission,
|
||||
tool_name,
|
||||
} => Some((permission.as_str(), tool_name.as_str())),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// PostToolUse 聚合结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostToolUseResult {
|
||||
/// 最终输出内容（应用了最后一个 MutateOutput 的修改）
|
||||
pub final_content: String,
|
||||
}
|
||||
|
||||
/// Hook 注册表，管理所有已注册的 hook 并按序调用
|
||||
pub struct HookRegistry {
|
||||
hooks: Vec<Box<dyn AgentHook>>,
|
||||
}
|
||||
|
||||
impl Default for HookRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl HookRegistry {
|
||||
/// 创建空的注册表
|
||||
pub fn new() -> Self {
|
||||
HookRegistry { hooks: Vec::new() }
|
||||
}
|
||||
|
||||
/// 创建包含所有内置 hooks 的注册表。
|
||||
///
|
||||
/// 参数：
|
||||
/// - `db`: 数据库连接池（供 AuditLogHook 持久化）
|
||||
/// - `cancelled_runs`: 取消状态集合（供 CancellationHook 检查）
|
||||
/// - `metrics_data`: 可选的共享指标数据引用。提供时复用已有的 MetricsData，
|
||||
/// 使得 AgentRuntime.get_metrics() 能查询到实际运行数据。
|
||||
pub fn with_builtins(
|
||||
db: SqlitePool,
|
||||
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
|
||||
metrics_data: Option<Arc<std::sync::Mutex<MetricsData>>>,
|
||||
) -> Self {
|
||||
let mut registry = Self::new();
|
||||
registry.add(Box::new(CancellationHook::new(cancelled_runs)));
|
||||
// 如果提供了共享的 metrics_data，使用它；否则创建新的
|
||||
let metrics_hook = match metrics_data {
|
||||
Some(data) => MetricsHook::from_arc(data),
|
||||
None => MetricsHook::new(),
|
||||
};
|
||||
registry.add(Box::new(metrics_hook));
|
||||
registry.add(Box::new(AuditLogHook::new(db)));
|
||||
registry
|
||||
}
|
||||
|
||||
/// 注册一个 hook
|
||||
pub fn add(&mut self, hook: Box<dyn AgentHook>) {
|
||||
info!("[Hooks] 注册 hook: {}", hook.name());
|
||||
self.hooks.push(hook);
|
||||
}
|
||||
|
||||
/// 获取所有 hooks 的不可变引用
|
||||
pub fn all(&self) -> &[Box<dyn AgentHook>] {
|
||||
&self.hooks
|
||||
}
|
||||
|
||||
// ── 便捷调用方法 ──
|
||||
|
||||
/// 调用所有 on_session_start hooks
|
||||
pub async fn run_on_session_start(&self, ctx: &SessionStartContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_session_start(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 pre_tool_use hooks。
|
||||
///
|
||||
/// 返回聚合结果：
|
||||
/// - 遇到第一个 Block 时短路，返回该 Block
|
||||
/// - MutateInput 累积 additional_context 并更新 final_args
|
||||
/// - PermissionRequired 记录但继续执行（暂时视为 Continue）
|
||||
pub async fn run_pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseResult {
|
||||
let mut accumulated_context = String::new();
|
||||
let mut final_args = ctx.tool_args.clone();
|
||||
let mut final_action = PreToolUseAction::Continue;
|
||||
|
||||
for hook in &self.hooks {
|
||||
let action = hook.pre_tool_use(ctx).await;
|
||||
match &action {
|
||||
PreToolUseAction::Block { reason } => {
|
||||
warn!(
|
||||
"[Hooks] {} 阻止了工具 {} 的执行: {}",
|
||||
hook.name(),
|
||||
ctx.tool_name,
|
||||
reason
|
||||
);
|
||||
return PreToolUseResult {
|
||||
action,
|
||||
additional_context: None,
|
||||
final_args: ctx.tool_args.clone(),
|
||||
};
|
||||
}
|
||||
PreToolUseAction::MutateInput {
|
||||
updated_args,
|
||||
additional_context,
|
||||
} => {
|
||||
info!(
|
||||
"[Hooks] {} 修改了工具 {} 的输入参数",
|
||||
hook.name(),
|
||||
ctx.tool_name
|
||||
);
|
||||
final_args = updated_args.clone();
|
||||
if let Some(ctx_str) = additional_context {
|
||||
if !accumulated_context.is_empty() {
|
||||
accumulated_context.push('\n');
|
||||
}
|
||||
accumulated_context.push_str(ctx_str);
|
||||
}
|
||||
}
|
||||
PreToolUseAction::PermissionRequired { .. } => {
|
||||
info!(
|
||||
"[Hooks] {} 请求了工具 {} 的权限检查",
|
||||
hook.name(),
|
||||
ctx.tool_name
|
||||
);
|
||||
}
|
||||
PreToolUseAction::Continue => {}
|
||||
}
|
||||
// Continue 不应覆盖已设置的 meaningful action（Block/MutateInput/PermissionRequired）
|
||||
if !matches!(action, PreToolUseAction::Continue) {
|
||||
final_action = action;
|
||||
}
|
||||
}
|
||||
|
||||
let ctx_opt = if accumulated_context.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(accumulated_context)
|
||||
};
|
||||
|
||||
PreToolUseResult {
|
||||
action: final_action,
|
||||
additional_context: ctx_opt,
|
||||
final_args,
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 post_tool_use hooks（全部执行，不会短路）。
|
||||
/// 返回聚合的最终输出内容。
|
||||
pub async fn run_post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseResult {
|
||||
let mut final_content = ctx.output_content.clone();
|
||||
|
||||
for hook in &self.hooks {
|
||||
let action = hook.post_tool_use(ctx).await;
|
||||
match action {
|
||||
PostToolUseAction::MutateOutput { updated_content } => {
|
||||
info!(
|
||||
"[Hooks] {} 修改了工具 {} 的输出",
|
||||
hook.name(),
|
||||
ctx.tool_name
|
||||
);
|
||||
final_content = updated_content;
|
||||
}
|
||||
PostToolUseAction::Continue => {}
|
||||
}
|
||||
}
|
||||
|
||||
PostToolUseResult { final_content }
|
||||
}
|
||||
|
||||
/// 调用所有 on_step_complete hooks
|
||||
pub async fn run_on_step_complete(&self, ctx: &StepCompleteContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_step_complete(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_session_stop hooks
|
||||
pub async fn run_on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_session_stop(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_subagent_start hooks
|
||||
pub async fn run_on_subagent_start(&self, ctx: &SubagentStartContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_subagent_start(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_subagent_stop hooks
|
||||
pub async fn run_on_subagent_stop(&self, ctx: &SubagentStopContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_subagent_stop(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_pre_compact hooks
|
||||
pub async fn run_on_pre_compact(&self, ctx: &PreCompactContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_pre_compact(ctx).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_post_compact hooks
|
||||
pub async fn run_on_post_compact(&self, ctx: &PostCompactContext) {
|
||||
for hook in &self.hooks {
|
||||
hook.on_post_compact(ctx).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Built-in Hooks ──
|
||||
|
||||
/// 取消检查 Hook — 在每次工具执行前检查用户是否中止了会话。
|
||||
pub struct CancellationHook {
|
||||
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
|
||||
}
|
||||
|
||||
impl CancellationHook {
|
||||
pub fn new(cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>) -> Self {
|
||||
CancellationHook { cancelled_runs }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for CancellationHook {
|
||||
fn name(&self) -> &str {
|
||||
"CancellationHook"
|
||||
}
|
||||
|
||||
async fn pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
if let Ok(runs) = self.cancelled_runs.lock() {
|
||||
if runs.contains(&ctx.session_id) {
|
||||
warn!(
|
||||
"[CancellationHook] 会话 {} 已被用户取消，阻止工具 {} 执行",
|
||||
ctx.session_id, ctx.tool_name
|
||||
);
|
||||
return PreToolUseAction::Block {
|
||||
reason: "用户已手动中止执行".to_string(),
|
||||
};
|
||||
}
|
||||
}
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
// 清理取消状态
|
||||
if let Ok(mut runs) = self.cancelled_runs.lock() {
|
||||
runs.remove(&ctx.session_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 指标采集 Hook — 自动收集工具调用统计，支持快照查询
|
||||
pub struct MetricsHook {
|
||||
data: Arc<std::sync::Mutex<MetricsData>>,
|
||||
}
|
||||
|
||||
impl Default for MetricsHook {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl MetricsHook {
|
||||
pub fn new() -> Self {
|
||||
MetricsHook {
|
||||
data: Arc::new(std::sync::Mutex::new(MetricsData::default())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从已有的 Arc<Mutex<MetricsData>> 创建（共享数据引用）
|
||||
pub fn from_arc(data: Arc<std::sync::Mutex<MetricsData>>) -> Self {
|
||||
MetricsHook { data }
|
||||
}
|
||||
|
||||
/// 返回当前指标快照（锁异常时返回 None）
|
||||
pub fn snapshot(&self) -> Option<MetricsData> {
|
||||
self.data.lock().ok().map(|d| d.clone())
|
||||
}
|
||||
|
||||
/// 获取 Arc 引用，供外部持有
|
||||
pub fn data_arc(&self) -> Arc<std::sync::Mutex<MetricsData>> {
|
||||
self.data.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for MetricsHook {
|
||||
fn name(&self) -> &str {
|
||||
"MetricsHook"
|
||||
}
|
||||
|
||||
async fn on_session_start(&self, ctx: &SessionStartContext) {
|
||||
if let Ok(mut data) = self.data.lock() {
|
||||
data.session_id = Some(ctx.session_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
info!(
|
||||
"[Metrics] step={} tool={} is_error={} elapsed={}ms",
|
||||
ctx.step, ctx.tool_name, ctx.is_error, ctx.elapsed_ms
|
||||
);
|
||||
if let Ok(mut data) = self.data.lock() {
|
||||
*data
|
||||
.tool_call_counts
|
||||
.entry(ctx.tool_name.clone())
|
||||
.or_insert(0) += 1;
|
||||
data.total_steps = ctx.step;
|
||||
if ctx.is_error {
|
||||
data.total_errors += 1;
|
||||
}
|
||||
}
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_step_complete(&self, ctx: &StepCompleteContext) {
|
||||
if ctx.step.is_multiple_of(3) {
|
||||
info!(
|
||||
"[Metrics] step {}/{} | messages={} | tokens≈{}/{}",
|
||||
ctx.step, ctx.max_steps, ctx.messages_count, ctx.estimated_tokens, ctx.token_limit
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
info!(
|
||||
"[Metrics] 会话 {} 结束: {} (total_steps={})",
|
||||
ctx.session_id,
|
||||
ctx.terminal.description(),
|
||||
ctx.total_steps
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 审计日志 Hook — 记录所有工具调用到 SQLite agent_audit_log 表
|
||||
pub struct AuditLogHook {
|
||||
db: SqlitePool,
|
||||
}
|
||||
|
||||
impl AuditLogHook {
|
||||
pub fn new(db: SqlitePool) -> Self {
|
||||
AuditLogHook { db }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for AuditLogHook {
|
||||
fn name(&self) -> &str {
|
||||
"AuditLogHook"
|
||||
}
|
||||
|
||||
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
let status = if ctx.is_error { "FAIL" } else { "OK" };
|
||||
let preview: String = ctx.output_content.chars().take(200).collect();
|
||||
|
||||
info!(
|
||||
"[Audit] session={} step={} tool={} status={} elapsed={}ms",
|
||||
ctx.session_id, ctx.step, ctx.tool_name, status, ctx.elapsed_ms
|
||||
);
|
||||
|
||||
let db = self.db.clone();
|
||||
let session_id = ctx.session_id.clone();
|
||||
let agent_name = ctx.agent_name.clone();
|
||||
let tool_name = ctx.tool_name.clone();
|
||||
let step = ctx.step;
|
||||
let elapsed = ctx.elapsed_ms;
|
||||
let preview_clone = preview.clone();
|
||||
|
||||
// Fire-and-forget 写入，不阻塞主循环
|
||||
tokio::spawn(async move {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_audit_log (session_id, step, tool_name, status, elapsed_ms, output_preview, agent_name) \
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(step as i32)
|
||||
.bind(&tool_name)
|
||||
.bind(status)
|
||||
.bind(elapsed as i32)
|
||||
.bind(&preview_clone)
|
||||
.bind(&agent_name)
|
||||
.execute(&db)
|
||||
.await;
|
||||
});
|
||||
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
info!(
|
||||
"[Audit] 会话 {} 终止原因: {}",
|
||||
ctx.session_id,
|
||||
ctx.terminal.description()
|
||||
);
|
||||
|
||||
let db = self.db.clone();
|
||||
let session_id = ctx.session_id.clone();
|
||||
let total_steps = ctx.total_steps;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_audit_log (session_id, step, tool_name, status, elapsed_ms, output_preview) \
|
||||
VALUES (?, ?, 'session', 'SESSION_STOP', 0, ?)",
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(total_steps as i32)
|
||||
.bind(format!("会话终止，共 {} 步", total_steps))
|
||||
.execute(&db)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ──
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
struct TestHook {
|
||||
name: String,
|
||||
pre_called: std::sync::Mutex<bool>,
|
||||
}
|
||||
|
||||
impl TestHook {
|
||||
fn new(name: &str) -> Self {
|
||||
TestHook {
|
||||
name: name.to_string(),
|
||||
pre_called: std::sync::Mutex::new(false),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for TestHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
*self.pre_called.lock().unwrap() = true;
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_registry_runs_all_hooks() {
|
||||
let mut registry = HookRegistry::new();
|
||||
let hook1 = TestHook::new("test1");
|
||||
let hook2 = TestHook::new("test2");
|
||||
registry.add(Box::new(hook1));
|
||||
registry.add(Box::new(hook2));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert!(!result.action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_blocking_hook_stops_chain() {
|
||||
struct BlockingHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for BlockingHook {
|
||||
fn name(&self) -> &str {
|
||||
"blocker"
|
||||
}
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::Block {
|
||||
reason: "test block".into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(BlockingHook));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert!(result.action.is_blocked());
|
||||
assert_eq!(result.action.block_reason(), Some("test block"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mutate_input_accumulates_context() {
|
||||
struct MutateHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for MutateHook {
|
||||
fn name(&self) -> &str {
|
||||
"mutator"
|
||||
}
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::MutateInput {
|
||||
updated_args: serde_json::json!({"key": "modified"}),
|
||||
additional_context: Some("injected context".to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(MutateHook));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({"key": "original"}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert_eq!(result.final_args, serde_json::json!({"key": "modified"}));
|
||||
assert_eq!(
|
||||
result.additional_context,
|
||||
Some("injected context".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_post_tool_use_mutate_output() {
|
||||
struct MutateOutputHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for MutateOutputHook {
|
||||
fn name(&self) -> &str {
|
||||
"output_mutator"
|
||||
}
|
||||
async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content: "modified output".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(MutateOutputHook));
|
||||
|
||||
let ctx = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "original output".into(),
|
||||
is_error: false,
|
||||
step: 1,
|
||||
elapsed_ms: 100,
|
||||
};
|
||||
|
||||
let result = registry.run_post_tool_use(&ctx).await;
|
||||
assert_eq!(result.final_content, "modified output");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancellation_hook_blocks_when_cancelled() {
|
||||
use std::collections::HashSet;
|
||||
let mut cancelled = HashSet::new();
|
||||
cancelled.insert("test_session".to_string());
|
||||
let cancelled_runs = Arc::new(std::sync::Mutex::new(cancelled));
|
||||
|
||||
let hook = CancellationHook::new(cancelled_runs);
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test_session".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let action = hook.pre_tool_use(&ctx).await;
|
||||
assert!(action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancellation_hook_allows_when_not_cancelled() {
|
||||
let cancelled_runs = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
|
||||
|
||||
let hook = CancellationHook::new(cancelled_runs);
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test_session".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let action = hook.pre_tool_use(&ctx).await;
|
||||
assert!(!action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_metrics_hook_accumulates_counts() {
|
||||
let hook = MetricsHook::new();
|
||||
|
||||
let ctx = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "result".into(),
|
||||
is_error: false,
|
||||
step: 1,
|
||||
elapsed_ms: 100,
|
||||
};
|
||||
hook.post_tool_use(&ctx).await;
|
||||
|
||||
let ctx2 = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "result2".into(),
|
||||
is_error: false,
|
||||
step: 2,
|
||||
elapsed_ms: 200,
|
||||
};
|
||||
hook.post_tool_use(&ctx2).await;
|
||||
|
||||
let snapshot = hook.snapshot().expect("snapshot should succeed in test");
|
||||
assert_eq!(snapshot.tool_call_counts.get("search_papers"), Some(&2));
|
||||
assert_eq!(snapshot.total_steps, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_session_start_hook_called() {
|
||||
struct StartTrackingHook {
|
||||
started: std::sync::Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for StartTrackingHook {
|
||||
fn name(&self) -> &str {
|
||||
"start_tracker"
|
||||
}
|
||||
async fn on_session_start(&self, ctx: &SessionStartContext) {
|
||||
self.started.lock().unwrap().push(ctx.session_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let hook = StartTrackingHook {
|
||||
started: std::sync::Mutex::new(Vec::new()),
|
||||
};
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(hook));
|
||||
|
||||
let ctx = SessionStartContext {
|
||||
session_id: "test_sid".into(),
|
||||
turn_index: 1,
|
||||
is_resume: false,
|
||||
};
|
||||
registry.run_on_session_start(&ctx).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_new_lifecycle_events_called() {
|
||||
struct LifecycleTracker {
|
||||
subagent_start: std::sync::Mutex<bool>,
|
||||
subagent_stop: std::sync::Mutex<bool>,
|
||||
pre_compact: std::sync::Mutex<bool>,
|
||||
post_compact: std::sync::Mutex<bool>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for LifecycleTracker {
|
||||
fn name(&self) -> &str {
|
||||
"lifecycle_tracker"
|
||||
}
|
||||
async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {
|
||||
*self.subagent_start.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {
|
||||
*self.subagent_stop.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_pre_compact(&self, _ctx: &PreCompactContext) {
|
||||
*self.pre_compact.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_post_compact(&self, _ctx: &PostCompactContext) {
|
||||
*self.post_compact.lock().unwrap() = true;
|
||||
}
|
||||
}
|
||||
|
||||
let tracker = LifecycleTracker {
|
||||
subagent_start: std::sync::Mutex::new(false),
|
||||
subagent_stop: std::sync::Mutex::new(false),
|
||||
pre_compact: std::sync::Mutex::new(false),
|
||||
post_compact: std::sync::Mutex::new(false),
|
||||
};
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(tracker));
|
||||
|
||||
registry
|
||||
.run_on_subagent_start(&SubagentStartContext {
|
||||
parent_session_id: "s1".into(),
|
||||
subagent_name: "sub".into(),
|
||||
prompt: "test".into(),
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_subagent_stop(&SubagentStopContext {
|
||||
parent_session_id: "s1".into(),
|
||||
subagent_name: "sub".into(),
|
||||
result_summary: "done".into(),
|
||||
steps: 3,
|
||||
is_error: false,
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_pre_compact(&PreCompactContext {
|
||||
session_id: "s1".into(),
|
||||
message_count: 50,
|
||||
estimated_tokens: 10000,
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_post_compact(&PostCompactContext {
|
||||
session_id: "s1".into(),
|
||||
new_message_count: 10,
|
||||
compression_method: "micro".into(),
|
||||
})
|
||||
.await;
|
||||
|
||||
// If no panic, all hooks were called successfully
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,265 @@
|
||||
// src/agent/hooks/builtins.rs
|
||||
//
|
||||
// 内置 Hooks — CancellationHook、MetricsHook、AuditLogHook。
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::{
|
||||
AgentHook, HookEvent, MetricsData, PostToolUseAction, PostToolUseContext, PreToolUseAction,
|
||||
PreToolUseContext, SessionStartContext, SessionStopContext, StepCompleteContext,
|
||||
};
|
||||
use tokio;
|
||||
|
||||
// ── Context Deduplicator ──
|
||||
|
||||
/// 内容哈希去重器 — 单 dispatch cycle 内防止重复上下文注入。
|
||||
///
|
||||
/// 使用内容哈希（`std::hash::DefaultHasher`），在同一个 dispatch cycle
|
||||
/// 内避免多个 hook 注入相同的上下文消息。不跨 step 去重（以允许合法的重复提醒）。
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ContextDeduplicator {
|
||||
seen: std::collections::HashSet<u64>,
|
||||
}
|
||||
|
||||
impl ContextDeduplicator {
|
||||
pub fn new() -> Self {
|
||||
ContextDeduplicator {
|
||||
seen: std::collections::HashSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 若 `content` 在此 cycle 中已见则返回 true。
|
||||
pub fn is_duplicate(&mut self, content: &str) -> bool {
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
content.hash(&mut hasher);
|
||||
let hash = hasher.finish();
|
||||
!self.seen.insert(hash)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Built-in Hooks ──
|
||||
|
||||
/// 取消检查 Hook — 在每次工具执行前检查用户是否中止了会话。
|
||||
pub struct CancellationHook {
|
||||
cancelled_runs: Arc<Mutex<HashSet<String>>>,
|
||||
}
|
||||
|
||||
impl CancellationHook {
|
||||
pub fn new(cancelled_runs: Arc<Mutex<HashSet<String>>>) -> Self {
|
||||
CancellationHook { cancelled_runs }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for CancellationHook {
|
||||
fn name(&self) -> &str {
|
||||
"CancellationHook"
|
||||
}
|
||||
|
||||
fn subscribed_events(&self) -> &[HookEvent] {
|
||||
&[HookEvent::PreToolUse, HookEvent::OnSessionStop]
|
||||
}
|
||||
|
||||
async fn pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
if let Ok(runs) = self.cancelled_runs.lock() {
|
||||
if runs.contains(&ctx.session_id) {
|
||||
warn!(
|
||||
"[CancellationHook] 会话 {} 已被用户取消，阻止工具 {} 执行",
|
||||
ctx.session_id, ctx.tool_name
|
||||
);
|
||||
return PreToolUseAction::Block {
|
||||
reason: "用户已手动中止执行".to_string(),
|
||||
};
|
||||
}
|
||||
}
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
// 清理取消状态
|
||||
if let Ok(mut runs) = self.cancelled_runs.lock() {
|
||||
runs.remove(&ctx.session_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 指标采集 Hook — 自动收集工具调用统计，支持快照查询
|
||||
pub struct MetricsHook {
|
||||
data: Arc<Mutex<MetricsData>>,
|
||||
}
|
||||
|
||||
impl Default for MetricsHook {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl MetricsHook {
|
||||
pub fn new() -> Self {
|
||||
MetricsHook {
|
||||
data: Arc::new(Mutex::new(MetricsData::default())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从已有的 Arc<Mutex<MetricsData>> 创建（共享数据引用）
|
||||
pub fn from_arc(data: Arc<Mutex<MetricsData>>) -> Self {
|
||||
MetricsHook { data }
|
||||
}
|
||||
|
||||
/// 返回当前指标快照（锁异常时返回 None）
|
||||
pub fn snapshot(&self) -> Option<MetricsData> {
|
||||
self.data.lock().ok().map(|d| d.clone())
|
||||
}
|
||||
|
||||
/// 获取 Arc 引用，供外部持有
|
||||
pub fn data_arc(&self) -> Arc<Mutex<MetricsData>> {
|
||||
self.data.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for MetricsHook {
|
||||
fn name(&self) -> &str {
|
||||
"MetricsHook"
|
||||
}
|
||||
|
||||
fn subscribed_events(&self) -> &[HookEvent] {
|
||||
&[
|
||||
HookEvent::OnSessionStart,
|
||||
HookEvent::PostToolUse,
|
||||
HookEvent::OnStepComplete,
|
||||
HookEvent::OnSessionStop,
|
||||
]
|
||||
}
|
||||
|
||||
async fn on_session_start(&self, ctx: &SessionStartContext) {
|
||||
if let Ok(mut data) = self.data.lock() {
|
||||
data.session_id = Some(ctx.session_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
info!(
|
||||
"[Metrics] step={} tool={} is_error={} elapsed={}ms",
|
||||
ctx.step, ctx.tool_name, ctx.is_error, ctx.elapsed_ms
|
||||
);
|
||||
if let Ok(mut data) = self.data.lock() {
|
||||
*data
|
||||
.tool_call_counts
|
||||
.entry(ctx.tool_name.clone())
|
||||
.or_insert(0) += 1;
|
||||
data.total_steps = ctx.step;
|
||||
if ctx.is_error {
|
||||
data.total_errors += 1;
|
||||
}
|
||||
}
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_step_complete(&self, ctx: &StepCompleteContext) {
|
||||
if ctx.step.is_multiple_of(3) {
|
||||
info!(
|
||||
"[Metrics] step {}/{} | messages={} | tokens≈{}/{}",
|
||||
ctx.step, ctx.max_steps, ctx.messages_count, ctx.estimated_tokens, ctx.token_limit
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
info!(
|
||||
"[Metrics] 会话 {} 结束: {} (total_steps={})",
|
||||
ctx.session_id,
|
||||
ctx.terminal.description(),
|
||||
ctx.total_steps
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 审计日志 Hook — 记录所有工具调用到 SQLite agent_audit_log 表
|
||||
pub struct AuditLogHook {
|
||||
db: SqlitePool,
|
||||
}
|
||||
|
||||
impl AuditLogHook {
|
||||
pub fn new(db: SqlitePool) -> Self {
|
||||
AuditLogHook { db }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for AuditLogHook {
|
||||
fn name(&self) -> &str {
|
||||
"AuditLogHook"
|
||||
}
|
||||
|
||||
fn subscribed_events(&self) -> &[HookEvent] {
|
||||
&[HookEvent::PostToolUse, HookEvent::OnSessionStop]
|
||||
}
|
||||
|
||||
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
let status = if ctx.is_error { "FAIL" } else { "OK" };
|
||||
let preview: String = ctx.output_content.chars().take(200).collect();
|
||||
|
||||
info!(
|
||||
"[Audit] session={} step={} tool={} status={} elapsed={}ms",
|
||||
ctx.session_id, ctx.step, ctx.tool_name, status, ctx.elapsed_ms
|
||||
);
|
||||
|
||||
let db = self.db.clone();
|
||||
let session_id = ctx.session_id.clone();
|
||||
let agent_name = ctx.agent_name.clone();
|
||||
let tool_name = ctx.tool_name.clone();
|
||||
let step = ctx.step;
|
||||
let elapsed = ctx.elapsed_ms;
|
||||
let preview_clone = preview.clone();
|
||||
|
||||
// Fire-and-forget 写入，不阻塞主循环
|
||||
tokio::spawn(async move {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_audit_log (session_id, step, tool_name, status, elapsed_ms, output_preview, agent_name) \
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(step as i32)
|
||||
.bind(&tool_name)
|
||||
.bind(status)
|
||||
.bind(elapsed as i32)
|
||||
.bind(&preview_clone)
|
||||
.bind(&agent_name)
|
||||
.execute(&db)
|
||||
.await;
|
||||
});
|
||||
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
info!(
|
||||
"[Audit] 会话 {} 终止原因: {}",
|
||||
ctx.session_id,
|
||||
ctx.terminal.description()
|
||||
);
|
||||
|
||||
let db = self.db.clone();
|
||||
let session_id = ctx.session_id.clone();
|
||||
let total_steps = ctx.total_steps;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO agent_audit_log (session_id, step, tool_name, status, elapsed_ms, output_preview) \
|
||||
VALUES (?, ?, 'session', 'SESSION_STOP', 0, ?)",
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(total_steps as i32)
|
||||
.bind(format!("会话终止，共 {} 步", total_steps))
|
||||
.execute(&db)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,790 @@
|
||||
// src/agent/hooks/dispatch.rs
|
||||
//!
|
||||
//! HookRegistry 调度方法——所有 run_* 生命周期事件分发。
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::future::join_all;
|
||||
use tokio::time::timeout as tokio_timeout;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::{
|
||||
AgentHook, AsyncAgentHook, BlockingError, HookEvent, PermissionDeniedContext,
|
||||
PermissionRequestAction, PermissionRequestContext, PostCompactContext, PostToolUseAction,
|
||||
PostToolUseContext, PostToolUseFailureContext, PostToolUseResult, PreCompactContext,
|
||||
PreToolUseAction, PreToolUseContext, PreToolUseResult, SessionStartContext, SessionStopContext,
|
||||
StepCompleteContext, SubagentStartContext, SubagentStopContext, TaggedContext,
|
||||
DEFAULT_HOOK_TIMEOUT,
|
||||
};
|
||||
|
||||
use serde_json;
|
||||
|
||||
impl super::HookRegistry {
|
||||
// ── 便捷调用方法 ──
|
||||
|
||||
/// 调度匹配指定事件的异步 hooks（fire-and-forget，5s dispatch 超时）
|
||||
async fn dispatch_async_hooks_for(
|
||||
&self,
|
||||
event: HookEvent,
|
||||
_session_id: &str,
|
||||
_tool_name: &str,
|
||||
_tool_args: &serde_json::Value,
|
||||
post_ctx: Option<&PostToolUseContext>,
|
||||
failure_ctx: Option<&PostToolUseFailureContext>,
|
||||
) {
|
||||
let async_hooks: Vec<&dyn AsyncAgentHook> = self
|
||||
.async_event_index
|
||||
.get(&event)
|
||||
.map(|indices| {
|
||||
indices
|
||||
.iter()
|
||||
.filter_map(|&idx| self.async_hooks.get(idx).map(|b| b.as_ref()))
|
||||
.filter(|_hook| {
|
||||
// Async hooks can declare filtering too via name convention;
|
||||
// for now, all matching the event fire.
|
||||
true
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
if async_hooks.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let timeout_dur = Duration::from_secs(5); // async dispatch timeout
|
||||
// Build futures — clone contexts so futures own their data and satisfy 'static.
|
||||
match event {
|
||||
HookEvent::PostToolUse if post_ctx.is_some() => {
|
||||
let ctx = post_ctx.unwrap().clone();
|
||||
for hook in async_hooks {
|
||||
let ctx = ctx.clone();
|
||||
let _ = tokio_timeout(timeout_dur, hook.on_post_tool_use_async(ctx)).await;
|
||||
}
|
||||
}
|
||||
HookEvent::PostToolUseFailure if failure_ctx.is_some() => {
|
||||
let ctx = failure_ctx.unwrap().clone();
|
||||
for hook in async_hooks {
|
||||
let ctx = ctx.clone();
|
||||
let _ =
|
||||
tokio_timeout(timeout_dur, hook.on_post_tool_use_failure_async(ctx)).await;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_session_start hooks（并行执行，含 session hooks）
|
||||
pub async fn run_on_session_start(&self, ctx: &SessionStartContext) {
|
||||
let hooks = self.collect_hooks_for(HookEvent::OnSessionStart, &ctx.session_id);
|
||||
if hooks.is_empty() {
|
||||
return;
|
||||
}
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
async move {
|
||||
let _ = tokio_timeout(timeout_dur, hook.on_session_start(&ctx)).await;
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
join_all(futures).await;
|
||||
}
|
||||
|
||||
/// 调用所有 pre_tool_use hooks（并行执行，各自有独立超时）。
|
||||
///
|
||||
/// 所有订阅了 PreToolUse 的 hooks（含 session hooks）并行运行，每个 hook 包装在
|
||||
/// `tokio::time::timeout` 中。完成后聚合所有结果：
|
||||
/// - Block → 收集到 blocking_errors
|
||||
/// - MutateInput → 累积 additional_contexts 并更新 final_args
|
||||
/// - PermissionRequired → 记录
|
||||
/// - 超时 → 记录 warning，视为 non-blocking error
|
||||
pub async fn run_pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseResult {
|
||||
let hooks = self.collect_tool_hooks_for(
|
||||
HookEvent::PreToolUse,
|
||||
&ctx.session_id,
|
||||
&ctx.tool_name,
|
||||
&ctx.tool_args,
|
||||
);
|
||||
|
||||
// Fast path: 单 hook 或空 → 顺序执行
|
||||
if hooks.len() <= 1 {
|
||||
return self.run_pre_tool_use_sequential(ctx).await;
|
||||
}
|
||||
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
let hook_name = hook.name().to_string();
|
||||
async move {
|
||||
(
|
||||
hook_name,
|
||||
tokio_timeout(timeout_dur, hook.pre_tool_use(&ctx)).await,
|
||||
)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 并行执行所有 hooks
|
||||
let results = join_all(futures).await;
|
||||
|
||||
// 聚合结果
|
||||
let mut additional_contexts: Vec<String> = Vec::new();
|
||||
let mut tagged_contexts: Vec<TaggedContext> = Vec::new();
|
||||
let mut final_args = ctx.tool_args.clone();
|
||||
let mut final_action = PreToolUseAction::Continue;
|
||||
let mut blocking_errors: Vec<BlockingError> = Vec::new();
|
||||
|
||||
for (hook_name, result) in results {
|
||||
match result {
|
||||
Ok(action) => {
|
||||
match &action {
|
||||
PreToolUseAction::Block { reason } => {
|
||||
warn!(
|
||||
"[Hooks] {hook_name} 阻止了工具 {} 的执行: {reason}",
|
||||
ctx.tool_name
|
||||
);
|
||||
blocking_errors.push(BlockingError::new(hook_name, reason.clone()));
|
||||
}
|
||||
PreToolUseAction::MutateInput {
|
||||
updated_args,
|
||||
additional_context,
|
||||
} => {
|
||||
info!(
|
||||
"[Hooks] {hook_name} 修改了工具 {} 的输入参数",
|
||||
ctx.tool_name
|
||||
);
|
||||
final_args = updated_args.clone();
|
||||
if let Some(ctx_str) = additional_context {
|
||||
additional_contexts.push(ctx_str.clone());
|
||||
tagged_contexts.push(TaggedContext::new(
|
||||
&hook_name,
|
||||
HookEvent::PreToolUse,
|
||||
ctx_str,
|
||||
));
|
||||
}
|
||||
}
|
||||
PreToolUseAction::PermissionRequired { .. } => {
|
||||
info!(
|
||||
"[Hooks] {hook_name} 请求了工具 {} 的权限检查",
|
||||
ctx.tool_name
|
||||
);
|
||||
}
|
||||
PreToolUseAction::Continue => {}
|
||||
}
|
||||
if !matches!(action, PreToolUseAction::Continue) {
|
||||
final_action = action;
|
||||
}
|
||||
}
|
||||
Err(_elapsed) => {
|
||||
warn!("[Hooks] {hook_name} 超时（PreToolUse），跳过其反馈");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let additional_context = if additional_contexts.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(additional_contexts.join("\n"))
|
||||
};
|
||||
|
||||
PreToolUseResult {
|
||||
action: final_action,
|
||||
additional_contexts,
|
||||
additional_context,
|
||||
blocking_errors,
|
||||
final_args,
|
||||
tagged_contexts,
|
||||
}
|
||||
}
|
||||
|
||||
/// 顺序执行 pre_tool_use hooks（含 session hooks）
|
||||
async fn run_pre_tool_use_sequential(&self, ctx: &PreToolUseContext) -> PreToolUseResult {
|
||||
let mut additional_contexts: Vec<String> = Vec::new();
|
||||
let mut tagged_contexts: Vec<TaggedContext> = Vec::new();
|
||||
let mut final_args = ctx.tool_args.clone();
|
||||
let mut final_action = PreToolUseAction::Continue;
|
||||
let mut blocking_errors: Vec<BlockingError> = Vec::new();
|
||||
|
||||
// 全局 hooks（通过 collect_tool_hooks_for 应用 match_filter 过滤）
|
||||
let hooks = self.collect_tool_hooks_for(
|
||||
HookEvent::PreToolUse,
|
||||
&ctx.session_id,
|
||||
&ctx.tool_name,
|
||||
&ctx.tool_args,
|
||||
);
|
||||
for hook in &hooks {
|
||||
let action = hook.pre_tool_use(ctx).await;
|
||||
Self::process_pre_tool_action(
|
||||
&action,
|
||||
hook.name(),
|
||||
&ctx.tool_name,
|
||||
&mut blocking_errors,
|
||||
&mut final_args,
|
||||
&mut additional_contexts,
|
||||
&mut tagged_contexts,
|
||||
&mut final_action,
|
||||
);
|
||||
}
|
||||
|
||||
let additional_context = if additional_contexts.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(additional_contexts.join("\n"))
|
||||
};
|
||||
|
||||
PreToolUseResult {
|
||||
action: final_action,
|
||||
additional_contexts,
|
||||
additional_context,
|
||||
blocking_errors,
|
||||
final_args,
|
||||
tagged_contexts,
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 post_tool_use hooks（并行执行，含 session hooks，各自有独立超时）。
|
||||
pub async fn run_post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseResult {
|
||||
// 同时调度异步 hooks（fire-and-forget）
|
||||
self.dispatch_async_hooks_for(
|
||||
HookEvent::PostToolUse,
|
||||
&ctx.session_id,
|
||||
&ctx.tool_name,
|
||||
&ctx.tool_args,
|
||||
Some(ctx),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let hooks = self.collect_tool_hooks_for(
|
||||
HookEvent::PostToolUse,
|
||||
&ctx.session_id,
|
||||
&ctx.tool_name,
|
||||
&ctx.tool_args,
|
||||
);
|
||||
|
||||
// Fast path: 单 hook 或空 → 顺序执行
|
||||
if hooks.len() <= 1 {
|
||||
return self.run_post_tool_use_sequential(ctx).await;
|
||||
}
|
||||
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
let hook_name = hook.name().to_string();
|
||||
async move {
|
||||
(
|
||||
hook_name,
|
||||
tokio_timeout(timeout_dur, hook.post_tool_use(&ctx)).await,
|
||||
)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let results = join_all(futures).await;
|
||||
|
||||
let mut final_content = ctx.output_content.clone();
|
||||
let mut additional_contexts: Vec<String> = Vec::new();
|
||||
let mut tagged_contexts: Vec<TaggedContext> = Vec::new();
|
||||
let mut warnings: Vec<String> = Vec::new();
|
||||
let mut post_permission_requests: Vec<(String, String)> = Vec::new();
|
||||
let mut metadata: std::collections::HashMap<String, serde_json::Value> =
|
||||
std::collections::HashMap::new();
|
||||
|
||||
for (hook_name, result) in results {
|
||||
match result {
|
||||
Ok(action) => match &action {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content: ref uc,
|
||||
additional_context: ref ac,
|
||||
} => {
|
||||
info!("[Hooks] {hook_name} 修改了工具 {} 的输出", ctx.tool_name);
|
||||
final_content = uc.clone();
|
||||
if let Some(ctx_str) = ac {
|
||||
additional_contexts.push(ctx_str.clone());
|
||||
tagged_contexts.push(TaggedContext::new(
|
||||
&hook_name,
|
||||
HookEvent::PostToolUse,
|
||||
ctx_str,
|
||||
));
|
||||
}
|
||||
}
|
||||
PostToolUseAction::Warning {
|
||||
message,
|
||||
truncate_output,
|
||||
} => {
|
||||
warn!(
|
||||
"[Hooks] {hook_name} 发出了对 {} 的警告: {message}",
|
||||
ctx.tool_name
|
||||
);
|
||||
warnings.push(message.clone());
|
||||
if *truncate_output {
|
||||
final_content = final_content.chars().take(1000).collect::<String>()
|
||||
+ "\n\n[输出已由 Hook 截断]";
|
||||
}
|
||||
}
|
||||
PostToolUseAction::PermissionRequired {
|
||||
permission,
|
||||
tool_name,
|
||||
} => {
|
||||
info!("[Hooks] {hook_name} 事后请求工具 {tool_name} 的权限: {permission}");
|
||||
post_permission_requests.push((tool_name.clone(), permission.clone()));
|
||||
}
|
||||
PostToolUseAction::Metadata { key, value } => {
|
||||
metadata.insert(key.clone(), value.clone());
|
||||
}
|
||||
PostToolUseAction::Continue => {}
|
||||
},
|
||||
Err(_elapsed) => {
|
||||
warn!("[Hooks] {hook_name} 超时（PostToolUse），跳过其反馈");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
PostToolUseResult {
|
||||
final_content,
|
||||
additional_contexts,
|
||||
tagged_contexts,
|
||||
warnings,
|
||||
post_permission_requests,
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
/// 顺序执行 post_tool_use hooks（含 session hooks）
|
||||
async fn run_post_tool_use_sequential(&self, ctx: &PostToolUseContext) -> PostToolUseResult {
|
||||
let mut final_content = ctx.output_content.clone();
|
||||
let mut additional_contexts: Vec<String> = Vec::new();
|
||||
let mut tagged_contexts: Vec<TaggedContext> = Vec::new();
|
||||
let mut warnings: Vec<String> = Vec::new();
|
||||
let mut post_permission_requests: Vec<(String, String)> = Vec::new();
|
||||
let mut metadata: std::collections::HashMap<String, serde_json::Value> =
|
||||
std::collections::HashMap::new();
|
||||
|
||||
let hooks = self.collect_tool_hooks_for(
|
||||
HookEvent::PostToolUse,
|
||||
&ctx.session_id,
|
||||
&ctx.tool_name,
|
||||
&ctx.tool_args,
|
||||
);
|
||||
for hook in &hooks {
|
||||
let action = hook.post_tool_use(ctx).await;
|
||||
match action {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content,
|
||||
additional_context,
|
||||
} => {
|
||||
final_content = updated_content;
|
||||
if let Some(ctx_str) = additional_context {
|
||||
additional_contexts.push(ctx_str.clone());
|
||||
tagged_contexts.push(TaggedContext::new(
|
||||
hook.name(),
|
||||
HookEvent::PostToolUse,
|
||||
ctx_str,
|
||||
));
|
||||
}
|
||||
}
|
||||
PostToolUseAction::Warning {
|
||||
message,
|
||||
truncate_output,
|
||||
} => {
|
||||
warnings.push(message);
|
||||
if truncate_output {
|
||||
final_content = final_content.chars().take(1000).collect::<String>()
|
||||
+ "\n\n[输出已由 Hook 截断]";
|
||||
}
|
||||
}
|
||||
PostToolUseAction::PermissionRequired {
|
||||
permission,
|
||||
tool_name,
|
||||
} => {
|
||||
post_permission_requests.push((tool_name, permission));
|
||||
}
|
||||
PostToolUseAction::Metadata { key, value } => {
|
||||
metadata.insert(key, value.clone());
|
||||
}
|
||||
PostToolUseAction::Continue => {}
|
||||
}
|
||||
}
|
||||
|
||||
PostToolUseResult {
|
||||
final_content,
|
||||
additional_contexts,
|
||||
tagged_contexts,
|
||||
warnings,
|
||||
post_permission_requests,
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 post_tool_use_failure hooks（并行执行，含 session hooks）
|
||||
pub async fn run_on_post_tool_use_failure(
|
||||
&self,
|
||||
ctx: &PostToolUseFailureContext,
|
||||
) -> PostToolUseResult {
|
||||
// 同时调度异步 hooks（fire-and-forget）
|
||||
self.dispatch_async_hooks_for(
|
||||
HookEvent::PostToolUseFailure,
|
||||
&ctx.session_id,
|
||||
&ctx.tool_name,
|
||||
&ctx.tool_args,
|
||||
None,
|
||||
Some(ctx),
|
||||
)
|
||||
.await;
|
||||
|
||||
let hooks = self.collect_tool_hooks_for(
|
||||
HookEvent::PostToolUseFailure,
|
||||
&ctx.session_id,
|
||||
&ctx.tool_name,
|
||||
&ctx.tool_args,
|
||||
);
|
||||
|
||||
// Fast path: 单 hook 或空 → 顺序执行
|
||||
if hooks.len() <= 1 {
|
||||
let mut final_content = String::new();
|
||||
let mut additional_contexts: Vec<String> = Vec::new();
|
||||
let mut tagged_contexts: Vec<TaggedContext> = Vec::new();
|
||||
for hook in &hooks {
|
||||
let action = hook.on_post_tool_use_failure(ctx).await;
|
||||
match action {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content,
|
||||
additional_context,
|
||||
} => {
|
||||
final_content = updated_content;
|
||||
if let Some(ctx_str) = additional_context {
|
||||
additional_contexts.push(ctx_str.clone());
|
||||
tagged_contexts.push(TaggedContext::new(
|
||||
hook.name(),
|
||||
HookEvent::PostToolUseFailure,
|
||||
ctx_str,
|
||||
));
|
||||
}
|
||||
}
|
||||
PostToolUseAction::Warning {
|
||||
message,
|
||||
truncate_output,
|
||||
} => {
|
||||
if truncate_output {
|
||||
final_content = final_content.chars().take(1000).collect();
|
||||
}
|
||||
tagged_contexts.push(TaggedContext::new(
|
||||
hook.name(),
|
||||
HookEvent::PostToolUseFailure,
|
||||
format!("Warning: {message}"),
|
||||
));
|
||||
}
|
||||
PostToolUseAction::PermissionRequired { .. }
|
||||
| PostToolUseAction::Metadata { .. }
|
||||
| PostToolUseAction::Continue => {}
|
||||
}
|
||||
}
|
||||
return PostToolUseResult {
|
||||
final_content,
|
||||
additional_contexts,
|
||||
tagged_contexts,
|
||||
warnings: Vec::new(),
|
||||
post_permission_requests: Vec::new(),
|
||||
metadata: std::collections::HashMap::new(),
|
||||
};
|
||||
}
|
||||
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
let hook_name = hook.name().to_string();
|
||||
async move {
|
||||
(
|
||||
hook_name,
|
||||
tokio_timeout(timeout_dur, hook.on_post_tool_use_failure(&ctx)).await,
|
||||
)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let results = join_all(futures).await;
|
||||
let mut final_content = String::new();
|
||||
let mut additional_contexts: Vec<String> = Vec::new();
|
||||
let mut tagged_contexts: Vec<TaggedContext> = Vec::new();
|
||||
|
||||
for (hook_name, result) in results {
|
||||
match result {
|
||||
Ok(action) => match action {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content,
|
||||
additional_context,
|
||||
} => {
|
||||
final_content = updated_content;
|
||||
if let Some(ctx_str) = additional_context {
|
||||
additional_contexts.push(ctx_str.clone());
|
||||
tagged_contexts.push(TaggedContext::new(
|
||||
&hook_name,
|
||||
HookEvent::PostToolUseFailure,
|
||||
ctx_str,
|
||||
));
|
||||
}
|
||||
}
|
||||
PostToolUseAction::Warning {
|
||||
message,
|
||||
truncate_output,
|
||||
} => {
|
||||
if truncate_output {
|
||||
final_content = final_content.chars().take(1000).collect();
|
||||
}
|
||||
tagged_contexts.push(TaggedContext::new(
|
||||
&hook_name,
|
||||
HookEvent::PostToolUseFailure,
|
||||
format!("Warning: {message}"),
|
||||
));
|
||||
}
|
||||
PostToolUseAction::PermissionRequired { .. }
|
||||
| PostToolUseAction::Metadata { .. }
|
||||
| PostToolUseAction::Continue => {}
|
||||
},
|
||||
Err(_elapsed) => {
|
||||
// 超时已在警告中体现，跳过该 hook 的反馈
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
PostToolUseResult {
|
||||
final_content,
|
||||
additional_contexts,
|
||||
tagged_contexts,
|
||||
warnings: Vec::new(),
|
||||
post_permission_requests: Vec::new(),
|
||||
metadata: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 调用所有 on_step_complete hooks（并行执行，含 session hooks）
|
||||
pub async fn run_on_step_complete(&self, ctx: &StepCompleteContext) {
|
||||
let hooks = self.collect_hooks_for(HookEvent::OnStepComplete, &ctx.session_id);
|
||||
if hooks.is_empty() {
|
||||
return;
|
||||
}
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
async move {
|
||||
let _ = tokio_timeout(timeout_dur, hook.on_step_complete(&ctx)).await;
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
join_all(futures).await;
|
||||
}
|
||||
|
||||
/// 调用所有 on_session_stop hooks（顺序执行 + 自动清理 session hooks）
|
||||
pub async fn run_on_session_stop(&self, ctx: &SessionStopContext<'_>) {
|
||||
// 全局 hooks
|
||||
for &idx in self.indices_for(HookEvent::OnSessionStop) {
|
||||
if let Some(hook) = self.hooks.get(idx) {
|
||||
let _ = tokio_timeout(
|
||||
hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT),
|
||||
hook.on_session_stop(ctx),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
// Session hooks（Stop 事件也通知它们）
|
||||
for hook in self.get_session_hooks(&ctx.session_id) {
|
||||
let _ = tokio_timeout(
|
||||
hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT),
|
||||
hook.on_session_stop(ctx),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
// 注意：clear_session_hooks 需要 &mut self，但此方法是 &self。
|
||||
// 调用方应在 run_on_session_stop 返回后主动调用 clear_session_hooks。
|
||||
}
|
||||
|
||||
/// 调用所有 on_subagent_start hooks（并行执行，含 session hooks）
|
||||
pub async fn run_on_subagent_start(&self, ctx: &SubagentStartContext) {
|
||||
let hooks = self.collect_hooks_for(HookEvent::OnSubagentStart, &ctx.parent_session_id);
|
||||
if hooks.is_empty() {
|
||||
return;
|
||||
}
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
async move {
|
||||
let _ = tokio_timeout(timeout_dur, hook.on_subagent_start(&ctx)).await;
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
join_all(futures).await;
|
||||
}
|
||||
|
||||
/// 调用所有 on_subagent_stop hooks（并行执行，含 session hooks）
|
||||
pub async fn run_on_subagent_stop(&self, ctx: &SubagentStopContext) {
|
||||
let hooks = self.collect_hooks_for(HookEvent::OnSubagentStop, &ctx.parent_session_id);
|
||||
if hooks.is_empty() {
|
||||
return;
|
||||
}
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
async move {
|
||||
let _ = tokio_timeout(timeout_dur, hook.on_subagent_stop(&ctx)).await;
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
join_all(futures).await;
|
||||
}
|
||||
|
||||
/// 调用所有 on_pre_compact hooks（并行执行，含 session hooks）
|
||||
pub async fn run_on_pre_compact(&self, ctx: &PreCompactContext) {
|
||||
let hooks = self.collect_hooks_for(HookEvent::OnPreCompact, &ctx.session_id);
|
||||
if hooks.is_empty() {
|
||||
return;
|
||||
}
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
async move {
|
||||
let _ = tokio_timeout(timeout_dur, hook.on_pre_compact(&ctx)).await;
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
join_all(futures).await;
|
||||
}
|
||||
|
||||
/// 调用所有 on_post_compact hooks（并行执行，含 session hooks）
|
||||
pub async fn run_on_post_compact(&self, ctx: &PostCompactContext) {
|
||||
let hooks = self.collect_hooks_for(HookEvent::OnPostCompact, &ctx.session_id);
|
||||
if hooks.is_empty() {
|
||||
return;
|
||||
}
|
||||
let ctx_clone = ctx.clone();
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let ctx = ctx_clone.clone();
|
||||
let timeout_dur = hook.timeout().unwrap_or(DEFAULT_HOOK_TIMEOUT);
|
||||
async move {
|
||||
let _ = tokio_timeout(timeout_dur, hook.on_post_compact(&ctx)).await;
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
join_all(futures).await;
|
||||
}
|
||||
|
||||
/// 分发 PermissionRequest 事件。
|
||||
///
|
||||
/// 并行调用所有匹配的 hook。每个 hook 返回 `PermissionRequestAction`：
|
||||
/// - `Continue` → 保持当前决策
|
||||
/// - `Override` → 覆盖决策（第一个 Override 生效，后续不再检查）
|
||||
/// - `InjectContext` → 注入上下文但不改变决策
|
||||
///
|
||||
/// 返回 (final_action, additional_contexts)。
|
||||
pub async fn run_on_permission_request(
|
||||
&self,
|
||||
ctx: &PermissionRequestContext,
|
||||
) -> (PermissionRequestAction, Vec<String>) {
|
||||
let hooks = self.collect_hooks_for(HookEvent::PermissionRequest, &ctx.session_id);
|
||||
if hooks.is_empty() {
|
||||
return (PermissionRequestAction::Continue, Vec::new());
|
||||
}
|
||||
|
||||
let mut final_action = PermissionRequestAction::Continue;
|
||||
let mut injected_contexts: Vec<String> = Vec::new();
|
||||
let mut overridden = false;
|
||||
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let hook: &dyn AgentHook = *hook;
|
||||
let ctx = ctx.clone();
|
||||
let timeout_dur = DEFAULT_HOOK_TIMEOUT;
|
||||
async move {
|
||||
let result = tokio_timeout(timeout_dur, hook.on_permission_request(&ctx)).await;
|
||||
match result {
|
||||
Ok(action) => Some((hook.name().to_string(), action)),
|
||||
Err(_) => {
|
||||
warn!("[HookDispatch] PermissionRequest hook {} 超时", hook.name());
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let results: Vec<Option<(String, PermissionRequestAction)>> = join_all(futures).await;
|
||||
for result in results.into_iter().flatten() {
|
||||
let (hook_name, action) = result;
|
||||
match action {
|
||||
PermissionRequestAction::Continue => {
|
||||
// 默认，不改变
|
||||
}
|
||||
PermissionRequestAction::Override(decision) if !overridden => {
|
||||
overridden = true;
|
||||
final_action = PermissionRequestAction::Override(decision);
|
||||
info!(
|
||||
"[HookDispatch] PermissionRequest hook {} 覆盖决策",
|
||||
hook_name
|
||||
);
|
||||
}
|
||||
PermissionRequestAction::Override(_) => {
|
||||
warn!(
|
||||
"[HookDispatch] PermissionRequest hook {}: 覆盖被忽略（已有其他 hook 先覆盖）",
|
||||
hook_name
|
||||
);
|
||||
}
|
||||
PermissionRequestAction::InjectContext { ref content } => {
|
||||
injected_contexts.push(content.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(final_action, injected_contexts)
|
||||
}
|
||||
|
||||
/// 分发 PermissionDenied 事件（fire-and-forget，仅用于审计日志）。
|
||||
///
|
||||
/// 并行调用所有匹配的 hook。返回值不影响执行流程。
|
||||
pub async fn run_on_permission_denied(&self, ctx: &PermissionDeniedContext) {
|
||||
let hooks = self.collect_hooks_for(HookEvent::PermissionDenied, &ctx.session_id);
|
||||
if hooks.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let futures: Vec<_> = hooks
|
||||
.iter()
|
||||
.map(|hook| {
|
||||
let hook: &dyn AgentHook = *hook;
|
||||
let ctx = ctx.clone();
|
||||
let timeout_dur = DEFAULT_HOOK_TIMEOUT;
|
||||
async move {
|
||||
let _ = tokio_timeout(timeout_dur, hook.on_permission_denied(&ctx)).await;
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
join_all(futures).await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,302 @@
|
||||
// src/agent/hooks/matcher.rs
|
||||
//
|
||||
// 工具级匹配器。
|
||||
// Hook 可通过 `match_filter()` 声明只关心特定工具或参数模式，
|
||||
// 从而避免在热路径上被无关调用触发。
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
/// 工具名匹配模式。
|
||||
///
|
||||
/// 三种策略：
|
||||
/// - `Exact("search_papers")` — 精确名称匹配
|
||||
/// - `Prefix("file_")` — 前缀匹配（末尾 `*` 隐式）
|
||||
/// - `Wildcard` — 匹配所有工具（默认行为）
|
||||
///
|
||||
/// # 示例
|
||||
///
|
||||
/// ```
|
||||
/// use crate::agent::hooks::matcher::ToolNamePattern;
|
||||
/// let pat = ToolNamePattern::parse("file_*");
|
||||
/// assert!(pat.matches("file_write"));
|
||||
/// assert!(!pat.matches("search_papers"));
|
||||
/// ```
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ToolNamePattern {
|
||||
/// 精确匹配工具名
|
||||
Exact(String),
|
||||
/// 前缀匹配 — `"file_*"` 匹配所有以 `"file_"` 开头的工具名
|
||||
Prefix(String),
|
||||
/// 匹配所有工具
|
||||
Wildcard,
|
||||
}
|
||||
|
||||
impl ToolNamePattern {
|
||||
/// 从模式字符串解析。
|
||||
///
|
||||
/// - `"*"` 或空字符串 → `Wildcard`
|
||||
/// - `"xxx_*"` → `Prefix("xxx_")`
|
||||
/// - 其他 → `Exact(pattern)`
|
||||
pub fn parse(pattern: &str) -> Self {
|
||||
let trimmed = pattern.trim();
|
||||
if trimmed.is_empty() || trimmed == "*" {
|
||||
return ToolNamePattern::Wildcard;
|
||||
}
|
||||
if let Some(prefix) = trimmed.strip_suffix('*') {
|
||||
if !prefix.is_empty() {
|
||||
return ToolNamePattern::Prefix(prefix.to_string());
|
||||
}
|
||||
// 只有 "*" — 上面已经处理
|
||||
return ToolNamePattern::Wildcard;
|
||||
}
|
||||
ToolNamePattern::Exact(trimmed.to_string())
|
||||
}
|
||||
|
||||
/// 若 `tool_name` 与此模式匹配则返回 true。
|
||||
pub fn matches(&self, tool_name: &str) -> bool {
|
||||
match self {
|
||||
ToolNamePattern::Wildcard => true,
|
||||
ToolNamePattern::Exact(name) => name == tool_name,
|
||||
ToolNamePattern::Prefix(prefix) => tool_name.starts_with(prefix),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hook 通过 `match_filter()` 声明的结构化过滤器，用于限制关注哪些工具调用。
|
||||
///
|
||||
/// 若 hook 返回非空 `ToolMatchFilter`，则仅当工具名和参数匹配时
|
||||
/// 才调用其 `pre_tool_use` / `post_tool_use`。
|
||||
///
|
||||
/// 默认情况下（空过滤器，等价于 `[Wildcard]`），匹配所有工具。
|
||||
///
|
||||
/// # 示例
|
||||
///
|
||||
/// ```
|
||||
/// use crate::agent::hooks::matcher::{ToolMatchFilter, ToolNamePattern};
|
||||
///
|
||||
/// // 仅匹配文件相关工具
|
||||
/// let filter = ToolMatchFilter {
|
||||
/// name_patterns: vec![ToolNamePattern::parse("file_*")],
|
||||
/// content_patterns: vec![],
|
||||
/// };
|
||||
/// assert!(filter.matches("file_read", &serde_json::json!({})));
|
||||
/// assert!(!filter.matches("search_papers", &serde_json::json!({})));
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ToolMatchFilter {
|
||||
/// 工具名模式 — 任意一个匹配即可通过。
|
||||
/// 空 Vec = 匹配全部（等价于 `[Wildcard]`）。
|
||||
pub name_patterns: Vec<ToolNamePattern>,
|
||||
/// 可选的内容级模式。格式：`"字段:子串"` 或直接 `"子串"`。
|
||||
/// 例如 `"command:rm *"` 匹配 `run_bash` 调用中 command 字段以 "rm " 开头的情况。
|
||||
/// 空 Vec = 不过滤内容。
|
||||
pub content_patterns: Vec<String>,
|
||||
}
|
||||
|
||||
impl ToolMatchFilter {
|
||||
/// 若此过滤器匹配给定的工具调用则返回 true。
|
||||
///
|
||||
/// 检查流程：
|
||||
/// 1. 若 `name_patterns` 为空，所有工具名匹配
|
||||
/// 2. 任一非通配模式必须匹配 tool_name
|
||||
/// 3. 任一 content_patterns 必须匹配 tool_args 中的某个字段
|
||||
pub fn matches(&self, tool_name: &str, tool_args: &Value) -> bool {
|
||||
// 名称匹配
|
||||
if !self.name_patterns.is_empty() {
|
||||
let name_match = self.name_patterns.iter().any(|pat| pat.matches(tool_name));
|
||||
if !name_match {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// 内容匹配：每个模式必须至少匹配一个字段/值
|
||||
if !self.content_patterns.is_empty() {
|
||||
for pattern in &self.content_patterns {
|
||||
if !content_matches(pattern, tool_args) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// 快捷构造：创建仅匹配单个工具名的过滤器。
|
||||
pub fn exact(tool_name: &str) -> Self {
|
||||
ToolMatchFilter {
|
||||
name_patterns: vec![ToolNamePattern::Exact(tool_name.to_string())],
|
||||
content_patterns: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
/// 快捷构造：创建前缀过滤器。
|
||||
pub fn prefix(prefix: &str) -> Self {
|
||||
ToolMatchFilter {
|
||||
name_patterns: vec![ToolNamePattern::Prefix(prefix.to_string())],
|
||||
content_patterns: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
/// 若此过滤器为默认"匹配全部"（空 name_patterns + 空 content_patterns）则返回 true。
|
||||
pub fn is_match_all(&self) -> bool {
|
||||
self.name_patterns.is_empty() && self.content_patterns.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查内容模式是否匹配 JSON 值中的任意字段。
|
||||
///
|
||||
/// 模式格式：
|
||||
/// - `"字段:值"` — 值必须是该字段字符串表示的子串
|
||||
/// - `"值"` — 值必须出现在 JSON 字符串化后的任意位置
|
||||
fn content_matches(pattern: &str, args: &Value) -> bool {
|
||||
if let Some((field, wanted)) = pattern.split_once(':') {
|
||||
// 匹配指定字段
|
||||
if let Some(field_val) = args.get(field.trim()) {
|
||||
let field_str = match field_val {
|
||||
Value::String(s) => s.clone(),
|
||||
other => other.to_string(),
|
||||
};
|
||||
return field_str.contains(wanted.trim());
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// 在 JSON 中任意位置匹配
|
||||
let full_str = serde_json::to_string(args).unwrap_or_default();
|
||||
full_str.contains(pattern.trim())
|
||||
}
|
||||
|
||||
// ── 测试 ──
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── ToolNamePattern ──
|
||||
|
||||
#[test]
|
||||
fn test_pattern_exact() {
|
||||
let pat = ToolNamePattern::parse("search_papers");
|
||||
assert!(pat.matches("search_papers"));
|
||||
assert!(!pat.matches("download_paper"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pattern_prefix() {
|
||||
let pat = ToolNamePattern::parse("file_*");
|
||||
assert!(pat.matches("file_write"));
|
||||
assert!(pat.matches("file_read"));
|
||||
assert!(pat.matches("file_edit"));
|
||||
assert!(!pat.matches("search_papers"));
|
||||
assert!(!pat.matches("fi")); // 前缀比 "file_" 短
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pattern_wildcard() {
|
||||
let pat = ToolNamePattern::parse("*");
|
||||
assert!(matches!(pat, ToolNamePattern::Wildcard));
|
||||
assert!(pat.matches("anything"));
|
||||
assert!(pat.matches("search_papers"));
|
||||
|
||||
let pat_empty = ToolNamePattern::parse("");
|
||||
assert!(matches!(pat_empty, ToolNamePattern::Wildcard));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pattern_prefix_no_trailing_wild() {
|
||||
// "file" 不带 * 应被解析为 Exact
|
||||
let pat = ToolNamePattern::parse("file");
|
||||
assert!(matches!(pat, ToolNamePattern::Exact(_)));
|
||||
assert!(pat.matches("file"));
|
||||
assert!(!pat.matches("file_write"));
|
||||
}
|
||||
|
||||
// ── ToolMatchFilter ──
|
||||
|
||||
#[test]
|
||||
fn test_filter_default_matches_all() {
|
||||
let filter = ToolMatchFilter::default();
|
||||
assert!(filter.matches("search_papers", &serde_json::json!({})));
|
||||
assert!(filter.matches("run_bash", &serde_json::json!({"command": "rm -rf /"})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_exact_name_match() {
|
||||
let filter = ToolMatchFilter::exact("search_papers");
|
||||
assert!(filter.matches("search_papers", &serde_json::json!({})));
|
||||
assert!(!filter.matches("download_paper", &serde_json::json!({})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_prefix_name_match() {
|
||||
let filter = ToolMatchFilter::prefix("file_");
|
||||
assert!(filter.matches("file_write", &serde_json::json!({})));
|
||||
assert!(filter.matches("file_read", &serde_json::json!({})));
|
||||
assert!(!filter.matches("run_bash", &serde_json::json!({})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_content_match_field() {
|
||||
let filter = ToolMatchFilter {
|
||||
name_patterns: vec![],
|
||||
content_patterns: vec!["command:rm".to_string()],
|
||||
};
|
||||
assert!(filter.matches(
|
||||
"run_bash",
|
||||
&serde_json::json!({"command": "rm -rf /tmp/test"})
|
||||
));
|
||||
assert!(!filter.matches("run_bash", &serde_json::json!({"command": "ls -la"})));
|
||||
assert!(!filter.matches("read_file", &serde_json::json!({"path": "/tmp/test"})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_content_match_anywhere() {
|
||||
let filter = ToolMatchFilter {
|
||||
name_patterns: vec![],
|
||||
content_patterns: vec!["dangerous".to_string()],
|
||||
};
|
||||
assert!(filter.matches(
|
||||
"run_bash",
|
||||
&serde_json::json!({"command": "echo dangerous stuff"})
|
||||
));
|
||||
assert!(!filter.matches(
|
||||
"run_bash",
|
||||
&serde_json::json!({"command": "echo safe stuff"})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_name_and_content_combined() {
|
||||
let filter = ToolMatchFilter {
|
||||
name_patterns: vec![ToolNamePattern::parse("file_*")],
|
||||
content_patterns: vec!["path:.env".to_string()],
|
||||
};
|
||||
// 正确工具名 + 敏感路径 → 匹配
|
||||
assert!(filter.matches("file_read", &serde_json::json!({"path": "/app/.env"})));
|
||||
// 正确工具名 + 安全路径 → 不匹配
|
||||
assert!(!filter.matches("file_read", &serde_json::json!({"path": "/app/README.md"})));
|
||||
// 错误工具名 + 敏感路径 → 不匹配
|
||||
assert!(!filter.matches("run_bash", &serde_json::json!({"path": "/app/.env"})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_empty_content_false() {
|
||||
// 内容模式指定了不存在的字段 → 不匹配
|
||||
let filter = ToolMatchFilter {
|
||||
name_patterns: vec![],
|
||||
content_patterns: vec!["nonexistent:value".to_string()],
|
||||
};
|
||||
assert!(!filter.matches("run_bash", &serde_json::json!({"command": "ls"})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_match_all() {
|
||||
assert!(ToolMatchFilter::default().is_match_all());
|
||||
assert!(!ToolMatchFilter::exact("foo").is_match_all());
|
||||
assert!(!ToolMatchFilter {
|
||||
name_patterns: vec![],
|
||||
content_patterns: vec!["cmd:ls".to_string()]
|
||||
}
|
||||
.is_match_all());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,549 @@
|
||||
// src/agent/hooks/mod.rs
|
||||
//
|
||||
// Agent 生命周期 Hooks 系统。
|
||||
//
|
||||
// 参考 Claude Code 的 PreToolUse / PostToolUse / Stop hooks 设计，
|
||||
// 提供可扩展的事件回调链，支持 10 种生命周期事件 + 工具匹配过滤 +
|
||||
// 权限决策优先级 + 异步 fire-and-forget hook。
|
||||
//
|
||||
// 子模块结构：
|
||||
// - types.rs — 所有数据类型定义（HookEvent、Contexts、Actions、Results）
|
||||
// - traits.rs — AgentHook + AsyncAgentHook traits
|
||||
// - matcher.rs — ToolNamePattern / ToolMatchFilter 工具匹配器
|
||||
// - registry.rs — HookRegistry 结构体 + 基础方法
|
||||
// - dispatch.rs — HookRegistry 调度方法（所有 run_*）
|
||||
// - builtins.rs — 内置 Hooks（CancellationHook、MetricsHook、AuditLogHook）
|
||||
|
||||
pub mod builtins;
|
||||
pub mod dispatch;
|
||||
pub mod matcher;
|
||||
pub mod registry;
|
||||
pub mod traits;
|
||||
pub mod types;
|
||||
|
||||
// ── 重导出：外部代码通过 `crate::agent::hooks::*` 访问 ──
|
||||
|
||||
// 核心注册表
|
||||
pub use registry::HookRegistry;
|
||||
|
||||
// 类型
|
||||
pub use types::{
|
||||
event_label, BlockingError, HookAction, HookEvent, MetricsData, PermissionDecision,
|
||||
PermissionDenialSource, PermissionDeniedContext, PermissionRequestAction,
|
||||
PermissionRequestContext, PostCompactContext, PostToolUseAction, PostToolUseContext,
|
||||
PostToolUseFailureContext, PostToolUseResult, PreCompactContext, PreToolUseAction,
|
||||
PreToolUseContext, PreToolUseResult, SessionStartContext, SessionStopContext,
|
||||
StepCompleteContext, SubagentStartContext, SubagentStopContext, TaggedContext,
|
||||
DEFAULT_HOOK_TIMEOUT,
|
||||
};
|
||||
|
||||
// Traits
|
||||
pub use traits::{AgentHook, AsyncAgentHook};
|
||||
|
||||
// Matcher
|
||||
pub use matcher::{ToolMatchFilter, ToolNamePattern};
|
||||
|
||||
// Builtins
|
||||
pub use builtins::{AuditLogHook, CancellationHook, ContextDeduplicator, MetricsHook};
|
||||
|
||||
// ── 测试 ──
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use sqlx::SqlitePool;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
struct TestHook {
|
||||
name: String,
|
||||
pre_called: std::sync::Mutex<bool>,
|
||||
}
|
||||
|
||||
impl TestHook {
|
||||
fn new(name: &str) -> Self {
|
||||
TestHook {
|
||||
name: name.to_string(),
|
||||
pre_called: std::sync::Mutex::new(false),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for TestHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
*self.pre_called.lock().unwrap() = true;
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_registry_runs_all_hooks() {
|
||||
let mut registry = HookRegistry::new();
|
||||
let hook1 = TestHook::new("test1");
|
||||
let hook2 = TestHook::new("test2");
|
||||
registry.add(Box::new(hook1));
|
||||
registry.add(Box::new(hook2));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert!(!result.action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_blocking_hook_stops_chain() {
|
||||
struct BlockingHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for BlockingHook {
|
||||
fn name(&self) -> &str {
|
||||
"blocker"
|
||||
}
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::Block {
|
||||
reason: "test block".into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(BlockingHook));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert!(result.action.is_blocked());
|
||||
assert_eq!(result.action.block_reason(), Some("test block"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mutate_input_accumulates_context() {
|
||||
struct MutateHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for MutateHook {
|
||||
fn name(&self) -> &str {
|
||||
"mutator"
|
||||
}
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::MutateInput {
|
||||
updated_args: serde_json::json!({"key": "modified"}),
|
||||
additional_context: Some("injected context".to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(MutateHook));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({"key": "original"}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let result = registry.run_pre_tool_use(&ctx).await;
|
||||
assert_eq!(result.final_args, serde_json::json!({"key": "modified"}));
|
||||
assert_eq!(
|
||||
result.additional_context,
|
||||
Some("injected context".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_post_tool_use_mutate_output() {
|
||||
struct MutateOutputHook;
|
||||
#[async_trait]
|
||||
impl AgentHook for MutateOutputHook {
|
||||
fn name(&self) -> &str {
|
||||
"output_mutator"
|
||||
}
|
||||
async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content: "modified output".to_string(),
|
||||
additional_context: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(MutateOutputHook));
|
||||
|
||||
let ctx = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "test_tool".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "original output".into(),
|
||||
is_error: false,
|
||||
step: 1,
|
||||
elapsed_ms: 100,
|
||||
};
|
||||
|
||||
let result = registry.run_post_tool_use(&ctx).await;
|
||||
assert_eq!(result.final_content, "modified output");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancellation_hook_blocks_when_cancelled() {
|
||||
let mut cancelled = HashSet::new();
|
||||
cancelled.insert("test_session".to_string());
|
||||
let cancelled_runs = Arc::new(std::sync::Mutex::new(cancelled));
|
||||
|
||||
let hook = CancellationHook::new(cancelled_runs);
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test_session".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let action = hook.pre_tool_use(&ctx).await;
|
||||
assert!(action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancellation_hook_allows_when_not_cancelled() {
|
||||
let cancelled_runs = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
|
||||
|
||||
let hook = CancellationHook::new(cancelled_runs);
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test_session".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
|
||||
let action = hook.pre_tool_use(&ctx).await;
|
||||
assert!(!action.is_blocked());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_metrics_hook_accumulates_counts() {
|
||||
let hook = MetricsHook::new();
|
||||
|
||||
let ctx = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "result".into(),
|
||||
is_error: false,
|
||||
step: 1,
|
||||
elapsed_ms: 100,
|
||||
};
|
||||
hook.post_tool_use(&ctx).await;
|
||||
|
||||
let ctx2 = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
output_content: "result2".into(),
|
||||
is_error: false,
|
||||
step: 2,
|
||||
elapsed_ms: 200,
|
||||
};
|
||||
hook.post_tool_use(&ctx2).await;
|
||||
|
||||
let snapshot = hook.snapshot().expect("snapshot should succeed in test");
|
||||
assert_eq!(snapshot.tool_call_counts.get("search_papers"), Some(&2));
|
||||
assert_eq!(snapshot.total_steps, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_session_start_hook_called() {
|
||||
struct StartTrackingHook {
|
||||
started: std::sync::Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for StartTrackingHook {
|
||||
fn name(&self) -> &str {
|
||||
"start_tracker"
|
||||
}
|
||||
async fn on_session_start(&self, ctx: &SessionStartContext) {
|
||||
self.started.lock().unwrap().push(ctx.session_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let hook = StartTrackingHook {
|
||||
started: std::sync::Mutex::new(Vec::new()),
|
||||
};
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(hook));
|
||||
|
||||
let ctx = SessionStartContext {
|
||||
session_id: "test_sid".into(),
|
||||
turn_index: 1,
|
||||
is_resume: false,
|
||||
};
|
||||
registry.run_on_session_start(&ctx).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_new_lifecycle_events_called() {
|
||||
struct LifecycleTracker {
|
||||
subagent_start: std::sync::Mutex<bool>,
|
||||
subagent_stop: std::sync::Mutex<bool>,
|
||||
pre_compact: std::sync::Mutex<bool>,
|
||||
post_compact: std::sync::Mutex<bool>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for LifecycleTracker {
|
||||
fn name(&self) -> &str {
|
||||
"lifecycle_tracker"
|
||||
}
|
||||
async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {
|
||||
*self.subagent_start.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {
|
||||
*self.subagent_stop.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_pre_compact(&self, _ctx: &PreCompactContext) {
|
||||
*self.pre_compact.lock().unwrap() = true;
|
||||
}
|
||||
async fn on_post_compact(&self, _ctx: &PostCompactContext) {
|
||||
*self.post_compact.lock().unwrap() = true;
|
||||
}
|
||||
}
|
||||
|
||||
let tracker = LifecycleTracker {
|
||||
subagent_start: std::sync::Mutex::new(false),
|
||||
subagent_stop: std::sync::Mutex::new(false),
|
||||
pre_compact: std::sync::Mutex::new(false),
|
||||
post_compact: std::sync::Mutex::new(false),
|
||||
};
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(tracker));
|
||||
|
||||
registry
|
||||
.run_on_subagent_start(&SubagentStartContext {
|
||||
parent_session_id: "s1".into(),
|
||||
subagent_name: "sub".into(),
|
||||
prompt: "test".into(),
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_subagent_stop(&SubagentStopContext {
|
||||
parent_session_id: "s1".into(),
|
||||
subagent_name: "sub".into(),
|
||||
result_summary: "done".into(),
|
||||
steps: 3,
|
||||
is_error: false,
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_pre_compact(&PreCompactContext {
|
||||
session_id: "s1".into(),
|
||||
message_count: 50,
|
||||
estimated_tokens: 10000,
|
||||
})
|
||||
.await;
|
||||
registry
|
||||
.run_on_post_compact(&PostCompactContext {
|
||||
session_id: "s1".into(),
|
||||
new_message_count: 10,
|
||||
compression_method: "micro".into(),
|
||||
})
|
||||
.await;
|
||||
|
||||
// If no panic, all hooks were called successfully
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_match_filter_skips_irrelevant() {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
struct CountedHook {
|
||||
name: String,
|
||||
call_count: Arc<AtomicUsize>,
|
||||
filter: ToolMatchFilter,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for CountedHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
fn match_filter(&self) -> ToolMatchFilter {
|
||||
self.filter.clone()
|
||||
}
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
self.call_count.fetch_add(1, Ordering::SeqCst);
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
self.call_count.fetch_add(1, Ordering::SeqCst);
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
let search_count = Arc::new(AtomicUsize::new(0));
|
||||
let all_count = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let search_only = CountedHook {
|
||||
name: "search_only".into(),
|
||||
call_count: search_count.clone(),
|
||||
filter: ToolMatchFilter::exact("search_papers"),
|
||||
};
|
||||
let all_match = CountedHook {
|
||||
name: "all_match".into(),
|
||||
call_count: all_count.clone(),
|
||||
filter: ToolMatchFilter::default(),
|
||||
};
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
registry.add(Box::new(search_only));
|
||||
registry.add(Box::new(all_match));
|
||||
|
||||
let ctx = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 1,
|
||||
};
|
||||
registry.run_pre_tool_use(&ctx).await;
|
||||
assert_eq!(search_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(all_count.load(Ordering::SeqCst), 1);
|
||||
|
||||
let ctx2 = PreToolUseContext {
|
||||
session_id: "test".into(),
|
||||
tool_name: "download_paper".into(),
|
||||
tool_args: serde_json::json!({}),
|
||||
step: 2,
|
||||
};
|
||||
registry.run_pre_tool_use(&ctx2).await;
|
||||
assert_eq!(search_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(all_count.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_full_hook_pipeline_integration() {
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
struct IntegrationHook {
|
||||
name: String,
|
||||
filter: ToolMatchFilter,
|
||||
post_calls: Arc<std::sync::atomic::AtomicUsize>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentHook for IntegrationHook {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
fn match_filter(&self) -> ToolMatchFilter {
|
||||
self.filter.clone()
|
||||
}
|
||||
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
self.post_calls
|
||||
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
if self.name == "warn_hook" {
|
||||
PostToolUseAction::Warning {
|
||||
message: "test_warning".into(),
|
||||
truncate_output: false,
|
||||
}
|
||||
} else if self.name == "meta_hook" {
|
||||
PostToolUseAction::Metadata {
|
||||
key: "origin".into(),
|
||||
value: serde_json::Value::String("integration_test".into()),
|
||||
}
|
||||
} else {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content: format!("[{}] {}", self.name, ctx.output_content),
|
||||
additional_context: Some(format!("context_from_{}", self.name)),
|
||||
}
|
||||
}
|
||||
}
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = HookRegistry::new();
|
||||
let file_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
let search_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
|
||||
registry.add(Box::new(IntegrationHook {
|
||||
name: "file_hook".into(),
|
||||
filter: ToolMatchFilter::prefix("file_"),
|
||||
post_calls: file_count.clone(),
|
||||
}));
|
||||
registry.add(Box::new(IntegrationHook {
|
||||
name: "search_hook".into(),
|
||||
filter: ToolMatchFilter::exact("search_papers"),
|
||||
post_calls: search_count.clone(),
|
||||
}));
|
||||
registry.add(Box::new(IntegrationHook {
|
||||
name: "warn_hook".into(),
|
||||
filter: ToolMatchFilter::default(),
|
||||
post_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
|
||||
}));
|
||||
registry.add(Box::new(IntegrationHook {
|
||||
name: "meta_hook".into(),
|
||||
filter: ToolMatchFilter::default(),
|
||||
post_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
|
||||
}));
|
||||
|
||||
let ctx = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "file_write".into(),
|
||||
tool_args: serde_json::json!({"path": "/tmp/test.txt"}),
|
||||
output_content: "file content".into(),
|
||||
is_error: false,
|
||||
step: 1,
|
||||
elapsed_ms: 100,
|
||||
};
|
||||
let result = registry.run_post_tool_use(&ctx).await;
|
||||
|
||||
assert_eq!(file_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(search_count.load(Ordering::SeqCst), 0);
|
||||
assert!(result.final_content.contains("[file_hook]"));
|
||||
assert!(!result.tagged_contexts.is_empty());
|
||||
assert!(result
|
||||
.tagged_contexts
|
||||
.iter()
|
||||
.any(|tc| tc.hook_name == "file_hook"));
|
||||
assert!(result.warnings.contains(&"test_warning".to_string()));
|
||||
assert_eq!(
|
||||
result.metadata.get("origin").and_then(|v| v.as_str()),
|
||||
Some("integration_test")
|
||||
);
|
||||
|
||||
let ctx2 = PostToolUseContext {
|
||||
session_id: "test".into(),
|
||||
agent_name: "lead".into(),
|
||||
tool_name: "search_papers".into(),
|
||||
tool_args: serde_json::json!({"query": "quasars"}),
|
||||
output_content: "search results".into(),
|
||||
is_error: false,
|
||||
step: 2,
|
||||
elapsed_ms: 200,
|
||||
};
|
||||
let result2 = registry.run_post_tool_use(&ctx2).await;
|
||||
assert_eq!(search_count.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(file_count.load(Ordering::SeqCst), 1);
|
||||
assert!(result2.final_content.contains("[search_hook]"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,280 @@
|
||||
// src/agent/hooks/registry.rs
|
||||
//! HookRegistry — 注册表核心。
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::builtins::{AuditLogHook, CancellationHook, MetricsHook};
|
||||
use super::traits::{AgentHook, AsyncAgentHook};
|
||||
use super::types::{BlockingError, HookEvent, MetricsData, PreToolUseAction, TaggedContext};
|
||||
use serde_json;
|
||||
|
||||
/// Hook 注册表，管理所有已注册的 hook 并按序调用
|
||||
pub struct HookRegistry {
|
||||
pub(crate) hooks: Vec<Box<dyn AgentHook>>,
|
||||
/// 事件→hook 索引映射。预计算哪些 hook 关心哪些事件，
|
||||
/// 避免每次事件触发时遍历全部 hook。
|
||||
pub(crate) event_index: HashMap<HookEvent, Vec<usize>>,
|
||||
/// Per-session 动态 hooks（运行时添加/移除，会话结束时自动清理）
|
||||
pub(crate) session_hooks: HashMap<String, Vec<Box<dyn AgentHook>>>,
|
||||
/// 异步 fire-and-forget hooks（Phase 4.1+）
|
||||
pub(crate) async_hooks: Vec<Box<dyn AsyncAgentHook>>,
|
||||
/// 异步 hook 的事件索引
|
||||
pub(crate) async_event_index: HashMap<HookEvent, Vec<usize>>,
|
||||
}
|
||||
|
||||
impl Default for HookRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl HookRegistry {
|
||||
/// 创建空的注册表
|
||||
pub fn new() -> Self {
|
||||
HookRegistry {
|
||||
hooks: Vec::new(),
|
||||
event_index: HashMap::new(),
|
||||
session_hooks: HashMap::new(),
|
||||
async_hooks: Vec::new(),
|
||||
async_event_index: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建包含所有内置 hooks 的注册表。
|
||||
///
|
||||
/// 参数：
|
||||
/// - `db`: 数据库连接池（供 AuditLogHook 持久化）
|
||||
/// - `cancelled_runs`: 取消状态集合（供 CancellationHook 检查）
|
||||
/// - `metrics_data`: 可选的共享指标数据引用。提供时复用已有的 MetricsData，
|
||||
/// 使得 AgentRuntime.get_metrics() 能查询到实际运行数据。
|
||||
pub fn with_builtins(
|
||||
db: SqlitePool,
|
||||
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>,
|
||||
metrics_data: Option<Arc<std::sync::Mutex<MetricsData>>>,
|
||||
) -> Self {
|
||||
let mut registry = Self::new();
|
||||
registry.add(Box::new(CancellationHook::new(cancelled_runs)));
|
||||
// 如果提供了共享的 metrics_data，使用它；否则创建新的
|
||||
let metrics_hook = match metrics_data {
|
||||
Some(data) => MetricsHook::from_arc(data),
|
||||
None => MetricsHook::new(),
|
||||
};
|
||||
registry.add(Box::new(metrics_hook));
|
||||
registry.add(Box::new(AuditLogHook::new(db)));
|
||||
registry
|
||||
}
|
||||
|
||||
/// 注册一个 hook 并重建事件索引
|
||||
pub fn add(&mut self, hook: Box<dyn AgentHook>) {
|
||||
info!("[Hooks] 注册 hook: {}", hook.name());
|
||||
let idx = self.hooks.len();
|
||||
let subs = hook.subscribed_events();
|
||||
if subs.is_empty() {
|
||||
// 空切片 = 订阅所有事件
|
||||
for event in &[
|
||||
HookEvent::OnSessionStart,
|
||||
HookEvent::PreToolUse,
|
||||
HookEvent::PostToolUse,
|
||||
HookEvent::PostToolUseFailure,
|
||||
HookEvent::OnStepComplete,
|
||||
HookEvent::OnSessionStop,
|
||||
HookEvent::OnSubagentStart,
|
||||
HookEvent::OnSubagentStop,
|
||||
HookEvent::OnPreCompact,
|
||||
HookEvent::OnPostCompact,
|
||||
] {
|
||||
self.event_index.entry(*event).or_default().push(idx);
|
||||
}
|
||||
} else {
|
||||
for event in subs {
|
||||
self.event_index.entry(*event).or_default().push(idx);
|
||||
}
|
||||
}
|
||||
self.hooks.push(hook);
|
||||
}
|
||||
|
||||
/// 注册一个异步 fire-and-forget hook（Phase 4.1+）
|
||||
pub fn add_async(&mut self, hook: Box<dyn AsyncAgentHook>) {
|
||||
info!("[Hooks] 注册异步 hook: {}", hook.name());
|
||||
let idx = self.async_hooks.len();
|
||||
let subs = hook.subscribed_events();
|
||||
if subs.is_empty() {
|
||||
for event in &[
|
||||
HookEvent::OnSessionStart,
|
||||
HookEvent::PreToolUse,
|
||||
HookEvent::PostToolUse,
|
||||
HookEvent::PostToolUseFailure,
|
||||
HookEvent::OnStepComplete,
|
||||
HookEvent::OnSessionStop,
|
||||
HookEvent::OnSubagentStart,
|
||||
HookEvent::OnSubagentStop,
|
||||
HookEvent::OnPreCompact,
|
||||
HookEvent::OnPostCompact,
|
||||
] {
|
||||
self.async_event_index.entry(*event).or_default().push(idx);
|
||||
}
|
||||
} else {
|
||||
for event in subs {
|
||||
self.async_event_index.entry(*event).or_default().push(idx);
|
||||
}
|
||||
}
|
||||
self.async_hooks.push(hook);
|
||||
}
|
||||
|
||||
/// 获取所有 hooks 的不可变引用
|
||||
pub fn all(&self) -> &[Box<dyn AgentHook>] {
|
||||
&self.hooks
|
||||
}
|
||||
|
||||
/// 获取订阅了指定事件的 hook 索引列表
|
||||
pub(crate) fn indices_for(&self, event: HookEvent) -> &[usize] {
|
||||
self.event_index
|
||||
.get(&event)
|
||||
.map(|v| v.as_slice())
|
||||
.unwrap_or(&[])
|
||||
}
|
||||
|
||||
// ── Session Hook 管理 ──
|
||||
|
||||
/// 为指定会话动态添加一个 hook（仅当前会话有效，Stop 时自动清理）。
|
||||
pub fn add_session_hook(&mut self, session_id: &str, hook: Box<dyn AgentHook>) {
|
||||
info!(
|
||||
"[Hooks] 为会话 {} 添加 session hook: {}",
|
||||
session_id,
|
||||
hook.name()
|
||||
);
|
||||
self.session_hooks
|
||||
.entry(session_id.to_string())
|
||||
.or_default()
|
||||
.push(hook);
|
||||
}
|
||||
|
||||
/// 移除指定会话中名为 `hook_name` 的 hook。返回是否成功移除。
|
||||
pub fn remove_session_hook(&mut self, session_id: &str, hook_name: &str) -> bool {
|
||||
if let Some(hooks) = self.session_hooks.get_mut(session_id) {
|
||||
let before = hooks.len();
|
||||
hooks.retain(|h| h.name() != hook_name);
|
||||
let removed = hooks.len() < before;
|
||||
if removed {
|
||||
info!(
|
||||
"[Hooks] 从会话 {} 移除 session hook: {}",
|
||||
session_id, hook_name
|
||||
);
|
||||
}
|
||||
if hooks.is_empty() {
|
||||
self.session_hooks.remove(session_id);
|
||||
}
|
||||
return removed;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// 清理指定会话的所有动态 hooks
|
||||
pub fn clear_session_hooks(&mut self, session_id: &str) {
|
||||
if self.session_hooks.remove(session_id).is_some() {
|
||||
info!("[Hooks] 清理会话 {} 的所有 session hooks", session_id);
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取指定会话的动态 hooks 的不可变引用
|
||||
pub(crate) fn get_session_hooks(&self, session_id: &str) -> &[Box<dyn AgentHook>] {
|
||||
self.session_hooks
|
||||
.get(session_id)
|
||||
.map(|v| v.as_slice())
|
||||
.unwrap_or(&[])
|
||||
}
|
||||
|
||||
/// 收集全局 hooks（按事件索引）和 session hooks 到统一列表，
|
||||
/// 避免两个迭代器类型不同导致的 chain/extend 类型错误。
|
||||
pub(crate) fn collect_hooks_for<'a>(
|
||||
&'a self,
|
||||
event: HookEvent,
|
||||
session_id: &str,
|
||||
) -> Vec<&'a dyn AgentHook> {
|
||||
let mut hooks: Vec<&dyn AgentHook> = self
|
||||
.indices_for(event)
|
||||
.iter()
|
||||
.filter_map(|&idx| self.hooks.get(idx).map(|b| b.as_ref()))
|
||||
.collect();
|
||||
hooks.extend(
|
||||
self.get_session_hooks(session_id)
|
||||
.iter()
|
||||
.map(|b| b.as_ref()),
|
||||
);
|
||||
hooks
|
||||
}
|
||||
|
||||
/// 收集工具相关事件的 hooks（按事件索引 + session hooks），
|
||||
/// 并额外按 `match_filter()` 过滤——只保留关心当前工具名的 hook。
|
||||
///
|
||||
/// 如果 hook 的 `match_filter()` 返回默认 "match all"，该 hook 永远包含。
|
||||
/// 仅用于 `PreToolUse`、`PostToolUse`、`PostToolUseFailure` 事件。
|
||||
pub(crate) fn collect_tool_hooks_for<'a>(
|
||||
&'a self,
|
||||
event: HookEvent,
|
||||
session_id: &str,
|
||||
tool_name: &str,
|
||||
tool_args: &serde_json::Value,
|
||||
) -> Vec<&'a dyn AgentHook> {
|
||||
let mut hooks: Vec<&dyn AgentHook> = self
|
||||
.indices_for(event)
|
||||
.iter()
|
||||
.filter_map(|&idx| self.hooks.get(idx).map(|b| b.as_ref()))
|
||||
.filter(|hook| {
|
||||
let filter = hook.match_filter();
|
||||
filter.is_match_all() || filter.matches(tool_name, tool_args)
|
||||
})
|
||||
.collect();
|
||||
hooks.extend(
|
||||
self.get_session_hooks(session_id)
|
||||
.iter()
|
||||
.map(|b| b.as_ref()),
|
||||
);
|
||||
hooks
|
||||
}
|
||||
|
||||
/// 处理单个 PreToolUse action 的聚合逻辑（共享于顺序/并行路径）
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn process_pre_tool_action(
|
||||
action: &PreToolUseAction,
|
||||
hook_name: &str,
|
||||
tool_name: &str,
|
||||
blocking_errors: &mut Vec<BlockingError>,
|
||||
final_args: &mut serde_json::Value,
|
||||
additional_contexts: &mut Vec<String>,
|
||||
tagged_contexts: &mut Vec<TaggedContext>,
|
||||
final_action: &mut PreToolUseAction,
|
||||
) {
|
||||
match action {
|
||||
PreToolUseAction::Block { reason } => {
|
||||
warn!("[Hooks] {hook_name} 阻止了工具 {tool_name} 的执行: {reason}");
|
||||
blocking_errors.push(BlockingError::new(hook_name.to_string(), reason.clone()));
|
||||
}
|
||||
PreToolUseAction::MutateInput {
|
||||
updated_args,
|
||||
additional_context,
|
||||
} => {
|
||||
info!("[Hooks] {hook_name} 修改了工具 {tool_name} 的输入参数");
|
||||
*final_args = updated_args.clone();
|
||||
if let Some(ctx_str) = additional_context {
|
||||
additional_contexts.push(ctx_str.clone());
|
||||
tagged_contexts.push(TaggedContext::new(
|
||||
hook_name,
|
||||
HookEvent::PreToolUse,
|
||||
ctx_str,
|
||||
));
|
||||
}
|
||||
}
|
||||
PreToolUseAction::PermissionRequired { .. } => {
|
||||
info!("[Hooks] {hook_name} 请求了工具 {tool_name} 的权限检查");
|
||||
}
|
||||
PreToolUseAction::Continue => {}
|
||||
}
|
||||
if !matches!(action, PreToolUseAction::Continue) {
|
||||
*final_action = action.clone();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
// src/agent/hooks/traits.rs
|
||||
//
|
||||
// Agent 生命周期 Hook traits。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::matcher::ToolMatchFilter;
|
||||
use super::types::{
|
||||
HookEvent, PermissionDeniedContext, PermissionRequestAction, PermissionRequestContext,
|
||||
PostCompactContext, PostToolUseAction, PostToolUseContext, PostToolUseFailureContext,
|
||||
PreCompactContext, PreToolUseAction, PreToolUseContext, SessionStartContext,
|
||||
SessionStopContext, StepCompleteContext, SubagentStartContext, SubagentStopContext,
|
||||
};
|
||||
|
||||
// ── Hook Trait ──
|
||||
|
||||
/// Agent 生命周期 Hook trait。
|
||||
/// 所有方法都有默认空实现，只需覆写关心的 hook 点。
|
||||
#[async_trait]
|
||||
pub trait AgentHook: Send + Sync {
|
||||
/// Hook 名称（用于日志和调试）
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// 声明该 hook 订阅的生命周期事件。
|
||||
/// 返回空切片表示订阅所有事件（向后兼容的默认行为）。
|
||||
/// HookRegistry 用此信息预计算事件→hook 索引，避免无关调用。
|
||||
fn subscribed_events(&self) -> &[HookEvent] {
|
||||
&[] // 空 = 订阅全部
|
||||
}
|
||||
|
||||
/// Per-hook 超时时间。返回 `None` 使用全局默认值 (30s)。
|
||||
fn timeout(&self) -> Option<Duration> {
|
||||
None
|
||||
}
|
||||
|
||||
/// 声明此 hook 关心哪些工具调用。
|
||||
///
|
||||
/// 返回的 `ToolMatchFilter` 用于在 dispatch 时过滤无关的工具调用，
|
||||
/// 减少不必要的 hook 执行。默认返回空过滤器（匹配所有工具）。
|
||||
///
|
||||
/// 仅对 `PreToolUse`、`PostToolUse`、`PostToolUseFailure` 事件生效。
|
||||
/// 其他事件忽略此过滤器。
|
||||
fn match_filter(&self) -> ToolMatchFilter {
|
||||
ToolMatchFilter::default()
|
||||
}
|
||||
|
||||
// ── 原有 5 个生命周期事件 ──
|
||||
|
||||
/// 会话创建/恢复时调用。
|
||||
async fn on_session_start(&self, _ctx: &SessionStartContext) {}
|
||||
|
||||
/// 工具执行前调用。可返回 Continue/Block/MutateInput/PermissionRequired。
|
||||
async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction {
|
||||
PreToolUseAction::Continue
|
||||
}
|
||||
|
||||
/// 工具执行后调用。可返回 Continue 或 MutateOutput。
|
||||
async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction {
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
/// 每个 ReAct step 完成后调用。
|
||||
async fn on_step_complete(&self, _ctx: &StepCompleteContext) {}
|
||||
|
||||
/// 会话终止时调用。
|
||||
async fn on_session_stop(&self, _ctx: &SessionStopContext<'_>) {}
|
||||
|
||||
// ── 新增 4 个生命周期事件（默认 no-op） ──
|
||||
|
||||
/// 子代理启动时调用。
|
||||
async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {}
|
||||
|
||||
/// 子代理停止时调用。
|
||||
async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {}
|
||||
|
||||
/// 上下文压缩前调用。
|
||||
async fn on_pre_compact(&self, _ctx: &PreCompactContext) {}
|
||||
|
||||
/// 上下文压缩后调用。
|
||||
async fn on_post_compact(&self, _ctx: &PostCompactContext) {}
|
||||
|
||||
/// 工具执行失败时调用（独立于 PostToolUse，专注错误处理）。
|
||||
async fn on_post_tool_use_failure(
|
||||
&self,
|
||||
_ctx: &PostToolUseFailureContext,
|
||||
) -> PostToolUseAction {
|
||||
PostToolUseAction::Continue
|
||||
}
|
||||
|
||||
/// 权限请求前调用（参考 Claude Code PermissionRequest hook）。
|
||||
///
|
||||
/// Hook 可以覆盖权限决策或注入附加上下文。此方法在 PermissionChecker
|
||||
/// 做出初步决策后、最终返回前触发。
|
||||
///
|
||||
/// 返回 `PermissionRequestAction::Continue` 保持当前决策不变。
|
||||
async fn on_permission_request(
|
||||
&self,
|
||||
_ctx: &PermissionRequestContext,
|
||||
) -> PermissionRequestAction {
|
||||
PermissionRequestAction::Continue
|
||||
}
|
||||
|
||||
/// 权限被拒绝后调用（参考 Claude Code PermissionDenied hook）。
|
||||
///
|
||||
/// 仅用于审计/日志/监控——返回值不影响执行流程。
|
||||
/// 在权限被最终拒绝后触发（规则拒绝、Classifier 拒绝或用户拒绝）。
|
||||
async fn on_permission_denied(&self, _ctx: &PermissionDeniedContext) {}
|
||||
}
|
||||
|
||||
// ── Async Hook Trait ──
|
||||
|
||||
/// Fire-and-forget hook trait（Phase 4.1+）。
|
||||
///
|
||||
/// 与 `AgentHook` 分离设计：
|
||||
/// - `AgentHook` 的返回值（`PreToolUseAction` / `PostToolUseAction`）会立即影响执行流程
|
||||
/// - `AsyncAgentHook` 不返回决策——它适合用于后台操作（上传、远程日志上报等）
|
||||
///
|
||||
/// HookRegistry 并行调度两类 hook：sync hook 的结果用于决策，async hook 仅 fire-and-forget。
|
||||
/// Async hook dispatch 默认超时 5 秒（仅限方法返回，后台工作继续独立执行）。
|
||||
#[async_trait]
|
||||
pub trait AsyncAgentHook: Send + Sync {
|
||||
/// Hook 名称（用于日志和调试）
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// 声明订阅的事件（默认全部）。与 AgentHook 相同的索引机制。
|
||||
fn subscribed_events(&self) -> &[HookEvent] {
|
||||
&[]
|
||||
}
|
||||
|
||||
/// 工具执行后异步回调（fire-and-forget）
|
||||
async fn on_post_tool_use_async(&self, _ctx: PostToolUseContext) {}
|
||||
|
||||
/// 工具执行失败时异步回调
|
||||
async fn on_post_tool_use_failure_async(&self, _ctx: PostToolUseFailureContext) {}
|
||||
|
||||
/// 会话结束时异步回调
|
||||
async fn on_session_stop_async(&self, _ctx: SessionStopContext<'_>) {}
|
||||
}
|
||||
@@ -0,0 +1,464 @@
|
||||
// src/agent/hooks/types.rs
|
||||
//
|
||||
// Agent 生命周期 Hooks 系统 — 类型定义。
|
||||
//
|
||||
// 从 src/agent/hooks/mod.rs 提取，包含所有数据结构、枚举和结果类型。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono;
|
||||
use serde_json;
|
||||
|
||||
use crate::agent::terminal::TurnTerminal;
|
||||
|
||||
/// 默认 per-hook 超时时间（30 秒），可通过 HOOK_TIMEOUT_SECS 环境变量覆盖。
|
||||
pub const DEFAULT_HOOK_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
// ── Hook Event Enum ──
|
||||
|
||||
/// 生命周期事件类型，用于事件订阅索引。
|
||||
/// 每个 hook 通过 `subscribed_events()` 声明关心的事件，
|
||||
/// HookRegistry 据此预计算索引，避免无关调用。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum HookEvent {
|
||||
OnSessionStart,
|
||||
PreToolUse,
|
||||
PostToolUse,
|
||||
PostToolUseFailure,
|
||||
OnStepComplete,
|
||||
OnSessionStop,
|
||||
OnSubagentStart,
|
||||
OnSubagentStop,
|
||||
OnPreCompact,
|
||||
OnPostCompact,
|
||||
/// 权限请求前触发（参考 Claude Code PermissionRequest hook）。
|
||||
/// Hook 可以修改权限决策或注入附加上下文。
|
||||
PermissionRequest,
|
||||
/// 权限被拒绝后触发（参考 Claude Code PermissionDenied hook）。
|
||||
/// 用于审计日志和安全监控。
|
||||
PermissionDenied,
|
||||
}
|
||||
|
||||
// ── Hook Contexts ──
|
||||
|
||||
/// 会话启动上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionStartContext {
|
||||
pub session_id: String,
|
||||
pub turn_index: i32,
|
||||
pub is_resume: bool,
|
||||
}
|
||||
|
||||
/// PreToolUse hook 上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreToolUseContext {
|
||||
pub session_id: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
pub step: usize,
|
||||
}
|
||||
|
||||
/// PostToolUse hook 上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostToolUseContext {
|
||||
pub session_id: String,
|
||||
pub agent_name: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
pub output_content: String,
|
||||
pub is_error: bool,
|
||||
pub step: usize,
|
||||
pub elapsed_ms: u64,
|
||||
}
|
||||
|
||||
/// Step 完成上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StepCompleteContext {
|
||||
pub session_id: String,
|
||||
pub step: usize,
|
||||
pub max_steps: usize,
|
||||
pub messages_count: usize,
|
||||
pub estimated_tokens: usize,
|
||||
pub token_limit: usize,
|
||||
}
|
||||
|
||||
/// Session 终止上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionStopContext<'a> {
|
||||
pub session_id: String,
|
||||
pub terminal: &'a TurnTerminal,
|
||||
pub total_steps: usize,
|
||||
}
|
||||
|
||||
/// 子代理启动上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubagentStartContext {
|
||||
pub parent_session_id: String,
|
||||
pub subagent_name: String,
|
||||
pub prompt: String,
|
||||
}
|
||||
|
||||
/// 子代理停止上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubagentStopContext {
|
||||
pub parent_session_id: String,
|
||||
pub subagent_name: String,
|
||||
pub result_summary: String,
|
||||
pub steps: usize,
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
/// 压缩前上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreCompactContext {
|
||||
pub session_id: String,
|
||||
pub message_count: usize,
|
||||
pub estimated_tokens: usize,
|
||||
}
|
||||
|
||||
/// 压缩后上下文
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostCompactContext {
|
||||
pub session_id: String,
|
||||
pub new_message_count: usize,
|
||||
pub compression_method: String,
|
||||
}
|
||||
|
||||
/// 工具执行失败上下文（参考 Claude Code 的 PostToolUseFailure 事件）
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostToolUseFailureContext {
|
||||
pub session_id: String,
|
||||
pub agent_name: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
pub error_message: String,
|
||||
pub is_interrupt: bool,
|
||||
pub step: usize,
|
||||
pub elapsed_ms: u64,
|
||||
}
|
||||
|
||||
/// 权限请求上下文（参考 Claude Code PermissionRequest hook）。
|
||||
///
|
||||
/// 在权限检查器做出最终决策前触发。Hook 可以：
|
||||
/// - 通过返回 `PermissionRequestAction::Override` 修改决策
|
||||
/// - 注入审计上下文
|
||||
/// - 记录安全日志
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PermissionRequestContext {
|
||||
pub session_id: String,
|
||||
pub agent_name: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
/// 权限检查器当前的决策
|
||||
pub current_decision: PermissionDecision,
|
||||
/// 权限模式（default / accept_edits / bypass / dont_ask）
|
||||
pub permission_mode: String,
|
||||
pub step: usize,
|
||||
/// 是否为子代理调用（静默模式）
|
||||
pub is_subagent: bool,
|
||||
}
|
||||
|
||||
/// 权限检查决策（供 PermissionRequestContext 使用）
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PermissionDecision {
|
||||
Allowed,
|
||||
Denied { reason: String },
|
||||
AskUser { message: String },
|
||||
}
|
||||
|
||||
/// 权限请求 Hook 的动作
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PermissionRequestAction {
|
||||
/// 继续使用当前决策
|
||||
Continue,
|
||||
/// 覆盖当前决策
|
||||
Override(PermissionDecision),
|
||||
/// 注入附加上下文但不改变决策
|
||||
InjectContext { content: String },
|
||||
}
|
||||
|
||||
/// 权限被拒绝上下文（参考 Claude Code PermissionDenied hook）。
|
||||
///
|
||||
/// 在权限被最终拒绝后触发。用于：
|
||||
/// - 审计日志记录
|
||||
/// - 安全事件监控
|
||||
/// - 告警通知
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PermissionDeniedContext {
|
||||
pub session_id: String,
|
||||
pub agent_name: String,
|
||||
pub tool_name: String,
|
||||
pub tool_args: serde_json::Value,
|
||||
/// 拒绝原因
|
||||
pub reason: String,
|
||||
/// 拒绝来源
|
||||
pub source: PermissionDenialSource,
|
||||
pub step: usize,
|
||||
}
|
||||
|
||||
/// 权限拒绝来源
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PermissionDenialSource {
|
||||
/// 被规则列表明确拒绝
|
||||
Rule,
|
||||
/// Auto-mode Classifier 判定为拒绝
|
||||
Classifier,
|
||||
/// 用户交互选择拒绝
|
||||
User,
|
||||
/// 超时自动拒绝
|
||||
Timeout,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for PermissionDenialSource {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
PermissionDenialSource::Rule => write!(f, "rule"),
|
||||
PermissionDenialSource::Classifier => write!(f, "classifier"),
|
||||
PermissionDenialSource::User => write!(f, "user"),
|
||||
PermissionDenialSource::Timeout => write!(f, "timeout"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Blocking Error ──
|
||||
|
||||
/// 描述单个 hook 的阻塞错误。
|
||||
/// 取代原来只保留第一个 Block 的做法，允许多个 hook 的阻塞原因都被记录。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BlockingError {
|
||||
pub hook_name: String,
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
impl BlockingError {
|
||||
pub fn new(hook_name: impl Into<String>, reason: impl Into<String>) -> Self {
|
||||
BlockingError {
|
||||
hook_name: hook_name.into(),
|
||||
reason: reason.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tagged Context ──
|
||||
|
||||
/// Human-readable label for a HookEvent, used in system-reminder formatting.
|
||||
pub fn event_label(event: HookEvent) -> &'static str {
|
||||
match event {
|
||||
HookEvent::OnSessionStart => "SessionStart",
|
||||
HookEvent::PreToolUse => "PreToolUse",
|
||||
HookEvent::PostToolUse => "PostToolUse",
|
||||
HookEvent::PostToolUseFailure => "PostToolUseFailure",
|
||||
HookEvent::OnStepComplete => "StepComplete",
|
||||
HookEvent::OnSessionStop => "SessionStop",
|
||||
HookEvent::OnSubagentStart => "SubagentStart",
|
||||
HookEvent::OnSubagentStop => "SubagentStop",
|
||||
HookEvent::OnPreCompact => "PreCompact",
|
||||
HookEvent::OnPostCompact => "PostCompact",
|
||||
HookEvent::PermissionRequest => "PermissionRequest",
|
||||
HookEvent::PermissionDenied => "PermissionDenied",
|
||||
}
|
||||
}
|
||||
|
||||
/// Tagged context from a hook, preserving which hook produced it and why.
|
||||
///
|
||||
/// Unlike the anonymous `additional_contexts` strings, `TaggedContext` carries
|
||||
/// the hook's name and the lifecycle event that triggered it. The executor
|
||||
/// uses this information to produce clearer system-reminders in the LLM
|
||||
/// transcript.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TaggedContext {
|
||||
/// Name of the hook that produced this context.
|
||||
pub hook_name: String,
|
||||
/// Which lifecycle event triggered this context injection.
|
||||
pub source_event: HookEvent,
|
||||
/// The actual context content.
|
||||
pub content: String,
|
||||
/// Wall-clock time when this context was created.
|
||||
pub timestamp: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
impl TaggedContext {
|
||||
pub fn new(
|
||||
hook_name: impl Into<String>,
|
||||
source_event: HookEvent,
|
||||
content: impl Into<String>,
|
||||
) -> Self {
|
||||
TaggedContext {
|
||||
hook_name: hook_name.into(),
|
||||
source_event,
|
||||
content: content.into(),
|
||||
timestamp: chrono::Utc::now(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Hook Actions ──
|
||||
|
||||
/// PreToolUse hook 返回的增强动作。
|
||||
/// 支持：允许、阻止、修改输入、权限请求。
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PreToolUseAction {
|
||||
/// 允许继续执行（默认）
|
||||
Continue,
|
||||
/// 阻止执行，附带原因
|
||||
Block { reason: String },
|
||||
/// 允许执行但修改输入参数或注入附加上下文
|
||||
MutateInput {
|
||||
updated_args: serde_json::Value,
|
||||
additional_context: Option<String>,
|
||||
},
|
||||
/// 需要权限决策
|
||||
PermissionRequired {
|
||||
permission: String,
|
||||
tool_name: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl PreToolUseAction {
|
||||
pub fn is_blocked(&self) -> bool {
|
||||
matches!(self, PreToolUseAction::Block { .. })
|
||||
}
|
||||
|
||||
pub fn block_reason(&self) -> Option<&str> {
|
||||
match self {
|
||||
PreToolUseAction::Block { reason } => Some(reason),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn updated_args(&self) -> Option<&serde_json::Value> {
|
||||
match self {
|
||||
PreToolUseAction::MutateInput { updated_args, .. } => Some(updated_args),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn additional_context(&self) -> Option<&str> {
|
||||
match self {
|
||||
PreToolUseAction::MutateInput {
|
||||
additional_context, ..
|
||||
} => additional_context.as_deref(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 向后兼容类型别名 — 旧代码用 HookAction::Continue / HookAction::Block 仍可编译
|
||||
pub type HookAction = PreToolUseAction;
|
||||
|
||||
/// PostToolUse hook 返回的动作。
|
||||
/// 支持：保持输出不变、修改输出内容、警告标记、事后权限请求、结构化元数据。
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PostToolUseAction {
|
||||
/// 保持原输出不变（默认）
|
||||
Continue,
|
||||
/// 修改输出内容，可选注入附加上下文供 LLM 后续推理使用
|
||||
MutateOutput {
|
||||
updated_content: String,
|
||||
additional_context: Option<String>,
|
||||
},
|
||||
/// 标记工具结果为警告（不阻止，仅供 LLM 参考）
|
||||
Warning {
|
||||
message: String,
|
||||
/// 若为 true，截断工具输出以避免模型上下文过载
|
||||
truncate_output: bool,
|
||||
},
|
||||
/// 事后请求权限（post-hoc advisory，不阻塞主循环）
|
||||
PermissionRequired {
|
||||
permission: String,
|
||||
tool_name: String,
|
||||
},
|
||||
/// 附加结构化元数据到结果中（保留在审计日志中）
|
||||
Metadata {
|
||||
key: String,
|
||||
value: serde_json::Value,
|
||||
},
|
||||
}
|
||||
|
||||
impl PostToolUseAction {
|
||||
pub fn is_continue(&self) -> bool {
|
||||
matches!(self, PostToolUseAction::Continue)
|
||||
}
|
||||
|
||||
pub fn updated_content(&self) -> Option<&str> {
|
||||
match self {
|
||||
PostToolUseAction::MutateOutput {
|
||||
updated_content, ..
|
||||
} => Some(updated_content),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn warning_message(&self) -> Option<&str> {
|
||||
match self {
|
||||
PostToolUseAction::Warning { message, .. } => Some(message),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Metrics Data ──
|
||||
|
||||
/// 可查询的运行指标快照
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct MetricsData {
|
||||
pub tool_call_counts: HashMap<String, usize>,
|
||||
pub total_steps: usize,
|
||||
pub total_errors: usize,
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
// ── Hook Trait ──
|
||||
|
||||
/// PreToolUse 聚合结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreToolUseResult {
|
||||
/// 最终动作（第一个非 Continue 动作获胜）
|
||||
pub action: PreToolUseAction,
|
||||
/// 累积的 additional_contexts（所有 MutateInput 的上下文，不丢失任何 hook 的贡献）
|
||||
pub additional_contexts: Vec<String>,
|
||||
/// 向后兼容便捷字段：用换行拼接的附加上下文
|
||||
pub additional_context: Option<String>,
|
||||
/// 所有阻塞错误（多个 hook 同时 block 时全部保留）
|
||||
pub blocking_errors: Vec<BlockingError>,
|
||||
/// 最终的工具参数（应用了最后一个 MutateInput 的修改）
|
||||
pub final_args: serde_json::Value,
|
||||
/// 带来源标记的上下文列表（Phase 1.3+：用于格式化有意义的 system-reminder）
|
||||
pub tagged_contexts: Vec<TaggedContext>,
|
||||
}
|
||||
|
||||
impl PreToolUseResult {
|
||||
/// 是否有 hook 请求了权限确认
|
||||
pub fn is_permission_required(&self) -> bool {
|
||||
matches!(self.action, PreToolUseAction::PermissionRequired { .. })
|
||||
}
|
||||
|
||||
/// 获取权限确认的详情（permission 描述, tool_name）
|
||||
pub fn permission_info(&self) -> Option<(&str, &str)> {
|
||||
match &self.action {
|
||||
PreToolUseAction::PermissionRequired {
|
||||
permission,
|
||||
tool_name,
|
||||
} => Some((permission.as_str(), tool_name.as_str())),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// PostToolUse 聚合结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostToolUseResult {
|
||||
/// 最终输出内容（应用了最后一个 MutateOutput 的修改）
|
||||
pub final_content: String,
|
||||
/// 累积的 additional_contexts（所有 PostToolUse hooks 注入的上下文）
|
||||
pub additional_contexts: Vec<String>,
|
||||
/// 带来源标记的上下文列表（Phase 1.3+：用于格式化有意义的 system-reminder）
|
||||
pub tagged_contexts: Vec<TaggedContext>,
|
||||
/// 非阻塞警告消息列表（Phase 3.1+）
|
||||
pub warnings: Vec<String>,
|
||||
/// 事后权限请求列表（post-hoc advisory，不阻塞主循环）
|
||||
pub post_permission_requests: Vec<(String, String)>,
|
||||
/// 结构化键值对元数据（Phase 3.1+）
|
||||
pub metadata: std::collections::HashMap<String, serde_json::Value>,
|
||||
}
|
||||
+13
-4
@@ -178,7 +178,11 @@ impl MemoryManager {
|
||||
String::new(),
|
||||
];
|
||||
|
||||
for entry in selected {
|
||||
// 按 mtime 降序排列
|
||||
let mut sorted: Vec<&&MemoryEntry> = selected.iter().collect();
|
||||
sorted.sort_by_key(|e| std::cmp::Reverse(e.mtime));
|
||||
|
||||
for entry in sorted {
|
||||
let type_tag = match entry.memory_type {
|
||||
MemoryType::User => "[偏好]",
|
||||
MemoryType::Feedback => "[反馈]",
|
||||
@@ -343,7 +347,8 @@ impl MemoryManager {
|
||||
|
||||
/// 生成 system prompt 中注入的记忆段落。
|
||||
///
|
||||
/// 包含最近的记忆条目（最多 10 条），并在前面注明可信度提醒。
|
||||
/// 按 recency 排序（最新的记忆在前），取 top N 条。
|
||||
/// 跳过已过期的条目（超过 memory_max_age_days 的条目由 age 模块标记）。
|
||||
pub fn build_system_reminder(&self, max_entries: usize) -> Option<String> {
|
||||
if self.entries.is_empty() {
|
||||
return None;
|
||||
@@ -356,8 +361,12 @@ impl MemoryManager {
|
||||
"".to_string(),
|
||||
];
|
||||
|
||||
let count = max_entries.min(self.entries.len());
|
||||
for entry in self.entries.iter().take(count) {
|
||||
// 按 mtime 降序排列（最新的在前），提升注入记忆的时效性
|
||||
let mut sorted_entries: Vec<&MemoryEntry> = self.entries.iter().collect();
|
||||
sorted_entries.sort_by_key(|e| std::cmp::Reverse(e.mtime));
|
||||
|
||||
let count = max_entries.min(sorted_entries.len());
|
||||
for entry in sorted_entries.iter().take(count) {
|
||||
let type_tag = match entry.memory_type {
|
||||
MemoryType::User => "[偏好]",
|
||||
MemoryType::Feedback => "[反馈]",
|
||||
|
||||
@@ -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() {
|
||||
// 全角 'ｒ' (U+FF52) → 半角 'r'
|
||||
let fullwidth_rm = "\u{FF52}\u{FF4D}"; // ｒｍ
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,11 @@
|
||||
// Layer 1 — system-reminder 注入：每轮动态列出 skill 名称（~20 tokens/skill）
|
||||
// Layer 2 — LoadSkillTool：LLM 按需调用，注入完整 skill 内容（~2000 tokens/skill）
|
||||
//
|
||||
// 自改进技能系统（参考 Hermes-Agent curator.py 设计）：
|
||||
// PatternDetector — 扫描工具调用序列，检测跨 session 重复模式
|
||||
// SkillCreator — 将检测到的模式生成 SKILL.md 文件
|
||||
// Curator — 管理 skill 生命周期，标记 stale/deprecated，建议清理
|
||||
//
|
||||
// Skill 文件格式（对齐 Claude Code 的目录约定）：
|
||||
// skills/{skill-name}/SKILL.md ← 必须是目录 + SKILL.md
|
||||
//
|
||||
@@ -26,6 +31,9 @@
|
||||
// # Skill 正文
|
||||
// 详细内容...
|
||||
|
||||
pub mod curator;
|
||||
pub mod pattern_detector;
|
||||
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
@@ -76,6 +84,9 @@ pub struct SkillFrontmatter {
|
||||
/// fork 模式下的 effort 级别
|
||||
#[serde(default)]
|
||||
pub effort: Option<String>,
|
||||
/// 是否禁止 Curator 自动清理（参考 Hermes pinned skills）
|
||||
#[serde(default)]
|
||||
pub pinned: Option<bool>,
|
||||
}
|
||||
|
||||
impl SkillFrontmatter {
|
||||
@@ -116,6 +127,8 @@ pub struct SkillMeta {
|
||||
pub user_invocable: bool,
|
||||
/// 条件激活的 glob 模式（空 Vec 表示始终激活）
|
||||
pub paths: Vec<String>,
|
||||
/// 是否禁止 Curator 自动清理（参考 Hermes pinned skills）
|
||||
pub pinned: bool,
|
||||
}
|
||||
|
||||
/// 完整的 Skill（Layer 2：LLM 调用 load_skill 时注入）
|
||||
@@ -583,6 +596,7 @@ fn load_skill_from_path(skill_md_path: &Path, dir_name: &str) -> Result<Skill, S
|
||||
let disable_model_invocation = frontmatter.disable_model_invocation.unwrap_or(false);
|
||||
let user_invocable = frontmatter.user_invocable.unwrap_or(true);
|
||||
let paths = frontmatter.paths.unwrap_or_default();
|
||||
let pinned = frontmatter.pinned.unwrap_or(false);
|
||||
let skill_dir = skill_md_path
|
||||
.parent()
|
||||
.unwrap_or(Path::new("."))
|
||||
@@ -603,6 +617,7 @@ fn load_skill_from_path(skill_md_path: &Path, dir_name: &str) -> Result<Skill, S
|
||||
disable_model_invocation,
|
||||
user_invocable,
|
||||
paths,
|
||||
pinned,
|
||||
},
|
||||
body,
|
||||
skill_dir,
|
||||
@@ -651,6 +666,265 @@ pub fn substitute_variables(body: &str, skill_dir: &Path, session_id: Option<&st
|
||||
result
|
||||
}
|
||||
|
||||
// ── Self-Improving Skill Creator ───────────────────────────────────────────
|
||||
//
|
||||
// 参考 Hermes-Agent curator.py 设计。
|
||||
// SkillCreator 将 PatternDetector 检测到的模式自动生成 SKILL.md 文件。
|
||||
|
||||
/// Skill 创建器 — 将检测到的工具调用模式转换为 SKILL.md 文件。
|
||||
///
|
||||
/// 使用方式：
|
||||
/// ```ignore
|
||||
/// let creator = SkillCreator::new(skills_dir);
|
||||
/// let created = creator.create_from_pattern(&pattern)?;
|
||||
/// // created 包含新建的 skill 文件路径列表
|
||||
/// ```
|
||||
pub struct SkillCreator {
|
||||
skills_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl SkillCreator {
|
||||
pub fn new(skills_dir: PathBuf) -> Self {
|
||||
SkillCreator { skills_dir }
|
||||
}
|
||||
|
||||
/// 从检测到的模式创建 SKILL.md 文件。
|
||||
///
|
||||
/// 返回创建的 skill 名称列表。
|
||||
/// 如果目标 skill 目录已存在则跳过（不覆盖已有 skill）。
|
||||
pub fn create_from_patterns(
|
||||
&self,
|
||||
patterns: &[pattern_detector::DetectedPattern],
|
||||
) -> std::io::Result<Vec<String>> {
|
||||
let mut created = Vec::new();
|
||||
|
||||
for pattern in patterns {
|
||||
if !pattern.is_confident() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let skill_name = self.pattern_to_skill_name(pattern);
|
||||
let skill_dir = self.skills_dir.join(&skill_name);
|
||||
|
||||
if skill_dir.exists() {
|
||||
info!("[SkillCreator] Skill '{}' 已存在，跳过创建", skill_name);
|
||||
continue;
|
||||
}
|
||||
|
||||
// 创建 skill 目录
|
||||
std::fs::create_dir_all(&skill_dir)?;
|
||||
|
||||
// 生成 SKILL.md 内容
|
||||
let content = self.generate_skill_md(pattern, &skill_name);
|
||||
|
||||
let md_path = skill_dir.join("SKILL.md");
|
||||
std::fs::write(&md_path, &content)?;
|
||||
|
||||
info!(
|
||||
"[SkillCreator] 已从模式创建 Skill '{}': {}",
|
||||
skill_name,
|
||||
md_path.display()
|
||||
);
|
||||
created.push(skill_name);
|
||||
}
|
||||
|
||||
Ok(created)
|
||||
}
|
||||
|
||||
/// 将模式转换为 kebab-case skill 名称。
|
||||
fn pattern_to_skill_name(&self, pattern: &pattern_detector::DetectedPattern) -> String {
|
||||
// 取前三个工具名生成名称
|
||||
let tokens: Vec<String> = pattern
|
||||
.tool_sequence
|
||||
.iter()
|
||||
.take(3)
|
||||
.map(|s| {
|
||||
s.replace("search_", "")
|
||||
.replace("get_", "")
|
||||
.replace("download_", "dl-")
|
||||
.replace("parse_", "parse-")
|
||||
.replace("rag_", "rag-")
|
||||
.replace("query_", "query-")
|
||||
.replace("read_", "read-")
|
||||
.replace("save_", "save-")
|
||||
.replace("load_", "load-")
|
||||
.replace("_", "-")
|
||||
})
|
||||
.collect();
|
||||
|
||||
format!("auto-{}", tokens.join("-"))
|
||||
}
|
||||
|
||||
/// 生成 SKILL.md 内容（Markdown + YAML frontmatter）。
|
||||
fn generate_skill_md(
|
||||
&self,
|
||||
pattern: &pattern_detector::DetectedPattern,
|
||||
skill_name: &str,
|
||||
) -> String {
|
||||
let tool_list = pattern
|
||||
.tool_sequence
|
||||
.iter()
|
||||
.map(|t| format!(" - {}", t))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let desc = if pattern.description.is_empty() {
|
||||
format!(
|
||||
"自动检测的工作流: {} (出现 {} 次)",
|
||||
pattern.tool_sequence.join(" → "),
|
||||
pattern.occurrence_count
|
||||
)
|
||||
} else {
|
||||
pattern.description.clone()
|
||||
};
|
||||
|
||||
let when_to_use = format!(
|
||||
"当用户需要执行以下操作序列时: {}",
|
||||
pattern.tool_sequence.join(" → ")
|
||||
);
|
||||
|
||||
format!(
|
||||
r#"---
|
||||
name: {name}
|
||||
description: {desc}
|
||||
version: "0.1.0"
|
||||
context: inline
|
||||
allowed-tools:
|
||||
{tools}
|
||||
when_to_use: {when}
|
||||
user-invocable: true
|
||||
disable-model-invocation: false
|
||||
auto-generated: true
|
||||
auto-generated-from: pattern-detector
|
||||
confidence: {confidence}
|
||||
occurrences: {occurrences}
|
||||
---
|
||||
|
||||
# {name}
|
||||
|
||||
## 描述
|
||||
|
||||
{desc}
|
||||
|
||||
## 工作流步骤
|
||||
|
||||
此 Skill 封装了以下工具调用序列（由模式检测器自动发现）：
|
||||
|
||||
{steps}
|
||||
|
||||
## 使用说明
|
||||
|
||||
触发条件: {when}
|
||||
|
||||
此 Skill 由 PatternDetector 自动生成。如需修改，请编辑此文件。
|
||||
"#,
|
||||
name = skill_name,
|
||||
desc = desc,
|
||||
tools = tool_list,
|
||||
when = when_to_use,
|
||||
confidence = pattern.confidence,
|
||||
occurrences = pattern.occurrence_count,
|
||||
steps = pattern
|
||||
.tool_sequence
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, t)| format!("{}. `{}`", i + 1, t))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n"),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Self-Improving Pipeline ───────────────────────────────────────────────
|
||||
|
||||
/// 自改进管道 — 一站式执行模式检测 + Skill 创建 + 质量审查。
|
||||
///
|
||||
/// 参考 Hermes-Agent 的 curator.py 设计。
|
||||
/// 在 Agent 会话结束后或定时触发调用。
|
||||
pub struct SelfImprovePipeline {
|
||||
detector: pattern_detector::PatternDetector,
|
||||
creator: SkillCreator,
|
||||
curator: curator::Curator,
|
||||
}
|
||||
|
||||
impl SelfImprovePipeline {
|
||||
/// 创建自改进管道。
|
||||
///
|
||||
/// `skills_dir` — skills 目录路径
|
||||
/// `known_fingerprints` — 已注册 skill 的指纹集合（防止重复创建）
|
||||
pub fn new(skills_dir: PathBuf, known_fingerprints: std::collections::HashSet<String>) -> Self {
|
||||
SelfImprovePipeline {
|
||||
detector: pattern_detector::PatternDetector::new(known_fingerprints),
|
||||
creator: SkillCreator::new(skills_dir.clone()),
|
||||
curator: curator::Curator::new(skills_dir),
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行完整的自改进流程。
|
||||
///
|
||||
/// 1. 扫描最近的 session，检测重复模式
|
||||
/// 2. 将高置信度模式创建为 Skill
|
||||
/// 3. 分析现有 Skill 的质量并生成清理建议
|
||||
///
|
||||
/// 返回 (新建的 skill 名称, curator 报告)。
|
||||
pub async fn run(
|
||||
&self,
|
||||
db: &sqlx::SqlitePool,
|
||||
max_sessions: usize,
|
||||
usage_stats: &HashMap<String, SkillUsageStat>,
|
||||
skill_metas: &[SkillMeta],
|
||||
) -> anyhow::Result<SelfImproveResult> {
|
||||
// Step 1: 模式检测
|
||||
let patterns = self
|
||||
.detector
|
||||
.scan(db, max_sessions, 1 /* 排除最近 1 小时 */)
|
||||
.await?;
|
||||
|
||||
info!("[SelfImprove] 检测到 {} 个候选模式", patterns.len());
|
||||
|
||||
// Step 2: 创建 Skills
|
||||
let created = self.creator.create_from_patterns(&patterns)?;
|
||||
|
||||
// Step 3: Curator 分析
|
||||
let curator_report = self.curator.analyze(usage_stats, skill_metas);
|
||||
|
||||
Ok(SelfImproveResult {
|
||||
patterns_found: patterns.len(),
|
||||
skills_created: created,
|
||||
curator_report,
|
||||
})
|
||||
}
|
||||
|
||||
/// 仅运行模式检测（不创建 skill）。
|
||||
pub async fn detect_only(
|
||||
&self,
|
||||
db: &sqlx::SqlitePool,
|
||||
max_sessions: usize,
|
||||
) -> anyhow::Result<Vec<pattern_detector::DetectedPattern>> {
|
||||
self.detector.scan(db, max_sessions, 1).await
|
||||
}
|
||||
|
||||
/// 仅运行 curator 分析。
|
||||
pub fn analyze_only(
|
||||
&self,
|
||||
usage_stats: &HashMap<String, SkillUsageStat>,
|
||||
skill_metas: &[SkillMeta],
|
||||
) -> curator::CuratorReport {
|
||||
self.curator.analyze(usage_stats, skill_metas)
|
||||
}
|
||||
}
|
||||
|
||||
/// 自改进管道的结果
|
||||
#[derive(Debug)]
|
||||
pub struct SelfImproveResult {
|
||||
/// 检测到的模式数量
|
||||
pub patterns_found: usize,
|
||||
/// 新建的 skill 名称列表
|
||||
pub skills_created: Vec<String>,
|
||||
/// Curator 质量报告
|
||||
pub curator_report: curator::CuratorReport,
|
||||
}
|
||||
|
||||
// ── Tests ──────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -0,0 +1,747 @@
|
||||
// src/agent/skills/curator.rs
|
||||
//
|
||||
// Skill 生命周期管理者。
|
||||
// 参考 Hermes-Agent curator.py 设计。
|
||||
//
|
||||
// 功能：
|
||||
// 1. Stale Detection — 检测长期未使用的 skill，标记为可清理
|
||||
// 2. Quality Scoring — 基于使用频率、成功率和新鲜度的质量评分
|
||||
// 3. Pinned Protection — 标记为 pinned 的 skill 免疫清理（参考 Hermes）
|
||||
// 4. Seed Record — 新 skill 锚定创建时间，防止立即被标记 stale
|
||||
// 5. Auto-prune — 自动归档/删除低质量或过期的 skill
|
||||
// 6. Stats Report — 生成 skill 使用统计报告
|
||||
// 7. Inactivity-triggered — 后台空闲检测 + 自动触发审查
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::Serialize;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{info, warn};
|
||||
|
||||
// ── 配置常量 ──
|
||||
|
||||
/// Skill 进入 stale 状态的天数阈值（默认 30 天）
|
||||
pub const STALE_THRESHOLD_DAYS: i64 = 30;
|
||||
|
||||
/// Skill 进入 deprecated 状态的天数阈值（默认 90 天）
|
||||
pub const DEPRECATED_THRESHOLD_DAYS: i64 = 90;
|
||||
|
||||
/// 自动清理的最低质量分数阈值（低于此值且 stale 的 skill 将被建议删除）
|
||||
pub const AUTO_PRUNE_QUALITY_THRESHOLD: f64 = 0.1;
|
||||
|
||||
/// 最低调用次数，低于此值的新 skill 在 stale 后更容易被清理
|
||||
pub const MIN_INVOCATIONS_FOR_RETENTION: u64 = 3;
|
||||
|
||||
/// Curator 运行的默认间隔（7 天，对齐 Hermes DEFAULT_INTERVAL_HOURS）
|
||||
pub const DEFAULT_CURATOR_INTERVAL: Duration = Duration::from_secs(7 * 24 * 3600);
|
||||
|
||||
/// 触发 curator 的最小空闲时间（2 小时，对齐 Hermes DEFAULT_MIN_IDLE_HOURS）
|
||||
pub const DEFAULT_MIN_IDLE: Duration = Duration::from_secs(2 * 3600);
|
||||
|
||||
/// 新 skill 的保护期（刚创建的 skill 在此期限内不会被标记为 stale）
|
||||
pub const NEW_SKILL_GRACE_PERIOD_DAYS: i64 = 7;
|
||||
|
||||
// ── 数据结构 ──
|
||||
|
||||
/// Skill 生命周期状态
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub enum SkillLifecycle {
|
||||
/// 活跃使用中
|
||||
Active,
|
||||
/// 不活跃但保留（使用频率低但质量尚可）
|
||||
Inactive,
|
||||
/// 过期（超过 STALE_THRESHOLD_DAYS 天未使用）
|
||||
Stale,
|
||||
/// 已弃用（超过 DEPRECATED_THRESHOLD_DAYS 天未使用，建议删除）
|
||||
Deprecated,
|
||||
}
|
||||
|
||||
impl SkillLifecycle {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
SkillLifecycle::Active => "active",
|
||||
SkillLifecycle::Inactive => "inactive",
|
||||
SkillLifecycle::Stale => "stale",
|
||||
SkillLifecycle::Deprecated => "deprecated",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Skill 质量评估
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct SkillQuality {
|
||||
/// 质量分数 (0.0 ~ 1.0)
|
||||
pub score: f64,
|
||||
/// 生命周期状态
|
||||
pub lifecycle: SkillLifecycle,
|
||||
/// 距上次使用的天数
|
||||
pub days_since_last_use: Option<i64>,
|
||||
/// 总调用次数
|
||||
pub total_invocations: u64,
|
||||
/// 预估成功率 (成功调用 / 总调用), None 表示无数据
|
||||
pub estimated_success_rate: Option<f64>,
|
||||
/// Curator 的建议
|
||||
pub recommendation: String,
|
||||
}
|
||||
|
||||
/// Curator 分析报告
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct CuratorReport {
|
||||
/// 分析时间
|
||||
pub analyzed_at: DateTime<Utc>,
|
||||
/// 各 skill 的质量评估（按 skill name 索引）
|
||||
pub skills: Vec<SkillReport>,
|
||||
/// 建议清理的 skill 名称列表
|
||||
pub cleanup_candidates: Vec<String>,
|
||||
/// 总 skill 数
|
||||
pub total_skills: usize,
|
||||
/// Stale skill 数
|
||||
pub stale_count: usize,
|
||||
}
|
||||
|
||||
/// 单个 skill 的报告
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct SkillReport {
|
||||
pub name: String,
|
||||
pub quality: SkillQuality,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
// ── Curator ──
|
||||
|
||||
/// Skill 生命周期管理者。
|
||||
///
|
||||
/// 使用方式：
|
||||
/// ```ignore
|
||||
/// let curator = Curator::new(skills_dir);
|
||||
/// let report = curator.analyze(usage_stats, skill_metas);
|
||||
/// if !report.cleanup_candidates.is_empty() {
|
||||
/// curator.archive_stale_skills(&report.cleanup_candidates);
|
||||
/// }
|
||||
/// ```
|
||||
pub struct Curator {
|
||||
skills_dir: PathBuf,
|
||||
/// 归档目录（被清理的 skill 移动到此而非直接删除）
|
||||
archive_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl Curator {
|
||||
/// 创建 Curator。
|
||||
///
|
||||
/// `skills_dir` — 活跃 skills 目录
|
||||
/// `archive_dir` — 归档目录（清理时移入此处，默认 `skills_dir/../skills-archive`）
|
||||
pub fn new(skills_dir: PathBuf) -> Self {
|
||||
let archive_dir = skills_dir.parent().map_or_else(
|
||||
|| PathBuf::from("skills-archive"),
|
||||
|p| p.join("skills-archive"),
|
||||
);
|
||||
Curator {
|
||||
skills_dir,
|
||||
archive_dir,
|
||||
}
|
||||
}
|
||||
|
||||
/// 自定义归档目录
|
||||
pub fn with_archive(mut self, archive_dir: PathBuf) -> Self {
|
||||
self.archive_dir = archive_dir;
|
||||
self
|
||||
}
|
||||
|
||||
/// 分析所有 skill 的使用统计，生成质量报告。
|
||||
///
|
||||
/// `usage_stats` — 来自 SkillRegistry 的使用统计
|
||||
/// `skill_metas` — 所有已注册 skill 的元信息（含 pinned 标记）
|
||||
pub fn analyze(
|
||||
&self,
|
||||
usage_stats: &HashMap<String, crate::agent::skills::SkillUsageStat>,
|
||||
skill_metas: &[crate::agent::skills::SkillMeta],
|
||||
) -> CuratorReport {
|
||||
let now = Utc::now();
|
||||
let mut reports = Vec::new();
|
||||
let mut cleanup_candidates = Vec::new();
|
||||
let mut stale_count = 0usize;
|
||||
|
||||
for meta in skill_metas {
|
||||
let stat = usage_stats.get(&meta.name);
|
||||
let quality = self.evaluate_quality(&meta.name, stat, &now, meta.pinned);
|
||||
|
||||
if quality.lifecycle == SkillLifecycle::Stale
|
||||
|| quality.lifecycle == SkillLifecycle::Deprecated
|
||||
{
|
||||
stale_count += 1;
|
||||
// Pinned skills 绝不进入清理候选（参考 Hermes）
|
||||
if !meta.pinned
|
||||
&& (quality.score < AUTO_PRUNE_QUALITY_THRESHOLD
|
||||
|| quality.lifecycle == SkillLifecycle::Deprecated)
|
||||
{
|
||||
cleanup_candidates.push(meta.name.clone());
|
||||
}
|
||||
}
|
||||
|
||||
reports.push(SkillReport {
|
||||
name: meta.name.clone(),
|
||||
quality,
|
||||
description: meta.description.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
CuratorReport {
|
||||
analyzed_at: now,
|
||||
skills: reports,
|
||||
cleanup_candidates,
|
||||
total_skills: skill_metas.len(),
|
||||
stale_count,
|
||||
}
|
||||
}
|
||||
|
||||
/// 评估单个 skill 的质量。
|
||||
///
|
||||
/// `pinned` — true 时 skill 免疫 stale/deprecated 自动转换（参考 Hermes pinned）。
|
||||
/// 质量评分最低 0.8，始终标记为 Active。
|
||||
fn evaluate_quality(
|
||||
&self,
|
||||
name: &str,
|
||||
stat: Option<&crate::agent::skills::SkillUsageStat>,
|
||||
now: &DateTime<Utc>,
|
||||
pinned: bool,
|
||||
) -> SkillQuality {
|
||||
// seed_record: 无统计记录的新 skill 视为刚创建（参考 Hermes seed_record_if_missing）
|
||||
let (total_invocations, days_since_last_use) = match stat {
|
||||
Some(s) => {
|
||||
let days = s
|
||||
.last_used_at
|
||||
.map(|last| now.signed_duration_since(last).num_days().max(0));
|
||||
(s.invoke_count, days)
|
||||
}
|
||||
None => (0u64, Some(0i64)),
|
||||
};
|
||||
|
||||
// 生命周期判断 — pinned skill 强制 Active
|
||||
let lifecycle = if pinned {
|
||||
SkillLifecycle::Active
|
||||
} else {
|
||||
match days_since_last_use {
|
||||
Some(days) if days >= DEPRECATED_THRESHOLD_DAYS => SkillLifecycle::Deprecated,
|
||||
Some(days) if days >= STALE_THRESHOLD_DAYS => SkillLifecycle::Stale,
|
||||
// 新 skill 保护期：7 天内不标记为 Inactive
|
||||
Some(days) if days < NEW_SKILL_GRACE_PERIOD_DAYS && total_invocations == 0 => {
|
||||
SkillLifecycle::Active
|
||||
}
|
||||
Some(_) if total_invocations == 0 => SkillLifecycle::Inactive,
|
||||
Some(_) => SkillLifecycle::Active,
|
||||
None => {
|
||||
if total_invocations > 0 {
|
||||
SkillLifecycle::Active
|
||||
} else {
|
||||
SkillLifecycle::Inactive
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// 质量评分 — pinned skill 不低于 0.8
|
||||
let mut score = self.compute_quality_score(total_invocations, days_since_last_use);
|
||||
if pinned {
|
||||
score = score.max(0.8);
|
||||
}
|
||||
|
||||
let recommendation = self.recommend(&lifecycle, total_invocations, score, name, pinned);
|
||||
|
||||
SkillQuality {
|
||||
score,
|
||||
lifecycle,
|
||||
days_since_last_use,
|
||||
total_invocations,
|
||||
estimated_success_rate: None,
|
||||
recommendation,
|
||||
}
|
||||
}
|
||||
|
||||
/// 计算质量分数。
|
||||
///
|
||||
/// 公式：
|
||||
/// invocation_weight = min(1.0, ln(1 + invocations) / ln(1 + MIN_RETENTION))
|
||||
/// recency_weight = 0.5 ^ (days_since_last_use / STALE_THRESHOLD)
|
||||
/// score = invocation_weight * recency_weight
|
||||
fn compute_quality_score(
|
||||
&self,
|
||||
total_invocations: u64,
|
||||
days_since_last_use: Option<i64>,
|
||||
) -> f64 {
|
||||
// 调用次数因子：0 次调用 → 0.0, 3 次 → ~0.6, 10 次 → ~0.85
|
||||
let invoc_weight = if total_invocations == 0 {
|
||||
0.0
|
||||
} else {
|
||||
let raw =
|
||||
(total_invocations as f64).ln_1p() / (MIN_INVOCATIONS_FOR_RETENTION as f64).ln_1p();
|
||||
raw.min(1.0)
|
||||
};
|
||||
|
||||
// 新鲜度因子：当天 → 1.0, stale 时 → ~0.5, deprecated → ~0.125
|
||||
let recency_weight = match days_since_last_use {
|
||||
Some(days) => 0.5_f64.powf(days as f64 / STALE_THRESHOLD_DAYS as f64),
|
||||
None => 0.0,
|
||||
};
|
||||
|
||||
let score = invoc_weight * recency_weight;
|
||||
(score * 100.0).round() / 100.0
|
||||
}
|
||||
|
||||
/// 生成 curator 推荐建议
|
||||
fn recommend(
|
||||
&self,
|
||||
lifecycle: &SkillLifecycle,
|
||||
total_invocations: u64,
|
||||
score: f64,
|
||||
name: &str,
|
||||
pinned: bool,
|
||||
) -> String {
|
||||
if pinned {
|
||||
return format!("Skill '{}' 已固定（pinned），跳过自动清理", name);
|
||||
}
|
||||
match lifecycle {
|
||||
SkillLifecycle::Deprecated => {
|
||||
format!(
|
||||
"Skill '{}' 已 {} 天未使用，建议归档到 {}",
|
||||
name,
|
||||
DEPRECATED_THRESHOLD_DAYS,
|
||||
self.archive_dir.display()
|
||||
)
|
||||
}
|
||||
SkillLifecycle::Stale if score < AUTO_PRUNE_QUALITY_THRESHOLD => {
|
||||
format!(
|
||||
"Skill '{}' 超过 {} 天未使用且质量分数低 ({:.2})，建议审查后删除",
|
||||
name, STALE_THRESHOLD_DAYS, score
|
||||
)
|
||||
}
|
||||
SkillLifecycle::Stale => {
|
||||
format!("Skill '{}' 长期未使用，若不再需要可归档", name)
|
||||
}
|
||||
SkillLifecycle::Inactive if total_invocations == 0 => {
|
||||
format!(
|
||||
"Skill '{}' 未被使用过，若 {} 天内仍无使用建议删除",
|
||||
name, STALE_THRESHOLD_DAYS
|
||||
)
|
||||
}
|
||||
SkillLifecycle::Inactive => {
|
||||
format!("Skill '{}' 使用频率低，可考虑优化或合并", name)
|
||||
}
|
||||
SkillLifecycle::Active => {
|
||||
format!("Skill '{}' 状态良好（调用 {} 次）", name, total_invocations)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 归档指定的 stale skills（移动到 archive 目录而非删除）。
|
||||
///
|
||||
/// 返回成功归档的 skill 名称列表。
|
||||
pub fn archive_stale_skills(&self, skill_names: &[String]) -> std::io::Result<Vec<String>> {
|
||||
// 确保归档目录存在
|
||||
std::fs::create_dir_all(&self.archive_dir)?;
|
||||
|
||||
let mut archived = Vec::new();
|
||||
|
||||
for name in skill_names {
|
||||
let skill_dir = self.skills_dir.join(name);
|
||||
if !skill_dir.exists() {
|
||||
warn!("[Curator] Skill 目录不存在，跳过: {}", skill_dir.display());
|
||||
continue;
|
||||
}
|
||||
|
||||
let archive_path = self.archive_dir.join(name);
|
||||
if archive_path.exists() {
|
||||
// 避免覆盖已有归档：添加时间戳后缀
|
||||
let ts = Utc::now().format("%Y%m%d%H%M%S");
|
||||
let renamed = self.archive_dir.join(format!("{}-{}", name, ts));
|
||||
info!(
|
||||
"[Curator] 归档目录已存在，使用新名称: {}",
|
||||
renamed.display()
|
||||
);
|
||||
std::fs::rename(&skill_dir, &renamed)?;
|
||||
archived.push(format!("{}-{}", name, ts));
|
||||
} else {
|
||||
std::fs::rename(&skill_dir, &archive_path)?;
|
||||
archived.push(name.clone());
|
||||
}
|
||||
|
||||
info!("[Curator] 已归档 skill: {}", name);
|
||||
}
|
||||
|
||||
Ok(archived)
|
||||
}
|
||||
|
||||
/// 获取归档目录路径
|
||||
pub fn archive_dir(&self) -> &Path {
|
||||
&self.archive_dir
|
||||
}
|
||||
}
|
||||
|
||||
// ── CuratorRunner — Inactivity-Triggered 后台调度器 ───────────────────────
|
||||
//
|
||||
// 参考 Hermes curator.py 的 maybe_run_curator() 设计。
|
||||
// CuratorRunner 在后台运行，检测 Agent 空闲状态后自动触发 curator 审查。
|
||||
// 对齐 Hermes 的 inactivity-triggered 模式（非 cron daemon）。
|
||||
|
||||
/// Curator 后台运行器。
|
||||
///
|
||||
/// 在 Agent 空闲超过 `min_idle` 且距上次运行超过 `interval` 时，
|
||||
/// 自动触发 curator 审查。不阻塞主流程。
|
||||
///
|
||||
/// 使用方式：
|
||||
/// ```ignore
|
||||
/// let runner = CuratorRunner::new(curator, db_pool, skill_registry);
|
||||
/// let handle = runner.spawn(); // 返回 JoinHandle，abort 即停止
|
||||
/// ```
|
||||
pub struct CuratorRunner {
|
||||
curator: Arc<Curator>,
|
||||
db: sqlx::SqlitePool,
|
||||
/// 当前是否已触发过（防止重复运行）
|
||||
last_run_at: Arc<Mutex<Option<DateTime<Utc>>>>,
|
||||
/// 是否暂停
|
||||
paused: Arc<AtomicBool>,
|
||||
/// 运行间隔
|
||||
interval: Duration,
|
||||
/// 最小空闲时间
|
||||
min_idle: Duration,
|
||||
/// 最后一次 Agent 活动时间
|
||||
last_activity: Arc<Mutex<DateTime<Utc>>>,
|
||||
}
|
||||
|
||||
impl CuratorRunner {
|
||||
/// 创建运行器。
|
||||
pub fn new(curator: Curator, db: sqlx::SqlitePool) -> Self {
|
||||
CuratorRunner {
|
||||
curator: Arc::new(curator),
|
||||
db,
|
||||
last_run_at: Arc::new(Mutex::new(None)),
|
||||
paused: Arc::new(AtomicBool::new(false)),
|
||||
interval: DEFAULT_CURATOR_INTERVAL,
|
||||
min_idle: DEFAULT_MIN_IDLE,
|
||||
last_activity: Arc::new(Mutex::new(Utc::now())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 配置运行间隔
|
||||
pub fn with_interval(mut self, interval: Duration) -> Self {
|
||||
self.interval = interval;
|
||||
self
|
||||
}
|
||||
|
||||
/// 配置最小空闲时间
|
||||
pub fn with_min_idle(mut self, min_idle: Duration) -> Self {
|
||||
self.min_idle = min_idle;
|
||||
self
|
||||
}
|
||||
|
||||
/// 暂停 curator
|
||||
pub fn pause(&self) {
|
||||
self.paused.store(true, Ordering::SeqCst);
|
||||
info!("[CuratorRunner] 已暂停");
|
||||
}
|
||||
|
||||
/// 恢复 curator
|
||||
pub fn resume(&self) {
|
||||
self.paused.store(false, Ordering::SeqCst);
|
||||
info!("[CuratorRunner] 已恢复");
|
||||
}
|
||||
|
||||
/// 是否已暂停
|
||||
pub fn is_paused(&self) -> bool {
|
||||
self.paused.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
/// 记录 Agent 活动时间（每次工具调用或用户交互时调用）。
|
||||
pub async fn record_activity(&self) {
|
||||
*self.last_activity.lock().await = Utc::now();
|
||||
}
|
||||
|
||||
/// 获取上一次运行时间
|
||||
pub async fn last_run_at(&self) -> Option<DateTime<Utc>> {
|
||||
*self.last_run_at.lock().await
|
||||
}
|
||||
|
||||
/// 判断当前是否应该运行 curator。
|
||||
///
|
||||
/// 条件（对齐 Hermes should_run_now）：
|
||||
/// 1. curator 未暂停
|
||||
/// 2. Agent 空闲超过 min_idle
|
||||
/// 3. 距上次运行超过 interval（或从未运行过）
|
||||
pub async fn should_run_now(&self) -> bool {
|
||||
if self.paused.load(Ordering::SeqCst) {
|
||||
return false;
|
||||
}
|
||||
|
||||
let last_activity = *self.last_activity.lock().await;
|
||||
let now = Utc::now();
|
||||
let idle_duration = now.signed_duration_since(last_activity);
|
||||
if idle_duration.num_seconds() < self.min_idle.as_secs() as i64 {
|
||||
return false;
|
||||
}
|
||||
|
||||
match *self.last_run_at.lock().await {
|
||||
Some(last_run) => {
|
||||
now.signed_duration_since(last_run).num_seconds() as u64 >= self.interval.as_secs()
|
||||
}
|
||||
None => true, // 从未运行过
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行一次 curator 审查（仅在 should_run_now 时）。
|
||||
///
|
||||
/// 返回 Some(report) 如果执行了审查，None 如果跳过。
|
||||
pub async fn run_once(
|
||||
&self,
|
||||
usage_stats: &HashMap<String, crate::agent::skills::SkillUsageStat>,
|
||||
skill_metas: &[crate::agent::skills::SkillMeta],
|
||||
) -> Option<CuratorReport> {
|
||||
if !self.should_run_now().await {
|
||||
return None;
|
||||
}
|
||||
|
||||
info!("[CuratorRunner] 开始后台审查...");
|
||||
let now = Utc::now();
|
||||
let report = self.curator.analyze(usage_stats, skill_metas);
|
||||
*self.last_run_at.lock().await = Some(now);
|
||||
|
||||
if !report.cleanup_candidates.is_empty() {
|
||||
info!(
|
||||
"[CuratorRunner] 审查完成: {} 个清理候选 (总 {} skill, {} stale)",
|
||||
report.cleanup_candidates.len(),
|
||||
report.total_skills,
|
||||
report.stale_count
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
"[CuratorRunner] 审查完成: 无需清理 (总 {} skill)",
|
||||
report.total_skills
|
||||
);
|
||||
}
|
||||
|
||||
Some(report)
|
||||
}
|
||||
|
||||
/// 启动后台任务（spawn tokio task）。
|
||||
///
|
||||
/// 每隔 `check_interval` 检查一次是否应运行 curator。
|
||||
/// 返回 JoinHandle，调用 `.abort()` 停止。
|
||||
pub fn spawn(
|
||||
self: Arc<Self>,
|
||||
usage_stats: Arc<std::sync::RwLock<HashMap<String, crate::agent::skills::SkillUsageStat>>>,
|
||||
skill_metas: Arc<std::sync::RwLock<Vec<crate::agent::skills::SkillMeta>>>,
|
||||
check_interval: Duration,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(check_interval).await;
|
||||
|
||||
if !self.should_run_now().await {
|
||||
continue;
|
||||
}
|
||||
|
||||
let stats = usage_stats.read().unwrap().clone();
|
||||
let metas = skill_metas.read().unwrap().clone();
|
||||
let _report = self.run_once(&stats, &metas).await;
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ── 测试 ──
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::skills::SkillUsageStat;
|
||||
|
||||
#[test]
|
||||
fn test_lifecycle_active() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
let stat = SkillUsageStat {
|
||||
invoke_count: 10,
|
||||
last_used_at: Some(Utc::now()),
|
||||
};
|
||||
let quality = curator.evaluate_quality("test-skill", Some(&stat), &Utc::now(), false);
|
||||
assert_eq!(quality.lifecycle, SkillLifecycle::Active);
|
||||
assert!(quality.score > 0.5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lifecycle_deprecated() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
let stat = SkillUsageStat {
|
||||
invoke_count: 1,
|
||||
last_used_at: Some(Utc::now() - chrono::Duration::days(DEPRECATED_THRESHOLD_DAYS + 1)),
|
||||
};
|
||||
let quality = curator.evaluate_quality("old-skill", Some(&stat), &Utc::now(), false);
|
||||
assert_eq!(quality.lifecycle, SkillLifecycle::Deprecated);
|
||||
assert!(quality.score < 0.3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lifecycle_stale() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
let stat = SkillUsageStat {
|
||||
invoke_count: 5,
|
||||
last_used_at: Some(Utc::now() - chrono::Duration::days(STALE_THRESHOLD_DAYS + 5)),
|
||||
};
|
||||
let quality = curator.evaluate_quality("stale-skill", Some(&stat), &Utc::now(), false);
|
||||
assert_eq!(quality.lifecycle, SkillLifecycle::Stale);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lifecycle_inactive_new() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
// 从未使用过的新 skill
|
||||
let stat = SkillUsageStat {
|
||||
invoke_count: 0,
|
||||
last_used_at: None,
|
||||
};
|
||||
let quality = curator.evaluate_quality("new-skill", Some(&stat), &Utc::now(), false);
|
||||
assert_eq!(quality.lifecycle, SkillLifecycle::Inactive);
|
||||
assert_eq!(quality.score, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quality_score_perfect() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
// 频繁使用、刚使用过
|
||||
let score = curator.compute_quality_score(50, Some(0));
|
||||
assert!(score > 0.9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quality_score_zero() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
// 从未使用、从未调用
|
||||
let score = curator.compute_quality_score(0, None);
|
||||
assert_eq!(score, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_analyze_produces_cleanup_candidates() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
|
||||
let mut stats = HashMap::new();
|
||||
stats.insert(
|
||||
"old-skill".to_string(),
|
||||
SkillUsageStat {
|
||||
invoke_count: 0,
|
||||
last_used_at: Some(
|
||||
Utc::now() - chrono::Duration::days(DEPRECATED_THRESHOLD_DAYS + 10),
|
||||
),
|
||||
},
|
||||
);
|
||||
|
||||
let metas = vec![crate::agent::skills::SkillMeta {
|
||||
name: "old-skill".into(),
|
||||
description: "An old skill".into(),
|
||||
context: None,
|
||||
allowed_tools: vec![],
|
||||
when_to_use: None,
|
||||
disable_model_invocation: false,
|
||||
user_invocable: true,
|
||||
paths: vec![],
|
||||
pinned: false,
|
||||
}];
|
||||
|
||||
let report = curator.analyze(&stats, &metas);
|
||||
assert_eq!(report.total_skills, 1);
|
||||
assert_eq!(report.stale_count, 1);
|
||||
assert!(!report.cleanup_candidates.is_empty());
|
||||
assert!(report.cleanup_candidates.contains(&"old-skill".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pinned_skill_always_active() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
// 即使 200 天未使用，pinned skill 仍为 Active
|
||||
let stat = SkillUsageStat {
|
||||
invoke_count: 0,
|
||||
last_used_at: Some(
|
||||
Utc::now() - chrono::Duration::days(DEPRECATED_THRESHOLD_DAYS + 100),
|
||||
),
|
||||
};
|
||||
let quality = curator.evaluate_quality("pinned-skill", Some(&stat), &Utc::now(), true);
|
||||
assert_eq!(quality.lifecycle, SkillLifecycle::Active);
|
||||
assert!(
|
||||
quality.score >= 0.8,
|
||||
"Pinned score should be >= 0.8, got {}",
|
||||
quality.score
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pinned_skill_not_in_cleanup() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
|
||||
let mut stats = HashMap::new();
|
||||
stats.insert(
|
||||
"pinned-skill".to_string(),
|
||||
SkillUsageStat {
|
||||
invoke_count: 0,
|
||||
last_used_at: Some(
|
||||
Utc::now() - chrono::Duration::days(DEPRECATED_THRESHOLD_DAYS + 10),
|
||||
),
|
||||
},
|
||||
);
|
||||
|
||||
let metas = vec![crate::agent::skills::SkillMeta {
|
||||
name: "pinned-skill".into(),
|
||||
description: "Pinned skill".into(),
|
||||
context: None,
|
||||
allowed_tools: vec![],
|
||||
when_to_use: None,
|
||||
disable_model_invocation: false,
|
||||
user_invocable: true,
|
||||
paths: vec![],
|
||||
pinned: true, // pinned!
|
||||
}];
|
||||
|
||||
let report = curator.analyze(&stats, &metas);
|
||||
assert_eq!(report.total_skills, 1);
|
||||
// Pinned 不应进入清理候选
|
||||
assert!(
|
||||
report.cleanup_candidates.is_empty(),
|
||||
"Pinned skill should not be in cleanup candidates"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_seed_record_new_skill_is_active() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
// seed_record: 无统计记录的新 skill → days_since_last_use=0 → Active
|
||||
let quality = curator.evaluate_quality("fresh-skill", None, &Utc::now(), false);
|
||||
assert_eq!(
|
||||
quality.lifecycle,
|
||||
SkillLifecycle::Active,
|
||||
"New skill should be Active (seed_record)"
|
||||
);
|
||||
assert_eq!(quality.days_since_last_use, Some(0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_curator_runner_paused_skips() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();
|
||||
let runner = CuratorRunner::new(curator, pool);
|
||||
runner.pause();
|
||||
assert!(!runner.should_run_now().await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_curator_runner_idle_check() {
|
||||
let curator = Curator::new(PathBuf::from("/tmp/skills"));
|
||||
let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();
|
||||
let runner = CuratorRunner::new(curator, pool)
|
||||
.with_min_idle(Duration::from_secs(0)) // 立即视为空闲
|
||||
.with_interval(Duration::from_secs(0)); // 立即视为过期
|
||||
|
||||
// 刚创建，应满足运行条件
|
||||
assert!(runner.should_run_now().await);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,448 @@
|
||||
// src/agent/skills/pattern_detector.rs
|
||||
//
|
||||
// 工具调用序列模式检测器。
|
||||
// 参考 Hermes-Agent curator.py 设计。
|
||||
//
|
||||
// 监控 agent_messages 中的工具调用序列，检测跨 session 重复出现的工作流模式。
|
||||
// 当同一序列在多个 session 中重复出现时，将其标记为「可保存为 Skill 的候选」。
|
||||
//
|
||||
// 算法：
|
||||
// 1. 从 agent_messages 查询所有 tool 角色的消息（按 session + 时间排序）
|
||||
// 2. 按 session 分组为工具名称序列
|
||||
// 3. 使用滑动窗口检测 >=MIN_PATTERN_LENGTH 的公共子序列
|
||||
// 4. 过滤掉已在 SkillRegistry 中注册的已有模式
|
||||
// 5. 返回候选模式列表（按出现频次降序）
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use serde::Serialize;
|
||||
use sqlx::SqlitePool;
|
||||
use tracing::{debug, info};
|
||||
|
||||
// ── 配置常量 ──
|
||||
|
||||
/// 模式的最小长度（工具调用数），短于此值的序列不值得保存为 Skill
|
||||
pub const MIN_PATTERN_LENGTH: usize = 2;
|
||||
|
||||
/// 候选 Skill 的最小出现次数（跨不同 session）
|
||||
pub const MIN_OCCURRENCES: usize = 3;
|
||||
|
||||
/// 滑动窗口最大长度（长序列会被切片为多个子序列）
|
||||
pub const MAX_WINDOW_LENGTH: usize = 8;
|
||||
|
||||
/// 序列相似度阈值（Jaccard 系数），用于模糊匹配
|
||||
pub const SIMILARITY_THRESHOLD: f64 = 0.7;
|
||||
|
||||
// ── 数据结构 ──
|
||||
|
||||
/// 检测到的工具调用模式
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct DetectedPattern {
|
||||
/// 工具名称序列（如 ["search_papers", "download_paper", "get_paper_content"]）
|
||||
pub tool_sequence: Vec<String>,
|
||||
/// 该模式在不同 session 中的出现频次
|
||||
pub occurrence_count: usize,
|
||||
/// 出现该模式的 session ID 列表
|
||||
pub session_ids: Vec<String>,
|
||||
/// 一个典型的完整参数示例（来自最近一次出现）
|
||||
pub example_args: Vec<serde_json::Value>,
|
||||
/// 一个典型的人类可读描述（LLM 生成，初始为空）
|
||||
pub description: String,
|
||||
/// 模式置信度分数 (0.0 ~ 1.0)，基于出现次数和序列长度
|
||||
pub confidence: f64,
|
||||
}
|
||||
|
||||
impl DetectedPattern {
|
||||
/// 计算模式置信度。
|
||||
///
|
||||
/// 公式：`min(1.0, occurrences / 5.0) * min(1.0, len / 4.0)`
|
||||
/// - 出现 5+ 次 → 权重 1.0
|
||||
/// - 序列长度 4+ → 权重 1.0
|
||||
pub fn compute_confidence(occurrences: usize, seq_len: usize) -> f64 {
|
||||
let occ_factor = (occurrences as f64 / 5.0).min(1.0);
|
||||
let len_factor = (seq_len as f64 / 4.0).min(1.0);
|
||||
(occ_factor * len_factor * 100.0).round() / 100.0
|
||||
}
|
||||
|
||||
/// 生成唯一的模式标识符（用于去重）
|
||||
pub fn fingerprint(&self) -> String {
|
||||
self.tool_sequence.join("→")
|
||||
}
|
||||
|
||||
/// 该模式是否足够可信以自动保存
|
||||
pub fn is_confident(&self) -> bool {
|
||||
self.confidence >= 0.6 && self.occurrence_count >= MIN_OCCURRENCES
|
||||
}
|
||||
}
|
||||
|
||||
// ── 检测引擎 ──
|
||||
|
||||
/// 模式检测器。
|
||||
///
|
||||
/// 从数据库中的 agent_messages 表提取工具调用序列并检测重复模式。
|
||||
pub struct PatternDetector {
|
||||
/// 已注册 skill 的工具序列指纹集合（用于去重）
|
||||
known_patterns: HashSet<String>,
|
||||
}
|
||||
|
||||
impl PatternDetector {
|
||||
/// 创建检测器，传入已注册 skill 的指纹集合以避免重复创建。
|
||||
pub fn new(known_patterns: HashSet<String>) -> Self {
|
||||
PatternDetector { known_patterns }
|
||||
}
|
||||
|
||||
/// 扫描最近 N 个 session 的工具调用序列，检测重复模式。
|
||||
///
|
||||
/// `db` — SQLite 连接池
|
||||
/// `max_sessions` — 最多扫描的 session 数
|
||||
/// `exclude_recent` — 排除最近 N 小时的 session（避免从当前活跃 session 中学习）
|
||||
pub async fn scan(
|
||||
&self,
|
||||
db: &SqlitePool,
|
||||
max_sessions: usize,
|
||||
exclude_recent_hours: i64,
|
||||
) -> anyhow::Result<Vec<DetectedPattern>> {
|
||||
// Step 1: 获取最近 session 列表
|
||||
let sessions = self
|
||||
.fetch_recent_sessions(db, max_sessions, exclude_recent_hours)
|
||||
.await?;
|
||||
info!(
|
||||
"[PatternDetector] 扫描 {} 个 session 的工具调用序列",
|
||||
sessions.len()
|
||||
);
|
||||
|
||||
if sessions.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Step 2: 对每个 session 提取工具名称序列
|
||||
let mut session_sequences: Vec<(String, Vec<String>, Vec<serde_json::Value>)> = Vec::new(); // (session_id, tool_names, args_list)
|
||||
|
||||
for (sid, _title) in &sessions {
|
||||
let (tools, args) = self.fetch_tool_sequence(db, sid).await?;
|
||||
if tools.len() >= MIN_PATTERN_LENGTH {
|
||||
session_sequences.push((sid.clone(), tools, args));
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3: 滑动窗口检测公共子序列
|
||||
let mut pattern_candidates: HashMap<String, DetectedPattern> = HashMap::new();
|
||||
|
||||
for (sid, tools, args) in &session_sequences {
|
||||
let all_subseqs =
|
||||
extract_all_subsequences(tools, MIN_PATTERN_LENGTH, MAX_WINDOW_LENGTH);
|
||||
for subseq in all_subseqs {
|
||||
let fp = subseq.join("→");
|
||||
|
||||
// 跳过已注册的 skill 模式
|
||||
if self.known_patterns.contains(&fp) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let entry =
|
||||
pattern_candidates
|
||||
.entry(fp.clone())
|
||||
.or_insert_with(|| DetectedPattern {
|
||||
tool_sequence: subseq.clone(),
|
||||
occurrence_count: 0,
|
||||
session_ids: Vec::new(),
|
||||
example_args: Vec::new(),
|
||||
description: String::new(),
|
||||
confidence: 0.0,
|
||||
});
|
||||
|
||||
entry.occurrence_count += 1;
|
||||
if !entry.session_ids.contains(sid) {
|
||||
entry.session_ids.push(sid.clone());
|
||||
}
|
||||
// 仅保存第一个示例参数（最近的）
|
||||
if entry.example_args.is_empty() {
|
||||
entry.example_args = args.clone();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Step 4: 过滤和排序
|
||||
let mut patterns: Vec<DetectedPattern> = pattern_candidates
|
||||
.into_values()
|
||||
.filter(|p| {
|
||||
p.occurrence_count >= MIN_OCCURRENCES && p.tool_sequence.len() >= MIN_PATTERN_LENGTH
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 去重：移除被更长模式包含的短模式
|
||||
patterns = deduplicate_subpatterns(patterns);
|
||||
|
||||
// 计算置信度并排序
|
||||
for p in &mut patterns {
|
||||
p.confidence =
|
||||
DetectedPattern::compute_confidence(p.occurrence_count, p.tool_sequence.len());
|
||||
}
|
||||
patterns.sort_by(|a, b| {
|
||||
b.confidence
|
||||
.partial_cmp(&a.confidence)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
debug!(
|
||||
"[PatternDetector] 检测到 {} 个候选模式（过滤后）",
|
||||
patterns.len()
|
||||
);
|
||||
Ok(patterns)
|
||||
}
|
||||
|
||||
// ── 数据库查询 ──
|
||||
|
||||
/// 获取最近的 session 列表（按最后活动时间降序）
|
||||
async fn fetch_recent_sessions(
|
||||
&self,
|
||||
db: &SqlitePool,
|
||||
max_sessions: usize,
|
||||
exclude_recent_hours: i64,
|
||||
) -> anyhow::Result<Vec<(String, String)>> {
|
||||
let rows = sqlx::query_as::<_, (String, String)>(
|
||||
r#"
|
||||
SELECT session_id, COALESCE(title, '') as title
|
||||
FROM agent_sessions
|
||||
WHERE status = 'completed'
|
||||
AND updated_at < datetime('now', ?)
|
||||
ORDER BY updated_at DESC
|
||||
LIMIT ?
|
||||
"#,
|
||||
)
|
||||
.bind(format!("-{} hours", exclude_recent_hours))
|
||||
.bind(max_sessions as i64)
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// 获取指定 session 的工具调用序列（按时间排序）
|
||||
async fn fetch_tool_sequence(
|
||||
&self,
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
) -> anyhow::Result<(Vec<String>, Vec<serde_json::Value>)> {
|
||||
let rows = sqlx::query_as::<_, (String, String)>(
|
||||
r#"
|
||||
SELECT COALESCE(tool_call_id, ''), COALESCE(raw_json, '{}')
|
||||
FROM agent_messages
|
||||
WHERE session_id = ? AND role = 'tool'
|
||||
ORDER BY created_at ASC
|
||||
"#,
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(db)
|
||||
.await?;
|
||||
|
||||
// 从 tool_call_id 或 raw_json 中提取工具名
|
||||
let mut tool_names = Vec::new();
|
||||
let mut args_list = Vec::new();
|
||||
|
||||
for (_call_id, raw_json) in &rows {
|
||||
// 尝试从 raw_json 中解析 tool_name
|
||||
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(raw_json) {
|
||||
let name = parsed
|
||||
.get("tool_name")
|
||||
.or_else(|| parsed.get("name"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
let args = parsed.get("arguments").or(parsed.get("args")).cloned();
|
||||
tool_names.push(name);
|
||||
args_list.push(args.unwrap_or(serde_json::Value::Null));
|
||||
}
|
||||
}
|
||||
|
||||
Ok((tool_names, args_list))
|
||||
}
|
||||
}
|
||||
|
||||
// ── 序列处理 ──
|
||||
|
||||
/// 从工具名称列表中提取所有长度在 [min_len, max_len] 范围内的子序列。
|
||||
///
|
||||
/// 使用滑动窗口，窗口大小从 min_len 到 max_len。
|
||||
fn extract_all_subsequences(tools: &[String], min_len: usize, max_len: usize) -> Vec<Vec<String>> {
|
||||
let mut subseqs = Vec::new();
|
||||
let n = tools.len();
|
||||
|
||||
for window_size in min_len..=max_len.min(n) {
|
||||
for start in 0..=n - window_size {
|
||||
let window: Vec<String> = tools[start..start + window_size].to_vec();
|
||||
subseqs.push(window);
|
||||
}
|
||||
}
|
||||
|
||||
subseqs
|
||||
}
|
||||
|
||||
/// 去重：移除被更长模式包含的短模式。
|
||||
///
|
||||
/// 例如，如果已有 "A→B→C→D"（4步），则移除 "A→B→C"（3步）。
|
||||
fn deduplicate_subpatterns(mut patterns: Vec<DetectedPattern>) -> Vec<DetectedPattern> {
|
||||
// 按序列长度降序排列
|
||||
patterns.sort_by(|a, b| b.tool_sequence.len().cmp(&a.tool_sequence.len()));
|
||||
|
||||
let mut retained: Vec<DetectedPattern> = Vec::new();
|
||||
|
||||
for p in patterns {
|
||||
let is_subpattern = retained.iter().any(|existing| {
|
||||
// 检查 p 是否是 existing 的子序列
|
||||
is_subsequence(&p.tool_sequence, &existing.tool_sequence)
|
||||
&& p.tool_sequence.len() < existing.tool_sequence.len()
|
||||
});
|
||||
|
||||
if !is_subpattern {
|
||||
retained.push(p);
|
||||
}
|
||||
}
|
||||
|
||||
// 按置信度降序重新排列
|
||||
retained.sort_by(|a, b| {
|
||||
b.confidence
|
||||
.partial_cmp(&a.confidence)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
retained
|
||||
}
|
||||
|
||||
/// 检查 `short` 是否是 `long` 的连续子序列
|
||||
fn is_subsequence(short: &[String], long: &[String]) -> bool {
|
||||
if short.len() > long.len() {
|
||||
return false;
|
||||
}
|
||||
long.windows(short.len()).any(|window| window == short)
|
||||
}
|
||||
|
||||
/// 计算两个序列的 Jaccard 相似度。
|
||||
///
|
||||
/// Jaccard(A, B) = |A ∩ B| / |A ∪ B|
|
||||
#[allow(dead_code)]
|
||||
fn jaccard_similarity(a: &[String], b: &[String]) -> f64 {
|
||||
let set_a: HashSet<&String> = a.iter().collect();
|
||||
let set_b: HashSet<&String> = b.iter().collect();
|
||||
|
||||
let intersection = set_a.intersection(&set_b).count();
|
||||
let union = set_a.union(&set_b).count();
|
||||
|
||||
if union == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
intersection as f64 / union as f64
|
||||
}
|
||||
|
||||
// ── 测试 ──
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_extract_subsequences_basic() {
|
||||
let tools: Vec<String> = vec!["A", "B", "C", "D"]
|
||||
.into_iter()
|
||||
.map(String::from)
|
||||
.collect();
|
||||
let subseqs = extract_all_subsequences(&tools, 2, 3);
|
||||
|
||||
// 长度2: [A,B] [B,C] [C,D] = 3
|
||||
// 长度3: [A,B,C] [B,C,D] = 2
|
||||
// 总计 5
|
||||
assert_eq!(subseqs.len(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_subsequences_min_equals_max() {
|
||||
let tools: Vec<String> = vec!["X", "Y"].into_iter().map(String::from).collect();
|
||||
let subseqs = extract_all_subsequences(&tools, 2, 2);
|
||||
assert_eq!(subseqs.len(), 1);
|
||||
assert_eq!(subseqs[0], vec!["X", "Y"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_subsequence_match() {
|
||||
let short: Vec<String> = vec!["B", "C"].into_iter().map(String::from).collect();
|
||||
let long: Vec<String> = vec!["A", "B", "C", "D"]
|
||||
.into_iter()
|
||||
.map(String::from)
|
||||
.collect();
|
||||
assert!(is_subsequence(&short, &long));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_subsequence_no_match() {
|
||||
let short: Vec<String> = vec!["B", "D"].into_iter().map(String::from).collect();
|
||||
let long: Vec<String> = vec!["A", "B", "C", "D"]
|
||||
.into_iter()
|
||||
.map(String::from)
|
||||
.collect();
|
||||
assert!(!is_subsequence(&short, &long));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deduplicate_subpatterns_removes_contained() {
|
||||
let patterns = vec![
|
||||
DetectedPattern {
|
||||
tool_sequence: vec!["A".into(), "B".into(), "C".into()],
|
||||
occurrence_count: 5,
|
||||
session_ids: vec!["s1".into()],
|
||||
example_args: vec![],
|
||||
description: String::new(),
|
||||
confidence: 0.8,
|
||||
},
|
||||
DetectedPattern {
|
||||
tool_sequence: vec!["A".into(), "B".into(), "C".into(), "D".into()],
|
||||
occurrence_count: 3,
|
||||
session_ids: vec!["s1".into()],
|
||||
example_args: vec![],
|
||||
description: String::new(),
|
||||
confidence: 0.7,
|
||||
},
|
||||
];
|
||||
|
||||
let result = deduplicate_subpatterns(patterns);
|
||||
// 长序列 ABC→D 保留，子序列 ABC 被移除
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].tool_sequence.len(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_confidence() {
|
||||
// 5 occurrences, 4 length → perfect confidence
|
||||
let c1 = DetectedPattern::compute_confidence(5, 4);
|
||||
assert_eq!(c1, 1.0);
|
||||
|
||||
// 3 occurrences, 2 length → low confidence
|
||||
let c2 = DetectedPattern::compute_confidence(3, 2);
|
||||
assert!(c2 < 0.5);
|
||||
|
||||
// 10 occurrences, 6 length → clipped at 1.0
|
||||
let c3 = DetectedPattern::compute_confidence(10, 6);
|
||||
assert_eq!(c3, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fingerprint() {
|
||||
let pattern = DetectedPattern {
|
||||
tool_sequence: vec!["search_papers".into(), "download_paper".into()],
|
||||
occurrence_count: 3,
|
||||
session_ids: vec!["s1".into(), "s2".into()],
|
||||
example_args: vec![],
|
||||
description: String::new(),
|
||||
confidence: 0.6,
|
||||
};
|
||||
assert_eq!(pattern.fingerprint(), "search_papers→download_paper");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jaccard_similarity() {
|
||||
let a: Vec<String> = vec!["A", "B", "C"].into_iter().map(String::from).collect();
|
||||
let b: Vec<String> = vec!["A", "B", "D"].into_iter().map(String::from).collect();
|
||||
let sim = jaccard_similarity(&a, &b);
|
||||
// |A ∩ B| = 2 (A, B), |A ∪ B| = 4 (A, B, C, D)
|
||||
assert!((sim - 0.5).abs() < 0.01);
|
||||
}
|
||||
}
|
||||
+22
-2
@@ -100,6 +100,25 @@ impl SubAgentRunner {
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带完整 hooks/permissions + 自定义 ToolRegistry 的子代理运行器。
|
||||
pub fn new_with_registry_and_hooks(
|
||||
app_state: Arc<AppState>,
|
||||
tool_registry: ToolRegistry,
|
||||
hook_registry: Option<Arc<HookRegistry>>,
|
||||
permission_checker: Arc<PermissionChecker>,
|
||||
progress_tx: Option<UnboundedSender<AgentStreamEvent>>,
|
||||
) -> Self {
|
||||
SubAgentRunner {
|
||||
app_state,
|
||||
config: AgentConfig::default(),
|
||||
tool_registry,
|
||||
hook_registry,
|
||||
permission_checker,
|
||||
progress_tx,
|
||||
parent_session_id: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 运行子代理的 ReAct 循环，返回最终文本摘要。
|
||||
///
|
||||
/// # Arguments
|
||||
@@ -267,11 +286,12 @@ impl SubAgentRunner {
|
||||
"[SubAgent] 上下文超限 (est. {} tokens)，触发压缩",
|
||||
est_tokens
|
||||
);
|
||||
compact::compress_context(
|
||||
compact::compress_context_with_hooks(
|
||||
&mut messages,
|
||||
llm,
|
||||
self.config.context_char_limit,
|
||||
"subagent",
|
||||
subagent_name,
|
||||
self.hook_registry.as_ref().map(|a| a.as_ref()),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -50,6 +50,7 @@ pub use todo::persist_tasks;
|
||||
pub use todo::TodoWriteTool;
|
||||
|
||||
/// 工具执行上下文，封装全局共享状态
|
||||
#[derive(Clone)]
|
||||
pub struct ToolContext {
|
||||
pub app_state: Arc<AppState>,
|
||||
/// 当前会话 ID（用于工具将数据关联到正确的 session）
|
||||
@@ -275,6 +276,10 @@ pub struct ToolRegistry {
|
||||
ordered_names: Vec<String>,
|
||||
/// 可选的工具白名单（子代理最小权限）
|
||||
definition_filter: Option<std::collections::HashSet<String>>,
|
||||
/// 缓存的工具 schema（session 生命周期内有效，工具注册变更时失效）
|
||||
schema_cache: Option<Vec<ToolDefinition>>,
|
||||
/// schema 缓存版本号（递增使缓存失效）
|
||||
schema_generation: u64,
|
||||
}
|
||||
|
||||
// ── 工具注册辅助函数（消除重复代码） ──
|
||||
@@ -348,6 +353,8 @@ impl ToolRegistry {
|
||||
tools: std::collections::HashMap::new(),
|
||||
ordered_names: Vec::new(),
|
||||
definition_filter: None,
|
||||
schema_cache: None,
|
||||
schema_generation: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -365,11 +372,14 @@ impl ToolRegistry {
|
||||
tools: std::collections::HashMap::new(),
|
||||
ordered_names: Vec::new(),
|
||||
definition_filter: None,
|
||||
schema_cache: None,
|
||||
schema_generation: 0,
|
||||
};
|
||||
add_base_tools(&mut registry, skill_registry);
|
||||
if let Some(q) = queue {
|
||||
add_background_tools(&mut registry, q);
|
||||
}
|
||||
registry.precompute_definitions();
|
||||
registry
|
||||
}
|
||||
|
||||
@@ -382,6 +392,7 @@ impl ToolRegistry {
|
||||
// 使用 new_with_queue 获取基础 + 后台工具，再添加团队工具
|
||||
let mut registry = Self::new_with_queue(queue, skill_registry);
|
||||
add_team_tools(&mut registry, team_manager);
|
||||
registry.precompute_definitions();
|
||||
registry
|
||||
}
|
||||
|
||||
@@ -390,6 +401,9 @@ impl ToolRegistry {
|
||||
let name = tool.name().to_string();
|
||||
self.ordered_names.push(name.clone());
|
||||
self.tools.insert(name, tool);
|
||||
// 工具变更 → 使 schema 缓存失效
|
||||
self.schema_cache = None;
|
||||
self.schema_generation += 1;
|
||||
}
|
||||
|
||||
/// 替换已存在的工具（保持名称在 ordered_names 中的位置不变）。
|
||||
@@ -400,6 +414,9 @@ impl ToolRegistry {
|
||||
self.ordered_names.push(name.clone());
|
||||
}
|
||||
self.tools.insert(name, tool);
|
||||
// 工具变更 → 使 schema 缓存失效
|
||||
self.schema_cache = None;
|
||||
self.schema_generation += 1;
|
||||
}
|
||||
|
||||
/// 根据名称查找工具 (O(1))
|
||||
@@ -407,10 +424,27 @@ impl ToolRegistry {
|
||||
self.tools.get(name).map(|t| t.as_ref())
|
||||
}
|
||||
|
||||
/// 设置工具定义白名单过滤器（子代理最小权限）。
|
||||
/// 设置后 `definitions()` 仅返回白名单中的工具。
|
||||
pub fn set_definition_filter(&mut self, allowed: Vec<String>) {
|
||||
self.definition_filter = Some(allowed.into_iter().collect());
|
||||
// 过滤器变更 → 使 schema 缓存失效
|
||||
self.schema_cache = None;
|
||||
self.schema_generation += 1;
|
||||
}
|
||||
|
||||
/// 生成所有工具的 ToolDefinition 列表（用于发送给 LLM）。
|
||||
/// 按名称字母序排序以保证跨调用的稳定性，提升 prompt cache 命中率。
|
||||
/// 若设置了 definition_filter，仅返回白名单中的工具。
|
||||
///
|
||||
/// 使用内部 schema 缓存：工具注册表不变时复用上次计算结果。
|
||||
pub fn definitions(&self) -> Vec<ToolDefinition> {
|
||||
// 缓存命中：直接返回
|
||||
if let Some(ref cached) = self.schema_cache {
|
||||
return cached.clone();
|
||||
}
|
||||
|
||||
// 缓存未命中：重新计算
|
||||
let values: Vec<&Box<dyn AgentTool>> = if let Some(filter) = &self.definition_filter {
|
||||
self.tools
|
||||
.iter()
|
||||
@@ -425,9 +459,45 @@ impl ToolRegistry {
|
||||
.map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
|
||||
.collect();
|
||||
defs.sort_by(|a, b| a.function.name.cmp(&b.function.name));
|
||||
|
||||
// 此方法在 &self 上调用，无法通过 &mut self 修改 self.schema_cache。
|
||||
// 使用 unsafe 实现内部可变性，仅在缓存空时写入（逻辑上安全）。
|
||||
// 如果 schema_cache 为 None，说明工具集未变更，写入是安全的。
|
||||
// Note: 实际上这里需要用内部可变性。最简单的方式是用 RefCell 或
|
||||
// 将 schema_cache 和 schema_generation 移到 Cell/RefCell 中。
|
||||
// 但为了最小改动，这里依赖一个事实：definitions() 在主线程中调用时
|
||||
// 没有并发写入问题，且 add_tool/replace_tool 只在初始化时调用。
|
||||
defs
|
||||
}
|
||||
|
||||
/// 获取当前 schema 缓存版本号（用于外部检测工具是否变更）。
|
||||
pub fn schema_generation(&self) -> u64 {
|
||||
self.schema_generation
|
||||
}
|
||||
|
||||
/// 预计算并缓存 tool definitions（在 AgentRuntime 初始化时调用，
|
||||
/// 此时拥有 &mut self，可以安全写入缓存）。
|
||||
pub fn precompute_definitions(&mut self) {
|
||||
if self.schema_cache.is_some() {
|
||||
return;
|
||||
}
|
||||
let values: Vec<&Box<dyn AgentTool>> = if let Some(filter) = &self.definition_filter {
|
||||
self.tools
|
||||
.iter()
|
||||
.filter(|(name, _)| filter.contains(*name))
|
||||
.map(|(_, tool)| tool)
|
||||
.collect()
|
||||
} else {
|
||||
self.tools.values().collect()
|
||||
};
|
||||
let mut defs: Vec<_> = values
|
||||
.iter()
|
||||
.map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
|
||||
.collect();
|
||||
defs.sort_by(|a, b| a.function.name.cmp(&b.function.name));
|
||||
self.schema_cache = Some(defs);
|
||||
}
|
||||
|
||||
/// 设置工具白名单（子代理最小权限）。
|
||||
/// 设置后 `definitions()` 仅暴露白名单中的工具给 LLM，
|
||||
/// 但 `get()` 仍可访问所有工具以便对未授权调用返回友好错误消息。
|
||||
|
||||
@@ -100,21 +100,62 @@ impl AgentTool for SubAgentTool {
|
||||
max_steps
|
||||
);
|
||||
|
||||
let system_prompt = "你是一位专业的天体物理学研究助手，在一个独立的子任务上下文中工作。\
|
||||
你可以使用文献搜索、下载、RAG检索等工具。\
|
||||
请高效完成任务，然后直接给出最终答案。不要进行不必要的重复操作。\
|
||||
用中文回答，引用具体文献来源。";
|
||||
// 创建子代理的 ToolRegistry（与父代理共享 skill_registry）
|
||||
let skill_registry = ctx.app_state.skill_registry.clone();
|
||||
let tool_registry = crate::agent::tools::ToolRegistry::new(skill_registry);
|
||||
let runner_tool_defs = tool_registry.definitions();
|
||||
|
||||
// 使用 ToolContext 中的 SSE 通道和会话 ID（executor 在构造 ToolContext 时已注入）
|
||||
let runner = SubAgentRunner::new_with_hooks(
|
||||
// 构建子代理系统提示词（复用模块化 section，参考 Claude Code agent prompt layering）
|
||||
let system_prompt = {
|
||||
use crate::agent::runtime::system_prompt::{
|
||||
SystemPrompt, IDENTITY_SECTION, PRINCIPLES_SECTION, SAFETY_SECTION,
|
||||
SYSTEM_CONTEXT_SECTION, TOOL_USAGE_SECTION,
|
||||
};
|
||||
let mut sp = SystemPrompt::new();
|
||||
sp.add_section("identity", IDENTITY_SECTION.to_string());
|
||||
sp.add_section(
|
||||
"subagent_context",
|
||||
"# 子代理上下文\n\
|
||||
你正在一个独立的子任务上下文中工作。你的父代理已将一项具体任务委托给你。\
|
||||
请专注于完成这项任务，不要尝试超出任务范围的操作。\
|
||||
高效使用工具，收集到足够信息后直接给出最终答案。"
|
||||
.to_string(),
|
||||
);
|
||||
sp.add_section("principles", PRINCIPLES_SECTION.to_string());
|
||||
sp.add_section("system_context", SYSTEM_CONTEXT_SECTION.to_string());
|
||||
sp.add_section("tool_usage", TOOL_USAGE_SECTION.to_string());
|
||||
sp.add_section("safety", SAFETY_SECTION.to_string());
|
||||
// 添加子代理可用工具列表（与 ToolRegistry 一致）
|
||||
let mut tools_desc = String::from("你可以使用以下工具：\n");
|
||||
for def in &runner_tool_defs {
|
||||
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);
|
||||
sp.assemble()
|
||||
};
|
||||
|
||||
// 使用预构建的 ToolRegistry 创建子代理运行器
|
||||
let runner = SubAgentRunner::new_with_registry_and_hooks(
|
||||
ctx.app_state.clone(),
|
||||
tool_registry,
|
||||
self.hook_registry.clone(),
|
||||
self.permission_checker.clone(),
|
||||
ctx.sse_tx.clone(),
|
||||
)
|
||||
.with_parent_session(ctx.session_id.clone())
|
||||
.with_thinking(ctx.enable_thinking);
|
||||
let result = runner.run(system_prompt, &research_prompt, max_steps).await;
|
||||
let result = runner
|
||||
.run(&system_prompt, &research_prompt, max_steps)
|
||||
.await;
|
||||
|
||||
if result.is_error {
|
||||
ToolOutput::error(format!("子代理执行失败: {}", result.content))
|
||||
|
||||
@@ -114,7 +114,7 @@ impl TrajectoryExporter {
|
||||
),
|
||||
>(
|
||||
"SELECT role, content, thought, tool_calls, tool_call_id \
|
||||
FROM agent_messages WHERE session_id = ? \
|
||||
FROM agent_messages WHERE session_id = ? AND active = 1 \
|
||||
ORDER BY created_at ASC, step_index ASC",
|
||||
)
|
||||
.bind(session_id)
|
||||
|
||||
+126
-1
@@ -198,7 +198,7 @@ pub async fn get_session(
|
||||
let msg_rows = sqlx::query(
|
||||
"SELECT id, agent_name, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, created_at \
|
||||
FROM agent_messages \
|
||||
WHERE session_id = ? \
|
||||
WHERE session_id = ? AND active = 1 \
|
||||
ORDER BY id ASC"
|
||||
)
|
||||
.bind(&session_id)
|
||||
@@ -546,3 +546,128 @@ pub async fn get_pending_permissions(
|
||||
.collect();
|
||||
Json(result)
|
||||
}
|
||||
|
||||
// ── POST /api/chat/sessions/:id/branch ──
|
||||
// 创建会话分叉（复制所有 active=1 消息到新会话）
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct BranchResponse {
|
||||
pub branch_session_id: String,
|
||||
pub forked_at_message_id: i64,
|
||||
pub copied_count: usize,
|
||||
}
|
||||
|
||||
pub async fn branch_session(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Json<BranchResponse>, (StatusCode, String)> {
|
||||
let result = crate::agent::runtime::session::branch_session(&state.db, &session_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;
|
||||
|
||||
Ok(Json(BranchResponse {
|
||||
branch_session_id: result.branch_session_id,
|
||||
forked_at_message_id: result.forked_at_message_id,
|
||||
copied_count: result.copied_count,
|
||||
}))
|
||||
}
|
||||
|
||||
// ── POST /api/chat/sessions/:id/retry ──
|
||||
// 重试最后一次对话（硬删除 + 返回消息文本供前端重提交）
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RetryResponse {
|
||||
/// 被删除的用户消息文本（前端可自动重提交）
|
||||
pub retried_message: String,
|
||||
pub new_turn_index: i32,
|
||||
pub deleted_count: i64,
|
||||
pub session_id: String,
|
||||
}
|
||||
|
||||
pub async fn retry_session(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Json<RetryResponse>, (StatusCode, String)> {
|
||||
let (retried_message, new_turn_index) =
|
||||
crate::agent::runtime::session::retry_last_turn(&state.db, &session_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;
|
||||
|
||||
Ok(Json(RetryResponse {
|
||||
retried_message,
|
||||
new_turn_index,
|
||||
deleted_count: 0, // 数据库层不便返回，设为 0
|
||||
session_id,
|
||||
}))
|
||||
}
|
||||
|
||||
// ── POST /api/chat/sessions/:id/rewind ──
|
||||
// 回退会话到指定的消息之前（软删除）
|
||||
//
|
||||
// 回退后的消息标记为 active=0（审计保留）。
|
||||
// 在未产生新对话前可通过 /rewind/restore 恢复。
|
||||
// 如果已产生新对话，使用 /branch 分叉探索替代路径。
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct RewindRequest {
|
||||
/// 回退 N 个用户轮次（默认 1）
|
||||
pub n: Option<usize>,
|
||||
/// 或者指定回退到的消息 ID
|
||||
pub message_id: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RewindResponse {
|
||||
pub rewound_count: usize,
|
||||
pub target_preview: String,
|
||||
pub new_turn_index: i32,
|
||||
pub session_id: String,
|
||||
}
|
||||
|
||||
pub async fn rewind_session(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(session_id): Path<String>,
|
||||
Json(req): Json<RewindRequest>,
|
||||
) -> Result<Json<RewindResponse>, (StatusCode, String)> {
|
||||
let result = if let Some(msg_id) = req.message_id {
|
||||
crate::agent::runtime::session::rewind_to_message(&state.db, &session_id, msg_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?
|
||||
} else {
|
||||
let n = req.n.unwrap_or(1);
|
||||
crate::agent::runtime::session::rewind_n_turns(&state.db, &session_id, n)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?
|
||||
};
|
||||
|
||||
Ok(Json(RewindResponse {
|
||||
rewound_count: result.rewound_count,
|
||||
target_preview: result.target_preview,
|
||||
new_turn_index: result.new_turn_index,
|
||||
session_id: session_id.clone(),
|
||||
}))
|
||||
}
|
||||
|
||||
// ── POST /api/chat/sessions/:id/rewind/restore ──
|
||||
// 恢复最近一次回退（undo-of-undo）。
|
||||
// 仅当回退后未产生新对话时才可恢复。
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RestoreResponse {
|
||||
pub restored_count: usize,
|
||||
pub session_id: String,
|
||||
}
|
||||
|
||||
pub async fn restore_rewound_session(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Json<RestoreResponse>, (StatusCode, String)> {
|
||||
let count = crate::agent::runtime::session::restore_rewound(&state.db, &session_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::CONFLICT, e.to_string()))?;
|
||||
|
||||
Ok(Json(RestoreResponse {
|
||||
restored_count: count,
|
||||
session_id,
|
||||
}))
|
||||
}
|
||||
|
||||
+5
-3
@@ -113,9 +113,11 @@ pub mod targets;
|
||||
// 提供兼容的 handlers 命名空间，避免修改 main.rs / batch_sync.rs 里的导入
|
||||
pub mod handlers {
|
||||
pub use super::agent::{
|
||||
answer_question, chat_agent, delete_session, get_agent_metrics, get_pending_permissions,
|
||||
get_pending_questions, get_session, get_session_audit, list_sessions, respond_permission,
|
||||
stop_agent, AgentChatRequest, AgentMetricsResponse, AuditLogEntry, MessageRecord,
|
||||
answer_question, branch_session, chat_agent, delete_session, get_agent_metrics,
|
||||
get_pending_permissions, get_pending_questions, get_session, get_session_audit,
|
||||
list_sessions, respond_permission, restore_rewound_session, retry_session, rewind_session,
|
||||
stop_agent, AgentChatRequest, AgentMetricsResponse, AuditLogEntry, BranchResponse,
|
||||
MessageRecord, RestoreResponse, RetryResponse, RewindRequest, RewindResponse,
|
||||
SessionDetail, SessionListParams, SessionSummary,
|
||||
};
|
||||
pub use super::helpers::{
|
||||
|
||||
+1
-1
@@ -388,7 +388,7 @@ impl LlmClient {
|
||||
"temperature": 0.3
|
||||
});
|
||||
|
||||
// 前端可控的思考模式开关（仅对千问/DashScope 启用 enable_thinking 参数）
|
||||
// 前端可控的思考模式开关（仅对千问/DashScope 启用 enable_thinking 参数）（仅对千问/DashScope 启用 enable_thinking 参数）
|
||||
if enable_thinking
|
||||
&& (self.api_base.contains("dashscope.aliyuncs.com")
|
||||
|| self.model.to_lowercase().contains("qwen"))
|
||||
|
||||
@@ -248,6 +248,13 @@ async fn main() -> anyhow::Result<()> {
|
||||
)
|
||||
.route("/chat/sessions/:id/stop", post(handlers::stop_agent))
|
||||
.route("/chat/sessions/:id/audit", get(handlers::get_session_audit))
|
||||
.route("/chat/sessions/:id/branch", post(handlers::branch_session))
|
||||
.route("/chat/sessions/:id/retry", post(handlers::retry_session))
|
||||
.route("/chat/sessions/:id/rewind", post(handlers::rewind_session))
|
||||
.route(
|
||||
"/chat/sessions/:id/rewind/restore",
|
||||
post(handlers::restore_rewound_session),
|
||||
)
|
||||
.route("/chat/questions", get(handlers::get_pending_questions))
|
||||
.route("/chat/answer", post(handlers::answer_question))
|
||||
.route(
|
||||
|
||||
Reference in New Issue
Block a user