storybook-onbook-plugin 0.2.2 → 0.2.4

package/dist/cli/index.d.ts

@@ -0,0 +1 @@
+#!/usr/bin/env node

package/dist/cli/index.js

@@ -0,0 +1,401 @@
+#!/usr/bin/env node
+import { command, string, run } from '@drizzle-team/brocli';
+import { spawn } from 'child_process';
+import path, { resolve, join, dirname } from 'path';
+import crypto from 'crypto';
+import fs, { existsSync, readFileSync } from 'fs';
+import { chromium } from 'playwright';
+
+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;
+var MIN_COMPONENT_HEIGHT = 280;
+
+// src/utils/fileSystem/fileSystem.ts
+function ensureCacheDirectories() {
+  if (!fs.existsSync(CACHE_DIR)) {
+    fs.mkdirSync(CACHE_DIR, { recursive: true });
+  }
+  if (!fs.existsSync(SCREENSHOTS_DIR)) {
+    fs.mkdirSync(SCREENSHOTS_DIR, { recursive: true });
+  }
+}
+function computeFileHash(filePath) {
+  if (!fs.existsSync(filePath)) {
+    return "";
+  }
+  const content = fs.readFileSync(filePath, "utf-8");
+  return crypto.createHash("sha256").update(content).digest("hex");
+}
+function loadManifest() {
+  if (fs.existsSync(MANIFEST_PATH)) {
+    const content = fs.readFileSync(MANIFEST_PATH, "utf-8");
+    return JSON.parse(content);
+  }
+  return { stories: {} };
+}
+function saveManifest(manifest) {
+  ensureCacheDirectories();
+  fs.writeFileSync(MANIFEST_PATH, JSON.stringify(manifest, null, 2));
+}
+function updateManifest(storyId, sourcePath, fileHash, boundingBox) {
+  const manifest = loadManifest();
+  manifest.stories[storyId] = {
+    fileHash,
+    lastGenerated: (/* @__PURE__ */ new Date()).toISOString(),
+    sourcePath,
+    screenshots: {
+      light: `screenshots/${storyId}/light.png`,
+      dark: `screenshots/${storyId}/dark.png`
+    },
+    boundingBox: boundingBox ?? void 0
+  };
+  saveManifest(manifest);
+}
+var browser = null;
+async function getBrowser() {
+  if (!browser) {
+    browser = await chromium.launch({
+      headless: true
+    });
+  }
+  return browser;
+}
+async function closeBrowser() {
+  if (browser) {
+    try {
+      await browser.close();
+    } catch (error) {
+      console.error("Error closing browser:", error);
+    } finally {
+      browser = null;
+    }
+  }
+}
+function getScreenshotPath(storyId, theme) {
+  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") {
+  const browser2 = await getBrowser();
+  const context = await browser2.newContext({
+    viewport: { width, height },
+    deviceScaleFactor: 2
+  });
+  const page = await context.newPage();
+  try {
+    const url = `${storybookUrl}/iframe.html?id=${storyId}&viewMode=story&globals=theme:${theme}`;
+    await page.goto(url, { timeout: 15e3 });
+    await page.waitForLoadState("domcontentloaded");
+    await page.waitForLoadState("load");
+    await page.waitForLoadState("networkidle");
+    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);
+          });
+        })
+      );
+    });
+    const contentBounds = await page.evaluate(() => {
+      const root = document.querySelector("#storybook-root");
+      if (!root) return null;
+      const children = root.querySelectorAll("*");
+      if (children.length === 0) return null;
+      let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;
+      children.forEach((child) => {
+        const rect = child.getBoundingClientRect();
+        if (rect.width === 0 || rect.height === 0) return;
+        minX = Math.min(minX, rect.left);
+        minY = Math.min(minY, rect.top);
+        maxX = Math.max(maxX, rect.right);
+        maxY = Math.max(maxY, rect.bottom);
+      });
+      if (minX === Infinity) return null;
+      return {
+        x: minX,
+        y: minY,
+        width: maxX - minX,
+        height: maxY - minY
+      };
+    });
+    let screenshotBuffer;
+    let resultBoundingBox = null;
+    if (contentBounds && contentBounds.width > 0 && contentBounds.height > 0) {
+      const PADDING = 20;
+      const clippedWidth = Math.min(width, contentBounds.width + PADDING);
+      const clippedHeight = Math.min(height, contentBounds.height + PADDING);
+      resultBoundingBox = {
+        width: Math.max(MIN_COMPONENT_WIDTH, Math.round(clippedWidth)),
+        height: Math.max(MIN_COMPONENT_HEIGHT, Math.round(clippedHeight))
+      };
+      screenshotBuffer = await page.screenshot({
+        type: "png",
+        clip: {
+          x: Math.max(0, contentBounds.x - 10),
+          y: Math.max(0, contentBounds.y - 10),
+          width: clippedWidth,
+          height: clippedHeight
+        }
+      });
+    } else {
+      screenshotBuffer = await page.screenshot({ type: "png" });
+    }
+    return { buffer: screenshotBuffer, boundingBox: resultBoundingBox };
+  } finally {
+    await context.close();
+  }
+}
+async function generateScreenshot(storyId, theme, storybookUrl = "http://localhost:6006") {
+  try {
+    ensureCacheDirectories();
+    const storyDir = path.join(SCREENSHOTS_DIR, storyId);
+    if (!fs.existsSync(storyDir)) {
+      fs.mkdirSync(storyDir, { recursive: true });
+    }
+    const screenshotPath = getScreenshotPath(storyId, theme);
+    const { buffer, boundingBox } = await captureScreenshotBuffer(
+      storyId,
+      theme,
+      VIEWPORT_WIDTH,
+      VIEWPORT_HEIGHT,
+      storybookUrl
+    );
+    fs.writeFileSync(screenshotPath, buffer);
+    return { path: screenshotPath, boundingBox };
+  } catch (error) {
+    console.error(`Error generating screenshot for ${storyId} (${theme}):`, error);
+    return null;
+  }
+}
+
+// src/screenshot-service/screenshot-service.ts
+async function generateAllScreenshots(stories, storybookUrl = "http://localhost:6006") {
+  console.log(`Generating screenshots for ${stories.length} stories...`);
+  const BATCH_SIZE = 10;
+  const batches = [];
+  for (let i = 0; i < stories.length; i += BATCH_SIZE) {
+    batches.push(stories.slice(i, i + BATCH_SIZE));
+  }
+  let completed = 0;
+  for (const batch of batches) {
+    await Promise.all(
+      batch.map(async (story) => {
+        const [lightResult, darkResult] = await Promise.all([
+          generateScreenshot(story.id, "light", storybookUrl),
+          generateScreenshot(story.id, "dark", storybookUrl)
+        ]);
+        if (lightResult && darkResult) {
+          const fileHash = computeFileHash(story.importPath);
+          updateManifest(story.id, story.importPath, fileHash, lightResult.boundingBox);
+        }
+        completed++;
+        console.log(
+          `[${completed}/${stories.length}] Generated screenshots for ${story.id}`
+        );
+      })
+    );
+  }
+  await closeBrowser();
+  console.log("Screenshot generation complete!");
+}
+var LOCKFILES = {
+  "bun.lockb": "bun",
+  "bun.lock": "bun",
+  "pnpm-lock.yaml": "pnpm",
+  "yarn.lock": "yarn",
+  "package-lock.json": "npm"
+};
+function findRepoRoot(startDir) {
+  let currentDir = resolve(startDir);
+  while (true) {
+    for (const lockfile of Object.keys(LOCKFILES)) {
+      if (existsSync(join(currentDir, lockfile))) {
+        return currentDir;
+      }
+    }
+    const parentDir = dirname(currentDir);
+    if (parentDir === currentDir) {
+      return null;
+    }
+    currentDir = parentDir;
+  }
+}
+function detectPackageManager(repoRoot) {
+  for (const [lockfile, pm] of Object.entries(LOCKFILES)) {
+    if (existsSync(join(repoRoot, lockfile))) {
+      return pm;
+    }
+  }
+  return "npm";
+}
+function getStorybookScript(storybookDir) {
+  const packageJsonPath = join(storybookDir, "package.json");
+  if (!existsSync(packageJsonPath)) {
+    return null;
+  }
+  const packageJson = JSON.parse(readFileSync(packageJsonPath, "utf-8"));
+  return packageJson.scripts?.storybook ?? null;
+}
+
+// src/cli/generate-screenshots/generate-screenshots.ts
+function getStorybookCommand(storybookDir) {
+  const script = getStorybookScript(storybookDir);
+  if (script === null) {
+    const packageJsonPath = join(storybookDir, "package.json");
+    throw new Error(
+      `No "storybook" script found in ${packageJsonPath}. Use --storybook-dir to specify the correct directory or add a storybook script.`
+    );
+  }
+  const repoRoot = findRepoRoot(storybookDir);
+  const pm = repoRoot ? detectPackageManager(repoRoot) : "npm";
+  console.log(
+    `\u{1F4E6} Detected package manager: ${pm}${repoRoot ? ` (from ${repoRoot})` : ""}`
+  );
+  const execCommand = {
+    bun: "bunx",
+    pnpm: "pnpm exec",
+    yarn: "yarn",
+    npm: "npx"
+  };
+  const scriptParts = script.split(" ").filter(Boolean);
+  return {
+    command: execCommand[pm] ?? "npx",
+    args: [...scriptParts, "--no-open"]
+  };
+}
+async function isStorybookRunning(url) {
+  try {
+    const response = await fetch(`${url}/index.json`, {
+      signal: AbortSignal.timeout(2e3)
+    });
+    return response.ok;
+  } catch {
+    return false;
+  }
+}
+async function fetchStoryIndex(url) {
+  const indexUrl = `${url}/index.json`;
+  const response = await fetch(indexUrl);
+  if (!response.ok) {
+    throw new Error(`Failed to fetch story index: ${response.statusText}`);
+  }
+  const data = await response.json();
+  return Object.values(data.entries);
+}
+async function startStorybook(command2, args, storybookDir) {
+  console.log(`\u{1F680} Starting Storybook with: ${command2} ${args.join(" ")} (in ${storybookDir})`);
+  const storybookProcess = spawn(command2, args, {
+    cwd: storybookDir,
+    stdio: "pipe",
+    shell: true
+  });
+  return new Promise((resolve3, reject) => {
+    let started = false;
+    const timeout = setTimeout(() => {
+      if (!started) {
+        storybookProcess.kill();
+        reject(new Error("Storybook failed to start within 60 seconds"));
+      }
+    }, 6e4);
+    storybookProcess.stdout?.on("data", (data) => {
+      const output = data.toString();
+      if (output.includes("Local:") || output.includes("localhost:")) {
+        if (!started) {
+          started = true;
+          clearTimeout(timeout);
+          console.log("\u2705 Storybook is ready!");
+          resolve3(storybookProcess);
+        }
+      }
+    });
+    storybookProcess.stderr?.on("data", (data) => {
+      const output = data.toString();
+      if (output.toLowerCase().includes("error")) {
+        console.error(output);
+      }
+    });
+    storybookProcess.on("error", (error) => {
+      clearTimeout(timeout);
+      reject(error);
+    });
+  });
+}
+async function warmupStorybook(url, firstStoryId) {
+  console.log("\u{1F525} Warming up Storybook...");
+  const browser2 = await getBrowser();
+  const context = await browser2.newContext();
+  const page = await context.newPage();
+  try {
+    const warmupUrl = `${url}/iframe.html?id=${firstStoryId}&viewMode=story`;
+    await page.goto(warmupUrl, { timeout: 15e3 });
+    await page.waitForLoadState("networkidle");
+    console.log("\u2705 Storybook warmed up");
+  } catch {
+    console.log("\u26A0\uFE0F Warmup had issues, proceeding anyway");
+  } finally {
+    await context.close();
+  }
+}
+async function generateScreenshots(options = {}) {
+  const url = "http://localhost:6006";
+  const storybookDir = resolve(options.storybookDir ?? process.cwd());
+  let storybookProcess = null;
+  let weStartedStorybook = false;
+  console.log("\u{1F4F8} Generating Storybook screenshots...");
+  console.log(`\u{1F4C2} Storybook directory: ${storybookDir}`);
+  try {
+    const alreadyRunning = await isStorybookRunning(url);
+    if (alreadyRunning) {
+      console.log("\u2705 Storybook is already running");
+    } else {
+      const { command: command2, args } = getStorybookCommand(storybookDir);
+      storybookProcess = await startStorybook(command2, args, storybookDir);
+      weStartedStorybook = true;
+    }
+    const stories = await fetchStoryIndex(url);
+    console.log(`Found ${stories.length} stories`);
+    const firstStory = stories[0];
+    if (!firstStory) {
+      throw new Error("No stories found");
+    }
+    await warmupStorybook(url, firstStory.id);
+    await generateAllScreenshots(
+      stories.map((story) => ({
+        id: story.id,
+        importPath: story.importPath
+      })),
+      url
+    );
+    console.log("\u2705 Screenshot generation complete!");
+  } finally {
+    if (storybookProcess && weStartedStorybook) {
+      console.log("\u{1F6D1} Stopping Storybook...");
+      storybookProcess.kill();
+    }
+  }
+}
+
+// src/cli/index.ts
+var generateScreenshotsCommand = command({
+  name: "generate-screenshots",
+  desc: "Generate screenshots for all Storybook stories",
+  options: {
+    storybookDir: string().alias("d").desc("Directory containing Storybook (defaults to current directory)")
+  },
+  handler: async (opts) => {
+    await generateScreenshots({ storybookDir: opts.storybookDir });
+  }
+});
+run([generateScreenshotsCommand], {
+  name: "storybook-onbook-plugin",
+  description: "Storybook plugin for Onbook - generate screenshots and more"
+});

package/dist/index.d.ts

@@ -0,0 +1,11 @@
+import { PluginOption } from 'vite';
+
+type OnlookPluginOptions = {
+    /** Storybook port (default: 6006) */
+    port?: number;
+    /** Additional allowed origins for CORS (merged with defaults) */
+    allowedOrigins?: string[];
+};
+declare function storybookOnlookPlugin(options?: OnlookPluginOptions): PluginOption[];
+
+export { type OnlookPluginOptions, storybookOnlookPlugin };

package/dist/index.js

@@ -0,0 +1,453 @@
+import fs, { existsSync } from 'fs';
+import path2, { dirname, join, relative } from 'path';
+import { fileURLToPath } from 'url';
+import generateModule from '@babel/generator';
+import { parse } from '@babel/parser';
+import traverseModule from '@babel/traverse';
+import * as t from '@babel/types';
+import crypto from 'crypto';
+import { chromium } from 'playwright';
+
+// src/storybook-onlook-plugin.ts
+function componentLocPlugin(options = {}) {
+  const include = options.include ?? /\.(jsx|tsx)$/;
+  const traverse = traverseModule.default ?? traverseModule;
+  const generate = generateModule.default ?? generateModule;
+  let root;
+  return {
+    name: "onbook-component-loc",
+    enforce: "pre",
+    apply: "serve",
+    configResolved(config) {
+      root = config.root;
+    },
+    transform(code, id) {
+      const filepath = id.split("?", 1)[0];
+      if (!filepath || filepath.includes("node_modules")) return null;
+      if (!include.test(filepath)) return null;
+      const ast = parse(code, {
+        sourceType: "module",
+        plugins: ["jsx", "typescript"],
+        sourceFilename: filepath
+      });
+      let mutated = false;
+      const relativePath = path2.relative(root, filepath);
+      traverse(ast, {
+        JSXElement(nodePath) {
+          const opening = nodePath.node.openingElement;
+          const element = nodePath.node;
+          if (!opening.loc || !element.loc) return;
+          const alreadyTagged = opening.attributes.some(
+            (attribute) => t.isJSXAttribute(attribute) && attribute.name.name === "data-component-file"
+          );
+          if (alreadyTagged) return;
+          opening.attributes.push(
+            t.jsxAttribute(
+              t.jsxIdentifier("data-component-file"),
+              t.stringLiteral(relativePath)
+            ),
+            t.jsxAttribute(
+              t.jsxIdentifier("data-component-start"),
+              t.stringLiteral(`${element.loc.start.line}:${element.loc.start.column}`)
+            ),
+            t.jsxAttribute(
+              t.jsxIdentifier("data-component-end"),
+              t.stringLiteral(`${element.loc.end.line}:${element.loc.end.column}`)
+            )
+          );
+          mutated = true;
+        }
+      });
+      if (!mutated) return null;
+      const output = generate(
+        ast,
+        { retainLines: true, sourceMaps: true, sourceFileName: id },
+        code
+      );
+      return { code: output.code, map: output.map };
+    }
+  };
+}
+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 VIEWPORT_WIDTH = 1920;
+var VIEWPORT_HEIGHT = 1080;
+var MIN_COMPONENT_WIDTH = 420;
+var MIN_COMPONENT_HEIGHT = 280;
+
+// src/utils/fileSystem/fileSystem.ts
+function ensureCacheDirectories() {
+  if (!fs.existsSync(CACHE_DIR)) {
+    fs.mkdirSync(CACHE_DIR, { recursive: true });
+  }
+  if (!fs.existsSync(SCREENSHOTS_DIR)) {
+    fs.mkdirSync(SCREENSHOTS_DIR, { recursive: true });
+  }
+}
+function computeFileHash(filePath) {
+  if (!fs.existsSync(filePath)) {
+    return "";
+  }
+  const content = fs.readFileSync(filePath, "utf-8");
+  return crypto.createHash("sha256").update(content).digest("hex");
+}
+function loadManifest() {
+  if (fs.existsSync(MANIFEST_PATH)) {
+    const content = fs.readFileSync(MANIFEST_PATH, "utf-8");
+    return JSON.parse(content);
+  }
+  return { stories: {} };
+}
+function saveManifest(manifest) {
+  ensureCacheDirectories();
+  fs.writeFileSync(MANIFEST_PATH, JSON.stringify(manifest, null, 2));
+}
+function updateManifest(storyId, sourcePath, fileHash, boundingBox) {
+  const manifest = loadManifest();
+  manifest.stories[storyId] = {
+    fileHash,
+    lastGenerated: (/* @__PURE__ */ new Date()).toISOString(),
+    sourcePath,
+    screenshots: {
+      light: `screenshots/${storyId}/light.png`,
+      dark: `screenshots/${storyId}/dark.png`
+    },
+    boundingBox: boundingBox ?? void 0
+  };
+  saveManifest(manifest);
+}
+var browser = null;
+async function getBrowser() {
+  if (!browser) {
+    browser = await chromium.launch({
+      headless: true
+    });
+  }
+  return browser;
+}
+function getScreenshotPath(storyId, theme) {
+  const storyDir = path2.join(SCREENSHOTS_DIR, storyId);
+  return path2.join(storyDir, `${theme}.png`);
+}
+async function captureScreenshotBuffer(storyId, theme, width = VIEWPORT_WIDTH, height = VIEWPORT_HEIGHT, storybookUrl = "http://localhost:6006") {
+  const browser2 = await getBrowser();
+  const context = await browser2.newContext({
+    viewport: { width, height },
+    deviceScaleFactor: 2
+  });
+  const page = await context.newPage();
+  try {
+    const url = `${storybookUrl}/iframe.html?id=${storyId}&viewMode=story&globals=theme:${theme}`;
+    await page.goto(url, { timeout: 15e3 });
+    await page.waitForLoadState("domcontentloaded");
+    await page.waitForLoadState("load");
+    await page.waitForLoadState("networkidle");
+    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);
+          });
+        })
+      );
+    });
+    const contentBounds = await page.evaluate(() => {
+      const root = document.querySelector("#storybook-root");
+      if (!root) return null;
+      const children = root.querySelectorAll("*");
+      if (children.length === 0) return null;
+      let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;
+      children.forEach((child) => {
+        const rect = child.getBoundingClientRect();
+        if (rect.width === 0 || rect.height === 0) return;
+        minX = Math.min(minX, rect.left);
+        minY = Math.min(minY, rect.top);
+        maxX = Math.max(maxX, rect.right);
+        maxY = Math.max(maxY, rect.bottom);
+      });
+      if (minX === Infinity) return null;
+      return {
+        x: minX,
+        y: minY,
+        width: maxX - minX,
+        height: maxY - minY
+      };
+    });
+    let screenshotBuffer;
+    let resultBoundingBox = null;
+    if (contentBounds && contentBounds.width > 0 && contentBounds.height > 0) {
+      const PADDING = 20;
+      const clippedWidth = Math.min(width, contentBounds.width + PADDING);
+      const clippedHeight = Math.min(height, contentBounds.height + PADDING);
+      resultBoundingBox = {
+        width: Math.max(MIN_COMPONENT_WIDTH, Math.round(clippedWidth)),
+        height: Math.max(MIN_COMPONENT_HEIGHT, Math.round(clippedHeight))
+      };
+      screenshotBuffer = await page.screenshot({
+        type: "png",
+        clip: {
+          x: Math.max(0, contentBounds.x - 10),
+          y: Math.max(0, contentBounds.y - 10),
+          width: clippedWidth,
+          height: clippedHeight
+        }
+      });
+    } else {
+      screenshotBuffer = await page.screenshot({ type: "png" });
+    }
+    return { buffer: screenshotBuffer, boundingBox: resultBoundingBox };
+  } finally {
+    await context.close();
+  }
+}
+async function generateScreenshot(storyId, theme, storybookUrl = "http://localhost:6006") {
+  try {
+    ensureCacheDirectories();
+    const storyDir = path2.join(SCREENSHOTS_DIR, storyId);
+    if (!fs.existsSync(storyDir)) {
+      fs.mkdirSync(storyDir, { recursive: true });
+    }
+    const screenshotPath = getScreenshotPath(storyId, theme);
+    const { buffer, boundingBox } = await captureScreenshotBuffer(
+      storyId,
+      theme,
+      VIEWPORT_WIDTH,
+      VIEWPORT_HEIGHT,
+      storybookUrl
+    );
+    fs.writeFileSync(screenshotPath, buffer);
+    return { path: screenshotPath, boundingBox };
+  } catch (error) {
+    console.error(`Error generating screenshot for ${storyId} (${theme}):`, error);
+    return null;
+  }
+}
+
+// src/handlers/handleStoryFileChange/handleStoryFileChange.ts
+var cachedIndex = null;
+var indexFetchPromise = null;
+var pendingFiles = /* @__PURE__ */ new Set();
+var debounceTimer = null;
+var DEBOUNCE_MS = 500;
+async function fetchStorybookIndex() {
+  try {
+    const response = await fetch("http://localhost:6006/index.json");
+    if (response.ok) {
+      cachedIndex = await response.json();
+      console.log("[Screenshots] Cached Storybook index");
+    }
+  } catch (error) {
+    console.error("[Screenshots] Failed to fetch Storybook index:", error);
+  }
+}
+function getStoriesForFile(filePath) {
+  if (!cachedIndex) return [];
+  const fileName = path2.basename(filePath);
+  return Object.values(cachedIndex.entries).filter((entry) => entry.type === "story" && entry.importPath.endsWith(fileName)).map((entry) => entry.id);
+}
+async function regenerateScreenshotsForFiles(files) {
+  await fetchStorybookIndex();
+  const allStoryIds = /* @__PURE__ */ new Set();
+  const fileToStories = /* @__PURE__ */ new Map();
+  for (const file of files) {
+    const storyIds = getStoriesForFile(file);
+    fileToStories.set(file, storyIds);
+    for (const id of storyIds) {
+      allStoryIds.add(id);
+    }
+  }
+  if (allStoryIds.size === 0) {
+    console.log("[Screenshots] No stories found for changed files");
+    return;
+  }
+  console.log(
+    `[Screenshots] Regenerating ${allStoryIds.size} stories from ${files.length} files`
+  );
+  const storybookUrl = "http://localhost:6006";
+  const storyBoundingBoxes = /* @__PURE__ */ new Map();
+  await Promise.all(
+    Array.from(allStoryIds).map(async (storyId) => {
+      const [lightResult, _darkResult] = await Promise.all([
+        generateScreenshot(storyId, "light", storybookUrl),
+        generateScreenshot(storyId, "dark", storybookUrl)
+      ]);
+      if (lightResult) {
+        storyBoundingBoxes.set(storyId, lightResult.boundingBox);
+      }
+    })
+  );
+  for (const [file, storyIds] of fileToStories) {
+    const fileHash = computeFileHash(file);
+    storyIds.forEach((storyId) => {
+      updateManifest(storyId, file, fileHash, storyBoundingBoxes.get(storyId));
+    });
+  }
+  console.log(`[Screenshots] \u2713 Regenerated ${allStoryIds.size} stories`);
+}
+function handleStoryFileChange({ file, modules }) {
+  if (file.endsWith(".stories.tsx") || file.endsWith(".stories.ts")) {
+    console.log(`[Screenshots] Story file changed: ${file}`);
+    if (!cachedIndex && !indexFetchPromise) {
+      indexFetchPromise = fetchStorybookIndex();
+    }
+    pendingFiles.add(file);
+    if (debounceTimer) {
+      clearTimeout(debounceTimer);
+    }
+    debounceTimer = setTimeout(async () => {
+      const files = Array.from(pendingFiles);
+      pendingFiles.clear();
+      debounceTimer = null;
+      try {
+        await regenerateScreenshotsForFiles(files);
+      } catch (error) {
+        console.error("[Screenshots] Error regenerating screenshots:", error);
+      }
+    }, DEBOUNCE_MS);
+    return modules;
+  }
+}
+function findGitRoot(startPath) {
+  let currentPath = startPath;
+  while (currentPath !== dirname(currentPath)) {
+    if (existsSync(join(currentPath, ".git"))) {
+      return currentPath;
+    }
+    currentPath = dirname(currentPath);
+  }
+  return null;
+}
+
+// src/storybook-onlook-plugin.ts
+var __dirname$1 = dirname(fileURLToPath(import.meta.url));
+var storybookDir = join(__dirname$1, "..");
+var gitRoot = findGitRoot(storybookDir);
+var storybookLocation = gitRoot ? relative(gitRoot, storybookDir) : "";
+var repoRoot = gitRoot || process.cwd();
+var DEFAULT_ALLOWED_ORIGINS = [
+  "https://app.onlook.ai",
+  "http://localhost:3000",
+  "http://localhost:6006"
+];
+var serveMetadataAndScreenshots = (req, res, next) => {
+  if (req.url === "/onbook-index.json") {
+    const manifestPath = path2.join(process.cwd(), ".storybook-cache", "manifest.json");
+    fetch("http://localhost:6006/index.json").then((response) => response.json()).then((indexData) => {
+      const manifest = fs.existsSync(manifestPath) ? JSON.parse(fs.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];
+        entry.boundingBox = manifestEntry?.boundingBox || defaultBoundingBox;
+      }
+      indexData.meta = { storybookLocation, repoRoot };
+      res.setHeader("Content-Type", "application/json");
+      res.setHeader("Access-Control-Allow-Origin", "*");
+      res.end(JSON.stringify(indexData));
+    }).catch((error) => {
+      console.error("Failed to fetch/extend index.json:", error);
+      res.statusCode = 500;
+      res.setHeader("Content-Type", "application/json");
+      res.end(
+        JSON.stringify({ error: "Failed to fetch index", details: String(error) })
+      );
+    });
+    return;
+  }
+  if (req.url?.startsWith("/api/capture-screenshot")) {
+    const url = new URL(req.url, "http://localhost");
+    const storyId = url.searchParams.get("storyId");
+    const theme = url.searchParams.get("theme") || "light";
+    const width = parseInt(url.searchParams.get("width") || "800", 10);
+    const height = parseInt(url.searchParams.get("height") || "600", 10);
+    if (!storyId) {
+      res.statusCode = 400;
+      res.setHeader("Content-Type", "application/json");
+      res.end(JSON.stringify({ error: "storyId is required" }));
+      return;
+    }
+    if (theme !== "light" && theme !== "dark") {
+      res.statusCode = 400;
+      res.setHeader("Content-Type", "application/json");
+      res.end(JSON.stringify({ error: 'theme must be "light" or "dark"' }));
+      return;
+    }
+    captureScreenshotBuffer(storyId, theme, width, height).then(({ buffer }) => {
+      res.setHeader("Content-Type", "image/png");
+      res.setHeader("Access-Control-Allow-Origin", "*");
+      res.setHeader("Cache-Control", "no-cache");
+      res.end(buffer);
+    }).catch((error) => {
+      console.error("Screenshot capture error:", error);
+      res.statusCode = 500;
+      res.setHeader("Content-Type", "application/json");
+      res.end(
+        JSON.stringify({
+          error: "Failed to capture screenshot",
+          details: String(error)
+        })
+      );
+    });
+    return;
+  }
+  if (req.url?.startsWith("/screenshots/")) {
+    const screenshotPath = path2.join(
+      process.cwd(),
+      ".storybook-cache",
+      req.url.replace("/screenshots/", "screenshots/")
+    );
+    if (fs.existsSync(screenshotPath)) {
+      res.setHeader("Content-Type", "image/png");
+      res.setHeader("Access-Control-Allow-Origin", "*");
+      res.setHeader("Cache-Control", "public, max-age=3600");
+      fs.createReadStream(screenshotPath).pipe(res);
+    } else {
+      res.statusCode = 404;
+      res.end("Screenshot not found");
+    }
+    return;
+  }
+  next();
+};
+function storybookOnlookPlugin(options = {}) {
+  if (process.env.CHROMATIC || process.env.CI) {
+    return [];
+  }
+  const port = options.port ?? 6006;
+  const allowedOrigins = [...DEFAULT_ALLOWED_ORIGINS, ...options.allowedOrigins ?? []];
+  const mainPlugin = {
+    name: "storybook-onlook-plugin",
+    config() {
+      return {
+        server: {
+          // 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
+          }
+        }
+      };
+    },
+    configureServer(server) {
+      server.middlewares.use(serveMetadataAndScreenshots);
+    },
+    configurePreviewServer(server) {
+      server.middlewares.use(serveMetadataAndScreenshots);
+    },
+    handleHotUpdate: handleStoryFileChange
+  };
+  return [componentLocPlugin(), mainPlugin];
+}
+
+export { storybookOnlookPlugin };

package/dist/screenshot-service/index.d.ts

@@ -0,0 +1,47 @@
+export { closeBrowser } from './utils/browser/index.js';
+import 'playwright';
+
+/**
+ * Generate screenshots for all stories (parallelized for speed)
+ */
+declare function generateAllScreenshots(stories: Array<{
+    id: string;
+    importPath: string;
+}>, storybookUrl?: string): Promise<void>;
+
+interface BoundingBox {
+    width: number;
+    height: number;
+}
+interface ScreenshotMetadata {
+    fileHash: string;
+    lastGenerated: string;
+    sourcePath: string;
+    screenshots: {
+        light: string;
+        dark: string;
+    };
+    boundingBox?: BoundingBox;
+}
+interface Manifest {
+    stories: Record<string, ScreenshotMetadata>;
+}
+
+interface GenerateScreenshotResult {
+    path: string;
+    boundingBox: BoundingBox | null;
+}
+/**
+ * Get screenshot file path
+ */
+declare function getScreenshotPath(storyId: string, theme: 'light' | 'dark'): string;
+/**
+ * Check if screenshot exists
+ */
+declare function screenshotExists(storyId: string, theme: 'light' | 'dark'): boolean;
+/**
+ * Generate a screenshot for a story and save to disk
+ */
+declare function generateScreenshot(storyId: string, theme: 'light' | 'dark', storybookUrl?: string): Promise<GenerateScreenshotResult | null>;
+
+export { type Manifest, type ScreenshotMetadata, generateAllScreenshots, generateScreenshot, getScreenshotPath, screenshotExists };

package/dist/screenshot-service/index.js

@@ -0,0 +1,213 @@
+import crypto from 'crypto';
+import fs from 'fs';
+import path from 'path';
+import { chromium } from 'playwright';
+
+// src/utils/fileSystem/fileSystem.ts
+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;
+var MIN_COMPONENT_HEIGHT = 280;
+
+// src/utils/fileSystem/fileSystem.ts
+function ensureCacheDirectories() {
+  if (!fs.existsSync(CACHE_DIR)) {
+    fs.mkdirSync(CACHE_DIR, { recursive: true });
+  }
+  if (!fs.existsSync(SCREENSHOTS_DIR)) {
+    fs.mkdirSync(SCREENSHOTS_DIR, { recursive: true });
+  }
+}
+function computeFileHash(filePath) {
+  if (!fs.existsSync(filePath)) {
+    return "";
+  }
+  const content = fs.readFileSync(filePath, "utf-8");
+  return crypto.createHash("sha256").update(content).digest("hex");
+}
+function loadManifest() {
+  if (fs.existsSync(MANIFEST_PATH)) {
+    const content = fs.readFileSync(MANIFEST_PATH, "utf-8");
+    return JSON.parse(content);
+  }
+  return { stories: {} };
+}
+function saveManifest(manifest) {
+  ensureCacheDirectories();
+  fs.writeFileSync(MANIFEST_PATH, JSON.stringify(manifest, null, 2));
+}
+function updateManifest(storyId, sourcePath, fileHash, boundingBox) {
+  const manifest = loadManifest();
+  manifest.stories[storyId] = {
+    fileHash,
+    lastGenerated: (/* @__PURE__ */ new Date()).toISOString(),
+    sourcePath,
+    screenshots: {
+      light: `screenshots/${storyId}/light.png`,
+      dark: `screenshots/${storyId}/dark.png`
+    },
+    boundingBox: boundingBox ?? void 0
+  };
+  saveManifest(manifest);
+}
+var browser = null;
+async function getBrowser() {
+  if (!browser) {
+    browser = await chromium.launch({
+      headless: true
+    });
+  }
+  return browser;
+}
+async function closeBrowser() {
+  if (browser) {
+    try {
+      await browser.close();
+    } catch (error) {
+      console.error("Error closing browser:", error);
+    } finally {
+      browser = null;
+    }
+  }
+}
+function getScreenshotPath(storyId, theme) {
+  const storyDir = path.join(SCREENSHOTS_DIR, storyId);
+  return path.join(storyDir, `${theme}.png`);
+}
+function screenshotExists(storyId, theme) {
+  const screenshotPath = getScreenshotPath(storyId, theme);
+  return fs.existsSync(screenshotPath);
+}
+async function captureScreenshotBuffer(storyId, theme, width = VIEWPORT_WIDTH, height = VIEWPORT_HEIGHT, storybookUrl = "http://localhost:6006") {
+  const browser2 = await getBrowser();
+  const context = await browser2.newContext({
+    viewport: { width, height },
+    deviceScaleFactor: 2
+  });
+  const page = await context.newPage();
+  try {
+    const url = `${storybookUrl}/iframe.html?id=${storyId}&viewMode=story&globals=theme:${theme}`;
+    await page.goto(url, { timeout: 15e3 });
+    await page.waitForLoadState("domcontentloaded");
+    await page.waitForLoadState("load");
+    await page.waitForLoadState("networkidle");
+    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);
+          });
+        })
+      );
+    });
+    const contentBounds = await page.evaluate(() => {
+      const root = document.querySelector("#storybook-root");
+      if (!root) return null;
+      const children = root.querySelectorAll("*");
+      if (children.length === 0) return null;
+      let minX = Infinity, minY = Infinity, maxX = -Infinity, maxY = -Infinity;
+      children.forEach((child) => {
+        const rect = child.getBoundingClientRect();
+        if (rect.width === 0 || rect.height === 0) return;
+        minX = Math.min(minX, rect.left);
+        minY = Math.min(minY, rect.top);
+        maxX = Math.max(maxX, rect.right);
+        maxY = Math.max(maxY, rect.bottom);
+      });
+      if (minX === Infinity) return null;
+      return {
+        x: minX,
+        y: minY,
+        width: maxX - minX,
+        height: maxY - minY
+      };
+    });
+    let screenshotBuffer;
+    let resultBoundingBox = null;
+    if (contentBounds && contentBounds.width > 0 && contentBounds.height > 0) {
+      const PADDING = 20;
+      const clippedWidth = Math.min(width, contentBounds.width + PADDING);
+      const clippedHeight = Math.min(height, contentBounds.height + PADDING);
+      resultBoundingBox = {
+        width: Math.max(MIN_COMPONENT_WIDTH, Math.round(clippedWidth)),
+        height: Math.max(MIN_COMPONENT_HEIGHT, Math.round(clippedHeight))
+      };
+      screenshotBuffer = await page.screenshot({
+        type: "png",
+        clip: {
+          x: Math.max(0, contentBounds.x - 10),
+          y: Math.max(0, contentBounds.y - 10),
+          width: clippedWidth,
+          height: clippedHeight
+        }
+      });
+    } else {
+      screenshotBuffer = await page.screenshot({ type: "png" });
+    }
+    return { buffer: screenshotBuffer, boundingBox: resultBoundingBox };
+  } finally {
+    await context.close();
+  }
+}
+async function generateScreenshot(storyId, theme, storybookUrl = "http://localhost:6006") {
+  try {
+    ensureCacheDirectories();
+    const storyDir = path.join(SCREENSHOTS_DIR, storyId);
+    if (!fs.existsSync(storyDir)) {
+      fs.mkdirSync(storyDir, { recursive: true });
+    }
+    const screenshotPath = getScreenshotPath(storyId, theme);
+    const { buffer, boundingBox } = await captureScreenshotBuffer(
+      storyId,
+      theme,
+      VIEWPORT_WIDTH,
+      VIEWPORT_HEIGHT,
+      storybookUrl
+    );
+    fs.writeFileSync(screenshotPath, buffer);
+    return { path: screenshotPath, boundingBox };
+  } catch (error) {
+    console.error(`Error generating screenshot for ${storyId} (${theme}):`, error);
+    return null;
+  }
+}
+
+// src/screenshot-service/screenshot-service.ts
+async function generateAllScreenshots(stories, storybookUrl = "http://localhost:6006") {
+  console.log(`Generating screenshots for ${stories.length} stories...`);
+  const BATCH_SIZE = 10;
+  const batches = [];
+  for (let i = 0; i < stories.length; i += BATCH_SIZE) {
+    batches.push(stories.slice(i, i + BATCH_SIZE));
+  }
+  let completed = 0;
+  for (const batch of batches) {
+    await Promise.all(
+      batch.map(async (story) => {
+        const [lightResult, darkResult] = await Promise.all([
+          generateScreenshot(story.id, "light", storybookUrl),
+          generateScreenshot(story.id, "dark", storybookUrl)
+        ]);
+        if (lightResult && darkResult) {
+          const fileHash = computeFileHash(story.importPath);
+          updateManifest(story.id, story.importPath, fileHash, lightResult.boundingBox);
+        }
+        completed++;
+        console.log(
+          `[${completed}/${stories.length}] Generated screenshots for ${story.id}`
+        );
+      })
+    );
+  }
+  await closeBrowser();
+  console.log("Screenshot generation complete!");
+}
+
+export { closeBrowser, generateAllScreenshots, generateScreenshot, getScreenshotPath, screenshotExists };

package/dist/screenshot-service/utils/browser/index.d.ts

@@ -0,0 +1,12 @@
+import { Browser } from 'playwright';
+
+/**
+ * Initialize browser instance
+ */
+declare function getBrowser(): Promise<Browser>;
+/**
+ * Close browser instance
+ */
+declare function closeBrowser(): Promise<void>;
+
+export { closeBrowser, getBrowser };

package/dist/screenshot-service/utils/browser/index.js

@@ -0,0 +1,25 @@
+import { chromium } from 'playwright';
+
+// src/screenshot-service/utils/browser/browser.ts
+var browser = null;
+async function getBrowser() {
+  if (!browser) {
+    browser = await chromium.launch({
+      headless: true
+    });
+  }
+  return browser;
+}
+async function closeBrowser() {
+  if (browser) {
+    try {
+      await browser.close();
+    } catch (error) {
+      console.error("Error closing browser:", error);
+    } finally {
+      browser = null;
+    }
+  }
+}
+
+export { closeBrowser, getBrowser };

package/package.json

@@ -1,22 +1,38 @@
 {
   "name": "storybook-onbook-plugin",
-  "version": "0.2.2",
+  "version": "0.2.4",
   "type": "module",
   "bin": {
     "storybook-onbook-plugin": "./dist/cli/index.js"
   },
   "exports": {
     ".": {
-      "types": "./dist/index.d.ts",
-      "import": "./dist/index.js"
+      "types": "./src/index.ts",
+      "default": "./src/index.ts"
     },
     "./screenshot-service": {
-      "types": "./dist/screenshot-service/index.d.ts",
-      "import": "./dist/screenshot-service/index.js"
+      "types": "./src/screenshot-service/index.ts",
+      "default": "./src/screenshot-service/index.ts"
     },
     "./screenshot-service/browser": {
-      "types": "./dist/screenshot-service/utils/browser/index.d.ts",
-      "import": "./dist/screenshot-service/utils/browser/index.js"
+      "types": "./src/screenshot-service/utils/browser/index.ts",
+      "default": "./src/screenshot-service/utils/browser/index.ts"
+    }
+  },
+  "publishConfig": {
+    "exports": {
+      ".": {
+        "types": "./dist/index.d.ts",
+        "import": "./dist/index.js"
+      },
+      "./screenshot-service": {
+        "types": "./dist/screenshot-service/index.d.ts",
+        "import": "./dist/screenshot-service/index.js"
+      },
+      "./screenshot-service/browser": {
+        "types": "./dist/screenshot-service/utils/browser/index.d.ts",
+        "import": "./dist/screenshot-service/utils/browser/index.js"
+      }
     }
   },
   "files": [
@@ -24,7 +40,7 @@
     "README.md"
   ],
   "scripts": {
-    "build": "tsc",
+    "build": "tsup",
     "clean": "git clean -xdf .cache .turbo dist node_modules",
     "typecheck": "tsc --noEmit",
     "prepublishOnly": "bun run build",
@@ -44,6 +60,7 @@
     "@types/babel__traverse": "^7.20.6",
     "@types/node": "^22.15.32",
     "bun-types": "^1.3.5",
+    "tsup": "^8.5.1",
     "typescript": "5.8.3",
     "vite": "^6.3.5"
   },