rightkit_hooks/
dispatch.rs

1//! Module fan-out under deadlines, with an exit-0 contract.
2//!
3//! Mirrors `@rightkit/hooks` `dispatchHookEvent` (serial modules in declared order, a typed
4//! per-module result, a per-module deadline that defaults to 3000 ms) and adds one total deadline,
5//! so two modules on one event share one budget. Nothing a module does can fail the host: an
6//! error, a panic or a timeout becomes a typed result, and the reply falls back to nothing.
7//!
8//! A module that times out is not killed: Rust cannot stop a thread. Its thread keeps running
9//! until the module returns (cooperative modules watch [`ModuleContext::is_cancelled`]) and any
10//! side effect it has, such as writing a receipt, can still land after the hook has replied.
11//!
12//! The only non-zero exit is a module's explicit [`Decision::DenyTool`] on Codex, which is exit 2
13//! with the reason on stderr (Legion's shipped shape).
14
15use crate::event::{EventKind, HookEvent, ParseError, Wire};
16use crate::output::Reply;
17use serde_json::Value;
18use std::io::{Read, Write};
19use std::panic::{catch_unwind, AssertUnwindSafe};
20use std::sync::atomic::{AtomicBool, Ordering};
21use std::sync::{mpsc, Arc};
22use std::time::{Duration, Instant};
23
24/// Per-module budget, the JS default.
25pub const DEFAULT_MODULE_DEADLINE: Duration = Duration::from_millis(3_000);
26/// Whole-dispatch budget, inside the 10 s timeout the Pulse plugins register.
27pub const DEFAULT_TOTAL_DEADLINE: Duration = Duration::from_millis(8_000);
28/// Merged context cap in bytes (Pulse's `TEXT_CAP`).
29pub const DEFAULT_CONTEXT_CAP_BYTES: usize = 60_000;
30
31/// A module's verdict on the event. Applied only where the event's schema takes it.
32#[derive(Clone, Debug, PartialEq, Eq)]
33pub enum Decision {
34    /// `Stop`: keep the run going with this reason.
35    BlockStop(String),
36    /// `PreToolUse`: refuse the call.
37    DenyTool(String),
38    /// `PreToolUse`: ask the person (Claude; Codex denies).
39    AskTool(String),
40}
41
42#[derive(Clone, Debug, Default, PartialEq)]
43pub struct ModuleOutput {
44    pub additional_context: Option<String>,
45    pub system_message: Option<String>,
46    pub decision: Option<Decision>,
47    /// Free-form, content-free detail for the result record.
48    pub detail: Value,
49}
50
51impl ModuleOutput {
52    pub fn empty() -> Self {
53        Self::default()
54    }
55
56    pub fn context(text: impl Into<String>) -> Self {
57        Self {
58            additional_context: Some(text.into()),
59            ..Self::default()
60        }
61    }
62
63    pub fn decide(decision: Decision) -> Self {
64        Self {
65            decision: Some(decision),
66            ..Self::default()
67        }
68    }
69}
70
71/// Passed to a module so it can stop early: the deadline, and a flag set once it has passed.
72#[derive(Clone, Debug)]
73pub struct ModuleContext {
74    pub deadline: Instant,
75    cancelled: Arc<AtomicBool>,
76}
77
78impl ModuleContext {
79    pub fn is_cancelled(&self) -> bool {
80        self.cancelled.load(Ordering::Relaxed)
81    }
82
83    pub fn remaining(&self) -> Duration {
84        self.deadline.saturating_duration_since(Instant::now())
85    }
86}
87
88pub trait HookModule: Send + Sync {
89    fn id(&self) -> &str;
90    fn handle(&self, event: &HookEvent, ctx: &ModuleContext) -> Result<ModuleOutput, String>;
91}
92
93/// A module from a closure.
94pub struct FnModule<F> {
95    id: String,
96    run: F,
97}
98
99impl<F> FnModule<F>
100where
101    F: Fn(&HookEvent, &ModuleContext) -> Result<ModuleOutput, String> + Send + Sync,
102{
103    pub fn new(id: impl Into<String>, run: F) -> Self {
104        Self { id: id.into(), run }
105    }
106}
107
108impl<F> HookModule for FnModule<F>
109where
110    F: Fn(&HookEvent, &ModuleContext) -> Result<ModuleOutput, String> + Send + Sync,
111{
112    fn id(&self) -> &str {
113        &self.id
114    }
115
116    fn handle(&self, event: &HookEvent, ctx: &ModuleContext) -> Result<ModuleOutput, String> {
117        (self.run)(event, ctx)
118    }
119}
120
121#[derive(Clone, Debug)]
122pub struct DispatchOptions {
123    pub module_deadline: Duration,
124    pub total_deadline: Duration,
125    pub context_cap_bytes: usize,
126    /// Printed when the payload cannot be parsed or no module produced anything while one failed.
127    pub fallback: Reply,
128}
129
130impl Default for DispatchOptions {
131    fn default() -> Self {
132        Self {
133            module_deadline: DEFAULT_MODULE_DEADLINE,
134            total_deadline: DEFAULT_TOTAL_DEADLINE,
135            context_cap_bytes: DEFAULT_CONTEXT_CAP_BYTES,
136            fallback: Reply::none(),
137        }
138    }
139}
140
141#[derive(Clone, Copy, Debug, PartialEq, Eq)]
142pub enum Status {
143    Ok,
144    Error,
145}
146
147#[derive(Clone, Debug, PartialEq)]
148pub struct ModuleResult {
149    pub id: String,
150    pub status: Status,
151    pub output: Option<ModuleOutput>,
152    pub error: Option<String>,
153    pub elapsed: Duration,
154}
155
156#[derive(Clone, Debug, PartialEq)]
157pub struct DispatchResult {
158    pub event: EventKind,
159    /// `Error` when any module failed, as in the JS result.
160    pub status: Status,
161    pub results: Vec<ModuleResult>,
162    pub reply: Reply,
163}
164
165fn run_module(
166    module: &Arc<dyn HookModule>,
167    event: &Arc<HookEvent>,
168    budget: Duration,
169) -> ModuleResult {
170    let started = Instant::now();
171    let id = module.id().to_owned();
172    let cancelled = Arc::new(AtomicBool::new(false));
173    let ctx = ModuleContext {
174        deadline: started + budget,
175        cancelled: cancelled.clone(),
176    };
177    let (tx, rx) = mpsc::channel();
178    let (module, event) = (module.clone(), event.clone());
179    let spawned = std::thread::Builder::new()
180        .name(format!("hook-module-{id}"))
181        .spawn(move || {
182            let result = catch_unwind(AssertUnwindSafe(|| module.handle(&event, &ctx)))
183                .unwrap_or_else(|_| Err("module panicked".to_owned()));
184            let _ = tx.send(result);
185        });
186    let outcome = match spawned {
187        Err(e) => Err(format!("module thread could not start: {e}")),
188        Ok(_) => match rx.recv_timeout(budget) {
189            Ok(result) => result,
190            Err(mpsc::RecvTimeoutError::Timeout) => {
191                cancelled.store(true, Ordering::Relaxed);
192                Err(format!("HookHost module timed out: {id}"))
193            }
194            Err(mpsc::RecvTimeoutError::Disconnected) => {
195                Err("module ended without a result".into())
196            }
197        },
198    };
199    let elapsed = started.elapsed();
200    match outcome {
201        Ok(output) => ModuleResult {
202            id,
203            status: Status::Ok,
204            output: Some(output),
205            error: None,
206            elapsed,
207        },
208        Err(error) => ModuleResult {
209            id,
210            status: Status::Error,
211            output: None,
212            error: Some(error),
213            elapsed,
214        },
215    }
216}
217
218fn cap_text(mut text: String, cap: usize) -> String {
219    if text.len() > cap {
220        let mut cut = cap;
221        while !text.is_char_boundary(cut) {
222            cut -= 1;
223        }
224        text.truncate(cut);
225    }
226    text
227}
228
229fn join(parts: Vec<String>) -> Option<String> {
230    let parts: Vec<String> = parts.into_iter().filter(|p| !p.is_empty()).collect();
231    (!parts.is_empty()).then(|| parts.join("\n\n"))
232}
233
234/// Merge module outputs into the one reply the event's schema allows.
235///
236/// A decision wins over context. Contexts are joined with a blank line and capped. A decision or
237/// context the event's schema does not take is dropped, never emitted.
238pub fn compose_reply(event: &HookEvent, results: &[ModuleResult], context_cap: usize) -> Reply {
239    let outputs = results.iter().filter_map(|r| r.output.as_ref());
240    let decision = outputs.clone().find_map(|o| o.decision.clone());
241    let system = join(
242        outputs
243            .clone()
244            .filter_map(|o| o.system_message.clone())
245            .collect(),
246    );
247    let context = join(
248        outputs
249            .filter_map(|o| o.additional_context.clone())
250            .collect(),
251    )
252    .map(|c| cap_text(c, context_cap));
253
254    let base = match (&event.kind, decision) {
255        (EventKind::Stop, Some(Decision::BlockStop(reason))) => Some(Reply::block_stop(&reason)),
256        (EventKind::PreToolUse, Some(Decision::DenyTool(reason))) => {
257            Some(Reply::deny_tool(event.wire, &reason))
258        }
259        (EventKind::PreToolUse, Some(Decision::AskTool(reason))) => {
260            Some(Reply::ask_tool(event.wire, &reason))
261        }
262        _ => None,
263    };
264    let base = base.or_else(|| {
265        context
266            .as_deref()
267            .and_then(|c| Reply::additional_context(&event.kind, c).ok())
268    });
269    match (base, system) {
270        (Some(reply), Some(message)) => reply.with_system_message(&message),
271        (Some(reply), None) => reply,
272        (None, _) => Reply::none(),
273    }
274}
275
276/// Run `modules` in order. Each gets `min(module_deadline, what is left of total_deadline)`;
277/// once the total is spent the rest are recorded as errors and not started.
278pub fn dispatch(
279    event: HookEvent,
280    modules: &[Arc<dyn HookModule>],
281    options: &DispatchOptions,
282) -> DispatchResult {
283    let started = Instant::now();
284    let event = Arc::new(event);
285    let mut results = Vec::with_capacity(modules.len());
286    for module in modules {
287        let left = options.total_deadline.saturating_sub(started.elapsed());
288        if left.is_zero() {
289            results.push(ModuleResult {
290                id: module.id().to_owned(),
291                status: Status::Error,
292                output: None,
293                error: Some("HookHost total deadline spent".into()),
294                elapsed: Duration::ZERO,
295            });
296            continue;
297        }
298        results.push(run_module(
299            module,
300            &event,
301            left.min(options.module_deadline),
302        ));
303    }
304    let failed = results.iter().any(|r| r.status == Status::Error);
305    let mut reply = compose_reply(&event, &results, options.context_cap_bytes);
306    if failed && reply.is_none() {
307        reply = options.fallback.clone();
308    }
309    DispatchResult {
310        event: event.kind.clone(),
311        status: if failed { Status::Error } else { Status::Ok },
312        results,
313        reply,
314    }
315}
316
317/// The whole hook process over arbitrary streams. Returns the exit code: 0 on every failure.
318pub fn run_with(
319    wire: Wire,
320    modules: &[Arc<dyn HookModule>],
321    options: &DispatchOptions,
322    input: impl Read,
323    output: impl Write,
324    errors: impl Write,
325) -> i32 {
326    let reply = match HookEvent::from_reader(wire, input) {
327        Err(error) => fallback_for(&error, options),
328        Ok(event) => catch_unwind(AssertUnwindSafe(|| dispatch(event, modules, options).reply))
329            .unwrap_or_else(|_| options.fallback.clone()),
330    };
331    reply.write_to(output, errors)
332}
333
334fn fallback_for(_error: &ParseError, options: &DispatchOptions) -> Reply {
335    options.fallback.clone()
336}
337
338/// [`run_with`] on the process's own stdin, stdout and stderr.
339pub fn run(wire: Wire, modules: &[Arc<dyn HookModule>], options: &DispatchOptions) -> i32 {
340    run_with(
341        wire,
342        modules,
343        options,
344        std::io::stdin().lock(),
345        std::io::stdout().lock(),
346        std::io::stderr().lock(),
347    )
348}
rightkit_hooks/dispatch.rs

rightkit_hooks/
dispatch.rs
