orbit/dev/wdshot.mjs

// One-shot headless-Firefox screenshot via geckodriver (W3C WebDriver).
//
//   node dev/wdshot.mjs <url> <out.png> [setupExpression] [timeoutMs=20000] [settleMs=0]
//
// The setup expression (optional) runs right before the capture — it may
// return a Promise; the runner waits for it to settle (so scene state /
// animations can be driven first). WebDriver's /screenshot endpoint is
// used (not canvas.toDataURL — a WebGL canvas without
// preserveDrawingBuffer captures black).
//
// settleMs (optional): a runner-side settle wait AFTER navigation, before
// the setup script runs. Pages with async module boot (top-level awaits —
// e.g. dev/popup-shot.mjs waiting for the web fonts) need a beat before
// their globals exist.
import { spawn, execFileSync } from 'node:child_process';
import { existsSync, readdirSync, writeFileSync } from 'node:fs';

const [url, out, setup, timeoutMs = '20000', settleMs = '0'] = process.argv.slice(2);
if (!url || !out) {
  console.error('usage: node dev/wdshot.mjs <url> <out.png> [setupExpression] [timeoutMs] [settleMs]');
  process.exit(2);
}

const sleep = (ms) => new Promise((r) => setTimeout(r, ms));
let geckodriver;
try {
  geckodriver = execFileSync('sh', ['-c', 'command -v geckodriver'], { encoding: 'utf8' }).trim();
} catch {
  console.error('need geckodriver on PATH');
  process.exit(2);
}
const firefoxCandidates = [];
try {
  const ff = execFileSync('sh', ['-c', 'command -v firefox'], { encoding: 'utf8' }).trim();
  if (ff) firefoxCandidates.push(ff);
} catch { /* not on PATH */ }
try {
  const snapRev = readdirSync('/snap/firefox', { withFileTypes: true }).map((d) => d.name);
  for (const rev of [null, ...snapRev]) {
    const base = rev ? `/snap/firefox/${rev}` : '/snap/firefox/current';
    const p = `${base}/usr/lib/firefox/firefox`;
    if (existsSync(p)) firefoxCandidates.push(p);
  }
} catch { /* non-snap install */ }
if (firefoxCandidates.length === 0) {
  console.error('no firefox binary found');
  process.exit(2);
}
const PORT = 4444 + (process.pid % 100);
const driver = spawn(geckodriver, ['--port', String(PORT), '--log', 'fatal'], { stdio: 'ignore' });
const base = `http://127.0.0.1:${PORT}`;

const fail = (msg, code = 1) => {
  try { driver.kill('SIGKILL'); } catch { /* dead */ }
  console.error(msg);
  process.exit(code);
};
const wd = async (method, path, body) => {
  const res = await fetch(base + path, {
    method,
    headers: { 'Content-Type': 'application/json' },
    body: body === undefined ? undefined : JSON.stringify(body),
  });
  const text = await res.text();
  let data = null;
  try { data = JSON.parse(text); } catch { /* empty */ }
  if (!res.ok) throw new Error(`WebDriver ${method} ${path} → ${res.status}: ${text.slice(0, 300)}`);
  return data;
};

try {
  let up = false;
  for (let i = 0; i < 60; i++) {
    try {
      const v = await (await fetch(`${base}/status`)).json();
      if (v?.value?.ready || v?.status === 'ok' || v?.value) { up = true; break; }
    } catch { /* not up */ }
    await sleep(250);
  }
  if (!up) fail('geckodriver never came up');

  let sess;
  let lastErr = null;
  for (const ff of firefoxCandidates) {
    try {
      sess = await wd('POST', '/session', {
        capabilities: {
          alwaysMatch: {
            browserName: 'firefox',
            'moz:firefoxOptions': { args: ['-headless'], binary: ff },
            // This box is CPU-starved: boot + animation can outlive the
            // driver's 30 s default. Raise the script budget.
            timeouts: { script: 180000 },
          },
        },
      });
      break;
    } catch (err) { lastErr = String(err?.message ?? err); }
  }
  if (!sess) fail(`no firefox candidate accepted: ${lastErr}`);
  const sid = sess.value?.sessionId;
  if (!sid) fail('no session id');

  // Belt & braces: also set the session script timeout explicitly.
  await wd('POST', `/session/${sid}/timeouts`, { script: 180000 }).catch((e) => {
    console.error('timeouts: ' + String(e?.message ?? e).slice(0, 120));
  });

  await wd('POST', `/session/${sid}/url`, { url });

  // Settle: let the page finish its async boot (module top-level awaits)
  // before the setup script runs — a runner-side wait, since the page's
  // own timers may be starved in headless.
  if (Number(settleMs) > 0) await sleep(Number(settleMs));

  if (setup) {
    // Multiple scripts, each held as a pending Promise under execute/sync
    // (geckodriver honours the session script timeout set above — 3 min).
    // setup may be a single script string or a JSON array of them.
    let steps = [setup];
    try { const parsed = JSON.parse(setup); if (Array.isArray(parsed)) steps = parsed; } catch { /* single script */ }
    for (let i = 0; i < steps.length; i++) {
      const script = `return (async () => { ${steps[i]} })();`;
      try {
        const out = await wd('POST', `/session/${sid}/execute/sync`, { script, args: [] });
        if (out?.value !== undefined && out.value !== null) {
          console.log('setup[' + i + '] → ' + JSON.stringify(out.value).slice(0, 400));
        }
      } catch (err) {
        console.error(`setup[${i}]: ` + String(err?.message ?? err).slice(0, 200));
      }
    }
  }

  const shot = await wd('GET', `/session/${sid}/screenshot`);
  if (!shot?.value) fail('no screenshot in response');
  const b64 = String(shot.value).replace(/^data:image\/\w+;base64,/, '');
  writeFileSync(out, Buffer.from(b64, 'base64'));
  console.log('saved ' + out);

  await wd('DELETE', `/session/${sid}`).catch(() => {});
} finally {
  try { driver.kill('SIGKILL'); } catch { /* dead */ }
}