feat(ai-sandbox): initial poc

package-lock.json

@@ -9154,6 +9154,43 @@
         "node": ">=8"
       }
     },
+    "node_modules/@jitl/quickjs-ffi-types": {
+      "version": "0.29.1",
+      "resolved": "https://registry.npmjs.org/@jitl/quickjs-ffi-types/-/quickjs-ffi-types-0.29.1.tgz",
+      "integrity": "sha512-dVt//PvFmwNJIwfyfPx7vR6cg7mLnAq6/THbhIeZ3NqXb66qXQATnW2gNqAnOcTh1D/f6HyrHFwRf3GdKUN3nw=="
+    },
+    "node_modules/@jitl/quickjs-wasmfile-debug-asyncify": {
+      "version": "0.29.1",
+      "resolved": "https://registry.npmjs.org/@jitl/quickjs-wasmfile-debug-asyncify/-/quickjs-wasmfile-debug-asyncify-0.29.1.tgz",
+      "integrity": "sha512-jhxRDq+9ZWT8ddRt6JLcjLLF168EvFrOceiaoy+GuYYXJkWVrWk5p2Ray2GHzh+DL7byszRKEUAKqA9B86Rd9g==",
+      "dependencies": {
+        "@jitl/quickjs-ffi-types": "0.29.1"
+      }
+    },
+    "node_modules/@jitl/quickjs-wasmfile-debug-sync": {
+      "version": "0.29.1",
+      "resolved": "https://registry.npmjs.org/@jitl/quickjs-wasmfile-debug-sync/-/quickjs-wasmfile-debug-sync-0.29.1.tgz",
+      "integrity": "sha512-qn6cxqg2fvobiJtQ/xHVSZH55djraIQ7g+9XGyJjf+ofsn1taMG8T5TIR2iZL7TMSdCRZkGiJeNOp9UaFCTN2w==",
+      "dependencies": {
+        "@jitl/quickjs-ffi-types": "0.29.1"
+      }
+    },
+    "node_modules/@jitl/quickjs-wasmfile-release-asyncify": {
+      "version": "0.29.1",
+      "resolved": "https://registry.npmjs.org/@jitl/quickjs-wasmfile-release-asyncify/-/quickjs-wasmfile-release-asyncify-0.29.1.tgz",
+      "integrity": "sha512-1PD+TEw5raGk5s7G1N89gaT81J5A1N6PT79ixbm70gaU0xdUJyAUJoiLiS3YwLjcbpiiM0f7Vnf85M3Li914aQ==",
+      "dependencies": {
+        "@jitl/quickjs-ffi-types": "0.29.1"
+      }
+    },
+    "node_modules/@jitl/quickjs-wasmfile-release-sync": {
+      "version": "0.29.1",
+      "resolved": "https://registry.npmjs.org/@jitl/quickjs-wasmfile-release-sync/-/quickjs-wasmfile-release-sync-0.29.1.tgz",
+      "integrity": "sha512-iVeMmZXIDg3gD0elfj5FscwrqGLmcglvpvWDIzs5tmzJ4AKiDAHXunGwd8X4gifeW6S+f0j681M8+tZuUIqnJA==",
+      "dependencies": {
+        "@jitl/quickjs-ffi-types": "0.29.1"
+      }
+    },
     "node_modules/@jridgewell/gen-mapping": {
       "version": "0.3.3",
       "license": "MIT",
@@ -21422,6 +21459,10 @@
       "resolved": "packages/ai-commands",
       "link": true
     },
+    "node_modules/ai-sandbox": {
+      "resolved": "packages/ai-sandbox",
+      "link": true
+    },
     "node_modules/ajv": {
       "version": "6.12.6",
       "license": "MIT",
@@ -24365,10 +24406,9 @@
       }
     },
     "node_modules/core-js": {
-      "version": "3.35.0",
-      "resolved": "https://registry.npmjs.org/core-js/-/core-js-3.35.0.tgz",
-      "integrity": "sha512-ntakECeqg81KqMueeGJ79Q5ZgQNR+6eaE8sxGCx62zMbAIj65q+uYvatToew3m6eAGdU4gNZwpZ34NMe4GYswg==",
-      "dev": true,
+      "version": "3.37.1",
+      "resolved": "https://registry.npmjs.org/core-js/-/core-js-3.37.1.tgz",
+      "integrity": "sha512-Xn6qmxrQZyB0FFY8E3bgRXei3lWDJHhvI+u0q9TKIYM49G8pAr0FgnnrFRAmsbptZL1yxRADVXn+x5AGsbBfyw==",
       "hasInstallScript": true,
       "funding": {
         "type": "opencollective",
@@ -42436,6 +42476,29 @@
       "integrity": "sha512-kJt5qhMxoszgU/62PLP1CJytzd2NKetjSRnyuj31fDd3Rlcz3fzlFdFLD1SItunPwyqEOkca6GbV612BWfaBag==",
       "dev": true
     },
+    "node_modules/quickjs-emscripten": {
+      "version": "0.29.1",
+      "resolved": "https://registry.npmjs.org/quickjs-emscripten/-/quickjs-emscripten-0.29.1.tgz",
+      "integrity": "sha512-9WX9cbbPW2/Wg/vedfAV6lL1kJB68xvE7kuILM24dkcIB/NCrIscZ94OdTGi9/77jYEJuajkzoQnHVn2BLzdaQ==",
+      "dependencies": {
+        "@jitl/quickjs-wasmfile-debug-asyncify": "0.29.1",
+        "@jitl/quickjs-wasmfile-debug-sync": "0.29.1",
+        "@jitl/quickjs-wasmfile-release-asyncify": "0.29.1",
+        "@jitl/quickjs-wasmfile-release-sync": "0.29.1",
+        "quickjs-emscripten-core": "0.29.1"
+      },
+      "engines": {
+        "node": ">=16.0.0"
+      }
+    },
+    "node_modules/quickjs-emscripten-core": {
+      "version": "0.29.1",
+      "resolved": "https://registry.npmjs.org/quickjs-emscripten-core/-/quickjs-emscripten-core-0.29.1.tgz",
+      "integrity": "sha512-urFIWrPgKfzLFe/IwE8px3Oznb7UGsXpMjHLe15PTPbWongXmy6KAyxRpYFW8K/dDELilg0H/voysHbWHKj6uA==",
+      "dependencies": {
+        "@jitl/quickjs-ffi-types": "0.29.1"
+      }
+    },
     "node_modules/raf-schd": {
       "version": "4.0.3",
       "license": "MIT"
@@ -52132,6 +52195,29 @@
         "sql-formatter": "bin/sql-formatter-cli.cjs"
       }
     },
+    "packages/ai-sandbox": {
+      "version": "0.0.0",
+      "license": "MIT",
+      "dependencies": {
+        "core-js": "^3.37.1",
+        "openai": "^4.26.1",
+        "quickjs-emscripten": "^0.29.1"
+      },
+      "devDependencies": {
+        "@babel/core": "^7.23.6",
+        "@babel/preset-env": "^7.23.6",
+        "@jest/globals": "^29.7.0",
+        "@types/common-tags": "^1.8.4",
+        "babel-jest": "^29.7.0",
+        "babel-plugin-transform-import-meta": "^2.2.1",
+        "common-tags": "^1.8.2",
+        "dotenv": "^16.3.1",
+        "jest": "^29.7.0",
+        "ts-jest": "^29.1.1",
+        "tsconfig": "*",
+        "typescript": "^5.4.3"
+      }
+    },
     "packages/api-types": {
       "version": "0.0.0",
       "license": "MIT",

packages/ai-sandbox/babel.config.js

@@ -0,0 +1,4 @@
+module.exports = {
+  presets: [['@babel/preset-env', { targets: { node: 'current' } }]],
+  plugins: ['babel-plugin-transform-import-meta'],
+}

packages/ai-sandbox/jest.config.js

@@ -0,0 +1,11 @@
+/** @type {import('ts-jest').JestConfigWithTsJest} */
+module.exports = {
+  preset: 'ts-jest',
+  testEnvironment: 'node',
+  transform: {
+    '^.+\\.ts?$': 'ts-jest',
+    '^.+\\.(js|jsx)$': 'babel-jest',
+  },
+  setupFiles: ['./test/setup.ts'],
+  transformIgnorePatterns: ['node_modules/(?!(@jitl/quickjs-wasmfile-.*)/)'],
+}

packages/ai-sandbox/package.json

@@ -0,0 +1,31 @@
+{
+  "name": "ai-sandbox",
+  "version": "0.0.0",
+  "main": "./src/index.ts",
+  "types": "./src/index.ts",
+  "license": "MIT",
+  "sideEffects": false,
+  "scripts": {
+    "typecheck": "tsc --noEmit",
+    "test": "jest"
+  },
+  "dependencies": {
+    "core-js": "^3.37.1",
+    "openai": "^4.26.1",
+    "quickjs-emscripten": "^0.29.1"
+  },
+  "devDependencies": {
+    "@babel/core": "^7.23.6",
+    "@babel/preset-env": "^7.23.6",
+    "@jest/globals": "^29.7.0",
+    "@types/common-tags": "^1.8.4",
+    "babel-jest": "^29.7.0",
+    "babel-plugin-transform-import-meta": "^2.2.1",
+    "common-tags": "^1.8.2",
+    "dotenv": "^16.3.1",
+    "jest": "^29.7.0",
+    "ts-jest": "^29.1.1",
+    "tsconfig": "*",
+    "typescript": "^5.4.3"
+  }
+}

packages/ai-sandbox/src/index.test.ts

@@ -0,0 +1,77 @@
+import { describe, expect, test } from '@jest/globals'
+import { executeJS } from '.'
+import { codeBlock } from 'common-tags'
+
+describe('execute js', () => {
+  test('exports', async () => {
+    const { exports, error } = await executeJS(codeBlock`
+      export const output = 'test';
+    `)
+
+    if (error) {
+      throw error
+    }
+
+    expect(exports.output).toBe('test')
+  })
+
+  test('imports', async () => {
+    const { exports, error } = await executeJS(
+      codeBlock`
+        import { value } from 'my-lib';
+        export const output = value;
+      `,
+      {
+        'my-lib': codeBlock`
+          export const value = 'test';
+        `,
+      }
+    )
+
+    if (error) {
+      throw error
+    }
+
+    expect(exports.output).toBe('test')
+  })
+
+  test('top-level await', async () => {
+    const { exports, error } = await executeJS(
+      codeBlock`
+        async function doStuff() {
+          await new Promise(r => setTimeout(r, 0));
+          return 'test';
+        }
+
+        export const output = await doStuff();
+      `
+    )
+
+    if (error) {
+      throw error
+    }
+
+    expect(exports.output).toBe('test')
+  })
+
+  test('top-level throw', async () => {
+    const { exports, error } = await executeJS(
+      codeBlock`
+        throw new Error('test')
+      `
+    )
+
+    expect(error?.message).toBe('test')
+  })
+
+  test('top-level async throw', async () => {
+    const { exports, error } = await executeJS(
+      codeBlock`
+        await Promise.resolve()
+        throw new Error('test')
+      `
+    )
+
+    expect(error?.message).toBe('test')
+  })
+})

packages/ai-sandbox/src/index.ts

@@ -0,0 +1,148 @@
+import 'core-js/actual/symbol/async-dispose'
+import 'core-js/actual/symbol/dispose'
+
+import { JSPromiseState, getQuickJS, shouldInterruptAfterDeadline } from 'quickjs-emscripten'
+
+export type SuccessResult = {
+  exports: Record<string, any>
+  error: undefined
+}
+
+export type FailureResult = {
+  exports: undefined
+  error: Error
+}
+
+export type ExecutionResult = SuccessResult | FailureResult
+
+/**
+ * Executes JavaScript source code in a secure sandbox via an
+ * embedded runtime (powered by QuickJS running in WebAssembly).
+ *
+ * The VM's runtime is completely isolated from the host's by default.
+ * The only interface between the VM and the host are the ones you create
+ * through modules and functions.
+ *
+ * The `source` is expected to be an ES module (with imports and exports).
+ * Async modules are supported.
+ *
+ * All exports will be returned after execution via the `exports` property.
+ * Any error that is thrown will be returned via the `error` property.
+ */
+export async function executeJS(
+  source: string,
+  modules: Record<string, string> = {}
+): Promise<ExecutionResult> {
+  const QuickJS = await getQuickJS()
+
+  using runtime = QuickJS.newRuntime()
+
+  runtime.setMemoryLimit(1024 * 640)
+  runtime.setMaxStackSize(1024 * 320)
+  runtime.setInterruptHandler(shouldInterruptAfterDeadline(Date.now() + 1000))
+  runtime.setModuleLoader((moduleName) => modules[moduleName])
+
+  using context = runtime.newContext()
+
+  using setTimeoutHandle = context.newFunction('setTimeout', (fnHandle, timeoutHandle) => {
+    // Make a copy because otherwise fnHandle does not live long enough to call after the timeout
+    const fnHandleCopy = fnHandle.dup()
+    const timeout = context.dump(timeoutHandle)
+
+    // Node.js implementations of `setTimeout` return a `Timer` instead of a `number`
+    // Wrapping the result in `Number()` implicitly calls `[Symbol.toPrimitive]()` on the `Timer`
+    const timeoutId = Number(
+      setTimeout(() => {
+        // callFunction() will call the function handle within the VM
+        // we pass context.undefined because we need to pass something for the "this" argument
+        context.callFunction(fnHandleCopy, context.undefined)
+        fnHandleCopy.dispose()
+      }, timeout)
+    )
+
+    return context.newNumber(timeoutId)
+  })
+
+  context.setProp(context.global, 'setTimeout', setTimeoutHandle)
+
+  const result = context.evalCode(source, 'index.js', { type: 'module' })
+
+  // This is normally a syntax error, but can also be any top-level error thrown in the module
+  if (result.error) {
+    using errorHandle = result.error
+    const error = context.dump(errorHandle)
+
+    return {
+      exports: undefined,
+      error: deserializeError(error),
+    }
+  }
+
+  using valueHandle = result.value
+  let promiseState: JSPromiseState
+
+  // Continuously execute pending jobs until the module has resolved (if async)
+  do {
+    // If we're proxying async functions through the host runtime (eg. fetch),
+    // we need to defer execution until those functions have completed.
+    // Choosing `setTimeout` over `queueMicrotask` since it runs after I/O.
+    await new Promise((r) => setTimeout(r, 0))
+
+    runtime.executePendingJobs()
+    promiseState = context.getPromiseState(valueHandle)
+  } while (promiseState.type === 'pending')
+
+  // Promise state will be 'fulfilled' even for non-promises, so
+  // this covers both sync and async modules
+  if (promiseState.type === 'fulfilled') {
+    const promiseValueHandle = promiseState.value
+    const moduleExports = context.dump(promiseValueHandle)
+
+    // These handles are equal when the module is sync, so only
+    // dispose when it's async (or else it will get disposed twice)
+    if (promiseValueHandle !== valueHandle) {
+      promiseValueHandle.dispose()
+    }
+
+    return {
+      exports: moduleExports,
+      error: undefined,
+    }
+  } else if (promiseState.type === 'rejected') {
+    using promiseErrorHandle = promiseState.error
+    const error = context.dump(promiseErrorHandle)
+
+    return {
+      exports: undefined,
+      error: deserializeError(error),
+    }
+  } else {
+    throw new Error('Something went wrong, please report this error')
+  }
+}
+
+/**
+ * Deserializes an error from the VM.
+ */
+function deserializeError(error: unknown): Error {
+  if (
+    error !== undefined &&
+    error !== null &&
+    typeof error === 'object' &&
+    'name' in error &&
+    'message' in error &&
+    'stack' in error &&
+    typeof error.name === 'string' &&
+    typeof error.message === 'string' &&
+    typeof error.stack === 'string'
+  ) {
+    const newError = new Error(error.message)
+    Object.assign(newError, error)
+    newError.stack = `${error.message}\n${error.stack}`
+    return newError
+  } else if (typeof error === 'string') {
+    return new Error(error)
+  }
+
+  return new Error(JSON.stringify(error))
+}

packages/ai-sandbox/test/setup.ts

@@ -0,0 +1,10 @@
+import { config } from 'dotenv'
+import { statSync } from 'fs'
+
+if (!process.env.CI) {
+  // Use keys from studio .env.local for local tests
+  const envPath = '../../apps/studio/.env.local'
+
+  statSync(envPath)
+  config({ path: envPath })
+}

packages/ai-sandbox/tsconfig.json

@@ -0,0 +1,11 @@
+{
+  "extends": "tsconfig/react-library.json",
+  "include": [
+    "."
+  ],
+  "exclude": [
+    "dist",
+    "build",
+    "node_modules"
+  ]
+}