suaveplan/packages/devtools/tooling/sandbox/src/worker-engine/worker-engine.ts

import { Worker } from "node:worker_threads";
import { recordDuration, recordExecution } from "../telemetry/telemetry.js";
import type {
    ExecutionContext,
    SandboxEngine,
    SandboxEngineConfig,
    SandboxModule,
    SandboxOptions,
    SandboxResult,
} from "../types/sandbox.types.js";

interface WorkerPayload {
    readonly code: string;
    readonly context: ExecutionContext;
    readonly modules: readonly SandboxModule[];
    readonly maxOutputSize: number;
}

interface WorkerReply {
    readonly success: boolean;
    readonly result: unknown;
    readonly error: string | undefined;
    readonly logs: readonly string[];
    readonly memoryUsed: number;
}

const WORKER_SOURCE = `
const { parentPort } = require('worker_threads');
parentPort.once('message', (payload) => {
    const logs = [];
    let totalLogSize = 0;
    const stringify = (v) => {
        if (typeof v === 'string') return v;
        if (v === null) return 'null';
        if (v === undefined) return 'undefined';
        try { return JSON.stringify(v); } catch { return String(v); }
    };
    const append = (line) => {
        totalLogSize += line.length;
        if (totalLogSize <= payload.maxOutputSize) logs.push(line);
    };
    const sandboxConsole = {
        log: (...a) => append(a.map(stringify).join(' ')),
        warn: (...a) => append('[warn] ' + a.map(stringify).join(' ')),
        error: (...a) => append('[error] ' + a.map(stringify).join(' ')),
    };
    const scope = { console: sandboxConsole, ...payload.context };
    for (const m of payload.modules) scope[m.name] = m.exports;
    const memBefore = process.memoryUsage().heapUsed;
    try {
        const keys = Object.keys(scope);
        const vals = keys.map(k => scope[k]);
        const factory = Reflect.construct(Function, [...keys, '"use strict";\\nreturn (' + payload.code + ')']);
        const result = factory(...vals);
        const memAfter = process.memoryUsage().heapUsed;
        parentPort.postMessage({
            success: true, result, error: undefined,
            logs, memoryUsed: Math.max(0, memAfter - memBefore),
        });
    } catch (err) {
        const memAfter = process.memoryUsage().heapUsed;
        parentPort.postMessage({
            success: false, result: undefined,
            error: err instanceof Error ? err.message : String(err),
            logs, memoryUsed: Math.max(0, memAfter - memBefore),
        });
    }
});
`;

/**
 * Real `node:worker_threads`-backed sandbox engine. Each `execute` call spawns
 * a fresh worker, posts the code + context + modules over the message channel,
 * receives the structured reply, and enforces the timeout via
 * `worker.terminate()`. The worker boundary is a true OS thread; intrinsics
 * are NOT shared with the host realm.
 */
export class WorkerEngine implements SandboxEngine {
    public readonly name = "worker";

    private readonly config: SandboxEngineConfig;
    private modules: SandboxModule[];
    private disposed = false;

    public constructor(config: SandboxEngineConfig) {
        this.config = config;
        this.modules = [...(config.modules ?? [])];
    }

    public async execute<T = unknown>(
        code: string,
        context: ExecutionContext = {},
        options: SandboxOptions = {}
    ): Promise<SandboxResult<T>> {
        if (this.disposed) {
            recordExecution("worker", "failure");
            return failure<T>("Engine has been disposed", 0);
        }
        const timeout = options.timeout ?? this.config.timeout;
        const maxOutputSize =
            options.maxOutputSize ?? this.config.maxOutputSize;
        const payload: WorkerPayload = {
            code,
            context,
            modules: this.modules,
            maxOutputSize,
        };
        const start = performance.now();
        const worker = new Worker(WORKER_SOURCE, { eval: true });
        try {
            const reply = await raceWorker<WorkerReply>(
                worker,
                payload,
                timeout
            );
            const duration = performance.now() - start;
            if (reply === "timeout") {
                recordExecution("worker", "failure");
                recordDuration("worker", "failure", duration);
                return failure<T>("Worker execution timed out", duration);
            }
            recordExecution("worker", reply.success ? "success" : "failure");
            recordDuration(
                "worker",
                reply.success ? "success" : "failure",
                duration
            );
            return {
                success: reply.success,
                result: reply.result as T,
                error: reply.error,
                duration,
                memoryUsed: reply.memoryUsed,
                logs: reply.logs,
            };
        } finally {
            await worker.terminate();
        }
    }

    public registerModule(mod: SandboxModule): void {
        this.modules = this.modules.filter((m) => m.name !== mod.name);
        this.modules.push(mod);
    }

    public unregisterModule(name: string): boolean {
        const before = this.modules.length;
        this.modules = this.modules.filter((m) => m.name !== name);
        return this.modules.length < before;
    }

    public reset(): void {
        this.modules = [...(this.config.modules ?? [])];
    }

    public destroy(): void {
        this.disposed = true;
        this.modules = [];
    }
}

function raceWorker<T>(
    worker: Worker,
    payload: WorkerPayload,
    timeout: number
): Promise<T | "timeout"> {
    return new Promise((resolve) => {
        const timer = setTimeout(() => {
            cleanup();
            resolve("timeout");
        }, timeout);
        const onMessage = (msg: unknown): void => {
            cleanup();
            resolve(msg as T);
        };
        const cleanup = (): void => {
            clearTimeout(timer);
            worker.off("message", onMessage);
        };
        worker.on("message", onMessage);
        // Worker eval boundary surfaces all sub-evaluation errors through
        // the structured reply payload — see WORKER_SOURCE's try/catch.
        worker.postMessage(payload);
    });
}

function failure<T>(error: string, duration: number): SandboxResult<T> {
    return {
        success: false,
        result: undefined,
        error,
        duration,
        memoryUsed: 0,
        logs: [],
    };
}