use crate::plugin::lifecycle::PluginLoader;
use crate::plugin::manifest::PluginManifest;
use crate::plugin::rich_output::normalize_plugin_tool_result;
use parking_lot::Mutex;
use serde_json::Value;
use std::collections::HashMap;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::Arc;
use tracing::{debug, info, warn};
pub type WasmResult<T> = Result<T, WasmError>;
pub type WasmExecutionFuture = Pin<Box<dyn Future<Output = WasmResult<Value>> + Send>>;
pub const DEFAULT_PLUGIN_MAX_MEMORY_BYTES: u64 = 128 * 1024 * 1024;
const WASM_PAGE_SIZE: u64 = 64 * 1024;
fn bytes_to_pages(bytes: u64) -> u32 {
let pages = bytes / WASM_PAGE_SIZE;
pages.clamp(1, u32::MAX as u64) as u32
}
struct LoadedPlugin {
artifact_path: PathBuf,
plugin: extism::Plugin,
manifest: Option<PluginManifest>,
}
#[derive(Clone)]
pub struct WasmHost {
inner: Arc<Mutex<WasmHostInner>>,
max_memory_bytes: u64,
}
struct WasmHostInner {
loaded: HashMap<String, LoadedPlugin>,
}
impl WasmHost {
pub fn new() -> Self {
Self::with_max_memory_bytes(DEFAULT_PLUGIN_MAX_MEMORY_BYTES)
}
pub fn with_max_memory_bytes(max_memory_bytes: u64) -> Self {
Self {
inner: Arc::new(Mutex::new(WasmHostInner {
loaded: HashMap::new(),
})),
max_memory_bytes,
}
}
pub fn max_memory_bytes(&self) -> u64 {
self.max_memory_bytes
}
pub fn load_plugin(&self, plugin_id: &str, artifact_path: &Path) -> WasmResult<()> {
self.load_plugin_with_limit(plugin_id, artifact_path, None)
}
pub fn load_plugin_with_limit(
&self,
plugin_id: &str,
artifact_path: &Path,
plugin_declared_max_bytes: Option<u64>,
) -> WasmResult<()> {
let wasm_bytes = std::fs::read(artifact_path).map_err(|e| {
WasmError::LoadFailed(format!(
"Failed to read artifact {}: {}",
artifact_path.display(),
e
))
})?;
let effective_bytes = plugin_declared_max_bytes
.map(|declared| declared.min(self.max_memory_bytes))
.unwrap_or(self.max_memory_bytes);
let max_pages = bytes_to_pages(effective_bytes);
let manifest = extism::Manifest::new([extism::Wasm::data(wasm_bytes)])
.with_timeout(std::time::Duration::from_secs(30))
.with_memory_max(max_pages);
let plugin = extism::Plugin::new(manifest, [], true)
.map_err(|e| WasmError::LoadFailed(format!("Extism plugin creation failed: {}", e)))?;
let mut inner = self.inner.lock();
inner.loaded.insert(
plugin_id.to_string(),
LoadedPlugin {
artifact_path: artifact_path.to_path_buf(),
plugin,
manifest: None, },
);
info!(
plugin_id = %plugin_id,
max_memory_bytes = effective_bytes,
max_pages,
"Plugin loaded into WASM host"
);
Ok(())
}
pub fn unload_plugin(&self, plugin_id: &str) -> WasmResult<()> {
let mut inner = self.inner.lock();
if inner.loaded.remove(plugin_id).is_some() {
info!(plugin_id = %plugin_id, "Plugin unloaded from WASM host");
} else {
debug!(plugin_id = %plugin_id, "Unload requested but plugin was not loaded");
}
Ok(())
}
pub fn execute_tool(
&self,
plugin_id: &str,
tool_name: &str,
arguments: Value,
) -> WasmExecutionFuture {
let plugin_id = plugin_id.to_string();
let tool_name = tool_name.to_string();
let inner = self.inner.clone();
Box::pin(async move {
let result = tokio::task::spawn_blocking(move || {
let mut guard = inner.lock();
let loaded = guard
.loaded
.get_mut(&plugin_id)
.ok_or_else(|| WasmError::NotFound(plugin_id.clone()))?;
let entrypoint = format!("tool_{}", tool_name);
if !loaded.plugin.function_exists(&entrypoint) {
return Err(WasmError::ExecutionFailed(format!(
"Plugin does not export function '{}'",
entrypoint
)));
}
let input = serde_json::to_string(&arguments).map_err(|e| {
WasmError::InvalidInput(format!("Failed to serialize arguments: {}", e))
})?;
let output: &str = loaded.plugin.call(&entrypoint, &input).map_err(|e| {
let msg = e.to_string();
if msg.contains("timeout") || msg.contains("timed out") {
WasmError::Timeout
} else if msg.contains("trap") || msg.contains("panic") {
WasmError::PluginPanicked(msg)
} else {
WasmError::ExecutionFailed(msg)
}
})?;
let response: Value = serde_json::from_str(output).map_err(|e| {
WasmError::ExecutionFailed(format!("Plugin returned invalid JSON: {}", e))
})?;
if let Some(error) = response.get("error") {
let msg = error.as_str().unwrap_or("Unknown plugin error");
return Err(WasmError::ExecutionFailed(msg.to_string()));
}
let result = response.get("ok").cloned().unwrap_or(Value::Null);
normalize_plugin_tool_result(&plugin_id, &tool_name, result)
})
.await;
match result {
Ok(inner_result) => inner_result,
Err(join_error) => Err(WasmError::ExecutionFailed(format!(
"Task join error: {}",
join_error
))),
}
})
}
pub fn execute_tool_sync(
&self,
plugin_id: &str,
tool_name: &str,
arguments: Value,
) -> WasmResult<Value> {
let mut guard = self.inner.lock();
let loaded = guard
.loaded
.get_mut(plugin_id)
.ok_or_else(|| WasmError::NotFound(plugin_id.to_string()))?;
let entrypoint = format!("tool_{}", tool_name);
if !loaded.plugin.function_exists(&entrypoint) {
return Err(WasmError::ExecutionFailed(format!(
"Plugin does not export function '{}'",
entrypoint
)));
}
let input = serde_json::to_string(&arguments).map_err(|e| {
WasmError::InvalidInput(format!("Failed to serialize arguments: {}", e))
})?;
let output: &str = loaded.plugin.call(&entrypoint, &input).map_err(|e| {
let msg = e.to_string();
if msg.contains("timeout") || msg.contains("timed out") {
WasmError::Timeout
} else if msg.contains("trap") || msg.contains("panic") {
WasmError::PluginPanicked(msg)
} else {
WasmError::ExecutionFailed(msg)
}
})?;
let response: Value = serde_json::from_str(output).map_err(|e| {
WasmError::ExecutionFailed(format!("Plugin returned invalid JSON: {}", e))
})?;
if let Some(error) = response.get("error") {
let msg = error.as_str().unwrap_or("Unknown plugin error");
return Err(WasmError::ExecutionFailed(msg.to_string()));
}
let result = response.get("ok").cloned().unwrap_or(Value::Null);
normalize_plugin_tool_result(plugin_id, tool_name, result)
}
pub fn is_plugin_loaded(&self, plugin_id: &str) -> bool {
let inner = self.inner.lock();
inner.loaded.contains_key(plugin_id)
}
pub fn is_plugin_healthy(&self, plugin_id: &str) -> bool {
let inner = self.inner.lock();
match inner.loaded.get(plugin_id) {
Some(loaded) => loaded.artifact_path.exists(),
None => false,
}
}
pub fn set_manifest(&self, plugin_id: &str, manifest: PluginManifest) {
let mut inner = self.inner.lock();
if let Some(loaded) = inner.loaded.get_mut(plugin_id) {
loaded.manifest = Some(manifest);
}
}
pub fn get_plugin_manifest(&self, plugin_id: &str) -> Option<PluginManifest> {
let inner = self.inner.lock();
inner.loaded.get(plugin_id).and_then(|l| l.manifest.clone())
}
pub fn loaded_plugins(&self) -> Vec<String> {
let inner = self.inner.lock();
inner.loaded.keys().cloned().collect()
}
}
impl Default for WasmHost {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for WasmHost {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let inner = self.inner.lock();
let ids: Vec<&String> = inner.loaded.keys().collect();
f.debug_struct("WasmHost")
.field("loaded_plugins", &ids)
.finish()
}
}
impl PluginLoader for WasmHost {
fn load(&self, plugin_id: &str, artifact_path: &Path) -> Result<(), String> {
self.load_plugin(plugin_id, artifact_path)
.map_err(|e| e.to_string())
}
fn load_with_limit(
&self,
plugin_id: &str,
artifact_path: &Path,
plugin_declared_max_bytes: Option<u64>,
) -> Result<(), String> {
self.load_plugin_with_limit(plugin_id, artifact_path, plugin_declared_max_bytes)
.map_err(|e| e.to_string())
}
fn unload(&self, plugin_id: &str) {
if let Err(e) = self.unload_plugin(plugin_id) {
warn!(
plugin_id = %plugin_id,
error = %e,
"Failed to unload plugin from WASM host during uninstall"
);
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WasmError {
NotFound(String),
LoadFailed(String),
ExecutionFailed(String),
InvalidInput(String),
PluginPanicked(String),
Timeout,
}
impl std::fmt::Display for WasmError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotFound(id) => write!(f, "Plugin not found: {}", id),
Self::LoadFailed(msg) => write!(f, "Failed to load plugin: {}", msg),
Self::ExecutionFailed(msg) => write!(f, "Tool execution failed: {}", msg),
Self::InvalidInput(msg) => write!(f, "Invalid input: {}", msg),
Self::PluginPanicked(msg) => write!(f, "Plugin panicked: {}", msg),
Self::Timeout => write!(f, "Plugin execution timed out"),
}
}
}
impl std::error::Error for WasmError {}
#[cfg(test)]
mod tests {
use super::*;
use crate::plugin::manifest::{
ExportedTool, PluginIdentity, PluginManifest, PluginPublisher, PresentationMetadata,
};
use crate::plugin::network::NetworkPolicy;
fn sample_manifest() -> PluginManifest {
PluginManifest {
identity: PluginIdentity {
id: "com.example.test".to_string(),
name: "Test".to_string(),
version: "1.0.0".to_string(),
},
publisher: PluginPublisher {
name: "Test".to_string(),
url: None,
contact: None,
},
presentation: PresentationMetadata {
description: "A test plugin".to_string(),
long_description: None,
icon: None,
category: None,
keywords: vec![],
},
network_policy: NetworkPolicy::Wildcard,
auth: None,
tools: vec![ExportedTool {
name: "greet".to_string(),
description: "Say hello".to_string(),
input_schema: serde_json::json!({"type": "object"}),
requires_approval: false,
auth_requirements: None,
}],
max_memory_bytes: None,
api_version: "1.0".to_string(),
}
}
#[test]
fn load_nonexistent_artifact_fails() {
let host = WasmHost::new();
let result = host.load_plugin("test", Path::new("/nonexistent/file.wasm"));
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), WasmError::LoadFailed(_)));
}
#[test]
fn load_minimal_wasm_succeeds_but_no_functions() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("minimal.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::new();
let result = host.load_plugin("test", &artifact);
assert!(
result.is_ok(),
"Extism should accept a minimal valid WASM module, got {:?}",
result
);
assert!(host.is_plugin_loaded("test"));
}
#[test]
fn unload_nonexistent_plugin_is_ok() {
let host = WasmHost::new();
assert!(host.unload_plugin("no-such-plugin").is_ok());
}
#[tokio::test]
async fn execute_tool_on_unloaded_plugin_returns_not_found() {
let host = WasmHost::new();
let result = host
.execute_tool("no-such-plugin", "tool", serde_json::json!({}))
.await;
assert!(matches!(result.unwrap_err(), WasmError::NotFound(_)));
}
#[test]
fn execute_tool_sync_on_unloaded_plugin_returns_not_found() {
let host = WasmHost::new();
let result = host.execute_tool_sync("no-such-plugin", "tool", serde_json::json!({}));
assert!(matches!(result.unwrap_err(), WasmError::NotFound(_)));
}
#[test]
fn is_plugin_loaded_initially_false() {
let host = WasmHost::new();
assert!(!host.is_plugin_loaded("any"));
}
#[test]
fn is_plugin_healthy_initially_false() {
let host = WasmHost::new();
assert!(!host.is_plugin_healthy("any"));
}
#[test]
fn get_plugin_manifest_initially_none() {
let host = WasmHost::new();
assert!(host.get_plugin_manifest("any").is_none());
}
#[test]
fn loaded_plugins_initially_empty() {
let host = WasmHost::new();
assert!(host.loaded_plugins().is_empty());
}
#[test]
fn plugin_loader_trait_load_accepts_minimal_wasm() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("real.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::new();
let result = host.load("test", &artifact);
assert!(
result.is_ok(),
"Extism should accept a minimal valid WASM module"
);
assert!(host.is_plugin_loaded("test"));
}
#[tokio::test]
async fn execute_tool_on_empty_plugin_returns_execution_failed() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("empty.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::new();
host.load_plugin("empty-plugin", &artifact).unwrap();
let result = host
.execute_tool("empty-plugin", "greet", serde_json::json!({}))
.await;
assert!(
matches!(result, Err(WasmError::ExecutionFailed(_))),
"Expected ExecutionFailed for missing function, got {:?}",
result
);
}
#[test]
fn plugin_loader_trait_unload_noop_for_unknown() {
let host = WasmHost::new();
host.unload("no-such-plugin");
assert!(!host.is_plugin_loaded("no-such-plugin"));
}
#[test]
fn debug_impl_works() {
let host = WasmHost::new();
let debug_str = format!("{:?}", host);
assert!(debug_str.contains("WasmHost"));
}
#[test]
fn clone_shares_state() {
let host = WasmHost::new();
let host2 = host.clone();
assert!(host2.loaded_plugins().is_empty());
assert!(host.loaded_plugins().is_empty());
}
#[test]
fn set_manifest_on_nonexistent_plugin_is_noop() {
let host = WasmHost::new();
host.set_manifest("no-such-plugin", sample_manifest());
assert!(host.get_plugin_manifest("no-such-plugin").is_none());
}
#[test]
fn load_and_unload_round_trip() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("rt.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::new();
host.load_plugin("rt", &artifact).unwrap();
assert!(host.is_plugin_loaded("rt"));
assert!(host.loaded_plugins().contains(&"rt".to_string()));
host.unload_plugin("rt").unwrap();
assert!(!host.is_plugin_loaded("rt"));
assert!(!host.loaded_plugins().contains(&"rt".to_string()));
}
#[test]
fn load_same_plugin_id_replaces() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("replace.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::new();
host.load_plugin("dup", &artifact).unwrap();
host.load_plugin("dup", &artifact).unwrap();
assert!(host.is_plugin_loaded("dup"));
}
#[test]
fn health_check_delegates_to_artifact_existence() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("health.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::new();
host.load_plugin("hp", &artifact).unwrap();
assert!(host.is_plugin_healthy("hp"));
std::fs::remove_file(&artifact).unwrap();
assert!(!host.is_plugin_healthy("hp"));
assert!(host.is_plugin_loaded("hp"));
}
#[test]
fn set_and_get_manifest_round_trip() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("manifest.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::new();
host.load_plugin("mp", &artifact).unwrap();
assert!(host.get_plugin_manifest("mp").is_none());
let m = sample_manifest();
host.set_manifest("mp", m.clone());
let retrieved = host.get_plugin_manifest("mp").unwrap();
assert_eq!(retrieved.identity.id, m.identity.id);
assert_eq!(retrieved.tools.len(), m.tools.len());
}
#[test]
fn sync_execution_on_unloaded_returns_not_found() {
let host = WasmHost::new();
let err = host
.execute_tool_sync("nope", "tool", serde_json::json!({}))
.unwrap_err();
assert!(matches!(err, WasmError::NotFound(_)));
}
#[test]
fn sync_execution_on_empty_plugin_returns_execution_failed() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("empty_sync.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::new();
host.load_plugin("empty-sync", &artifact).unwrap();
let err = host
.execute_tool_sync("empty-sync", "greet", serde_json::json!({}))
.unwrap_err();
assert!(
matches!(err, WasmError::ExecutionFailed(ref msg) if msg.contains("does not export")),
"expected ExecutionFailed for missing export, got: {:?}",
err
);
}
#[test]
fn default_constructor_applies_default_memory_ceiling() {
let host = WasmHost::new();
assert_eq!(host.max_memory_bytes(), DEFAULT_PLUGIN_MAX_MEMORY_BYTES);
}
#[test]
fn explicit_memory_ceiling_is_stored() {
let host = WasmHost::with_max_memory_bytes(4 * 1024 * 1024);
assert_eq!(host.max_memory_bytes(), 4 * 1024 * 1024);
}
#[test]
fn plugin_declared_limit_lower_than_host_wins() {
let host_bytes: u64 = 128 * 1024 * 1024;
let plugin_bytes: u64 = 4 * 1024 * 1024;
let effective = plugin_bytes.min(host_bytes);
assert_eq!(effective, plugin_bytes);
assert_eq!(bytes_to_pages(effective), 64);
}
#[test]
fn host_ceiling_overrides_plugin_declared_higher_limit() {
let host_bytes: u64 = 128 * 1024 * 1024;
let plugin_bytes: u64 = 512 * 1024 * 1024;
let effective = plugin_bytes.min(host_bytes);
assert_eq!(effective, host_bytes);
}
#[test]
fn bytes_to_pages_has_minimum_of_one() {
assert_eq!(bytes_to_pages(0), 1);
assert_eq!(bytes_to_pages(1), 1);
assert_eq!(bytes_to_pages(WASM_PAGE_SIZE - 1), 1);
assert_eq!(bytes_to_pages(WASM_PAGE_SIZE), 1);
assert_eq!(bytes_to_pages(WASM_PAGE_SIZE * 2), 2);
}
#[test]
fn load_plugin_with_limit_applies_smaller_of_two() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("cap.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host = WasmHost::with_max_memory_bytes(16 * 1024 * 1024);
host.load_plugin_with_limit("cap", &artifact, Some(4 * 1024 * 1024))
.unwrap();
assert!(host.is_plugin_loaded("cap"));
}
#[test]
fn clone_independent_lifecycle() {
let dir = tempfile::tempdir().unwrap();
let artifact = dir.path().join("clone.wasm");
std::fs::write(&artifact, b"\x00asm\x01\x00\x00\x00").unwrap();
let host1 = WasmHost::new();
let host2 = host1.clone();
host1.load_plugin("shared", &artifact).unwrap();
assert!(host2.is_plugin_loaded("shared"));
host2.unload_plugin("shared").unwrap();
assert!(!host1.is_plugin_loaded("shared"));
}
#[test]
fn error_display_contains_useful_info() {
let err = WasmError::NotFound("my-plugin".to_string());
assert!(err.to_string().contains("my-plugin"));
let err = WasmError::LoadFailed("bad wasm".to_string());
assert!(err.to_string().contains("bad wasm"));
let err = WasmError::Timeout;
assert!(err.to_string().contains("timed out"));
let err = WasmError::InvalidInput("bad json".to_string());
assert!(err.to_string().contains("bad json"));
let err = WasmError::PluginPanicked("trap".to_string());
assert!(err.to_string().contains("trap"));
}
#[test]
fn wasm_error_is_std_error() {
let err = WasmError::LoadFailed("test".to_string());
let _: &dyn std::error::Error = &err;
}
}