Machine Learning One

We have an engine that compiles WAT, a bump allocator that manages linear memory, and a wire protocol that serializes Rust types into bytes. But WASM programs that can only do arithmetic on their own memory are useless. They need to interact with the host: fetch URLs, query AI models, store data. The question is how to let them do this without blowing a hole in the sandbox.

The answer is capabilities. A capability is a host function that the guest can call — and the guest can only call capabilities that have been explicitly provided. No import, no access. This is the inverse of a blocklist: instead of denying specific dangerous operations, we grant specific safe ones. Everything else is impossible by construction.

This article builds the trait that defines a capability, the type-keyed map that gives each capability its own state, the runtime state struct that ties it all together, and the first capability implementation: CoreCapability, which provides the fundamental ABI for argument passing and result storage.

The Capability Trait

A capability needs to do four things:

Identify itself — for logging, error messages, and LLM prompts
Declare its WAT imports — so the template can assemble valid modules
Register host functions — so the linker can resolve those imports
Initialize per-instance state — so each WASM execution gets fresh state

Here is the trait:

// rt/src/capability.rs

pub trait Capability: Send + Sync + 'static {
    /// Human-readable name (e.g. "core", "http", "ai").
    fn name(&self) -> &str;

    /// WAT import declarations for this capability.
    /// Example: `(import "http" "get" (func $http.get (param i32) (result i32 i32)))`
    fn wat_imports(&self) -> &str;

    /// Human-readable description of each import for LLM prompts.
    fn describe(&self) -> &str {
        ""
    }

    /// Register host functions on the linker.
    fn register(&self, linker: &mut wasmtime::Linker<State>) -> Result<()>;

    /// Initialize per-instance state in the ExtensionMap.
    fn init_state(&self, _ext: &mut ExtensionMap) -> Result<()> {
        Ok(())
    }
}

The Send + Sync + 'static bounds are required because capabilities are shared across threads — they are registered once on a Linker and used to instantiate modules on any tokio task.

Two methods have default implementations. describe() returns an empty string because not every capability needs to advertise itself to the LLM. init_state() is a no-op because not every capability has per-instance state — CoreCapability, for example, uses fields already present on State.

The register() method takes a &mut wasmtime::Linker<State>, not our wrapper Linker. This is deliberate — it gives capability implementors direct access to wasmtime's linker API, so they can use func_wrap, func_wrap_async, or any other registration method.

ExtensionMap: Type-Keyed State

Different capabilities need different kinds of state. The HTTP capability needs an allowlist of domains. The KV capability needs a HashMap<String, Vec<u8>>. The filesystem capability needs a jail root path. We could add fields to State for each one, but that would couple the runtime to every capability that will ever exist.

Instead, we use a type-keyed map. Each capability stores its own state type, and retrieves it by type. Two different capabilities can never collide because they use different Rust types as keys.

// rt/src/capability.rs

pub struct ExtensionMap {
    map: HashMap<TypeId, Box<dyn Any + Send>>,
}

impl ExtensionMap {
    pub fn new() -> Self {
        ExtensionMap { map: HashMap::new() }
    }

    pub fn insert<T: Any + Send>(&mut self, val: T) {
        self.map.insert(TypeId::of::<T>(), Box::new(val));
    }

    pub fn get<T: Any + Send>(&self) -> Option<&T> {
        self.map.get(&TypeId::of::<T>()).and_then(|b| b.downcast_ref())
    }

    pub fn get_mut<T: Any + Send>(&mut self) -> Option<&mut T> {
        self.map.get_mut(&TypeId::of::<T>()).and_then(|b| b.downcast_mut())
    }
}

The API is minimal: insert, get, get_mut. The TypeId::of::<T>() call produces a unique identifier for each concrete type, and downcast_ref()/downcast_mut() recover the concrete type from the Box<dyn Any>. If you insert a KvState, you get it back with extensions.get::<KvState>(). If you ask for a type that was never inserted, you get None.

This pattern is well-known in the Rust ecosystem — Axum uses it for request extensions, hyper uses it for HTTP extensions, and actix-web uses it for app data. It trades compile-time type safety for runtime flexibility, which is exactly the tradeoff we need for a plugin system.

ExtensionMap Lifecycle

The ExtensionMap lives inside State, which is owned by a wasmtime::Store<State>. Each WASM execution gets its own Store, which means each execution gets its own ExtensionMap. The lifecycle is:

Creation: Before instantiation, LinkedModule.instantiate() creates a fresh ExtensionMap and calls init_state() on each capability, giving each one a chance to insert its state.
Active use: During WASM execution, host functions access the map through caller.data().extensions (read) or caller.data_mut().extensions (write). Since a Store is single-threaded (only one host function runs at a time per instance), no locking is needed inside the map itself.
Destruction: When the Instance is dropped, the Store is dropped, which drops the State, which drops the ExtensionMap and all its contents.

One important exception: capabilities that need state to survive across multiple WASM executions (like the KV store in a ReAct loop) store an Arc in the map. The Arc is cloned into each new ExtensionMap, so the underlying data is shared while the map itself is per-instance. We will see this pattern in Article 7.

How Capabilities Use It

A capability initializes its state during init_state() and accesses it during host function execution via the Caller:

// In a capability's init_state():
fn init_state(&self, ext: &mut ExtensionMap) -> Result<()> {
    ext.insert(KvState { store: HashMap::new() });
    Ok(())
}

// In a host function registered by that capability:
linker.func_wrap("kv", "get", |mut caller: Caller<'_, State>, key_ptr: u32| {
    let kv = caller.data().extensions.get::<KvState>().unwrap();
    // ... use kv.store ...
});

The State Struct

Every WASM execution gets its own State. This struct is what wasmtime's Store<State> holds — it's accessible from any host function via caller.data() and caller.data_mut().

// rt/src/instance.rs

pub struct State {
    pub memory: WasmMemory,
    pub args: Vec<Vec<u8>>,
    pub results: Vec<Vec<u8>>,
    pub arg_pointers: Vec<Option<u32>>,
    pub nil_pointer: Option<u32>,
    pub extensions: ExtensionMap,
}

Each field has a specific role:

memory: The bump allocator from Article 4. Manages the top pointer and enforces the memory limit.
args: Pre-serialized argument bytes. The orchestrator serializes arguments before execution; sys.argv writes them into WASM memory on demand.
results: Collected result bytes. Each sys.resv call reads data from WASM memory and pushes it here. After execution, the orchestrator deserializes these.
arg_pointers: Cache of already-materialized argument pointers. sys.argv(0) writes argument 0 into memory the first time; subsequent calls return the cached pointer.
nil_pointer: Cached pointer to a serialized None value. Many programs need to return "nothing" — caching it avoids repeated allocation.
extensions: The ExtensionMap for capability-specific state.

The args and results vectors form a one-way data pipeline: arguments flow into the WASM program (host writes to memory, guest reads), and results flow out (guest writes to memory, host reads back). This is the only data path between host and guest.

Why Lazy Argument Materialization?

Arguments are stored as Vec<Vec<u8>> — pre-serialized byte vectors — not as pointers into WASM memory. They are only written into WASM memory when the guest calls sys.argv(index). This is deliberate.

A program might receive three arguments but only use one. Writing all three into WASM memory upfront wastes allocator space and memory pages. Lazy materialization means we only pay for arguments the guest actually reads. The arg_pointers cache ensures we never write the same argument twice.

The Instance Struct

An Instance wraps a running WASM execution: the Store<State> plus the typed function handle for the run export.

// rt/src/instance.rs

pub struct Instance {
    store: wasmtime::Store<State>,
    run_fn: wasmtime::TypedFunc<(), u32>,
}

impl Instance {
    pub(crate) fn new(
        store: wasmtime::Store<State>,
        run_fn: wasmtime::TypedFunc<(), u32>,
    ) -> Self {
        Instance { store, run_fn }
    }

    /// Execute the WASM program. Returns the exit code.
    pub async fn run(&mut self) -> Result<u32> {
        let result = self
            .run_fn
            .call_async(&mut self.store, ())
            .await
            .map_err(Error::from)?;
        Ok(result)
    }

    /// Get a specific result by index.
    pub fn result<T: DeserializeOwned>(&self, index: usize) -> Result<T> {
        let raw = self.store.data().results.get(index).ok_or_else(|| {
            Error::ArgOutOfBounds {
                index: index as u32,
                count: self.store.data().results.len() as u32,
            }
        })?;
        Ok(protocol::from_bytes(raw)?)
    }

    /// Get all results.
    pub fn results<T: DeserializeOwned>(&self) -> Result<Vec<T>> {
        self.store
            .data()
            .results
            .iter()
            .map(|raw| protocol::from_bytes(raw).map_err(Error::from))
            .collect()
    }

    /// Number of results stored.
    pub fn result_count(&self) -> usize {
        self.store.data().results.len()
    }

    /// Access raw result bytes.
    pub fn raw_results(&self) -> &[Vec<u8>] {
        &self.store.data().results
    }
}

The run() method uses call_async because some host functions (HTTP requests, AI inference) are async. Even for synchronous-only programs, we use async execution — the overhead is negligible, and it means we don't need separate sync/async code paths.

The result<T>() method deserializes a result at a given index using the wire protocol from Article 4. The caller specifies the expected type via the turbofish: instance.result::<String>(0). If the bytes don't match the expected type, deserialization fails with a protocol error.

The result_count() and raw_results() methods give the orchestrator access to raw result data without forcing deserialization, which is useful when the orchestrator needs to pass results as-is to the next program.

CoreCapability: The Foundational ABI

Every WASM program needs four things: memory allocation, argument access, result storage, and a way to represent "no value." These are provided by CoreCapability, which is always registered.

The Import Declarations

fn wat_imports(&self) -> &str {
    r#"(import "sys" "alloc" (func $sys.alloc (param i32) (result i32)))
(import "sys" "argv" (func $sys.argv (param i32) (result i32 i32)))
(import "sys" "resv" (func $sys.resv (param i32)))
(import "sys" "nil" (func $sys.nil (result i32)))"#
}

Note the signatures:

sys.alloc(size) -> ptr — one parameter, one result
sys.argv(idx) -> (ptr, err) — one parameter, two results (multi-value return)
sys.resv(ptr) — one parameter, no result (side effect: pushes to State.results)
sys.nil() -> ptr — no parameters, one result

sys.argv uses WASM's multi-value return to give both a pointer and an error code. The guest checks the error code first: if it's non-zero, the pointer is meaningless.

The Host Function Implementations

Here is the full register method:

// capability/src/core.rs

fn register(&self, linker: &mut Linker<State>) -> rt::Result<()> {
    // sys.alloc — bump allocate `size` bytes, return pointer
    linker
        .func_wrap("sys", "alloc", |mut caller: Caller<'_, State>, size: u32| -> u32 {
            caller_alloc(&mut caller, size).unwrap_or(0)
        })
        .map_err(rt::Error::Other)?;

    // sys.argv — lazy-materialize argument into WASM memory, cache pointer
    linker
        .func_wrap(
            "sys",
            "argv",
            |mut caller: Caller<'_, State>, idx: u32| -> (u32, u32) {
                let count = caller.data().args.len() as u32;
                if idx >= count {
                    return (0, rt::EBOUND);
                }
                // Return cached pointer if already materialized
                if let Some(ptr) = caller.data().arg_pointers[idx as usize] {
                    return (ptr, rt::SUCCESS);
                }
                // Materialize: write serialized arg bytes into WASM memory
                let arg_data = caller.data().args[idx as usize].clone();
                match caller_write_raw(&mut caller, &arg_data) {
                    Ok(ptr) => {
                        caller.data_mut().arg_pointers[idx as usize] = Some(ptr);
                        (ptr, rt::SUCCESS)
                    }
                    Err(_) => (0, rt::ENOMEM),
                }
            },
        )
        .map_err(rt::Error::Other)?;

    // sys.resv — store a result from WASM memory
    linker
        .func_wrap(
            "sys",
            "resv",
            |mut caller: Caller<'_, State>, ptr: u32| {
                if let Ok(bytes) = caller_read_bytes(&mut caller, ptr) {
                    caller.data_mut().results.push(bytes);
                }
            },
        )
        .map_err(rt::Error::Other)?;

    // sys.nil — return pointer to serialized None::<u8>, cached after first call
    linker
        .func_wrap(
            "sys",
            "nil",
            |mut caller: Caller<'_, State>| -> u32 {
                if let Some(ptr) = caller.data().nil_pointer {
                    return ptr;
                }
                let nil_bytes: [u8; 5] = [1, 0, 0, 0, 0];
                match caller_write_raw(&mut caller, &nil_bytes) {
                    Ok(ptr) => {
                        caller.data_mut().nil_pointer = Some(ptr);
                        ptr
                    }
                    Err(_) => 0,
                }
            },
        )
        .map_err(rt::Error::Other)?;

    Ok(())
}

Walking Through Each Function

sys.alloc(size) -> ptr

The simplest function. Delegates to caller_alloc() (the borrow-splitting helper from Article 4), which bumps the allocator and grows pages if needed. Returns 0 on failure — the guest should treat pointer 0 as an error, though in practice allocation failures are rare because we set generous memory limits.

sys.argv(idx) -> (ptr, err)

This is the most interesting function. The flow:

Check idx < args.len(). If out of bounds, return (0, EBOUND).
Check the arg_pointers cache. If this argument was already materialized, return the cached pointer with SUCCESS.
Clone the argument bytes (we need owned data to pass to caller_write_raw).
Write the bytes into WASM memory. This allocates space and copies the pre-serialized data.
Cache the resulting pointer in arg_pointers[idx].
Return (ptr, SUCCESS).

The clone in step 3 is necessary because caller.data().args borrows the Caller, and caller_write_raw needs &mut Caller. We can't hold both borrows simultaneously. The cost is one allocation per argument per execution, which is negligible.

sys.resv(ptr)

Reads length-prefixed data from WASM memory at ptr and pushes it onto State.results. This is how the guest "returns" data to the host. A program can call sys.resv multiple times to store multiple results.

The read uses caller_read_bytes, which reads the 4-byte length prefix at ptr, then reads length more bytes. The entire slice (including the length prefix) is stored in results, so the host can later deserialize it with protocol::from_bytes().

sys.nil() -> ptr

Returns a pointer to the serialized representation of None. The byte pattern is [1, 0, 0, 0, 0] — a 4-byte length prefix (value: 1) followed by a single discriminant byte (value: 0, meaning None). This matches the wire protocol's encoding of Option::None.

The pointer is cached in State.nil_pointer after the first call. This matters because many programs call sys.nil() in error-handling branches, and we don't want to allocate 5 bytes of memory for every error check.

Tests

argv/resv Roundtrip

The fundamental test: pass a string argument, have the WASM program read it with sys.argv and store it with sys.resv, then verify the host can deserialize it back.

#[tokio::test]
async fn argv_resv_roundtrip() {
    let engine = Engine::new(EngineConfig::default()).unwrap();
    let template = build_template();

    let body = r#"
(local $ptr i32)
(local $err i32)
(call $sys.argv (i32.const 0))
(local.set $err)
(local.set $ptr)
(if (i32.ne (local.get $err) (i32.const 0))
    (then (return (local.get $err)))
)
(call $sys.resv (local.get $ptr))
(i32.const 0)
"#;
    let ar = template.assemble(body).unwrap();
    let module = engine.compile_wat(&ar.wat).unwrap();
    let mut linker = Linker::new(&engine);
    linker.register(&CoreCapability).unwrap();
    let linked = linker.pre_link(&module).unwrap();

    let args = vec![rt::serialize_arg(&"hello world").unwrap()];
    let mut instance = linked
        .instantiate(&engine, ExtensionMap::new(), args, vec![], 0)
        .await
        .unwrap();

    let code = instance.run().await.unwrap();
    assert_eq!(code, 0);
    assert_eq!(instance.result_count(), 1);
    let result: String = instance.result(0).unwrap();
    assert_eq!(result, "hello world");
}

The WAT body reads argv[0] (which pushes (ptr, err) onto the WASM stack), checks the error code, stores the result via sys.resv, and returns 0 (success). On the Rust side, we serialize "hello world" as an argument, run the program, and verify the result deserializes back to the same string.

argv Out of Bounds

Requesting an argument index that doesn't exist should return EBOUND (5):

#[tokio::test]
async fn argv_out_of_bounds() {
    let engine = Engine::new(EngineConfig::default()).unwrap();
    let template = build_template();

    let body = r#"
(local $ptr i32)
(local $err i32)
(call $sys.argv (i32.const 5))
(local.set $err)
(local.set $ptr)
(local.get $err)
"#;
    let ar = template.assemble(body).unwrap();
    let module = engine.compile_wat(&ar.wat).unwrap();
    let mut linker = Linker::new(&engine);
    linker.register(&CoreCapability).unwrap();
    let linked = linker.pre_link(&module).unwrap();

    let mut instance = linked
        .instantiate(&engine, ExtensionMap::new(), vec![], vec![], 0)
        .await
        .unwrap();
    let code = instance.run().await.unwrap();
    assert_eq!(code, rt::EBOUND);
}

No arguments are provided, but the WAT body requests argv[5]. The exit code is EBOUND (5), telling the orchestrator that the program tried to access an argument that doesn't exist.

nil Returns a Valid Pointer

sys.nil() should return a pointer that deserializes to None:

#[tokio::test]
async fn nil_returns_valid_pointer() {
    let engine = Engine::new(EngineConfig::default()).unwrap();
    let template = build_template();

    let body = r#"
(call $sys.resv (call $sys.nil))
(i32.const 0)
"#;
    let ar = template.assemble(body).unwrap();
    let module = engine.compile_wat(&ar.wat).unwrap();
    let mut linker = Linker::new(&engine);
    linker.register(&CoreCapability).unwrap();
    let linked = linker.pre_link(&module).unwrap();

    let mut instance = linked
        .instantiate(&engine, ExtensionMap::new(), vec![], vec![], 0)
        .await
        .unwrap();
    let code = instance.run().await.unwrap();
    assert_eq!(code, 0);
    assert_eq!(instance.result_count(), 1);
    let result: Option<u8> = instance.result(0).unwrap();
    assert_eq!(result, None);
}

The WAT body calls sys.nil() to get a pointer, passes it directly to sys.resv to store it as a result, and returns success. The host deserializes the result as Option<u8> and verifies it's None.

The Data Flow

Here is how data moves through a single program execution:

Host (Rust)                           Guest (WASM)
────────────                          ────────────
serialize_arg("hello")
    │
    ▼
State.args = [serialized bytes]
                                      call $sys.argv (i32.const 0)
                                          │
                                          ▼
                            ┌─── caller_write_raw(arg_data) ───┐
                            │   writes into linear memory      │
                            │   caches pointer                 │
                            └──────────────────────────────────┘
                                          │
                                          ▼
                                      (ptr, SUCCESS)
                                          │
                                      ... compute ...
                                          │
                                      call $sys.resv (local.get $ptr)
                                          │
                                          ▼
                            ┌─── caller_read_bytes(ptr) ───────┐
                            │   reads from linear memory       │
                            │   pushes to State.results        │
                            └──────────────────────────────────┘
                                          │
instance.result::<String>(0)              │
    │                                     │
    ▼                                     │
protocol::from_bytes(results[0])          │
    │                                     │
    ▼                                     │
"hello" ◄─────────────────────────────────┘

Arguments go in through sys.argv, results come out through sys.resv. The wire protocol handles serialization in both directions. The bump allocator manages the memory. The capability trait defines the interface. Each layer does one thing.

What We Have So Far

After this article, our system has:

A trait for defining host function plugins (Capability)
A type-keyed extension map for per-instance capability state (ExtensionMap)
A runtime state struct that holds memory, arguments, results, and extensions (State)
A wrapper for running WASM instances and extracting results (Instance)
A concrete implementation of the core ABI: sys.alloc, sys.argv, sys.resv, sys.nil
Tests proving the argument-result roundtrip works end-to-end

What it can't do yet: assemble a complete WAT module from a function body and capability imports, or link capabilities into a reusable pre-compiled module. For that, we need the Template and Linker.

Capability-Based Security: The Plugin System