add support for kv host function

src/internal/kv.ts

@@ -0,0 +1,46 @@
+import { KVAction } from "@streamdal/protos/protos/shared/sp_shared";
+import { KVCommand } from "@streamdal/protos/protos/sp_command";
+import { KVInstruction } from "@streamdal/protos/protos/sp_kv";
+
+import { internal } from "./register.js";
+
+export const kvInstruction = (instruction: KVInstruction) => {
+  switch (instruction.action) {
+    case KVAction.KV_ACTION_DELETE_ALL: {
+      internal.kv.clear();
+      break;
+    }
+    case KVAction.KV_ACTION_DELETE: {
+      instruction.object?.key && internal.kv.delete(instruction.object.key);
+      break;
+    }
+    case KVAction.KV_ACTION_EXISTS: {
+      return instruction.object?.key && internal.kv.has(instruction.object.key);
+    }
+    case KVAction.KV_ACTION_GET: {
+      return instruction.object?.key && internal.kv.get(instruction.object.key);
+    }
+    case KVAction.KV_ACTION_UPDATE:
+    case KVAction.KV_ACTION_CREATE: {
+      instruction.object?.key &&
+        internal.kv.set(instruction.object.key, instruction.object.value);
+      break;
+    }
+  }
+};
+
+export const kvCommand = (command: KVCommand) => {
+  for (const instruction of command.instructions) {
+    kvInstruction(instruction);
+  }
+};
+
+export const kvExists = (
+  memory: WebAssembly.Memory,
+  pointer: number,
+  length: number
+): boolean => {
+  const bytes = new Uint8Array(memory.buffer, pointer, length);
+  const key = new TextDecoder().decode(bytes);
+  return internal.kv.has(key);
+};

src/internal/pipeline.ts

@@ -4,6 +4,7 @@ import { GetAttachCommandsByServiceResponse } from "@streamdal/protos/protos/sp_
 import { Pipeline, PipelineStep } from "@streamdal/protos/protos/sp_pipeline";
 
 import { Configs } from "../streamdal.js";
+import { kvCommand } from "./kv.js";
 import { audienceKey, internal, TailStatus } from "./register.js";
 import { instantiateWasm } from "./wasm.js";
 
@@ -40,6 +41,11 @@ export const initPipelines = async (configs: Configs) => {
 };
 
 export const processResponse = async (response: Command) => {
+  if (response.command.oneofKind === "kv") {
+    kvCommand(response.command.kv);
+    return;
+  }
+
   if (!response.audience) {
     response.command.oneofKind !== "keepAlive" &&
       console.debug("command response has no audience, ignoring");

src/internal/register.ts

@@ -28,6 +28,7 @@ export const internal = {
   registered: false,
   pipelineInitialized: false,
   pipelines: new Map<string, Map<string, InternalPipeline>>(),
+  kv: new Map<string, Uint8Array>(),
   audiences: new Map<
     string,
     { audience: Audience; tails: Map<string, TailStatus> }

src/internal/wasm.ts

@@ -7,6 +7,7 @@ import {
 // eslint-disable-next-line import/no-unresolved
 import { WASI } from "wasi";
 
+import { kvExists } from "./kv.js";
 import { MAX_PAYLOAD_SIZE } from "./process.js";
 import { internal } from "./register.js";
 
@@ -34,8 +35,25 @@ export const instantiateWasm = async (
   }
 
   const wasm = await WebAssembly.compile(wasmBytes);
-  const importObject = { wasi_snapshot_preview1: wasi.wasiImport };
-  const instantiated = await WebAssembly.instantiate(wasm, importObject);
+  const instantiated = await WebAssembly.instantiate(wasm, {
+    wasi_snapshot_preview1: wasi.wasiImport,
+    env: {
+      kvExists: (pointer: number, length: number): number => {
+        const result = instantiated.exports.memory
+          ? kvExists(
+              instantiated.exports.memory as WebAssembly.Memory,
+              pointer,
+              length
+            )
+          : false;
+        //
+        // Wasm function expects a bigint but our runtime type expects a number
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-expect-error
+        return BigInt(result ? 1 : 0);
+      },
+    },
+  });
   internal.wasmModules.set(wasmId, instantiated);
 };
 

src/sandbox/billing.ts

@@ -3,7 +3,7 @@ import { readFileSync } from "node:fs";
 import { Audience } from "@streamdal/protos/protos/sp_common";
 
 import { OperationType, Streamdal, StreamdalConfigs } from "../streamdal.js";
-import { QUIET, runPipeline } from "./index.js";
+import { runPipeline } from "./index.js";
 
 const serviceBillingConfig: StreamdalConfigs = {
   streamdalUrl: "localhost:8082",