diff --git a/apps/host-selfhost/src/execution.ts b/apps/host-selfhost/src/execution.ts index 8dab577b9..e4523124f 100644 --- a/apps/host-selfhost/src/execution.ts +++ b/apps/host-selfhost/src/execution.ts @@ -68,9 +68,10 @@ export const SelfHostCodeExecutorProvider: Layer.Layer = L CodeExecutorProvider, () => { const { sandboxTimeoutMs } = loadConfig(); - return makeQuickJsExecutor( - sandboxTimeoutMs === undefined ? {} : { timeoutMs: sandboxTimeoutMs }, - ); + return makeQuickJsExecutor({ + ...(sandboxTimeoutMs === undefined ? {} : { timeoutMs: sandboxTimeoutMs }), + worker: true, + }); }, ); diff --git a/packages/kernel/runtime-quickjs/package.json b/packages/kernel/runtime-quickjs/package.json index 5ea156acc..60d693ee1 100644 --- a/packages/kernel/runtime-quickjs/package.json +++ b/packages/kernel/runtime-quickjs/package.json @@ -31,6 +31,7 @@ }, "scripts": { "build": "tsup && (tsc --declaration --emitDeclarationOnly --outDir dist --rootDir src || true)", + "pretest": "bun run build", "typecheck": "tsgo --noEmit", "test": "vitest run", "test:watch": "vitest", diff --git a/packages/kernel/runtime-quickjs/src/index.test.ts b/packages/kernel/runtime-quickjs/src/index.test.ts index a55f022ac..c4137148c 100644 --- a/packages/kernel/runtime-quickjs/src/index.test.ts +++ b/packages/kernel/runtime-quickjs/src/index.test.ts @@ -433,3 +433,72 @@ describe("quickjs executor", () => { }), ); }); + +describe("quickjs worker executor (worker_threads)", () => { + const workerExecutor = makeQuickJsExecutor({ worker: true, timeoutMs: 5_000 }); + + it.effect("runs plain code and returns results in a worker thread", () => + Effect.gen(function* () { + const result = yield* workerExecutor.execute(`return 10 + 20;`, makeTestInvoker({})); + expect(result.result).toBe(30); + expect(result.error).toBeUndefined(); + }), + ); + + it.effect("handles bidirectional tool invocations across the worker thread boundary", () => + Effect.gen(function* () { + const result = yield* workerExecutor.execute( + ` + const user = await tools.users.get({ id: 42 }); + return { name: user.name, doubled: user.id * 2 }; + `, + makeTestInvoker({ + "users.get": (args: any) => ({ id: args.id, name: "Alice" }), + }), + ); + + expect(result.error).toBeUndefined(); + expect(result.result).toEqual({ name: "Alice", doubled: 84 }); + }), + ); + + it.effect("handles tool errors across the worker thread boundary", () => + Effect.gen(function* () { + const result = yield* workerExecutor.execute( + ` + try { + await tools.db.query({ sql: "SELECT * FROM secrets" }); + return "unexpected"; + } catch (err) { + return err instanceof Error ? err.message : String(err); + } + `, + makeTestInvoker({ + "db.query": () => { + throw new Error("connection refused"); + }, + }), + ); + + expect(result.error).toBeUndefined(); + expect(result.result).toBe("Internal tool error"); + }), + ); + + it.effect("preempts a CPU-bound loop in the worker thread without hanging the main thread", () => + Effect.gen(function* () { + const timedWorkerExecutor = makeQuickJsExecutor({ worker: true, timeoutMs: 1_000 }); + const result = yield* timedWorkerExecutor.execute( + ` + const t0 = Date.now(); + while (true) {} + return "unreachable"; + `, + makeTestInvoker({}), + ); + + expect(result.result).toBeNull(); + expect(result.error).toContain("QuickJS execution timed out after 1000ms"); + }), + ); +}); diff --git a/packages/kernel/runtime-quickjs/src/index.ts b/packages/kernel/runtime-quickjs/src/index.ts index 90de8461b..a5dc9543c 100644 --- a/packages/kernel/runtime-quickjs/src/index.ts +++ b/packages/kernel/runtime-quickjs/src/index.ts @@ -1,3 +1,6 @@ +import { existsSync } from "node:fs"; +import { fileURLToPath } from "node:url"; + import { recoverExecutionBody, stripTypeScript, @@ -7,6 +10,7 @@ import { type SandboxToolInvoker, } from "@executor-js/codemode-core"; import * as Data from "effect/Data"; +import * as Deferred from "effect/Deferred"; import * as Effect from "effect/Effect"; import { executeSealedBundle as runSealedBundleWith, @@ -26,6 +30,7 @@ export type QuickJsExecutorOptions = { timeoutMs?: number; memoryLimitBytes?: number; maxStackSizeBytes?: number; + worker?: boolean; }; // Allow pre-loading a QuickJS module (e.g. with custom WASM bytes for compiled binaries) @@ -434,7 +439,7 @@ const drainAsync = async ( drainJobs(context, runtime, deadline, timeoutMs); }; -const evaluateInQuickJs = async ( +export const evaluateInQuickJs = async ( options: QuickJsExecutorOptions, code: string, toolInvoker: SandboxToolInvoker, @@ -565,9 +570,136 @@ const runInQuickJs = ( }), ); +export const resolveWorkerScriptPath = (): URL => { + const currentUrl = new URL(import.meta.url); + const devDistWorker = new URL("../dist/worker.js", currentUrl); + if (existsSync(fileURLToPath(devDistWorker))) { + return devDistWorker; + } + return new URL("./worker.js", currentUrl); +}; + +type NodeWorkerInstance = { + on(event: "message", listener: (value: unknown) => void): void; + on(event: "error", listener: (err: Error) => void): void; + on(event: "exit", listener: (exitCode: number) => void): void; + postMessage(value: unknown): void; + terminate(): Promise | void; +}; + +const runInQuickJsWorker = ( + options: QuickJsExecutorOptions, + code: string, + toolInvoker: SandboxToolInvoker, +): Effect.Effect => + Effect.gen(function* () { + const workerModule = yield* Effect.tryPromise({ + try: () => import("node:worker_threads"), + catch: (err) => + new QuickJsExecutionError({ + message: `worker_threads module is not available in this runtime: ${String(err)}`, + }), + }); + + const context = yield* Effect.context(); + const runPromise = Effect.runPromiseWith(context); + const resultDeferred = yield* Deferred.make(); + + const scriptPath = resolveWorkerScriptPath(); + const worker: NodeWorkerInstance = yield* Effect.try({ + try: () => + // oxlint-disable-next-line executor/no-double-cast -- boundary: dynamic worker_threads Worker instance bridged across heterogeneous runtime type environments + new workerModule.Worker(scriptPath, { + workerData: { code, options }, + }) as unknown as NodeWorkerInstance, + catch: (cause) => new QuickJsExecutionError({ message: String(cause) }), + }); + + worker.on("message", (msg: unknown) => { + if (!msg || typeof msg !== "object") return; + const payload = msg as { + type?: string; + id?: number; + path?: string; + args?: Record; + result?: ExecuteResult; + error?: string; + }; + if ( + payload.type === "tool_call" && + typeof payload.id === "number" && + typeof payload.path === "string" + ) { + void runPromise(toolInvoker.invoke({ path: payload.path, args: payload.args ?? {} })).then( + (value) => { + worker.postMessage({ type: "tool_result", id: payload.id, ok: true, value }); + }, + (cause) => { + worker.postMessage({ + type: "tool_result", + id: payload.id, + ok: false, + error: cause instanceof Error ? cause.message : String(cause), + }); + }, + ); + } else if (payload.type === "done" && payload.result) { + Effect.runSync(Deferred.complete(resultDeferred, Effect.succeed(payload.result))); + void worker.terminate(); + } else if (payload.type === "error" && typeof payload.error === "string") { + Effect.runSync( + Deferred.complete( + resultDeferred, + Effect.fail(new QuickJsExecutionError({ message: payload.error })), + ), + ); + void worker.terminate(); + } + }); + + worker.on("error", (err: Error) => { + Effect.runSync( + Deferred.complete( + resultDeferred, + Effect.fail(new QuickJsExecutionError({ message: err.message })), + ), + ); + void worker.terminate(); + }); + + worker.on("exit", (exitCode: number) => { + if (exitCode !== 0) { + Effect.runSync( + Deferred.complete( + resultDeferred, + Effect.fail( + new QuickJsExecutionError({ + message: `QuickJS worker thread exited unexpectedly with code ${exitCode}`, + }), + ), + ), + ); + } + }); + + return yield* Deferred.await(resultDeferred).pipe( + Effect.ensuring( + Effect.sync(() => { + void worker.terminate(); + }), + ), + ); + }).pipe( + Effect.withSpan("executor.code.exec.quickjs_worker", { + attributes: { "executor.runtime": "quickjs-worker" }, + }), + ); + export const makeQuickJsExecutor = ( options: QuickJsExecutorOptions = {}, ): CodeExecutor => ({ execute: (code: string, toolInvoker: SandboxToolInvoker) => - runInQuickJs(options, code, toolInvoker), + options.worker + ? runInQuickJsWorker(options, code, toolInvoker) + : runInQuickJs(options, code, toolInvoker), }); diff --git a/packages/kernel/runtime-quickjs/src/worker.ts b/packages/kernel/runtime-quickjs/src/worker.ts new file mode 100644 index 000000000..dd8aaf29b --- /dev/null +++ b/packages/kernel/runtime-quickjs/src/worker.ts @@ -0,0 +1,57 @@ +import { parentPort, workerData } from "node:worker_threads"; +import type { SandboxToolInvoker } from "@executor-js/codemode-core"; +import * as Effect from "effect/Effect"; +import { evaluateInQuickJs, type QuickJsExecutorOptions } from "./index"; + +if (parentPort) { + const port = parentPort; + const { code, options } = workerData as { + code: string; + options: QuickJsExecutorOptions; + }; + + let nextCallId = 1; + const pendingCalls = new Map< + number, + { + resolve: (value: unknown) => void; + reject: (error: unknown) => void; + } + >(); + + port.on("message", (msg) => { + if (msg && typeof msg === "object" && msg.type === "tool_result") { + const pending = pendingCalls.get(msg.id); + if (pending) { + pendingCalls.delete(msg.id); + if (msg.ok) { + pending.resolve(msg.value); + } else { + pending.reject(msg.error); + } + } + } + }); + + const workerToolInvoker: SandboxToolInvoker = { + invoke: ({ path, args }) => + Effect.promise( + () => + new Promise((resolve, reject) => { + const id = nextCallId++; + pendingCalls.set(id, { resolve, reject }); + port.postMessage({ type: "tool_call", id, path, args }); + }), + ), + }; + + const runPromise = (effect: Effect.Effect): Promise => Effect.runPromise(effect); + + evaluateInQuickJs(options, code, workerToolInvoker, runPromise) + .then((result) => { + port.postMessage({ type: "done", result }); + }) + .catch((cause) => { + port.postMessage({ type: "error", error: String(cause) }); + }); +} diff --git a/packages/kernel/runtime-quickjs/tsup.config.ts b/packages/kernel/runtime-quickjs/tsup.config.ts index 58569cf04..a9e977a7f 100644 --- a/packages/kernel/runtime-quickjs/tsup.config.ts +++ b/packages/kernel/runtime-quickjs/tsup.config.ts @@ -3,10 +3,11 @@ import { defineConfig } from "tsup"; export default defineConfig({ entry: { index: "src/index.ts", + worker: "src/worker.ts", }, format: ["esm"], dts: false, sourcemap: true, - clean: true, - external: [/^@executor-js\//, /^effect/, /^@effect\//, "quickjs-emscripten"], + noExternal: [/^@executor-js\//], + external: [/^effect/, /^@effect\//, "quickjs-emscripten"], });