WIP: NodeJS SDK working in Vercel Edge Runtime

lib/shared/bucketing-assembly-script/__tests__/protobufVariableHelper.ts

@@ -1,6 +1,13 @@
 import { SDKVariable, VariableType as VariableTypeStr } from '@devcycle/types'
 import { variableForUser_PB, VariableType } from './bucketingImportHelper'
 import { VariableForUserParams_PB, SDKVariable_PB } from '../protobuf/compiled'
+import {
+    BinaryWriter,
+    BinaryReader,
+    BinaryWriteOptions,
+    BinaryReadOptions,
+    MessageType,
+} from '@protobuf-ts/runtime'
 
 type SDKVariable_PB_Type = {
     _id: string
@@ -96,6 +103,31 @@ export type VariableForUserArgs = {
     variableType: VariableType
 }
 
+export function encodeProtobufMessage<T extends object>(
+    message: T,
+    messageType: MessageType<T>,
+): Uint8Array {
+    const writer = new BinaryWriter()
+    const writeOptions: BinaryWriteOptions = {
+        writeUnknownFields: true,
+        writerFactory: () => new BinaryWriter(),
+    }
+    messageType.internalBinaryWrite(message, writer, writeOptions)
+    return writer.finish()
+}
+
+export function decodeProtobufMessage<T extends object>(
+    buffer: Uint8Array,
+    messageType: MessageType<T>,
+): T {
+    const reader = new BinaryReader(buffer)
+    const readOptions: BinaryReadOptions = {
+        readUnknownField: true,
+        readerFactory: () => new BinaryReader(new Uint8Array()),
+    }
+    return messageType.internalBinaryRead(reader, 0, readOptions)
+}
+
 export const variableForUserPB = ({
     sdkKey,
     user,
@@ -138,14 +170,11 @@ export const variableForUserPB = ({
         },
         shouldTrackEvent: true,
     }
-    const err = VariableForUserParams_PB.verify(params)
-    if (err) throw new Error(err)
-
     const pbMsg = VariableForUserParams_PB.create(params)
-    const buffer = VariableForUserParams_PB.encode(pbMsg).finish()
+    const buffer = encodeProtobufMessage(pbMsg, VariableForUserParams_PB)
     const resultBuffer = variableForUser_PB(buffer)
 
     return !resultBuffer
         ? null
-        : pbSDKVariableToJS(SDKVariable_PB.decode(resultBuffer))
+        : pbSDKVariableToJS(decodeProtobufMessage(resultBuffer, SDKVariable_PB))
 }

lib/shared/bucketing-assembly-script/index.js

@@ -1,45 +1,105 @@
 const ProtobufTypes = require('./protobuf/compiled')
 const path = require('path')
 
-// This approach avoids using a named global variable that conflicts with itself
+const isEdgeRuntime = () => {
+    return process.env.NEXT_RUNTIME === 'edge'
+}
+
+const base64ToUint8Array = (base64) => {
+    const binaryString = atob(base64)
+    const bytes = new Uint8Array(binaryString.length)
+    for (let i = 0; i < binaryString.length; i++) {
+        bytes[i] = binaryString.charCodeAt(i)
+    }
+    return bytes
+}
+
+// Import WASM modules at the top level for Edge Runtime
+import releaseWasmModule from './build/bucketing-lib.release.wasm?module'
+import debugWasmModule from './build/bucketing-lib.debug.wasm?module'
+
 const instantiate = async (debug = false, imports = {}) => {
     const releaseStr = debug ? 'debug' : 'release'
-    const { instantiate } = require(`./build/bucketing-lib.${releaseStr}.js`)
-    // using path.resolve here to fix an issue caused by webpack in the example apps
-    // In normal operation __dirname is an absolute path, so path.resolve does nothing to it
-    // When running from the example app webpack bundle, __dirname becomes a relative path (relative to the repo root)
-    // We therefore need to 'path.resolve' it relative to the current working directory (which must be the repo root
-    // when running any Nx commands like serving the example app)
-    const url = new URL(
-        `file://${path.resolve(
-            __dirname,
-        )}/build/bucketing-lib.${releaseStr}.wasm`,
-    )
-
-    // asynchronously compile the webassembly source
-    const compiled = await (async () => {
+    let compiled
+    console.log(`isEdgeRuntime: ${isEdgeRuntime()}`)
+
+    if (isEdgeRuntime()) {
+        console.log('Using Edge Runtime WASM binary')
         try {
-            const source = await globalThis.fetch(url)
-            return await globalThis.WebAssembly.compileStreaming(source)
-        } catch {
-            let fs
+            // Use the appropriate WASM module based on debug flag
+            const wasmModule = debug ? debugWasmModule : releaseWasmModule
+            console.log('WASM module loaded:', {
+                type: typeof wasmModule,
+                isArray: Array.isArray(wasmModule),
+                isBuffer: Buffer.isBuffer(wasmModule),
+                length: wasmModule?.length,
+                keys: Object.keys(wasmModule)
+            })
+
+            // If the module is a string with numeric keys, convert it to a Uint8Array
+            let wasmBinary
+            if (typeof wasmModule === 'string' && Object.keys(wasmModule).every(key => !isNaN(key))) {
+                // Convert string to Uint8Array
+                wasmBinary = new Uint8Array(wasmModule.length)
+                for (let i = 0; i < wasmModule.length; i++) {
+                    wasmBinary[i] = wasmModule.charCodeAt(i)
+                }
+            } else {
+                // Try to use the default export if available
+                wasmBinary = base64ToUint8Array(wasmModule.default)
+            }
+
+            console.log('Converted to Uint8Array:', {
+                length: wasmBinary.length,
+                byteLength: wasmBinary.byteLength,
+                firstBytes: Array.from(wasmBinary.slice(0, 4))
+            })
+
+            // Instantiate with the buffer
+            const { instance } = await WebAssembly.instantiate(wasmBinary, imports)
+            return instance.exports
+        } catch (e) {
+            console.error('Failed to instantiate WASM in Edge Runtime:', e)
+            throw e
+        }
+    } else {
+        const { instantiate } = require(`./build/bucketing-lib.${releaseStr}.js`)
+        /**
+         * using path.resolve here to fix an issue caused by webpack in the example apps
+         * In normal operation __dirname is an absolute path, so path.resolve does nothing to it
+         * When running from the example app webpack bundle, __dirname becomes a relative path
+         * (relative to the repo root)
+         * We therefore need to 'path.resolve' it relative to the current working directory
+         * (which must be the repo root when running any Nx commands like serving the example app)
+         */
+        const url = new URL(
+            `file://${path.resolve(
+                __dirname,
+            )}/build/bucketing-lib.${releaseStr}.wasm`,
+        )
+
+        compiled = await (async () => {
             try {
-                // Use dynamic import with string concatenation to prevent webpack
-                const prefix = 'node:'
-                const suffix = 'fs/promises'
-                fs = require(prefix + suffix)
+                const source = await globalThis.fetch(url)
+                return await globalThis.WebAssembly.compileStreaming(source)
             } catch {
-                fs = require('fs/promises')
+                let fs
+                try {
+                    const prefix = 'node:'
+                    const suffix = 'fs/promises'
+                    fs = require(prefix + suffix)
+                } catch {
+                    fs = require('fs/promises')
+                }
+                return globalThis.WebAssembly.compile(await fs.readFile(url))
             }
-            return globalThis.WebAssembly.compile(await fs.readFile(url))
-        }
-    })()
+        })()
+    }
 
-    // call the instantiate function with the compiled source to execute the WASM and set up the bindings to native code
     return await instantiate(compiled, imports)
 }
 
-module.exports = {
+export {
     instantiate,
     ProtobufTypes,
 }

lib/shared/bucketing-assembly-script/package.json

@@ -12,9 +12,10 @@
         "README.md"
     ],
     "dependencies": {
-        "protobufjs": "^7.4.0"
+        "@protobuf-ts/runtime": "^2.9.3"
     },
     "devDependencies": {
-        "@devcycle/bucketing-test-data": "^1.27.0"
+        "@devcycle/bucketing-test-data": "^1.27.0",
+        "@protobuf-ts/plugin": "^2.9.6"
     }
 }

lib/shared/bucketing-assembly-script/project.json

@@ -77,16 +77,16 @@
             "options": {
                 "commands": [
                     {
-                        "command": "rm -rf compiled.js compiled.d.ts"
+                        "command": "rm -rf ./protobuf/compiled.js ./protobuf/compiled.d.ts"
                     },
                     {
-                        "command": "pbjs -t static-module --keep-case -w commonjs -o compiled.js ./*.proto"
+                        "command": "protoc --plugin=../../../node_modules/@protobuf-ts/plugin/bin/protoc-gen-ts --ts_out=./protobuf --ts_opt=eslint_disable,output_javascript --proto_path=./protobuf ./protobuf/*.proto"
                     },
-                    {
-                        "command": "pbts -o compiled.d.ts compiled.js"
+                    { 
+                        "command": "cd protobuf/ && mv variableForUserParams.d.ts compiled.d.ts && mv variableForUserParams.js compiled.js"
                     }
                 ],
-                "cwd": "lib/shared/bucketing-assembly-script/protobuf",
+                "cwd": "lib/shared/bucketing-assembly-script",
                 "parallel": false
             }
         },