How I ported a Rust framework to TS with WASM

During one of our internal weekly meetings at work, we decided that we wanted to get Rig to JavaScript so that we can run it in the browser and in Node applications. Of course, there was a potential tonne of code rewriting to do if we were to write it entirely from TypeScript. Therefore, we ended up using WASM to port the inner core to TypeScript and then using manual glue code from TypeScript to be able to finish off the parts that couldn't be done directly through the Rust-WASM bridge.

What needs to get ported?

Here is a short list of things we were able to port from Rust:

• All model providers
• Functionality for users to create their own vector store and tool integrations from JS/TS
• Wrapper types and conversion impls to/from non-WASM Rust types

Thankfully, a large portion of the WASM implementation is mostly just wrappers around the original structs. We just need to move them all over to WASM-friendly bindings by attaching the #[wasm_bindgen] macro onto newtype-tuple pattern structs. Since we didn't need to export all of the model behaviour from Rust, the implementation was made significantly more simple.

However, there is a large portion of the library that uses generics and traits. Unfortunately as JavaScript doesn't understand lifetimes, this means we have to create a struct that represents each generic variant. For example, we have the Agent struct (which represents, in short, an LLM wrapper):

struct Agent<M> {
    model: M,
    // .. some other fields here
}

impl Agent<M> where M: CompletionModel {
    // methods, etc
}

The following would have to be created in our new rig-wasm crate:

struct OpenAIAgent {
    // .. fields here
}

struct AnthropicAgent {
    // .. fields here
}

There is also the need for users to be able to generate their own types and use it with the library. For example, when you're using an agent you can attach some tools to it which then get called by the language model. You might also want to attach a vector store onto it for RAG (eg Qdrant or Pinecone).

Finally, we should ensure the final resulting library uses idiomatic TypeScript. This is technically optional. However to ensure that users don't go through hell and back trying to use our library, we'll be doing it anyway. This means that we will have to write some TypeScript code ourselves. No worries - at least it's not raw JavaScript. It could always be worse.

Regenerating provider modules

One of the issues that we had originally was that we support about 20 different providers to date - that's a lot of code to re-write. One of the solutions that I proposed was using a build script to regenerate all of the original provider modules based on whatever traits they implement and then use a Jinja template.

It worked for the most part - the primary concern is that the build script can be quite brittle because of the Jinja templating. You can see the build file, which is pretty much just a load of string parsing (and then conditionally adding more code to the final file that gets produced).

One thing you will notice if you check the provider modules for the relevant Agent struct(s) is that we actually take a JavaScript value as the parameter for fn new(), then check it for the required shape and ensure it satisfies all the requirements where required - and then (and only then!) do we create the relevant struct. Fortunately with types that are JSON-compatible, we can simply deserialize it into the relevant Rust struct. However, this is not possible with types in JS/TS that have function signatures for obvious reasons - so we parse each field manually.

User integrations

Alright, so here's where most of the technical complexity starts. As mentioned before we have a couple of integration types that hook directly into our Agent type, and they are as follows:

• Vector stores (used in dynamic tool RAG-ing and dynamic document context)
• Tools (called by the LLM)

Both traits are mostly all-async functions and require Send + Sync futures to be returned. While this has been useful in Rust to ensure runtime properties, it also can be a huge nuisance sometimes. This is one of those times. Fortunately, we have an escape hatch in the form of wasm_bindgen_futures::spawn_blocking - which allows us to run said JavaScript (and Rust!) functions. Unfortunately, it doesn't allow us to return a raw value.

A solution to this that I ended up using was to essentially use oneshot channels to return the error (or the raw result), then use a tokio::select! macro:

fn call(
        &self,
        args: Self::Args,
    ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + Sync {
        let (result_tx, result_rx) = futures::channel::oneshot::channel();
        let (error_tx, error_rx) = futures::channel::oneshot::channel();
        let js_args: JsValue = serde_wasm_bindgen::to_value(&args).expect("This should be a JSON object!");

        let func = self.inner.clone();

        spawn_local(async move {
            // .. do some work here

            let res = match do_some_work() {
                Ok(res) => res,
                Err(e) => {
                    error_tx
                        .send(e)
                        .expect("shouldn't fail while sending to a oneshot channel");
                return;
                }

            }

            result_tx
                .send(res)
                .expect("sending a message to a oneshot channel shouldn't fail");
        });

        async {
            tokio::select! {
                res = result_rx => {
                    Ok(res.unwrap())
                },
                err = error_rx => {
                    Err(ToolError::ToolCallError(err.inspect_err(|x| println!("Future was cancelled: {x}")).unwrap().to_string().into()))
                }
            }
        }
    }

This essentially means that whichever channel receives a value first is the channel that a returned value will be pulled from (ie, if I have an error channel & a result channel, if the error channel receives an error first, it'll return the error from the function - if no error gets sent, ie the function fully completed, it returns the result). It's not what I would consider perfect considering that sending to a oneshot channel can technically fail and I had to include an initPanicHook() function so I can track panics from the JS/TS side, but it works and does so well enough that I haven't had to worry about it yet.

Here is an example of using a tool in rig-wasm:

import { Agent } from "@0xplaygrounds/rig-wasm/openai"

const counter = {
  counter: 5324,
  name: function () {
    return "counter";
  },
  definition: function (_prompt: string) {
    return {
      name: "counter",
      description: "a counter that can only be incremented",
      parameters: {
        $schema: "https://json-schema.org/draft/2020-12/schema",
        title: "ToolDefinition",
        type: "object",
        properties: {},
        required: [],
      },
    };
  },
  call: async function (args: any) {
    this.counter += 1;
    return { result: this.counter };
  },
};

let prompt = `Please increment the counter by 1 and let me know what the resulting number is, as well as the original number.`;
let key = process.env.OPENAI_API_KEY;

try {
  const agent = new Agent({
    apiKey: key,
    model: "gpt-4o",
    tools: [counter],
  });

  console.log(`Prompt: ${prompt}`);
  let res = await agent.prompt(prompt);
  console.log(`GPT-4o: ${res}`);
} catch (e) {
  if (e instanceof Error) {
    console.log(`Error while prompting: ${e.message}`);
  }
}

As you can see, we don't really need to do anything out of the ordinary here. Calling the agent should just tell us that the counter has increased by 1, and that the new total is 5325. Internally Rig already supports multi-turn calling.

Unfortunately while we have a plethora of vector database integrations, we basically have to write them in TS manually. For this, we had to write an interface that looks like this:

export interface VectorStore {
  topN: (req: VectorSearchOpts) => Promise<TopNResult[]>;
  topNIds: (req: VectorSearchOpts) => Promise<TopNIdsResult[]>;
}

Then we had to write classes that basically had those function identifiers (and matched the typedefinitions). The Qdrant integration for example looks like this:

export class QdrantAdapter {
  private client: QdrantClient;
  private collectionName: string;
  private params: QdrantClientParams;
  private embeddingModel: CanEmbed;

  constructor(
    collectionName: string,
    embeddingModel: CanEmbed,
    params: QdrantClientParams,
  ) {
    this.params = params;
    this.embeddingModel = embeddingModel;
    this.collectionName = collectionName;
  }

  async loadClient() {
    if (!this.client) {
      try {
        this.client = new QdrantClient(this.params);
      } catch (err) {
        throw new Error("Failed to load Qdrant client: " + err);
      }
    }
  }

  async init(dimensions: number) {
    await this.loadClient();

    const collections = await this.client.getCollections();
    const exists = collections.collections.some(
      (c: any) => c.name === this.collectionName,
    );

    if (!exists) {
      await this.client.createCollection(this.collectionName, {
        vectors: {
          size: dimensions,
          distance: "Cosine",
        },
      });
    }
  }

  async topN(opts: VectorSearchOpts): Promise<SearchResult[]> {
    await this.loadClient();

    const embedding = await this.embeddingModel.embedText(opts.query);

    const result = await this.client.search(this.collectionName, {
      vector: embedding.vec,
      limit: opts.samples,
    });

    return result.map((res: any) => ({
      id: res.id,
      score: res.score,
      payload: res.payload,
    }));
  }

  async topNIds(opts: VectorSearchOpts): Promise<SearchResult[]> {
    await this.loadClient();

    const embedding = await this.embeddingModel.embedText(opts.query);

    const result = await this.client.search(this.collectionName, {
      vector: embedding.vec,
      limit: opts.samples,
    });

    return result.map((res: any) => ({
      id: res.id,
      score: res.score,
    }));
  }
}

Basically speaking, as long as it implements topN and topNIds, it should satisfy the criteria to be considered a VectorStore.

Build pipelines

Another aspect of actually preparing the library for shipping to production was figuring out how to get the build tooling working. Admittedly, this was my first time ever shipping my own JS/TS library at all, so I ended up having more friction than perhaps someone who's more experienced with the JS/TS build tooling ecosystem might be. I made judicious use of LLM prompting here to solve specific issues, and then checked the documentation to ensure that everything was totally correct and figure out why some things did work (or not) after the fact.

Fortunately, the Rust part was simple: build the binary for the target then use the wasm-bindgen CLI. At the time of writing the build pipeline (roughly a couple of weeks ago) I wanted to experiment with going much more level and basically just using wasm-bindgen-cli itself. However, if you wanted to you could definitely do this with something like wasm-pack. Perhaps at some point, it would be ideal to revisit this aspect of the build pipeline.

In terms of JS build tooling, I chose Rollup as it seemed quite minimal, is good for making smaller bundles and has a lot of different plugins you can play around with. I also didn't really feel like I needed all of the features that Webpack has. However, if your requirements differ and you aren't just building a library, you may want to use Webpack. It's much more popular.

There were quite a few gotchas that I encountered along the way while trying to get the thing to build (mostly skill issues, but perhaps some valuable advice for anyone who's bad at TypeScript trying to do the same thing):

• The node target for wasm-bindgen-cli actually compiles to CommonJS, but most modern projects actually use ESM. We ended up using the experimental-nodejs-module target. Not very ideal to be honest, but it did work and I'm quite thankful for that. Hopefully, this target will stabilise at some point.
• I tried to use .d.ts files in my library alongside my .ts files because I needed to declare some typedefs. It turns out that if you import anything, they are now not legitimate .d.ts files anymore. I realised after several Google searches that this is in fact an extremely bad idea, not to mention that this actually caused the compilation to fail. I ended up just writing TS files.
• Setting up tsconfig.json was quite a task in itself (as someone who previously interacted very rarely with the tsconfig file, even when I primarily used TS as a hobbyist).
• rollup does not automatically add the wasm files to your dist directory (or wherever you have the final compiled output), so you've got to copy them manually. Fortunately this is probably the least of your issues here - you can just run cp <foo> <bar> on each of the required WASM files and it should do the job.

The other major issue was trying to add all of the exports/inputs to the package.json file and the rollup.config.ts file. Fortunately, the rollup file issue was simple. Iterate through a directory, check the file extensions and then do some string operations:

const readDir = (inputDir: string): string[] => {
  const files = fs
    .readdirSync(inputDir)
    .filter((file: string) => file.endsWith(".ts")) 
    .map((file: string) => path.join(inputDir, file));

  return files;
};

const input: Record<string, string> = {};

const dirs = [providersDir, vectorStoresDir, coreDir];

for (const dir of dirs) {
  for (const file of readDir(dir)) {
    const name = path.basename(file, ".ts");
    input[name] = file;
  }
}

Unfortunately, the package.json file issue was not quite so simple. I ended up basically creating a Bash script that iterated through every file/directory in src, then operated on the file/folder names and created raw JSON. At the end, it then used jq to edit the package.json file so that it would set the exports as the newly created list.

A short excerpt of the Bash script:

echo '{' > "$EXPORTS_FILE"

# Add the root export manually
echo '  ".": {' >> "$EXPORTS_FILE"
echo '    "import": "./dist/esm/index.js",' >> "$EXPORTS_FILE"
echo '    "require": "./dist/cjs/index.cjs",' >> "$EXPORTS_FILE"
echo '    "types": "./dist/esm/index.d.ts"' >> "$EXPORTS_FILE"
echo '  },' >> "$EXPORTS_FILE"

# Gather all .js files (excluding index.js)
mapfile -t core_files < <(find "$CORE_MODULE" -maxdepth 1 -type f -name "*.ts" | sort)
echo "Found ${#core_files[@]} core modules. Building exports..."

for i in "${!core_files[@]}"; do
  file="${core_files[$i]}"
  base=$(basename "$file" .js)

  echo "  \"./${base%.ts}\": {" >> "$EXPORTS_FILE"
  echo "    \"import\": \"./$OUT_DIR/${base%.ts}.js\"," >> "$EXPORTS_FILE"
  echo "    \"types\": \"./$OUT_DIR/${base%.ts}.d.ts\"" >> "$EXPORTS_FILE"
  if [[ $i -lt $((${#providers_files[@]} - 1)) ]]; then
    echo "  }," >> "$EXPORTS_FILE"
  else
    echo "  }" >> "$EXPORTS_FILE"
  fi
done

echo "}" >> "$EXPORTS_FILE"

echo "✅ Wrote updated exports to $EXPORTS_FILE"

jq --slurpfile exports exports.json '.exports = $exports[0]' package.json > package.new.json && mv package.new.json package.json

echo "✅ Updated exports in package.json"

rm $EXPORTS_FILE

I will admit that I'm not a Bash expert and had an LLM spin up the skeleton of the script, then I finished off what parts I couldn't get the LLM to figure out in a reasonable amount of time (~5 minutes). It was mostly quite quick and painless. I am sure there is probably a more reasonable (and maintainable!) way to write this, but it works for now.

Now that I'd written the script, the build pipeline was actually mostly done. I just had to set up the final build script:

#!/usr/bin/env bash
cargo build -p rig-wasm --release --target wasm32-unknown-unknown
wasm-bindgen \
    --target experimental-nodejs-module \
    --out-dir rig-wasm/pkg/src/generated \
    target/wasm32-unknown-unknown/release/rig_wasm.wasm

./scripts/exports.sh 
npm ci
rollup -c 
cp src/generated/rig_wasm_bg.wasm dist/esm/generated 
cp src/generated/rig_wasm_bg.wasm.d.ts dist/esm/generated

Then all that was left to do was to run the script and then use npm publish.

Stuff we couldn't port

Of course, now that we've ported pretty much everything we could, there are a number of things that we couldn't port:

• the DynClientBuilder that basically allows you to create whatever client you want based on the provider/model input (as strings). Unfortunately, JS/TS doesn't understand lifetimes, so this was impossible.
• the PromptRequest<'a> struct that we use internally to control multi-turn usage also, unfortunately, can't be turned into JS because of lifetimes
• The vector database integrations, as they're not WASM friendly because they use libraries that are not compatible with WASM. Thankfully, basically every vector database that exists has a TypeScript SDK. Those integrations will have to be re-created on the TS side. Thankfully as SDKs are often (mostly!) well documented, it is not difficult to do so.