1import { Worker } from "node:worker_threads";

2import { toCodemodeIdentifier } from "../identifier.ts";

3import type {

4 CodemodeCall,

5 CodemodeCallStatus,

6 CodemodeError,

7 CodemodeExecuteOptions,

8 CodemodeOutputItem,

9 CodemodeResult,

10 CodemodeSandboxOptions,

11 CodemodeStoreWrites,

12 CodemodeTool,

13} from "../types.ts";

14import { type CodemodeWasmModule, loadQuickJSWasm } from "../wasm.ts";

15import {

16 type HostToWorkerMessage,

17 isWorkerToHostMessage,

18 type WorkerData,

19 type WorkerToHostMessage,

20} from "./protocol.ts";

21

22const DEFAULT_TIMEOUT_MS = 300_000;

23const IDENTIFIER = /^[A-Za-z_$][A-Za-z0-9_$]*$/;

24const RESERVED_GLOBALS: ReadonlySet<string> = new Set([

25 "tools",

26 "ALL_TOOLS",

27 "console",

28 "text",

29 "image",

30 "exit",

31 "globalThis",

32 "store",

33 "load",

34]);

35

36function errorMessage(error: unknown): string {

37 return error instanceof Error ? error.message : String(error);

38}

39

40function serializeStore(store: Readonly<Record<string, unknown>> | undefined): Record<string, string> {

41 const serialized: Record<string, string> = {};

42 for (const [key, value] of Object.entries(store ?? {})) {

43 const json = JSON.stringify(value);

44 if (json !== undefined) serialized[key] = json;

45 }

46 return serialized;

47}

48

49function parseStoreWrites(json: string): CodemodeStoreWrites {

50 const writes: CodemodeStoreWrites = { set: {}, delete: [] };

51 for (const [key, value] of JSON.parse(json) as [string, string?][]) {

52 if (value === undefined) writes.delete.push(key);

53 else writes.set[key] = JSON.parse(value);

54 }

55 return writes;

56}

57

58function defaultWorkerUrl(): URL {

59 // `.ts` when running from source (tests, tsx), `.js` from the published dist.

60 return new URL(import.meta.url.endsWith(".ts") ? "./worker.ts" : "./worker.js", import.meta.url);

61}

62

63interface PendingCall {

64 record: CodemodeCall | undefined;

65 startedAt: number;

66 controller: AbortController;

67}

68

69interface ExecutionOptions {

70 code: string;

71 tools: ReadonlyMap<string, CodemodeTool>;

72 globals: ReadonlyMap<string, CodemodeTool>;

73 timeoutMs: number;

74 signal: AbortSignal | undefined;

75 memoryLimitBytes: number | undefined;

76 store: Record<string, string>;

77 wasm: Promise<CodemodeWasmModule>;

78 workerUrl: string | URL;

79}

80

81/**

82 * One script run in its own worker and QuickJS VM. A fresh worker per run keeps

83 * termination simple: a runaway script, including one that only spins the

84 * microtask queue, is killed with `terminate()` and cannot poison a later run.

85 */

86class Execution {

87 readonly promise: Promise<CodemodeResult>;

88 private resolveResult!: (result: CodemodeResult) => void;

89 private worker: Worker | undefined;

90 private readonly interrupt = new SharedArrayBuffer(4);

91 private readonly tools: ReadonlyMap<string, CodemodeTool>;

92 private readonly globals: ReadonlyMap<string, CodemodeTool>;

93 private readonly signal: AbortSignal | undefined;

94 private readonly timer: NodeJS.Timeout | undefined;

95 private readonly output: CodemodeOutputItem[] = [];

96 private readonly calls: CodemodeCall[] = [];

97 private readonly pending = new Map<number, PendingCall>();

98 private finished = false;

99

100 constructor(options: ExecutionOptions) {

101 this.promise = new Promise<CodemodeResult>((resolve) => {

102 this.resolveResult = resolve;

103 });

104 this.tools = options.tools;

105 this.globals = options.globals;

106 this.signal = options.signal;

107

108 if (Number.isFinite(options.timeoutMs)) {

109 this.timer = setTimeout(() => {

110 this.finish({ kind: "timeout", message: `Execution timed out after ${options.timeoutMs} ms` });

111 }, options.timeoutMs);

112 }

113

114 if (options.signal) {

115 if (options.signal.aborted) {

116 this.onAbort();

117 } else {

118 options.signal.addEventListener("abort", this.onAbort, { once: true });

119 }

120 }

121

122 options.wasm.then(

123 (wasm) => this.start(options, wasm),

124 (error: unknown) => {

125 this.finish({ kind: "sandbox", message: `Failed to load QuickJS: ${errorMessage(error)}` });

126 },

127 );

128 }

129

130 abort(message: string): Promise<CodemodeResult> {

131 this.finish({ kind: "aborted", message });

132 return this.promise;

133 }

134

135 private start(options: ExecutionOptions, wasm: CodemodeWasmModule): void {

136 if (this.finished) return;

137 const workerData: WorkerData = {

138 code: options.code,

139 tools: [...options.tools.values()].map((tool) => ({

140 name: tool.name,

141 jsName: toCodemodeIdentifier(tool.name),

142 description: tool.description ?? "",

143 })),

144 globals: [...options.globals.values()].map((global) => ({

145 name: global.name,

146 spread: global.spread === true,

147 })),

148 wasm,

149 memoryLimitBytes: options.memoryLimitBytes,

150 store: options.store,

151 interrupt: this.interrupt,

152 };

153 let worker: Worker;

154 try {

155 worker = new Worker(options.workerUrl, { workerData });

156 } catch (error) {

157 this.finish({ kind: "sandbox", message: `Failed to start worker: ${errorMessage(error)}` });

158 return;

159 }

160 this.worker = worker;

161 worker.on("message", (message: unknown) => this.handleMessage(message));

162 worker.on("error", (error: unknown) => {

163 this.finish({

164 kind: "sandbox",

165 name: error instanceof Error ? error.name : undefined,

166 message: errorMessage(error),

167 });

168 });

169 worker.on("exit", (code) => {

170 this.finish({ kind: "sandbox", message: `Worker exited with code ${code} before the script settled` });

171 });

172 }

173

174 private readonly onAbort = (): void => {

175 const reason: unknown = this.signal?.reason;

176 this.finish({ kind: "aborted", message: reason instanceof Error ? reason.message : "Execution aborted" });

177 };

178

179 private post(message: HostToWorkerMessage): void {

180 this.worker?.postMessage(message);

181 }

182

183 private handleMessage(message: unknown): void {

184 if (this.finished || !isWorkerToHostMessage(message)) return;

185 switch (message.type) {

186 case "output":

187 this.output.push(message.item);

188 break;

189 case "call":

190 void this.handleCall(message);

191 break;

192 case "done":

193 this.handleDone(message);

194 break;

195 case "crash":

196 this.finish({ kind: "sandbox", message: message.message });

197 break;

198 }

199 }

200

201 private handleDone(message: Extract<WorkerToHostMessage, { type: "done" }>): void {

202 if (!message.ok) {

203 const parsed = JSON.parse(message.error) as Omit<CodemodeError, "kind">;

204 this.finish({ kind: "script", ...parsed });

205 return;

206 }

207 this.finish(undefined, message.value === undefined ? undefined : JSON.parse(message.value), message.writes);

208 }

209

210 private async handleCall(message: Extract<WorkerToHostMessage, { type: "call" }>): Promise<void> {

211 const { id, name } = message;

212 const isTool = message.target === "tool";

213 const record: CodemodeCall | undefined = isTool ? { name, status: "cancelled", durationMs: 0 } : undefined;

214 if (record) this.calls.push(record);

215 const pending: PendingCall = { record, startedAt: performance.now(), controller: new AbortController() };

216 this.pending.set(id, pending);

217

218 let status: CodemodeCallStatus;

219 let reply: HostToWorkerMessage;

220 try {

221 const tool = (isTool ? this.tools : this.globals).get(name);

222 if (!tool) throw new Error(`Unknown ${isTool ? "tool" : "global"} "${name}"`);

223 const args: unknown = message.args === undefined ? undefined : JSON.parse(message.args);

224 const value = await tool.execute(args, { signal: pending.controller.signal });

225 reply = { type: "result", id, ok: true, payload: value === undefined ? undefined : JSON.stringify(value) };

226 status = "ok";

227 } catch (error) {

228 reply = { type: "result", id, ok: false, payload: errorMessage(error) };

229 status = "error";

230 }

231

232 // Already cancelled by finish(): the record keeps "cancelled" and the

233 // worker is gone or going.

234 if (!this.pending.delete(id)) return;

235 if (record) {

236 record.status = status;

237 record.durationMs = performance.now() - pending.startedAt;

238 }

239 this.post(reply);

240 }

241

242 private finish(error: CodemodeError | undefined, value?: unknown, writes?: string): void {

243 if (this.finished) return;

244 this.finished = true;

245 clearTimeout(this.timer);

246 this.signal?.removeEventListener("abort", this.onAbort);

247

248 const now = performance.now();

249 for (const pending of this.pending.values()) {

250 if (pending.record) pending.record.durationMs = now - pending.startedAt;

251 pending.controller.abort();

252 }

253 this.pending.clear();

254

255 const result: CodemodeResult = error

256 ? { ok: false, error, output: this.output, calls: this.calls }

257 : {

258 ok: true,

259 value,

260 output: this.output,

261 calls: this.calls,

262 storeWrites: writes === undefined ? { set: {}, delete: [] } : parseStoreWrites(writes),

263 };

264 if (!this.worker) {

265 this.resolveResult(result);

266 return;

267 }

268 Atomics.store(new Int32Array(this.interrupt), 0, 1);

269 this.worker

270 .terminate()

271 .catch(() => undefined)

272 .then(() => this.resolveResult(result));

273 }

274}

275

276/**

277 * Runs JavaScript in a QuickJS VM (a separate wasm instance) inside a worker

278 * thread. The script sees `tools.<name>(args)` for every registered tool, `ALL_TOOLS`,

279 * the output helpers `text`, `image`, `exit`, and `console.*`, `store`/`load`, and the

280 * configured globals; nothing else (no timers, `fetch`, `process`, `require`, modules).

281 *

282 * Each `execute()` gets its own worker and VM; the sandbox only holds the tool

283 * table and defaults. `close()` aborts in-flight executions.

284 */

285export class CodemodeSandbox {

286 private readonly toolsByName = new Map<string, CodemodeTool>();

287 private readonly globalsByName = new Map<string, CodemodeTool>();

288 private readonly timeoutMs: number;

289 private readonly memoryLimitBytes: number | undefined;

290 private readonly wasm: CodemodeWasmModule | Promise<CodemodeWasmModule> | undefined;

291 private readonly workerUrl: string | URL;

292 private readonly running = new Set<Execution>();

293 private closed = false;

294

295 constructor(options: CodemodeSandboxOptions = {}) {

296 this.timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;

297 this.memoryLimitBytes = options.memoryLimitBytes;

298 this.wasm = options.wasm;

299 this.workerUrl = options.workerUrl ?? defaultWorkerUrl();

300 for (const tool of options.tools ?? []) this.registerTool(tool);

301 const namespaces = new Set<string>();

302 for (const global of options.globals ?? []) {

303 const parts = global.name.split(".");

304 if (parts.length > 2 || !parts.every((part) => IDENTIFIER.test(part)) || RESERVED_GLOBALS.has(parts[0])) {

305 throw new Error(`Invalid global name "${global.name}"`);

306 }

307 if (this.globalsByName.has(global.name)) throw new Error(`Global "${global.name}" is already registered`);

308 if (parts.length === 2) namespaces.add(parts[0]);

309 this.globalsByName.set(global.name, global);

310 }

311 for (const name of namespaces) {

312 if (this.globalsByName.has(name)) throw new Error(`Global "${name}" conflicts with the namespace "${name}"`);

313 }

314 }

315

316 /** Throws if a tool with the same name is already registered. */

317 registerTool(tool: CodemodeTool): void {

318 if (this.toolsByName.has(tool.name)) throw new Error(`Tool "${tool.name}" is already registered`);

319 this.toolsByName.set(tool.name, tool);

320 }

321

322 unregisterTool(name: string): boolean {

323 return this.toolsByName.delete(name);

324 }

325

326 get tools(): CodemodeTool[] {

327 return [...this.toolsByName.values()];

328 }

329

330 get globals(): CodemodeTool[] {

331 return [...this.globalsByName.values()];

332 }

333

334 /**

335 * `code` is an async function body: `return` and top-level `await` work.

336 * Never rejects for script failures; those come back as `{ ok: false }`.

337 * The script can use `store(key, value)` and `load(key)` on `options.store`.

338 */

339 execute(code: string, options: CodemodeExecuteOptions = {}): Promise<CodemodeResult> {

340 if (this.closed) return Promise.reject(new Error("Sandbox is closed"));

341 const execution = new Execution({

342 code,

343 tools: new Map(this.toolsByName),

344 globals: this.globalsByName,

345 timeoutMs: options.timeoutMs ?? this.timeoutMs,

346 signal: options.signal,

347 memoryLimitBytes: this.memoryLimitBytes,

348 store: serializeStore(options.store),

349 wasm: this.wasm === undefined ? loadQuickJSWasm() : Promise.resolve(this.wasm),

350 workerUrl: this.workerUrl,

351 });

352 this.running.add(execution);

353 return execution.promise.finally(() => this.running.delete(execution));

354 }

355

356 /** Aborts in-flight executions (they resolve with `kind: "aborted"`) and rejects new ones. */

357 async close(): Promise<void> {

358 this.closed = true;

359 await Promise.all([...this.running].map((execution) => execution.abort("Sandbox closed")));

360 }

361}