Skip to main content

pleiades_agent_plugins/
hooks.rs

1use std::ffi::OsStr;
2use std::io::Write;
3use std::process::Command;
4
5use crate::manifest::PluginHooks;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum HookEvent {
9    PreToolUse,
10    PostToolUse,
11    PostToolUseFailure,
12}
13
14impl HookEvent {
15    fn as_str(self) -> &'static str {
16        match self {
17            Self::PreToolUse => "PreToolUse",
18            Self::PostToolUse => "PostToolUse",
19            Self::PostToolUseFailure => "PostToolUseFailure",
20        }
21    }
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct HookRunResult {
26    denied: bool,
27    failed: bool,
28    messages: Vec<String>,
29}
30
31impl HookRunResult {
32    pub fn allow(messages: Vec<String>) -> Self {
33        Self {
34            denied: false,
35            failed: false,
36            messages,
37        }
38    }
39
40    pub fn is_denied(&self) -> bool {
41        self.denied
42    }
43
44    pub fn is_failed(&self) -> bool {
45        self.failed
46    }
47
48    pub fn messages(&self) -> &[String] {
49        &self.messages
50    }
51
52    pub fn into_messages(self) -> Vec<String> {
53        self.messages
54    }
55}
56
57#[derive(Debug, Clone, PartialEq, Eq)]
58pub struct HookRunner {
59    hooks: PluginHooks,
60}
61
62impl HookRunner {
63    pub fn new(hooks: PluginHooks) -> Self {
64        Self { hooks }
65    }
66
67    pub fn pre_tool_use(&self, tool_name: &str, tool_input: &str) -> HookRunResult {
68        Self::run_commands(
69            HookEvent::PreToolUse,
70            &self.hooks.pre_tool_use,
71            tool_name,
72            tool_input,
73            None,
74            false,
75        )
76    }
77
78    pub fn post_tool_use(
79        &self,
80        tool_name: &str,
81        tool_input: &str,
82        tool_output: &str,
83        is_error: bool,
84    ) -> HookRunResult {
85        Self::run_commands(
86            HookEvent::PostToolUse,
87            &self.hooks.post_tool_use,
88            tool_name,
89            tool_input,
90            Some(tool_output),
91            is_error,
92        )
93    }
94
95    pub fn post_tool_use_failure(
96        &self,
97        tool_name: &str,
98        tool_input: &str,
99        tool_error: &str,
100    ) -> HookRunResult {
101        Self::run_commands(
102            HookEvent::PostToolUseFailure,
103            &self.hooks.post_tool_use_failure,
104            tool_name,
105            tool_input,
106            Some(tool_error),
107            true,
108        )
109    }
110
111    fn run_commands(
112        event: HookEvent,
113        commands: &[String],
114        tool_name: &str,
115        tool_input: &str,
116        tool_output: Option<&str>,
117        is_error: bool,
118    ) -> HookRunResult {
119        if commands.is_empty() {
120            return HookRunResult::allow(Vec::new());
121        }
122
123        let payload = hook_payload(event, tool_name, tool_input, tool_output, is_error).to_string();
124        let mut messages = Vec::new();
125
126        for command in commands {
127            match run_single_command(
128                command,
129                event,
130                tool_name,
131                tool_input,
132                tool_output,
133                is_error,
134                &payload,
135            ) {
136                HookCommandOutcome::Allow { message } => {
137                    if let Some(msg) = message {
138                        messages.push(msg);
139                    }
140                }
141                HookCommandOutcome::Deny { message } => {
142                    messages.push(message.unwrap_or_else(|| {
143                        format!("{} hook denied tool `{tool_name}`", event.as_str())
144                    }));
145                    return HookRunResult {
146                        denied: true,
147                        failed: false,
148                        messages,
149                    };
150                }
151                HookCommandOutcome::Failed { message } => {
152                    messages.push(message);
153                    return HookRunResult {
154                        denied: false,
155                        failed: true,
156                        messages,
157                    };
158                }
159            }
160        }
161
162        HookRunResult::allow(messages)
163    }
164}
165
166enum HookCommandOutcome {
167    Allow { message: Option<String> },
168    Deny { message: Option<String> },
169    Failed { message: String },
170}
171
172fn run_single_command(
173    command: &str,
174    event: HookEvent,
175    tool_name: &str,
176    tool_input: &str,
177    tool_output: Option<&str>,
178    is_error: bool,
179    payload: &str,
180) -> HookCommandOutcome {
181    let mut child = build_shell_command(command);
182    child.stdin(std::process::Stdio::piped());
183    child.stdout(std::process::Stdio::piped());
184    child.stderr(std::process::Stdio::piped());
185
186    child.env("HOOK_EVENT", event.as_str());
187    child.env("HOOK_TOOL_NAME", tool_name);
188    child.env("HOOK_TOOL_INPUT", tool_input);
189    child.env("HOOK_TOOL_IS_ERROR", if is_error { "1" } else { "0" });
190    if let Some(output) = tool_output {
191        child.env("HOOK_TOOL_OUTPUT", output);
192    }
193
194    let mut spawned = match child.spawn() {
195        Ok(c) => c,
196        Err(e) => {
197            return HookCommandOutcome::Failed {
198                message: format!(
199                    "{} hook `{command}` failed to start for `{tool_name}`: {e}",
200                    event.as_str()
201                ),
202            };
203        }
204    };
205
206    if let Some(mut stdin) = spawned.stdin.take() {
207        let _ = stdin.write_all(payload.as_bytes());
208    }
209
210    match spawned.wait_with_output() {
211        Ok(output) => {
212            let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
213            let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
214            let message = (!stdout.is_empty()).then_some(stdout);
215
216            match output.status.code() {
217                Some(0) => HookCommandOutcome::Allow { message },
218                Some(2) => HookCommandOutcome::Deny { message },
219                Some(code) => {
220                    let detail = message
221                        .as_deref()
222                        .filter(|s| !s.is_empty())
223                        .or_else(|| (!stderr.is_empty()).then_some(stderr.as_str()))
224                        .unwrap_or("");
225                    HookCommandOutcome::Failed {
226                        message: if detail.is_empty() {
227                            format!("Hook `{command}` exited with status {code}")
228                        } else {
229                            format!("Hook `{command}` exited with status {code}: {detail}")
230                        },
231                    }
232                }
233                None => HookCommandOutcome::Failed {
234                    message: format!(
235                        "{} hook `{command}` was terminated by signal",
236                        event.as_str()
237                    ),
238                },
239            }
240        }
241        Err(e) => HookCommandOutcome::Failed {
242            message: format!(
243                "{} hook `{command}` failed for `{tool_name}`: {e}",
244                event.as_str()
245            ),
246        },
247    }
248}
249
250fn hook_payload(
251    event: HookEvent,
252    tool_name: &str,
253    tool_input: &str,
254    tool_output: Option<&str>,
255    is_error: bool,
256) -> serde_json::Value {
257    match event {
258        HookEvent::PostToolUseFailure => serde_json::json!({
259            "hook_event_name": event.as_str(),
260            "tool_name": tool_name,
261            "tool_input": parse_tool_input(tool_input),
262            "tool_input_json": tool_input,
263            "tool_error": tool_output,
264            "tool_result_is_error": true,
265        }),
266        _ => serde_json::json!({
267            "hook_event_name": event.as_str(),
268            "tool_name": tool_name,
269            "tool_input": parse_tool_input(tool_input),
270            "tool_input_json": tool_input,
271            "tool_output": tool_output,
272            "tool_result_is_error": is_error,
273        }),
274    }
275}
276
277fn parse_tool_input(tool_input: &str) -> serde_json::Value {
278    serde_json::from_str(tool_input).unwrap_or_else(|_| serde_json::json!({"raw": tool_input}))
279}
280
281fn build_shell_command(command: &str) -> Command {
282    let (shell, flag) = if cfg!(windows) {
283        ("cmd", "/C")
284    } else {
285        ("sh", "-c")
286    };
287    let mut cmd = Command::new(shell);
288    cmd.arg(flag);
289    cmd.arg(command);
290    cmd
291}
292
293pub trait HookCommandExt {
294    fn env<K, V>(&mut self, key: K, value: V) -> &mut Self
295    where
296        K: AsRef<OsStr>,
297        V: AsRef<OsStr>;
298}
299
300impl HookCommandExt for Command {
301    fn env<K, V>(&mut self, key: K, value: V) -> &mut Self
302    where
303        K: AsRef<OsStr>,
304        V: AsRef<OsStr>,
305    {
306        self.env(key, value);
307        self
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314
315    fn exit_command(code: u8) -> String {
316        if cfg!(windows) {
317            format!("exit /b {code}")
318        } else {
319            format!("exit {code}")
320        }
321    }
322
323    fn echo_and_exit(message: &str, code: u8) -> String {
324        if cfg!(windows) {
325            format!("echo {message} & exit /b {code}")
326        } else {
327            format!("printf '{message}'; exit {code}")
328        }
329    }
330
331    #[test]
332    fn pre_tool_use_allows_when_hooks_succeed() {
333        let runner = HookRunner::new(PluginHooks {
334            pre_tool_use: vec!["echo all good".to_string()],
335            post_tool_use: Vec::new(),
336            post_tool_use_failure: Vec::new(),
337        });
338
339        let result = runner.pre_tool_use("Read", r#"{"path":"test"}"#);
340
341        assert!(!result.is_denied());
342        assert!(!result.is_failed());
343        assert_eq!(result.messages(), &["all good"]);
344    }
345
346    #[test]
347    fn pre_tool_use_denies_when_hook_exits_two() {
348        let runner = HookRunner::new(PluginHooks {
349            pre_tool_use: vec![echo_and_exit("blocked", 2)],
350            post_tool_use: Vec::new(),
351            post_tool_use_failure: Vec::new(),
352        });
353
354        let result = runner.pre_tool_use("Bash", r#"{"command":"pwd"}"#);
355
356        assert!(result.is_denied());
357        assert_eq!(result.messages(), &["blocked"]);
358    }
359
360    #[test]
361    fn pre_tool_use_fails_on_nonzero_exit() {
362        let runner = HookRunner::new(PluginHooks {
363            pre_tool_use: vec![exit_command(1)],
364            post_tool_use: Vec::new(),
365            post_tool_use_failure: Vec::new(),
366        });
367
368        let result = runner.pre_tool_use("Bash", r#"{"command":"pwd"}"#);
369
370        assert!(result.is_failed());
371    }
372
373    #[test]
374    fn pre_tool_use_first_failure_stops_later_hooks() {
375        let runner = HookRunner::new(PluginHooks {
376            pre_tool_use: vec![exit_command(2), "echo should not run".to_string()],
377            post_tool_use: Vec::new(),
378            post_tool_use_failure: Vec::new(),
379        });
380
381        let result = runner.pre_tool_use("Read", "{}");
382
383        assert!(result.is_denied());
384        assert!(
385            !result
386                .messages()
387                .iter()
388                .any(|m| m.contains("should not run"))
389        );
390    }
391
392    #[test]
393    fn post_tool_use_runs_after_successful_tool() {
394        let runner = HookRunner::new(PluginHooks {
395            pre_tool_use: Vec::new(),
396            post_tool_use: vec!["echo tool completed".to_string()],
397            post_tool_use_failure: Vec::new(),
398        });
399
400        let result = runner.post_tool_use("Write", r#"{"path":"f"}"#, "wrote 1 file", false);
401
402        assert!(!result.is_denied());
403        assert!(
404            result
405                .messages()
406                .iter()
407                .any(|m| m.contains("tool completed"))
408        );
409    }
410
411    #[test]
412    fn post_tool_use_failure_runs_on_tool_error() {
413        let runner = HookRunner::new(PluginHooks {
414            pre_tool_use: Vec::new(),
415            post_tool_use: Vec::new(),
416            post_tool_use_failure: vec!["echo tool errored".to_string()],
417        });
418
419        let result =
420            runner.post_tool_use_failure("Bash", r#"{"command":"invalid"}"#, "command not found");
421
422        assert!(!result.is_denied());
423        assert!(result.messages().iter().any(|m| m.contains("tool errored")));
424    }
425
426    #[test]
427    fn empty_hooks_returns_allow_with_no_messages() {
428        let runner = HookRunner::new(PluginHooks::default());
429
430        let result = runner.pre_tool_use("Read", "{}");
431
432        assert!(!result.is_denied());
433        assert!(!result.is_failed());
434        assert!(result.messages().is_empty());
435    }
436}