Skip to main content

agentzero_tools/
wasm_tools.rs

1use agentzero_core::{Tool, ToolContext, ToolResult};
2use anyhow::{anyhow, Context};
3use async_trait::async_trait;
4use serde::Deserialize;
5use serde_json::json;
6use std::path::Path;
7
8// --- wasm_module ---
9
10#[derive(Debug, Deserialize)]
11struct WasmModuleInput {
12    op: String,
13    #[serde(default)]
14    path: Option<String>,
15}
16
17/// Load and inspect WASM modules.
18///
19/// Operations:
20/// - `inspect`: Read a .wasm file and return basic module info (size, header validation)
21/// - `list`: List .wasm files in the workspace plugins directory
22#[derive(Debug, Default, Clone, Copy)]
23pub struct WasmModuleTool;
24
25#[async_trait]
26impl Tool for WasmModuleTool {
27    fn name(&self) -> &'static str {
28        "wasm_module"
29    }
30
31    fn description(&self) -> &'static str {
32        "Inspect or list WASM modules in the workspace plugins directory."
33    }
34
35    fn input_schema(&self) -> Option<serde_json::Value> {
36        Some(json!({
37            "type": "object",
38            "properties": {
39                "op": { "type": "string", "enum": ["inspect", "list"], "description": "The WASM module operation" },
40                "path": { "type": "string", "description": "Path to the WASM file (for inspect)" }
41            },
42            "required": ["op"],
43            "additionalProperties": false
44        }))
45    }
46
47    async fn execute(&self, input: &str, ctx: &ToolContext) -> anyhow::Result<ToolResult> {
48        let req: WasmModuleInput =
49            serde_json::from_str(input).context("wasm_module expects JSON: {\"op\", ...}")?;
50
51        match req.op.as_str() {
52            "inspect" => {
53                let path_str = req
54                    .path
55                    .as_deref()
56                    .ok_or_else(|| anyhow!("inspect requires a `path` field"))?;
57
58                if path_str.trim().is_empty() {
59                    return Err(anyhow!("path must not be empty"));
60                }
61
62                let full_path = resolve_wasm_path(&ctx.workspace_root, path_str);
63                if !full_path.exists() {
64                    return Err(anyhow!("WASM file not found: {}", full_path.display()));
65                }
66
67                let metadata = tokio::fs::metadata(&full_path)
68                    .await
69                    .context("failed to read WASM file metadata")?;
70
71                let bytes = tokio::fs::read(&full_path)
72                    .await
73                    .context("failed to read WASM file")?;
74
75                let valid_magic = bytes.len() >= 4 && bytes[..4] == [0x00, 0x61, 0x73, 0x6D];
76                let version = if bytes.len() >= 8 {
77                    Some(u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]))
78                } else {
79                    None
80                };
81
82                let output = json!({
83                    "path": full_path.display().to_string(),
84                    "size_bytes": metadata.len(),
85                    "valid_wasm_header": valid_magic,
86                    "wasm_version": version,
87                })
88                .to_string();
89
90                Ok(ToolResult { output })
91            }
92            "list" => {
93                let plugins_dir = Path::new(&ctx.workspace_root).join(".agentzero/plugins");
94                let mut wasm_files = Vec::new();
95
96                if plugins_dir.exists() {
97                    let mut entries = tokio::fs::read_dir(&plugins_dir)
98                        .await
99                        .context("failed to read plugins directory")?;
100
101                    while let Some(entry) = entries.next_entry().await? {
102                        let path = entry.path();
103                        if path.extension().is_some_and(|ext| ext == "wasm") {
104                            let meta = tokio::fs::metadata(&path).await.ok();
105                            wasm_files.push(json!({
106                                "name": path.file_name().unwrap_or_default().to_string_lossy(),
107                                "path": path.display().to_string(),
108                                "size_bytes": meta.map(|m| m.len()).unwrap_or(0),
109                            }));
110                        }
111                    }
112                }
113
114                if wasm_files.is_empty() {
115                    return Ok(ToolResult {
116                        output: "no WASM modules found".to_string(),
117                    });
118                }
119
120                Ok(ToolResult {
121                    output: serde_json::to_string_pretty(&wasm_files)
122                        .unwrap_or_else(|_| "[]".to_string()),
123                })
124            }
125            other => Ok(ToolResult {
126                output: json!({ "error": format!("unknown op: {other}") }).to_string(),
127            }),
128        }
129    }
130}
131
132// --- wasm_tool ---
133
134#[derive(Debug, Deserialize)]
135struct WasmToolInput {
136    module: String,
137    #[serde(default)]
138    function: Option<String>,
139    #[serde(default)]
140    args: Option<serde_json::Value>,
141}
142
143/// Execute WASM-based tools via plugin runtime.
144///
145/// This tool loads a WASM module and invokes a function within it.
146/// Currently validates the module path and reports that WASM execution
147/// requires the `wasmtime` runtime (future integration).
148#[derive(Debug, Default, Clone, Copy)]
149pub struct WasmToolExecTool;
150
151#[async_trait]
152impl Tool for WasmToolExecTool {
153    fn name(&self) -> &'static str {
154        "wasm_tool"
155    }
156
157    fn description(&self) -> &'static str {
158        "Execute a function within a WASM module."
159    }
160
161    fn input_schema(&self) -> Option<serde_json::Value> {
162        Some(json!({
163            "type": "object",
164            "properties": {
165                "module": { "type": "string", "description": "Path to the WASM module file" },
166                "function": { "type": "string", "description": "Function to invoke (default: _start)" },
167                "args": { "type": "object", "description": "Arguments to pass to the function" }
168            },
169            "required": ["module"],
170            "additionalProperties": false
171        }))
172    }
173
174    async fn execute(&self, input: &str, ctx: &ToolContext) -> anyhow::Result<ToolResult> {
175        let req: WasmToolInput =
176            serde_json::from_str(input).context("wasm_tool expects JSON: {\"module\", ...}")?;
177
178        if req.module.trim().is_empty() {
179            return Err(anyhow!("module must not be empty"));
180        }
181
182        let full_path = resolve_wasm_path(&ctx.workspace_root, &req.module);
183        if !full_path.exists() {
184            return Err(anyhow!("WASM module not found: {}", full_path.display()));
185        }
186
187        let bytes = tokio::fs::read(&full_path)
188            .await
189            .context("failed to read WASM module")?;
190
191        let valid_magic = bytes.len() >= 4 && bytes[..4] == [0x00, 0x61, 0x73, 0x6D];
192        if !valid_magic {
193            return Err(anyhow!(
194                "file is not a valid WASM module (invalid magic bytes)"
195            ));
196        }
197
198        let function = req.function.as_deref().unwrap_or("_start");
199
200        // WASM runtime execution is a future integration point.
201        // For now, validate the module and report readiness.
202        let output = json!({
203            "module": full_path.display().to_string(),
204            "function": function,
205            "args": req.args,
206            "size_bytes": bytes.len(),
207            "valid": true,
208            "status": "validated",
209            "note": "WASM execution requires wasmtime runtime (not yet integrated)"
210        })
211        .to_string();
212
213        Ok(ToolResult { output })
214    }
215}
216
217fn resolve_wasm_path(workspace_root: &str, path: &str) -> std::path::PathBuf {
218    let p = Path::new(path);
219    if p.is_absolute() {
220        p.to_path_buf()
221    } else {
222        Path::new(workspace_root).join(path)
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use agentzero_core::ToolContext;
230    use std::fs;
231    use std::path::PathBuf;
232    use std::sync::atomic::{AtomicU64, Ordering};
233    use std::time::{SystemTime, UNIX_EPOCH};
234
235    static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
236
237    fn temp_dir() -> PathBuf {
238        let nanos = SystemTime::now()
239            .duration_since(UNIX_EPOCH)
240            .expect("clock")
241            .as_nanos();
242        let seq = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
243        let dir = std::env::temp_dir().join(format!(
244            "agentzero-wasm-tools-{}-{nanos}-{seq}",
245            std::process::id()
246        ));
247        fs::create_dir_all(&dir).expect("temp dir should be created");
248        dir
249    }
250
251    fn write_minimal_wasm(dir: &Path, name: &str) -> PathBuf {
252        // Minimal valid WASM: magic + version only (8 bytes)
253        let path = dir.join(name);
254        let wasm_header: [u8; 8] = [0x00, 0x61, 0x73, 0x6D, 0x01, 0x00, 0x00, 0x00];
255        fs::write(&path, wasm_header).expect("write wasm");
256        path
257    }
258
259    // --- wasm_module tests ---
260
261    #[tokio::test]
262    async fn wasm_module_inspect_valid() {
263        let dir = temp_dir();
264        let wasm_path = write_minimal_wasm(&dir, "test.wasm");
265
266        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
267        let result = WasmModuleTool
268            .execute(
269                &format!(r#"{{"op": "inspect", "path": "{}"}}"#, wasm_path.display()),
270                &ctx,
271            )
272            .await
273            .expect("inspect should succeed");
274        let v: serde_json::Value = serde_json::from_str(&result.output).unwrap();
275        assert_eq!(v["valid_wasm_header"], true);
276        assert_eq!(v["wasm_version"], 1);
277        assert_eq!(v["size_bytes"], 8);
278
279        fs::remove_dir_all(dir).ok();
280    }
281
282    #[tokio::test]
283    async fn wasm_module_inspect_not_found() {
284        let ctx = ToolContext::new("/tmp".to_string());
285        let err = WasmModuleTool
286            .execute(r#"{"op": "inspect", "path": "nonexistent.wasm"}"#, &ctx)
287            .await
288            .expect_err("missing file should fail");
289        assert!(err.to_string().contains("not found"));
290    }
291
292    #[tokio::test]
293    async fn wasm_module_inspect_invalid_header() {
294        let dir = temp_dir();
295        let path = dir.join("bad.wasm");
296        fs::write(&path, b"not a wasm file").unwrap();
297
298        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
299        let result = WasmModuleTool
300            .execute(
301                &format!(r#"{{"op": "inspect", "path": "{}"}}"#, path.display()),
302                &ctx,
303            )
304            .await
305            .expect("inspect should succeed even for invalid wasm");
306        let v: serde_json::Value = serde_json::from_str(&result.output).unwrap();
307        assert_eq!(v["valid_wasm_header"], false);
308
309        fs::remove_dir_all(dir).ok();
310    }
311
312    #[tokio::test]
313    async fn wasm_module_list_empty() {
314        let dir = temp_dir();
315        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
316
317        let result = WasmModuleTool
318            .execute(r#"{"op": "list"}"#, &ctx)
319            .await
320            .expect("list should succeed");
321        assert!(result.output.contains("no WASM modules found"));
322
323        fs::remove_dir_all(dir).ok();
324    }
325
326    #[tokio::test]
327    async fn wasm_module_list_finds_files() {
328        let dir = temp_dir();
329        let plugins_dir = dir.join(".agentzero/plugins");
330        fs::create_dir_all(&plugins_dir).unwrap();
331        write_minimal_wasm(&plugins_dir, "plugin1.wasm");
332        write_minimal_wasm(&plugins_dir, "plugin2.wasm");
333
334        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
335        let result = WasmModuleTool
336            .execute(r#"{"op": "list"}"#, &ctx)
337            .await
338            .expect("list should succeed");
339        assert!(result.output.contains("plugin1.wasm"));
340        assert!(result.output.contains("plugin2.wasm"));
341
342        fs::remove_dir_all(dir).ok();
343    }
344
345    // --- wasm_tool tests ---
346
347    #[tokio::test]
348    async fn wasm_tool_validates_module() {
349        let dir = temp_dir();
350        let wasm_path = write_minimal_wasm(&dir, "tool.wasm");
351
352        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
353        let result = WasmToolExecTool
354            .execute(
355                &format!(
356                    r#"{{"module": "{}", "function": "run"}}"#,
357                    wasm_path.display()
358                ),
359                &ctx,
360            )
361            .await
362            .expect("should succeed");
363        let v: serde_json::Value = serde_json::from_str(&result.output).unwrap();
364        assert_eq!(v["valid"], true);
365        assert_eq!(v["status"], "validated");
366        assert_eq!(v["function"], "run");
367
368        fs::remove_dir_all(dir).ok();
369    }
370
371    #[tokio::test]
372    async fn wasm_tool_rejects_missing_module() {
373        let ctx = ToolContext::new("/tmp".to_string());
374        let err = WasmToolExecTool
375            .execute(r#"{"module": "missing.wasm"}"#, &ctx)
376            .await
377            .expect_err("missing module should fail");
378        assert!(err.to_string().contains("not found"));
379    }
380
381    #[tokio::test]
382    async fn wasm_tool_rejects_invalid_wasm() {
383        let dir = temp_dir();
384        let path = dir.join("bad.wasm");
385        fs::write(&path, b"not wasm").unwrap();
386
387        let ctx = ToolContext::new(dir.to_string_lossy().to_string());
388        let err = WasmToolExecTool
389            .execute(&format!(r#"{{"module": "{}"}}"#, path.display()), &ctx)
390            .await
391            .expect_err("invalid wasm should fail");
392        assert!(err.to_string().contains("invalid magic bytes"));
393
394        fs::remove_dir_all(dir).ok();
395    }
396
397    #[tokio::test]
398    async fn wasm_tool_empty_module_fails() {
399        let ctx = ToolContext::new("/tmp".to_string());
400        let err = WasmToolExecTool
401            .execute(r#"{"module": ""}"#, &ctx)
402            .await
403            .expect_err("empty module should fail");
404        assert!(err.to_string().contains("module must not be empty"));
405    }
406}