Skip to main content

walrus_core/runtime/
hook.rs

1//! Hook trait — lifecycle backend for agent building, event observation,
2//! and tool registration.
3//!
4//! All hook crates implement this trait. [`Runtime`](crate) calls these
5//! methods at the appropriate lifecycle points. `DaemonHook` composes
6//! multiple Hook implementations by delegating to each.
7
8use crate::{AgentConfig, AgentEvent, Memory, agent::tool::ToolRegistry, memory::tools};
9use std::{future::Future, sync::Arc};
10
11/// Lifecycle backend for agent building, event observation, and tool registration.
12///
13/// Default implementations are no-ops so implementors only override what they need.
14pub trait Hook: Send + Sync {
15    /// Called by `Runtime::add_agent()` before building the `Agent`.
16    ///
17    /// Enriches the agent config: appends skill instructions, injects memory
18    /// into the system prompt, etc. The returned config is passed to `AgentBuilder`.
19    ///
20    /// Default: returns config unchanged.
21    fn on_build_agent(&self, config: AgentConfig) -> AgentConfig {
22        config
23    }
24
25    /// Called by Runtime after each agent step during execution.
26    ///
27    /// Receives every `AgentEvent` produced during `send_to` and `stream_to`.
28    /// Use for logging, metrics, persistence, or forwarding.
29    ///
30    /// Default: no-op.
31    fn on_event(&self, _agent: &str, _event: &AgentEvent) {}
32
33    /// Called by `Runtime::new()` to register tools into the shared registry.
34    ///
35    /// Implementations insert `(Tool, Handler)` pairs via `tools.insert()`.
36    /// `DaemonHook` delegates to each sub-hook's `on_register_tools`.
37    ///
38    /// Default: no-op async.
39    fn on_register_tools(&self, _tools: &mut ToolRegistry) -> impl Future<Output = ()> + Send {
40        async {}
41    }
42}
43
44impl Hook for () {}
45
46/// Blanket Hook impl for all Memory types that are Clone + 'static.
47///
48/// Injects compiled memory into the system prompt via `on_build_agent`
49/// and registers `remember`/`recall` tools via `on_register_tools`.
50impl<M: Memory + Clone + 'static> Hook for M {
51    fn on_build_agent(&self, mut config: AgentConfig) -> AgentConfig {
52        let compiled = self.compile();
53        if !compiled.is_empty() {
54            config.system_prompt = format!("{}\n\n{compiled}", config.system_prompt);
55        }
56        config
57    }
58
59    fn on_register_tools(&self, registry: &mut ToolRegistry) -> impl Future<Output = ()> + Send {
60        let mem = Arc::new(self.clone());
61        let (tool, handler) = tools::remember(Arc::clone(&mem));
62        registry.insert(tool, handler);
63        let (tool, handler) = tools::recall(mem);
64        registry.insert(tool, handler);
65        async {}
66    }
67}