jan-cli 0.27.0

#!/usr/bin/env node
"use strict";
/**
 * Jan warm Node worker: one long-lived worker_thread, jobs via postMessage.
 *
 * Creating a Worker per job was as slow as spawning node. Keeping a single
 * isolate warm and running each script in a fresh Module context gives the
 * speedup. process.chdir() is unavailable in workers — process.cwd is overridden.
 * If the worker dies, the listener respawns it.
 */
const net = require("net");
const fs = require("fs");
const path = require("path");
const {
  Worker,
  isMainThread,
  parentPort,
  workerData,
} = require("worker_threads");

if (!isMainThread) {
  persistentJobWorker();
} else {
  mainListener();
}

function persistentJobWorker() {
  parentPort.on("message", (job) => {
    const chunksOut = [];
    const chunksErr = [];
    const origStdoutWrite = process.stdout.write.bind(process.stdout);
    const origStderrWrite = process.stderr.write.bind(process.stderr);

    // Capture console/stdio into buffers for the parent (parent also pipes
    // worker stdout, but capturing here keeps ordering with thrown errors).
    process.stdout.write = (chunk, enc, cb) => {
      chunksOut.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(String(chunk)));
      if (typeof enc === "function") enc();
      else if (typeof cb === "function") cb();
      return true;
    };
    process.stderr.write = (chunk, enc, cb) => {
      chunksErr.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(String(chunk)));
      if (typeof enc === "function") enc();
      else if (typeof cb === "function") cb();
      return true;
    };

    let exitCode = 0;
    let errExtra = "";
    try {
      runOneJob(job);
    } catch (err) {
      exitCode = 1;
      errExtra = err && err.stack ? String(err.stack) : String(err);
    } finally {
      process.stdout.write = origStdoutWrite;
      process.stderr.write = origStderrWrite;
    }

    const stdout = Buffer.concat(chunksOut);
    const stderr = Buffer.concat([
      Buffer.concat(chunksErr),
      errExtra ? Buffer.from(errExtra) : Buffer.alloc(0),
    ]);
    parentPort.postMessage({
      ok: true,
      exit_code: exitCode,
      stdout_b64: stdout.toString("base64"),
      stderr_b64: stderr.toString("base64"),
    });
  });

  parentPort.postMessage({ type: "ready" });
}

function runOneJob(job) {
  const cwd = job.cwd || process.cwd();
  process.cwd = () => cwd;

  // Reset env to the job map (isolation between runs).
  for (const k of Object.keys(process.env)) {
    delete process.env[k];
  }
  if (job.env && typeof job.env === "object") {
    for (const [k, v] of Object.entries(job.env)) {
      if (v != null) process.env[k] = String(v);
    }
  }

  // Evict prior user entrypoints only — keep node_modules warm in require.cache.
  const cwdNorm = path.resolve(cwd);
  for (const id of Object.keys(require.cache)) {
    if (id === __filename) continue;
    if (id.includes(`${path.sep}node_modules${path.sep}`)) continue;
    if (id.startsWith(cwdNorm) || id.includes("<jan-inline>")) {
      delete require.cache[id];
    }
  }

  const vm = require("vm");
  const Module = require("module");
  const src = job.source || {};
  const argv = job.argv || [];

  let code;
  let filename;
  if (src.kind === "path") {
    filename = path.isAbsolute(src.value)
      ? src.value
      : path.resolve(cwd, src.value);
    code = fs.readFileSync(filename, "utf8");
    // Match `node script.js a b`.
    process.argv = [process.execPath, filename, ...argv];
  } else if (src.kind === "inline") {
    // `filename` stays a synthetic path for module resolution and stack
    // traces, but must not reach argv: `node -e code a b` puts no script
    // path there, so argv[1] is the first user arg.
    filename = path.join(cwd, "<jan-inline>");
    code = src.value || "";
    process.argv = [process.execPath, ...argv];
  } else {
    throw new Error("unknown source kind: " + src.kind);
  }

  const dirname = src.kind === "path" ? path.dirname(filename) : cwd;
  const m = new Module(filename);
  m.filename = filename;
  m.paths = Module._nodeModulePaths(dirname);
  const wrapper = Module.wrap(code);
  const compiled = vm.runInThisContext(wrapper, {
    filename,
    lineOffset: 0,
    displayErrors: true,
  });
  compiled(m.exports, makeJobRequire(dirname, Module), m, filename, dirname);
}

/** require() rooted at the script dir, without sticky global pollution. */
function makeJobRequire(dirname, Module) {
  const m = new Module(path.join(dirname, "<jan-require>"));
  m.filename = m.id;
  m.paths = Module._nodeModulePaths(dirname);
  return function jobRequire(id) {
    return m.require(id);
  };
}

function mainListener() {
  const sockPath = process.argv[2];
  if (!sockPath) {
    console.error("usage: node_worker.js <sock-path>");
    process.exit(2);
  }

  try {
    fs.unlinkSync(sockPath);
  } catch (_) {
    /* missing is fine */
  }

  /** @type {Worker | null} */
  let worker = null;
  /** @type {Promise<Worker> | null} */
  let starting = null;
  let busy = false;
  /** @type {{ job: any, resolve: Function, reject: Function }[]} */
  const queue = [];

  function spawnWorker() {
    if (starting) return starting;
    starting = new Promise((resolve, reject) => {
      const w = new Worker(__filename, {
        workerData: { role: "pool" },
        stdout: true,
        stderr: true,
      });
      // Drain pipes so they do not fill (job payload carries captured stdio).
      if (w.stdout) w.stdout.resume();
      if (w.stderr) w.stderr.resume();

      const onReady = (msg) => {
        if (msg && msg.type === "ready") {
          w.off("message", onReady);
          worker = w;
          starting = null;
          resolve(w);
          pump();
        }
      };
      w.on("message", onReady);
      w.on("error", (err) => {
        worker = null;
        starting = null;
        reject(err);
        failAll(err);
      });
      w.on("exit", (code) => {
        if (worker === w) worker = null;
        starting = null;
        if (busy) {
          failAll(new Error(`worker exited during job (code ${code})`));
        }
        // Respawn lazily on next job.
      });
    });
    return starting;
  }

  function failAll(err) {
    busy = false;
    while (queue.length) {
      queue.shift().reject(err);
    }
  }

  async function ensureWorker() {
    if (worker) return worker;
    return spawnWorker();
  }

  function pump() {
    if (busy || queue.length === 0) return;
    busy = true;
    const item = queue.shift();
    ensureWorker()
      .then((w) => {
        const onMsg = (msg) => {
          if (msg && msg.type === "ready") return;
          w.off("message", onMsg);
          busy = false;
          item.resolve(msg);
          pump();
        };
        w.on("message", onMsg);
        w.postMessage(item.job);
      })
      .catch((err) => {
        busy = false;
        item.reject(err);
        pump();
      });
  }

  function runJob(job) {
    return new Promise((resolve, reject) => {
      queue.push({ job, resolve, reject });
      pump();
    });
  }

  // Pre-warm the isolate so the first client job is fast.
  spawnWorker().catch(() => {});

  const server = net.createServer((socket) => {
    let buf = Buffer.alloc(0);
    socket.on("data", async (chunk) => {
      buf = Buffer.concat([buf, chunk]);
      const idx = buf.indexOf(0x0a);
      if (idx < 0) {
        if (buf.length > 16 * 1024 * 1024) {
          socket.end(JSON.stringify({ ok: false, error: "request too large" }) + "\n");
        }
        return;
      }
      const line = buf.slice(0, idx).toString("utf8");
      buf = buf.slice(idx + 1);
      let job;
      try {
        job = JSON.parse(line);
      } catch (e) {
        socket.end(JSON.stringify({ ok: false, error: "invalid json: " + e }) + "\n");
        return;
      }
      try {
        const resp = await runJob(job);
        socket.end(JSON.stringify(resp) + "\n");
      } catch (e) {
        // Hard failure — respawn for next time.
        try {
          if (worker) worker.terminate();
        } catch (_) {}
        worker = null;
        spawnWorker().catch(() => {});
        socket.end(
          JSON.stringify({ ok: false, error: String(e && e.stack ? e.stack : e) }) +
            "\n"
        );
      }
    });
  });

  server.listen(sockPath, () => {
    try {
      fs.chmodSync(sockPath, 0o600);
    } catch (_) {}
    process.stdout.write(`ready ${sockPath}\n`);
  });
}