@pygmalionjs/pygmalion 0.2.44 → 0.4.0

package/node/preview-capture-worker.mjs

@@ -0,0 +1,529 @@
+import { spawn } from 'node:child_process';
+import fs from 'node:fs/promises';
+import { validateRoutePreviewArtifactBundle } from './route-preview-artifact-v3.mjs';
+
+const DEFAULT_IDLE_MS = 300_000;
+const DEFAULT_JOB_TIMEOUT_MS = 1_200_000;
+const DEFAULT_MAX_RESTARTS = 3;
+const DEFAULT_RESTART_WINDOW_MS = 600_000;
+const DEFAULT_PING_TIMEOUT_MS = 10_000;
+const MAX_RESTART_BACKOFF_MS = 8_000;
+const STOP_KILL_GRACE_MS = 5_000;
+const STDERR_TAIL_LIMIT = 8_000;
+
+function assertPositiveNumber(value, name, fallback) {
+  if (value == null) return fallback;
+  if (typeof value !== 'number' || !Number.isFinite(value) || value <= 0) {
+    throw new TypeError(`Preview capture worker ${name} must be a positive number.`);
+  }
+  return value;
+}
+
+function assertNonNegativeInteger(value, name, fallback) {
+  if (value == null) return fallback;
+  if (!Number.isInteger(value) || value < 0) {
+    throw new TypeError(
+      `Preview capture worker ${name} must be a non-negative integer.`,
+    );
+  }
+  return value;
+}
+
+function normalizeArgs(args) {
+  if (args == null) return [];
+  if (
+    !Array.isArray(args) ||
+    args.some((entry) => typeof entry !== 'string')
+  ) {
+    throw new TypeError(
+      'Preview capture worker args must be an array of strings.',
+    );
+  }
+  return [...args];
+}
+
+function normalizeRequest(request) {
+  if (!request || typeof request !== 'object') {
+    throw new TypeError('Preview capture request must be an object.');
+  }
+  const { namespace, sourceRevision, frames, captureBaseUrl } = request;
+  if (typeof namespace !== 'string' || !namespace.trim()) {
+    throw new TypeError('Preview capture request requires a namespace.');
+  }
+  if (typeof sourceRevision !== 'string' || !sourceRevision.trim()) {
+    throw new TypeError('Preview capture request requires a source revision.');
+  }
+  if (frames != null && !Array.isArray(frames)) {
+    throw new TypeError('Preview capture request frames must be an array.');
+  }
+  if (captureBaseUrl != null && typeof captureBaseUrl !== 'string') {
+    throw new TypeError(
+      'Preview capture request captureBaseUrl must be a string.',
+    );
+  }
+  return {
+    namespace,
+    sourceRevision,
+    ...(frames == null ? {} : { frames: frames.map((frame) => ({ ...frame })) }),
+    ...(captureBaseUrl == null ? {} : { captureBaseUrl }),
+  };
+}
+
+function attachLineParser(stream, onMessage) {
+  let buffer = '';
+  stream.setEncoding('utf8');
+  stream.on('data', (chunk) => {
+    buffer += chunk;
+    let newline = buffer.indexOf('\n');
+    while (newline >= 0) {
+      const line = buffer.slice(0, newline).trim();
+      buffer = buffer.slice(newline + 1);
+      newline = buffer.indexOf('\n');
+      if (!line) continue;
+      let message;
+      try {
+        message = JSON.parse(line);
+      } catch {
+        continue; // Not a protocol line; the worker owns its own noise.
+      }
+      if (message && typeof message === 'object') onMessage(message);
+    }
+  });
+}
+
+function attachStderrTail(stream, worker) {
+  stream.setEncoding('utf8');
+  stream.on('data', (chunk) => {
+    worker.stderrTail = `${worker.stderrTail}${chunk}`.slice(-STDERR_TAIL_LIMIT);
+  });
+}
+
+async function consumeArtifactFile(response) {
+  const file = response?.artifactFile;
+  if (typeof file !== 'string' || !file.trim()) {
+    throw new Error('Preview capture worker returned no artifact file.');
+  }
+  try {
+    const bundle = JSON.parse(await fs.readFile(file, 'utf8'));
+    const validation = validateRoutePreviewArtifactBundle(bundle);
+    if (!validation.valid) {
+      throw new Error(
+        `Preview capture worker artifact is invalid: ${(
+          validation.errors ?? ['unknown validation error']
+        ).join(', ')}`,
+      );
+    }
+    return bundle;
+  } finally {
+    // The hand-off file is single-use: the bundle now lives in memory and the
+    // plugin persists it through the artifact store, so the temp file would
+    // only accumulate.
+    await fs.rm(file, { force: true }).catch(() => undefined);
+  }
+}
+
+function workerFailure(worker, fallbackMessage) {
+  const tail = worker.stderrTail.trim();
+  return new Error(tail ? `${fallbackMessage}\n${tail}` : fallbackMessage);
+}
+
+/**
+ * Long-lived capture worker channel: one child process that keeps its expensive
+ * boot (runtime preflight, case bundling, browser launch) warm across artifact
+ * requests, speaking JSON lines over stdio.
+ *
+ * Protocol: the channel writes `{ id, type: 'capture', payload }` — payload is
+ * `{ namespace, sourceRevision, frames?, captureBaseUrl? }` — and the worker
+ * answers `{ id, ok: true, artifactFile }` with the bundle in a temp file the
+ * channel reads, validates, and deletes, or `{ id, ok: false, error }`.
+ * `{ id, type: 'ping' }` must be answered immediately even while booting; the
+ * channel uses it to detect a wedged worker before reusing it.
+ *
+ * Failure policy: a crash rejects the in-flight job and the worker restarts on
+ * the next job with capped backoff. After `maxRestarts` crashes inside
+ * `restartWindowMs` the channel degrades permanently to spawning the script
+ * per request with `--once` — slower, but never dark.
+ */
+export function createPreviewCaptureWorkerChannel({
+  script,
+  args = [],
+  env,
+  cwd,
+  idleMs = DEFAULT_IDLE_MS,
+  jobTimeoutMs = DEFAULT_JOB_TIMEOUT_MS,
+  maxRestarts = DEFAULT_MAX_RESTARTS,
+  restartWindowMs = DEFAULT_RESTART_WINDOW_MS,
+  pingTimeoutMs = DEFAULT_PING_TIMEOUT_MS,
+} = {}) {
+  if (typeof script !== 'string' || !script.trim()) {
+    throw new TypeError(
+      'Preview capture worker script must be a non-empty string.',
+    );
+  }
+  const workerArgs = normalizeArgs(args);
+  if (env != null && (typeof env !== 'object' || Array.isArray(env))) {
+    throw new TypeError('Preview capture worker env must be an object.');
+  }
+  if (cwd != null && (typeof cwd !== 'string' || !cwd.trim())) {
+    throw new TypeError('Preview capture worker cwd must be a non-empty string.');
+  }
+  const resolvedIdleMs = assertPositiveNumber(idleMs, 'idleMs', DEFAULT_IDLE_MS);
+  const resolvedJobTimeoutMs = assertPositiveNumber(
+    jobTimeoutMs,
+    'jobTimeoutMs',
+    DEFAULT_JOB_TIMEOUT_MS,
+  );
+  const resolvedMaxRestarts = assertNonNegativeInteger(
+    maxRestarts,
+    'maxRestarts',
+    DEFAULT_MAX_RESTARTS,
+  );
+  const resolvedRestartWindowMs = assertPositiveNumber(
+    restartWindowMs,
+    'restartWindowMs',
+    DEFAULT_RESTART_WINDOW_MS,
+  );
+  const resolvedPingTimeoutMs = assertPositiveNumber(
+    pingTimeoutMs,
+    'pingTimeoutMs',
+    DEFAULT_PING_TIMEOUT_MS,
+  );
+
+  const spawnEnv = env ? { ...process.env, ...env } : undefined;
+  let worker = null;
+  let inFlight = false;
+  let degraded = false;
+  let disposed = false;
+  let idleTimer = null;
+  let requestCounter = 0;
+  const crashTimestamps = [];
+
+  function nextRequestId() {
+    requestCounter += 1;
+    return `capture-${process.pid}-${requestCounter}`;
+  }
+
+  function recordCrash(now = Date.now()) {
+    while (
+      crashTimestamps.length > 0 &&
+      now - crashTimestamps[0] > resolvedRestartWindowMs
+    ) {
+      crashTimestamps.shift();
+    }
+    crashTimestamps.push(now);
+    if (crashTimestamps.length > resolvedMaxRestarts) degraded = true;
+  }
+
+  function restartBackoffMs() {
+    const recent = crashTimestamps.length;
+    if (recent === 0) return 0;
+    return Math.min(500 * 2 ** (recent - 1), MAX_RESTART_BACKOFF_MS);
+  }
+
+  function clearIdleTimer() {
+    if (idleTimer) clearTimeout(idleTimer);
+    idleTimer = null;
+  }
+
+  function scheduleIdleShutdown() {
+    clearIdleTimer();
+    if (!worker || disposed) return;
+    idleTimer = setTimeout(() => {
+      // Idle means no job is in flight, so the worker drains instantly.
+      void stopWorker();
+    }, resolvedIdleMs);
+    idleTimer.unref?.();
+  }
+
+  function settlePending(current, settle) {
+    const pending = current.pending;
+    if (!pending) return;
+    current.pending = null;
+    clearTimeout(pending.timer);
+    settle(pending);
+  }
+
+  function spawnWorker() {
+    const proc = spawn(process.execPath, [script, ...workerArgs], {
+      ...(cwd ? { cwd } : {}),
+      ...(spawnEnv ? { env: spawnEnv } : {}),
+      stdio: ['pipe', 'pipe', 'pipe'],
+    });
+    const current = {
+      proc,
+      pending: null,
+      stopping: false,
+      stderrTail: '',
+      exited: new Promise((resolve) => {
+        proc.once('exit', resolve);
+        proc.once('error', resolve);
+      }),
+    };
+    attachLineParser(proc.stdout, (message) => {
+      const pending = current.pending;
+      if (!pending || message.id !== pending.id) return;
+      if (pending.kind === 'ping') {
+        settlePending(current, ({ resolve }) => resolve(message));
+        return;
+      }
+      if (message.ok === true) {
+        settlePending(current, ({ resolve }) => resolve(message));
+      } else {
+        // A refused job is the worker doing its job — the process stays warm.
+        settlePending(current, ({ reject }) =>
+          reject(
+            new Error(
+              typeof message.error === 'string' && message.error
+                ? message.error
+                : 'Preview capture worker reported a failure.',
+            ),
+          ),
+        );
+      }
+    });
+    attachStderrTail(proc.stderr, current);
+    const onGone = () => {
+      if (worker === current) worker = null;
+      if (current.stopping) return;
+      current.stopping = true;
+      recordCrash();
+      settlePending(current, ({ reject }) =>
+        reject(workerFailure(current, 'Preview capture worker exited unexpectedly.')),
+      );
+    };
+    proc.once('exit', onGone);
+    proc.once('error', onGone);
+    return current;
+  }
+
+  async function stopWorker() {
+    const current = worker;
+    if (!current) return;
+    worker = null;
+    current.stopping = true;
+    settlePending(current, ({ reject }) =>
+      reject(new Error('Preview capture worker channel is shutting down.')),
+    );
+    try {
+      current.proc.stdin.end();
+    } catch {
+      // Already closed; the signal below still applies.
+    }
+    current.proc.kill('SIGTERM');
+    const killTimer = setTimeout(() => {
+      current.proc.kill('SIGKILL');
+    }, STOP_KILL_GRACE_MS);
+    killTimer.unref?.();
+    await current.exited;
+    clearTimeout(killTimer);
+  }
+
+  function killWorker(current) {
+    if (worker === current) worker = null;
+    current.stopping = true;
+    current.proc.kill('SIGKILL');
+  }
+
+  function dispatch(current, message, { kind, timeoutMs, onTimeout }) {
+    return new Promise((resolve, reject) => {
+      const timer = setTimeout(() => {
+        current.pending = null;
+        onTimeout();
+        reject(
+          workerFailure(
+            current,
+            kind === 'ping'
+              ? 'Preview capture worker did not answer a health check.'
+              : `Preview capture job timed out after ${timeoutMs}ms.`,
+          ),
+        );
+      }, timeoutMs);
+      timer.unref?.();
+      current.pending = { id: message.id, kind, resolve, reject, timer };
+      try {
+        current.proc.stdin.write(`${JSON.stringify(message)}\n`);
+      } catch (error) {
+        settlePending(current, () => undefined);
+        reject(
+          workerFailure(
+            current,
+            `Preview capture worker is not accepting requests: ${
+              error?.message ?? error
+            }`,
+          ),
+        );
+      }
+    });
+  }
+
+  async function pingWorker(current) {
+    await dispatch(
+      current,
+      { id: nextRequestId(), type: 'ping' },
+      {
+        kind: 'ping',
+        timeoutMs: resolvedPingTimeoutMs,
+        onTimeout: () => {
+          // A worker that cannot answer a ping cannot run a job either.
+          recordCrash();
+          killWorker(current);
+        },
+      },
+    );
+  }
+
+  /** Returns a live worker, or null when the channel degraded on the way. */
+  async function ensureWorker() {
+    if (worker) {
+      const current = worker;
+      try {
+        await pingWorker(current);
+        return current;
+      } catch {
+        // The wedged worker was killed and the crash recorded. When that
+        // crash exhausted the restart budget, the caller takes the degraded
+        // path; otherwise fall through to a fresh spawn below.
+        if (degraded) return null;
+      }
+    }
+    const backoff = restartBackoffMs();
+    if (backoff > 0) {
+      await new Promise((resolve) => {
+        const timer = setTimeout(resolve, backoff);
+        timer.unref?.();
+      });
+    }
+    if (disposed) {
+      throw new Error('Preview capture worker channel is disposed.');
+    }
+    worker = spawnWorker();
+    return worker;
+  }
+
+  function runWorkerJob(current, payload) {
+    return dispatch(
+      current,
+      { id: nextRequestId(), type: 'capture', payload },
+      {
+        kind: 'capture',
+        timeoutMs: resolvedJobTimeoutMs,
+        onTimeout: () => {
+          recordCrash();
+          killWorker(current);
+        },
+      },
+    );
+  }
+
+  /**
+   * Degraded path: one spawn per request with `--once`, the same protocol over
+   * the child's own stdio. Slow again, but on-demand capture keeps answering.
+   */
+  function runOnceJob(payload) {
+    return new Promise((resolve, reject) => {
+      const proc = spawn(
+        process.execPath,
+        [script, ...workerArgs, '--once'],
+        {
+          ...(cwd ? { cwd } : {}),
+          ...(spawnEnv ? { env: spawnEnv } : {}),
+          stdio: ['pipe', 'pipe', 'pipe'],
+        },
+      );
+      const once = { stderrTail: '' };
+      let settled = false;
+      const settle = (fn) => {
+        if (settled) return;
+        settled = true;
+        clearTimeout(timer);
+        fn();
+      };
+      const timer = setTimeout(() => {
+        settle(() => {
+          proc.kill('SIGKILL');
+          reject(
+            new Error(`Preview capture job timed out after ${resolvedJobTimeoutMs}ms.`),
+          );
+        });
+      }, resolvedJobTimeoutMs);
+      timer.unref?.();
+      const id = nextRequestId();
+      attachLineParser(proc.stdout, (message) => {
+        if (message.id !== id) return;
+        if (message.ok === true) {
+          settle(() => resolve(message));
+        } else {
+          settle(() =>
+            reject(
+              new Error(
+                typeof message.error === 'string' && message.error
+                  ? message.error
+                  : 'Preview capture worker reported a failure.',
+              ),
+            ),
+          );
+        }
+        proc.stdin.end();
+      });
+      attachStderrTail(proc.stderr, once);
+      const onGone = () => {
+        settle(() =>
+          reject(
+            workerFailure(once, 'Preview capture worker exited without a response.'),
+          ),
+        );
+      };
+      proc.once('exit', onGone);
+      proc.once('error', onGone);
+      try {
+        proc.stdin.write(
+          `${JSON.stringify({ id, type: 'capture', payload })}\n`,
+        );
+      } catch (error) {
+        settle(() =>
+          reject(
+            new Error(
+              `Preview capture worker is not accepting requests: ${
+                error?.message ?? error
+              }`,
+            ),
+          ),
+        );
+      }
+    });
+  }
+
+  async function generateArtifact(request) {
+    if (disposed) {
+      throw new Error('Preview capture worker channel is disposed.');
+    }
+    // The artifact plugin serializes per identity; overlap here means a caller
+    // bypassed that queue, and interleaving two captures on one warm browser
+    // would let them poison each other.
+    if (inFlight) {
+      throw new Error('Preview capture worker accepts one job at a time.');
+    }
+    const payload = normalizeRequest(request);
+    inFlight = true;
+    clearIdleTimer();
+    try {
+      const current = degraded ? null : await ensureWorker();
+      const response = current
+        ? await runWorkerJob(current, payload)
+        : await runOnceJob(payload);
+      return await consumeArtifactFile(response);
+    } finally {
+      inFlight = false;
+      scheduleIdleShutdown();
+    }
+  }
+
+  async function dispose() {
+    if (disposed) return;
+    disposed = true;
+    clearIdleTimer();
+    await stopWorker();
+  }
+
+  return { generateArtifact, dispose };
+}