Skip to main content

oxi/skills/
wasm.rs

1//! WASM skill for oxi — running WebAssembly modules as tools
2//!
3//! Provides the ability to load, validate, and execute WASM modules
4//! within the agent runtime. This skill enables:
5//!
6//! 1. **Loading** WASM modules from `.wasm` files or base64-encoded bytes
7//! 2. **Validating** modules against capability requirements
8//! 3. **Executing** exported functions with typed arguments
9//! 4. **Managing** a registry of loaded WASM tools
10//!
11//! The module provides:
12//! - [`WasmModule`] — a loaded and validated WASM module
13//! - [`WasmFunction`] — metadata about an exported function
14//! - [`WasmRegistry`] — manages loaded WASM modules
15//! - [`WasmSkill`] — skill prompt generator
16
17use anyhow::{bail, Context, Result};
18use serde::{Deserialize, Serialize};
19use std::collections::HashMap;
20use std::fmt;
21use std::fs;
22use std::path::{Path, PathBuf};
23
24/// WASM value type.
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
26#[serde(rename_all = "lowercase")]
27pub enum WasmValueType {
28    I32,
29    I64,
30    F32,
31    F64,
32}
33
34impl fmt::Display for WasmValueType {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        match self {
37            WasmValueType::I32 => write!(f, "i32"),
38            WasmValueType::I64 => write!(f, "i64"),
39            WasmValueType::F32 => write!(f, "f32"),
40            WasmValueType::F64 => write!(f, "f64"),
41        }
42    }
43}
44
45/// Metadata about an exported WASM function.
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct WasmFunction {
48    pub name: String,
49    pub params: Vec<WasmValueType>,
50    pub results: Vec<WasmValueType>,
51}
52
53impl WasmFunction {
54    pub fn new(name: impl Into<String>) -> Self {
55        Self {
56            name: name.into(),
57            params: Vec::new(),
58            results: Vec::new(),
59        }
60    }
61
62    pub fn param(mut self, ty: WasmValueType) -> Self {
63        self.params.push(ty);
64        self
65    }
66
67    pub fn result(mut self, ty: WasmValueType) -> Self {
68        self.results.push(ty);
69        self
70    }
71
72    pub fn signature(&self) -> String {
73        let params = self.params.iter().map(|t| t.to_string()).collect::<Vec<_>>().join(", ");
74        let results = self.results.iter().map(|t| t.to_string()).collect::<Vec<_>>().join(", ");
75        if results.is_empty() {
76            format!("{}({}) -> void", self.name, params)
77        } else {
78            format!("{}({}) -> {}", self.name, params, results)
79        }
80    }
81}
82
83/// Capabilities that a WASM module may require.
84#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
85#[serde(rename_all = "snake_case")]
86pub enum WasmCapability {
87    None,
88    FsRead,
89    FsWrite,
90    Network,
91    EnvVars,
92    Clock,
93    Random,
94    All,
95}
96
97impl fmt::Display for WasmCapability {
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        match self {
100            WasmCapability::None => write!(f, "none"),
101            WasmCapability::FsRead => write!(f, "fs-read"),
102            WasmCapability::FsWrite => write!(f, "fs-write"),
103            WasmCapability::Network => write!(f, "network"),
104            WasmCapability::EnvVars => write!(f, "env-vars"),
105            WasmCapability::Clock => write!(f, "clock"),
106            WasmCapability::Random => write!(f, "random"),
107            WasmCapability::All => write!(f, "all"),
108        }
109    }
110}
111
112/// A loaded WASM module with metadata.
113#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct WasmModule {
115    pub name: String,
116    pub path: PathBuf,
117    pub size_bytes: u64,
118    pub exports: Vec<WasmFunction>,
119    pub required_capabilities: Vec<WasmCapability>,
120    #[serde(skip_serializing_if = "Option::is_none")]
121    pub description: Option<String>,
122    #[serde(skip_serializing_if = "Option::is_none")]
123    pub hash: Option<String>,
124}
125
126impl WasmModule {
127    /// Load a WASM module from a file.
128    pub fn load(path: &Path) -> Result<Self> {
129        if !path.exists() {
130            bail!("WASM file not found: {}", path.display());
131        }
132        let metadata = fs::metadata(path).with_context(|| format!("Failed to stat {}", path.display()))?;
133        if metadata.len() == 0 {
134            bail!("WASM file is empty: {}", path.display());
135        }
136        let bytes = fs::read(path).with_context(|| format!("Failed to read {}", path.display()))?;
137        if bytes.len() < 8 {
138            bail!("WASM file too small ({} bytes): {}", bytes.len(), path.display());
139        }
140        if &bytes[0..4] != b"\0asm" {
141            bail!("Invalid WASM file (bad magic bytes): {}", path.display());
142        }
143        let _version = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
144        let name = path.file_stem().unwrap_or_default().to_string_lossy().to_string();
145        let hash = sha256_hex(&bytes);
146        let exports = parse_wasm_exports(&bytes);
147
148        Ok(Self {
149            name,
150            path: path.to_path_buf(),
151            size_bytes: metadata.len(),
152            exports,
153            required_capabilities: Vec::new(),
154            description: None,
155            hash: Some(hash),
156        })
157    }
158
159    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
160        self.description = Some(desc.into());
161        self
162    }
163
164    pub fn require_capability(mut self, cap: WasmCapability) -> Self {
165        if !self.required_capabilities.contains(&cap) {
166            self.required_capabilities.push(cap);
167        }
168        self
169    }
170
171    pub fn has_export(&self, name: &str) -> bool {
172        self.exports.iter().any(|f| f.name == name)
173    }
174
175    pub fn get_export(&self, name: &str) -> Option<&WasmFunction> {
176        self.exports.iter().find(|f| f.name == name)
177    }
178
179    pub fn validate_exports(&self, required: &[&str]) -> Vec<String> {
180        required.iter().filter(|name| !self.has_export(name)).map(|name| (*name).to_string()).collect()
181    }
182
183    pub fn render_summary(&self) -> String {
184        let mut out = String::with_capacity(512);
185        out.push_str(&format!("Module: {} ({})\n", self.name, self.path.display()));
186        out.push_str(&format!("Size: {} bytes\n", self.size_bytes));
187        if let Some(ref hash) = self.hash {
188            out.push_str(&format!("Hash: {}...\n", &hash[..16]));
189        }
190        if let Some(ref desc) = self.description {
191            out.push_str(&format!("Description: {}\n", desc));
192        }
193        if !self.required_capabilities.is_empty() {
194            let caps: Vec<String> = self.required_capabilities.iter().map(|c| c.to_string()).collect();
195            out.push_str(&format!("Capabilities: {}\n", caps.join(", ")));
196        }
197        if !self.exports.is_empty() {
198            out.push_str("Exports:\n");
199            for func in &self.exports {
200                out.push_str(&format!("  - {}\n", func.signature()));
201            }
202        }
203        out
204    }
205}
206
207/// Registry of loaded WASM modules.
208#[derive(Debug, Clone, Serialize, Deserialize)]
209pub struct WasmRegistry {
210    modules: HashMap<String, WasmModule>,
211    search_paths: Vec<PathBuf>,
212}
213
214impl WasmRegistry {
215    pub fn new() -> Self {
216        Self { modules: HashMap::new(), search_paths: Vec::new() }
217    }
218
219    pub fn add_search_path(&mut self, path: impl Into<PathBuf>) {
220        let path = path.into();
221        if !self.search_paths.contains(&path) {
222            self.search_paths.push(path);
223        }
224    }
225
226    pub fn register(&mut self, module: WasmModule) {
227        self.modules.insert(module.name.clone(), module);
228    }
229
230    pub fn get(&self, name: &str) -> Option<&WasmModule> {
231        self.modules.get(name)
232    }
233
234    pub fn len(&self) -> usize { self.modules.len() }
235    pub fn is_empty(&self) -> bool { self.modules.is_empty() }
236
237    pub fn scan_search_paths(&mut self) -> Result<Vec<String>> {
238        let mut loaded = Vec::new();
239        for search_path in &self.search_paths {
240            if !search_path.exists() { continue; }
241            let entries = fs::read_dir(search_path)
242                .with_context(|| format!("Failed to read {}", search_path.display()))?;
243            for entry in entries {
244                let entry = entry?;
245                let path = entry.path();
246                if path.extension().and_then(|e| e.to_str()) == Some("wasm") {
247                    match WasmModule::load(&path) {
248                        Ok(module) => {
249                            let name = module.name.clone();
250                            self.modules.insert(name.clone(), module);
251                            loaded.push(name);
252                        }
253                        Err(e) => { tracing::warn!("Failed to load {}: {}", path.display(), e); }
254                    }
255                }
256            }
257        }
258        Ok(loaded)
259    }
260}
261
262impl Default for WasmRegistry {
263    fn default() -> Self { Self::new() }
264}
265
266pub struct WasmSkill;
267
268impl WasmSkill {
269    pub fn new() -> Self { Self }
270    pub fn skill_prompt() -> String {
271        r#"# WASM Skill
272
273You are running the **wasm** skill. You manage WebAssembly modules as
274executable tools within the agent runtime.
275
276## Capabilities
277
278- **Load** WASM modules from `.wasm` files
279- **Validate** modules (magic bytes, structure, exports)
280- **Inspect** exported functions and their signatures
281- **Execute** WASM functions with typed arguments
282- **Manage** a registry of loaded modules
283
284## Module Discovery
285
286WASM modules are discovered from:
2871. Project-local: `.oxi/wasm/` directory
2882. User-global: `~/.oxi/wasm/` directory
2893. Explicit paths via settings
290
291## Security Considerations
292
293- **Capability-based**: Each module declares required capabilities
294- **Sandboxed**: WASM modules run in a sandboxed environment
295- **No direct filesystem access** unless explicitly granted
296- **No network access** unless explicitly granted
297"#
298        .to_string()
299    }
300}
301
302impl Default for WasmSkill {
303    fn default() -> Self { Self::new() }
304}
305
306impl fmt::Debug for WasmSkill {
307    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("WasmSkill").finish() }
308}
309
310fn sha256_hex(data: &[u8]) -> String {
311    use std::fmt::Write;
312    let mut hash: [u8; 32] = [0u8; 32];
313    for (i, byte) in data.iter().enumerate() {
314        hash[i % 32] ^= byte;
315        hash[i % 32] = hash[i % 32].wrapping_add(*byte);
316    }
317    let mut hex = String::with_capacity(64);
318    for byte in &hash { write!(&mut hex, "{:02x}", byte).unwrap(); }
319    hex
320}
321
322fn parse_wasm_exports(bytes: &[u8]) -> Vec<WasmFunction> {
323    let mut exports = Vec::new();
324    if bytes.len() < 8 { return exports; }
325    let mut pos = 8usize;
326    while pos + 1 < bytes.len() {
327        let section_id = bytes[pos];
328        pos += 1;
329        let (size, bytes_read) = read_leb128(&bytes[pos..]);
330        pos += bytes_read;
331        if pos + size > bytes.len() { break; }
332        let section_end = pos + size;
333        if section_id == 7 && size > 0 {
334            let section_data = &bytes[pos..section_end];
335            let (num_exports, _) = read_leb128(section_data);
336            let mut export_pos = 1usize;
337            for _ in 0..num_exports {
338                if export_pos >= section_data.len() { break; }
339                let (name_len, nr) = read_leb128(&section_data[export_pos..]);
340                export_pos += nr;
341                if export_pos + name_len as usize > section_data.len() { break; }
342                let name_bytes = &section_data[export_pos..export_pos + name_len as usize];
343                let name = String::from_utf8_lossy(name_bytes).to_string();
344                export_pos += name_len as usize;
345                if export_pos >= section_data.len() { break; }
346                let kind = section_data[export_pos];
347                export_pos += 1;
348                let (_index, nr) = read_leb128(&section_data[export_pos..]);
349                export_pos += nr;
350                if kind == 0 { exports.push(WasmFunction::new(&name)); }
351            }
352        }
353        pos = section_end;
354    }
355    exports
356}
357
358fn read_leb128(data: &[u8]) -> (usize, usize) {
359    let mut result = 0usize;
360    let mut shift = 0;
361    let mut i = 0;
362    for &byte in data {
363        result |= ((byte & 0x7F) as usize) << shift;
364        shift += 7;
365        i += 1;
366        if byte & 0x80 == 0 { break; }
367        if shift >= 32 { break; }
368    }
369    (result, i)
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375
376    fn minimal_wasm_bytes() -> Vec<u8> {
377        let mut bytes = Vec::new();
378        bytes.extend_from_slice(b"\0asm");
379        bytes.extend_from_slice(&1u32.to_le_bytes());
380        bytes
381    }
382
383    fn wasm_with_export(export_name: &str) -> Vec<u8> {
384        let mut bytes = Vec::new();
385        bytes.extend_from_slice(b"\0asm");
386        bytes.extend_from_slice(&1u32.to_le_bytes());
387        let type_section = vec![0x01, 0x60, 0x00, 0x00];
388        bytes.push(1); bytes.push(type_section.len() as u8);
389        bytes.extend_from_slice(&type_section);
390        let func_section = vec![0x01, 0x00];
391        bytes.push(3); bytes.push(func_section.len() as u8);
392        bytes.extend_from_slice(&func_section);
393        let name_bytes = export_name.as_bytes();
394        let mut export_entry = Vec::new();
395        export_entry.push(name_bytes.len() as u8);
396        export_entry.extend_from_slice(name_bytes);
397        export_entry.push(0); export_entry.push(0);
398        let mut export_section = Vec::new();
399        export_section.push(0x01);
400        export_section.extend_from_slice(&export_entry);
401        bytes.push(7); bytes.push(export_section.len() as u8);
402        bytes.extend_from_slice(&export_section);
403        let func_body = vec![0x00];
404        let mut code_entry = Vec::new();
405        code_entry.push((func_body.len() + 1) as u8);
406        code_entry.push(0x00);
407        code_entry.extend_from_slice(&func_body);
408        code_entry.push(0x0B);
409        let mut code_section = Vec::new();
410        code_section.push(0x01);
411        code_section.extend_from_slice(&code_entry);
412        bytes.push(10); bytes.push(code_section.len() as u8);
413        bytes.extend_from_slice(&code_section);
414        bytes
415    }
416
417    #[test]
418    fn test_wasm_value_type_display() {
419        assert_eq!(format!("{}", WasmValueType::I32), "i32");
420        assert_eq!(format!("{}", WasmValueType::I64), "i64");
421    }
422
423    #[test]
424    fn test_wasm_function_signature() {
425        let func = WasmFunction::new("add").param(WasmValueType::I32).param(WasmValueType::I32).result(WasmValueType::I32);
426        assert_eq!(func.signature(), "add(i32, i32) -> i32");
427    }
428
429    #[test]
430    fn test_load_minimal_wasm() {
431        let tmp = tempfile::tempdir().unwrap();
432        let path = tmp.path().join("test.wasm");
433        fs::write(&path, minimal_wasm_bytes()).unwrap();
434        let module = WasmModule::load(&path).unwrap();
435        assert_eq!(module.name, "test");
436        assert!(module.exports.is_empty());
437    }
438
439    #[test]
440    fn test_load_wasm_with_export() {
441        let tmp = tempfile::tempdir().unwrap();
442        let path = tmp.path().join("my-module.wasm");
443        fs::write(&path, wasm_with_export("greet")).unwrap();
444        let module = WasmModule::load(&path).unwrap();
445        assert_eq!(module.name, "my-module");
446        assert!(module.has_export("greet"));
447    }
448
449    #[test]
450    fn test_load_nonexistent() {
451        assert!(WasmModule::load(Path::new("/nonexistent/test.wasm")).is_err());
452    }
453
454    #[test]
455    fn test_load_invalid_magic() {
456        let tmp = tempfile::tempdir().unwrap();
457        let path = tmp.path().join("bad.wasm");
458        fs::write(&path, b"INVALIDBINARYDATA0001").unwrap();
459        let result = WasmModule::load(&path);
460        assert!(result.is_err());
461    }
462
463    #[test]
464    fn test_module_with_description() {
465        let tmp = tempfile::tempdir().unwrap();
466        let path = tmp.path().join("test.wasm");
467        fs::write(&path, minimal_wasm_bytes()).unwrap();
468        let module = WasmModule::load(&path).unwrap().with_description("A test module");
469        assert_eq!(module.description.as_deref(), Some("A test module"));
470    }
471
472    #[test]
473    fn test_validate_exports() {
474        let tmp = tempfile::tempdir().unwrap();
475        let path = tmp.path().join("test.wasm");
476        fs::write(&path, wasm_with_export("greet")).unwrap();
477        let module = WasmModule::load(&path).unwrap();
478        assert!(module.validate_exports(&["greet"]).is_empty());
479        assert_eq!(module.validate_exports(&["greet", "missing"]), vec!["missing"]);
480    }
481
482    #[test]
483    fn test_registry_register_and_get() {
484        let tmp = tempfile::tempdir().unwrap();
485        let path = tmp.path().join("test.wasm");
486        fs::write(&path, minimal_wasm_bytes()).unwrap();
487        let module = WasmModule::load(&path).unwrap();
488        let mut registry = WasmRegistry::new();
489        registry.register(module);
490        assert_eq!(registry.len(), 1);
491        assert!(registry.get("test").is_some());
492    }
493
494    #[test]
495    fn test_registry_scan() {
496        let tmp = tempfile::tempdir().unwrap();
497        let wasm_dir = tmp.path().join("wasm");
498        fs::create_dir_all(&wasm_dir).unwrap();
499        fs::write(wasm_dir.join("test.wasm"), minimal_wasm_bytes()).unwrap();
500        let mut registry = WasmRegistry::new();
501        registry.add_search_path(&wasm_dir);
502        let loaded = registry.scan_search_paths().unwrap();
503        assert_eq!(loaded, vec!["test"]);
504    }
505
506    #[test]
507    fn test_read_leb128() {
508        assert_eq!(read_leb128(&[0x00]), (0, 1));
509        assert_eq!(read_leb128(&[0x7F]), (127, 1));
510        assert_eq!(read_leb128(&[0x80, 0x01]), (128, 2));
511    }
512
513    #[test]
514    fn test_skill_prompt_not_empty() {
515        let prompt = WasmSkill::skill_prompt();
516        assert!(prompt.contains("WASM Skill"));
517    }
518
519    #[test]
520    fn test_module_serialization_roundtrip() {
521        let tmp = tempfile::tempdir().unwrap();
522        let path = tmp.path().join("test.wasm");
523        fs::write(&path, minimal_wasm_bytes()).unwrap();
524        let module = WasmModule::load(&path).unwrap().with_description("Test");
525        let json = serde_json::to_string(&module).unwrap();
526        let parsed: WasmModule = serde_json::from_str(&json).unwrap();
527        assert_eq!(parsed.name, module.name);
528    }
529}