1use futures_util::future::BoxFuture;
12use serde_json::Value;
13
14pub const HOOK_PAYLOAD_CAP: usize = 2048;
16
17#[derive(Debug, Clone, PartialEq)]
21pub enum PreToolDecision {
22 Continue,
23 Deny { message: String },
24 Rewrite { input: Value },
25}
26
27pub trait Hooks: Send + Sync {
28 fn pre_tool<'a>(&'a self, name: &'a str, input: &'a Value) -> BoxFuture<'a, PreToolDecision>;
30 fn post_tool<'a>(&'a self, name: &'a str, result: &'a str) -> BoxFuture<'a, Option<String>>;
33}
34
35pub fn cap_payload(s: &str) -> &str {
37 if s.len() <= HOOK_PAYLOAD_CAP {
38 return s;
39 }
40 let mut end = HOOK_PAYLOAD_CAP;
41 while !s.is_char_boundary(end) {
42 end -= 1;
43 }
44 &s[..end]
45}
46
47#[derive(Default)]
50pub struct InProcessHooks {
51 #[allow(clippy::type_complexity)]
52 pre: Vec<Box<dyn Fn(&str, &Value) -> PreToolDecision + Send + Sync>>,
53 #[allow(clippy::type_complexity)]
54 post: Vec<Box<dyn Fn(&str, &str) -> Option<String> + Send + Sync>>,
55}
56
57impl InProcessHooks {
58 pub fn new() -> Self {
59 Self::default()
60 }
61 pub fn on_pre_tool(
62 mut self,
63 f: impl Fn(&str, &Value) -> PreToolDecision + Send + Sync + 'static,
64 ) -> Self {
65 self.pre.push(Box::new(f));
66 self
67 }
68 pub fn on_post_tool(
69 mut self,
70 f: impl Fn(&str, &str) -> Option<String> + Send + Sync + 'static,
71 ) -> Self {
72 self.post.push(Box::new(f));
73 self
74 }
75 pub fn is_empty(&self) -> bool {
76 self.pre.is_empty() && self.post.is_empty()
77 }
78}
79
80impl Hooks for InProcessHooks {
81 fn pre_tool<'a>(&'a self, name: &'a str, input: &'a Value) -> BoxFuture<'a, PreToolDecision> {
82 Box::pin(async move {
83 for hook in &self.pre {
85 match hook(name, input) {
86 PreToolDecision::Continue => {}
87 decision => return decision,
88 }
89 }
90 PreToolDecision::Continue
91 })
92 }
93 fn post_tool<'a>(&'a self, name: &'a str, result: &'a str) -> BoxFuture<'a, Option<String>> {
94 Box::pin(async move {
95 let capped = cap_payload(result);
96 let mut current: Option<String> = None;
97 for hook in &self.post {
98 let view = current.as_deref().unwrap_or(capped);
99 if let Some(replacement) = hook(name, view) {
100 current = Some(replacement);
101 }
102 }
103 current
104 })
105 }
106}
107
108#[cfg(test)]
109mod tests {
110 use super::*;
111 use serde_json::json;
112
113 #[tokio::test]
114 async fn pre_hook_denies_rewrites_or_continues() {
115 let hooks = InProcessHooks::new().on_pre_tool(|name, input| {
116 if name == "bash" && input.get("command").and_then(Value::as_str) == Some("rm -rf /") {
117 PreToolDecision::Deny {
118 message: "no destructive commands".into(),
119 }
120 } else if name == "write" {
121 PreToolDecision::Rewrite {
122 input: json!({"path": "safe.txt", "content": "x"}),
123 }
124 } else {
125 PreToolDecision::Continue
126 }
127 });
128 assert_eq!(
129 hooks
130 .pre_tool("bash", &json!({"command": "rm -rf /"}))
131 .await,
132 PreToolDecision::Deny {
133 message: "no destructive commands".into()
134 }
135 );
136 assert!(matches!(
137 hooks
138 .pre_tool("write", &json!({"path": "x", "content": "y"}))
139 .await,
140 PreToolDecision::Rewrite { .. }
141 ));
142 assert_eq!(
143 hooks.pre_tool("read", &json!({})).await,
144 PreToolDecision::Continue
145 );
146 }
147
148 #[tokio::test]
149 async fn post_hook_caps_and_replaces() {
150 let hooks = InProcessHooks::new()
151 .on_post_tool(|_n, result| Some(format!("[annotated] {} chars", result.len())));
152 let big = "z".repeat(HOOK_PAYLOAD_CAP * 2);
153 let out = hooks.post_tool("read", &big).await.unwrap();
154 assert!(out.contains(&format!("{} chars", HOOK_PAYLOAD_CAP)));
156 }
157}