soma_provider_adapters/
wasm.rs

1//! The generic WASM provider kind: runs a drop-in `.wasm` module's exported
2//! `soma_call` ABI in a fuel-bounded wasmtime sandbox. Ported unchanged
3//! (beyond product-type swaps) from `soma-service::providers::wasm`.
4
5use std::{fs, path::PathBuf, sync::Arc, time::Duration};
6
7use async_trait::async_trait;
8use serde_json::Value;
9use soma_provider_core::{
10    Provider, ProviderCall, ProviderCatalog, ProviderError, ProviderOutput, ProviderTool,
11};
12use tokio::time::timeout;
13use wasmtime::{Config, Engine, Instance, Memory, Module, Store, TypedFunc};
14
15use crate::sidecar::execution_payload;
16
17#[derive(Clone)]
18pub struct WasmProvider {
19    path: PathBuf,
20    catalog: ProviderCatalog,
21}
22
23impl WasmProvider {
24    pub fn new(path: PathBuf, catalog: ProviderCatalog) -> Self {
25        Self { path, catalog }
26    }
27
28    pub fn arc(path: PathBuf, catalog: ProviderCatalog) -> Arc<Self> {
29        Arc::new(Self::new(path, catalog))
30    }
31}
32
33#[async_trait]
34impl Provider for WasmProvider {
35    fn catalog(&self) -> ProviderCatalog {
36        self.catalog.clone()
37    }
38
39    async fn call(&self, call: ProviderCall) -> Result<ProviderOutput, ProviderError> {
40        let tool = self.tool(&call)?.clone();
41        let provider = self.catalog.provider.name.clone();
42        let action = call.action.clone();
43        let source = self.path.display().to_string();
44        let path = self.path.clone();
45        let input = execution_payload(&call).map_err(|error| {
46            ProviderError::execution(&provider, call.action.clone(), error)
47                .with_provider_kind("wasm")
48                .with_source(source.clone())
49                .with_phase("input-serialization")
50        })?;
51        let limits = WasmRuntimeLimits::from_tool(&tool);
52        if input.len() > limits.max_input_bytes {
53            return Err(ProviderError::validation(
54                provider,
55                call.action,
56                "wasm_input_too_large",
57                format!("WASM input exceeds {} bytes", limits.max_input_bytes),
58            )
59            .with_provider_kind("wasm")
60            .with_source(source)
61            .with_phase("input-validation"));
62        }
63
64        let timeout_ms = limits.timeout_ms;
65        let task = tokio::task::spawn_blocking(move || run_wasm(&path, &input, limits));
66        let output = timeout(Duration::from_millis(timeout_ms), task)
67            .await
68            .map_err(|_| {
69                ProviderError::new(
70                    "wasm_provider_timeout",
71                    &provider,
72                    Some(action.clone()),
73                    format!("WASM provider exceeded {timeout_ms}ms timeout"),
74                    "Increase tool.limits.timeout_ms or fix the WASM provider.",
75                )
76                .with_provider_kind("wasm")
77                .with_source(source.clone())
78                .with_phase("execution")
79            })?
80            .map_err(|error| {
81                ProviderError::execution(&provider, action.clone(), error)
82                    .with_provider_kind("wasm")
83                    .with_source(source.clone())
84                    .with_phase("execution")
85            })?
86            .map_err(|error| {
87                ProviderError::execution(&provider, action.clone(), error)
88                    .with_provider_kind("wasm")
89                    .with_source(source.clone())
90                    .with_phase("execution")
91            })?;
92
93        let value = serde_json::from_slice(&output).map_err(|error| {
94            ProviderError::validation(
95                &provider,
96                &action,
97                "wasm_invalid_json_output",
98                error.to_string(),
99            )
100            .with_provider_kind("wasm")
101            .with_source(source)
102            .with_phase("output-validation")
103        })?;
104        Ok(ProviderOutput::json(value))
105    }
106}
107
108impl WasmProvider {
109    fn tool(&self, call: &ProviderCall) -> Result<&ProviderTool, ProviderError> {
110        self.catalog
111            .tools
112            .iter()
113            .find(|tool| tool.name == call.action)
114            .ok_or_else(|| {
115                ProviderError::validation(
116                    &self.catalog.provider.name,
117                    &call.action,
118                    "unknown_wasm_action",
119                    format!("WASM provider has no action `{}`", call.action),
120                )
121            })
122    }
123}
124
125#[derive(Debug, Clone, Copy)]
126struct WasmRuntimeLimits {
127    timeout_ms: u64,
128    max_input_bytes: usize,
129    max_output_bytes: usize,
130    fuel: u64,
131}
132
133impl WasmRuntimeLimits {
134    fn from_tool(tool: &ProviderTool) -> Self {
135        let meta = tool.meta.get("wasm");
136        Self {
137            timeout_ms: tool
138                .limits
139                .as_ref()
140                .and_then(|limits| limits.timeout_ms)
141                .or_else(|| {
142                    meta.and_then(|value| value.get("timeout_ms"))
143                        .and_then(Value::as_u64)
144                })
145                .unwrap_or(5_000),
146            max_input_bytes: tool
147                .limits
148                .as_ref()
149                .and_then(|limits| limits.max_input_bytes)
150                .unwrap_or(64 * 1024),
151            max_output_bytes: tool
152                .limits
153                .as_ref()
154                .and_then(|limits| limits.max_response_bytes)
155                .unwrap_or(256 * 1024),
156            fuel: meta
157                .and_then(|value| value.get("fuel"))
158                .and_then(Value::as_u64)
159                .unwrap_or(1_000_000),
160        }
161    }
162}
163
164fn run_wasm(
165    path: &std::path::Path,
166    input: &[u8],
167    limits: WasmRuntimeLimits,
168) -> Result<Vec<u8>, String> {
169    let bytes = fs::read(path).map_err(|error| error.to_string())?;
170    let mut config = Config::new();
171    config.consume_fuel(true);
172    let engine = Engine::new(&config).map_err(|error| error.to_string())?;
173    let module = Module::from_binary(&engine, &bytes).map_err(|error| error.to_string())?;
174    let mut store = Store::new(&engine, ());
175    store
176        .set_fuel(limits.fuel)
177        .map_err(|error| error.to_string())?;
178    let instance = Instance::new(&mut store, &module, &[]).map_err(|error| error.to_string())?;
179    let memory = instance
180        .get_memory(&mut store, "memory")
181        .ok_or_else(|| "WASM provider must export memory".to_owned())?;
182    let input_alloc = typed::<i32, i32>(&instance, &mut store, "soma_input_alloc")?;
183    let input_ptr_fn = typed::<(), i32>(&instance, &mut store, "soma_input_ptr")?;
184    let call_fn = typed::<(), i32>(&instance, &mut store, "soma_call")?;
185    let output_ptr_fn = typed::<(), i32>(&instance, &mut store, "soma_output_ptr")?;
186    let output_len_fn = typed::<(), i32>(&instance, &mut store, "soma_output_len")?;
187
188    let ptr = input_alloc
189        .call(&mut store, input.len() as i32)
190        .map_err(|error| error.to_string())? as usize;
191    let input_ptr = input_ptr_fn
192        .call(&mut store, ())
193        .map_err(|error| error.to_string())? as usize;
194    if ptr != input_ptr {
195        return Err("WASM provider input pointer mismatch".to_owned());
196    }
197    write_memory(&memory, &mut store, ptr, input)?;
198    let status = call_fn
199        .call(&mut store, ())
200        .map_err(|error| error.to_string())?;
201    if status != 0 {
202        return Err(format!("WASM provider returned non-zero status {status}"));
203    }
204    let output_ptr = output_ptr_fn
205        .call(&mut store, ())
206        .map_err(|error| error.to_string())? as usize;
207    let output_len = output_len_fn
208        .call(&mut store, ())
209        .map_err(|error| error.to_string())? as usize;
210    if output_len > limits.max_output_bytes {
211        return Err(format!(
212            "WASM provider output exceeds {} bytes",
213            limits.max_output_bytes
214        ));
215    }
216    read_memory(&memory, &mut store, output_ptr, output_len)
217}
218
219fn typed<Params, Results>(
220    instance: &Instance,
221    store: &mut Store<()>,
222    name: &str,
223) -> Result<TypedFunc<Params, Results>, String>
224where
225    Params: wasmtime::WasmParams,
226    Results: wasmtime::WasmResults,
227{
228    instance
229        .get_typed_func(store, name)
230        .map_err(|error| error.to_string())
231}
232
233fn write_memory(
234    memory: &Memory,
235    store: &mut Store<()>,
236    offset: usize,
237    bytes: &[u8],
238) -> Result<(), String> {
239    memory
240        .write(store, offset, bytes)
241        .map_err(|error| error.to_string())?;
242    Ok(())
243}
244
245fn read_memory(
246    memory: &Memory,
247    store: &mut Store<()>,
248    offset: usize,
249    len: usize,
250) -> Result<Vec<u8>, String> {
251    let mut bytes = vec![0; len];
252    memory
253        .read(store, offset, &mut bytes)
254        .map_err(|error| error.to_string())?;
255    Ok(bytes)
256}
257
258#[cfg(test)]
259#[path = "wasm_tests.rs"]
260mod tests;
soma_provider_adapters/wasm.rs

soma_provider_adapters/
wasm.rs
