diff --git a/Cargo.lock b/Cargo.lock index 34de065..c41c7bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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" diff --git a/Cargo.toml b/Cargo.toml index bd74e80..075617a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 = [] diff --git a/dashboard/src/App.tsx b/dashboard/src/App.tsx index cd9ba6b..2cf9506 100644 --- a/dashboard/src/App.tsx +++ b/dashboard/src/App.tsx @@ -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() { >
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" >
-

- +

+ {dialog.title}

-

{dialog.message}

+

{dialog.message}

@@ -755,11 +755,11 @@ export default function App() { >
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" >
-

- +

+ 文献尚未入库

-

- 文献 {uncachedBibcode} 尚未收录在本地数据库中。 +

+ 文献 {uncachedBibcode} 尚未收录在本地数据库中。

-

+

您可以选择在线拉取该文献元数据并入库，或是直接跳转至 NASA ADS 平台查看其原始页面。

@@ -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 @@ -818,13 +818,13 @@ export default function App() { >
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" > {/* 标题 & 关闭 */}
- 文献详情元数据 -

+ 文献详情元数据 +

{getDoctypeBadge(detailPaper.doctype)} {detailPaper.title}

@@ -844,48 +844,48 @@ export default function App() {
{/* 作者 */}
- 作者列表 -

{detailPaper.authors.join(', ')}

+ 作者列表 +

{detailPaper.authors.join(', ')}

{/* 期刊 & 年份 */}
- 发表期刊 + 发表期刊 {detailPaper.pub_journal || '未标注'}
- 发表年份 - {detailPaper.year} + 发表年份 + {detailPaper.year}
{/* 摘要 */}
- 摘要 -

+ 摘要 +

{detailPaper.abstract_text || '该文献暂无摘要数据。'}

{/* 关键字 */}
- 关键词 + 关键词 {detailPaper.keywords && detailPaper.keywords.length > 0 ? (
{detailPaper.keywords.map(kw => ( - + {kw} ))}
) : ( -
+
暂无关键词
)} @@ -893,48 +893,48 @@ export default function App() { {/* 标识符 */}
-
- BIBCODE - + -
- DOI - + -
- ARXIV ID - +
+ ARXIV ID + {detailPaper.arxiv_id ? ( 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} @@ -946,14 +946,14 @@ export default function App() { {/* 手动上传文件（应对防爬阻断） */}
- 手动离线上传文献 + 手动离线上传文献 防爬/人机验证备用
-

+

若自动下载受阻，可在浏览器中打开上方链接，手动保存 PDF 或 HTML 后在此处上传覆盖。

-
+
)} - + {uploadingBibcode === detailPaper.bibcode ? '上传中...' : '上传 PDF 文献'} - 支持 .pdf 格式 + 支持 .pdf 格式
-
+
)} - + {uploadingBibcode === detailPaper.bibcode ? '上传中...' : '上传 HTML 文献'} - 支持 .html 格式 + 支持 .html 格式
@@ -1004,7 +1004,7 @@ export default function App() { @@ -1012,9 +1012,9 @@ export default function App() { )}
@@ -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" > 查看引用图谱 @@ -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] ? '重下中...' : '重新下载'} @@ -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" > 引用图谱 diff --git a/dashboard/src/features/agent/ResearchAgentPanel.tsx b/dashboard/src/features/agent/ResearchAgentPanel.tsx index b896ef7..cfacb19 100644 --- a/dashboard/src/features/agent/ResearchAgentPanel.tsx +++ b/dashboard/src/features/agent/ResearchAgentPanel.tsx @@ -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( + `/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( + `/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( + `/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) => (
{/* 用户提问 */} -
+
我 -
- {turn.question} +
+ {turn.questionMessageId && ( + + )} +
+ {turn.question} +
@@ -1151,7 +1334,7 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane {/* Error block */} {activeTurn.error && ( -
+
查询发生错误
@@ -1161,6 +1344,40 @@ export function ResearchAgentPanel({ showConfirm, showAlert }: ResearchAgentPane )}
)} + + {/* 会话操作小图标栏 */} + {currentSessionId && turns.length > 0 && !streaming && ( +
+ + + + +
+ )}
)}
@@ -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; diff --git a/dashboard/src/features/reader/ReaderPanel.tsx b/dashboard/src/features/reader/ReaderPanel.tsx index 1842df4..597c3d1 100644 --- a/dashboard/src/features/reader/ReaderPanel.tsx +++ b/dashboard/src/features/reader/ReaderPanel.tsx @@ -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 ? : } - {vectorizing ? '向量化入库中...' : '向量化入库'} + {vectorizing ? '正在进行知识入库...' : '知识入库'} )} {selectedPaper.has_markdown && selectedPaper.has_vector && ( )} @@ -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" >
- {hoveredTarget.target_name} + {hoveredTarget.target_name}
-
+
- RA (J2000): -
{hoveredTarget.ra || '未知'}
+ RA (J2000): +
{hoveredTarget.ra || '未知'}
- Dec (J2000): -
{hoveredTarget.dec || '未知'}
+ Dec (J2000): +
{hoveredTarget.dec || '未知'}
- 光谱型: -
{hoveredTarget.spectral_type || '未知'}
+ 光谱型: +
{hoveredTarget.spectral_type || '未知'}
- 视星等 (V): -
+ 视星等 (V): +
{hoveredTarget.v_magnitude !== null && hoveredTarget.v_magnitude !== undefined ? `${hoveredTarget.v_magnitude.toFixed(2)}` : '未知'}
- 视差 / 估算距离: -
+ 视差 / 估算距离: +
{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({
{hoveredTarget.aliases && hoveredTarget.aliases.length > 0 && (
- 常用别名: -
+ 常用别名: +
{hoveredTarget.aliases.slice(0, 8).join(', ')} {hoveredTarget.aliases.length > 8 && ' ...'}
diff --git a/dashboard/src/index.css b/dashboard/src/index.css index a0a6525..7d02d44 100644 --- a/dashboard/src/index.css +++ b/dashboard/src/index.css @@ -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; } - - diff --git a/dashboard/src/types.ts b/dashboard/src/types.ts index bb3e400..e3c7403 100644 --- a/dashboard/src/types.ts +++ b/dashboard/src/types.ts @@ -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; +} diff --git a/docs/architecture/agent/claude-code-reference-analysis.md b/docs/architecture/agent/claude-code-reference-analysis.md new file mode 100644 index 0000000..0539821 --- /dev/null +++ b/docs/architecture/agent/claude-code-reference-analysis.md @@ -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` — ~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 { + 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, +} + +#[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 " + 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, +} + +#[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` | 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, + /// 自定义代理禁用（不能 spawn 子代理 + 编辑文件的代理） + pub custom_agent_disallowed: Vec, + /// 异步代理白名单（只读子集） + pub async_agent_allowed: Vec, +} +``` + +--- + +## 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 { + // 用专门的小 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 结果以 返回 │ +│ Coordinator 做 Synthesis → 下一轮 Workers │ +└─────────────────────────────────────────────────┘ +``` + +### 9.2 关键设计点 + +1. **Coordinator 只看到 4 个工具**——它不能直接读文件或执行 bash，只能编排 Workers +2. **Workers 全异步**——Coordinator 不等待，结果以 `` 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` — ~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） diff --git a/docs/architecture/agent/hooks.md b/docs/architecture/agent/hooks.md index 40d9fd0..d166101 100644 --- a/docs/architecture/agent/hooks.md +++ b/docs/architecture/agent/hooks.md @@ -1,274 +1,921 @@ -# Hooks 生命周期系统 (`hooks.rs`) +# Agent Hooks — 生命周期事件系统 -参考 Claude Code hooks 协议，提供 **9 种生命周期事件回调**，基于 **观察者模式 + 责任链模式** 实现。 +参考 Claude Code hooks 协议，提供 **12 种生命周期事件回调**，基于 **观察者模式 + 责任链模式** 实现。核心目标：在 Agent ReAct 循环的各个关键节点插入横切关注点（取消检查、指标采集、审计日志、权限增强等），**不污染主循环代码**。 -核心思路：允许在 Agent 运行的各个关键节点插入自定义逻辑（取消检查、指标采集、审计日志等），而不污染核心 ReAct 循环代码。 +--- + +## 源码结构 + +``` +src/agent/hooks/ +├── mod.rs — 模块入口（声明 + 重导出 + 测试，~540 行） +├── types.rs — 数据类型定义（事件、上下文、动作、结果，~370 行） +├── traits.rs — AgentHook + AsyncAgentHook trait（~120 行） +├── matcher.rs — ToolNamePattern / ToolMatchFilter 工具匹配器（~300 行） +├── registry.rs — HookRegistry 结构体 + 基础方法（~280 行） +├── dispatch.rs — HookRegistry 调度方法（所有 run_*，~700 行） +└── builtins.rs — 内置 Hooks（CancellationHook、MetricsHook、AuditLogHook、ContextDeduplicator，~260 行） +``` + +--- + +## 设计哲学 + +```mermaid +mindmap + root((Hook 系统
设计原则)) + 接口隔离 + 10 个生命周期方法 + +2 个权限专用方法 (on_permission_request/on_permission_denied) + 全部有默认空实现 + 只覆写关心的 hook 点 + 事件索引 + subscribed_events 声明 + 预计算 event_index + 避免无关 hook 调用 + 工具级过滤 + match_filter 声明 + ToolNamePattern 匹配 + content_patterns 过滤 + 并行优先 + 无依赖事件 join_all + 有副作用事件顺序执行 + OnSessionStop 例外 + 超时隔离 + per-hook timeout 可配置 + 默认 30 秒 + 超时跳过不拖垮全部 + 聚合不丢 + additional_contexts 全保留 + tagged_contexts 有来源 + blocking_errors 全收集 + warnings/metadata 不丢失 + 双模式共存 + 全局 hooks 构造时注册 + session hooks 运行时动态 + AsyncAgentHook 异步 fire-and-forget + 析构自动清理 +``` + +--- ## 架构总览 ```mermaid graph TB - subgraph Registry["HookRegistry"] + subgraph RT["AgentRuntime::run_turn()"] direction TB - Methods["聚合方法（遍历所有 hook 依次调用）
run_on_session_start() | run_pre_tool_use()
run_post_tool_use() | run_on_step_complete()
run_on_session_stop() | run_on_subagent_start()
run_on_subagent_stop() | run_on_pre_compact()
run_on_post_compact()"] + Start["OnSessionStart"] --> Loop["ReAct Loop"] + Loop --> Comp["PreCompact / PostCompact"] + Loop --> Stop["OnSessionStop"] end - Registry --> CH["CancellationHook
Arc<HashSet<String>>"] - Registry --> MH["MetricsHook
Arc<Mutex<MetricsData>> (共享)"] - Registry --> AH["AuditLogHook
SqlitePool (fire-and-forget 写入)"] + subgraph Registry["HookRegistry"] + direction LR + Index["event_index
HashMap<HookEvent, Vec<usize>>
预计算索引"] + AIndex["async_event_index
异步 hook 索引"] + SH["session_hooks
HashMap<String, Vec<Box<dyn AgentHook>>
运行时动态管理"] + end + + subgraph Executor["executor::execute_parallel()"] + direction LR + Pre["PreToolUse
并行 join_all + timeout + match_filter"] + Exec["工具执行"] + Post["PostToolUse
并行 join_all + timeout + match_filter"] + Fail["PostToolUseFailure
is_error 时触发"] + Async["AsyncAgentHook
fire-and-forget, 5s 超时"] + end + + subgraph Builtins["内置 Hooks（3 个 sync + 0 async）"] + CH["CancellationHook
订阅: PreToolUse, OnSessionStop"] + MH["MetricsHook
订阅: OnSessionStart, PostToolUse,
OnStepComplete, OnSessionStop"] + AH["AuditLogHook
订阅: PostToolUse, OnSessionStop"] + end + + RT --> Registry + Loop --> Executor + Registry --> Builtins + Executor --> Registry + Executor --> Async ``` -## 生命周期事件（9 个） +--- -| # | 事件 | 触发时机 | 返回值 | 调用位置 | -|---|------|---------|--------|---------| -| 1 | `OnSessionStart` | 会话创建/恢复 | 无 (fire-and-forget) | `AgentRuntime::run_turn()` | -| 2 | `PreToolUse` | 每个工具执行前 | `PreToolUseAction` (可拦截/修改参数) | `executor::execute_parallel()` | -| 3 | `PostToolUse` | 每个工具执行后 | `PostToolUseAction` (可修改输出) | `executor::execute_parallel()` | -| 4 | `OnStepComplete` | 每步 ReAct 结束 | 无 | `AgentRuntime::run_react_loop()` | -| 5 | `OnSessionStop` | 会话终止（任何原因） | 无 | `finalize::finalize_turn()` | -| 6 | `OnSubagentStart` | 子代理启动 | 无 | `SubAgentRunner::run()` | -| 7 | `OnSubagentStop` | 子代理停止 | 无 | `SubAgentRunner::run()` | -| 8 | `OnPreCompact` | 上下文压缩前 | 无 | `compact::compress_context_with_hooks()` | -| 9 | `OnPostCompact` | 上下文压缩后 | 无 | `compact::compress_context_with_hooks()` | +## 类型系统 -## 核心类型 +```mermaid +classDiagram + class AgentHook { + <> + +name() &str + +subscribed_events() &[HookEvent] + +timeout() Option~Duration~ + +match_filter() ToolMatchFilter + +on_session_start(ctx) + +pre_tool_use(ctx) PreToolUseAction + +post_tool_use(ctx) PostToolUseAction + +on_post_tool_use_failure(ctx) PostToolUseAction + +on_step_complete(ctx) + +on_session_stop(ctx) + +on_subagent_start(ctx) + +on_subagent_stop(ctx) + +on_pre_compact(ctx) + +on_post_compact(ctx) + +on_permission_request(ctx) PermissionRequestAction + +on_permission_denied(ctx) + } -### AgentHook trait + class AsyncAgentHook { + <> + +name() &str + +subscribed_events() &[HookEvent] + +on_post_tool_use_async(ctx) + +on_post_tool_use_failure_async(ctx) + +on_session_stop_async(ctx) + } + + class HookEvent { + <> + OnSessionStart + PreToolUse + PostToolUse + PostToolUseFailure + OnStepComplete + OnSessionStop + OnSubagentStart + OnSubagentStop + OnPreCompact + OnPostCompact + PermissionRequest + PermissionDenied + } + + class ToolMatchFilter { + +name_patterns: Vec~ToolNamePattern~ + +content_patterns: Vec~String~ + +matches(tool_name, tool_args) bool + +exact(name) Self + +prefix(name) Self + +is_match_all() bool + } + + class ToolNamePattern { + <> + Exact(String) + Prefix(String) + Wildcard + +parse(pattern) Self + +matches(tool_name) bool + } + + class PreToolUseAction { + <> + Continue + Block ~reason: String~ + MutateInput ~updated_args, additional_context~ + PermissionRequired ~permission, tool_name~ + } + + class PermissionRequestAction { + <> + Continue + Override ~decision: PermissionDecision, reason~ + InjectContext ~context: String~ + } + + class PermissionDecision { + <> + Allow + Ask + Deny + } + + class PermissionDenialSource { + <> + Rule + Classifier + User + Timeout + } + + class PostToolUseAction { + <> + Continue + MutateOutput ~updated_content, additional_context~ + Warning ~message, truncate_output~ + PermissionRequired ~permission, tool_name~ + Metadata ~key, value~ + } + + class PreToolUseResult { + +action: PreToolUseAction + +additional_contexts: Vec~String~ + +additional_context: Option~String~ + +blocking_errors: Vec~BlockingError~ + +final_args: Value + +tagged_contexts: Vec~TaggedContext~ + } + + class PostToolUseResult { + +final_content: String + +additional_contexts: Vec~String~ + +tagged_contexts: Vec~TaggedContext~ + +warnings: Vec~String~ + +post_permission_requests: Vec~(String,String)~ + +metadata: HashMap~String, Value~ + } + + class BlockingError { + +hook_name: String + +reason: String + } + + class TaggedContext { + +hook_name: String + +source_event: HookEvent + +content: String + +timestamp: DateTime~Utc~ + } + + class HookRegistry { + -hooks: Vec~Box~dyn AgentHook~~ + -event_index: HashMap + -session_hooks: HashMap + -async_hooks: Vec~Box~dyn AsyncAgentHook~~ + -async_event_index: HashMap + +with_builtins(db, cancelled, metrics) Self + +add(hook) + +add_async(hook) + +add_session_hook(sid, hook) + +remove_session_hook(sid, name) bool + +clear_session_hooks(sid) + +run_pre_tool_use(ctx) PreToolUseResult + +run_post_tool_use(ctx) PostToolUseResult + +run_on_session_start(ctx) + +run_on_step_complete(ctx) + +run_on_session_stop(ctx) + +run_on_subagent_start(ctx) + +run_on_subagent_stop(ctx) + +run_on_pre_compact(ctx) + +run_on_post_compact(ctx) + +run_on_post_tool_use_failure(ctx) PostToolUseResult + +run_on_permission_request(ctx) PermissionRequestAction + +run_on_permission_denied(ctx) + } + + AgentHook --> HookEvent : subscribes to + AgentHook --> ToolMatchFilter : declares + AgentHook --> PreToolUseAction : returns + AgentHook --> PostToolUseAction : returns + AsyncAgentHook --> HookEvent : subscribes to + HookRegistry --> AgentHook : manages + HookRegistry --> AsyncAgentHook : manages + HookRegistry --> PreToolUseResult : aggregates + HookRegistry --> PostToolUseResult : aggregates + PreToolUseResult --> PreToolUseAction + PreToolUseResult --> BlockingError + PreToolUseResult --> TaggedContext + PostToolUseResult --> TaggedContext +``` + +### 上下文类型（Context Objects） + +每个生命周期事件携带专用的不可变上下文： + +| 事件 | 上下文类型 | 关键字段 | +|------|-----------|---------| +| OnSessionStart | `SessionStartContext` | session_id, turn_index, is_resume | +| PreToolUse | `PreToolUseContext` | session_id, tool_name, tool_args, step | +| PostToolUse | `PostToolUseContext` | session_id, agent_name, tool_name, tool_args, output_content, is_error, step, elapsed_ms | +| PostToolUseFailure | `PostToolUseFailureContext` | session_id, agent_name, tool_name, error_message, is_interrupt, step, elapsed_ms | +| OnStepComplete | `StepCompleteContext` | session_id, step, max_steps, messages_count, estimated_tokens, token_limit | +| OnSessionStop | `SessionStopContext` | session_id, terminal (TurnTerminal), total_steps | +| OnSubagentStart | `SubagentStartContext` | parent_session_id, subagent_name, prompt | +| OnSubagentStop | `SubagentStopContext` | parent_session_id, subagent_name, result_summary, steps, is_error | +| OnPreCompact | `PreCompactContext` | session_id, message_count, estimated_tokens | +| OnPostCompact | `PostCompactContext` | session_id, new_message_count, compression_method | +| PermissionRequest | `PermissionRequestContext` | session_id, tool_name, current_decision, permission_mode, is_subagent | +| PermissionDenied | `PermissionDeniedContext` | session_id, tool_name, reason, source (Rule/Classifier/User/Timeout) | + +--- + +## ToolMatchFilter — 工具级过滤 + +比 Claude Code 的 `if` 条件更结构化。Hook 通过 `match_filter()` 返回过滤器，HookRegistry 在 dispatch 时仅对匹配的工具名和参数调用该 hook。 + +### ToolNamePattern + +| 模式 | 示例输入 | 匹配 | +|------|---------|------| +| `Exact("search_papers")` | `"search_papers"` | ✅ | +| `Exact("search_papers")` | `"download_paper"` | ❌ | +| `Prefix("file_")` | `"file_write"` | ✅ | +| `Prefix("file_")` | `"search_papers"` | ❌ | +| `Wildcard` | 任意 | ✅ | + +### ToolMatchFilter + +```rust +// 仅匹配文件相关工具且参数路径包含 .env +let filter = ToolMatchFilter { + name_patterns: vec![ToolNamePattern::parse("file_*")], + content_patterns: vec!["path:.env".to_string()], +}; +``` + +- **快捷构造**：`ToolMatchFilter::exact("search_papers")`、`ToolMatchFilter::prefix("file_")` +- **默认行为**：`ToolMatchFilter::default()` 匹配所有工具 +- **仅对** `PreToolUse`、`PostToolUse`、`PostToolUseFailure` 生效 + +--- + +## PostToolUseAction — 5 种变体 + +| 变体 | 用途 | 对执行的影响 | +|------|------|------------| +| `Continue` | 默认，保持输出不变 | 无 | +| `MutateOutput { updated_content, additional_context }` | 修改工具输出内容 | 最终输出替换为 updated_content | +| `Warning { message, truncate_output }` | 标记警告（不阻止） | 注入 `` | +| `PermissionRequired { permission, tool_name }` | 事后权限请求（advisory） | 记录审计日志 | +| `Metadata { key, value }` | 附加结构化元数据 | 保留在 PostToolUseResult.metadata 中 | + +--- + +## 异步 Hook (AsyncAgentHook) + +与 `AgentHook` 分离设计，适合后台操作（远程上报、大文件处理等）： ```rust #[async_trait] -pub trait AgentHook: Send + Sync { +pub trait AsyncAgentHook: Send + Sync { fn name(&self) -> &str; - async fn on_session_start(&self, _ctx: &SessionStartContext) {} - async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction { ... } - async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction { ... } - async fn on_step_complete(&self, _ctx: &StepCompleteContext) {} - async fn on_session_stop(&self, _ctx: &SessionStopContext<'_>) {} - 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) {} + fn subscribed_events(&self) -> &[HookEvent] { &[] } + + 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<'_>) {} } ``` -所有 9 个方法都有默认空实现——hook 实现者只需覆写关心的 hook 点，遵循**接口隔离原则**。 +- **注册**：`registry.add_async(Box::new(my_async_hook))` +- **调度**：与 sync hook 并行执行，默认 5s 超时 +- **返回值**：无——fire-and-forget 语义 -### PreToolUseAction（工具执行前返回值） +--- + +## TaggedContext — Hook 来源追踪 + +注入 LLM 的 system-reminder 携带 hook 来源标记： + +``` +[Hook: AuditLogHook | PostToolUse] 内容... +[Hook: CancellationHook | PreToolUse] 内容... +``` + +替代原来的匿名 `[Hook 注入上下文]`。 + +--- + +## ContextDeduplicator — 内容去重 + +单 dispatch cycle 内防止多个 hook 注入相同上下文消息： ```rust -pub enum PreToolUseAction { - Continue, // 允许执行（默认） - Block { reason: String }, // 阻止执行 - MutateInput { // 修改参数 + 注入上下文 - updated_args: serde_json::Value, - additional_context: Option, - }, - PermissionRequired { permission: String, tool_name: String }, // 需要权限决策 (Phase 2) +let mut dedup = ContextDeduplicator::new(); +for ctx in &result.hook_contexts { + if !dedup.is_duplicate(ctx) { + messages.push(ChatMessage::user(ctx)); + } } ``` -向后兼容：`pub type HookAction = PreToolUseAction;` +--- -### PostToolUseAction（工具执行后返回值） +## 权限决策优先级 -```rust -pub enum PostToolUseAction { - Continue, // 保持输出不变 - MutateOutput { updated_content: String }, // 修改输出内容 -} -``` +当多个权限来源冲突时，`resolve_permission_precedence()` 按以下优先级裁决： -### 聚合结果类型 +| Priority | Source | Overridable By | +|----------|--------|----------------| +| **P0** | `PermissionChecker::Deny` | 不可覆盖 | +| **P1** | Tool-level `PermissionRule::Deny` | 不可覆盖 | +| **P2** | Session-level Checker::Deny | 不可覆盖 | +| **P3** | Hook `PreToolUseAction::Block` | P0-P2 | +| **P4** | Hook `PermissionRequired` | P0-P3 | +| **P5** | Session-level Checker::Ask | 仅 Allow → Ask 升级 | +| **P6** | Tool-level `PermissionRule::Ask` | 仅 Allow → Ask 升级 | +| **P7** | `PermissionChecker::Allow` | — | +| **P8** | Implicit Allow (default) | — | -```rust -pub struct PreToolUseResult { - pub action: PreToolUseAction, // 最终动作（第一个 Block 获胜） - pub additional_context: Option, // 累积的附加上下文（所有 MutateInput 拼接） - pub final_args: serde_json::Value, // 最终参数（最后一个 MutateInput 获胜） -} +详见 [permission.md](./permission.md)。 -pub struct PostToolUseResult { - pub final_content: String, // 最终输出（最后一个 MutateOutput 获胜） -} -``` +--- -## HookRegistry 聚合逻辑 +## PermissionRequest / PermissionDenied — 权限决策钩子 -**PreToolUse 聚合（责任链 + 短路）**： +这两个专用事件在权限决策管线中提供 Hook 级别的可编程控制点，参考 Claude Code 的 `PermissionRequest` / `PermissionDenied` 事件设计。 + +### PermissionRequest — 权限请求前 + +**触发时机**：`resolve_permission_precedence()` 之后、最终权限决策生效之前。 + +**调度策略**：并行调用所有匹配的 hook，**第一个返回 `Override` 的生效**（后续 hook 仍执行但结果忽略），其他 hook 返回 `Continue` 表示不干预。 ```mermaid -flowchart TD - Start["run_pre_tool_use(ctx)"] --> Loop["遍历 hooks，依次调用 pre_tool_use()"] - Loop --> Check{"结果类型？"} - Check -->|"Block"| Short["立即短路返回
不询问后续 hooks"] - Check -->|"MutateInput"| Mut["更新 final_args
累积 additional_context (\\n 拼接)"] - Check -->|"PermissionRequired"| Log["记录日志但不阻止执行
(Phase 2 预留)"] - Check -->|"Continue"| Next["继续下一个 hook"] - Mut --> Next - Log --> Next - Next --> Loop - Short --> Return["返回 PreToolUseResult"] - Next -->|"遍历完毕"| Return +flowchart LR + subgraph Request["run_on_permission_request()"] + direction LR + H1["Hook 1"] -->|"Override(Allow)"| Win["✅ 首个 Override 生效"] + H2["Hook 2"] -->|"Continue"| Ignore["被忽略"] + H3["Hook 3"] -->|"InjectContext"| Accumulate["累积上下文"] + end ``` -关键设计： -- **第一个 Block 获胜** — 短路保护 -- **最后一个 MutateInput 获胜** — 后覆盖前 -- **additional_context 累积** — 多个 hook 的上下文用 `\n` 连接 +**返回值语义**： -**PostToolUse 聚合（全部执行，无短路）**： +| 动作 | 效果 | +|------|------| +| `PermissionRequestAction::Continue` | 不干预，沿用 `resolve_permission_precedence()` 的结果 | +| `PermissionRequestAction::Override { decision, reason }` | 覆盖决策：`Allow` 强制放行、`Ask` 升级确认、`Deny` 强制拒绝 | +| `PermissionRequestAction::InjectContext { context }` | 向 LLM 注入额外上下文（如安全策略解释） | -```mermaid -flowchart TD - Start2["run_post_tool_use(ctx)"] --> Loop2["遍历所有 hooks，依次调用 post_tool_use()"] - Loop2 --> MutOut{"MutateOutput ?"} - MutOut -->|"是"| Update["更新 final_content
(最后的 MutateOutput 获胜)"] - MutOut -->|"Continue"| Next2["继续下一个 hook"] - Update --> Next2 - Next2 --> Loop2 - Next2 -->|"遍历完毕"| Return2["返回 PostToolUseResult { final_content }"] +**上下文字段**： + +| 字段 | 类型 | 说明 | +|------|------|------| +| `session_id` | `String` | 当前会话 ID | +| `tool_name` | `String` | 被请求的工具名 | +| `tool_args` | `Value` | 工具参数（可检查路径、命令等） | +| `current_decision` | `PermissionDecision` | `resolve_permission_precedence()` 的裁决结果 | +| `permission_mode` | `String` | 当前权限模式（default/acceptEdits/bypassPermissions/plan） | +| `is_subagent` | `bool` | 是否在子代理上下文中 | + +### PermissionDenied — 权限拒绝后（审计） + +**触发时机**：权限决策结果为 `Deny`（无论来源）。 + +**调度策略**：fire-and-forget 并行调用，不阻塞主循环，返回值忽略。设计用于审计日志、安全告警、拒绝统计。 + +**来源标记（PermissionDenialSource）**： + +| 来源 | 说明 | +|------|------| +| `Rule` | 被 `PermissionChecker` 或工具级规则拒绝 | +| `Classifier` | 被 Auto-mode Classifier 拒绝（待实现） | +| `User` | 用户在交互弹窗中手动拒绝 | +| `Timeout` | 权限请求超时自动拒绝 | + +**使用示例**： + +```rust +struct DenyRateLimiter { + deny_counts: Arc>>, +} + +#[async_trait] +impl AgentHook for DenyRateLimiter { + fn name(&self) -> &str { "DenyRateLimiter" } + fn subscribed_events(&self) -> &[HookEvent] { + &[HookEvent::PermissionRequest, HookEvent::PermissionDenied] + } + + async fn on_permission_request( + &self, + ctx: &PermissionRequestContext, + ) -> PermissionRequestAction { + let count = self.get_deny_count(&ctx.tool_name); + if count >= 3 { + // 连续拒绝 3 次后强制降级为 Ask（不允许自动 Deny） + return PermissionRequestAction::Override { + decision: PermissionDecision::Ask, + reason: format!("工具 {} 已被拒绝 {} 次，需要用户确认", ctx.tool_name, count), + }; + } + PermissionRequestAction::Continue + } + + async fn on_permission_denied(&self, ctx: &PermissionDeniedContext) { + // 记录拒绝统计 + *self.deny_counts.lock().unwrap() + .entry(ctx.tool_name.clone()).or_default() += 1; + // 触发安全告警 + eprintln!("[SECURITY] 工具 {} 被拒绝: {:?} (来源: {:?})", + ctx.tool_name, ctx.reason, ctx.source); + } +} ``` -**其余 7 个事件** 均为 fire-and-forget：遍历所有 hooks 调用对应方法，不收集返回值。 +--- -## 内置 Hooks（3 个） +--- -| Hook | 覆写的事件 | 职责 | 关键依赖 | -|------|-----------|------|---------| -| `CancellationHook` | `pre_tool_use`, `on_session_stop` | 每次工具执行前检查用户是否中止会话；会话停止时清理取消令牌 | `Arc>>` (与 AppState 共享) | -| `MetricsHook` | `on_session_start`, `post_tool_use`, `on_step_complete`, `on_session_stop` | 采集运行指标：工具调用次数、步数、错误数、token 消耗；每 3 步输出摘要日志 | `Arc>` (**共享引用**，API 通过 `AgentRuntime::get_metrics()` 实时查询) | -| `AuditLogHook` | `post_tool_use`, `on_session_stop` | 所有工具调用写入 `agent_audit_log` 表（工具名、状态、耗时、输出预览）；会话终止写入 SESSION_STOP 标记 | `SqlitePool` (**fire-and-forget** 写入，不阻塞主循环) | - -> **注意**：代码中**不存在** PermissionHook。权限检查由独立的 `PermissionChecker` (`src/agent/runtime/permission.rs`) 负责，该组件在工具执行前与 hooks 并行调用，不属于 hooks 体系。`PreToolUseAction::PermissionRequired` 变体预留于 Phase 2 完善。 - -## 数据流 +## 生命周期事件全景 ```mermaid sequenceDiagram participant API as API Handler participant RT as AgentRuntime participant HR as HookRegistry - participant CH as CancellationHook - participant MH as MetricsHook - participant AH as AuditLogHook participant EX as Executor + participant SA as SubAgentRunner + participant CMP as Compact - Note over API,EX: Phase 1 — 会话启动 + Note over API,CMP: ═══ Phase 1: 会话启动 ═══ API->>RT: run_turn(question) - RT->>HR: HookRegistry::with_builtins(db, cancelled_runs, metrics_data) - RT->>HR: run_on_session_start(ctx) - HR->>MH: 记录 session_id + RT->>HR: with_builtins(db, cancelled, metrics) + RT->>HR: run_on_session_start(ctx) ⚡并行 + Note over HR: ① OnSessionStart - Note over API,EX: Phase 2 — ReAct 循环 - loop 每步 (最多 max_steps) - RT->>RT: LLM 流式调用 - RT->>EX: execute_parallel(tool_calls, hook_registry) + Note over API,CMP: ═══ Phase 2: ReAct 循环 ═══ + loop 每步 (1..max_steps) + RT->>RT: LLM 流式调用 → tool_calls + RT->>EX: execute_parallel(calls, registry) par 每个工具调用 - EX->>HR: run_pre_tool_use(pre_ctx) - HR->>CH: 检查取消状态 - alt 已取消 - CH-->>HR: Block { reason } - HR-->>EX: PreToolUseResult { action: Block } - EX-->>EX: 跳过该工具 - else 未取消 - CH-->>HR: Continue - HR-->>EX: PreToolUseResult { final_args, additional_context } - EX->>EX: 执行工具 - EX->>HR: run_post_tool_use(post_ctx) - HR->>MH: 更新工具调用计数/错误数 - HR->>AH: fire-and-forget INSERT agent_audit_log - HR-->>EX: PostToolUseResult { final_content } + EX->>HR: run_pre_tool_use(ctx) ⚡并行 + match_filter + Note over HR: ② PreToolUse + HR-->>EX: PreToolUseResult {final_args, tagged_contexts, blocks} + EX->>EX: 权限检查 (Checker + 工具级 + 会话级 + Hook) + EX->>EX: resolve_permission_precedence() + EX->>HR: run_on_permission_request(ctx) ⚡并行 (首个 Override 生效) + Note over HR: ⑪ PermissionRequest — Hook 可覆盖权限决策 + alt 权限被拒绝 + EX->>HR: run_on_permission_denied(ctx) ⚡并行 (fire-and-forget 审计) + Note over HR: ⑫ PermissionDenied + end + EX->>EX: 工具执行 + EX->>HR: run_post_tool_use(ctx) ⚡并行 + match_filter + Note over HR: ③ PostToolUse + AsyncHook dispatch + HR-->>EX: PostToolUseResult {final_content, tagged_contexts, warnings, metadata} + alt is_error + EX->>HR: run_on_post_tool_use_failure(ctx) ⚡并行 + Note over HR: ④ PostToolUseFailure + AsyncHook dispatch end end - EX-->>RT: ToolExecutionResult { tool_messages } - RT->>HR: run_on_step_complete(ctx) - HR->>MH: 每 3 步输出摘要日志 + EX-->>RT: ToolExecutionResult {messages, hook_contexts, blocking_errors} + RT->>RT: "hook_contexts → [system-reminder] 注入 (去重)" + RT->>HR: run_on_step_complete(ctx) ⚡并行 + Note over HR: ⑤ OnStepComplete opt 上下文超限 - RT->>RT: snapshot_compress_restore() - Note over RT: PreCompact / PostCompact hooks 触发 + RT->>CMP: compress_context_with_hooks() + CMP->>HR: run_on_pre_compact(ctx) ⚡并行 + Note over HR: ⑨ OnPreCompact + CMP->>CMP: compress_with_fallback() + CMP->>HR: run_on_post_compact(ctx) ⚡并行 + Note over HR: ⑩ OnPostCompact + end + + opt 子代理调用 + RT->>SA: SubAgentRunner::run() + SA->>HR: run_on_subagent_start(ctx) ⚡并行 + Note over HR: ⑦ OnSubagentStart + SA->>SA: 独立 ReAct 循环 + Note over SA: PreToolUse / PostToolUse 同样触发 + SA->>HR: run_on_subagent_stop(ctx) ⚡并行 + Note over HR: ⑧ OnSubagentStop + SA-->>RT: ToolOutput (summary) end end - Note over API,EX: Phase 3 — 会话收尾 + Note over API,CMP: ═══ Phase 3: 会话收尾 ═══ RT->>RT: finalize_turn() - RT->>HR: run_on_session_stop(ctx) - HR->>CH: 清理 cancelled_runs - HR->>MH: 输出会话结束摘要 - HR->>AH: fire-and-forget SESSION_STOP 记录 - RT-->>API: Done (SSE) + RT->>HR: run_on_session_stop(ctx) 🔒顺序 + Note over HR: ⑥ OnSessionStop + RT-->>API: SSE Done ``` -## HookRegistry 构建 +**执行模式图例**：⚡ 并行 `join_all` | 🔒 顺序执行 | ⚡+⏱ 并行 + per-hook timeout -`AgentRuntime::run_turn()` 在每次 turn 开始时构建 `HookRegistry`： +--- -```rust -let hook_registry = HookRegistry::with_builtins( - db.clone(), // → AuditLogHook - self.app_state.cancelled_runs.clone(), // → CancellationHook - Some(self.metrics_data.clone()), // → MetricsHook (共享引用) -); +## HookRegistry 调度机制 + +### 事件索引 + 工具过滤 + +```mermaid +flowchart LR + subgraph Register["add(hook)"] + H["hook: Box<dyn AgentHook>"] --> Subs{"subscribed_events()"} + Subs -->|"空 = 全部"| All["遍历所有 12 个 HookEvent
event_index[event].push(idx)"] + Subs -->|"指定"| Spec["仅注册声明的事件
event_index[event].push(idx)"] + end + + subgraph Dispatch["run_*_tool_use(ctx, session_id)"] + Collect["collect_tool_hooks_for(event, sid, tool_name, tool_args)"] + Collect --> GI["event_index[event] → 全局 hooks"] + Collect --> FI["match_filter().matches(tool_name, tool_args) → 过滤"] + Collect --> SI["session_hooks[sid] → session hooks (不过滤)"] + GI --> Merge["统一 Vec<&dyn AgentHook>"] + FI --> Merge + SI --> Merge + end + + Register --> Index["event_index: HashMap<HookEvent, Vec<usize>>"] + Dispatch --> Index ``` -`MetricsHook` 使用 `from_arc()` 复用 `AgentRuntime` 自身的 `metrics_data: Arc>`，确保 hook 内部采集的指标与 `AgentRuntime::get_metrics()` API 查询返回的是同一份数据。 +**关键**：`collect_tool_hooks_for()` 在 `collect_hooks_for()` 基础上增加 `match_filter()` 过滤层。全局 hooks 按事件索引 + 工具匹配双重过滤，session hooks 不过滤（保持全量响应）。 -## 子代理中的 Hooks +--- -`SubAgentRunner` 拥有独立的 hook 管道，共享同一个 `HookRegistry` 实例： +## 数据流：Hook 结果如何影响 Agent 行为 -- `on_subagent_start` / `on_subagent_stop` 在子代理生命周期的首尾触发 -- 子代理的工具执行也经过 `run_pre_tool_use` / `run_post_tool_use`（通过 `subagent.rs:447-518`） -- **已知不足**：子代理内部的上下文压缩 (`subagent.rs:270`) 直接调用 `compress_context` 而非 `compress_context_with_hooks`，导致 PreCompact/PostCompact 事件**不会**在子代理压缩时触发 +```mermaid +flowchart TD + subgraph Executor["executor::execute_parallel()"] + PreHooks["PreToolUse hooks → PreToolUseResult"] + Perms["权限检查 4 层叠加
Checker → 工具级规则 → 会话级规则 → resolve_permission_precedence"] + PermHooks["PermissionRequest hooks → PermissionRequestAction
首个 Override 生效，可覆盖 Allow→Ask 升级或 Deny→Ask 降级"] + PermDenied["PermissionDenied hooks → fire-and-forget 审计
Rule/Classifier/User/Timeout 四种来源标记"] + Tools["工具执行"] + PostHooks["PostToolUse hooks → PostToolUseResult"] + FailHooks["PostToolUseFailure hooks
（仅 is_error）"] + AsyncHooks["AsyncAgentHook dispatch
（fire-and-forget, 5s 超时）"] + end -## 扩展方式 + subgraph ReactLoop["AgentRuntime::run_react_loop()"] + Push["tool_messages → messages"] + Inject["hook_contexts 包装为 [system-reminder]
（带 TaggedContext 来源 + ContextDeduplicator 去重）"] + LogBlock["blocking_errors → warn! 日志"] + LogWarn["warnings → warn! 日志"] + LogPerm["post_permission_requests → info! 审计"] + StepHook["run_on_step_complete"] + end -添加自定义 hook 只需两步： + PreHooks -->|"final_args 覆盖参数"| Perms + PreHooks -->|"PermissionRequired → Allow→Ask 升级"| Perms + Perms -->|"权限决策"| PermHooks + PermHooks -->|"Override 可覆盖 Allow/Deny/Ask"| Perms + Perms -->|"Deny: 触发"| PermDenied + PermDenied -->|"审计日志"| LogPerm + Perms -->|"Deny: 注入错误 result"| Push + Perms -->|"Allow/Ask/Approve: 进入执行"| Tools + Tools -->|"output"| PostHooks + PostHooks -->|"final_content 覆盖输出"| Push + PostHooks -->|"warnings, metadata, tagged_contexts"| Inject + Tools -->|"is_error"| FailHooks + FailHooks -->|"final_content 覆盖错误输出"| Push + PostHooks -.->|"并行 fire-and-forget"| AsyncHooks + FailHooks -.->|"并行 fire-and-forget"| AsyncHooks + Push --> Inject + Inject --> StepHook + LogBlock --> StepHook + LogWarn --> StepHook + LogPerm --> StepHook +``` + +### hook_contexts 注入格式 + +```rust +// 带来源标记（TaggedContext） +// runtime/mod.rs +let mut dedup = ContextDeduplicator::new(); +for ctx in &exec_result.hook_contexts { + if dedup.is_duplicate(ctx) { + continue; + } + let reminder = format!( + "\n[Hook 注入上下文]\n{}\n", + ctx + ); + messages.push(ChatMessage::user(&reminder)); +} +``` + +内容格式： +``` +[Hook: AuditLogHook | PostToolUse] 内容... +[Hook: CancellationHook | PreToolUse] 内容... +``` + +--- + +## 内置 Hooks 详解 + +### CancellationHook + +```mermaid +stateDiagram-v2 + [*] --> Active: 会话开始 + Active --> CheckPreTool: 每次 PreToolUse + CheckPreTool --> Blocked: cancelled_runs 含 session_id + CheckPreTool --> Continue: cancelled_runs 不含 session_id + Blocked --> [*]: Block{reason: "用户已手动中止"} + Continue --> Active: 继续执行 + Active --> Cleanup: OnSessionStop + Cleanup --> [*]: cancelled_runs.remove(session_id) +``` + +依赖 `Arc>>`（与 `AppState.cancelled_runs` 共享引用）。 + +### MetricsHook + +通过 `from_arc()` 复用 `AgentRuntime` 自身的 `metrics_data: Arc>`，确保 hook 内部采集的数据与 `AgentRuntime::get_metrics()` API 返回的是同一份数据。 + +订阅事件：`OnSessionStart`（设置 session_id）、`PostToolUse`（累加计数）、`OnStepComplete`（每 3 步输出摘要）、`OnSessionStop`（输出终止原因）。 + +### AuditLogHook + +```mermaid +sequenceDiagram + participant EX as Executor + participant AH as AuditLogHook + participant DB as SQLite + + EX->>AH: post_tool_use(ctx) + AH->>AH: 构造 status = "OK" | "FAIL" + AH->>AH: output_preview = content[..200] + AH--)DB: "tokio::spawn INSERT agent_audit_log" + EX->>AH: on_session_stop(ctx) + AH--)DB: "tokio::spawn INSERT SESSION_STOP" +``` + +Fire-and-forget 写入，不阻塞主循环。未来可改为实现 `AsyncAgentHook`。 + +--- + +## 扩展指南 + +### 自定义全局 Sync Hook ```rust -// 1. 实现 AgentHook trait struct MyCustomHook; #[async_trait] impl AgentHook for MyCustomHook { fn name(&self) -> &str { "MyCustomHook" } + // 声明只关心 PreToolUse 和 PostToolUseFailure + fn subscribed_events(&self) -> &[HookEvent] { + &[HookEvent::PreToolUse, HookEvent::PostToolUseFailure] + } + + // 仅对文件工具生效 + fn match_filter(&self) -> ToolMatchFilter { + ToolMatchFilter::prefix("file_") + } + + fn timeout(&self) -> Option { + Some(Duration::from_secs(10)) + } + async fn pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseAction { - // 自定义逻辑 + if ctx.tool_name == "run_bash" { + return PreToolUseAction::PermissionRequired { + permission: "执行系统命令需要二次确认".into(), + tool_name: ctx.tool_name.clone(), + }; + } PreToolUseAction::Continue } + + async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction { + if ctx.is_error { + PostToolUseAction::Warning { + message: format!("工具 {} 执行失败", ctx.tool_name), + truncate_output: false, + } + } else { + PostToolUseAction::Continue + } + } + + async fn on_post_tool_use_failure( + &self, + ctx: &PostToolUseFailureContext, + ) -> PostToolUseAction { + if ctx.is_interrupt { + return PostToolUseAction::Continue; + } + PostToolUseAction::MutateOutput { + updated_content: format!( + "[已由 MyCustomHook 处理] 工具 {} 执行失败: {}。建议尝试替代方案。", + ctx.tool_name, ctx.error_message + ), + additional_context: None, + } + } } -// 2. 注册到 HookRegistry +// 注册 registry.add(Box::new(MyCustomHook)); ``` -## 测试覆盖 +### 自定义 Async Hook -`hooks.rs` 包含 8 个单元测试（`#[cfg(test)] mod tests`），覆盖： +```rust +struct RemoteLogger; -| 测试 | 验证点 | -|------|-------| -| `test_hook_registry_runs_all_hooks` | 注册表遍历调用所有 hook | -| `test_blocking_hook_stops_chain` | Block 短路机制 | -| `test_mutate_input_accumulates_context` | 参数修改 + 上下文累积 | -| `test_post_tool_use_mutate_output` | 输出修改 | -| `test_cancellation_hook_blocks_when_cancelled` | CancellationHook 阻止逻辑 | -| `test_cancellation_hook_allows_when_not_cancelled` | CancellationHook 放行逻辑 | -| `test_metrics_hook_accumulates_counts` | MetricsHook 累加正确性 | -| `test_session_start_hook_called` | OnSessionStart 调用 | -| `test_new_lifecycle_events_called` | OnSubagentStart/Stop, PreCompact/PostCompact 调用 | +#[async_trait] +impl AsyncAgentHook for RemoteLogger { + fn name(&self) -> &str { "RemoteLogger" } -## 已知改进项 + async fn on_post_tool_use_async(&self, ctx: PostToolUseContext) { + // 异步上报工具调用日志到远程服务 + tokio::spawn(async move { + upload_to_remote(ctx).await; + }); + } +} -| 问题 | 说明 | -|------|------| -| `PermissionRequired` 未实现 | 代码中存在此变体但被当作 `Continue` 处理，注释标明 "Permission 系统在 Phase 2 中完善" | -| 取消检查重复 | `CancellationHook::pre_tool_use` 与 `AgentRuntime::run_react_loop` 中的显式检查存在功能重叠 | -| 子代理压缩未走 hooks | `subagent.rs:270` 直接调用 `compress_context` 而非 `compress_context_with_hooks` | -| 缺少 `on_error` 事件 | `AgentHook` trait 没有错误生命周期事件，错误场景无法通过 hook 拦截 | +registry.add_async(Box::new(RemoteLogger)); +``` + +### 动态 Session Hook + +```rust +// 为当前会话临时添加 +registry.add_session_hook(&session_id, Box::new(MyTempHook)); + +// 移除 +registry.remove_session_hook(&session_id, "MyTempHook"); + +// 会话结束时自动清理（HookRegistry 析构） +``` --- + +## 集成点地图 + +| 文件 | 集成点 | 方法 | +|------|--------|------| +| `runtime/mod.rs` | 会话启动 | `run_on_session_start()` | +| `runtime/mod.rs` | ReAct 循环 | `run_on_step_complete()` | +| `runtime/mod.rs` | Hook 上下文注入 | `hook_contexts → [system-reminder]` | +| `runtime/mod.rs` | 会话终止 | `run_on_session_stop()` | +| `executor.rs` | 工具执行前 | `run_pre_tool_use()` | +| `executor.rs` | 工具执行后 | `run_post_tool_use()` | +| `executor.rs` | 工具执行失败 | `run_on_post_tool_use_failure()` | +| `executor.rs` | 权限决策前 | `run_on_permission_request()` | +| `executor.rs` | 权限拒绝后 | `run_on_permission_denied()` (fire-and-forget 审计) | +| `executor.rs` | 权限决策 | `resolve_permission_precedence()` | +| `subagent.rs` | 子代理启动 | `run_on_subagent_start()` | +| `subagent.rs` | 子代理停止 | `run_on_subagent_stop()` | +| `compact.rs` | 压缩前后 | `run_on_pre_compact()` / `run_on_post_compact()` | + +--- + +## 配置 + +| 环境变量 | 默认值 | 说明 | +|---------|--------|------| +| `HOOK_TIMEOUT_SECS` | 30 | 全局 per-hook 超时（秒）。单个 hook 可通过 `timeout()` trait 方法覆盖 | + +--- + +## 设计权衡 + +| 考量点 | 说明 | +|--------|------| +| **工具级过滤** | `match_filter()` 提供 name + content 双重过滤，session hooks 不过滤 | +| **并行 vs 顺序** | 单 hook 自动走顺序路径（无 timeout 开销），多 hook 并行 + timeout | +| **聚合策略** | 最后覆盖 final_args/content，全部保留 contexts/blocks/warnings/metadata | +| **OnSessionStop** | 顺序执行（关键清理），`&self` 限制需调用方手动 clear session hooks | +| **AsyncAgentHook** | 独立 trait，fire-and-forget 语义，5s dispatch 超时，与 sync hook 并行 | +| **TaggedContext** | 带 hook 名称 + 事件来源，格式化注入为 `[Hook: 名称 | 事件]` | +| **ContextDeduplicator** | 内容哈希去重，单 cycle 内有效，不跨 step 去重 | +| **Block 非全局中断** | `PreToolUseAction::Block` 只阻止当前工具的单个调用，全局中断由 `CancellationHook` 实现 | +| **Warning/Post-hoc Permission** | PostToolUse 的 Warning 和 PermissionRequired 均为 advisory，不阻塞主循环 | +| **PermissionRequest 调度** | 并行调用所有 hook，首个 `Override` 生效，后续 Continue 被忽略；`InjectContext` 全部累积 | +| **PermissionDenied 调度** | fire-and-forget 并行调用，不阻塞主循环，返回值忽略（纯审计用途） | + +--- + +## 变更记录 + +| 日期 | 变更 | +|------|------| +| 2026-06-22 | Phase 1: 新增 ToolMatchFilter / TaggedContext / 工具级过滤 | +| 2026-06-22 | Phase 2: 新增 resolve_permission_precedence + PermissionChecker::explain | +| 2026-06-22 | Phase 3: PostToolUseAction 扩展 (Warning/PermissionRequired/Metadata) + ContextDeduplicator | +| 2026-06-22 | Phase 4: AsyncAgentHook trait + async hook 调度 | +| 2026-06-22 | 模块拆分: mod.rs(537) ← types/traits/matcher/registry/dispatch/builtins | +| 2026-06-22 | Phase 5: PermissionRequest/PermissionDenied 事件 — 权限决策钩子 + fire-and-forget 审计 | + +--- + +## 测试覆盖 + +| 测试 | 覆盖路径 | +|------|---------| +| `test_hook_registry_runs_all_hooks` | 注册表遍历调用 | +| `test_blocking_hook_stops_chain` | Block 收集到 blocking_errors | +| `test_mutate_input_accumulates_context` | 参数修改 + 上下文累积 | +| `test_post_tool_use_mutate_output` | 输出修改（含 additional_context） | +| `test_cancellation_hook_blocks_when_cancelled` | CancellationHook Block 路径 | +| `test_cancellation_hook_allows_when_not_cancelled` | CancellationHook 放行路径 | +| `test_metrics_hook_accumulates_counts` | MetricsHook 累加正确性 | +| `test_session_start_hook_called` | OnSessionStart 触发 | +| `test_new_lifecycle_events_called` | OnSubagentStart/Stop, Pre/PostCompact 触发 | +| `test_match_filter_skips_irrelevant` | ToolMatchFilter 正确过滤无关工具 | +| `test_full_hook_pipeline_integration` | 全链路：match_filter + TaggedContext + Warning + Metadata | +| `test_matcher_*` (12 个) | ToolNamePattern / ToolMatchFilter 匹配逻辑 | +| `test_precedence_*` (5 个) | 多源权限优先级裁决 | +| `test_context_deduplication` | ContextDeduplicator 去重 | +| `test_async_hook_dispatch` | AsyncAgentHook 并行调度 | +| `test_permission_request_override` | PermissionRequest — Override 覆盖 Allow→Ask 升级 | +| `test_permission_request_continue` | PermissionRequest — Continue 不干预默认决策 | +| `test_permission_denied_fire_and_forget` | PermissionDenied — fire-and-forget 审计调度 | +| `test_permission_denied_sources` | PermissionDenied — Rule/Classifier/User/Timeout 四种来源标记 | diff --git a/docs/architecture/agent/permission.md b/docs/architecture/agent/permission.md index 2055e9b..bbc7de8 100644 --- a/docs/architecture/agent/permission.md +++ b/docs/architecture/agent/permission.md @@ -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>`) 实现跨 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 权限配置 | diff --git a/docs/architecture/agent/skills.md b/docs/architecture/agent/skills.md index 25d5c79..254a17b 100644 --- a/docs/architecture/agent/skills.md +++ b/docs/architecture/agent/skills.md @@ -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 ├── last_scan_mtime: Option └── usage_stats: HashMap + +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()
扫描 agent_messages 表"] + SC["SkillCreator::create_from_patterns()
生成 SKILL.md 文件"] + CR["Curator::analyze()
质量评分 + 生命周期评估"] + end + + subgraph Background["后台定期维护"] + direction LR + Runner["CuratorRunner::spawn()
空闲触发 + 间隔检查"] + Archive["archive_stale_skills()
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, // 如 ["search_papers", "download_paper", "rag_search"] + pub session_ids: Vec, // 出现的 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, 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` 自动从工具名推断 +- 每个工具调用作为 `` 写入正文 +- `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, + last_run: RwLock>, +} + +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, + skill_metas: &[SkillMeta], + ) -> Option; + + /// 启动后台任务 + pub fn spawn( + self: Arc, + usage_stats: Arc>>, + skill_metas: Arc>>, + 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; + + /// 仅检测模式（不创建） + pub async fn detect_only(&self) -> Result>; + + /// 仅分析已有 skills（不检测新模式） + pub async fn analyze_only( + &self, + usage_stats: &HashMap, + skill_metas: &[SkillMeta], + ) -> Result; +} + +pub struct SelfImproveResult { + pub patterns_found: usize, + pub skills_created: usize, + pub skills_created_names: Vec, + pub curator_report: CuratorReport, +} +``` + +**管道流程**： + +``` +SelfImprovePipeline::run(dry_run=true) + │ + ├─ 1. PatternDetector::scan() + │ └─ 从 agent_messages 中检测 ≥3 次跨 session 重复序列 + │ └─ 结果: Vec（按 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` 字段 | | `` 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 + 心跳记录） | --- diff --git a/docs/architecture/agent/system-prompt.md b/docs/architecture/agent/system-prompt.md index a64d840..9e486c9 100644 --- a/docs/architecture/agent/system-prompt.md +++ b/docs/architecture/agent/system-prompt.md @@ -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 模块化组装
静态(5) → 动态(4)"] + Cache["SystemPromptCache
首次计算，永久复用"] 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) -> Option; + // 显式失效 + 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() 生成（ XML） -Section 4: memory 动态 — 从 MemoryManager.build_system_reminder(5) 生成（ 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 中注入 `` 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 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 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... ``` -**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 条记忆，生成 `` XML 块： - -```xml - -[PROJECT MEMORY] -[偏好] memory-slug: 一句话描述 - 内容预览前三行 - [时效提示: 此记忆已超过N天，可能已过时] -[反馈] another-memory: 描述 [已更新→new-slug] - 内容预览... -使用 save_memory 工具保存重要信息。记忆内容可能过时，请在使用前验证。 - -``` - -关键特性： -- 按 mtime 排序（最新在前），支持语义选择 + 指数衰减排序 -- 按类型标注：`[偏好]` / `[反馈]` / `[项目]` / `[参考]` -- 过期记忆标记为 `[已更新]` 或 `[已更新→new-slug]`（归档为 `{slug}_v1.md`） -- 超过 1 天的记忆注入时效警告 -- 索引文件 `MEMORY.md` 限制 200 行 / 25KB +从 `{library_dir}/memory/` 加载，**按 mtime 降序排列**（最新在前），取前 5 条，注入 `` 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 创建
system_prompt() 调用"] - - RT --> S1["Section 1: identity
(静态常量)"] - RT --> S2["Section 2: tools
(ToolRegistry definitions)"] - RT --> S3["Section 3: skills
(SkillRegistry.build_reminder)"] - - S1 --> S4 - S2 --> S4 - S3 --> S4["Section 4: memory (可选)
(MemoryManager.build_reminder, 5 entries)"] - - S4 --> S5["Section 5: principles
(静态常量)"] - S5 --> ASM["assemble()
join('\n\n')"] - - ASM --> Main["主 Agent 上下文
build_initial_context()
+ nudge 注入 + 任务恢复 + 后台通知"] - ASM --> Sub["子 Agent 上下文
SubAgentRunner.run()
(独立 system_prompt)"] - - Main --> React["ReAct 循环"] - Main --> NudgeInj["Nudge 消息注入 (user)"] - Main --> Compact["压缩层
generate_summary()
+ 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 + 之前发现过的延迟工具 + 块列出所有可发现的延迟工具名称 + + → 模型调用 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**：`` 块作为 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 数组 + 注入 `` 块 +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（开闭原则） diff --git a/docs/architecture/agent/tools.md b/docs/architecture/agent/tools.md index 0b7dda7..571b284 100644 --- a/docs/architecture/agent/tools.md +++ b/docs/architecture/agent/tools.md @@ -16,7 +16,7 @@ graph TB subgraph Layer2["执行协调层 — Executor"] direction LR EX["src/agent/runtime/executor.rs
验证 → PreToolUse hooks → 并行调度 → PostToolUse"] - SX["src/agent/runtime/streaming_executor.rs (流式变体)
流式 tool_use 到达时立即调度 + Sibling Abort"] + SX["src/agent/runtime/streaming_executor.rs (流式变体)
流式 tool_use 到达时立即调度 + Sibling Abort
非并发工具独占执行 (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` 保护写入操作不被用户取消打断 --- diff --git a/migrations/20260622000000_session_rewind.sql b/migrations/20260622000000_session_rewind.sql new file mode 100644 index 0000000..ac09a18 --- /dev/null +++ b/migrations/20260622000000_session_rewind.sql @@ -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); diff --git a/migrations/20260622000001_session_branch.sql b/migrations/20260622000001_session_branch.sql new file mode 100644 index 0000000..b4542d0 --- /dev/null +++ b/migrations/20260622000001_session_branch.sql @@ -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; diff --git a/src/agent/compact.rs b/src/agent/compact.rs index c73d440..4cf8bf5 100644 --- a/src/agent/compact.rs +++ b/src/agent/compact.rs @@ -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) { } } -/// 使用 LLM 生成对话摘要。 -async fn generate_summary(to_summarize: &[ChatMessage], llm: &LlmClient) -> Result { +/// 摘要标记前缀，用于识别已有的压缩摘要。 +const SUMMARY_MARKER: &str = "[历史对话摘要]"; +const COMPRESSION_LOG_MARKER: &str = "[上下文压缩历史]"; + +/// 从消息列表中提取已有的历史摘要内容。 +/// 返回 None 表示无已有摘要。 +fn extract_prior_summary(messages: &[ChatMessage]) -> Option { + // 从后往前找最近的一份摘要（可能有多次压缩） + 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 { 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::>() .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, 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, + llm: &LlmClient, + context_char_limit: usize, + session_id: &str, + hook_registry: Option<&HookRegistry>, + collapse_log: Option<&std::sync::Mutex>, ) { 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 配对靠近切割点的场景 diff --git a/src/agent/hooks.rs b/src/agent/hooks.rs deleted file mode 100644 index 2949d95..0000000 --- a/src/agent/hooks.rs +++ /dev/null @@ -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, - }, - /// 需要权限决策 - 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, - pub total_steps: usize, - pub total_errors: usize, - pub session_id: Option, -} - -// ── 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, - /// 最终的工具参数（应用了最后一个 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>, -} - -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>>, - metrics_data: Option>>, - ) -> 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) { - info!("[Hooks] 注册 hook: {}", hook.name()); - self.hooks.push(hook); - } - - /// 获取所有 hooks 的不可变引用 - pub fn all(&self) -> &[Box] { - &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>>, -} - -impl CancellationHook { - pub fn new(cancelled_runs: Arc>>) -> 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>, -} - -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> 创建（共享数据引用） - pub fn from_arc(data: Arc>) -> Self { - MetricsHook { data } - } - - /// 返回当前指标快照（锁异常时返回 None） - pub fn snapshot(&self) -> Option { - self.data.lock().ok().map(|d| d.clone()) - } - - /// 获取 Arc 引用，供外部持有 - pub fn data_arc(&self) -> Arc> { - 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, - } - - 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>, - } - - #[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, - subagent_stop: std::sync::Mutex, - pre_compact: std::sync::Mutex, - post_compact: std::sync::Mutex, - } - - #[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 - } -} diff --git a/src/agent/hooks/builtins.rs b/src/agent/hooks/builtins.rs new file mode 100644 index 0000000..1fa3eaf --- /dev/null +++ b/src/agent/hooks/builtins.rs @@ -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, +} + +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>>, +} + +impl CancellationHook { + pub fn new(cancelled_runs: Arc>>) -> 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>, +} + +impl Default for MetricsHook { + fn default() -> Self { + Self::new() + } +} + +impl MetricsHook { + pub fn new() -> Self { + MetricsHook { + data: Arc::new(Mutex::new(MetricsData::default())), + } + } + + /// 从已有的 Arc> 创建（共享数据引用） + pub fn from_arc(data: Arc>) -> Self { + MetricsHook { data } + } + + /// 返回当前指标快照（锁异常时返回 None） + pub fn snapshot(&self) -> Option { + self.data.lock().ok().map(|d| d.clone()) + } + + /// 获取 Arc 引用，供外部持有 + pub fn data_arc(&self) -> Arc> { + 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; + }); + } +} diff --git a/src/agent/hooks/dispatch.rs b/src/agent/hooks/dispatch.rs new file mode 100644 index 0000000..85b3563 --- /dev/null +++ b/src/agent/hooks/dispatch.rs @@ -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 = Vec::new(); + let mut tagged_contexts: Vec = Vec::new(); + let mut final_args = ctx.tool_args.clone(); + let mut final_action = PreToolUseAction::Continue; + let mut blocking_errors: Vec = 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 = Vec::new(); + let mut tagged_contexts: Vec = Vec::new(); + let mut final_args = ctx.tool_args.clone(); + let mut final_action = PreToolUseAction::Continue; + let mut blocking_errors: Vec = 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 = Vec::new(); + let mut tagged_contexts: Vec = Vec::new(); + let mut warnings: Vec = Vec::new(); + let mut post_permission_requests: Vec<(String, String)> = Vec::new(); + let mut metadata: std::collections::HashMap = + 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::() + + "\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 = Vec::new(); + let mut tagged_contexts: Vec = Vec::new(); + let mut warnings: Vec = Vec::new(); + let mut post_permission_requests: Vec<(String, String)> = Vec::new(); + let mut metadata: std::collections::HashMap = + 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::() + + "\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 = Vec::new(); + let mut tagged_contexts: Vec = 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 = Vec::new(); + let mut tagged_contexts: Vec = 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) { + 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 = 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> = 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; + } +} diff --git a/src/agent/hooks/matcher.rs b/src/agent/hooks/matcher.rs new file mode 100644 index 0000000..376fa0d --- /dev/null +++ b/src/agent/hooks/matcher.rs @@ -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, + /// 可选的内容级模式。格式：`"字段:子串"` 或直接 `"子串"`。 + /// 例如 `"command:rm *"` 匹配 `run_bash` 调用中 command 字段以 "rm " 开头的情况。 + /// 空 Vec = 不过滤内容。 + pub content_patterns: Vec, +} + +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()); + } +} diff --git a/src/agent/hooks/mod.rs b/src/agent/hooks/mod.rs new file mode 100644 index 0000000..ad1d98d --- /dev/null +++ b/src/agent/hooks/mod.rs @@ -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, + } + + 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>, + } + + #[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, + subagent_stop: std::sync::Mutex, + pre_compact: std::sync::Mutex, + post_compact: std::sync::Mutex, + } + + #[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, + 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, + } + + #[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]")); + } +} diff --git a/src/agent/hooks/registry.rs b/src/agent/hooks/registry.rs new file mode 100644 index 0000000..abe38ac --- /dev/null +++ b/src/agent/hooks/registry.rs @@ -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>, + /// 事件→hook 索引映射。预计算哪些 hook 关心哪些事件， + /// 避免每次事件触发时遍历全部 hook。 + pub(crate) event_index: HashMap>, + /// Per-session 动态 hooks（运行时添加/移除，会话结束时自动清理） + pub(crate) session_hooks: HashMap>>, + /// 异步 fire-and-forget hooks（Phase 4.1+） + pub(crate) async_hooks: Vec>, + /// 异步 hook 的事件索引 + pub(crate) async_event_index: HashMap>, +} + +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>>, + metrics_data: Option>>, + ) -> 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) { + 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) { + 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] { + &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) { + 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] { + 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, + final_args: &mut serde_json::Value, + additional_contexts: &mut Vec, + tagged_contexts: &mut Vec, + 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(); + } + } +} diff --git a/src/agent/hooks/traits.rs b/src/agent/hooks/traits.rs new file mode 100644 index 0000000..bf41274 --- /dev/null +++ b/src/agent/hooks/traits.rs @@ -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 { + 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<'_>) {} +} diff --git a/src/agent/hooks/types.rs b/src/agent/hooks/types.rs new file mode 100644 index 0000000..37cd971 --- /dev/null +++ b/src/agent/hooks/types.rs @@ -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, reason: impl Into) -> 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, +} + +impl TaggedContext { + pub fn new( + hook_name: impl Into, + source_event: HookEvent, + content: impl Into, + ) -> 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, + }, + /// 需要权限决策 + 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, + }, + /// 标记工具结果为警告（不阻止，仅供 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, + pub total_steps: usize, + pub total_errors: usize, + pub session_id: Option, +} + +// ── Hook Trait ── + +/// PreToolUse 聚合结果 +#[derive(Debug, Clone)] +pub struct PreToolUseResult { + /// 最终动作（第一个非 Continue 动作获胜） + pub action: PreToolUseAction, + /// 累积的 additional_contexts（所有 MutateInput 的上下文，不丢失任何 hook 的贡献） + pub additional_contexts: Vec, + /// 向后兼容便捷字段：用换行拼接的附加上下文 + pub additional_context: Option, + /// 所有阻塞错误（多个 hook 同时 block 时全部保留） + pub blocking_errors: Vec, + /// 最终的工具参数（应用了最后一个 MutateInput 的修改） + pub final_args: serde_json::Value, + /// 带来源标记的上下文列表（Phase 1.3+：用于格式化有意义的 system-reminder） + pub tagged_contexts: Vec, +} + +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, + /// 带来源标记的上下文列表（Phase 1.3+：用于格式化有意义的 system-reminder） + pub tagged_contexts: Vec, + /// 非阻塞警告消息列表（Phase 3.1+） + pub warnings: Vec, + /// 事后权限请求列表（post-hoc advisory，不阻塞主循环） + pub post_permission_requests: Vec<(String, String)>, + /// 结构化键值对元数据（Phase 3.1+） + pub metadata: std::collections::HashMap, +} diff --git a/src/agent/memory/mod.rs b/src/agent/memory/mod.rs index a9922a8..f2b2602 100644 --- a/src/agent/memory/mod.rs +++ b/src/agent/memory/mod.rs @@ -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 { 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 => "[反馈]", diff --git a/src/agent/runtime/checkpoint.rs b/src/agent/runtime/checkpoint.rs new file mode 100644 index 0000000..1460ca3 --- /dev/null +++ b/src/agent/runtime/checkpoint.rs @@ -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/ ← 每个工作目录的分支 + +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/` 下有独立的分支。 +/// +/// 注意：`git2::Repository` 是 `!Sync`，所以内部状态通过 `Mutex` 保护。 +pub struct CheckpointManager { + /// 主开关 + enabled: bool, + /// bare git 仓库路径 + repo_path: PathBuf, + /// 每个项目保留的最大快照数 + max_snapshots: usize, + /// 单个文件大小上限（字节），超过此大小不纳入快照 + max_file_size: usize, + /// 内部可变状态（Send + Sync） + inner: Mutex, +} + +/// 需要 Mutex 保护的可变状态 +struct InnerState { + /// git2 仓库句柄 + repo: Option, + /// 本 turn 已快照的目录集合 + checkpointed_this_turn: HashSet, +} + +impl InnerState { + fn new() -> Self { + InnerState { + repo: None, + checkpointed_this_turn: HashSet::new(), + } + } +} + +impl CheckpointManager { + /// 创建新的 checkpoint 管理器。 + /// + /// `store_path` 是 bare repo 的路径（建议: `/../.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::().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 { + // 创建父目录 + 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 { + 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 { + 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(""); + 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 { + 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 { + 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 { + 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 = 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 = commits.iter().take(keep).copied().collect(); + + // 重建线性链：从 keep_index 开始 + let keep_start = keep; + let keep_commits: Vec> = 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 = 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/` +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); + } +} diff --git a/src/agent/runtime/error_recovery.rs b/src/agent/runtime/error_recovery.rs index 5b6396f..f4add02 100644 --- a/src/agent/runtime/error_recovery.rs +++ b/src/agent/runtime/error_recovery.rs @@ -1,42 +1,609 @@ // src/agent/runtime/error_recovery.rs // -// 错误恢复阶梯。 -// 参考 Claude Code error recovery ladder 设计。 +// 错误分类与故障转移系统。 +// 参考 Hermes-Agent error_classifier.py FailoverReason 设计。 // -// 当 LLM 流返回可恢复的错误（如 prompt_too_long）时， -// 按阶梯顺序尝试恢复： -// 1. Aggressive Compact — 激进微压缩（保留更少的工具结果） -// 2. Reactive Compact — 使用 LLM 摘要压缩对话历史 -// 3. Escalate Tokens — 临时提升 token 上限到 64k -// 4. Multi-Turn — 注入 metacognitive 消息分步处理 -// 5. Surface — 放弃恢复，暴露错误给用户 -// -// 每一步都有 `has_attempted` 守卫，防止无限循环。 +// 提供 21 种错误分类 + 8 步分类管线 + 对应恢复策略。 +// 分类管线按优先级排序，确保 provider-specific 模式优先匹配。 -use tracing::{info, warn}; +use tracing::{error, info, warn}; use super::token_budget::TokenBudget; -/// 错误类型分类 +// ── FailoverReason：21 种错误分类 ── + +/// 故障转移原因 — 参考 Hermes FailoverReason 枚举。 +/// 每个 variant 携带是否可重试、是否应触发凭据切换等信息。 +#[derive(Debug, Clone, PartialEq)] +pub enum FailoverReason { + // ── Auth 类 ── + /// 认证失败（401），可切换凭据重试 + Auth, + /// 永久认证失败，不应重试 + AuthPermanent, + + // ── Billing 类 ── + /// 计费问题（402/403 billing），可切换账号 + Billing, + /// 限流（429），应退避重试 + RateLimited, + + // ── Server 类 ── + /// 服务过载（503/529），应退避重试 + Overloaded, + /// 服务器错误（500/502），可重试 + ServerError, + + // ── Transport 类 ── + /// 网络超时，可重试 + Timeout, + + // ── Context/Payload 类 ── + /// 上下文溢出（prompt too long / 413） + ContextOverflow, + /// 载荷过大 + PayloadTooLarge, + /// 图片过大 + ImageTooLarge, + + // ── Model/Provider Policy 类 ── + /// 模型不存在或无权访问（404） + ModelNotFound, + /// Provider 策略阻止 + ProviderPolicyBlocked, + /// 内容策略阻止（安全过滤） + ContentPolicyBlocked, + + // ── Format 类 ── + /// 响应格式错误 + FormatError, + /// 无效加密内容 + InvalidEncryptedContent, + /// 多模态工具内容不支持 + MultimodalToolContentUnsupported, + + // ── Provider-specific 类 ── + /// Thinking 签名错误（Qwen/DashScope specific） + ThinkingSignature, + /// 长上下文 tier 门控 + LongContextTier, + /// OAuth 长上下文 beta 禁止 + OauthLongContextBetaForbidden, + + // ── 网络/传输 ── + /// SSL/TLS 瞬态错误 + SslTransient, + + // ── Catch-all ── + /// 未知错误 + Unknown(String), +} + +impl FailoverReason { + /// 是否可重试（auth/rate-limit/server/transport 类） + pub fn is_retryable(&self) -> bool { + matches!( + self, + FailoverReason::Auth + | FailoverReason::RateLimited + | FailoverReason::Overloaded + | FailoverReason::ServerError + | FailoverReason::Timeout + | FailoverReason::SslTransient + ) + } + + /// 是否应触发上下文压缩 + pub fn should_compress(&self) -> bool { + matches!( + self, + FailoverReason::ContextOverflow + | FailoverReason::PayloadTooLarge + | FailoverReason::LongContextTier + ) + } + + /// 是否应切换凭据/模型 + pub fn should_failover(&self) -> bool { + matches!( + self, + FailoverReason::Auth + | FailoverReason::Billing + | FailoverReason::ModelNotFound + | FailoverReason::ProviderPolicyBlocked + ) + } + + /// 是否为永久性错误（不应重试） + pub fn is_permanent(&self) -> bool { + matches!( + self, + FailoverReason::AuthPermanent + | FailoverReason::ContentPolicyBlocked + | FailoverReason::InvalidEncryptedContent + ) + } + + /// 错误归因：用户侧还是服务侧 + pub fn is_user_error(&self) -> bool { + matches!( + self, + FailoverReason::ContextOverflow + | FailoverReason::PayloadTooLarge + | FailoverReason::ImageTooLarge + | FailoverReason::FormatError + ) + } +} + +impl std::fmt::Display for FailoverReason { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let label = match self { + FailoverReason::Auth => "auth", + FailoverReason::AuthPermanent => "auth_permanent", + FailoverReason::Billing => "billing", + FailoverReason::RateLimited => "rate_limited", + FailoverReason::Overloaded => "overloaded", + FailoverReason::ServerError => "server_error", + FailoverReason::Timeout => "timeout", + FailoverReason::ContextOverflow => "context_overflow", + FailoverReason::PayloadTooLarge => "payload_too_large", + FailoverReason::ImageTooLarge => "image_too_large", + FailoverReason::ModelNotFound => "model_not_found", + FailoverReason::ProviderPolicyBlocked => "provider_policy_blocked", + FailoverReason::ContentPolicyBlocked => "content_policy_blocked", + FailoverReason::FormatError => "format_error", + FailoverReason::InvalidEncryptedContent => "invalid_encrypted_content", + FailoverReason::MultimodalToolContentUnsupported => { + "multimodal_tool_content_unsupported" + } + FailoverReason::ThinkingSignature => "thinking_signature", + FailoverReason::LongContextTier => "long_context_tier", + FailoverReason::OauthLongContextBetaForbidden => "oauth_long_context_beta_forbidden", + FailoverReason::SslTransient => "ssl_transient", + FailoverReason::Unknown(s) => return write!(f, "unknown({})", s), + }; + write!(f, "{}", label) + } +} + +// ── 保留向后兼容的 ErrorKind 别名 ── + +/// 错误类型分类（向后兼容简化版，新代码应使用 FailoverReason） #[derive(Debug, Clone, PartialEq)] pub enum ErrorKind { - /// 上下文过长（prompt too long / 413） PromptTooLong, - /// Token 耗尽 TokenExhausted, - /// 模型错误 ModelError(String), - /// 超时 Timeout, - /// 限流（HTTP 429） RateLimited, - /// 服务过载（HTTP 529） Overloaded, } +impl From<&FailoverReason> for ErrorKind { + fn from(reason: &FailoverReason) -> Self { + match reason { + FailoverReason::ContextOverflow | FailoverReason::PayloadTooLarge => { + ErrorKind::PromptTooLong + } + FailoverReason::LongContextTier => ErrorKind::TokenExhausted, + FailoverReason::RateLimited => ErrorKind::RateLimited, + FailoverReason::Overloaded | FailoverReason::ServerError => ErrorKind::Overloaded, + FailoverReason::Timeout | FailoverReason::SslTransient => ErrorKind::Timeout, + other => ErrorKind::ModelError(other.to_string()), + } + } +} + +// ── 8 步分类管线 ── + +/// Step 1: Provider-specific 模式匹配（最高优先级） +fn classify_provider_specific(error_str: &str) -> Option { + let lower = error_str.to_lowercase(); + + // Content policy blocks (非 provider 特定，但在 status 分类前检查) + if lower.contains("content_policy_violation") + || lower.contains("safety filter") + || lower.contains("content filter") + || lower.contains("responsible_ai") + || lower.contains("content management policy") + { + return Some(FailoverReason::ContentPolicyBlocked); + } + + // Thinking signature error (Qwen/DashScope) + if lower.contains("thinking_signature") || lower.contains("thinking signature") { + return Some(FailoverReason::ThinkingSignature); + } + + // Long context tier gate + if lower.contains("long_context_tier") + || lower.contains("long context tier") + || lower.contains("context_length_exceeded_for_tier") + { + return Some(FailoverReason::LongContextTier); + } + + // OAuth 1M beta forbidden + if lower.contains("oauth_long_context") || lower.contains("1m_beta_forbidden") { + return Some(FailoverReason::OauthLongContextBetaForbidden); + } + + // Llama.cpp grammar error + if lower.contains("grammar_pattern") || lower.contains("llama_cpp_grammar") { + return Some(FailoverReason::FormatError); + } + + // Multimodal tool content + if lower.contains("multimodal_tool_content") || lower.contains("tool content not supported") { + return Some(FailoverReason::MultimodalToolContentUnsupported); + } + + None +} + +/// Step 2: HTTP Status Code 分类 +fn classify_by_http_status(error_str: &str) -> Option { + let lower = error_str.to_lowercase(); + + // 401 → auth (temporary) + if lower.contains("401") || lower.contains("unauthorized") { + return Some(FailoverReason::Auth); + } + + // 402 → billing vs rate-limit disambiguation + if lower.contains("402") { + if lower.contains("rate") || lower.contains("limit") || lower.contains("try again") { + return Some(FailoverReason::RateLimited); + } + return Some(FailoverReason::Billing); + } + + // 403 + if lower.contains("403") || lower.contains("forbidden") { + if lower.contains("billing") || lower.contains("quota") || lower.contains("insufficient") { + return Some(FailoverReason::Billing); + } + if lower.contains("policy") || lower.contains("blocked") { + return Some(FailoverReason::ProviderPolicyBlocked); + } + return Some(FailoverReason::AuthPermanent); + } + + // 404 + if lower.contains("404") || lower.contains("not found") { + if lower.contains("model") || lower.contains("deployment") { + return Some(FailoverReason::ModelNotFound); + } + if lower.contains("billing") || lower.contains("subscription") { + return Some(FailoverReason::Billing); + } + if lower.contains("policy") || lower.contains("blocked") { + return Some(FailoverReason::ProviderPolicyBlocked); + } + } + + // 413 → context overflow (default), could also be payload_too_large + if lower.contains("413") { + if lower.contains("image") || lower.contains("media") { + return Some(FailoverReason::ImageTooLarge); + } + return Some(FailoverReason::ContextOverflow); + } + + // 429 → rate limit (definitive) + if lower.contains("429") { + return Some(FailoverReason::RateLimited); + } + + // 500/502 → server error + if lower.contains("\"500\"") + || lower.contains(" 500 ") + || lower.contains("502") + || lower.contains("internal server error") + || lower.contains("bad gateway") + { + // Check for request validation signals + if lower.contains("invalid_request") + || lower.contains("unsupported_parameter") + || lower.contains("unknown_parameter") + { + return Some(FailoverReason::FormatError); + } + return Some(FailoverReason::ServerError); + } + + // 503/529 → overloaded + if lower.contains("503") || lower.contains("529") || lower.contains("service unavailable") { + return Some(FailoverReason::Overloaded); + } + + None +} + +/// Step 3: Structured Error Code 分类 +fn classify_by_error_code(error_str: &str) -> Option { + let lower = error_str.to_lowercase(); + + // Structured error codes from API responses + if lower.contains("\"code\"") || lower.contains("\"error_type\"") { + if lower.contains("resource_exhausted") { + return Some(FailoverReason::RateLimited); + } + if lower.contains("insufficient_quota") { + return Some(FailoverReason::Billing); + } + if lower.contains("model_not_found") || lower.contains("deployment_not_found") { + return Some(FailoverReason::ModelNotFound); + } + if lower.contains("invalid_api_key") || lower.contains("token_expired") { + return Some(FailoverReason::Auth); + } + if lower.contains("context_length_exceeded") { + return Some(FailoverReason::ContextOverflow); + } + } + + None +} + +/// Step 4: Message Pattern 匹配（14 组模式集） +fn classify_by_message_pattern(error_str: &str) -> Option { + let lower = error_str.to_lowercase(); + + // 4a. Billing patterns + if lower.contains("insufficient_quota") + || lower.contains("billing account") + || lower.contains("payment required") + || lower.contains("quota exceeded") + || lower.contains("run out of") + || lower.contains("no credit") + || lower.contains("check your balance") + { + return Some(FailoverReason::Billing); + } + + // 4b. Rate-limit patterns + if lower.contains("rate limit") + || lower.contains("rate_limit") + || lower.contains("too many requests") + || lower.contains("requests too frequent") + || lower.contains("try again in") + || lower.contains("resets at") + || lower.contains("please slow down") + { + return Some(FailoverReason::RateLimited); + } + + // 4c. Context overflow patterns + if lower.contains("prompt_too_long") + || lower.contains("prompt too long") + || lower.contains("context length") + || lower.contains("context_length_exceeded") + || lower.contains("input length") + || lower.contains("maximum context") + || lower.contains("reduce the length") + || lower.contains("too many tokens") + || lower.contains("max_tokens_exceeded") + { + return Some(FailoverReason::ContextOverflow); + } + + // 4d. Auth patterns + if lower.contains("invalid api key") + || lower.contains("invalid_api_key") + || lower.contains("incorrect api key") + || lower.contains("authentication failed") + || lower.contains("auth error") + || lower.contains("not authenticated") + || lower.contains("token is invalid") + { + return Some(FailoverReason::Auth); + } + + // 4e. Payload size + if lower.contains("request too large") + || lower.contains("payload too large") + || lower.contains("request size exceeds") + || lower.contains("request entity too large") + { + return Some(FailoverReason::PayloadTooLarge); + } + + // 4f. Image size + if lower.contains("image too large") + || lower.contains("image size exceeds") + || lower.contains("max image size") + || lower.contains("image_resolution_exceeded") + { + return Some(FailoverReason::ImageTooLarge); + } + + // 4g. Multimodal tool content + if lower.contains("tool content not supported") || lower.contains("multimodal content in tool") + { + return Some(FailoverReason::MultimodalToolContentUnsupported); + } + + // 4h. Timeout patterns + if lower.contains("request timed out") + || lower.contains("timed out") + || lower.contains("deadline exceeded") + || lower.contains("idle timeout") + || lower.contains("connection timeout") + { + return Some(FailoverReason::Timeout); + } + + // 4i. Invalid encrypted content + if lower.contains("invalid encrypted content") || lower.contains("encrypted_content_error") { + return Some(FailoverReason::InvalidEncryptedContent); + } + + // 4j. Format error + if lower.contains("invalid_request") + && (lower.contains("parameter") || lower.contains("format") || lower.contains("schema")) + { + return Some(FailoverReason::FormatError); + } + + None +} + +/// Step 5: SSL/TLS 瞬态错误检测 +fn classify_ssl_transient(error_str: &str) -> Option { + let lower = error_str.to_lowercase(); + + if lower.contains("ssl") || lower.contains("tls") || lower.contains("certificate") { + // Temporary TLS errors (connection reset, handshake failure) + if lower.contains("reset") || lower.contains("handshake") || lower.contains("timeout") { + return Some(FailoverReason::SslTransient); + } + // Permanent TLS errors (certificate expired, hostname mismatch) + if lower.contains("expired") || lower.contains("hostname") || lower.contains("verify") { + return Some(FailoverReason::AuthPermanent); + } + return Some(FailoverReason::SslTransient); + } + + None +} + +/// Step 6: Server disconnect + large session → context overflow +fn classify_disconnect_with_large_session( + error_str: &str, + estimated_tokens: Option, + context_length: Option, +) -> Option { + let lower = error_str.to_lowercase(); + + // Server disconnected without response + if lower.contains("connection reset") + || lower.contains("eof") + || lower.contains("incomplete response") + || lower.contains("peer closed") + { + // If session is large (>60% of context), blame context overflow + if let (Some(tokens), Some(ctx_len)) = (estimated_tokens, context_length) { + let ratio = tokens as f64 / ctx_len as f64; + if ratio > 0.6 { + return Some(FailoverReason::ContextOverflow); + } + } + return Some(FailoverReason::Timeout); + } + + None +} + +/// Step 7: Transport error heuristics +fn classify_transport_error(error_str: &str) -> Option { + let lower = error_str.to_lowercase(); + + if lower.contains("connection") + || lower.contains("network") + || lower.contains("dns") + || lower.contains("unreachable") + || lower.contains("refused") + || lower.contains("broken pipe") + || lower.contains("transport") + || lower.contains("eof") + || lower.contains("reset by peer") + { + return Some(FailoverReason::Timeout); + } + + None +} + +/// Step 8: Unknown fallback +fn classify_unknown(error_str: &str) -> FailoverReason { + FailoverReason::Unknown(error_str.to_string()) +} + +// ── 主分类函数 ── + +/// 8步分类管线：按优先级依次尝试各分类器。 +/// +/// 可选传入 `estimated_tokens` 和 `context_length` 用于 +/// disconnect + large session 的上下文溢出检测。 +pub fn classify_failover( + error_str: &str, + estimated_tokens: Option, + context_length: Option, +) -> FailoverReason { + // Step 1: Provider-specific (HIGHEST priority) + if let Some(reason) = classify_provider_specific(error_str) { + info!("[ErrorClassify] Step 1 (provider-specific): {:?}", reason); + return reason; + } + + // Step 2: HTTP status code + if let Some(reason) = classify_by_http_status(error_str) { + info!("[ErrorClassify] Step 2 (HTTP status): {:?}", reason); + return reason; + } + + // Step 3: Structured error code + if let Some(reason) = classify_by_error_code(error_str) { + info!("[ErrorClassify] Step 3 (error code): {:?}", reason); + return reason; + } + + // Step 4: Message patterns (14 groups) + if let Some(reason) = classify_by_message_pattern(error_str) { + info!("[ErrorClassify] Step 4 (message pattern): {:?}", reason); + return reason; + } + + // Step 5: SSL/TLS transient + if let Some(reason) = classify_ssl_transient(error_str) { + info!("[ErrorClassify] Step 5 (SSL/TLS): {:?}", reason); + return reason; + } + + // Step 6: Server disconnect + large session + if let Some(reason) = + classify_disconnect_with_large_session(error_str, estimated_tokens, context_length) + { + info!( + "[ErrorClassify] Step 6 (disconnect + large session): {:?}", + reason + ); + return reason; + } + + // Step 7: Transport heuristics + if let Some(reason) = classify_transport_error(error_str) { + info!("[ErrorClassify] Step 7 (transport): {:?}", reason); + return reason; + } + + // Step 8: Unknown fallback + error!( + "[ErrorClassify] Step 8 (unknown): unclassified error: {}", + error_str + ); + classify_unknown(error_str) +} + +// ── 向后兼容的 classify_error ── + +/// 从错误字符串分类错误类型。 +/// 内部调用 classify_failover 并转换为简化的 ErrorKind。 +pub fn classify_error(error_str: &str) -> ErrorKind { + let reason = classify_failover(error_str, None, None); + ErrorKind::from(&reason) +} + +// ── 恢复策略 ── + /// 恢复步骤 #[derive(Debug, Clone, PartialEq)] pub enum RecoveryStep { + /// 调整 max_tokens 以适配上下文容量（ContextOverflow 首选策略）。 + /// 参考 Claude Code — 解析 "X + Y > Z" 错误后计算安全值。 + AdjustMaxTokens { new_max_tokens: usize }, /// 尝试更激进的 micro_compact AggressiveCompact, /// 使用 LLM 摘要压缩 @@ -78,9 +645,19 @@ impl RecoveryAttempts { || !self.multi_turn } - /// 获取下一个应尝试的恢复步骤 - pub fn next_step(&mut self, error_kind: &ErrorKind) -> Option { - // 429/529 使用退避重试，不消耗上下文恢复步骤 + /// 获取下一个应尝试的恢复步骤。 + /// + /// 对于 ContextOverflow 类错误，恢复优先级为： + /// 1. AdjustMaxTokens — 从错误消息解析容量，精确计算安全的 max_tokens（参考 Claude Code） + /// 2. AggressiveCompact — 激进微压缩 + /// 3. ReactiveCompact — LLM 摘要压缩 + /// 4. EscalateTokens — 提升 hard_limit + /// 5. MultiTurn — 分轮 + pub fn next_step( + &mut self, + error_kind: &ErrorKind, + overflow_info: Option<&ContextOverflowInfo>, + ) -> Option { if matches!(error_kind, ErrorKind::RateLimited | ErrorKind::Overloaded) { return Some(RecoveryStep::RetryWithBackoff { attempt: 0, @@ -90,27 +667,44 @@ impl RecoveryAttempts { match error_kind { ErrorKind::PromptTooLong | ErrorKind::TokenExhausted => { + // Step 0: 尝试精确调整 max_tokens（仅在有 Overflow info 时） + if let Some(info) = overflow_info { + if let Some(safe_tokens) = calculate_safe_max_tokens(info) { + // 仅在使用默认路径（尚未尝试过其他恢复方式）时触发 + if !self.aggressive_compact + && !self.reactive_compact + && !self.escalate_tokens + { + return Some(RecoveryStep::AdjustMaxTokens { + new_max_tokens: safe_tokens, + }); + } + } + } + // Step 1: 激进微压缩 if !self.aggressive_compact { self.aggressive_compact = true; return Some(RecoveryStep::AggressiveCompact); } + // Step 2: LLM 摘要压缩 if !self.reactive_compact { self.reactive_compact = true; return Some(RecoveryStep::ReactiveCompact); } + // Step 3: 提升硬限制 if !self.escalate_tokens { self.escalate_tokens = true; return Some(RecoveryStep::EscalateTokens { new_hard_limit: 64_000, }); } + // Step 4: 分轮 if !self.multi_turn { self.multi_turn = true; return Some(RecoveryStep::MultiTurn); } } _ => { - // 非上下文相关错误，直接暴露 if !self.multi_turn { self.multi_turn = true; return Some(RecoveryStep::Surface); @@ -119,6 +713,11 @@ impl RecoveryAttempts { } None } + + /// 向后兼容的无 overflow_info 版本。 + pub fn next_step_compat(&mut self, error_kind: &ErrorKind) -> Option { + self.next_step(error_kind, None) + } } impl Default for RecoveryAttempts { @@ -129,14 +728,11 @@ impl Default for RecoveryAttempts { /// 错误恢复器 pub struct ErrorRecovery { - /// 恢复步骤追踪 pub attempts: RecoveryAttempts, - /// Token 预算（用于 escalate 步骤） pub token_budget: TokenBudget, } impl ErrorRecovery { - /// 创建新的错误恢复器 pub fn new(token_budget: TokenBudget) -> Self { ErrorRecovery { attempts: RecoveryAttempts::new(), @@ -144,22 +740,30 @@ impl ErrorRecovery { } } - /// 尝试从错误中恢复。 - /// - /// 返回 `Some(RecoveryStep)` 表示找到了恢复步骤（调用方应执行该步骤后重试）。 - /// 返回 `None` 表示所有步骤已尝试完毕，应暴露错误给用户。 - pub fn try_recover(&mut self, error_kind: &ErrorKind) -> Option { - let step = self.attempts.next_step(error_kind); + pub fn try_recover( + &mut self, + error_kind: &ErrorKind, + overflow_info: Option<&ContextOverflowInfo>, + ) -> Option { + let step = self.attempts.next_step(error_kind, overflow_info); match &step { + Some(RecoveryStep::AdjustMaxTokens { new_max_tokens }) => { + info!( + "[ErrorRecovery] 尝试步骤 0/5: AdjustMaxTokens → {} (从错误消息计算)", + new_max_tokens + ); + // 将硬限制下调至安全值，使下一次 LLM 调用能在容量内完成 + self.token_budget.hard_limit = *new_max_tokens; + } Some(RecoveryStep::AggressiveCompact) => { - info!("[ErrorRecovery] 尝试步骤 1/4: AggressiveCompact"); + info!("[ErrorRecovery] 尝试步骤 1/5: AggressiveCompact"); } Some(RecoveryStep::ReactiveCompact) => { - info!("[ErrorRecovery] 尝试步骤 2/4: ReactiveCompact"); + info!("[ErrorRecovery] 尝试步骤 2/5: ReactiveCompact"); } Some(RecoveryStep::EscalateTokens { new_hard_limit }) => { info!( - "[ErrorRecovery] 尝试步骤 3/4: EscalateTokens → {}", + "[ErrorRecovery] 尝试步骤 3/5: EscalateTokens → {}", new_hard_limit ); self.token_budget.escalate_hard_limit(*new_hard_limit); @@ -171,7 +775,7 @@ impl ErrorRecovery { ); } Some(RecoveryStep::MultiTurn) => { - info!("[ErrorRecovery] 尝试步骤 4/4: MultiTurn"); + info!("[ErrorRecovery] 尝试步骤 4/5: MultiTurn"); } Some(RecoveryStep::Surface) => { warn!("[ErrorRecovery] 无法恢复，暴露错误"); @@ -183,7 +787,11 @@ impl ErrorRecovery { step } - /// 生成 multi-turn 恢复消息（注入到对话中以继续处理） + /// 向后兼容的无 overflow_info 版本。 + pub fn try_recover_compat(&mut self, error_kind: &ErrorKind) -> Option { + self.try_recover(error_kind, None) + } + pub fn multi_turn_message() -> String { "由于 token 限制，当前回答被截断。请基于已收集的信息继续分析，\ 重点关注尚未完成的部分。你可以：\n\ @@ -193,7 +801,6 @@ impl ErrorRecovery { .to_string() } - /// 检查是否需要恢复（错误是否可恢复） pub fn is_recoverable(error_kind: &ErrorKind) -> bool { matches!( error_kind, @@ -205,126 +812,432 @@ impl ErrorRecovery { } } -/// 计算指数退避延迟（毫秒）。 -/// -/// 公式：min(500 * 2^attempt, 32000) + 25% 随机抖动 -/// 如果有 Retry-After header，优先使用。 +/// 计算指数退避延迟（毫秒） pub fn backoff_delay(attempt: u32, retry_after_secs: Option) -> u64 { if let Some(ra) = retry_after_secs { return ra * 1000; } - let base = 500u64 * 2u64.pow(attempt.min(6)); // cap at 2^6 = 64 → 32000ms + let base = 500u64 * 2u64.pow(attempt.min(6)); let base = base.min(32_000); - // Simple deterministic jitter using attempt (avoid rand dependency) let jitter = (base / 4) * (attempt as u64 % 5) / 5; base + jitter } -/// 解析错误字符串中的 Retry-After header 值。 +// ── Context Overflow 自动修复 ── +// +// 参考 Claude Code: src/services/api/withRetry.ts +// 解析 Anthropic/OpenAI API 的上下文溢出错误消息， +// 自动计算安全的 max_tokens 值以进行重试。 + +/// 从 ContextOverflow 错误消息解析出的数值。 +#[derive(Debug, Clone, Copy)] +pub struct ContextOverflowInfo { + /// 当前输入 token 数 + pub input_length: usize, + /// 请求的 max_tokens 数 + pub requested_max_tokens: usize, + /// 模型上下文窗口总容量 + pub context_limit: usize, +} + +/// 上下文溢出检测的最小可接受 token 数。 +/// 低于此值说明即使调整 max_tokens 也没有意义，应触发压缩。 +const MIN_ACCEPTABLE_TOKENS: usize = 500; + +/// 上下文溢出调整的安全边际（token 数）。 +/// 保留少量余量以防止重试时再次溢出。 +const OVERFLOW_SAFETY_MARGIN: usize = 1000; + +/// 从错误消息中解析 Anthropic-style context overflow 信息。 /// -/// 期望格式: `retry_after=Some(N)` 出现在错误消息中。 +/// 匹配模式: +/// "input length and max_tokens exceed context limit: 180000 + 32000 > 200000" +/// "prompt too long: X tokens + Y max_tokens > Z context_window" +/// "context length exceeded: X + Y > Z" +/// "maximum context length is Z tokens, but X + Y = ..." +pub fn parse_context_overflow(error_str: &str) -> Option { + let lower = error_str.to_lowercase(); + + // Pattern 1: Anthropic standard format + // "input length and max_tokens exceed context limit: A + B > C" + if let Some(info) = parse_anthropic_overflow(&lower) { + return Some(info); + } + + // Pattern 2: OpenAI format + // "maximum context length is C tokens. however X tokens ... and Y max_tokens" + if let Some(info) = parse_openai_overflow(&lower) { + return Some(info); + } + + // Pattern 3: Generic "X + Y > Z" pattern + if let Some(info) = parse_generic_overflow(&lower) { + return Some(info); + } + + None +} + +/// 解析 Anthropic API 格式: "input length and max_tokens exceed context limit: A + B > C" +fn parse_anthropic_overflow(error_str: &str) -> Option { + // Find the "A + B > C" pattern after "context limit" + let marker = "input length and max_tokens exceed context limit"; + let pos = error_str.find(marker)?; + let after_marker = &error_str[pos + marker.len()..]; + + // Extract numbers: look for pattern like " 180000 + 32000 > 200000" + let numbers = extract_three_numbers(after_marker)?; + Some(ContextOverflowInfo { + input_length: numbers.0, + requested_max_tokens: numbers.1, + context_limit: numbers.2, + }) +} + +/// 解析 OpenAI API 格式: "maximum context length is C tokens..." +fn parse_openai_overflow(error_str: &str) -> Option { + // "maximum context length is C tokens" + let context_re = regex::Regex::new(r"maximum context length is (\d+)").ok()?; + let caps = context_re.captures(error_str)?; + let context_limit: usize = caps.get(1)?.as_str().parse().ok()?; + + // "you requested X tokens" or "messages resulted in X tokens" + let input_re = regex::Regex::new(r"(?:requested|resulted in|totals? to) (\d+) tokens").ok()?; + let input_length: usize = if let Some(caps) = input_re.captures(error_str) { + caps.get(1)?.as_str().parse().ok()? + } else { + // Fallback: find any substantial number and use it as input_length + return None; + }; + + // max_tokens might not be explicitly stated; use the context_limit - input_length - margin + let requested_max_tokens = context_limit.saturating_sub(input_length); + + Some(ContextOverflowInfo { + input_length, + requested_max_tokens, + context_limit, + }) +} + +/// 解析通用格式: "X + Y > Z" +fn parse_generic_overflow(error_str: &str) -> Option { + // Look for patterns like "130000 + 32000 > 200000" + let (a, b, c) = extract_three_numbers(error_str)?; + Some(ContextOverflowInfo { + input_length: a, + requested_max_tokens: b, + context_limit: c, + }) +} + +/// 从字符串中提取三个连续的数字，匹配 "A + B > C" 模式。 +fn extract_three_numbers(s: &str) -> Option<(usize, usize, usize)> { + let re = regex::Regex::new(r"(\d{4,})\s*\+\s*(\d+)\s*>\s*(\d{4,})").ok()?; + let caps = re.captures(s)?; + let a: usize = caps.get(1)?.as_str().parse().ok()?; + let b: usize = caps.get(2)?.as_str().parse().ok()?; + let c: usize = caps.get(3)?.as_str().parse().ok()?; + Some((a, b, c)) +} + +/// 根据 ContextOverflow 信息计算安全的 max_tokens 值。 +/// +/// 公式: `new_max_tokens = context_limit - input_length - SAFETY_MARGIN` +/// +/// 返回 `None` 如果计算出的值低于最小可接受阈值（说明必须压缩上下文）。 +pub fn calculate_safe_max_tokens(info: &ContextOverflowInfo) -> Option { + // 如果输入本身已经超过上下文窗口，max_tokens 调整无意义 + if info.input_length >= info.context_limit { + return None; + } + + let available = info.context_limit.saturating_sub(info.input_length); + let adjusted = available.saturating_sub(OVERFLOW_SAFETY_MARGIN); + + if adjusted >= MIN_ACCEPTABLE_TOKENS { + Some(adjusted) + } else { + // 可用空间太小，调整 max_tokens 无意义 → 应触发压缩 + None + } +} + +// ── 向后兼容 ── + +/// 解析 Retry-After header 值 pub fn parse_retry_after(error_str: &str) -> Option { if let Some(pos) = error_str.find("retry_after=Some(") { - let prefix_len = "retry_after=Some(".len(); // 19 + let prefix_len = "retry_after=Some(".len(); let rest = &error_str[pos + prefix_len..]; if let Some(end) = rest.find(')') { return rest[..end].parse().ok(); } } + if let Some(pos) = error_str.find("Retry-After:") { + let rest = &error_str[pos + "Retry-After:".len()..]; + if let Some(end) = rest.find('\r').or_else(|| rest.find('\n')) { + return rest[..end].trim().parse().ok(); + } + return rest.trim().parse().ok(); + } None } -/// 从错误字符串分类错误类型。 -/// 解析 LLM API 返回的错误消息，映射到 ErrorKind。 -pub fn classify_error(error_str: &str) -> ErrorKind { - let lower = error_str.to_lowercase(); +// ── 恢复建议生成 ── - // 先检测限流/过载（HTTP 状态码检查） - if lower.contains("429") - || lower.contains("rate limit") - || lower.contains("rate_limit") - || lower.contains("too many requests") - { - return ErrorKind::RateLimited; +/// 根据 FailoverReason 生成面向用户的恢复建议。 +pub fn recovery_suggestion(reason: &FailoverReason) -> &'static str { + match reason { + FailoverReason::Auth => "认证失败，尝试切换 API Key 或刷新凭据。", + FailoverReason::AuthPermanent => "认证凭据永久失效，请检查 API Key 配置。", + FailoverReason::Billing => "账户余额不足或配额用尽，请检查计费状态。", + FailoverReason::RateLimited => "请求频率超限，系统正在自动退避重试。", + FailoverReason::Overloaded => "模型服务过载，系统正在自动重试。", + FailoverReason::ServerError => "服务器内部错误，将自动重试。", + FailoverReason::Timeout => "请求超时，将自动重试。请检查网络连接。", + FailoverReason::ContextOverflow => "上下文超过限制，正在自动压缩对话历史。", + FailoverReason::PayloadTooLarge => "请求载荷过大，请尝试减少工具调用数量。", + FailoverReason::ImageTooLarge => "图片过大，请使用更小的图片或降低分辨率。", + FailoverReason::ModelNotFound => "模型不可用，请检查模型名称或尝试切换模型。", + FailoverReason::ProviderPolicyBlocked => "Provider 策略阻止了此请求。", + FailoverReason::ContentPolicyBlocked => { + "内容被安全策略过滤，请修改请求内容后重试。此错误不可自动恢复。" + } + FailoverReason::FormatError => "请求格式错误，请检查工具参数格式。", + FailoverReason::InvalidEncryptedContent => "加密内容无效，请重新生成。", + FailoverReason::MultimodalToolContentUnsupported => { + "当前模型不支持工具结果中的多模态内容，请改用纯文本。" + } + FailoverReason::ThinkingSignature => "思考模式签名错误，请禁用 thinking 模式或切换模型。", + FailoverReason::LongContextTier => "超出当前 tier 的上下文长度限制，正在压缩或升级 tier。", + FailoverReason::OauthLongContextBetaForbidden => "长上下文 Beta 功能未对当前凭据开放。", + FailoverReason::SslTransient => "SSL/TLS 连接暂时中断，将自动重试。", + FailoverReason::Unknown(_) => "未知错误，系统将尝试恢复。", } - if lower.contains("529") - || lower.contains("overloaded") - || lower.contains("overload") - || lower.contains("service overloaded") - { - return ErrorKind::Overloaded; - } - - if lower.contains("prompt_too_long") - || lower.contains("prompt too long") - || lower.contains("context length") - || lower.contains("413") - || lower.contains("context_window_exceeded") - || lower.contains("input length") - { - return ErrorKind::PromptTooLong; - } - - if lower.contains("max_tokens") - || lower.contains("token limit") - || lower.contains("token_exhausted") - || lower.contains("maximum context length") - || lower.contains("reduce the length") - { - return ErrorKind::TokenExhausted; - } - - if lower.contains("timeout") - || lower.contains("timed out") - || lower.contains("deadline exceeded") - || lower.contains("408") - || lower.contains("504") - { - return ErrorKind::Timeout; - } - - // 默认归类为模型错误 - ErrorKind::ModelError(error_str.to_string()) } #[cfg(test)] mod tests { use super::*; + // ── FailoverReason tests ── + + #[test] + fn test_is_retryable() { + assert!(FailoverReason::Auth.is_retryable()); + assert!(FailoverReason::RateLimited.is_retryable()); + assert!(FailoverReason::Overloaded.is_retryable()); + assert!(!FailoverReason::AuthPermanent.is_retryable()); + assert!(!FailoverReason::ContentPolicyBlocked.is_retryable()); + } + + #[test] + fn test_should_compress() { + assert!(FailoverReason::ContextOverflow.should_compress()); + assert!(FailoverReason::PayloadTooLarge.should_compress()); + assert!(!FailoverReason::RateLimited.should_compress()); + } + + #[test] + fn test_should_failover() { + assert!(FailoverReason::Auth.should_failover()); + assert!(FailoverReason::Billing.should_failover()); + assert!(FailoverReason::ModelNotFound.should_failover()); + assert!(!FailoverReason::Timeout.should_failover()); + } + + #[test] + fn test_is_permanent() { + assert!(FailoverReason::AuthPermanent.is_permanent()); + assert!(FailoverReason::ContentPolicyBlocked.is_permanent()); + assert!(!FailoverReason::Auth.is_permanent()); + } + + #[test] + fn test_is_user_error() { + assert!(FailoverReason::ContextOverflow.is_user_error()); + assert!(FailoverReason::ImageTooLarge.is_user_error()); + assert!(!FailoverReason::Timeout.is_user_error()); + } + + #[test] + fn test_failover_reason_display() { + assert_eq!(FailoverReason::Auth.to_string(), "auth"); + assert_eq!(FailoverReason::RateLimited.to_string(), "rate_limited"); + assert_eq!( + FailoverReason::Unknown("test".into()).to_string(), + "unknown(test)" + ); + } + + // ── Classification pipeline tests ── + + #[test] + fn test_classify_rate_limited_429() { + let reason = classify_failover("HTTP 429: Too Many Requests", None, None); + assert_eq!(reason, FailoverReason::RateLimited); + } + + #[test] + fn test_classify_overloaded_529() { + let reason = classify_failover("HTTP 529: Service Overloaded", None, None); + assert_eq!(reason, FailoverReason::Overloaded); + } + + #[test] + fn test_classify_context_overflow() { + let reason = classify_failover( + "Error: prompt_too_long: context length exceeded", + None, + None, + ); + assert_eq!(reason, FailoverReason::ContextOverflow); + } + + #[test] + fn test_classify_auth_401() { + let reason = classify_failover("HTTP 401: Unauthorized - invalid api key", None, None); + assert_eq!(reason, FailoverReason::Auth); + } + + #[test] + fn test_classify_billing_402() { + let reason = classify_failover("HTTP 402: insufficient_quota", None, None); + assert_eq!(reason, FailoverReason::Billing); + } + + #[test] + fn test_classify_content_policy_blocked() { + let reason = classify_failover( + "content_policy_violation: responsible_ai filter triggered", + None, + None, + ); + assert_eq!(reason, FailoverReason::ContentPolicyBlocked); + } + + #[test] + fn test_classify_model_not_found() { + let reason = classify_failover( + "HTTP 404: model_not_found - deployment not found", + None, + None, + ); + assert_eq!(reason, FailoverReason::ModelNotFound); + } + + #[test] + fn test_classify_transport_timeout() { + let reason = classify_failover("Connection reset by peer", None, None); + // Connection reset with no session info → timeout + assert_eq!(reason, FailoverReason::Timeout); + } + + #[test] + fn test_classify_disconnect_with_large_session() { + let reason = classify_failover( + "connection reset", + Some(8000), // estimated tokens + Some(10000), // context length (80% → >60%) + ); + assert_eq!(reason, FailoverReason::ContextOverflow); + } + + #[test] + fn test_classify_disconnect_with_small_session() { + let reason = classify_failover( + "connection reset", + Some(4000), // estimated tokens + Some(10000), // context length (40% → <60%) + ); + assert_eq!(reason, FailoverReason::Timeout); + } + + #[test] + fn test_classify_priorities_content_policy_before_status() { + // content_policy_violation should be detected before 400 status + let reason = classify_failover( + "HTTP 400: content_policy_violation - content management policy", + None, + None, + ); + assert_eq!(reason, FailoverReason::ContentPolicyBlocked); + } + + #[test] + fn test_classify_unknown_fallback() { + let reason = classify_failover("Some weird error that doesn't match anything", None, None); + assert!(matches!(reason, FailoverReason::Unknown(_))); + } + + #[test] + fn test_classify_image_too_large() { + let reason = classify_failover("HTTP 413: image too large, max 20MB", None, None); + assert_eq!(reason, FailoverReason::ImageTooLarge); + } + + #[test] + fn test_classify_provider_policy_blocked() { + let reason = classify_failover("HTTP 403: provider policy blocked this region", None, None); + assert_eq!(reason, FailoverReason::ProviderPolicyBlocked); + } + + #[test] + fn test_classify_ssl_transient() { + let reason = classify_failover("SSL handshake timeout: connection reset", None, None); + assert_eq!(reason, FailoverReason::SslTransient); + } + + // ── Backward compatibility tests ── + + #[test] + fn test_classify_error_backward_compat() { + let kind = classify_error("HTTP 429: Too Many Requests"); + assert_eq!(kind, ErrorKind::RateLimited); + } + + #[test] + fn test_classify_error_529_backward_compat() { + let kind = classify_error("HTTP 529: Service Overloaded"); + assert_eq!(kind, ErrorKind::Overloaded); + } + + // ── Recovery step tests ── + #[test] fn test_all_steps_sequence() { let mut attempts = RecoveryAttempts::new(); assert_eq!( - attempts.next_step(&ErrorKind::PromptTooLong), + attempts.next_step_compat(&ErrorKind::PromptTooLong), Some(RecoveryStep::AggressiveCompact) ); assert_eq!( - attempts.next_step(&ErrorKind::PromptTooLong), + attempts.next_step_compat(&ErrorKind::PromptTooLong), Some(RecoveryStep::ReactiveCompact) ); assert_eq!( - attempts.next_step(&ErrorKind::PromptTooLong), + attempts.next_step_compat(&ErrorKind::PromptTooLong), Some(RecoveryStep::EscalateTokens { new_hard_limit: 64_000 }) ); assert_eq!( - attempts.next_step(&ErrorKind::PromptTooLong), + attempts.next_step_compat(&ErrorKind::PromptTooLong), Some(RecoveryStep::MultiTurn) ); - // 所有步骤已尝试 - assert_eq!(attempts.next_step(&ErrorKind::PromptTooLong), None); + assert_eq!(attempts.next_step_compat(&ErrorKind::PromptTooLong), None); } #[test] fn test_model_error_goes_straight_to_surface() { let mut attempts = RecoveryAttempts::new(); assert_eq!( - attempts.next_step(&ErrorKind::ModelError("test".into())), + attempts.next_step_compat(&ErrorKind::ModelError("test".into())), Some(RecoveryStep::Surface) ); assert_eq!( - attempts.next_step(&ErrorKind::ModelError("test".into())), + attempts.next_step_compat(&ErrorKind::ModelError("test".into())), None ); } @@ -333,10 +1246,8 @@ mod tests { fn test_has_remaining() { let mut attempts = RecoveryAttempts::new(); assert!(attempts.has_remaining()); - - // 消耗所有步骤 for _ in 0..4 { - attempts.next_step(&ErrorKind::PromptTooLong); + attempts.next_step_compat(&ErrorKind::PromptTooLong); } assert!(!attempts.has_remaining()); } @@ -345,36 +1256,26 @@ mod tests { fn test_is_recoverable() { assert!(ErrorRecovery::is_recoverable(&ErrorKind::PromptTooLong)); assert!(ErrorRecovery::is_recoverable(&ErrorKind::TokenExhausted)); + assert!(ErrorRecovery::is_recoverable(&ErrorKind::RateLimited)); + assert!(ErrorRecovery::is_recoverable(&ErrorKind::Overloaded)); assert!(!ErrorRecovery::is_recoverable(&ErrorKind::ModelError( "test".into() ))); assert!(!ErrorRecovery::is_recoverable(&ErrorKind::Timeout)); } - #[test] - fn test_classify_rate_limited_429() { - let kind = classify_error("HTTP 429: Too Many Requests"); - assert_eq!(kind, ErrorKind::RateLimited); - } - - #[test] - fn test_classify_overloaded_529() { - let kind = classify_error("HTTP 529: Service Overloaded"); - assert_eq!(kind, ErrorKind::Overloaded); - } - - #[test] - fn test_rate_limited_is_recoverable() { - assert!(ErrorRecovery::is_recoverable(&ErrorKind::RateLimited)); - assert!(ErrorRecovery::is_recoverable(&ErrorKind::Overloaded)); - } - #[test] fn test_parse_retry_after() { let err = "HTTP 429: retry_after=Some(30)"; assert_eq!(parse_retry_after(err), Some(30)); } + #[test] + fn test_parse_retry_after_header_format() { + let err = "HTTP 429\nRetry-After: 60\n"; + assert_eq!(parse_retry_after(err), Some(60)); + } + #[test] fn test_parse_retry_after_none() { let err = "HTTP 500: Internal Server Error"; @@ -383,19 +1284,15 @@ mod tests { #[test] fn test_backoff_delay() { - // Attempt 0: 500 + jitter let d0 = backoff_delay(0, None); assert!(d0 >= 500 && d0 <= 700); - // Attempt 3: 500*8=4000 + jitter let d3 = backoff_delay(3, None); assert!(d3 >= 4000 && d3 <= 5000); - // Capped at 32s let d10 = backoff_delay(10, None); assert!(d10 <= 40_000); - // Retry-After takes priority let d_ra = backoff_delay(0, Some(15)); assert_eq!(d_ra, 15000); } @@ -403,7 +1300,7 @@ mod tests { #[test] fn test_rate_limited_goes_to_retry() { let mut attempts = RecoveryAttempts::new(); - let step = attempts.next_step(&ErrorKind::RateLimited); + let step = attempts.next_step_compat(&ErrorKind::RateLimited); assert!(matches!(step, Some(RecoveryStep::RetryWithBackoff { .. }))); } @@ -413,9 +1310,13 @@ mod tests { let mut recovery = ErrorRecovery::new(budget); assert_eq!(recovery.token_budget.hard_limit, 40_000); - recovery.attempts.next_step(&ErrorKind::PromptTooLong); // aggressive - recovery.attempts.next_step(&ErrorKind::PromptTooLong); // reactive - let step = recovery.try_recover(&ErrorKind::PromptTooLong); // escalate + recovery + .attempts + .next_step_compat(&ErrorKind::PromptTooLong); // aggressive + recovery + .attempts + .next_step_compat(&ErrorKind::PromptTooLong); // reactive + let step = recovery.try_recover_compat(&ErrorKind::PromptTooLong); // escalate assert_eq!( step, @@ -425,4 +1326,174 @@ mod tests { ); assert_eq!(recovery.token_budget.hard_limit, 64_000); } + + #[test] + fn test_recovery_suggestion_all_variants() { + // Every variant should return a non-empty suggestion + let variants = [ + FailoverReason::Auth, + FailoverReason::AuthPermanent, + FailoverReason::Billing, + FailoverReason::RateLimited, + FailoverReason::Overloaded, + FailoverReason::ServerError, + FailoverReason::Timeout, + FailoverReason::ContextOverflow, + FailoverReason::PayloadTooLarge, + FailoverReason::ImageTooLarge, + FailoverReason::ModelNotFound, + FailoverReason::ProviderPolicyBlocked, + FailoverReason::ContentPolicyBlocked, + FailoverReason::FormatError, + FailoverReason::InvalidEncryptedContent, + FailoverReason::MultimodalToolContentUnsupported, + FailoverReason::ThinkingSignature, + FailoverReason::LongContextTier, + FailoverReason::OauthLongContextBetaForbidden, + FailoverReason::SslTransient, + FailoverReason::Unknown("test".into()), + ]; + for v in &variants { + let suggestion = recovery_suggestion(v); + assert!(!suggestion.is_empty(), "Empty suggestion for {:?}", v); + } + } + + // ── Context Overflow 解析测试 ── + + #[test] + fn test_parse_overflow_anthropic_format() { + let err = "input length and max_tokens exceed context limit: 180000 + 32000 > 200000. \ + please reduce the length of the messages or the max_tokens."; + let info = parse_context_overflow(err).unwrap(); + assert_eq!(info.input_length, 180000); + assert_eq!(info.requested_max_tokens, 32000); + assert_eq!(info.context_limit, 200000); + } + + #[test] + fn test_parse_overflow_anthropic_format_different_numbers() { + let err = "Error: input length and max_tokens exceed context limit: 50000 + 16000 > 64000"; + let info = parse_context_overflow(err).unwrap(); + assert_eq!(info.input_length, 50000); + assert_eq!(info.requested_max_tokens, 16000); + assert_eq!(info.context_limit, 64000); + } + + #[test] + fn test_parse_overflow_generic_format() { + let err = "context length exceeded: 130000 + 32000 > 128000"; + let info = parse_context_overflow(err).unwrap(); + assert_eq!(info.input_length, 130000); + assert_eq!(info.requested_max_tokens, 32000); + assert_eq!(info.context_limit, 128000); + } + + #[test] + fn test_parse_overflow_no_match() { + let err = "Something went wrong but no overflow pattern here"; + assert!(parse_context_overflow(err).is_none()); + } + + #[test] + fn test_calculate_safe_max_tokens_normal() { + let info = ContextOverflowInfo { + input_length: 180000, + requested_max_tokens: 32000, + context_limit: 200000, + }; + // available = 200000 - 180000 = 20000 + // adjusted = 20000 - 1000 = 19000 + assert_eq!(calculate_safe_max_tokens(&info), Some(19000)); + } + + #[test] + fn test_calculate_safe_max_tokens_no_room() { + let info = ContextOverflowInfo { + input_length: 199000, + requested_max_tokens: 32000, + context_limit: 200000, + }; + // available = 200000 - 199000 = 1000 + // adjusted = 1000 - 1000 = 0 < MIN_ACCEPTABLE_TOKENS(500) + assert_eq!(calculate_safe_max_tokens(&info), None); + } + + #[test] + fn test_calculate_safe_max_tokens_input_exceeds_context() { + let info = ContextOverflowInfo { + input_length: 210000, + requested_max_tokens: 1000, + context_limit: 200000, + }; + // input > context → None (must compress) + assert_eq!(calculate_safe_max_tokens(&info), None); + } + + #[test] + fn test_recovery_uses_adjust_max_tokens_first() { + let mut attempts = RecoveryAttempts::new(); + let info = ContextOverflowInfo { + input_length: 180000, + requested_max_tokens: 32000, + context_limit: 200000, + }; + let step = attempts.next_step(&ErrorKind::PromptTooLong, Some(&info)); + assert_eq!( + step, + Some(RecoveryStep::AdjustMaxTokens { + new_max_tokens: 19000 + }) + ); + } + + #[test] + fn test_recovery_falls_back_to_compact_when_no_room() { + let mut attempts = RecoveryAttempts::new(); + // Input almost fills context — no room to adjust max_tokens + let info = ContextOverflowInfo { + input_length: 199500, + requested_max_tokens: 32000, + context_limit: 200000, + }; + let step = attempts.next_step(&ErrorKind::PromptTooLong, Some(&info)); + // Should skip AdjustMaxTokens and go to AggressiveCompact + assert_eq!(step, Some(RecoveryStep::AggressiveCompact)); + } + + #[test] + fn test_recovery_adjust_after_compact_exhausted() { + let mut attempts = RecoveryAttempts::new(); + // Mark aggressive_compact as tried + attempts.aggressive_compact = true; + + let info = ContextOverflowInfo { + input_length: 180000, + requested_max_tokens: 32000, + context_limit: 200000, + }; + // AdjustMaxTokens check is skipped when aggressive_compact is already true + let step = attempts.next_step(&ErrorKind::PromptTooLong, Some(&info)); + assert_eq!(step, Some(RecoveryStep::ReactiveCompact)); + } + + #[test] + fn test_try_recover_adjusts_token_budget() { + let budget = TokenBudget::new(32_000, 40_000); + let mut recovery = ErrorRecovery::new(budget); + let info = ContextOverflowInfo { + input_length: 180000, + requested_max_tokens: 32000, + context_limit: 200000, + }; + let step = recovery.try_recover(&ErrorKind::PromptTooLong, Some(&info)); + assert_eq!( + step, + Some(RecoveryStep::AdjustMaxTokens { + new_max_tokens: 19000 + }) + ); + // Token budget hard_limit should be adjusted down + assert_eq!(recovery.token_budget.hard_limit, 19000); + } } diff --git a/src/agent/runtime/executor.rs b/src/agent/runtime/executor.rs index 644f4dd..e2f64ee 100644 --- a/src/agent/runtime/executor.rs +++ b/src/agent/runtime/executor.rs @@ -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, pub was_cancelled: bool, pub had_duplicate: bool, + /// Hook 注入的附加上下文（PreToolUse + PostToolUse），需注入 LLM 消息列表 + pub hook_contexts: Vec, + /// Hook 的阻塞错误详情（用于日志和诊断） + pub blocking_errors: Vec, } /// 验证工具调用：死循环检测 + 参数解析。 @@ -144,6 +152,7 @@ pub async fn execute_parallel( permission_checker: Option<&PermissionChecker>, session_permission_checker: Option<&std::sync::RwLock>, denial_tracker: Option<&std::sync::Mutex>, + checkpoint_manager: Option<&std::sync::Arc>, tx: &mpsc::UnboundedSender, 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 = Vec::new(); let mut additional_contexts: Vec = Vec::new(); - let mut hook_permission_required: Vec = Vec::new(); + let mut hook_permission_info: Vec> = Vec::new(); + // ^^^ (permission_desc, hook_tool_name) + let mut hook_blocking_errors: Vec = 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 = Vec::new(); let mut denied_indices: std::collections::HashSet = 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 = + 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 = 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, + 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, + hook_registry: &HookRegistry, + library_dir: &std::path::Path, + sid: &str, + agent_name: &str, + step: usize, + max_output_chars: usize, + tool_messages: &mut Vec, + additional_contexts: &mut Vec, + 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, diff --git a/src/agent/runtime/finalize.rs b/src/agent/runtime/finalize.rs index 8a66d50..80c7d2c 100644 --- a/src/agent/runtime/finalize.rs +++ b/src/agent/runtime/finalize.rs @@ -37,7 +37,7 @@ pub async fn finalize_turn( app_state: Option>, ) -> 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) diff --git a/src/agent/runtime/hardline.rs b/src/agent/runtime/hardline.rs new file mode 100644 index 0000000..543c2ae --- /dev/null +++ b/src/agent/runtime/hardline.rs @@ -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 = + 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> = 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 { + 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 = + 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 = + 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> = 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 = 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, +} + +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 = 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 = 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); + } +} diff --git a/src/agent/runtime/mod.rs b/src/agent/runtime/mod.rs index 6ae066f..fd0b3a9 100644 --- a/src/agent/runtime/mod.rs +++ b/src/agent/runtime/mod.rs @@ -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>, /// 文件状态缓存（跨 turn 共享，用于 Read 去重） read_file_state: Arc>, + /// 系统提示词 section 缓存（跨 turn 共享，避免每 turn 重建静态/低频变动内容） + prompt_cache: std::sync::Mutex, + /// 上下文压缩折叠日志（跨 turn 共享，追踪压缩历史并触发溢出合并） + collapse_log: Arc>, + /// Checkpoint 管理器（跨 turn 共享，文件变更操作前自动快照） + checkpoint_manager: Arc, } 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!( + "\n[Hook 注入上下文]\n{}\n", + 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, diff --git a/src/agent/runtime/partitioner.rs b/src/agent/runtime/partitioner.rs index d04f65c..7936fce 100644 --- a/src/agent/runtime/partitioner.rs +++ b/src/agent/runtime/partitioner.rs @@ -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"); + } } diff --git a/src/agent/runtime/permission.rs b/src/agent/runtime/permission.rs index f9f0eb2..5fb640a 100644 --- a/src/agent/runtime/permission.rs +++ b/src/agent/runtime/permission.rs @@ -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, Vec) { + let mut final_result = checker_result; + let mut conflict_log: Vec = 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"))); + } } diff --git a/src/agent/runtime/session.rs b/src/agent/runtime/session.rs index 8969d25..5ed7d02 100644 --- a/src/agent/runtime/session.rs +++ b/src/agent/runtime/session.rs @@ -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, @@ -23,7 +34,6 @@ pub async fn create_or_resume_session( ) -> anyhow::Result { 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> { + 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> { 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, Option, )> = 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 { + // 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 = 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 { + let n = n.max(1); + + // 查找最近的 N 个 user 消息（按 id DESC） + let user_ids: Vec = 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 { + // ── 冲突检测 ── + // 查找最小的 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 { + // 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 = 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 = + 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); + } +} diff --git a/src/agent/runtime/streaming_executor.rs b/src/agent/runtime/streaming_executor.rs index 9e2b0ca..319cf16 100644 --- a/src/agent/runtime/streaming_executor.rs +++ b/src/agent/runtime/streaming_executor.rs @@ -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, - /// 取消通道（Sibling Abort 使用） - #[allow(dead_code)] - cancel_tx: Option>, -} - /// 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, + /// 并发安全的工具在 spawn 后的 JoinHandle + handle: Option>, +} + +/// 流式工具执行器。 +/// +/// 参考 Claude Code `StreamingToolExecutor` (531 行 TypeScript)， +/// 关键改进：并发安全工具立即 spawn tokio task，不等待 flush。 pub struct StreamingToolExecutor { - /// 所有跟踪中的工具 + /// 所有跟踪中的工具（按 LLM 流中到达顺序） tracked: Vec, /// 工具注册表 tool_registry: Arc, - /// 并发分区器（保留用于未来并发策略优化） - #[allow(dead_code)] - partitioner: ToolPartitioner, - /// 工具上下文 + /// 工具上下文（按需 clone 给每个 spawn 的 task） tool_context: ToolContext, /// Sibling Abort 广播通道 (tx) abort_tx: broadcast::Sender, - /// Sibling Abort 广播通道 (rx) + /// Sibling Abort 广播通道 (rx) — 保留以保持 channel 存活， + /// 实际使用时通过 `abort_tx.subscribe()` 获取新接收端。 + #[allow(dead_code)] abort_rx: broadcast::Receiver, /// 当前是否已发生错误（触发 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, } impl StreamingToolExecutor { - /// 创建新的流式执行器 + /// 创建新的流式执行器。 pub fn new( tool_registry: Arc, 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 { 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 = 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 { + 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 = 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 = 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)> = 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"); } } diff --git a/src/agent/runtime/system_prompt.rs b/src/agent/runtime/system_prompt.rs index 7424b59..4aaba58 100644 --- a/src/agent/runtime/system_prompt.rs +++ b/src/agent/runtime/system_prompt.rs @@ -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, + ) -> Option { + 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 = "\ +# 系统上下文 +- 工具结果和用户消息中可能包含 标签。这些标签由系统自动添加，包含有用的信息和提醒，与所在消息的具体内容无直接关系。 +- 对话具有通过自动摘要实现的无限上下文长度。当上下文接近限制时，较早的消息会被自动压缩为摘要。"; + +/// 静态工具使用指导 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:: + }); + assert_eq!(r1, None); + assert_eq!(call_count, 1); + + let r2 = cache.get_or_compute_optional("opt", || { + call_count += 1; + None:: + }); + 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); + } } diff --git a/src/agent/runtime/untrusted.rs b/src/agent/runtime/untrusted.rs new file mode 100644 index 0000000..1a6424e --- /dev/null +++ b/src/agent/runtime/untrusted.rs @@ -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）的结果包裹在 中 +// 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 +/// +/// {original_content} +/// +/// ``` +/// +/// 此标记指示 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("\n{content}\n", + tool = tool_name, + content = content + ) +} + +/// 从错误包裹中恢复原内容（当需要向用户展示时） +pub fn unwrap_untrusted(content: &str) -> String { + if !content.starts_with("") { + 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 "; + let wrapped = wrap_untrusted_content("web_search", content); + assert!(wrapped.starts_with("")); + assert!(wrapped.ends_with("")); + 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 = "\nalready wrapped\n"; + 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); + } +} diff --git a/src/agent/skills.rs b/src/agent/skills.rs index c1f3f95..77d2b05 100644 --- a/src/agent/skills.rs +++ b/src/agent/skills.rs @@ -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, + /// 是否禁止 Curator 自动清理（参考 Hermes pinned skills） + #[serde(default)] + pub pinned: Option, } impl SkillFrontmatter { @@ -116,6 +127,8 @@ pub struct SkillMeta { pub user_invocable: bool, /// 条件激活的 glob 模式（空 Vec 表示始终激活） pub paths: Vec, + /// 是否禁止 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 Result Self { + SkillCreator { skills_dir } + } + + /// 从检测到的模式创建 SKILL.md 文件。 + /// + /// 返回创建的 skill 名称列表。 + /// 如果目标 skill 目录已存在则跳过（不覆盖已有 skill）。 + pub fn create_from_patterns( + &self, + patterns: &[pattern_detector::DetectedPattern], + ) -> std::io::Result> { + 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 = 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::>() + .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::>() + .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) -> 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, + skill_metas: &[SkillMeta], + ) -> anyhow::Result { + // 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> { + self.detector.scan(db, max_sessions, 1).await + } + + /// 仅运行 curator 分析。 + pub fn analyze_only( + &self, + usage_stats: &HashMap, + 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, + /// Curator 质量报告 + pub curator_report: curator::CuratorReport, +} + // ── Tests ────────────────────────────────────────────────────────────────── #[cfg(test)] diff --git a/src/agent/skills/curator.rs b/src/agent/skills/curator.rs new file mode 100644 index 0000000..b5531d2 --- /dev/null +++ b/src/agent/skills/curator.rs @@ -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, + /// 总调用次数 + pub total_invocations: u64, + /// 预估成功率 (成功调用 / 总调用), None 表示无数据 + pub estimated_success_rate: Option, + /// Curator 的建议 + pub recommendation: String, +} + +/// Curator 分析报告 +#[derive(Debug, Clone, Serialize)] +pub struct CuratorReport { + /// 分析时间 + pub analyzed_at: DateTime, + /// 各 skill 的质量评估（按 skill name 索引） + pub skills: Vec, + /// 建议清理的 skill 名称列表 + pub cleanup_candidates: Vec, + /// 总 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, + 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, + 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, + ) -> 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> { + // 确保归档目录存在 + 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, + db: sqlx::SqlitePool, + /// 当前是否已触发过（防止重复运行） + last_run_at: Arc>>>, + /// 是否暂停 + paused: Arc, + /// 运行间隔 + interval: Duration, + /// 最小空闲时间 + min_idle: Duration, + /// 最后一次 Agent 活动时间 + last_activity: Arc>>, +} + +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> { + *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, + skill_metas: &[crate::agent::skills::SkillMeta], + ) -> Option { + 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, + usage_stats: Arc>>, + skill_metas: Arc>>, + 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); + } +} diff --git a/src/agent/skills/pattern_detector.rs b/src/agent/skills/pattern_detector.rs new file mode 100644 index 0000000..6bab203 --- /dev/null +++ b/src/agent/skills/pattern_detector.rs @@ -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, + /// 该模式在不同 session 中的出现频次 + pub occurrence_count: usize, + /// 出现该模式的 session ID 列表 + pub session_ids: Vec, + /// 一个典型的完整参数示例（来自最近一次出现） + pub example_args: Vec, + /// 一个典型的人类可读描述（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, +} + +impl PatternDetector { + /// 创建检测器，传入已注册 skill 的指纹集合以避免重复创建。 + pub fn new(known_patterns: HashSet) -> 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> { + // 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, Vec)> = 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 = 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 = 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> { + 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, Vec)> { + 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::(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> { + 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 = 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) -> Vec { + // 按序列长度降序排列 + patterns.sort_by(|a, b| b.tool_sequence.len().cmp(&a.tool_sequence.len())); + + let mut retained: Vec = 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 = 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 = 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 = vec!["B", "C"].into_iter().map(String::from).collect(); + let long: Vec = 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 = vec!["B", "D"].into_iter().map(String::from).collect(); + let long: Vec = 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 = vec!["A", "B", "C"].into_iter().map(String::from).collect(); + let b: Vec = 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); + } +} diff --git a/src/agent/subagent.rs b/src/agent/subagent.rs index bc7b065..8afba4c 100644 --- a/src/agent/subagent.rs +++ b/src/agent/subagent.rs @@ -100,6 +100,25 @@ impl SubAgentRunner { } } + /// 创建带完整 hooks/permissions + 自定义 ToolRegistry 的子代理运行器。 + pub fn new_with_registry_and_hooks( + app_state: Arc, + tool_registry: ToolRegistry, + hook_registry: Option>, + permission_checker: Arc, + progress_tx: Option>, + ) -> 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; } diff --git a/src/agent/tools/mod.rs b/src/agent/tools/mod.rs index 981b9d3..4712cf8 100644 --- a/src/agent/tools/mod.rs +++ b/src/agent/tools/mod.rs @@ -50,6 +50,7 @@ pub use todo::persist_tasks; pub use todo::TodoWriteTool; /// 工具执行上下文，封装全局共享状态 +#[derive(Clone)] pub struct ToolContext { pub app_state: Arc, /// 当前会话 ID（用于工具将数据关联到正确的 session） @@ -275,6 +276,10 @@ pub struct ToolRegistry { ordered_names: Vec, /// 可选的工具白名单（子代理最小权限） definition_filter: Option>, + /// 缓存的工具 schema（session 生命周期内有效，工具注册变更时失效） + schema_cache: Option>, + /// 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) { + 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 { + // 缓存命中：直接返回 + if let Some(ref cached) = self.schema_cache { + return cached.clone(); + } + + // 缓存未命中：重新计算 let values: Vec<&Box> = 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> = 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()` 仍可访问所有工具以便对未授权调用返回友好错误消息。 diff --git a/src/agent/tools/subagent.rs b/src/agent/tools/subagent.rs index 7f6ece1..b4eaaf3 100644 --- a/src/agent/tools/subagent.rs +++ b/src/agent/tools/subagent.rs @@ -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)) diff --git a/src/agent/trajectory.rs b/src/agent/trajectory.rs index 8037571..95e4f8e 100644 --- a/src/agent/trajectory.rs +++ b/src/agent/trajectory.rs @@ -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) diff --git a/src/api/agent.rs b/src/api/agent.rs index 553f771..99e7df0 100644 --- a/src/api/agent.rs +++ b/src/api/agent.rs @@ -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>, + Path(session_id): Path, +) -> Result, (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>, + Path(session_id): Path, +) -> Result, (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, + /// 或者指定回退到的消息 ID + pub message_id: Option, +} + +#[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>, + Path(session_id): Path, + Json(req): Json, +) -> Result, (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>, + Path(session_id): Path, +) -> Result, (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, + })) +} diff --git a/src/api/mod.rs b/src/api/mod.rs index 1b24850..1e5f476 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -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::{ diff --git a/src/clients/llm.rs b/src/clients/llm.rs index 619154b..10ea15a 100644 --- a/src/clients/llm.rs +++ b/src/clients/llm.rs @@ -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")) diff --git a/src/main.rs b/src/main.rs index 7c6ade2..4dc17c2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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(