clawft_kernel/wasm_runner/
registry.rs

1//! Tool registry: trait, hierarchical lookup, WASM adapter, and signing.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use super::types::*;
7use super::runner::WasmToolRunner;
8use crate::governance::EffectVector;
9
10// ---------------------------------------------------------------------------
11// Built-in tool trait
12// ---------------------------------------------------------------------------
13
14/// Trait for built-in kernel tools.
15///
16/// Each tool has a spec and an execute method. Tools hold their own
17/// dependencies (e.g. `Arc<ProcessTable>` for agent tools).
18pub trait BuiltinTool: Send + Sync {
19    /// Return the tool name (e.g. "fs.read_file").
20    fn name(&self) -> &str;
21    /// Return the tool specification.
22    fn spec(&self) -> &BuiltinToolSpec;
23    /// Execute the tool with the given JSON arguments.
24    fn execute(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError>;
25}
26
27// ---------------------------------------------------------------------------
28// Tool Registry
29// ---------------------------------------------------------------------------
30
31/// Registry of available tools for dispatch.
32///
33/// Supports hierarchical lookup: a child registry can overlay a parent.
34/// The parent chain is walked when a tool is not found locally.
35pub struct ToolRegistry {
36    tools: HashMap<String, Arc<dyn BuiltinTool>>,
37    /// Optional parent registry for hierarchical lookup.
38    parent: Option<Arc<ToolRegistry>>,
39    /// When true, only signed tools may be registered.
40    require_signatures: bool,
41    /// Trusted public keys for signature verification (32-byte Ed25519 keys).
42    trusted_keys: Vec<[u8; 32]>,
43    /// Signatures for registered tools (tool_name -> ToolSignature).
44    signatures: HashMap<String, ToolSignature>,
45}
46
47impl ToolRegistry {
48    /// Create an empty registry with no parent.
49    pub fn new() -> Self {
50        Self {
51            tools: HashMap::new(),
52            parent: None,
53            require_signatures: false,
54            trusted_keys: Vec::new(),
55            signatures: HashMap::new(),
56        }
57    }
58
59    /// Create a child registry that delegates to `parent` for missing tools.
60    pub fn with_parent(parent: Arc<ToolRegistry>) -> Self {
61        Self {
62            tools: HashMap::new(),
63            parent: Some(parent),
64            require_signatures: false,
65            trusted_keys: Vec::new(),
66            signatures: HashMap::new(),
67        }
68    }
69
70    /// Register a tool (local to this registry level).
71    ///
72    /// When `require_signatures` is enabled, this rejects unsigned tools
73    /// with [`ToolError::SignatureRequired`]. Use [`register_signed`]
74    /// to supply a signature, or disable the requirement.
75    pub fn register(&mut self, tool: Arc<dyn BuiltinTool>) {
76        if self.require_signatures {
77            tracing::warn!(
78                tool = tool.name(),
79                "unsigned tool registration rejected (require_signatures=true)"
80            );
81            return;
82        }
83        self.tools.insert(tool.name().to_string(), tool);
84    }
85
86    /// Register a tool, checking signatures when required.
87    ///
88    /// Returns `Err(SignatureRequired)` when `require_signatures` is on
89    /// and no signature is provided. Returns `Ok(())` otherwise.
90    pub fn try_register(&mut self, tool: Arc<dyn BuiltinTool>) -> Result<(), ToolError> {
91        if self.require_signatures {
92            return Err(ToolError::SignatureRequired(tool.name().to_string()));
93        }
94        self.tools.insert(tool.name().to_string(), tool);
95        Ok(())
96    }
97
98    /// Register a tool with a cryptographic signature.
99    ///
100    /// Verifies the signature against trusted keys before allowing registration.
101    /// The signature is stored and the tool is chain-logged if ExoChain is available.
102    pub fn register_signed(
103        &mut self,
104        tool: Arc<dyn BuiltinTool>,
105        signature: ToolSignature,
106    ) -> Result<(), ToolError> {
107        // Verify the signature against at least one trusted key.
108        if !self.verify_tool_signature(&signature) {
109            return Err(ToolError::InvalidSignature(format!(
110                "no trusted key verified signature for tool '{}'",
111                signature.tool_name,
112            )));
113        }
114        let name = tool.name().to_string();
115        self.tools.insert(name.clone(), tool);
116        self.signatures.insert(name, signature);
117        Ok(())
118    }
119
120    /// Check whether a tool signature is valid against any trusted key.
121    pub fn verify_tool_signature(&self, signature: &ToolSignature) -> bool {
122        self.trusted_keys.iter().any(|key| signature.verify(key))
123    }
124
125    /// Enable or disable mandatory signature verification for tool registration.
126    pub fn set_require_signatures(&mut self, require: bool) {
127        self.require_signatures = require;
128    }
129
130    /// Whether signatures are required for tool registration.
131    pub fn requires_signatures(&self) -> bool {
132        self.require_signatures
133    }
134
135    /// Add a trusted Ed25519 public key for signature verification.
136    pub fn add_trusted_key(&mut self, key: [u8; 32]) {
137        self.trusted_keys.push(key);
138    }
139
140    /// Get the signature for a registered tool, if any.
141    pub fn get_signature(&self, tool_name: &str) -> Option<&ToolSignature> {
142        self.signatures.get(tool_name)
143    }
144
145    /// Look up a tool by name, walking the parent chain.
146    pub fn get(&self, name: &str) -> Option<&Arc<dyn BuiltinTool>> {
147        self.tools.get(name).or_else(|| {
148            // Walk parent chain -- returns &Arc from parent, which is valid
149            // because `self` borrows the parent via `Arc`.
150            self.parent.as_ref().and_then(|p| p.get(name))
151        })
152    }
153
154    /// Execute a tool by name, walking the parent chain.
155    pub fn execute(
156        &self,
157        name: &str,
158        args: serde_json::Value,
159    ) -> Result<serde_json::Value, ToolError> {
160        let tool = self
161            .get(name)
162            .ok_or_else(|| ToolError::NotFound(name.to_string()))?;
163        tool.execute(args)
164    }
165
166    /// List all registered tool names (merges parent + local, local wins).
167    pub fn list(&self) -> Vec<String> {
168        let mut seen = std::collections::HashSet::new();
169        // Local tools first
170        for name in self.tools.keys() {
171            seen.insert(name.clone());
172        }
173        // Parent tools (only if not overridden locally)
174        if let Some(ref parent) = self.parent {
175            for name in parent.list() {
176                seen.insert(name);
177            }
178        }
179        let mut result: Vec<String> = seen.into_iter().collect();
180        result.sort();
181        result
182    }
183
184    /// Number of registered tools (parent + local, deduplicated).
185    pub fn len(&self) -> usize {
186        if self.parent.is_none() {
187            return self.tools.len();
188        }
189        self.list().len()
190    }
191
192    /// Whether the registry has no tools (including parent).
193    pub fn is_empty(&self) -> bool {
194        self.tools.is_empty() && self.parent.as_ref().is_none_or(|p| p.is_empty())
195    }
196
197    /// Get a reference to the parent registry, if any.
198    pub fn parent(&self) -> Option<&Arc<ToolRegistry>> {
199        self.parent.as_ref()
200    }
201
202    /// Register a WASM tool that executes through a [`WasmToolRunner`].
203    ///
204    /// The WASM bytes are stored inside the adapter and compiled on each
205    /// execution (K3). Compiled module caching is deferred to K4.
206    ///
207    /// The tool is dispatched synchronously via [`BuiltinTool::execute`],
208    /// which spawns a blocking thread internally to run the async Wasmtime
209    /// execution. For fully async dispatch, call
210    /// [`WasmToolRunner::execute_bytes`] directly.
211    #[cfg(feature = "wasm-sandbox")]
212    pub fn register_wasm_tool(
213        &mut self,
214        name: &str,
215        description: &str,
216        wasm_bytes: Vec<u8>,
217        runner: Arc<WasmToolRunner>,
218    ) -> Result<(), WasmError> {
219        // Validate the module by attempting compilation (handles both
220        // binary WASM and WAT text format).
221        wasmtime::Module::new(runner.engine(), &wasm_bytes)
222            .map_err(|e| WasmError::InvalidModule(e.to_string()))?;
223
224        let adapter = WasmToolAdapter {
225            tool_name: name.to_owned(),
226            spec: BuiltinToolSpec {
227                name: name.to_owned(),
228                category: ToolCategory::User,
229                description: description.to_owned(),
230                parameters: serde_json::json!({
231                    "type": "object",
232                    "properties": {
233                        "input": {"type": "object", "description": "JSON input passed to WASM stdin"}
234                    }
235                }),
236                gate_action: format!("tool.wasm.{name}"),
237                effect: EffectVector {
238                    risk: 0.5,
239                    ..Default::default()
240                },
241                native: false,
242            },
243            wasm_bytes: Arc::new(wasm_bytes),
244            runner,
245        };
246        self.register(Arc::new(adapter));
247        Ok(())
248    }
249}
250
251// ---------------------------------------------------------------------------
252// WASM tool adapter (bridges BuiltinTool to WasmToolRunner)
253// ---------------------------------------------------------------------------
254
255/// Adapter that wraps WASM bytes + a [`WasmToolRunner`] as a [`BuiltinTool`].
256///
257/// When [`BuiltinTool::execute`] is called, this adapter spawns a blocking
258/// thread to run [`WasmToolRunner::execute_bytes`] asynchronously. The JSON
259/// args are passed as stdin to the WASM module.
260#[cfg(feature = "wasm-sandbox")]
261struct WasmToolAdapter {
262    tool_name: String,
263    spec: BuiltinToolSpec,
264    wasm_bytes: Arc<Vec<u8>>,
265    runner: Arc<WasmToolRunner>,
266}
267
268#[cfg(feature = "wasm-sandbox")]
269impl BuiltinTool for WasmToolAdapter {
270    fn name(&self) -> &str {
271        &self.tool_name
272    }
273
274    fn spec(&self) -> &BuiltinToolSpec {
275        &self.spec
276    }
277
278    fn execute(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
279        // Extract input from args, defaulting to the full args object
280        let input = args
281            .get("input")
282            .cloned()
283            .unwrap_or(args.clone());
284
285        let runner = self.runner.clone();
286        let wasm_bytes = self.wasm_bytes.clone();
287        let name = self.tool_name.clone();
288
289        // Run async execute_bytes on a blocking thread with its own runtime
290        let result = std::thread::spawn(move || {
291            let rt = tokio::runtime::Builder::new_current_thread()
292                .enable_all()
293                .build()
294                .map_err(|e| ToolError::ExecutionFailed(format!("runtime: {e}")))?;
295            rt.block_on(runner.execute_bytes(&name, &wasm_bytes, input))
296                .map_err(ToolError::Wasm)
297        })
298        .join()
299        .map_err(|_| ToolError::ExecutionFailed("WASM execution thread panicked".into()))??;
300
301        Ok(serde_json::json!({
302            "stdout": result.stdout,
303            "stderr": result.stderr,
304            "exit_code": result.exit_code,
305            "fuel_consumed": result.fuel_consumed,
306            "execution_time_ms": result.execution_time.as_millis() as u64,
307        }))
308    }
309}
310
311// Safety: WasmToolAdapter is Send+Sync because all its fields are Send+Sync.
312// - tool_name/spec: plain data
313// - wasm_bytes: Arc<Vec<u8>>
314// - runner: Arc<WasmToolRunner> (Engine is Send+Sync)
315#[cfg(feature = "wasm-sandbox")]
316unsafe impl Send for WasmToolAdapter {}
317#[cfg(feature = "wasm-sandbox")]
318unsafe impl Sync for WasmToolAdapter {}
319
320// Safety: ToolRegistry is Send+Sync because it contains Send+Sync fields.
321// The `parent` is behind an Arc, and `tools` contains Arc<dyn BuiltinTool>
322// which requires Send+Sync.
323unsafe impl Send for ToolRegistry {}
324unsafe impl Sync for ToolRegistry {}
325
326impl Default for ToolRegistry {
327    fn default() -> Self {
328        Self::new()
329    }
330}
clawft_kernel/wasm_runner/registry.rs

clawft_kernel/wasm_runner/
registry.rs
