@onlook/storybook-plugin 0.3.2 → 0.4.0-beta.0

package/dist/cli/index.js

@@ -1,9 +1,9 @@
 #!/usr/bin/env node
 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';
 
 var CACHE_DIR = path.join(process.cwd(), ".storybook-cache");
@@ -95,20 +95,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;
@@ -199,23 +206,48 @@ async function generateAllScreenshots(stories, storybookUrl = "http://localhost:
       batch.map(async (story) => {
         if (skipExisting && screenshotExists(story.id, "light") && screenshotExists(story.id, "dark")) {
           completed++;
-          const absoluteIndex2 = offset + completed;
-          console.log(`[${absoluteIndex2}/${displayTotal}] Skipped (exists) ${story.id}`);
+          const absoluteIndex = offset + completed;
+          console.log(`[${absoluteIndex}/${displayTotal}] Skipped (exists) ${story.id}`);
           return;
         }
-        const [lightResult, darkResult] = await Promise.all([
-          generateScreenshot(story.id, "light", storybookUrl, timeoutMs),
-          generateScreenshot(story.id, "dark", storybookUrl, timeoutMs)
-        ]);
-        if (lightResult && darkResult) {
-          const fileHash = computeFileHash(story.importPath);
-          updateManifest(story.id, story.importPath, fileHash, lightResult.boundingBox);
+        const storyTimeout = timeoutMs * 2 + 1e4;
+        let timer;
+        try {
+          const result = await Promise.race([
+            Promise.all([
+              generateScreenshot(story.id, "light", storybookUrl, timeoutMs),
+              generateScreenshot(story.id, "dark", storybookUrl, timeoutMs)
+            ]),
+            new Promise((_, reject) => {
+              timer = setTimeout(
+                () => reject(
+                  new Error(
+                    `Story ${story.id} timed out after ${storyTimeout / 1e3}s`
+                  )
+                ),
+                storyTimeout
+              );
+            })
+          ]);
+          clearTimeout(timer);
+          const [lightResult, darkResult] = result;
+          if (lightResult && darkResult) {
+            const fileHash = computeFileHash(story.importPath);
+            updateManifest(story.id, story.importPath, fileHash, lightResult.boundingBox);
+          }
+          completed++;
+          const absoluteIndex = offset + completed;
+          console.log(
+            `[${absoluteIndex}/${displayTotal}] Generated screenshots for ${story.id}`
+          );
+        } catch (error) {
+          completed++;
+          const absoluteIndex = offset + completed;
+          console.error(
+            `[${absoluteIndex}/${displayTotal}] \u26A0\uFE0F Failed ${story.id}:`,
+            error instanceof Error ? error.message : error
+          );
         }
-        completed++;
-        const absoluteIndex = offset + completed;
-        console.log(
-          `[${absoluteIndex}/${displayTotal}] Generated screenshots for ${story.id}`
-        );
       })
     );
   }
@@ -260,10 +292,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 +325,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 +374,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 +408,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 +431,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,24 @@
 import { PluginOption } from 'vite';
 
+/** Folder name for auto-generated stories (created next to each component) */
+declare const AUTO_STORIES_FOLDER = ".onlook-stories";
 type OnlookPluginOptions = {
     /** Storybook port (default: 6006) */
     port?: number;
     /** Additional allowed origins for CORS (merged with defaults) */
     allowedOrigins?: string[];
+    /**
+     * Auto-generate stories for all components.
+     * Glob patterns for component file discovery.
+     * Set to false to disable. (default: ['src\/\*\*\/*.tsx'])
+     */
+    autoStories?: string[] | false;
+    /**
+     * Glob patterns to exclude from auto-story generation.
+     * (default: stories, tests, and node_modules)
+     */
+    autoStoriesIgnore?: string[];
 };
 declare function storybookOnlookPlugin(options?: OnlookPluginOptions): PluginOption[];
 
-export { type OnlookPluginOptions, storybookOnlookPlugin };
+export { AUTO_STORIES_FOLDER, type OnlookPluginOptions, storybookOnlookPlugin };

package/dist/index.js

@@ -1,6 +1,8 @@
-import fs4, { existsSync } from 'fs';
-import path2, { dirname, join, relative } from 'path';
+import fs5, { existsSync } from 'fs';
+import path, { dirname, join, relative } from 'path';
 import { fileURLToPath } from 'url';
+import autoStoryGenerator from '@takuma-ru/auto-story-generator';
+import { withDefaultConfig } from 'react-docgen-typescript';
 import generateModule from '@babel/generator';
 import { parse } from '@babel/parser';
 import traverseModule from '@babel/traverse';
@@ -9,6 +11,88 @@ import crypto from 'crypto';
 import { chromium } from 'playwright';
 
 // src/storybook-onlook-plugin.ts
+var FIXED_MARKER = "// @onlook-fixed";
+var parser = withDefaultConfig({
+  shouldExtractLiteralValuesFromEnum: true,
+  shouldRemoveUndefinedFromOptional: true,
+  propFilter: (prop) => {
+    if (prop.parent?.fileName.includes("node_modules")) return false;
+    return true;
+  }
+});
+function resolveComponentPath(storyFilePath) {
+  const dir = path.dirname(storyFilePath);
+  const parentDir = path.dirname(dir);
+  const storyName = path.basename(storyFilePath);
+  const componentName = storyName.replace(".stories.tsx", ".tsx");
+  const componentPath = path.join(parentDir, componentName);
+  return fs5.existsSync(componentPath) ? componentPath : null;
+}
+function generateArgTypes(componentPath) {
+  try {
+    const docs = parser.parse(componentPath);
+    if (docs.length === 0) return null;
+    const doc = docs[0];
+    if (!doc) return null;
+    const argTypes = {};
+    for (const [name, prop] of Object.entries(doc.props)) {
+      if (prop.type.name.includes("=>") || prop.type.name === "Function") {
+        continue;
+      }
+      const argType = {};
+      if (prop.description) {
+        argType.description = prop.description;
+      }
+      if (prop.type.name === "enum" && prop.type.value) {
+        argType.control = { type: "select" };
+        argType.options = prop.type.value.map(
+          (v) => v.value
+        );
+      }
+      if (prop.type.name === "boolean") {
+        argType.control = { type: "boolean" };
+      }
+      if (prop.type.name === "number") {
+        argType.control = { type: "number" };
+      }
+      if (Object.keys(argType).length > 0) {
+        argTypes[name] = argType;
+      }
+    }
+    return Object.keys(argTypes).length > 0 ? argTypes : null;
+  } catch (err) {
+    console.error(`[AutoStories] Failed to parse component: ${componentPath}`, err);
+    return null;
+  }
+}
+function enrichStoryFile(storyFilePath) {
+  try {
+    const content = fs5.readFileSync(storyFilePath, "utf-8");
+    if (content.includes(FIXED_MARKER)) {
+      return;
+    }
+    if (content.includes("argTypes:")) {
+      return;
+    }
+    const componentPath = resolveComponentPath(storyFilePath);
+    if (!componentPath) return;
+    const argTypes = generateArgTypes(componentPath);
+    if (!argTypes) return;
+    const argTypesStr = JSON.stringify(argTypes, null, 2).split("\n").map((line, i) => i === 0 ? line : `  ${line}`).join("\n");
+    const enriched = content.replace(
+      /};\s*\nexport default meta;/,
+      `  argTypes: ${argTypesStr},
+};
+export default meta;`
+    );
+    if (enriched !== content) {
+      fs5.writeFileSync(storyFilePath, enriched);
+      console.log(`[AutoStories] Enriched ${path.basename(storyFilePath)} with argTypes`);
+    }
+  } catch (err) {
+    console.error(`[AutoStories] Failed to enrich story: ${storyFilePath}`, err);
+  }
+}
 function componentLocPlugin(options = {}) {
   const include = options.include ?? /\.(jsx|tsx)$/;
   const traverse = traverseModule.default ?? traverseModule;
@@ -31,7 +115,7 @@ function componentLocPlugin(options = {}) {
         sourceFilename: filepath
       });
       let mutated = false;
-      const relativePath = path2.relative(root, filepath);
+      const relativePath = path.relative(root, filepath);
       traverse(ast, {
         JSXElement(nodePath) {
           const opening = nodePath.node.openingElement;
@@ -68,9 +152,9 @@ function componentLocPlugin(options = {}) {
     }
   };
 }
-var CACHE_DIR = path2.join(process.cwd(), ".storybook-cache");
-var SCREENSHOTS_DIR = path2.join(CACHE_DIR, "screenshots");
-var MANIFEST_PATH = path2.join(CACHE_DIR, "manifest.json");
+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");
 var VIEWPORT_WIDTH = 1920;
 var VIEWPORT_HEIGHT = 1080;
 var MIN_COMPONENT_WIDTH = 420;
@@ -78,30 +162,30 @@ var MIN_COMPONENT_HEIGHT = 280;
 
 // src/utils/fileSystem/fileSystem.ts
 function ensureCacheDirectories() {
-  if (!fs4.existsSync(CACHE_DIR)) {
-    fs4.mkdirSync(CACHE_DIR, { recursive: true });
+  if (!fs5.existsSync(CACHE_DIR)) {
+    fs5.mkdirSync(CACHE_DIR, { recursive: true });
   }
-  if (!fs4.existsSync(SCREENSHOTS_DIR)) {
-    fs4.mkdirSync(SCREENSHOTS_DIR, { recursive: true });
+  if (!fs5.existsSync(SCREENSHOTS_DIR)) {
+    fs5.mkdirSync(SCREENSHOTS_DIR, { recursive: true });
   }
 }
 function computeFileHash(filePath) {
-  if (!fs4.existsSync(filePath)) {
+  if (!fs5.existsSync(filePath)) {
     return "";
   }
-  const content = fs4.readFileSync(filePath, "utf-8");
+  const content = fs5.readFileSync(filePath, "utf-8");
   return crypto.createHash("sha256").update(content).digest("hex");
 }
 function loadManifest() {
-  if (fs4.existsSync(MANIFEST_PATH)) {
-    const content = fs4.readFileSync(MANIFEST_PATH, "utf-8");
+  if (fs5.existsSync(MANIFEST_PATH)) {
+    const content = fs5.readFileSync(MANIFEST_PATH, "utf-8");
     return JSON.parse(content);
   }
   return { stories: {} };
 }
 function saveManifest(manifest) {
   ensureCacheDirectories();
-  fs4.writeFileSync(MANIFEST_PATH, JSON.stringify(manifest, null, 2));
+  fs5.writeFileSync(MANIFEST_PATH, JSON.stringify(manifest, null, 2));
 }
 function updateManifest(storyId, sourcePath, fileHash, boundingBox) {
   const manifest = loadManifest();
@@ -127,8 +211,8 @@ async function getBrowser() {
   return browser;
 }
 function getScreenshotPath(storyId, theme) {
-  const storyDir = path2.join(SCREENSHOTS_DIR, storyId);
-  return path2.join(storyDir, `${theme}.png`);
+  const storyDir = path.join(SCREENSHOTS_DIR, storyId);
+  return path.join(storyDir, `${theme}.png`);
 }
 async function captureScreenshotBuffer(storyId, theme, width = VIEWPORT_WIDTH, height = VIEWPORT_HEIGHT, storybookUrl = "http://localhost:6006", timeoutMs = 3e4) {
   const browser2 = await getBrowser();
@@ -142,20 +226,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((resolve) => {
-            img.addEventListener("load", resolve);
-            img.addEventListener("error", resolve);
-          });
-        })
-      );
-    });
+    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((resolve) => {
+              img.addEventListener("load", resolve);
+              img.addEventListener("error", resolve);
+              setTimeout(resolve, 3e3);
+            });
+          })
+        );
+      });
+    } catch {
+    }
     const contentBounds = await page.evaluate(() => {
       const root = document.querySelector("#storybook-root");
       if (!root) return null;
@@ -208,9 +299,9 @@ async function captureScreenshotBuffer(storyId, theme, width = VIEWPORT_WIDTH, h
 async function generateScreenshot(storyId, theme, storybookUrl = "http://localhost:6006", timeoutMs = 3e4) {
   try {
     ensureCacheDirectories();
-    const storyDir = path2.join(SCREENSHOTS_DIR, storyId);
-    if (!fs4.existsSync(storyDir)) {
-      fs4.mkdirSync(storyDir, { recursive: true });
+    const storyDir = path.join(SCREENSHOTS_DIR, storyId);
+    if (!fs5.existsSync(storyDir)) {
+      fs5.mkdirSync(storyDir, { recursive: true });
     }
     const screenshotPath = getScreenshotPath(storyId, theme);
     const { buffer, boundingBox } = await captureScreenshotBuffer(
@@ -221,7 +312,7 @@ async function generateScreenshot(storyId, theme, storybookUrl = "http://localho
       storybookUrl,
       timeoutMs
     );
-    fs4.writeFileSync(screenshotPath, buffer);
+    fs5.writeFileSync(screenshotPath, buffer);
     return { path: screenshotPath, boundingBox };
   } catch (error) {
     console.error(`Error generating screenshot for ${storyId} (${theme}):`, error);
@@ -248,7 +339,7 @@ async function fetchStorybookIndex() {
 }
 function getStoriesForFile(filePath) {
   if (!cachedIndex) return [];
-  const fileName = path2.basename(filePath);
+  const fileName = path.basename(filePath);
   return Object.values(cachedIndex.entries).filter((entry) => entry.type === "story" && entry.importPath.endsWith(fileName)).map((entry) => entry.id);
 }
 async function regenerateScreenshotsForFiles(files) {
@@ -335,10 +426,82 @@ var DEFAULT_ALLOWED_ORIGINS = [
   "http://localhost:3000",
   "http://localhost:6006"
 ];
+var AUTO_STORIES_FOLDER = ".onlook-stories";
+var storyRuntimeErrors = /* @__PURE__ */ new Map();
 var serveMetadataAndScreenshots = (req, res, next) => {
+  if (req.url === "/onbook-health.json") {
+    const cacheBuster = Date.now();
+    fetch(`http://localhost:6006/index.json?_t=${cacheBuster}`, {
+      cache: "no-store"
+    }).then((response) => response.json()).then((indexData) => {
+      const entries = indexData.entries || {};
+      const autoStories = [];
+      const userStories = [];
+      for (const [storyId, entry] of Object.entries(entries)) {
+        const importPath = entry.importPath || "";
+        if (importPath.includes(AUTO_STORIES_FOLDER)) {
+          autoStories.push(storyId);
+        } else {
+          userStories.push(storyId);
+        }
+      }
+      const healthy = autoStories.filter((id) => !storyRuntimeErrors.has(id));
+      const broken = autoStories.filter((id) => storyRuntimeErrors.has(id)).map((id) => ({
+        storyId: id,
+        // biome-ignore lint/style/noNonNullAssertion: filtered above
+        error: storyRuntimeErrors.get(id)
+      }));
+      res.setHeader("Content-Type", "application/json");
+      res.setHeader("Access-Control-Allow-Origin", "*");
+      res.setHeader("Cache-Control", "no-store, no-cache, must-revalidate");
+      res.end(
+        JSON.stringify({
+          autoStoryCount: autoStories.length,
+          userStoryCount: userStories.length,
+          healthy,
+          broken
+        })
+      );
+    }).catch((error) => {
+      res.statusCode = 500;
+      res.setHeader("Content-Type", "application/json");
+      res.end(
+        JSON.stringify({
+          error: "Failed to fetch story index",
+          details: String(error)
+        })
+      );
+    });
+    return;
+  }
+  if (req.url === "/onbook-report-error" && req.method === "POST") {
+    let body = "";
+    req.on("data", (chunk) => {
+      body += chunk.toString();
+    });
+    req.on("end", () => {
+      try {
+        const { storyId, error } = JSON.parse(body);
+        if (storyId && error) {
+          storyRuntimeErrors.set(storyId, error);
+          console.log(
+            `[STORYBOOK_PLUGIN] Story runtime error reported: ${storyId}`,
+            error
+          );
+        }
+        res.setHeader("Access-Control-Allow-Origin", "*");
+        res.statusCode = 200;
+        res.end("ok");
+      } catch {
+        res.statusCode = 400;
+        res.end("Invalid JSON");
+      }
+    });
+    return;
+  }
   if (req.url === "/onbook-index.json") {
     console.log("[STORYBOOK_PLUGIN] Serving /onbook-index.json endpoint");
-    const manifestPath = path2.join(process.cwd(), ".storybook-cache", "manifest.json");
+    const manifestPath = path.join(process.cwd(), ".storybook-cache", "manifest.json");
     const cacheBuster = Date.now();
     console.log("[STORYBOOK_PLUGIN] Fetching http://localhost:6006/index.json");
     fetch(`http://localhost:6006/index.json?_t=${cacheBuster}`, {
@@ -355,7 +518,7 @@ var serveMetadataAndScreenshots = (req, res, next) => {
       });
       return response.json();
     }).then((indexData) => {
-      const manifest = fs4.existsSync(manifestPath) ? JSON.parse(fs4.readFileSync(manifestPath, "utf-8")) : { stories: {} };
+      const manifest = fs5.existsSync(manifestPath) ? JSON.parse(fs5.readFileSync(manifestPath, "utf-8")) : { stories: {} };
       const defaultBoundingBox = { width: 1920, height: 1080 };
       for (const [storyId, entry] of Object.entries(indexData.entries || {})) {
         const manifestEntry = manifest.stories?.[storyId];
@@ -423,7 +586,7 @@ var serveMetadataAndScreenshots = (req, res, next) => {
     return;
   }
   if (req.url?.startsWith("/screenshots/")) {
-    const screenshotPath = path2.join(
+    const screenshotPath = path.join(
       process.cwd(),
       ".storybook-cache",
       req.url.replace("/screenshots/", "screenshots/")
@@ -432,11 +595,11 @@ var serveMetadataAndScreenshots = (req, res, next) => {
     const storyId = urlParts[0];
     const themeFile = urlParts[1];
     const theme = themeFile?.replace(".png", "");
-    if (fs4.existsSync(screenshotPath)) {
+    if (fs5.existsSync(screenshotPath)) {
       res.setHeader("Content-Type", "image/png");
       res.setHeader("Access-Control-Allow-Origin", "*");
       res.setHeader("Cache-Control", "public, max-age=3600");
-      fs4.createReadStream(screenshotPath).pipe(res);
+      fs5.createReadStream(screenshotPath).pipe(res);
       return;
     }
     if (storyId && theme && (theme === "light" || theme === "dark")) {
@@ -444,16 +607,16 @@ var serveMetadataAndScreenshots = (req, res, next) => {
         `[STORYBOOK_PLUGIN] Generating screenshot on-demand: ${storyId}/${theme}`
       );
       captureScreenshotBuffer(storyId, theme).then(({ buffer }) => {
-        const storyDir = path2.join(
+        const storyDir = path.join(
           process.cwd(),
           ".storybook-cache",
           "screenshots",
           storyId
         );
-        if (!fs4.existsSync(storyDir)) {
-          fs4.mkdirSync(storyDir, { recursive: true });
+        if (!fs5.existsSync(storyDir)) {
+          fs5.mkdirSync(storyDir, { recursive: true });
         }
-        fs4.writeFileSync(screenshotPath, buffer);
+        fs5.writeFileSync(screenshotPath, buffer);
         res.setHeader("Content-Type", "image/png");
         res.setHeader("Access-Control-Allow-Origin", "*");
         res.setHeader("Cache-Control", "public, max-age=3600");
@@ -496,17 +659,16 @@ function storybookOnlookPlugin(options = {}) {
   const mainPlugin = {
     name: "storybook-onlook-plugin",
     config() {
-      const isE2B = !!process.env.E2B_SANDBOX;
       return {
         server: {
-          // Only override HMR for E2B sandboxes (wss + port 443)
-          // Local dev uses Vite's default HMR which works out of the box
-          ...isE2B && {
-            hmr: {
-              protocol: "wss",
-              clientPort: 443,
-              port
-            }
+          // E2B sandbox HMR configuration
+          hmr: {
+            // E2B sandboxes use HTTPS, so we need secure WebSocket
+            protocol: "wss",
+            // E2B routes through standard HTTPS port 443
+            clientPort: 443,
+            // The actual Storybook server port inside the sandbox
+            port
           },
           cors: {
             origin: allowedOrigins
@@ -528,15 +690,49 @@ function storybookOnlookPlugin(options = {}) {
       console.log("[STORYBOOK_PLUGIN] Configuring preview server middleware");
       server.middlewares.use(serveMetadataAndScreenshots);
     },
-    // Only auto-regenerate screenshots in E2B sandboxes where Onbook
-    // consumes them. In local dev this blocks HMR with 36 Playwright
-    // page loads per save. Screenshots are still available on-demand
-    // via /api/capture-screenshot and the CLI.
-    ...process.env.E2B_SANDBOX && {
-      handleHotUpdate: handleStoryFileChange
+    handleHotUpdate(ctx) {
+      if (ctx.file.includes(AUTO_STORIES_FOLDER) && ctx.file.endsWith(".stories.tsx")) {
+        enrichStoryFile(ctx.file);
+      }
+      return handleStoryFileChange(ctx);
     }
   };
-  return [componentLocPlugin(), mainPlugin];
+  const plugins = [componentLocPlugin(), mainPlugin];
+  if (options.autoStories !== false) {
+    const imports = options.autoStories ?? ["src/**/*.tsx"];
+    const ignores = options.autoStoriesIgnore ?? [
+      "src/**/*.stories.tsx",
+      "src/**/*.stories.ts",
+      "src/**/*.test.tsx",
+      "src/**/*.test.ts",
+      "src/**/*.spec.tsx",
+      "src/**/*.spec.ts",
+      "node_modules/**",
+      "**/.onlook-stories/**"
+    ];
+    console.log("[STORYBOOK_PLUGIN] Auto-story generation enabled", {
+      imports,
+      ignores,
+      storiesFolder: AUTO_STORIES_FOLDER
+    });
+    try {
+      plugins.push(
+        autoStoryGenerator.vite({
+          preset: "react",
+          imports,
+          ignores,
+          storiesFolder: AUTO_STORIES_FOLDER,
+          isGenerateStoriesFileAtBuild: true
+        })
+      );
+    } catch (err) {
+      console.error(
+        "[STORYBOOK_PLUGIN] ASG plugin failed to initialize, continuing without auto-stories",
+        err
+      );
+    }
+  }
+  return plugins;
 }
 
-export { storybookOnlookPlugin };
+export { AUTO_STORIES_FOLDER, storybookOnlookPlugin };

package/dist/screenshot-service/index.js

@@ -93,20 +93,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((resolve) => {
-            img.addEventListener("load", resolve);
-            img.addEventListener("error", resolve);
-          });
-        })
-      );
-    });
+    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((resolve) => {
+              img.addEventListener("load", resolve);
+              img.addEventListener("error", resolve);
+              setTimeout(resolve, 3e3);
+            });
+          })
+        );
+      });
+    } catch {
+    }
     const contentBounds = await page.evaluate(() => {
       const root = document.querySelector("#storybook-root");
       if (!root) return null;
@@ -197,23 +204,48 @@ async function generateAllScreenshots(stories, storybookUrl = "http://localhost:
       batch.map(async (story) => {
         if (skipExisting && screenshotExists(story.id, "light") && screenshotExists(story.id, "dark")) {
           completed++;
-          const absoluteIndex2 = offset + completed;
-          console.log(`[${absoluteIndex2}/${displayTotal}] Skipped (exists) ${story.id}`);
+          const absoluteIndex = offset + completed;
+          console.log(`[${absoluteIndex}/${displayTotal}] Skipped (exists) ${story.id}`);
           return;
         }
-        const [lightResult, darkResult] = await Promise.all([
-          generateScreenshot(story.id, "light", storybookUrl, timeoutMs),
-          generateScreenshot(story.id, "dark", storybookUrl, timeoutMs)
-        ]);
-        if (lightResult && darkResult) {
-          const fileHash = computeFileHash(story.importPath);
-          updateManifest(story.id, story.importPath, fileHash, lightResult.boundingBox);
+        const storyTimeout = timeoutMs * 2 + 1e4;
+        let timer;
+        try {
+          const result = await Promise.race([
+            Promise.all([
+              generateScreenshot(story.id, "light", storybookUrl, timeoutMs),
+              generateScreenshot(story.id, "dark", storybookUrl, timeoutMs)
+            ]),
+            new Promise((_, reject) => {
+              timer = setTimeout(
+                () => reject(
+                  new Error(
+                    `Story ${story.id} timed out after ${storyTimeout / 1e3}s`
+                  )
+                ),
+                storyTimeout
+              );
+            })
+          ]);
+          clearTimeout(timer);
+          const [lightResult, darkResult] = result;
+          if (lightResult && darkResult) {
+            const fileHash = computeFileHash(story.importPath);
+            updateManifest(story.id, story.importPath, fileHash, lightResult.boundingBox);
+          }
+          completed++;
+          const absoluteIndex = offset + completed;
+          console.log(
+            `[${absoluteIndex}/${displayTotal}] Generated screenshots for ${story.id}`
+          );
+        } catch (error) {
+          completed++;
+          const absoluteIndex = offset + completed;
+          console.error(
+            `[${absoluteIndex}/${displayTotal}] \u26A0\uFE0F Failed ${story.id}:`,
+            error instanceof Error ? error.message : error
+          );
         }
-        completed++;
-        const absoluteIndex = offset + completed;
-        console.log(
-          `[${absoluteIndex}/${displayTotal}] Generated screenshots for ${story.id}`
-        );
       })
     );
   }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@onlook/storybook-plugin",
-  "version": "0.3.2",
+  "version": "0.4.0-beta.0",
   "type": "module",
   "bin": {
     "onlook-storybook": "./dist/cli/index.js"
@@ -37,7 +37,9 @@
     "@babel/traverse": "^7.26.9",
     "@babel/types": "^7.26.9",
     "@drizzle-team/brocli": "^0.11.0",
-    "playwright": "^1.52.0"
+    "@takuma-ru/auto-story-generator": "^0.4.0",
+    "playwright": "^1.52.0",
+    "react-docgen-typescript": "^2.4.0"
   },
   "devDependencies": {
     "@onbook/tsconfig": "workspace:*",