forge_sdk/hooks/
registry.rs

1//! Hook registry - stores callbacks, mints opaque `callback_id`s, and
2//! renders the `hooks` field of the `initialize` `control_request` payload.
3
4use std::marker::PhantomData;
5use std::sync::Arc;
6
7use super::callback::{ErasedCallback, ErasedHookCallback, HookCallback};
8use forge_primitives::hooks::HookKind;
9use forge_primitives::hooks::inputs::{
10    NotificationInput, PermissionRequestInput, PostToolUseFailureInput, PostToolUseInput,
11    PreCompactInput, PreToolUseInput, StopInput, SubagentStartInput, SubagentStopInput,
12    UserPromptSubmitInput,
13};
14
15/// Registry of hook callbacks. Construct with [`HooksBuilder`], attach to
16/// `OptionsBuilder` via `.hooks(...)`.
17#[derive(Default, Clone)]
18pub struct Hooks {
19    pub(crate) pre_tool_use: Vec<(String, Arc<dyn ErasedHookCallback>)>,
20    pub(crate) post_tool_use: Vec<(String, Arc<dyn ErasedHookCallback>)>,
21    pub(crate) post_tool_use_failure: Vec<(String, Arc<dyn ErasedHookCallback>)>,
22    pub(crate) user_prompt_submit: Vec<Arc<dyn ErasedHookCallback>>,
23    pub(crate) stop: Vec<Arc<dyn ErasedHookCallback>>,
24    pub(crate) subagent_stop: Vec<Arc<dyn ErasedHookCallback>>,
25    pub(crate) subagent_start: Vec<Arc<dyn ErasedHookCallback>>,
26    pub(crate) pre_compact: Vec<Arc<dyn ErasedHookCallback>>,
27    pub(crate) notification: Vec<Arc<dyn ErasedHookCallback>>,
28    pub(crate) permission_request: Vec<(String, Arc<dyn ErasedHookCallback>)>,
29    /// Timeout (seconds) the CLI should apply to every hook callback.
30    /// `None` means "use the default forge-sdk emits" (120 seconds).
31    /// Overridable via
32    /// [`HooksBuilder::default_timeout_secs`] for scenarios that need
33    /// to provoke `control_cancel_request` on slow callbacks.
34    pub(crate) default_timeout_secs: Option<u64>,
35}
36
37impl std::fmt::Debug for Hooks {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        f.debug_struct("Hooks")
40            .field("pre_tool_use_count", &self.pre_tool_use.len())
41            .field("post_tool_use_count", &self.post_tool_use.len())
42            .field("post_tool_use_failure_count", &self.post_tool_use_failure.len())
43            .field("user_prompt_submit_count", &self.user_prompt_submit.len())
44            .field("stop_count", &self.stop.len())
45            .field("subagent_stop_count", &self.subagent_stop.len())
46            .field("subagent_start_count", &self.subagent_start.len())
47            .field("pre_compact_count", &self.pre_compact.len())
48            .field("notification_count", &self.notification.len())
49            .field("permission_request_count", &self.permission_request.len())
50            .field("default_timeout_secs", &self.default_timeout_secs)
51            .finish()
52    }
53}
54
55impl Hooks {
56    /// Mint opaque `callback_id`s (e.g. `hook_0`, `hook_1`, …) for every
57    /// registered callback and return them bundled for dispatch-time
58    /// lookup. Returns a map `id → callback` plus a parallel metadata
59    /// vector describing which event/matcher each id belongs to (used to
60    /// populate the initialize `control_request` payload).
61    pub(crate) fn mint_registry(&self) -> HookRegistry {
62        let mut registry = HookRegistry {
63            default_timeout_secs: self.default_timeout_secs,
64            ..HookRegistry::default()
65        };
66        let mut counter: u64 = 0;
67
68        let mut mint =
69            |kind: HookKind, matcher: Option<String>, cb: Arc<dyn ErasedHookCallback>| {
70                let id = format!("hook_{counter}");
71                counter += 1;
72                registry.metadata.push(HookRegistryEntry { id: id.clone(), kind, matcher });
73                registry.by_id.insert(id, cb);
74            };
75
76        for (matcher, cb) in &self.pre_tool_use {
77            mint(HookKind::PreToolUse, Some(matcher.clone()), cb.clone());
78        }
79        for (matcher, cb) in &self.post_tool_use {
80            mint(HookKind::PostToolUse, Some(matcher.clone()), cb.clone());
81        }
82        for (matcher, cb) in &self.post_tool_use_failure {
83            mint(HookKind::PostToolUseFailure, Some(matcher.clone()), cb.clone());
84        }
85        for cb in &self.user_prompt_submit {
86            mint(HookKind::UserPromptSubmit, None, cb.clone());
87        }
88        for cb in &self.stop {
89            mint(HookKind::Stop, None, cb.clone());
90        }
91        for cb in &self.subagent_stop {
92            mint(HookKind::SubagentStop, None, cb.clone());
93        }
94        for cb in &self.subagent_start {
95            mint(HookKind::SubagentStart, None, cb.clone());
96        }
97        for cb in &self.pre_compact {
98            mint(HookKind::PreCompact, None, cb.clone());
99        }
100        for cb in &self.notification {
101            mint(HookKind::Notification, None, cb.clone());
102        }
103        for (matcher, cb) in &self.permission_request {
104            mint(HookKind::PermissionRequest, Some(matcher.clone()), cb.clone());
105        }
106
107        registry
108    }
109}
110
111/// Internal bundle mapping opaque ids to erased callbacks, with parallel
112/// metadata for the initialize payload.
113#[derive(Default)]
114pub(crate) struct HookRegistry {
115    pub(crate) by_id: std::collections::HashMap<String, Arc<dyn ErasedHookCallback>>,
116    pub(crate) metadata: Vec<HookRegistryEntry>,
117    pub(crate) default_timeout_secs: Option<u64>,
118}
119
120impl HookRegistry {
121    /// Render the `hooks` field of the `initialize` `control_request`:
122    /// `{"PreToolUse": [{"matcher": "...", "hookCallbackIds": ["hook_0"], "timeout": 120}, ...], ...}`.
123    pub(crate) fn to_initialize_payload(&self) -> serde_json::Value {
124        use std::collections::BTreeMap;
125
126        // Group ids by (kind, matcher). BTreeMap for deterministic output.
127        let mut by_kind: BTreeMap<&'static str, BTreeMap<String, Vec<String>>> = BTreeMap::new();
128        for entry in &self.metadata {
129            let kind_name = entry.kind.as_str();
130            let matcher_key = entry.matcher.clone().unwrap_or_default();
131            by_kind
132                .entry(kind_name)
133                .or_default()
134                .entry(matcher_key)
135                .or_default()
136                .push(entry.id.clone());
137        }
138
139        let mut map = serde_json::Map::new();
140        for (kind_name, matcher_group) in by_kind {
141            let specs: Vec<serde_json::Value> = matcher_group
142                .into_iter()
143                .map(|(matcher, ids)| {
144                    let mut spec = serde_json::Map::new();
145                    if !matcher.is_empty() {
146                        spec.insert("matcher".into(), serde_json::Value::String(matcher));
147                    }
148                    spec.insert(
149                        "hookCallbackIds".into(),
150                        serde_json::Value::Array(ids.into_iter().map(Into::into).collect()),
151                    );
152                    spec.insert(
153                        "timeout".into(),
154                        serde_json::json!(
155                            self.default_timeout_secs.unwrap_or(DEFAULT_HOOK_TIMEOUT_SECS)
156                        ),
157                    );
158                    serde_json::Value::Object(spec)
159                })
160                .collect();
161            map.insert(kind_name.into(), serde_json::Value::Array(specs));
162        }
163        serde_json::Value::Object(map)
164    }
165}
166
167/// Timeout (seconds) emitted per hook matcher when
168/// [`HooksBuilder::default_timeout_secs`] is unset. Raised from the
169/// CLI's per-matcher default of 30 so a hook response already written
170/// by forge survives a CLI stdin pump wedged on an inline
171/// `get_context_usage` computation.
172pub(crate) const DEFAULT_HOOK_TIMEOUT_SECS: u64 = 120;
173
174/// One entry describing a minted hook id.
175#[derive(Debug, Clone)]
176pub(crate) struct HookRegistryEntry {
177    pub(crate) id: String,
178    pub(crate) kind: HookKind,
179    pub(crate) matcher: Option<String>,
180}
181
182/// Builder for [`Hooks`].
183#[derive(Default)]
184pub struct HooksBuilder {
185    inner: Hooks,
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use crate::hooks::HookDecision;
192
193    /// Pins the deliberate hook-budget raise: the initialize payload
194    /// must emit the raised default, not the CLI's per-matcher 30.
195    #[test]
196    fn initialize_payload_emits_the_deliberate_hook_timeout() {
197        let hooks = HooksBuilder::new()
198            .pre_tool_use("*", |_input: PreToolUseInput, _ctx| async {
199                HookDecision::passthrough()
200            })
201            .build();
202        let payload = hooks.mint_registry().to_initialize_payload();
203        assert_eq!(
204            payload["PreToolUse"][0]["timeout"],
205            serde_json::json!(120),
206            "the raised hook budget must reach the wire"
207        );
208    }
209}
210
211impl std::fmt::Debug for HooksBuilder {
212    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
213        f.debug_struct("HooksBuilder").field("inner", &self.inner).finish()
214    }
215}
216
217impl HooksBuilder {
218    /// Start empty.
219    pub fn new() -> Self {
220        Self::default()
221    }
222
223    /// Override the per-hook timeout (seconds) emitted in the
224    /// initialize `control_request`. The CLI uses this to decide how
225    /// long to wait on each `hook_callback` reply before giving up and
226    /// emitting a `control_cancel_request`. Default when unset is
227    /// 120 seconds. Lower it for scenarios
228    /// that deliberately provoke cancellation; raise it for
229    /// long-running callbacks (e.g. LLM sub-calls inside a hook).
230    pub fn default_timeout_secs(mut self, secs: u64) -> Self {
231        self.inner.default_timeout_secs = Some(secs);
232        self
233    }
234
235    /// Register a `PreToolUse` hook. `matcher` is a glob against tool names
236    /// (`"*"` matches all); pass `"Bash"` to match only the Bash tool.
237    pub fn pre_tool_use<C>(mut self, matcher: impl Into<String>, callback: C) -> Self
238    where
239        C: HookCallback<PreToolUseInput> + 'static,
240    {
241        self.inner.pre_tool_use.push((
242            matcher.into(),
243            Arc::new(ErasedCallback::<PreToolUseInput, C> {
244                inner: callback,
245                _marker: PhantomData,
246            }),
247        ));
248        self
249    }
250
251    /// Register a `PostToolUse` hook.
252    pub fn post_tool_use<C>(mut self, matcher: impl Into<String>, callback: C) -> Self
253    where
254        C: HookCallback<PostToolUseInput> + 'static,
255    {
256        self.inner.post_tool_use.push((
257            matcher.into(),
258            Arc::new(ErasedCallback::<PostToolUseInput, C> {
259                inner: callback,
260                _marker: PhantomData,
261            }),
262        ));
263        self
264    }
265
266    /// Register a `UserPromptSubmit` hook.
267    pub fn user_prompt_submit<C>(mut self, callback: C) -> Self
268    where
269        C: HookCallback<UserPromptSubmitInput> + 'static,
270    {
271        self.inner.user_prompt_submit.push(Arc::new(ErasedCallback::<UserPromptSubmitInput, C> {
272            inner: callback,
273            _marker: PhantomData,
274        }));
275        self
276    }
277
278    /// Register a `Stop` hook.
279    pub fn stop<C>(mut self, callback: C) -> Self
280    where
281        C: HookCallback<StopInput> + 'static,
282    {
283        self.inner.stop.push(Arc::new(ErasedCallback::<StopInput, C> {
284            inner: callback,
285            _marker: PhantomData,
286        }));
287        self
288    }
289
290    /// Register a `SubagentStop` hook.
291    pub fn subagent_stop<C>(mut self, callback: C) -> Self
292    where
293        C: HookCallback<SubagentStopInput> + 'static,
294    {
295        self.inner.subagent_stop.push(Arc::new(ErasedCallback::<SubagentStopInput, C> {
296            inner: callback,
297            _marker: PhantomData,
298        }));
299        self
300    }
301
302    /// Register a `PreCompact` hook.
303    pub fn pre_compact<C>(mut self, callback: C) -> Self
304    where
305        C: HookCallback<PreCompactInput> + 'static,
306    {
307        self.inner.pre_compact.push(Arc::new(ErasedCallback::<PreCompactInput, C> {
308            inner: callback,
309            _marker: PhantomData,
310        }));
311        self
312    }
313
314    /// Register a `PostToolUseFailure` hook. `matcher` follows the same
315    /// tool-name glob semantics as [`Self::pre_tool_use`] / [`Self::post_tool_use`].
316    pub fn post_tool_use_failure<C>(mut self, matcher: impl Into<String>, callback: C) -> Self
317    where
318        C: HookCallback<PostToolUseFailureInput> + 'static,
319    {
320        self.inner.post_tool_use_failure.push((
321            matcher.into(),
322            Arc::new(ErasedCallback::<PostToolUseFailureInput, C> {
323                inner: callback,
324                _marker: PhantomData,
325            }),
326        ));
327        self
328    }
329
330    /// Register a `Notification` hook.
331    pub fn notification<C>(mut self, callback: C) -> Self
332    where
333        C: HookCallback<NotificationInput> + 'static,
334    {
335        self.inner.notification.push(Arc::new(ErasedCallback::<NotificationInput, C> {
336            inner: callback,
337            _marker: PhantomData,
338        }));
339        self
340    }
341
342    /// Register a `SubagentStart` hook.
343    pub fn subagent_start<C>(mut self, callback: C) -> Self
344    where
345        C: HookCallback<SubagentStartInput> + 'static,
346    {
347        self.inner.subagent_start.push(Arc::new(ErasedCallback::<SubagentStartInput, C> {
348            inner: callback,
349            _marker: PhantomData,
350        }));
351        self
352    }
353
354    /// Register a `PermissionRequest` hook (observational; `matcher` globs
355    /// against tool names the same way as [`Self::pre_tool_use`]).
356    pub fn permission_request<C>(mut self, matcher: impl Into<String>, callback: C) -> Self
357    where
358        C: HookCallback<PermissionRequestInput> + 'static,
359    {
360        self.inner.permission_request.push((
361            matcher.into(),
362            Arc::new(ErasedCallback::<PermissionRequestInput, C> {
363                inner: callback,
364                _marker: PhantomData,
365            }),
366        ));
367        self
368    }
369
370    /// Finalise.
371    pub fn build(self) -> Hooks {
372        self.inner
373    }
374}
375
376#[cfg(test)]
377mod tests_hooks_registration {
378    // Test-mod `use super::*;` brings the parent's full surface in; not every test consumes every item.
379    #[allow(unused_imports)]
380    use super::*;
381
382    use crate::{HookContext, HookDecision, HooksBuilder, OptionsBuilder, PreToolUseInput};
383    use crate::{
384        NotificationInput, PermissionRequestInput, PostToolUseFailureInput, SubagentStartInput,
385    };
386
387    #[test]
388    fn hooks_attach_to_options() {
389        let hooks = HooksBuilder::new()
390            .pre_tool_use("*", |_input: PreToolUseInput, _ctx: HookContext| async move {
391                HookDecision::allow()
392            })
393            .pre_tool_use("Bash", |_input: PreToolUseInput, _ctx: HookContext| async move {
394                HookDecision::deny("no bash")
395            })
396            .build();
397
398        let opts = OptionsBuilder::new().hooks(hooks).build();
399        let desc = format!("{opts:?}");
400        assert!(
401            desc.contains("pre_tool_use_count: 2"),
402            "expected pre_tool_use_count: 2, got: {desc}"
403        );
404    }
405
406    #[test]
407    fn all_ten_hook_kinds_register_and_count() {
408        // Confirm the builder exposes a registration method for every HookKind
409        // variant the CLI emits, and that each ends up in the initialize payload.
410        let hooks = HooksBuilder::new()
411            .pre_tool_use("Bash", |_i: PreToolUseInput, _c| async move { HookDecision::allow() })
412            .post_tool_use("Bash", |_i: crate::PostToolUseInput, _c| async move {
413                HookDecision::passthrough()
414            })
415            .post_tool_use_failure("Bash", |_i: PostToolUseFailureInput, _c| async move {
416                HookDecision::passthrough()
417            })
418            .user_prompt_submit(|_i: crate::UserPromptSubmitInput, _c| async move {
419                HookDecision::allow()
420            })
421            .stop(|_i: crate::StopInput, _c| async move { HookDecision::passthrough() })
422            .subagent_stop(
423                |_i: crate::SubagentStopInput, _c| async move { HookDecision::passthrough() },
424            )
425            .subagent_start(|_i: SubagentStartInput, _c| async move { HookDecision::passthrough() })
426            .pre_compact(
427                |_i: crate::PreCompactInput, _c| async move { HookDecision::passthrough() },
428            )
429            .notification(|_i: NotificationInput, _c| async move { HookDecision::passthrough() })
430            .permission_request("*", |_i: PermissionRequestInput, _c| async move {
431                HookDecision::passthrough()
432            })
433            .build();
434
435        let desc = format!("{hooks:?}");
436        // All counts should be 1.
437        for field in [
438            "pre_tool_use_count: 1",
439            "post_tool_use_count: 1",
440            "post_tool_use_failure_count: 1",
441            "user_prompt_submit_count: 1",
442            "stop_count: 1",
443            "subagent_stop_count: 1",
444            "subagent_start_count: 1",
445            "pre_compact_count: 1",
446            "notification_count: 1",
447            "permission_request_count: 1",
448        ] {
449            assert!(desc.contains(field), "missing {field} in debug: {desc}");
450        }
451    }
452
453    #[test]
454    fn initialize_payload_lists_every_registered_kind() {
455        // The initialize control_request's hooks payload groups by event name.
456        // Every registered kind should appear as a key.
457        let hooks = HooksBuilder::new()
458            .pre_tool_use("Bash", |_i: PreToolUseInput, _c| async move { HookDecision::allow() })
459            .post_tool_use_failure("*", |_i: PostToolUseFailureInput, _c| async move {
460                HookDecision::passthrough()
461            })
462            .subagent_start(|_i: SubagentStartInput, _c| async move { HookDecision::passthrough() })
463            .notification(|_i: NotificationInput, _c| async move { HookDecision::passthrough() })
464            .permission_request("Bash", |_i: PermissionRequestInput, _c| async move {
465                HookDecision::passthrough()
466            })
467            .build();
468
469        let payload = hooks.mint_registry().to_initialize_payload();
470        for key in [
471            "PreToolUse",
472            "PostToolUseFailure",
473            "SubagentStart",
474            "Notification",
475            "PermissionRequest",
476        ] {
477            assert!(
478                payload.get(key).is_some(),
479                "expected `{key}` in initialize payload: {payload}"
480            );
481        }
482    }
483}
forge_sdk/hooks/registry.rs

forge_sdk/hooks/
registry.rs
