@onlook/storybook-plugin 0.3.1 → 0.3.2-beta.1

package/README.md

@@ -44,7 +44,7 @@ Configures Vite HMR for E2B sandboxes (WSS protocol, port 443 routing).
 
 ### CORS
 
-Pre-configured for `https://app.onlook.ai`, `localhost:3000`, and `localhost:6006`.
+Pre-configured for `https://app.onlook.com`, `localhost:3000`, and `localhost:6006`.
 
 ## Options
 

package/dist/cli/index.js

@@ -1,11 +1,13 @@
 #!/usr/bin/env node
+import { createRequire } from 'node:module';
 import { command, string, boolean, run } from '@drizzle-team/brocli';
 import { spawn } from 'child_process';
+import fs, { readFileSync, existsSync } from 'fs';
 import path, { resolve, join, dirname } from 'path';
 import crypto from 'crypto';
-import fs, { existsSync } from 'fs';
 import { chromium } from 'playwright';
 
+globalThis.require = createRequire(import.meta.url);
 var CACHE_DIR = path.join(process.cwd(), ".storybook-cache");
 var SCREENSHOTS_DIR = path.join(CACHE_DIR, "screenshots");
 var MANIFEST_PATH = path.join(CACHE_DIR, "manifest.json");
@@ -95,20 +97,27 @@ async function captureScreenshotBuffer(storyId, theme, width = VIEWPORT_WIDTH, h
     await page.goto(url, { timeout: timeoutMs });
     await page.waitForLoadState("domcontentloaded", { timeout: timeoutMs });
     await page.waitForLoadState("load", { timeout: timeoutMs });
-    await page.waitForLoadState("networkidle", { timeout: timeoutMs });
+    try {
+      await page.waitForLoadState("networkidle", { timeout: 5e3 });
+    } catch {
+    }
     await page.evaluate(() => document.fonts.ready);
-    await page.evaluate(async () => {
-      const images = document.querySelectorAll("img");
-      await Promise.all(
-        Array.from(images).map((img) => {
-          if (img.complete) return Promise.resolve();
-          return new Promise((resolve3) => {
-            img.addEventListener("load", resolve3);
-            img.addEventListener("error", resolve3);
-          });
-        })
-      );
-    });
+    try {
+      await page.evaluate(async () => {
+        const images = document.querySelectorAll("img");
+        await Promise.all(
+          Array.from(images).map((img) => {
+            if (img.complete) return Promise.resolve();
+            return new Promise((resolve3) => {
+              img.addEventListener("load", resolve3);
+              img.addEventListener("error", resolve3);
+              setTimeout(resolve3, 3e3);
+            });
+          })
+        );
+      });
+    } catch {
+    }
     const contentBounds = await page.evaluate(() => {
       const root = document.querySelector("#storybook-root");
       if (!root) return null;
@@ -183,6 +192,33 @@ async function generateScreenshot(storyId, theme, storybookUrl = "http://localho
 }
 
 // src/screenshot-service/screenshot-service.ts
+var StoryTimeoutError = class extends Error {
+  constructor(storyId, timeoutMs) {
+    super(`Story ${storyId} timed out after ${timeoutMs / 1e3}s`);
+    this.storyId = storyId;
+    this.timeoutMs = timeoutMs;
+    this.name = "StoryTimeoutError";
+  }
+};
+async function captureBothThemes(storyId, storybookUrl, timeoutMs, storyTimeoutMs) {
+  let timer;
+  try {
+    return await Promise.race([
+      Promise.all([
+        generateScreenshot(storyId, "light", storybookUrl, timeoutMs),
+        generateScreenshot(storyId, "dark", storybookUrl, timeoutMs)
+      ]),
+      new Promise((_, reject) => {
+        timer = setTimeout(
+          () => reject(new StoryTimeoutError(storyId, storyTimeoutMs)),
+          storyTimeoutMs
+        );
+      })
+    ]);
+  } finally {
+    if (timer) clearTimeout(timer);
+  }
+}
 async function generateAllScreenshots(stories, storybookUrl = "http://localhost:6006", concurrency = 10, offset = 0, total, skipExisting = false, timeoutMs = 3e4) {
   const displayTotal = total ?? offset + stories.length;
   console.log(
@@ -193,6 +229,7 @@ async function generateAllScreenshots(stories, storybookUrl = "http://localhost:
   for (let i = 0; i < stories.length; i += BATCH_SIZE) {
     batches.push(stories.slice(i, i + BATCH_SIZE));
   }
+  const storyTimeout = timeoutMs * 2 + 1e4;
   let completed = 0;
   for (const batch of batches) {
     await Promise.all(
@@ -203,19 +240,64 @@ async function generateAllScreenshots(stories, storybookUrl = "http://localhost:
           console.log(`[${absoluteIndex2}/${displayTotal}] Skipped (exists) ${story.id}`);
           return;
         }
-        const [lightResult, darkResult] = await Promise.all([
-          generateScreenshot(story.id, "light", storybookUrl, timeoutMs),
-          generateScreenshot(story.id, "dark", storybookUrl, timeoutMs)
-        ]);
+        let lightResult = null;
+        let darkResult = null;
+        let lastError;
+        const startedAt = Date.now();
+        let attempts = 0;
+        const maxAttempts = 2;
+        while (attempts < maxAttempts) {
+          attempts++;
+          try {
+            [lightResult, darkResult] = await captureBothThemes(
+              story.id,
+              storybookUrl,
+              timeoutMs,
+              storyTimeout
+            );
+            lastError = void 0;
+            break;
+          } catch (error) {
+            lastError = error;
+            if (!(error instanceof StoryTimeoutError) || attempts >= maxAttempts) {
+              break;
+            }
+            console.warn(
+              `[screenshot] Retrying ${story.id} after timeout (attempt ${attempts + 1}/${maxAttempts})`
+            );
+          }
+        }
+        completed++;
+        const absoluteIndex = offset + completed;
+        const durationMs = Date.now() - startedAt;
         if (lightResult && darkResult) {
           const fileHash = computeFileHash(story.importPath);
           updateManifest(story.id, story.importPath, fileHash, lightResult.boundingBox);
+          console.log(
+            `[${absoluteIndex}/${displayTotal}] Generated screenshots for ${story.id}`,
+            { storyId: story.id, attempts, durationMs, outcome: "ok" }
+          );
+        } else if (lastError) {
+          console.error(
+            `[${absoluteIndex}/${displayTotal}] \u26A0\uFE0F Failed ${story.id}:`,
+            lastError instanceof Error ? lastError.message : lastError,
+            {
+              storyId: story.id,
+              attempts,
+              durationMs,
+              outcome: lastError instanceof StoryTimeoutError ? "timeout" : "error"
+            }
+          );
+        } else {
+          console.warn(`[${absoluteIndex}/${displayTotal}] Partial result ${story.id}`, {
+            storyId: story.id,
+            attempts,
+            durationMs,
+            outcome: "partial",
+            hasLight: !!lightResult,
+            hasDark: !!darkResult
+          });
         }
-        completed++;
-        const absoluteIndex = offset + completed;
-        console.log(
-          `[${absoluteIndex}/${displayTotal}] Generated screenshots for ${story.id}`
-        );
       })
     );
   }
@@ -260,10 +342,27 @@ function getStorybookCommand(storybookDir) {
   console.log(
     `\u{1F4E6} Detected package manager: ${pm}${repoRoot ? ` (from ${repoRoot})` : ""}`
   );
-  const separator = pm === "npm" || pm === "pnpm" ? ["--"] : [];
+  let extraArgs = [];
+  try {
+    const pkgPath = resolve(storybookDir, "package.json");
+    const pkg = JSON.parse(readFileSync(pkgPath, "utf-8"));
+    const script = pkg.scripts?.storybook ?? "";
+    const hasPort = script.includes("-p ") || script.includes("--port");
+    const hasNoOpen = script.includes("--no-open");
+    if (!hasPort || !hasNoOpen) {
+      const needsSeparator = pm === "npm" || pm === "pnpm";
+      const flags = [];
+      if (!hasPort) flags.push("-p", "6006");
+      if (!hasNoOpen) flags.push("--no-open");
+      extraArgs = needsSeparator && flags.length > 0 ? ["--", ...flags] : flags;
+    }
+  } catch {
+    const separator = pm === "npm" || pm === "pnpm" ? ["--"] : [];
+    extraArgs = [...separator, "-p", "6006", "--no-open"];
+  }
   return {
     command: pm,
-    args: ["run", "storybook", ...separator, "-p", "6006", "--no-open"]
+    args: ["run", "storybook", ...extraArgs]
   };
 }
 async function isStorybookRunning(url) {
@@ -276,14 +375,38 @@ async function isStorybookRunning(url) {
     return false;
   }
 }
-async function fetchStoryIndex(url) {
+async function waitForStorybookReady(url, timeoutMs = 12e4) {
+  const start = Date.now();
+  const pollInterval = 2e3;
+  while (Date.now() - start < timeoutMs) {
+    if (await isStorybookRunning(url)) return;
+    await new Promise((r) => setTimeout(r, pollInterval));
+  }
+  throw new Error(`Storybook did not become ready within ${timeoutMs / 1e3}s`);
+}
+async function fetchStoryIndex(url, retries = 3) {
   const indexUrl = `${url}/index.json`;
-  const response = await fetch(indexUrl);
-  if (!response.ok) {
-    throw new Error(`Failed to fetch story index: ${response.statusText}`);
+  for (let attempt = 1; attempt <= retries; attempt++) {
+    try {
+      const response = await fetch(indexUrl, {
+        signal: AbortSignal.timeout(1e4)
+      });
+      if (!response.ok) {
+        throw new Error(
+          `Failed to fetch story index: ${response.status} ${response.statusText}`
+        );
+      }
+      const data = await response.json();
+      return Object.values(data.entries);
+    } catch (error) {
+      if (attempt === retries) throw error;
+      console.log(
+        `\u26A0\uFE0F Fetch story index failed (attempt ${attempt}/${retries}), retrying in 3s...`
+      );
+      await new Promise((r) => setTimeout(r, 3e3));
+    }
   }
-  const data = await response.json();
-  return Object.values(data.entries);
+  throw new Error("Unreachable");
 }
 async function startStorybook(command2, args, storybookDir) {
   console.log(
@@ -301,12 +424,12 @@ async function startStorybook(command2, args, storybookDir) {
     const timeout = setTimeout(() => {
       if (!started) {
         storybookProcess.kill();
-        reject(new Error("Storybook failed to start within 60 seconds"));
+        reject(new Error("Storybook failed to start within 180 seconds"));
       }
-    }, 6e4);
+    }, 18e4);
     storybookProcess.stdout?.on("data", (data) => {
       const output = data.toString();
-      if (output.includes("Local:") || output.includes("localhost:")) {
+      if (output.includes("Local:") || output.includes("localhost:") || output.includes("local:")) {
         if (!started) {
           started = true;
           clearTimeout(timeout);
@@ -335,7 +458,7 @@ async function warmupStorybook(url, firstStoryId) {
   try {
     const warmupUrl = `${url}/iframe.html?id=${firstStoryId}&viewMode=story`;
     await page.goto(warmupUrl, { timeout: 15e3 });
-    await page.waitForLoadState("networkidle");
+    await page.waitForLoadState("domcontentloaded", { timeout: 15e3 });
     console.log("\u2705 Storybook warmed up");
   } catch {
     console.log("\u26A0\uFE0F Warmup had issues, proceeding anyway");
@@ -358,6 +481,8 @@ async function generateScreenshots(options = {}) {
       const { command: command2, args } = getStorybookCommand(storybookDir);
       storybookProcess = await startStorybook(command2, args, storybookDir);
       weStartedStorybook = true;
+      console.log("\u23F3 Waiting for Storybook to finish compiling...");
+      await waitForStorybookReady(url);
     }
     const allStories = await fetchStoryIndex(url);
     const total = allStories.length;

package/dist/index.d.ts

@@ -1,11 +1,21 @@
 import { PluginOption } from 'vite';
+import { Indexer } from 'storybook/internal/types';
 
 type OnlookPluginOptions = {
     /** Storybook port (default: 6006) */
     port?: number;
     /** Additional allowed origins for CORS (merged with defaults) */
     allowedOrigins?: string[];
+    /**
+     * Set when this plugin runs inside an E2B sandbox. Enables the
+     * WSS:443 HMR override required for Vite's HMR client to talk back
+     * through E2B's HTTPS port-proxy. Leave false (default) for local
+     * `bun dev:storybook` so HMR uses Vite defaults and actually fires.
+     */
+    e2b?: boolean;
 };
 declare function storybookOnlookPlugin(options?: OnlookPluginOptions): PluginOption[];
 
-export { type OnlookPluginOptions, storybookOnlookPlugin };
+declare const tolerantCsfIndexer: Indexer;
+
+export { type OnlookPluginOptions, storybookOnlookPlugin, tolerantCsfIndexer };