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

clawft_kernel/wasm_runner/
registry.rs
