From 3e331661830f85bb3cb38ec5102b7707a25ea5a8 Mon Sep 17 00:00:00 2001 From: Greg Richardson Date: Sat, 1 Jun 2024 00:50:00 -0600 Subject: [PATCH] feat(ai-sandbox): initial poc --- package-lock.json | 94 +++++++++++++++- packages/ai-sandbox/babel.config.js | 4 + packages/ai-sandbox/jest.config.js | 11 ++ packages/ai-sandbox/package.json | 31 ++++++ packages/ai-sandbox/src/index.test.ts | 77 ++++++++++++++ packages/ai-sandbox/src/index.ts | 148 ++++++++++++++++++++++++++ packages/ai-sandbox/test/setup.ts | 10 ++ packages/ai-sandbox/tsconfig.json | 11 ++ 8 files changed, 382 insertions(+), 4 deletions(-) create mode 100644 packages/ai-sandbox/babel.config.js create mode 100644 packages/ai-sandbox/jest.config.js create mode 100644 packages/ai-sandbox/package.json create mode 100644 packages/ai-sandbox/src/index.test.ts create mode 100644 packages/ai-sandbox/src/index.ts create mode 100644 packages/ai-sandbox/test/setup.ts create mode 100644 packages/ai-sandbox/tsconfig.json diff --git a/package-lock.json b/package-lock.json index 30964ca48e5..37ed1efb049 100644 --- a/package-lock.json +++ b/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", diff --git a/packages/ai-sandbox/babel.config.js b/packages/ai-sandbox/babel.config.js new file mode 100644 index 00000000000..7a52c7fe867 --- /dev/null +++ b/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'], +} diff --git a/packages/ai-sandbox/jest.config.js b/packages/ai-sandbox/jest.config.js new file mode 100644 index 00000000000..11024e9a1b5 --- /dev/null +++ b/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-.*)/)'], +} diff --git a/packages/ai-sandbox/package.json b/packages/ai-sandbox/package.json new file mode 100644 index 00000000000..07c6daf6c79 --- /dev/null +++ b/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" + } +} diff --git a/packages/ai-sandbox/src/index.test.ts b/packages/ai-sandbox/src/index.test.ts new file mode 100644 index 00000000000..a8d04769ca5 --- /dev/null +++ b/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') + }) +}) diff --git a/packages/ai-sandbox/src/index.ts b/packages/ai-sandbox/src/index.ts new file mode 100644 index 00000000000..9ca8edca6bc --- /dev/null +++ b/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 + 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 = {} +): Promise { + 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)) +} diff --git a/packages/ai-sandbox/test/setup.ts b/packages/ai-sandbox/test/setup.ts new file mode 100644 index 00000000000..a11d7ac2258 --- /dev/null +++ b/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 }) +} diff --git a/packages/ai-sandbox/tsconfig.json b/packages/ai-sandbox/tsconfig.json new file mode 100644 index 00000000000..892c33438ec --- /dev/null +++ b/packages/ai-sandbox/tsconfig.json @@ -0,0 +1,11 @@ +{ + "extends": "tsconfig/react-library.json", + "include": [ + "." + ], + "exclude": [ + "dist", + "build", + "node_modules" + ] +} \ No newline at end of file