1use crate::infra::hook::definition::*;
2use crate::infra::hook::executor::execute_hook_with_provider;
3use crate::infra::hook::types::*;
4use crate::util::log::{write_error_log, write_info_log};
5use std::collections::HashMap;
6use std::sync::{Arc, Mutex};
7use std::thread;
8use std::time::{Duration, Instant};
9
10#[derive(Debug, Clone, Default, PartialEq)]
14pub struct HookMetrics {
15 pub executions: u64,
17 pub successes: u64,
19 pub failures: u64,
21 pub skipped: u64,
23 pub total_duration_ms: u64,
25}
26
27const HOOK_SOURCE_BUILTIN: &str = "builtin";
29const HOOK_SOURCE_USER: &str = "user";
30const HOOK_SOURCE_PROJECT: &str = "project";
31const HOOK_SOURCE_SESSION: &str = "session";
32
33pub struct HookEntry {
35 pub name: Option<String>,
37 pub event: HookEvent,
38 pub source: &'static str,
39 pub hook_type: &'static str,
41 pub label: String,
43 pub timeout: Option<u64>,
45 pub on_error: Option<OnError>,
47 pub session_index: Option<usize>,
49 pub filter: Option<HookFilter>,
51 pub metrics: Option<HookMetrics>,
53 pub unique_id: String,
55}
56
57#[derive(Debug, Default)]
64pub struct HookManager {
65 builtin_hooks: HashMap<HookEvent, Vec<HookKind>>,
66 user_hooks: HashMap<HookEvent, Vec<HookKind>>,
67 project_hooks: HashMap<HookEvent, Vec<HookKind>>,
68 session_hooks: HashMap<HookEvent, Vec<HookKind>>,
69 pub(crate) metrics: Mutex<HashMap<String, HookMetrics>>,
71 pub(crate) provider: Option<Arc<Mutex<crate::storage::ModelProvider>>>,
73}
74
75impl Clone for HookManager {
76 fn clone(&self) -> Self {
77 HookManager {
78 builtin_hooks: self.builtin_hooks.clone(),
79 user_hooks: self.user_hooks.clone(),
80 project_hooks: self.project_hooks.clone(),
81 session_hooks: self.session_hooks.clone(),
82 metrics: Mutex::new(self.metrics.lock().map(|m| m.clone()).unwrap_or_default()),
83 provider: self.provider.clone(),
84 }
85 }
86}
87
88impl HookManager {
89 pub fn load() -> Self {
91 let mut manager = HookManager::default();
92
93 let user_dir = hooks_dir();
95 if user_dir.is_dir() {
96 for (name, dir_def, dir_path) in load_hooks_from_dir(&user_dir, "用户级") {
97 match dir_def.into_hook_kinds(&name, &dir_path) {
98 Ok(pairs) => {
99 for (event, kind) in pairs {
100 manager.user_hooks.entry(event).or_default().push(kind);
101 }
102 }
103 Err(e) => write_error_log("HookManager::load", &e),
104 }
105 }
106 write_info_log(
107 "HookManager::load",
108 &format!("已加载用户级 hooks: {}", user_dir.display()),
109 );
110 }
111
112 if let Some(proj_dir) = project_hooks_dir() {
114 for (name, dir_def, dir_path) in load_hooks_from_dir(&proj_dir, "项目级") {
115 match dir_def.into_hook_kinds(&name, &dir_path) {
116 Ok(pairs) => {
117 for (event, kind) in pairs {
118 manager.project_hooks.entry(event).or_default().push(kind);
119 }
120 }
121 Err(e) => write_error_log("HookManager::load", &e),
122 }
123 }
124 write_info_log(
125 "HookManager::load",
126 &format!("已加载项目级 hooks: {}", proj_dir.display()),
127 );
128 }
129
130 manager
131 }
132
133 pub fn register_builtin(
135 &mut self,
136 event: HookEvent,
137 name: impl Into<String>,
138 handler: impl Fn(&HookContext) -> Option<HookResult> + Send + Sync + 'static,
139 ) {
140 self.builtin_hooks
141 .entry(event)
142 .or_default()
143 .push(HookKind::Builtin(BuiltinHook {
144 name: name.into(),
145 handler: Arc::new(handler),
146 }));
147 }
148
149 pub fn register_session_hook(&mut self, event: HookEvent, def: HookDef) {
151 match def.into_hook_kind() {
152 Ok(kind) => {
153 self.session_hooks.entry(event).or_default().push(kind);
154 }
155 Err(e) => {
156 write_error_log("HookManager::register_session_hook", &e);
157 }
158 }
159 }
160
161 pub fn session_hooks_snapshot(&self) -> Vec<crate::storage::SessionHookPersist> {
164 let mut result = Vec::new();
165 for (event, hooks) in &self.session_hooks {
166 for kind in hooks {
167 match kind {
168 HookKind::Shell(sh) => {
169 result.push(crate::storage::SessionHookPersist {
170 event: *event,
171 definition: HookDef {
172 r#type: HookType::Bash,
173 command: Some(sh.command.clone()),
174 prompt: None,
175 model: None,
176 timeout: sh.timeout,
177 retry: sh.retry,
178 on_error: sh.on_error,
179 filter: sh.filter.clone(),
180 },
181 });
182 }
183 HookKind::Llm(lh) => {
184 result.push(crate::storage::SessionHookPersist {
185 event: *event,
186 definition: HookDef {
187 r#type: HookType::Llm,
188 command: None,
189 prompt: Some(lh.prompt.clone()),
190 model: lh.model.clone(),
191 timeout: lh.timeout,
192 retry: lh.retry,
193 on_error: lh.on_error,
194 filter: lh.filter.clone(),
195 },
196 });
197 }
198 HookKind::Builtin(_) => {
199 }
201 }
202 }
203 }
204 result
205 }
206
207 pub fn clear_session_hooks(&mut self) {
209 self.session_hooks.clear();
210 }
211
212 pub fn restore_session_hooks(&mut self, hooks: &[crate::storage::SessionHookPersist]) {
214 self.session_hooks.clear();
215 for hook in hooks {
216 self.register_session_hook(hook.event, hook.definition.clone());
217 }
218 }
219
220 #[allow(dead_code)]
222 pub fn register_session_hook_kind(&mut self, event: HookEvent, kind: HookKind) {
223 self.session_hooks.entry(event).or_default().push(kind);
224 }
225
226 pub fn set_provider(&mut self, provider: Arc<Mutex<crate::storage::ModelProvider>>) {
228 self.provider = Some(provider);
229 }
230
231 pub fn remove_session_hook(&mut self, event: HookEvent, index: usize) -> bool {
233 if let Some(hooks) = self.session_hooks.get_mut(&event)
234 && index < hooks.len()
235 {
236 hooks.remove(index);
237 return true;
238 }
239 false
240 }
241
242 pub fn list_hooks(&self) -> Vec<HookEntry> {
244 let mut result = Vec::new();
245 let metrics_map = self.metrics.lock().ok();
246 let empty_metrics = HashMap::new();
247 let metrics_ref = metrics_map.as_deref().unwrap_or(&empty_metrics);
248 let make_entry = |event: HookEvent,
249 source: &'static str,
250 hook: &HookKind,
251 session_index: Option<usize>,
252 metrics: &HashMap<String, HookMetrics>| {
253 let label = hook_label(hook);
254 let uid = hook_unique_id(source, hook, session_index);
255 HookEntry {
256 name: hook_name(hook).map(|s| s.to_string()),
257 event,
258 source,
259 hook_type: hook_type_str(hook),
260 timeout: hook_timeout(hook),
261 on_error: hook_on_error(hook),
262 filter: hook_filter(hook).cloned(),
263 metrics: metrics.get(&label).cloned(),
264 session_index,
265 label,
266 unique_id: uid,
267 }
268 };
269 for event in HookEvent::all() {
270 if let Some(hooks) = self.builtin_hooks.get(event) {
271 for hook in hooks {
272 result.push(make_entry(
273 *event,
274 HOOK_SOURCE_BUILTIN,
275 hook,
276 None,
277 metrics_ref,
278 ));
279 }
280 }
281 if let Some(hooks) = self.user_hooks.get(event) {
282 for hook in hooks {
283 result.push(make_entry(
284 *event,
285 HOOK_SOURCE_USER,
286 hook,
287 None,
288 metrics_ref,
289 ));
290 }
291 }
292 if let Some(hooks) = self.project_hooks.get(event) {
293 for hook in hooks {
294 result.push(make_entry(
295 *event,
296 HOOK_SOURCE_PROJECT,
297 hook,
298 None,
299 metrics_ref,
300 ));
301 }
302 }
303 if let Some(hooks) = self.session_hooks.get(event) {
304 for (idx, hook) in hooks.iter().enumerate() {
305 result.push(make_entry(
306 *event,
307 HOOK_SOURCE_SESSION,
308 hook,
309 Some(idx),
310 metrics_ref,
311 ));
312 }
313 }
314 }
315 result
316 }
317
318 pub fn has_hooks_for(&self, event: HookEvent) -> bool {
323 self.builtin_hooks
324 .get(&event)
325 .is_some_and(|h| !h.is_empty())
326 || self.user_hooks.get(&event).is_some_and(|h| !h.is_empty())
327 || self
328 .project_hooks
329 .get(&event)
330 .is_some_and(|h| !h.is_empty())
331 || self
332 .session_hooks
333 .get(&event)
334 .is_some_and(|h| !h.is_empty())
335 }
336
337 pub fn execute_fire_and_forget(
340 manager: Arc<Mutex<HookManager>>,
341 event: HookEvent,
342 context: HookContext,
343 disabled_hooks: Vec<String>,
344 ) {
345 thread::spawn(move || {
346 if let Ok(m) = manager.lock() {
347 let _ = m.execute(event, context, &disabled_hooks);
348 }
349 });
350 }
351
352 pub fn execute(
361 &self,
362 event: HookEvent,
363 mut context: HookContext,
364 disabled_hooks: &[String],
365 ) -> Option<HookResult> {
366 let all_hooks = collect_hooks_for_event(self, event);
367 if all_hooks.is_empty() {
368 return None;
369 }
370
371 write_info_log(
372 "HookManager::execute",
373 &format!(
374 "执行 {} 个 hook (事件: {})",
375 all_hooks.len(),
376 event.as_str()
377 ),
378 );
379
380 let mut had_modification = false;
381 let mut final_result = HookResult::default();
382 let chain_start = Instant::now();
383 let chain_timeout = Duration::from_secs(MAX_CHAIN_DURATION_SECS);
384
385 for hook_ref in &all_hooks {
386 if chain_start.elapsed() > chain_timeout {
388 write_error_log(
389 "HookManager::execute",
390 &format!(
391 "Hook 链总超时 ({}s),中止剩余 hook (事件: {})",
392 MAX_CHAIN_DURATION_SECS,
393 event.as_str()
394 ),
395 );
396 break;
397 }
398
399 let label = hook_label(hook_ref.kind);
400
401 let uid = hook_unique_id(hook_ref.source, hook_ref.kind, hook_ref.session_index);
403 if disabled_hooks.contains(&uid) {
404 if let Ok(mut metrics) = self.metrics.lock() {
405 let m = metrics.entry(label).or_default();
406 m.skipped += 1;
407 }
408 continue;
409 }
410
411 if !hook_should_execute(hook_ref.kind, &context) {
413 if let Ok(mut metrics) = self.metrics.lock() {
414 let m = metrics.entry(label).or_default();
415 m.skipped += 1;
416 }
417 continue;
418 }
419
420 let max_attempts = 1 + hook_retry_count(hook_ref.kind); let mut last_outcome = None;
422
423 for attempt in 0..max_attempts {
424 if chain_start.elapsed() > chain_timeout {
426 write_error_log(
427 "HookManager::execute",
428 &format!(
429 "Hook 链总超时,中止 {} 的重试 (事件: {})",
430 label,
431 event.as_str()
432 ),
433 );
434 last_outcome = Some(HookOutcome::Err(format!(
435 "链总超时,第 {} 次尝试中止",
436 attempt + 1
437 )));
438 break;
439 }
440
441 let hook_start = Instant::now();
442 let result = execute_hook_with_provider(hook_ref.kind, &context, &self.provider);
443
444 let elapsed_ms = hook_start.elapsed().as_millis() as u64;
445
446 match result {
447 Ok(hook_result) => {
448 if let Ok(mut metrics) = self.metrics.lock() {
449 let m = metrics.entry(label.clone()).or_default();
450 m.executions += 1;
451 m.successes += 1;
452 m.total_duration_ms += elapsed_ms;
453 }
454
455 if hook_result.is_halt() {
456 let action_str = if hook_result.is_stop() {
457 "stop"
458 } else {
459 "skip"
460 };
461 write_info_log(
462 "HookManager::execute",
463 &format!("Hook {} ({})", action_str, label),
464 );
465 return Some(HookResult {
466 action: Some(if hook_result.is_stop() {
467 HookAction::Stop
468 } else {
469 HookAction::Skip
470 }),
471 retry_feedback: hook_result.retry_feedback.clone(),
472 system_message: hook_result.system_message.clone(),
473 ..Default::default()
474 });
475 }
476
477 merge_hook_result_into(&hook_result, &mut context, &mut final_result);
479 had_modification = true;
480
481 last_outcome = Some(HookOutcome::Success(hook_result));
482 break; }
484 Err(e) => {
485 if let Ok(mut metrics) = self.metrics.lock() {
486 let m = metrics.entry(label.clone()).or_default();
487 m.executions += 1;
488 m.failures += 1;
489 m.total_duration_ms += elapsed_ms;
490 }
491
492 let attempts_left = max_attempts - attempt - 1;
493 if attempts_left > 0 {
494 write_info_log(
495 "HookManager::execute",
496 &format!(
497 "Hook 执行失败 ({}), 第 {}/{} 次尝试, 剩余重试 {}: {}",
498 label,
499 attempt + 1,
500 max_attempts,
501 attempts_left,
502 e
503 ),
504 );
505 last_outcome = Some(HookOutcome::Retry {
506 error: e,
507 attempts_left,
508 });
509 } else {
511 write_error_log(
512 "HookManager::execute",
513 &format!("Hook 执行失败 ({}), 重试耗尽: {}", label, e),
514 );
515 last_outcome = Some(HookOutcome::Err(e));
516 break; }
518 }
519 }
520 }
521
522 match last_outcome {
524 Some(HookOutcome::Success(_)) => {
525 }
527 Some(HookOutcome::Retry { error, .. }) => {
528 write_error_log(
530 "HookManager::execute",
531 &format!("Hook 重试未完成 ({}): {}", label, error),
532 );
533 if let Some(action) = handle_hook_error(hook_ref.kind, &label) {
534 return Some(action);
535 }
536 }
537 Some(HookOutcome::Err(e)) => {
538 write_error_log(
539 "HookManager::execute",
540 &format!("Hook 最终失败 ({}): {}", label, e),
541 );
542 if let Some(action) = handle_hook_error(hook_ref.kind, &label) {
543 return Some(action);
544 }
545 }
546 None => {
547 continue;
548 }
549 }
550 }
551
552 if had_modification {
553 Some(final_result)
554 } else {
555 None
556 }
557 }
558}
559
560struct HookRef<'a> {
564 kind: &'a HookKind,
565 source: &'static str,
566 session_index: Option<usize>,
567}
568
569fn collect_hooks_for_event(manager: &HookManager, event: HookEvent) -> Vec<HookRef<'_>> {
571 let mut all_hooks: Vec<HookRef<'_>> = Vec::new();
572
573 if let Some(hooks) = manager.builtin_hooks.get(&event) {
574 for h in hooks.iter() {
575 all_hooks.push(HookRef {
576 kind: h,
577 source: HOOK_SOURCE_BUILTIN,
578 session_index: None,
579 });
580 }
581 }
582 if let Some(hooks) = manager.user_hooks.get(&event) {
583 for h in hooks.iter() {
584 all_hooks.push(HookRef {
585 kind: h,
586 source: HOOK_SOURCE_USER,
587 session_index: None,
588 });
589 }
590 }
591 if let Some(hooks) = manager.project_hooks.get(&event) {
592 for h in hooks.iter() {
593 all_hooks.push(HookRef {
594 kind: h,
595 source: HOOK_SOURCE_PROJECT,
596 session_index: None,
597 });
598 }
599 }
600 if let Some(hooks) = manager.session_hooks.get(&event) {
601 for (idx, h) in hooks.iter().enumerate() {
602 all_hooks.push(HookRef {
603 kind: h,
604 source: HOOK_SOURCE_SESSION,
605 session_index: Some(idx),
606 });
607 }
608 }
609
610 all_hooks
611}
612
613fn merge_hook_result_into(
615 hook_result: &HookResult,
616 context: &mut HookContext,
617 final_result: &mut HookResult,
618) {
619 if let Some(ref msgs) = hook_result.messages {
620 context.messages = Some(msgs.clone());
621 final_result.messages = context.messages.clone();
622 }
623 if let Some(ref sp) = hook_result.system_prompt {
624 context.system_prompt = Some(sp.clone());
625 final_result.system_prompt = context.system_prompt.clone();
626 }
627 if let Some(ref ui) = hook_result.user_input {
628 context.user_input = Some(ui.clone());
629 final_result.user_input = context.user_input.clone();
630 }
631 if let Some(ref ao) = hook_result.assistant_output {
632 context.assistant_output = Some(ao.clone());
633 final_result.assistant_output = context.assistant_output.clone();
634 }
635 if let Some(ref ta) = hook_result.tool_arguments {
636 context.tool_arguments = Some(ta.clone());
637 final_result.tool_arguments = context.tool_arguments.clone();
638 }
639 if let Some(ref tr) = hook_result.tool_result {
640 context.tool_result = Some(tr.clone());
641 final_result.tool_result = context.tool_result.clone();
642 }
643 if let Some(ref inject) = hook_result.inject_messages {
644 let existing = final_result.inject_messages.get_or_insert_with(Vec::new);
645 existing.extend(inject.clone());
646 }
647 if let Some(ref rf) = hook_result.retry_feedback {
648 final_result.retry_feedback = Some(rf.clone());
649 }
650 if let Some(ref ac) = hook_result.additional_context {
651 final_result.additional_context = Some(ac.clone());
652 }
653 if let Some(ref sm) = hook_result.system_message {
654 final_result.system_message = Some(sm.clone());
655 }
656 if let Some(ref te) = hook_result.tool_error {
657 final_result.tool_error = Some(te.clone());
658 }
659}
660
661fn handle_hook_error(kind: &HookKind, _label: &str) -> Option<HookResult> {
664 match hook_on_error_strategy(kind) {
665 OnError::Stop => Some(HookResult {
666 action: Some(HookAction::Stop),
667 ..Default::default()
668 }),
669 OnError::Skip => None,
670 }
671}
672
673pub fn hook_unique_id(source: &str, kind: &HookKind, session_index: Option<usize>) -> String {
675 let key = match kind {
676 HookKind::Builtin(b) => b.name.clone(),
677 HookKind::Shell(s) => s
678 .name
679 .clone()
680 .unwrap_or_else(|| s.command.chars().take(40).collect()),
681 HookKind::Llm(l) => l
682 .name
683 .clone()
684 .unwrap_or_else(|| l.prompt.chars().take(40).collect()),
685 };
686 match session_index {
687 Some(idx) => format!("{}:{}", source, idx),
688 None => format!("{}:{}", source, key),
689 }
690}
691
692pub(crate) fn hook_name(kind: &HookKind) -> Option<&str> {
694 match kind {
695 HookKind::Shell(shell) => shell.name.as_deref(),
696 HookKind::Llm(llm) => llm.name.as_deref(),
697 HookKind::Builtin(builtin) => Some(&builtin.name),
698 }
699}
700
701pub(crate) fn hook_label(kind: &HookKind) -> String {
703 match kind {
704 HookKind::Shell(shell) => {
705 if let Some(ref name) = shell.name {
706 format!("{}: {}", name, shell.command)
707 } else {
708 shell.command.clone()
709 }
710 }
711 HookKind::Llm(llm) => {
712 let first_line = llm
714 .prompt
715 .lines()
716 .find(|l| !l.trim().is_empty())
717 .unwrap_or(&llm.prompt);
718 let prompt_preview = if first_line.len() > crate::constants::HOOK_PROMPT_PREVIEW_MAX_LEN
719 {
720 format!(
721 "{}...",
722 &first_line[..crate::constants::HOOK_PROMPT_PREVIEW_MAX_LEN]
723 )
724 } else {
725 first_line.to_string()
726 };
727 if let Some(ref name) = llm.name {
728 format!("[llm: {}] {}", name, prompt_preview)
729 } else {
730 format!("[llm: {}]", prompt_preview)
731 }
732 }
733 HookKind::Builtin(builtin) => format!("[builtin: {}]", builtin.name),
734 }
735}
736
737pub(crate) fn hook_type_str(kind: &HookKind) -> &'static str {
739 match kind {
740 HookKind::Shell(_) => "bash",
741 HookKind::Llm(_) => "llm",
742 HookKind::Builtin(_) => "builtin",
743 }
744}
745
746pub(crate) fn hook_timeout(kind: &HookKind) -> Option<u64> {
748 match kind {
749 HookKind::Shell(shell) => Some(shell.timeout),
750 HookKind::Llm(llm) => Some(llm.timeout),
751 HookKind::Builtin(_) => None,
752 }
753}
754
755pub(crate) fn hook_retry_count(kind: &HookKind) -> u32 {
757 match kind {
758 HookKind::Shell(shell) => shell.retry,
759 HookKind::Llm(llm) => llm.retry,
760 HookKind::Builtin(_) => 0,
761 }
762}
763
764pub(crate) fn hook_on_error(kind: &HookKind) -> Option<OnError> {
766 match kind {
767 HookKind::Shell(shell) => Some(shell.on_error),
768 HookKind::Llm(llm) => Some(llm.on_error),
769 HookKind::Builtin(_) => None,
770 }
771}
772
773pub(crate) fn hook_on_error_strategy(kind: &HookKind) -> OnError {
775 match kind {
776 HookKind::Shell(shell) => shell.on_error,
777 HookKind::Llm(llm) => llm.on_error,
778 HookKind::Builtin(_) => OnError::Stop,
779 }
780}
781
782pub(crate) fn hook_filter(kind: &HookKind) -> Option<&HookFilter> {
784 match kind {
785 HookKind::Shell(shell) if !shell.filter.is_empty() => Some(&shell.filter),
786 HookKind::Llm(llm) if !llm.filter.is_empty() => Some(&llm.filter),
787 HookKind::Shell(_) | HookKind::Llm(_) | HookKind::Builtin(_) => None,
788 }
789}
790
791pub(crate) fn hook_should_execute(kind: &HookKind, context: &HookContext) -> bool {
793 match kind {
794 HookKind::Shell(shell) => shell.filter.matches(context),
795 HookKind::Llm(llm) => llm.filter.matches(context),
796 HookKind::Builtin(_) => true,
797 }
798}