Skip to content

[Draft] Add AI Capabilities to SpacetimeDB#4589

Draft
definenoob wants to merge 6 commits intoclockworklabs:masterfrom
definenoob:ai
Draft

[Draft] Add AI Capabilities to SpacetimeDB#4589
definenoob wants to merge 6 commits intoclockworklabs:masterfrom
definenoob:ai

Conversation

@definenoob
Copy link

@definenoob definenoob commented Mar 8, 2026

With this PR, STDB becomes the first AI-enabled database. Seeking feedback from CWL on this.

Summary

  • Adds native ONNX inference as a host-side capability exposed to WASM modules via the SpacetimeDB ABI (spacetime_10.5)
  • Models are loaded by name from {server_data_dir}/models/{name}.onnx on the host filesystem and cached after first use — model bytes never enter WASM memory, only input/output
    tensor data crosses the boundary
  • Uses tract-onnx (pure Rust) for inference, no C++ dependencies or WASI required

Design

Single ABI call (spacetime_10.5):

  • onnx_run(name_ptr, name_len, input_ptr, input_len, out) -> errno — loads/caches model by name, runs inference, returns output tensors via BytesSource

Crate changes:

  • spacetimedb-lib — shared Tensor type (Vec shape + Vec data) with BSATN serialization
  • spacetimedb-primitives — ONNX_ERROR errno code
  • spacetimedb-bindings-sys — raw ABI declaration + safe wrapper
  • spacetimedb-bindings — high-level OnnxClient with single run(name, inputs) method, exposed via ReducerContext::onnx and ProcedureContext::onnx
  • spacetimedb-core — host-side OnnxModel using tract-onnx, per-instance model cache (HashMap<String, OnnxModel>) in WasmInstanceEnv, path traversal validation, models_dir wired
    from HostController.data_dir

Usage from a module:

let input = vec![Tensor { shape: vec![1, 10], data: vec![0.0; 10] }];
let output = ctx.onnx.run("bot_brain", &input)?;

Test plan

  • Place a .onnx model file in {data_dir}/models/ and verify ctx.onnx.run("model_name", &input) succeeds
  • Verify inference produces correct output tensors for known model inputs
  • Verify second call to same model uses cache (no re-parse)
  • Verify invalid model names (path traversal, empty) return errors
  • Verify loading a nonexistent model returns a clear error
  • Verify compilation on both native and wasm32-unknown-unknown targets

The onnx feature flag works across all three crates:

  • spacetimedb (bindings): features = ["onnx"] — enables ctx.onnx.run(...) API
  • spacetimedb-bindings-sys: features = ["onnx"] — enables raw ABI declaration
  • spacetimedb-core: features = ["onnx"] — enables host-side inference + tract-onnx dependency

Without the feature, tract-onnx isn't compiled, no ONNX code exists in the binary, and the ABI function isn't registered.

This is ACID-compliant AI inference.

run the model inside the reducer to decide what to write. The inference result feeds directly into the transactional state change — no round-trip to an external service, no eventual consistency, just one atomic unit of "think and act."

@CLAassistant
Copy link

CLAassistant commented Mar 8, 2026

CLA assistant check
All committers have signed the CLA.

@definenoob definenoob marked this pull request as draft March 8, 2026 11:24
@cloutiertyler
Copy link
Contributor

@definenoob have you tested this?

@definenoob
Copy link
Author

definenoob commented Mar 9, 2026

@cloutiertyler Not yet, I was wondering if you thought it would be worth pursuing before diving too much deeper. Training a neural network to do something interesting in a whole game isn't exactly light work!

@definenoob
Copy link
Author

definenoob commented Mar 9, 2026

Testing the ONNX feature

Unit tests

The host-side inference logic has unit tests that run without any external dependencies — models are built programmatically using tract_onnx::pb protobuf types:

cargo test -p spacetimedb-core --features onnx -- onnx::tests

This covers: model loading from bytes, loading from filesystem, Add/Relu inference, run_multi batching, invalid model handling, shape mismatch errors, and path traversal rejection.

Integration tests (smoketests)

End-to-end tests verify the full WASM module → FFI → host path. The server must be built with the onnx feature:

  1. Build binaries with onnx support
    cargo build -p spacetimedb-cli -p spacetimedb-standalone --release \
    --features "spacetimedb-standalone/allow_loopback_http_for_tests,spacetimedb-standalone/onnx"

  2. Run the ONNX smoketests
    cargo test -p spacetimedb-smoketests -- onnx

The smoketests compile a WASM module with onnx bindings at runtime, place a test .onnx model file in the server's data directory, then call reducers that invoke ctx.onnx.run() and ctx.onnx.run_multi() and verify the inference results in the
module logs.

What's tested across the two levels

Layer Unit tests Smoketests
OnnxModel load + inference (tract)
run_multi batching
BSATN tensor serialization across WASM boundary
onnx_run / onnx_run_multi FFI host functions
Model caching in WasmInstanceEnv
Filesystem model loading (models_dir)
Error handling (bad bytes, shape mismatch)
Path traversal validation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants