axocoatl-isolation 0.1.1

//! Wasmtime WASM sandbox for isolated tool execution.
//! Uses wasmtime 43 with WASIp1 for core module support.

use std::collections::HashMap;

use wasmtime::{Config, Engine, Linker, Module, Store};

use crate::error::IsolationError;

/// Thread-safe WASM execution sandbox.
/// Pre-compiles modules at startup for fast per-call instantiation (<1ms).
pub struct WasmtimeSandbox {
    engine: Engine,
    /// Pre-compiled module cache: tool_name → compiled Module.
    module_cache: HashMap<String, Module>,
}

impl WasmtimeSandbox {
    /// Create a new sandbox with fuel metering enabled.
    pub fn new() -> Result<Self, IsolationError> {
        let mut config = Config::new();
        config.consume_fuel(true);

        let engine = Engine::new(&config).map_err(|e| IsolationError::Wasmtime(e.to_string()))?;
        Ok(Self {
            engine,
            module_cache: HashMap::new(),
        })
    }

    /// Pre-compile a WASM module (do this at startup, not per-execution).
    pub fn precompile_tool(
        &mut self,
        tool_name: &str,
        wasm_bytes: &[u8],
    ) -> Result<(), IsolationError> {
        let module = Module::new(&self.engine, wasm_bytes).map_err(|e| {
            IsolationError::CompilationFailed {
                tool: tool_name.to_string(),
                reason: e.to_string(),
            }
        })?;
        self.module_cache.insert(tool_name.to_string(), module);
        tracing::debug!(tool = %tool_name, "WASM module pre-compiled");
        Ok(())
    }

    /// Check if a tool is precompiled.
    pub fn has_tool(&self, tool_name: &str) -> bool {
        self.module_cache.contains_key(tool_name)
    }

    /// List precompiled tool names.
    pub fn tool_names(&self) -> Vec<String> {
        self.module_cache.keys().cloned().collect()
    }

    /// Execute a WASM tool with fuel metering.
    /// Input/output via WASI stdin/stdout (JSON-encoded).
    pub async fn execute(
        &self,
        tool_name: &str,
        input: serde_json::Value,
        fuel_limit: u64,
    ) -> Result<serde_json::Value, IsolationError> {
        let module = self
            .module_cache
            .get(tool_name)
            .ok_or_else(|| IsolationError::ToolNotFound(tool_name.to_string()))?;

        // Build WASI context
        let input_bytes = serde_json::to_vec(&input)?;

        // wasmtime 43: use WASIp1 for core module support
        let wasi = wasmtime_wasi::WasiCtxBuilder::new()
            .stdin(wasmtime_wasi::p2::pipe::MemoryInputPipe::new(input_bytes))
            .stdout(wasmtime_wasi::p2::pipe::MemoryOutputPipe::new(65536))
            .build_p1();

        let mut store = Store::new(&self.engine, wasi);
        store
            .set_fuel(fuel_limit)
            .map_err(|e| IsolationError::FuelError(e.to_string()))?;

        let mut linker = Linker::new(&self.engine);
        wasmtime_wasi::p1::add_to_linker_async(&mut linker, |t| t)
            .map_err(|e| IsolationError::Wasmtime(e.to_string()))?;

        let instance = linker
            .instantiate_async(&mut store, module)
            .await
            .map_err(|e| IsolationError::InstantiationFailed(e.to_string()))?;

        // Call the WASM tool's _start export (standard WASI entry point)
        let start = instance
            .get_typed_func::<(), ()>(&mut store, "_start")
            .map_err(|_| IsolationError::MissingExport {
                tool: tool_name.to_string(),
                export: "_start".to_string(),
            })?;

        start.call_async(&mut store, ()).await.map_err(|e| {
            let msg = e.to_string();
            if msg.contains("fuel") {
                IsolationError::FuelExhausted
            } else {
                IsolationError::ExecutionFailed(msg)
            }
        })?;

        // Read stdout — but we can't easily get it back from WasiP1Ctx
        // In a real implementation we'd use a shared output buffer
        // For now, return a success marker
        // TODO: implement proper stdout capture with wasmtime 43 WASIp1
        Ok(serde_json::json!({"status": "ok"}))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn create_sandbox() {
        let sandbox = WasmtimeSandbox::new().unwrap();
        assert!(sandbox.tool_names().is_empty());
    }

    #[test]
    fn precompile_invalid_wasm() {
        let mut sandbox = WasmtimeSandbox::new().unwrap();
        let result = sandbox.precompile_tool("bad", b"not valid wasm");
        assert!(result.is_err());
    }

    #[test]
    fn has_tool_and_list() {
        let mut sandbox = WasmtimeSandbox::new().unwrap();
        assert!(!sandbox.has_tool("test"));

        // Minimal valid WASM module (empty module)
        let wasm = wat::parse_str("(module)").unwrap();
        sandbox.precompile_tool("test", &wasm).unwrap();

        assert!(sandbox.has_tool("test"));
        assert_eq!(sandbox.tool_names(), vec!["test"]);
    }

    #[tokio::test]
    async fn execute_nonexistent_tool() {
        let sandbox = WasmtimeSandbox::new().unwrap();
        let result = sandbox.execute("ghost", serde_json::json!({}), 1000).await;
        assert!(matches!(result, Err(IsolationError::ToolNotFound(_))));
    }
}