@tpitre/story-ui 3.4.2 → 3.5.0

package/dist/mcp-server/routes/storyHelpers.d.ts

@@ -0,0 +1,51 @@
+/**
+ * Shared helper functions for story generation routes.
+ * These utilities are used by both generateStory.ts and generateStoryStream.ts
+ * to eliminate code duplication and ensure consistent behavior.
+ */
+import { ImageContent } from '../../story-generator/llm-providers/types.js';
+/**
+ * Slugify a string for use in filenames and identifiers.
+ */
+export declare function slugify(str: string): string;
+/**
+ * Extract code block from LLM response text.
+ * Accepts various language identifiers (tsx, jsx, typescript, etc.)
+ */
+export declare function extractCodeBlock(text: string): string | null;
+/**
+ * Call the LLM service with optional vision support.
+ * Automatically handles provider configuration and image attachment.
+ */
+export declare function callLLM(messages: {
+    role: 'user' | 'assistant';
+    content: string;
+}[], images?: ImageContent[]): Promise<string>;
+/**
+ * Clean a user prompt to create a readable title.
+ * Removes common leading phrases and normalizes formatting.
+ */
+export declare function cleanPromptForTitle(prompt: string): string;
+/**
+ * Generate a title using the LLM service.
+ * Falls back to empty string on failure.
+ */
+export declare function getLLMTitle(userPrompt: string): Promise<string>;
+/**
+ * Escape a title string for use in TypeScript string literals.
+ */
+export declare function escapeTitleForTS(title: string): string;
+/**
+ * Extract component imports from code for a specific import path.
+ */
+export declare function extractImportsFromCode(code: string, importPath: string): string[];
+/**
+ * Generate a filename from a story title and hash.
+ */
+export declare function fileNameFromTitle(title: string, hash: string): string;
+/**
+ * Find a similar icon name from the allowed icons set.
+ * Used for providing suggestions when an invalid icon is used.
+ */
+export declare function findSimilarIcon(iconName: string, allowedIcons: Set<string>): string | null;
+//# sourceMappingURL=storyHelpers.d.ts.map

package/dist/mcp-server/routes/storyHelpers.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"storyHelpers.d.ts","sourceRoot":"","sources":["../../../mcp-server/routes/storyHelpers.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAWH,OAAO,EAAE,YAAY,EAAE,MAAM,8CAA8C,CAAC;AAE5E;;GAEG;AACH,wBAAgB,OAAO,CAAC,GAAG,EAAE,MAAM,GAAG,MAAM,CAQ3C;AAED;;;GAGG;AACH,wBAAgB,gBAAgB,CAAC,IAAI,EAAE,MAAM,GAAG,MAAM,GAAG,IAAI,CAG5D;AAED;;;GAGG;AACH,wBAAsB,OAAO,CAC3B,QAAQ,EAAE;IAAE,IAAI,EAAE,MAAM,GAAG,WAAW,CAAC;IAAC,OAAO,EAAE,MAAM,CAAA;CAAE,EAAE,EAC3D,MAAM,CAAC,EAAE,YAAY,EAAE,GACtB,OAAO,CAAC,MAAM,CAAC,CA6BjB;AAED;;;GAGG;AACH,wBAAgB,mBAAmB,CAAC,MAAM,EAAE,MAAM,GAAG,MAAM,CAkC1D;AAED;;;GAGG;AACH,wBAAsB,WAAW,CAAC,UAAU,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC,CAOrE;AAED;;GAEG;AACH,wBAAgB,gBAAgB,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM,CAStD;AAED;;GAEG;AACH,wBAAgB,sBAAsB,CAAC,IAAI,EAAE,MAAM,EAAE,UAAU,EAAE,MAAM,GAAG,MAAM,EAAE,CAajF;AAED;;GAEG;AACH,wBAAgB,iBAAiB,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,GAAG,MAAM,CAgBrE;AAED;;;GAGG;AACH,wBAAgB,eAAe,CAAC,QAAQ,EAAE,MAAM,EAAE,YAAY,EAAE,GAAG,CAAC,MAAM,CAAC,GAAG,MAAM,GAAG,IAAI,CAwB1F"}

package/dist/mcp-server/routes/storyHelpers.js

@@ -0,0 +1,185 @@
+/**
+ * Shared helper functions for story generation routes.
+ * These utilities are used by both generateStory.ts and generateStoryStream.ts
+ * to eliminate code duplication and ensure consistent behavior.
+ */
+import { logger } from '../../story-generator/logger.js';
+import { chatCompletion, generateTitle as llmGenerateTitle, isProviderConfigured, getProviderInfo, chatCompletionWithImages, buildMessageWithImages } from '../../story-generator/llm-providers/story-llm-service.js';
+/**
+ * Slugify a string for use in filenames and identifiers.
+ */
+export function slugify(str) {
+    if (!str || typeof str !== 'string') {
+        return 'untitled';
+    }
+    return str
+        .toLowerCase()
+        .replace(/[^a-z0-9]+/g, '-')
+        .replace(/^-+|-+$/g, '');
+}
+/**
+ * Extract code block from LLM response text.
+ * Accepts various language identifiers (tsx, jsx, typescript, etc.)
+ */
+export function extractCodeBlock(text) {
+    const codeBlock = text.match(/```(?:tsx|jsx|typescript|ts|js|javascript)?([\s\S]*?)```/i);
+    return codeBlock ? codeBlock[1].trim() : null;
+}
+/**
+ * Call the LLM service with optional vision support.
+ * Automatically handles provider configuration and image attachment.
+ */
+export async function callLLM(messages, images) {
+    if (!isProviderConfigured()) {
+        throw new Error('No LLM provider configured. Please set ANTHROPIC_API_KEY, OPENAI_API_KEY, or GEMINI_API_KEY.');
+    }
+    const providerInfo = getProviderInfo();
+    logger.debug(`Using ${providerInfo.currentProvider} (${providerInfo.currentModel}) for story generation`);
+    if (images && images.length > 0) {
+        if (!providerInfo.supportsVision) {
+            throw new Error(`${providerInfo.currentProvider} does not support vision. Please configure a vision-capable provider.`);
+        }
+        logger.log(`🖼️ Using vision-capable chat with ${images.length} image(s)`);
+        const messagesWithImages = messages.map((msg, index) => {
+            if (msg.role === 'user' && index === 0) {
+                return {
+                    role: msg.role,
+                    content: buildMessageWithImages(msg.content, images)
+                };
+            }
+            return msg;
+        });
+        return await chatCompletionWithImages(messagesWithImages, { maxTokens: 8192 });
+    }
+    return await chatCompletion(messages, { maxTokens: 8192 });
+}
+/**
+ * Clean a user prompt to create a readable title.
+ * Removes common leading phrases and normalizes formatting.
+ */
+export function cleanPromptForTitle(prompt) {
+    if (!prompt || typeof prompt !== 'string') {
+        return 'Untitled Story';
+    }
+    const leadingPhrases = [
+        /^generate (a|an|the)? /i,
+        /^build (a|an|the)? /i,
+        /^create (a|an|the)? /i,
+        /^make (a|an|the)? /i,
+        /^design (a|an|the)? /i,
+        /^show (me )?(a|an|the)? /i,
+        /^write (a|an|the)? /i,
+        /^produce (a|an|the)? /i,
+        /^construct (a|an|the)? /i,
+        /^draft (a|an|the)? /i,
+        /^compose (a|an|the)? /i,
+        /^implement (a|an|the)? /i,
+        /^build out (a|an|the)? /i,
+        /^add (a|an|the)? /i,
+        /^render (a|an|the)? /i,
+        /^display (a|an|the)? /i,
+    ];
+    let cleaned = prompt.trim();
+    for (const regex of leadingPhrases) {
+        cleaned = cleaned.replace(regex, '');
+    }
+    return cleaned
+        .replace(/[^\w\s'"?!-]/g, ' ')
+        .replace(/\s+/g, ' ')
+        .trim()
+        .replace(/\b\w/g, c => c.toUpperCase());
+}
+/**
+ * Generate a title using the LLM service.
+ * Falls back to empty string on failure.
+ */
+export async function getLLMTitle(userPrompt) {
+    try {
+        return await llmGenerateTitle(userPrompt);
+    }
+    catch (error) {
+        logger.warn('Failed to generate title via LLM, using fallback', { error });
+        return '';
+    }
+}
+/**
+ * Escape a title string for use in TypeScript string literals.
+ */
+export function escapeTitleForTS(title) {
+    return title
+        .replace(/\\/g, '\\\\')
+        .replace(/"/g, '\\"')
+        .replace(/'/g, "\\'")
+        .replace(/`/g, '\\`')
+        .replace(/\n/g, '\\n')
+        .replace(/\r/g, '\\r')
+        .replace(/\t/g, '\\t');
+}
+/**
+ * Extract component imports from code for a specific import path.
+ */
+export function extractImportsFromCode(code, importPath) {
+    const imports = [];
+    const escapedPath = importPath.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+    const importRegex = new RegExp(`import\\s*{([^}]+)}\\s*from\\s*['"]${escapedPath}['"]`, 'g');
+    let match;
+    while ((match = importRegex.exec(code)) !== null) {
+        const importList = match[1];
+        const components = importList.split(',').map(comp => comp.trim());
+        imports.push(...components);
+    }
+    return imports;
+}
+/**
+ * Generate a filename from a story title and hash.
+ */
+export function fileNameFromTitle(title, hash) {
+    if (!title || typeof title !== 'string') {
+        title = 'untitled';
+    }
+    if (!hash || typeof hash !== 'string') {
+        hash = 'default';
+    }
+    const base = title
+        .toLowerCase()
+        .replace(/[^a-z0-9]+/g, '-')
+        .replace(/^-+|-+$/g, '')
+        .replace(/"|'/g, '')
+        .slice(0, 60);
+    return `${base}-${hash}.stories.tsx`;
+}
+/**
+ * Find a similar icon name from the allowed icons set.
+ * Used for providing suggestions when an invalid icon is used.
+ */
+export function findSimilarIcon(iconName, allowedIcons) {
+    if (!iconName || typeof iconName !== 'string') {
+        return null;
+    }
+    const iconLower = iconName.toLowerCase();
+    for (const allowed of allowedIcons) {
+        const allowedLower = allowed.toLowerCase();
+        // Check if core words match
+        if (iconLower.includes('commit') && allowedLower.includes('commit'))
+            return allowed;
+        if (iconLower.includes('branch') && allowedLower.includes('branch'))
+            return allowed;
+        if (iconLower.includes('merge') && allowedLower.includes('merge'))
+            return allowed;
+        if (iconLower.includes('pull') && allowedLower.includes('pull'))
+            return allowed;
+        if (iconLower.includes('push') && allowedLower.includes('push'))
+            return allowed;
+        if (iconLower.includes('star') && allowedLower.includes('star'))
+            return allowed;
+        if (iconLower.includes('user') && allowedLower.includes('user'))
+            return allowed;
+        if (iconLower.includes('search') && allowedLower.includes('search'))
+            return allowed;
+        if (iconLower.includes('settings') && allowedLower.includes('settings'))
+            return allowed;
+        if (iconLower.includes('home') && allowedLower.includes('home'))
+            return allowed;
+    }
+    return null;
+}

package/dist/story-generator/generateStory.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"generateStory.d.ts","sourceRoot":"","sources":["../../story-generator/generateStory.ts"],"names":[],"mappings":"AAEA,OAAO,EAAE,aAAa,EAAE,MAAM,uBAAuB,CAAC;AA6BtD,wBAAgB,aAAa,CAAC,EAC5B,YAAY,EACZ,QAAQ,EACR,MAAM,EACP,EAAE;IACD,YAAY,EAAE,MAAM,CAAC;IACrB,QAAQ,EAAE,MAAM,CAAC;IACjB,MAAM,EAAE,aAAa,CAAC;CACvB,UAmBA"}
+{"version":3,"file":"generateStory.d.ts","sourceRoot":"","sources":["../../story-generator/generateStory.ts"],"names":[],"mappings":"AAEA,OAAO,EAAE,aAAa,EAAE,MAAM,uBAAuB,CAAC;AAsBtD,wBAAgB,aAAa,CAAC,EAC5B,YAAY,EACZ,QAAQ,EACR,MAAM,EACP,EAAE;IACD,YAAY,EAAE,MAAM,CAAC;IACrB,QAAQ,EAAE,MAAM,CAAC;IACjB,MAAM,EAAE,aAAa,CAAC;CACvB,UAmBA"}

package/dist/story-generator/generateStory.js

@@ -1,11 +1,5 @@
 import fs from 'fs';
 import path from 'path';
-function slugify(str) {
-    return str
-        .toLowerCase()
-        .replace(/[^a-z0-9]+/g, '-')
-        .replace(/^-+|-+$/g, '');
-}
 /**
  * Check if the current working directory is the Story UI package itself.
  * This prevents accidentally generating stories in the package source code.
@@ -41,8 +35,3 @@ export function generateStory({ fileContents, fileName, config }) {
     fs.writeFileSync(outPath, fileContents, 'utf-8');
     return outPath;
 }
-// Mock usage:
-// generateStory({
-//   title: 'Login Form',
-//   jsx: '<al-input label="Email"></al-input>\n<al-input label="Password" type="password"></al-input>\n<al-button>Login</al-button>'
-// });

package/dist/templates/StoryUI/StoryUIPanel.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"StoryUIPanel.d.ts","sourceRoot":"","sources":["../../../templates/StoryUI/StoryUIPanel.tsx"],"names":[],"mappings":"AAiTA,OAAO,CAAC,MAAM,CAAC;IACb,UAAU,MAAM;QACd,2BAA2B,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;KACtD;CACF;AA6oCD,iBAAS,YAAY,4CA8sCpB;AAED,eAAe,YAAY,CAAC;AAC5B,OAAO,EAAE,YAAY,EAAE,CAAC"}
+{"version":3,"file":"StoryUIPanel.d.ts","sourceRoot":"","sources":["../../../templates/StoryUI/StoryUIPanel.tsx"],"names":[],"mappings":"AAkTA,OAAO,CAAC,MAAM,CAAC;IACb,UAAU,MAAM;QACd,2BAA2B,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;KACtD;CACF;AA6oCD,iBAAS,YAAY,4CAmtCpB;AAED,eAAe,YAAY,CAAC;AAC5B,OAAO,EAAE,YAAY,EAAE,CAAC"}

package/dist/templates/StoryUI/StoryUIPanel.js

@@ -1104,6 +1104,7 @@ function StoryUIPanel() {
                     file,
                     preview,
                     base64,
+                    mediaType: file.type || 'image/png',
                 });
             }
             catch (err) {
@@ -1182,6 +1183,7 @@ function StoryUIPanel() {
                     file,
                     preview,
                     base64,
+                    mediaType: file.type || 'image/png',
                 });
             }
             catch (err) {
@@ -1236,6 +1238,7 @@ function StoryUIPanel() {
                     file: renamedFile,
                     preview,
                     base64,
+                    mediaType: file.type || 'image/png',
                 });
             }
             catch (err) {
@@ -1874,7 +1877,9 @@ function StoryUIPanel() {
                                                     maxWidth: '200px'
                                                 }, children: availableProviders
                                                     .find(p => p.type === selectedProvider)
-                                                    ?.models.map(model => (_jsx("option", { value: model, children: model }, model))) })] })] }))] }), _jsxs("div", { style: STYLES.chatContainer, children: [error && (_jsx("div", { style: STYLES.errorMessage, children: error })), conversation.length === 0 && !loading && (_jsxs("div", { style: STYLES.emptyState, children: [_jsx("div", { style: STYLES.emptyStateTitle, children: "Start a new conversation" }), _jsx("div", { style: STYLES.emptyStateSubtitle, children: "Describe the UI component you'd like to create" })] })), conversation.map((msg, i) => (_jsx("div", { style: STYLES.messageContainer, children: _jsxs("div", { style: msg.role === 'user' ? STYLES.userMessage : STYLES.aiMessage, children: [msg.role === 'ai' ? renderMarkdown(msg.content) : msg.content, msg.role === 'user' && msg.attachedImages && msg.attachedImages.length > 0 && (_jsx("div", { style: STYLES.userMessageImages, children: msg.attachedImages.map((img) => (_jsx("img", { src: img.preview, alt: "attached", style: STYLES.userMessageImage }, img.id))) }))] }) }, i))), loading && (_jsx("div", { style: STYLES.messageContainer, children: streamingState ? (_jsx(StreamingProgressMessage, { streamingData: streamingState })) : (_jsxs("div", { style: STYLES.loadingMessage, children: [_jsx("span", { children: "Generating story" }), _jsx("span", { className: "loading-dots" })] })) })), _jsx("div", { ref: chatEndRef })] }), _jsx("input", { ref: fileInputRef, type: "file", accept: "image/*", multiple: true, style: { display: 'none' }, onChange: handleFileSelect }), attachedImages.length > 0 && (_jsxs("div", { style: STYLES.imagePreviewContainer, children: [_jsxs("span", { style: STYLES.imagePreviewLabel, children: [_jsx("svg", { width: 14, height: 14, viewBox: "0 0 24 24", fill: "currentColor", children: _jsx("path", { d: "M21 19V5c0-1.1-.9-2-2-2H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2zM8.5 13.5l2.5 3.01L14.5 12l4.5 6H5l3.5-4.5z" }) }), attachedImages.length, " image", attachedImages.length > 1 ? 's' : '', " attached"] }), attachedImages.map((img) => (_jsxs("div", { style: STYLES.imagePreviewItem, children: [_jsx("img", { src: img.preview, alt: "preview", style: STYLES.imagePreviewImg }), _jsx("button", { type: "button", style: STYLES.imageRemoveButton, onClick: () => removeAttachedImage(img.id), title: "Remove image", children: "\u00D7" })] }, img.id)))] })), _jsxs("form", { onSubmit: handleSend, style: {
+                                                    ?.models.map(model => (_jsx("option", { value: model, children: model }, model))) })] })] }))] }), _jsxs("div", { style: STYLES.chatContainer, children: [error && (_jsx("div", { style: STYLES.errorMessage, children: error })), conversation.length === 0 && !loading && (_jsxs("div", { style: STYLES.emptyState, children: [_jsx("div", { style: STYLES.emptyStateTitle, children: "Start a new conversation" }), _jsx("div", { style: STYLES.emptyStateSubtitle, children: "Describe the UI component you'd like to create" })] })), conversation.map((msg, i) => (_jsx("div", { style: STYLES.messageContainer, children: _jsxs("div", { style: msg.role === 'user' ? STYLES.userMessage : STYLES.aiMessage, children: [msg.role === 'ai' ? renderMarkdown(msg.content) : msg.content, msg.role === 'user' && msg.attachedImages && msg.attachedImages.length > 0 && (_jsx("div", { style: STYLES.userMessageImages, children: msg.attachedImages.map((img) => (_jsx("img", { src: img.base64
+                                                    ? `data:${img.mediaType || 'image/png'};base64,${img.base64}`
+                                                    : img.preview, alt: "attached", style: STYLES.userMessageImage }, img.id))) }))] }) }, i))), loading && (_jsx("div", { style: STYLES.messageContainer, children: streamingState ? (_jsx(StreamingProgressMessage, { streamingData: streamingState })) : (_jsxs("div", { style: STYLES.loadingMessage, children: [_jsx("span", { children: "Generating story" }), _jsx("span", { className: "loading-dots" })] })) })), _jsx("div", { ref: chatEndRef })] }), _jsx("input", { ref: fileInputRef, type: "file", accept: "image/*", multiple: true, style: { display: 'none' }, onChange: handleFileSelect }), attachedImages.length > 0 && (_jsxs("div", { style: STYLES.imagePreviewContainer, children: [_jsxs("span", { style: STYLES.imagePreviewLabel, children: [_jsx("svg", { width: 14, height: 14, viewBox: "0 0 24 24", fill: "currentColor", children: _jsx("path", { d: "M21 19V5c0-1.1-.9-2-2-2H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2zM8.5 13.5l2.5 3.01L14.5 12l4.5 6H5l3.5-4.5z" }) }), attachedImages.length, " image", attachedImages.length > 1 ? 's' : '', " attached"] }), attachedImages.map((img) => (_jsxs("div", { style: STYLES.imagePreviewItem, children: [_jsx("img", { src: img.preview, alt: "preview", style: STYLES.imagePreviewImg }), _jsx("button", { type: "button", style: STYLES.imageRemoveButton, onClick: () => removeAttachedImage(img.id), title: "Remove image", children: "\u00D7" })] }, img.id)))] })), _jsxs("form", { onSubmit: handleSend, style: {
                             ...STYLES.inputForm,
                             ...(attachedImages.length > 0 ? {
                                 marginTop: 0,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tpitre/story-ui",
-  "version": "3.4.2",
+  "version": "3.5.0",
   "description": "AI-powered Storybook story generator with dynamic component discovery",
   "type": "module",
   "main": "dist/index.js",

package/templates/StoryUI/StoryUIPanel.tsx

@@ -138,6 +138,7 @@ interface AttachedImage {
   file: File;
   preview: string;
   base64?: string;
+  mediaType?: string; // Store MIME type for data URL reconstruction after localStorage restore
 }
 
 // Provider info from /story-ui/providers API
@@ -1544,6 +1545,7 @@ function StoryUIPanel() {
           file,
           preview,
           base64,
+          mediaType: file.type || 'image/png',
         });
       } catch (err) {
         errors.push(`${file.name}: Failed to process`);
@@ -1638,6 +1640,7 @@ function StoryUIPanel() {
           file,
           preview,
           base64,
+          mediaType: file.type || 'image/png',
         });
       } catch (err) {
         errors.push(`${file.name}: Failed to process`);
@@ -1703,6 +1706,7 @@ function StoryUIPanel() {
           file: renamedFile,
           preview,
           base64,
+          mediaType: file.type || 'image/png',
         });
       } catch (err) {
         errors.push('Failed to process pasted image');
@@ -2548,7 +2552,9 @@ function StoryUIPanel() {
                     {msg.attachedImages.map((img) => (
                       <img
                         key={img.id}
-                        src={img.preview}
+                        src={img.base64
+                          ? `data:${img.mediaType || 'image/png'};base64,${img.base64}`
+                          : img.preview}
                         alt="attached"
                         style={STYLES.userMessageImage}
                       />

package/templates/StoryUI/manager.tsx

@@ -38,7 +38,101 @@ declare global {
  * Into:
  *   <Button>Click</Button>
  */
-const extractUsageCode = (fullStoryCode: string): string => {
+const extractUsageCode = (fullStoryCode: string, variantName?: string): string => {
+  // Helper function to generate JSX from args
+  const generateJsxFromArgs = (argsStr: string, componentName: string): string | null => {
+    try {
+      // Extract children if present
+      const childrenMatch = argsStr.match(/children:\s*['"`]([^'"`]+)['"`]/);
+      const children = childrenMatch ? childrenMatch[1] : '';
+
+      // Extract other props (simplified)
+      const propsStr = argsStr
+        .replace(/children:\s*['"`][^'"`]*['"`],?/, '') // Remove children
+        .replace(/^\{|\}$/g, '') // Remove braces
+        .trim();
+
+      if (children) {
+        if (propsStr) {
+          return `<${componentName} ${propsStr.replace(/,\s*$/, '')}>${children}</${componentName}>`;
+        }
+        return `<${componentName}>${children}</${componentName}>`;
+      } else if (propsStr) {
+        return `<${componentName} ${propsStr.replace(/,\s*$/, '')} />`;
+      }
+      return `<${componentName} />`;
+    } catch {
+      return null;
+    }
+  };
+
+  // Get the component name from meta
+  const componentMatch = fullStoryCode.match(/component:\s*([A-Z][A-Za-z0-9]*)/);
+  const componentName = componentMatch ? componentMatch[1] : null;
+
+  // If we have a variant name, try to find that specific variant's args or render
+  if (variantName) {
+    // Normalize variant name for matching:
+    // - "primary" -> "Primary"
+    // - "full-width" -> "FullWidth" (kebab-case to PascalCase)
+    const normalizedVariant = variantName
+      .split('-')
+      .map(word => word.charAt(0).toUpperCase() + word.slice(1).toLowerCase())
+      .join('');
+
+    // Pattern A: export const Primary: Story = { args: {...} }
+    // Match the specific variant's export block
+    const variantExportRegex = new RegExp(
+      `export\\s+const\\s+${normalizedVariant}\\s*(?::\\s*Story)?\\s*=\\s*\\{([\\s\\S]*?)\\}\\s*;`,
+      'i'
+    );
+    const variantExportMatch = fullStoryCode.match(variantExportRegex);
+
+    if (variantExportMatch) {
+      const variantBlock = variantExportMatch[1];
+
+      // Try to extract render function from this variant
+      const renderWithParensMatch = variantBlock.match(/render:\s*\([^)]*\)\s*=>\s*\(\s*([\s\S]*?)\s*\)\s*[,}]/);
+      if (renderWithParensMatch) {
+        return renderWithParensMatch[1].trim().replace(/,\s*$/, '');
+      }
+
+      const renderNoParensMatch = variantBlock.match(/render:\s*\([^)]*\)\s*=>\s*(<[A-Z][^,}]*(?:\/>|<\/[A-Za-z.]+>))/s);
+      if (renderNoParensMatch) {
+        return renderNoParensMatch[1].trim();
+      }
+
+      // Try to extract args from this variant
+      const argsMatch = variantBlock.match(/args:\s*(\{[\s\S]*?\})\s*[,}]/);
+      if (argsMatch && componentName) {
+        const result = generateJsxFromArgs(argsMatch[1], componentName);
+        if (result) return result;
+      }
+    }
+
+    // Pattern B: Arrow function variant: export const Primary = () => <Component...>
+    const arrowVariantRegex = new RegExp(
+      `export\\s+const\\s+${normalizedVariant}\\s*=\\s*\\(\\)\\s*=>\\s*\\(\\s*([\\s\\S]*?)\\s*\\)\\s*;`,
+      'i'
+    );
+    const arrowVariantMatch = fullStoryCode.match(arrowVariantRegex);
+    if (arrowVariantMatch) {
+      return arrowVariantMatch[1].trim().replace(/,\s*$/, '');
+    }
+
+    // Pattern C: Arrow function without parens: export const Primary = () => <Component...>;
+    const arrowNoParensRegex = new RegExp(
+      `export\\s+const\\s+${normalizedVariant}\\s*=\\s*\\(\\)\\s*=>\\s*(<[A-Z][^;]*(?:\\/>|<\\/[A-Za-z.]+>))\\s*;`,
+      'is'
+    );
+    const arrowNoParensMatch = fullStoryCode.match(arrowNoParensRegex);
+    if (arrowNoParensMatch) {
+      return arrowNoParensMatch[1].trim();
+    }
+  }
+
+  // Fallback: Try generic patterns (for Default or when variant not specified)
+
   // Try to extract JSX from render function: render: () => (<JSX>) or render: () => <JSX>
   // Pattern 1: render: () => (\n  <Component...>\n)
   const renderWithParensMatch = fullStoryCode.match(/render:\s*\([^)]*\)\s*=>\s*\(\s*([\s\S]*?)\s*\)\s*[,}]/);
@@ -70,37 +164,9 @@ const extractUsageCode = (fullStoryCode: string): string => {
   // Pattern 5: Look for args-based stories with component prop spreading
   // e.g., args: { children: 'Click me', color: 'blue' }
   const argsMatch = fullStoryCode.match(/args:\s*(\{[\s\S]*?\})\s*[,}]/);
-  if (argsMatch) {
-    // Try to find the component from the meta
-    const componentMatch = fullStoryCode.match(/component:\s*([A-Z][A-Za-z0-9]*)/);
-    if (componentMatch) {
-      const componentName = componentMatch[1];
-      try {
-        // Parse the args to generate JSX
-        const argsStr = argsMatch[1];
-        // Extract children if present
-        const childrenMatch = argsStr.match(/children:\s*['"`]([^'"`]+)['"`]/);
-        const children = childrenMatch ? childrenMatch[1] : '';
-
-        // Extract other props (simplified)
-        const propsStr = argsStr
-          .replace(/children:\s*['"`][^'"`]*['"`],?/, '') // Remove children
-          .replace(/^\{|\}$/g, '') // Remove braces
-          .trim();
-
-        if (children) {
-          if (propsStr) {
-            return `<${componentName} ${propsStr.replace(/,\s*$/, '')}>${children}</${componentName}>`;
-          }
-          return `<${componentName}>${children}</${componentName}>`;
-        } else if (propsStr) {
-          return `<${componentName} ${propsStr.replace(/,\s*$/, '')} />`;
-        }
-        return `<${componentName} />`;
-      } catch {
-        // Fall through to return full code
-      }
-    }
+  if (argsMatch && componentName) {
+    const result = generateJsxFromArgs(argsMatch[1], componentName);
+    if (result) return result;
   }
 
   // Pattern 6: Look for any JSX block starting with < and ending with /> or </Component>
@@ -297,14 +363,33 @@ const SourceCodePanel: React.FC<{ active?: boolean }> = ({ active }) => {
   const [copied, setCopied] = useState(false);
   const [showFullCode, setShowFullCode] = useState(false);
 
-  // Memoize the usage code extraction
+  // Get the current story ID
+  const currentStoryId = state?.storyId;
+
+  // Extract variant name from story ID (e.g., "generated-button--primary" -> "primary", "generated-button--full-width" -> "full-width")
+  const currentVariant = useMemo(() => {
+    if (!currentStoryId) return undefined;
+    // Match everything after the last -- (variant can contain hyphens like "full-width")
+    const variantMatch = currentStoryId.match(/--([a-z0-9-]+)$/i);
+    return variantMatch ? variantMatch[1] : undefined;
+  }, [currentStoryId]);
+
+  // Memoize the usage code extraction with variant awareness
+  const usageCode = useMemo(() => {
+    if (!sourceCode) return '';
+    return extractUsageCode(sourceCode, currentVariant);
+  }, [sourceCode, currentVariant]);
+
   const displayCode = useMemo(() => {
     if (!sourceCode) return '';
-    return showFullCode ? sourceCode : extractUsageCode(sourceCode);
-  }, [sourceCode, showFullCode]);
+    return showFullCode ? sourceCode : usageCode;
+  }, [sourceCode, showFullCode, usageCode]);
 
-  // Get the current story ID
-  const currentStoryId = state?.storyId;
+  // Check if there's different usage code (for showing toggle button)
+  // This should remain true even when showing full code
+  const hasUsageCode = useMemo(() => {
+    return sourceCode && usageCode && usageCode !== sourceCode;
+  }, [sourceCode, usageCode]);
 
   // Try to get source code from the story
   useEffect(() => {
@@ -340,11 +425,81 @@ const SourceCodePanel: React.FC<{ active?: boolean }> = ({ active }) => {
       let cachedCode = topWindow.__STORY_UI_GENERATED_CODE__?.[currentStoryId] ||
                         window.__STORY_UI_GENERATED_CODE__?.[currentStoryId];
 
+      console.log('[Source Code Panel DEBUG] Looking for code:', {
+        currentStoryId,
+        foundInWindowCache: !!cachedCode,
+        windowCacheKeys: Object.keys(topWindow.__STORY_UI_GENERATED_CODE__ || {}),
+      });
+
       // If not in memory cache, check localStorage (survives page navigation)
       if (!cachedCode) {
         try {
           const stored = JSON.parse(localStorage.getItem('storyui_generated_code') || '{}');
-          cachedCode = stored[currentStoryId];
+
+          console.log('[Source Code Panel DEBUG] localStorage lookup:', {
+            localStorageKeys: Object.keys(stored),
+            localStorageKeyCount: Object.keys(stored).length,
+          });
+
+          // Try multiple key formats since Storybook IDs differ from our storage keys
+          // Storybook ID format: "generated-componentname--variant" or "generated/componentname--variant"
+          // Our storage keys: "ComponentName", "ComponentName.stories.tsx", "story-hash123", etc.
+          const keysToTry: string[] = [currentStoryId];
+
+          // Extract component name and base story ID from Storybook ID
+          // e.g., "generated-simple-test-button--primary" -> baseId: "generated-simple-test-button--default", component: "simpletestbutton"
+          const match = currentStoryId.match(/^(generated[-\/]?.+?)(?:--(.*))?$/i);
+          if (match) {
+            const baseId = match[1];
+            const variant = match[2];
+
+            // Try base story ID with --default variant (this is what we store)
+            if (variant && variant !== 'default') {
+              keysToTry.push(`${baseId}--default`);
+            }
+            // Also try just the base ID without any variant
+            keysToTry.push(baseId);
+
+            // Extract component name (e.g., "generated-simple-test-button" -> "simpletestbutton")
+            const componentMatch = baseId.match(/^generated[-\/]?(.+)$/i);
+            if (componentMatch) {
+              const componentNameLower = componentMatch[1].replace(/-/g, '');
+              keysToTry.push(componentNameLower);
+              // Try PascalCase version (e.g., "simpletestbutton" -> "Simpletestbutton")
+              const pascalCase = componentNameLower.charAt(0).toUpperCase() + componentNameLower.slice(1);
+              keysToTry.push(pascalCase);
+              keysToTry.push(`${pascalCase}.stories.tsx`);
+
+              // Try with spaces converted to title case (e.g., "Simple Test Button" -> "SimpleTestButton")
+              const words = componentMatch[1].split('-');
+              if (words.length > 1) {
+                const titleCase = words.map(w => w.charAt(0).toUpperCase() + w.slice(1)).join('');
+                keysToTry.push(titleCase);
+                // Also try with space-separated title (what we store)
+                const spacedTitle = words.map(w => w.charAt(0).toUpperCase() + w.slice(1)).join(' ');
+                keysToTry.push(spacedTitle);
+              }
+            }
+          }
+
+          // Also try the story title if available
+          if (storyTitle) {
+            keysToTry.push(storyTitle);
+            keysToTry.push(storyTitle.replace(/\s+/g, ''));
+            keysToTry.push(`${storyTitle.replace(/\s+/g, '')}.stories.tsx`);
+          }
+
+          console.log('[Source Code Panel DEBUG] trying keys:', keysToTry);
+
+          // Try each key format
+          for (const key of keysToTry) {
+            if (stored[key]) {
+              cachedCode = stored[key];
+              console.log('[Source Code Panel DEBUG] found code with key:', key, 'codeLength:', cachedCode?.length);
+              break;
+            }
+          }
+
           // Restore to memory cache if found
           if (cachedCode) {
             if (!topWindow.__STORY_UI_GENERATED_CODE__) {
@@ -357,6 +512,11 @@ const SourceCodePanel: React.FC<{ active?: boolean }> = ({ active }) => {
         }
       }
 
+      console.log('[Source Code Panel DEBUG] final result:', {
+        foundCode: !!cachedCode,
+        codeLength: cachedCode?.length,
+      });
+
       if (cachedCode) {
         setSourceCode(cachedCode);
       } else {
@@ -441,9 +601,6 @@ const SourceCodePanel: React.FC<{ active?: boolean }> = ({ active }) => {
     );
   }
 
-  // Check if extraction was successful (displayCode is different from sourceCode)
-  const hasUsageCode = displayCode !== sourceCode;
-
   return (
     <div style={styles.container}>
       <div style={styles.header}>