Harden run-code-block server readiness and screenshot navigation retries

src/lib/plugins/run-code-blocks/directives/screenshot.ts

@@ -57,11 +57,24 @@ function compile(steps: string, path: `${string}.png`, args: Args): string {
 
   let script = [
 `const puppeteer = require('puppeteer');
+const NAVIGATION_TIMEOUT = 180000;
+const MAX_NAVIGATION_RETRIES = 6;
+const RETRYABLE_NAVIGATION_ERRORS = [
+  'Navigation timeout',
+  'net::ERR_CONNECTION_REFUSED',
+  'net::ERR_CONNECTION_RESET',
+  'net::ERR_ABORTED',
+  'ERR_HTTP_RESPONSE_CODE_FAILURE'
+];
+
+function retryDelay(attempt) {
+  return Math.min(500 * Math.pow(2, attempt), 5000);
+}
 
 async function main() {
   let browser = await puppeteer.launch();
   let page = await browser.newPage();
-  page.setDefaultNavigationTimeout(120000);
+  page.setDefaultNavigationTimeout(NAVIGATION_TIMEOUT);
   await page.setViewport(${js(viewport)});
 `
   ];
@@ -92,13 +105,24 @@ async function main() {
         script.push(`  await page.evaluate(${js(params[0])});`);
         break;
       case 'visit':
-        script.push(`  for (let _attempt = 0; _attempt < 3; _attempt++) {`);
+        script.push(`  for (let _attempt = 0; _attempt < MAX_NAVIGATION_RETRIES; _attempt++) {`);
         script.push(`    try {`);
-        script.push(`      await page.goto(${js(params[0])}, { waitUntil: 'domcontentloaded', timeout: 120000 });`);
+        script.push(`      await page.goto(${js(params[0])}, { waitUntil: 'domcontentloaded', timeout: NAVIGATION_TIMEOUT });`);
         script.push(`      break;`);
         script.push(`    } catch (e) {`);
-        script.push(`      if (_attempt === 2) throw e;`);
-        script.push(`      await new Promise(r => setTimeout(r, 2000));`);
+        script.push(`      let message = e instanceof Error ? e.message : String(e);`);
+        script.push(`      let shouldRetry = RETRYABLE_NAVIGATION_ERRORS.some(pattern => message.includes(pattern));`);
+        script.push(`      if (_attempt === MAX_NAVIGATION_RETRIES - 1 || !shouldRetry) throw e;`);
+        script.push(`      try {`);
+        script.push(`        await page.close();`);
+        script.push(`      } catch (closeError) {`);
+        script.push(`        let closeMessage = closeError instanceof Error ? closeError.message : String(closeError);`);
+        script.push(`        if (!closeMessage.includes('Target closed')) throw closeError;`);
+        script.push(`      }`);
+        script.push(`      page = await browser.newPage();`);
+        script.push(`      page.setDefaultNavigationTimeout(NAVIGATION_TIMEOUT);`);
+        script.push(`      await page.setViewport(${js(viewport)});`);
+        script.push(`      await new Promise(r => setTimeout(r, retryDelay(_attempt)));`);
         script.push(`    }`);
         script.push(`  }`);
         break;

src/lib/plugins/run-code-blocks/directives/server/start.ts

@@ -1,11 +1,15 @@
 import { Code } from 'mdast';
+import { createConnection } from 'net';
 import { join } from 'path';
 import { Option, assert } from 'ts-std';
 import { parseCommand } from '../../commands';
 import Options from '../../options';
 import parseArgs, { ToBool, optional } from '../../parse-args';
 import Servers from '../../servers';
 
+const DEFAULT_READY_TIMEOUT = 30000;
+const READY_RETRY_INTERVAL = 500;
+
 interface Args {
   id?: string;
   lang?: string;
@@ -17,6 +21,68 @@ interface Args {
   captureOutput?: boolean;
 }
 
+function extractURL(expect: Option<string> | undefined): Option<string> {
+  if (!expect) {
+    return null;
+  }
+
+  let match = expect.match(/https?:\/\/[^\s"']+/);
+
+  if (match) {
+    try {
+      return new URL(match[0]!).toString();
+    } catch {
+      return null;
+    }
+  } else {
+    return null;
+  }
+}
+
+async function waitForServerReady(url: string, timeout: number): Promise<void> {
+  let parsed = new URL(url);
+  let host = parsed.hostname;
+  let port = parsed.port ? Number(parsed.port) : (parsed.protocol === 'https:' ? 443 : 80);
+  let startedAt = Date.now();
+  let lastError: Option<Error> = null;
+
+  while (Date.now() - startedAt < timeout) {
+    try {
+      await new Promise<void>((resolve, reject) => {
+        let socket = createConnection({ host, port });
+
+        socket.once('connect', () => {
+          socket.end();
+          resolve();
+        });
+
+        socket.once('error', e => {
+          socket.destroy();
+
+          if (e instanceof Error) {
+            reject(e);
+          } else {
+            reject(new Error(String(e)));
+          }
+        });
+      });
+
+      return;
+    } catch (e) {
+      if (e instanceof Error) {
+        lastError = e;
+      } else {
+        lastError = new Error(String(e));
+      }
+    }
+
+    await new Promise(resolve => setTimeout(resolve, READY_RETRY_INTERVAL));
+  }
+
+  let details = lastError ? ` Last error: ${lastError.message}` : '';
+  throw new Error(`Timed out while waiting for ${url} to accept connections after ${timeout}ms.${details}`);
+}
+
 export default async function startServer(node: Code, options: Options, servers: Servers): Promise<Option<Code>> {
   let args = parseArgs<Args>(node, [
     optional('id', String),
@@ -72,7 +138,36 @@ export default async function startServer(node: Code, options: Options, servers:
     output.push(`$ ${display}`);
   }
 
-  let stdout = await server.start(args.expect);
+  let stdout = await server.start(args.expect, args.timeout);
+
+  let readyURL = extractURL(args.expect);
+
+  if (readyURL) {
+    let timeout = args.timeout || DEFAULT_READY_TIMEOUT;
+
+    try {
+      await waitForServerReady(readyURL, timeout);
+    } catch (e) {
+      await server.kill();
+
+      let message = (e instanceof Error) ? e.message : String(e);
+
+      throw new Error(
+`${message}
+
+====== STDOUT ======
+
+${server.stdout || '(No output)'}
+
+====== STDERR ======
+
+${server.stderr || '(No output)'}
+
+====================
+`
+      );
+    }
+  }
 
   if (args.captureOutput && stdout) {
     output.push(stdout);