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
10static HOOK_HELP_CONTENT: OnceLock<String> = OnceLock::new();
12
13pub 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#[derive(Deserialize, JsonSchema)]
22struct RegisterHookParams {
23 #[serde(default = "default_action")]
25 action: String,
26 #[serde(default)]
28 event: Option<String>,
29 #[serde(default)]
31 r#type: Option<String>,
32 #[serde(default)]
34 command: Option<String>,
35 #[serde(default)]
37 prompt: Option<String>,
38 #[serde(default)]
40 model: Option<String>,
41 #[serde(default)]
43 timeout: Option<u64>,
44 #[serde(default)]
46 retry: Option<u32>,
47 #[serde(default)]
49 index: Option<usize>,
50 #[serde(default)]
52 on_error: Option<String>,
53}
54
55fn default_action() -> String {
56 "register".to_string()
57}
58
59#[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(¶ms),
97 _ => self.handle_register(¶ms),
98 }
99 }
100
101 fn requires_confirmation(&self) -> bool {
102 true }
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 #[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 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 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}