xz_skill/runtime/
wasm.rs

1#[cfg(feature = "wasm-runtime")]
2use std::collections::HashMap;
3#[cfg(feature = "wasm-runtime")]
4use std::sync::RwLock;
5
6#[cfg(feature = "wasm-runtime")]
7use crate::error::SkillError;
8
9#[cfg(feature = "wasm-runtime")]
10#[derive(Debug)]
11pub struct WasmConfig {
12    pub memory_limit_mb: u64,
13    pub default_timeout_ms: u64,
14    pub max_instances: usize,
15}
16
17#[cfg(feature = "wasm-runtime")]
18impl Default for WasmConfig {
19    fn default() -> Self {
20        Self {
21            memory_limit_mb: 64,
22            default_timeout_ms: 5000,
23            max_instances: 10,
24        }
25    }
26}
27
28#[cfg(feature = "wasm-runtime")]
29pub struct WasmRuntime {
30    engine: wasmtime::Engine,
31    module_cache: RwLock<HashMap<String, wasmtime::Module>>,
32    config: WasmConfig,
33}
34
35#[cfg(feature = "wasm-runtime")]
36impl std::fmt::Debug for WasmRuntime {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        f.debug_struct("WasmRuntime")
39            .field("engine", &"<wasmtime::Engine>")
40            .field("module_cache", &self.module_cache)
41            .field("config", &self.config)
42            .finish()
43    }
44}
45
46#[cfg(feature = "wasm-runtime")]
47impl WasmRuntime {
48    pub fn new(config: WasmConfig) -> Result<Self, SkillError> {
49        let mut engine_config = wasmtime::Config::default();
50        engine_config
51            .consume_fuel(true)
52            .epoch_interruption(true);
53
54        let engine = wasmtime::Engine::new(&engine_config)
55            .map_err(|e| SkillError::Wasm(e.to_string()))?;
56
57        Ok(Self {
58            engine,
59            module_cache: RwLock::new(HashMap::new()),
60            config,
61        })
62    }
63
64    /// Compile and execute a WASM module.
65    pub async fn execute(
66        &self,
67        module_bytes: &[u8],
68        tool: &str,
69        _args: serde_json::Value,
70    ) -> Result<serde_json::Value, SkillError> {
71        let cache_key = format!("{}:{}", tool, hex::encode(&module_bytes[..module_bytes.len().min(32)]));
72
73        let module = {
74            let cache = self.module_cache.read().unwrap();
75            cache.get(&cache_key).cloned()
76        };
77
78        let module = match module {
79            Some(m) => m,
80            None => {
81                let m = wasmtime::Module::new(&self.engine, module_bytes)
82                    .map_err(|e| SkillError::Wasm(e.to_string()))?;
83                self.module_cache
84                    .write()
85                    .unwrap()
86                    .insert(cache_key.clone(), m.clone());
87                m
88            }
89        };
90
91        let mut store = wasmtime::Store::new(&self.engine, ());
92        store.set_fuel(1_000_000)
93            .map_err(|e| SkillError::Wasm(e.to_string()))?;
94
95        let linker = wasmtime::Linker::new(&self.engine);
96        let instance = linker
97            .instantiate_async(&mut store, &module)
98            .await
99            .map_err(|e| SkillError::Wasm(e.to_string()))?;
100
101        let result = if let Ok(func) = instance.get_typed_func::<(), i32>(&mut store, tool) {
102            let ret = func.call_async(&mut store, ())
103                .await
104                .map_err(|e| SkillError::Wasm(e.to_string()))?;
105            serde_json::json!({"exit_code": ret})
106        } else if let Ok(func) = instance.get_typed_func::<(), ()>(&mut store, tool) {
107            func.call_async(&mut store, ())
108                .await
109                .map_err(|e| SkillError::Wasm(e.to_string()))?;
110            serde_json::json!({"status": "ok"})
111        } else if let Ok(func) = instance.get_typed_func::<(), ()>(&mut store, "_start") {
112            func.call_async(&mut store, ())
113                .await
114                .map_err(|e| SkillError::Wasm(e.to_string()))?;
115            serde_json::json!({"status": "started"})
116        } else {
117            return Err(SkillError::Wasm(format!(
118                "Function '{}' not found in WASM module", tool
119            )));
120        };
121
122        Ok(result)
123    }
124}
125
126// Simple hex encoding helper for cache key
127#[cfg(feature = "wasm-runtime")]
128mod hex {
129    pub fn encode(bytes: &[u8]) -> String {
130        bytes.iter().map(|b| format!("{:02x}", b)).collect()
131    }
132}
xz_skill/runtime/wasm.rs

xz_skill/runtime/
wasm.rs
