Skip to main content

j_agent/tools/
hook.rs

1use std::sync::OnceLock;
2
3use crate::constants::HOOK_LOG_DESC_MAX_LEN;
4use crate::infra::hook::{HookDef, HookEvent, HookFilter, HookManager, HookType, OnError};
5use crate::tools::{PlanDecision, Tool, ToolResult, parse_tool_args, schema_to_tool_params};
6use schemars::JsonSchema;
7use serde::Deserialize;
8use serde_json::Value;
9
10/// 全局 Hook 帮助内容(由 j-cli 在启动时通过 set_hook_help_content 注入)
11static HOOK_HELP_CONTENT: OnceLock<String> = OnceLock::new();
12
13/// 设置 Hook 帮助文档内容(j-cli 启动时调用)
14pub fn set_hook_help_content(content: String) {
15    let _ = HOOK_HELP_CONTENT.set(content);
16}
17use std::borrow::Cow;
18use std::sync::{Arc, Mutex, atomic::AtomicBool};
19
20/// RegisterHookTool 参数
21#[derive(Deserialize, JsonSchema)]
22struct RegisterHookParams {
23    /// Action type: register (default), list, remove, help
24    #[serde(default = "default_action")]
25    action: String,
26    /// Hook event name (required for register/remove)
27    #[serde(default)]
28    event: Option<String>,
29    /// Hook type: "bash" (default) or "llm"
30    #[serde(default)]
31    r#type: Option<String>,
32    /// Shell command to execute (required for type=bash)
33    #[serde(default)]
34    command: Option<String>,
35    /// LLM prompt template (required for type=llm, supports {{variable}} template vars)
36    #[serde(default)]
37    prompt: Option<String>,
38    /// LLM model name override (optional for type=llm)
39    #[serde(default)]
40    model: Option<String>,
41    /// Timeout in seconds (default 10 for bash, 30 for llm)
42    #[serde(default)]
43    timeout: Option<u64>,
44    /// Retry count on error (default 0 for bash, 1 for llm; only applies to Err path)
45    #[serde(default)]
46    retry: Option<u32>,
47    /// Index of the session hook to remove (required for remove). Use session_idx from list output.
48    #[serde(default)]
49    index: Option<usize>,
50    /// Error handling strategy: "skip" (default, log and continue) or "stop" (stop hook chain)
51    #[serde(default)]
52    on_error: Option<String>,
53}
54
55fn default_action() -> String {
56    "register".to_string()
57}
58
59/// register_hook 工具:让 LLM 动态注册/管理 session 级 hook
60#[derive(Debug)]
61pub struct RegisterHookTool {
62    pub hook_manager: Arc<Mutex<HookManager>>,
63}
64
65impl RegisterHookTool {
66    pub const NAME: &'static str = "RegisterHook";
67}
68
69impl Tool for RegisterHookTool {
70    fn name(&self) -> &str {
71        Self::NAME
72    }
73
74    fn description(&self) -> Cow<'_, str> {
75        r#"
76        Register, list, remove session-level hooks, or view the full protocol documentation.
77        Actions: register (requires event+command or event+prompt), list, remove (requires event+index), help (view stdin/stdout JSON schema and script examples).
78        Supports two hook types: "bash" (shell command, default) and "llm" (LLM prompt template).
79        Call action="help" first to learn the script protocol before registering hooks.
80        "#.into()
81    }
82
83    fn parameters_schema(&self) -> Value {
84        schema_to_tool_params::<RegisterHookParams>()
85    }
86
87    fn execute(&self, arguments: &str, _cancelled: &Arc<AtomicBool>) -> ToolResult {
88        let params: RegisterHookParams = match parse_tool_args(arguments) {
89            Ok(p) => p,
90            Err(e) => return e,
91        };
92
93        match params.action.as_str() {
94            "help" => Self::handle_help(),
95            "list" => self.handle_list(),
96            "remove" => self.handle_remove(&params),
97            _ => self.handle_register(&params),
98        }
99    }
100
101    fn requires_confirmation(&self) -> bool {
102        true // 注册 hook 需要用户确认
103    }
104
105    fn confirmation_message(&self, arguments: &str) -> String {
106        if let Ok(params) = serde_json::from_str::<RegisterHookParams>(arguments) {
107            match params.action.as_str() {
108                "help" => "View Hook protocol documentation".to_string(),
109                "list" => "List all registered hooks".to_string(),
110                "remove" => {
111                    let event = params.event.as_deref().unwrap_or("?");
112                    let index = params.index.unwrap_or(0);
113                    format!("Remove hook: event={}, index={}", event, index)
114                }
115                _ => {
116                    let event = params.event.as_deref().unwrap_or("?");
117                    let hook_type = params.r#type.as_deref().unwrap_or("bash");
118                    let desc = if hook_type == "llm" {
119                        let prompt_preview = params
120                            .prompt
121                            .as_deref()
122                            .map(|p| if p.len() > 60 { &p[..60] } else { p })
123                            .unwrap_or("?");
124                        format!("type=llm, prompt={}", prompt_preview)
125                    } else {
126                        let cmd = params.command.as_deref().unwrap_or("?");
127                        format!("type=bash, command={}", cmd)
128                    };
129                    let on_error = params.on_error.as_deref().unwrap_or("skip");
130                    format!(
131                        "Register hook: event={}, {}, on_error={}",
132                        event, desc, on_error
133                    )
134                }
135            }
136        } else {
137            "RegisterHook operation".to_string()
138        }
139    }
140}
141
142impl RegisterHookTool {
143    fn handle_help() -> ToolResult {
144        let content = HOOK_HELP_CONTENT
145            .get()
146            .map(|s| Self::strip_frontmatter(s).to_string())
147            .unwrap_or_else(|| {
148                "Hook 文档加载失败(需要调用 set_hook_help_content 注入)".to_string()
149            });
150
151        ToolResult {
152            output: content,
153            is_error: false,
154            images: vec![],
155            plan_decision: PlanDecision::None,
156        }
157    }
158
159    /// 去掉 YAML frontmatter(`---` ... `---`),返回 body 部分
160    #[allow(dead_code)]
161    fn strip_frontmatter(content: &str) -> &str {
162        let trimmed = content.trim_start();
163        if !trimmed.starts_with("---") {
164            return trimmed;
165        }
166        let after_first = &trimmed[3..];
167        if let Some(end) = after_first.find("\n---") {
168            return after_first[end + 4..].trim_start();
169        }
170        trimmed
171    }
172
173    fn handle_register(&self, params: &RegisterHookParams) -> ToolResult {
174        let event_str = match params.event.as_deref() {
175            Some(e) => e,
176            None => {
177                return ToolResult {
178                    output: "缺少 event 参数".to_string(),
179                    is_error: true,
180                    images: vec![],
181                    plan_decision: PlanDecision::None,
182                };
183            }
184        };
185
186        let event = match HookEvent::parse(event_str) {
187            Some(e) => e,
188            None => {
189                return ToolResult {
190                    output: format!("未知事件: {}", event_str),
191                    is_error: true,
192                    images: vec![],
193                    plan_decision: PlanDecision::None,
194                };
195            }
196        };
197
198        let (hook_def, detail, on_error_str) = match Self::build_hook_def_from_params(params) {
199            Ok(result) => result,
200            Err(err_msg) => {
201                return ToolResult {
202                    output: err_msg,
203                    is_error: true,
204                    images: vec![],
205                    plan_decision: PlanDecision::None,
206                };
207            }
208        };
209
210        match self.hook_manager.lock() {
211            Ok(mut manager) => {
212                manager.register_session_hook(event, hook_def);
213                let type_str = params.r#type.as_deref().unwrap_or("bash").to_string();
214                ToolResult {
215                    output: format!(
216                        "已注册 session hook: event={}, type={}, {}, timeout={}s, retry={}, on_error={}",
217                        event_str,
218                        type_str,
219                        detail,
220                        params
221                            .timeout
222                            .unwrap_or(if params.r#type.as_deref() == Some("llm") {
223                                30
224                            } else {
225                                10
226                            }),
227                        params
228                            .retry
229                            .unwrap_or(if params.r#type.as_deref() == Some("llm") {
230                                1
231                            } else {
232                                0
233                            }),
234                        on_error_str
235                    ),
236                    is_error: false,
237                    images: vec![],
238                    plan_decision: PlanDecision::None,
239                }
240            }
241            Err(e) => ToolResult {
242                output: format!("获取 HookManager 锁失败: {}", e),
243                is_error: true,
244                images: vec![],
245                plan_decision: PlanDecision::None,
246            },
247        }
248    }
249
250    /// 从参数构建 HookDef,返回 (hook_def, 日志描述, on_error标签) 或错误消息
251    fn build_hook_def_from_params(
252        params: &RegisterHookParams,
253    ) -> Result<(HookDef, String, &'static str), String> {
254        let hook_type = match params.r#type.as_deref() {
255            Some("llm") => HookType::Llm,
256            _ => HookType::Bash,
257        };
258
259        // 校验必填字段
260        match hook_type {
261            HookType::Bash => {
262                if params.command.is_none() {
263                    return Err("bash hook 缺少 command 参数".to_string());
264                }
265            }
266            HookType::Llm => {
267                if params.prompt.is_none() {
268                    return Err("llm hook 缺少 prompt 参数".to_string());
269                }
270            }
271        }
272
273        let timeout = params.timeout.unwrap_or(match hook_type {
274            HookType::Bash => 10,
275            HookType::Llm => 30,
276        });
277
278        let retry = params.retry.unwrap_or(match hook_type {
279            HookType::Bash => 0,
280            HookType::Llm => 1,
281        });
282
283        let on_error = match params.on_error.as_deref() {
284            Some("stop") => OnError::Stop,
285            _ => OnError::Skip,
286        };
287
288        let on_error_str = match on_error {
289            OnError::Skip => "skip",
290            OnError::Stop => "stop",
291        };
292
293        let hook_def = HookDef {
294            r#type: hook_type,
295            command: params.command.clone(),
296            prompt: params.prompt.clone(),
297            model: params.model.clone(),
298            timeout,
299            retry,
300            on_error,
301            filter: HookFilter::default(),
302        };
303
304        let detail = match hook_type {
305            HookType::Bash => {
306                format!("command={}", params.command.as_deref().unwrap_or("?"))
307            }
308            HookType::Llm => {
309                let prompt_preview = params
310                    .prompt
311                    .as_deref()
312                    .map(|p| {
313                        if p.len() > HOOK_LOG_DESC_MAX_LEN {
314                            &p[..HOOK_LOG_DESC_MAX_LEN]
315                        } else {
316                            p
317                        }
318                    })
319                    .unwrap_or("?");
320                format!("prompt={}", prompt_preview)
321            }
322        };
323
324        Ok((hook_def, detail, on_error_str))
325    }
326
327    fn handle_list(&self) -> ToolResult {
328        match self.hook_manager.lock() {
329            Ok(manager) => {
330                let hooks = manager.list_hooks();
331                if hooks.is_empty() {
332                    return ToolResult {
333                        output: "当前没有已注册的 hook".to_string(),
334                        is_error: false,
335                        images: vec![],
336                        plan_decision: PlanDecision::None,
337                    };
338                }
339
340                let mut output = String::from("已注册的 hook:\n");
341                for (i, entry) in hooks.iter().enumerate() {
342                    let timeout_str = entry
343                        .timeout
344                        .map(|t| format!("{}s", t))
345                        .unwrap_or_else(|| "-".to_string());
346                    let on_error_str = entry
347                        .on_error
348                        .map(|e| match e {
349                            OnError::Skip => "skip",
350                            OnError::Stop => "stop",
351                        })
352                        .unwrap_or("-");
353                    let session_idx_str = entry
354                        .session_index
355                        .map(|idx| format!(", session_idx={}", idx))
356                        .unwrap_or_default();
357                    let filter_str = entry
358                        .filter
359                        .as_ref()
360                        .map(|f| {
361                            let mut parts = Vec::new();
362                            if let Some(ref t) = f.tool_name {
363                                parts.push(format!("tool={}", t));
364                            }
365                            if let Some(ref m) = f.model_prefix {
366                                parts.push(format!("model={}*", m));
367                            }
368                            if parts.is_empty() {
369                                String::new()
370                            } else {
371                                format!(", filter=[{}]", parts.join(","))
372                            }
373                        })
374                        .unwrap_or_default();
375                    let metrics_str = entry
376                        .metrics
377                        .as_ref()
378                        .map(|m| {
379                            format!(
380                                ", runs={}/ok={}/fail={}/skip={}/{}ms",
381                                m.executions,
382                                m.successes,
383                                m.failures,
384                                m.skipped,
385                                m.total_duration_ms
386                            )
387                        })
388                        .unwrap_or_default();
389                    let name_str = entry.name.as_deref().unwrap_or("");
390                    let name_display = if name_str.is_empty() {
391                        String::new()
392                    } else {
393                        format!(", name={}", name_str)
394                    };
395                    output.push_str(&format!(
396                        "  [{}] event={}, source={}, type={}{}, label={}, timeout={}, on_error={}{}{}{}\n",
397                        i,
398                        entry.event.as_str(),
399                        entry.source,
400                        entry.hook_type,
401                        session_idx_str,
402                        entry.label,
403                        timeout_str,
404                        on_error_str,
405                        filter_str,
406                        metrics_str,
407                        name_display,
408                    ));
409                }
410                ToolResult {
411                    output,
412                    is_error: false,
413                    images: vec![],
414                    plan_decision: PlanDecision::None,
415                }
416            }
417            Err(e) => ToolResult {
418                output: format!("获取 HookManager 锁失败: {}", e),
419                is_error: true,
420                images: vec![],
421                plan_decision: PlanDecision::None,
422            },
423        }
424    }
425
426    fn handle_remove(&self, params: &RegisterHookParams) -> ToolResult {
427        let event_str = match params.event.as_deref() {
428            Some(e) => e,
429            None => {
430                return ToolResult {
431                    output: "缺少 event 参数".to_string(),
432                    is_error: true,
433                    images: vec![],
434                    plan_decision: PlanDecision::None,
435                };
436            }
437        };
438
439        let event = match HookEvent::parse(event_str) {
440            Some(e) => e,
441            None => {
442                return ToolResult {
443                    output: format!("未知事件: {}", event_str),
444                    is_error: true,
445                    images: vec![],
446                    plan_decision: PlanDecision::None,
447                };
448            }
449        };
450
451        let index = params.index.unwrap_or(0);
452
453        match self.hook_manager.lock() {
454            Ok(mut manager) => {
455                if manager.remove_session_hook(event, index) {
456                    ToolResult {
457                        output: format!(
458                            "已移除 session hook: event={}, index={}",
459                            event_str, index
460                        ),
461                        is_error: false,
462                        images: vec![],
463                        plan_decision: PlanDecision::None,
464                    }
465                } else {
466                    ToolResult {
467                        output: format!(
468                            "移除失败:event={} 的 session hook 索引 {} 不存在",
469                            event_str, index
470                        ),
471                        is_error: true,
472                        images: vec![],
473                        plan_decision: PlanDecision::None,
474                    }
475                }
476            }
477            Err(e) => ToolResult {
478                output: format!("获取 HookManager 锁失败: {}", e),
479                is_error: true,
480                images: vec![],
481                plan_decision: PlanDecision::None,
482            },
483        }
484    }
485}