@tpitre/story-ui 3.0.0 → 3.2.0

package/dist/story-generator/postgresStoryService.js

@@ -0,0 +1,240 @@
+import pg from 'pg';
+const { Pool } = pg;
+/**
+ * PostgreSQL Story Service
+ * Persistent story storage using PostgreSQL database
+ * Designed for Railway PostgreSQL deployments
+ */
+export class PostgresStoryService {
+    constructor(connectionString) {
+        this.initialized = false;
+        this.pool = new Pool({
+            connectionString,
+            ssl: process.env.NODE_ENV === 'production' ? { rejectUnauthorized: false } : false,
+            max: 10,
+            idleTimeoutMillis: 30000,
+            connectionTimeoutMillis: 10000,
+        });
+    }
+    /**
+     * Initialize database schema
+     */
+    async initialize() {
+        if (this.initialized)
+            return;
+        const client = await this.pool.connect();
+        try {
+            await client.query(`
+        CREATE TABLE IF NOT EXISTS stories (
+          id VARCHAR(255) PRIMARY KEY,
+          title VARCHAR(500) NOT NULL,
+          description TEXT,
+          content TEXT NOT NULL,
+          prompt TEXT,
+          components TEXT[],
+          created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
+          last_accessed TIMESTAMP WITH TIME ZONE DEFAULT NOW()
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_stories_created_at ON stories(created_at DESC);
+        CREATE INDEX IF NOT EXISTS idx_stories_last_accessed ON stories(last_accessed DESC);
+      `);
+            this.initialized = true;
+            console.log('✅ PostgreSQL story service initialized');
+        }
+        finally {
+            client.release();
+        }
+    }
+    /**
+     * Store a generated story
+     */
+    async storeStory(story) {
+        await this.initialize();
+        const client = await this.pool.connect();
+        try {
+            await client.query(`INSERT INTO stories (id, title, description, content, prompt, components, created_at, last_accessed)
+         VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+         ON CONFLICT (id) DO UPDATE SET
+           title = EXCLUDED.title,
+           description = EXCLUDED.description,
+           content = EXCLUDED.content,
+           prompt = EXCLUDED.prompt,
+           components = EXCLUDED.components,
+           last_accessed = EXCLUDED.last_accessed`, [
+                story.id,
+                story.title,
+                story.description,
+                story.content,
+                story.prompt || null,
+                story.components || [],
+                story.createdAt || new Date(),
+                new Date()
+            ]);
+            console.log(`✅ Stored story in PostgreSQL: ${story.id}`);
+        }
+        finally {
+            client.release();
+        }
+    }
+    /**
+     * Retrieve a story by ID
+     */
+    async getStory(id) {
+        await this.initialize();
+        const client = await this.pool.connect();
+        try {
+            // Update last_accessed and return the story
+            const result = await client.query(`UPDATE stories SET last_accessed = NOW()
+         WHERE id = $1
+         RETURNING id, title, description, content, prompt, components, created_at, last_accessed`, [id]);
+            if (result.rows.length === 0) {
+                return null;
+            }
+            const row = result.rows[0];
+            return {
+                id: row.id,
+                title: row.title,
+                description: row.description,
+                content: row.content,
+                prompt: row.prompt,
+                components: row.components || [],
+                createdAt: new Date(row.created_at),
+                lastAccessed: new Date(row.last_accessed)
+            };
+        }
+        finally {
+            client.release();
+        }
+    }
+    /**
+     * Get all stored stories
+     */
+    async getAllStories() {
+        await this.initialize();
+        const client = await this.pool.connect();
+        try {
+            const result = await client.query(`SELECT id, title, description, content, prompt, components, created_at, last_accessed
+         FROM stories
+         ORDER BY created_at DESC`);
+            return result.rows.map(row => ({
+                id: row.id,
+                title: row.title,
+                description: row.description,
+                content: row.content,
+                prompt: row.prompt,
+                components: row.components || [],
+                createdAt: new Date(row.created_at),
+                lastAccessed: new Date(row.last_accessed)
+            }));
+        }
+        finally {
+            client.release();
+        }
+    }
+    /**
+     * Delete a story by ID
+     */
+    async deleteStory(id) {
+        await this.initialize();
+        const client = await this.pool.connect();
+        try {
+            const result = await client.query(`DELETE FROM stories WHERE id = $1 RETURNING id`, [id]);
+            const deleted = result.rowCount !== null && result.rowCount > 0;
+            if (deleted) {
+                console.log(`✅ Deleted story from PostgreSQL: ${id}`);
+            }
+            return deleted;
+        }
+        finally {
+            client.release();
+        }
+    }
+    /**
+     * Clear all stories
+     */
+    async clearAllStories() {
+        await this.initialize();
+        const client = await this.pool.connect();
+        try {
+            await client.query('TRUNCATE TABLE stories');
+            console.log('✅ Cleared all stories from PostgreSQL');
+        }
+        finally {
+            client.release();
+        }
+    }
+    /**
+     * Get story content for Storybook integration
+     */
+    async getStoryContent(id) {
+        const story = await this.getStory(id);
+        return story ? story.content : null;
+    }
+    /**
+     * Get story metadata for listing
+     */
+    async getStoryMetadata() {
+        await this.initialize();
+        const client = await this.pool.connect();
+        try {
+            const result = await client.query(`SELECT id, title, description, content, created_at, last_accessed
+         FROM stories
+         ORDER BY created_at DESC`);
+            return result.rows.map(row => ({
+                id: row.id,
+                title: row.title,
+                description: row.description,
+                createdAt: new Date(row.created_at),
+                lastAccessed: new Date(row.last_accessed),
+                componentCount: this.countComponents(row.content)
+            }));
+        }
+        finally {
+            client.release();
+        }
+    }
+    /**
+     * Get storage usage statistics
+     */
+    async getStorageStats() {
+        await this.initialize();
+        const client = await this.pool.connect();
+        try {
+            const result = await client.query(`
+        SELECT
+          COUNT(*)::int as story_count,
+          COALESCE(SUM(LENGTH(content)), 0)::bigint as total_size,
+          COALESCE(AVG(LENGTH(content)), 0)::int as avg_size,
+          MIN(created_at) as oldest,
+          MAX(created_at) as newest
+        FROM stories
+      `);
+            const row = result.rows[0];
+            return {
+                storyCount: row.story_count,
+                totalSizeBytes: parseInt(row.total_size) || 0,
+                averageSizeBytes: row.avg_size || 0,
+                oldestStory: row.oldest ? new Date(row.oldest) : null,
+                newestStory: row.newest ? new Date(row.newest) : null
+            };
+        }
+        finally {
+            client.release();
+        }
+    }
+    /**
+     * Count unique components in story content
+     */
+    countComponents(content) {
+        const componentMatches = content.match(/<[A-Z][A-Za-z0-9]*\s/g);
+        return componentMatches ? new Set(componentMatches).size : 0;
+    }
+    /**
+     * Close the connection pool
+     */
+    async close() {
+        await this.pool.end();
+        console.log('PostgreSQL connection pool closed');
+    }
+}

package/dist/story-generator/storyServiceFactory.d.ts

@@ -0,0 +1,22 @@
+import type { IStoryService } from './storyServiceInterface.js';
+import { StoryUIConfig } from '../story-ui.config.js';
+/**
+ * Get or create the story service based on environment configuration
+ *
+ * - If DATABASE_URL is set, uses PostgreSQL for persistent storage
+ * - Otherwise, falls back to in-memory storage
+ */
+export declare function getStoryService(config: StoryUIConfig): Promise<IStoryService>;
+/**
+ * Check if PostgreSQL storage is configured
+ */
+export declare function isPostgresConfigured(): boolean;
+/**
+ * Get storage type description
+ */
+export declare function getStorageType(): 'postgresql' | 'memory';
+/**
+ * Reset the global service (for testing)
+ */
+export declare function resetStoryService(): Promise<void>;
+//# sourceMappingURL=storyServiceFactory.d.ts.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"storyServiceFactory.d.ts","sourceRoot":"","sources":["../../story-generator/storyServiceFactory.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,aAAa,EAA+C,MAAM,4BAA4B,CAAC;AAG7G,OAAO,EAAE,aAAa,EAAE,MAAM,uBAAuB,CAAC;AAkEtD;;;;;GAKG;AACH,wBAAsB,eAAe,CAAC,MAAM,EAAE,aAAa,GAAG,OAAO,CAAC,aAAa,CAAC,CAmBnF;AAED;;GAEG;AACH,wBAAgB,oBAAoB,IAAI,OAAO,CAE9C;AAED;;GAEG;AACH,wBAAgB,cAAc,IAAI,YAAY,GAAG,QAAQ,CAExD;AAED;;GAEG;AACH,wBAAsB,iBAAiB,IAAI,OAAO,CAAC,IAAI,CAAC,CAKvD"}

package/dist/story-generator/storyServiceFactory.js

@@ -0,0 +1,97 @@
+import { getInMemoryStoryService } from './inMemoryStoryService.js';
+import { PostgresStoryService } from './postgresStoryService.js';
+/**
+ * Async wrapper for InMemoryStoryService
+ * Makes the synchronous in-memory service conform to the async IStoryService interface
+ */
+class AsyncInMemoryStoryService {
+    constructor(config) {
+        this.innerService = getInMemoryStoryService(config);
+    }
+    async initialize() {
+        // No-op for in-memory
+    }
+    async storeStory(story) {
+        this.innerService.storeStory(story);
+    }
+    async getStory(id) {
+        return this.innerService.getStory(id);
+    }
+    async getAllStories() {
+        return this.innerService.getAllStories();
+    }
+    async deleteStory(id) {
+        return this.innerService.deleteStory(id);
+    }
+    async clearAllStories() {
+        this.innerService.clearAllStories();
+    }
+    async getStoryContent(id) {
+        return this.innerService.getStoryContent(id);
+    }
+    async getStoryMetadata() {
+        return this.innerService.getStoryMetadata();
+    }
+    async getStorageStats() {
+        const stats = this.innerService.getMemoryStats();
+        return {
+            storyCount: stats.storyCount,
+            totalSizeBytes: stats.totalSizeBytes,
+            averageSizeBytes: stats.averageSizeBytes,
+            oldestStory: stats.oldestStory,
+            newestStory: stats.newestStory
+        };
+    }
+    async close() {
+        // No-op for in-memory
+    }
+}
+/**
+ * Global story service instance
+ */
+let globalStoryService = null;
+/**
+ * Get or create the story service based on environment configuration
+ *
+ * - If DATABASE_URL is set, uses PostgreSQL for persistent storage
+ * - Otherwise, falls back to in-memory storage
+ */
+export async function getStoryService(config) {
+    if (globalStoryService) {
+        return globalStoryService;
+    }
+    const databaseUrl = process.env.DATABASE_URL;
+    if (databaseUrl) {
+        console.log('🗄️  Using PostgreSQL for story persistence');
+        const pgService = new PostgresStoryService(databaseUrl);
+        await pgService.initialize();
+        globalStoryService = pgService;
+    }
+    else {
+        console.log('💾 Using in-memory storage (no DATABASE_URL configured)');
+        globalStoryService = new AsyncInMemoryStoryService(config);
+        await globalStoryService.initialize();
+    }
+    return globalStoryService;
+}
+/**
+ * Check if PostgreSQL storage is configured
+ */
+export function isPostgresConfigured() {
+    return !!process.env.DATABASE_URL;
+}
+/**
+ * Get storage type description
+ */
+export function getStorageType() {
+    return process.env.DATABASE_URL ? 'postgresql' : 'memory';
+}
+/**
+ * Reset the global service (for testing)
+ */
+export async function resetStoryService() {
+    if (globalStoryService) {
+        await globalStoryService.close();
+        globalStoryService = null;
+    }
+}

package/dist/story-generator/storyServiceInterface.d.ts

@@ -0,0 +1,85 @@
+/**
+ * Story Service Interface
+ * Common interface for both in-memory and PostgreSQL story storage
+ */
+/**
+ * Generated story interface
+ */
+export interface GeneratedStory {
+    id: string;
+    title: string;
+    description: string;
+    content: string;
+    createdAt: Date;
+    lastAccessed: Date;
+    prompt?: string;
+    components?: string[];
+}
+/**
+ * Story metadata for listing
+ */
+export interface StoryMetadata {
+    id: string;
+    title: string;
+    description: string;
+    createdAt: Date;
+    lastAccessed: Date;
+    componentCount: number;
+}
+/**
+ * Memory/storage usage statistics
+ */
+export interface StorageStats {
+    storyCount: number;
+    totalSizeBytes: number;
+    averageSizeBytes: number;
+    oldestStory: Date | null;
+    newestStory: Date | null;
+}
+/**
+ * Story Service Interface
+ * All methods are async to support both in-memory and database backends
+ */
+export interface IStoryService {
+    /**
+     * Initialize the service (create tables, etc.)
+     */
+    initialize(): Promise<void>;
+    /**
+     * Store a generated story
+     */
+    storeStory(story: GeneratedStory): Promise<void>;
+    /**
+     * Retrieve a story by ID
+     */
+    getStory(id: string): Promise<GeneratedStory | null>;
+    /**
+     * Get all stored stories
+     */
+    getAllStories(): Promise<GeneratedStory[]>;
+    /**
+     * Delete a story by ID
+     */
+    deleteStory(id: string): Promise<boolean>;
+    /**
+     * Clear all stories
+     */
+    clearAllStories(): Promise<void>;
+    /**
+     * Get story content for Storybook integration
+     */
+    getStoryContent(id: string): Promise<string | null>;
+    /**
+     * Get story metadata for listing
+     */
+    getStoryMetadata(): Promise<StoryMetadata[]>;
+    /**
+     * Get storage usage statistics
+     */
+    getStorageStats(): Promise<StorageStats>;
+    /**
+     * Close any connections (for cleanup)
+     */
+    close(): Promise<void>;
+}
+//# sourceMappingURL=storyServiceInterface.d.ts.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"storyServiceInterface.d.ts","sourceRoot":"","sources":["../../story-generator/storyServiceInterface.ts"],"names":[],"mappings":"AAAA;;;GAGG;AAEH;;GAEG;AACH,MAAM,WAAW,cAAc;IAC7B,EAAE,EAAE,MAAM,CAAC;IACX,KAAK,EAAE,MAAM,CAAC;IACd,WAAW,EAAE,MAAM,CAAC;IACpB,OAAO,EAAE,MAAM,CAAC;IAChB,SAAS,EAAE,IAAI,CAAC;IAChB,YAAY,EAAE,IAAI,CAAC;IACnB,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,UAAU,CAAC,EAAE,MAAM,EAAE,CAAC;CACvB;AAED;;GAEG;AACH,MAAM,WAAW,aAAa;IAC5B,EAAE,EAAE,MAAM,CAAC;IACX,KAAK,EAAE,MAAM,CAAC;IACd,WAAW,EAAE,MAAM,CAAC;IACpB,SAAS,EAAE,IAAI,CAAC;IAChB,YAAY,EAAE,IAAI,CAAC;IACnB,cAAc,EAAE,MAAM,CAAC;CACxB;AAED;;GAEG;AACH,MAAM,WAAW,YAAY;IAC3B,UAAU,EAAE,MAAM,CAAC;IACnB,cAAc,EAAE,MAAM,CAAC;IACvB,gBAAgB,EAAE,MAAM,CAAC;IACzB,WAAW,EAAE,IAAI,GAAG,IAAI,CAAC;IACzB,WAAW,EAAE,IAAI,GAAG,IAAI,CAAC;CAC1B;AAED;;;GAGG;AACH,MAAM,WAAW,aAAa;IAC5B;;OAEG;IACH,UAAU,IAAI,OAAO,CAAC,IAAI,CAAC,CAAC;IAE5B;;OAEG;IACH,UAAU,CAAC,KAAK,EAAE,cAAc,GAAG,OAAO,CAAC,IAAI,CAAC,CAAC;IAEjD;;OAEG;IACH,QAAQ,CAAC,EAAE,EAAE,MAAM,GAAG,OAAO,CAAC,cAAc,GAAG,IAAI,CAAC,CAAC;IAErD;;OAEG;IACH,aAAa,IAAI,OAAO,CAAC,cAAc,EAAE,CAAC,CAAC;IAE3C;;OAEG;IACH,WAAW,CAAC,EAAE,EAAE,MAAM,GAAG,OAAO,CAAC,OAAO,CAAC,CAAC;IAE1C;;OAEG;IACH,eAAe,IAAI,OAAO,CAAC,IAAI,CAAC,CAAC;IAEjC;;OAEG;IACH,eAAe,CAAC,EAAE,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,GAAG,IAAI,CAAC,CAAC;IAEpD;;OAEG;IACH,gBAAgB,IAAI,OAAO,CAAC,aAAa,EAAE,CAAC,CAAC;IAE7C;;OAEG;IACH,eAAe,IAAI,OAAO,CAAC,YAAY,CAAC,CAAC;IAEzC;;OAEG;IACH,KAAK,IAAI,OAAO,CAAC,IAAI,CAAC,CAAC;CACxB"}

package/dist/story-generator/storyServiceInterface.js

@@ -0,0 +1,5 @@
+/**
+ * Story Service Interface
+ * Common interface for both in-memory and PostgreSQL story storage
+ */
+export {};

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

@@ -1,3 +1,8 @@
+declare global {
+    interface Window {
+        __STORY_UI_GENERATED_CODE__?: Record<string, string>;
+    }
+}
 declare function StoryUIPanel(): import("react/jsx-runtime").JSX.Element;
 export default StoryUIPanel;
 export { StoryUIPanel };

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

@@ -1 +1 @@
-{"version":3,"file":"StoryUIPanel.d.ts","sourceRoot":"","sources":["../../../templates/StoryUI/StoryUIPanel.tsx"],"names":[],"mappings":"AA85CA,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":"AAsSA,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"}

package/dist/templates/StoryUI/StoryUIPanel.js

@@ -121,12 +121,27 @@ const titleToStoryPath = (title) => {
     const kebabTitle = title.toLowerCase().replace(/[^a-z0-9]+/g, '-').replace(/^-|-$/g, '');
     return `generated-${kebabTitle}--default`;
 };
+// Helper to store generated code for the Source Code panel to display
+const storeGeneratedCode = (title, code) => {
+    const storyPath = titleToStoryPath(title);
+    const topWindow = window.top || window;
+    // Store code in the top window so it's accessible from manager frame
+    if (!topWindow.__STORY_UI_GENERATED_CODE__) {
+        topWindow.__STORY_UI_GENERATED_CODE__ = {};
+    }
+    topWindow.__STORY_UI_GENERATED_CODE__[storyPath] = code;
+    console.log(`[Story UI] Stored code for story "${storyPath}" in window cache`);
+};
 // Helper to navigate to a newly created story after generation completes
 // In dev mode with HMR, this prevents the "Couldn't find story after HMR" error
 // In all modes, this provides a better UX by auto-navigating to the new story
-const navigateToNewStory = (title, delayMs = 4000) => {
+const navigateToNewStory = (title, code, delayMs = 1500) => {
     const storyPath = titleToStoryPath(title);
     console.log(`[Story UI] Will navigate to story "${storyPath}" in ${delayMs}ms...`);
+    // Store the code for the Source Code panel if provided
+    if (code) {
+        storeGeneratedCode(title, code);
+    }
     setTimeout(() => {
         // Navigate the TOP window (parent Storybook UI), not the iframe
         // The Story UI panel runs inside an iframe, so we need window.top to escape it
@@ -1438,7 +1453,7 @@ function StoryUIPanel() {
             // Auto-navigate to the newly created story after HMR processes the file
             // This prevents the "Couldn't find story after HMR" error by refreshing
             // after the file system has been updated and HMR has processed the change
-            navigateToNewStory(chatTitle);
+            navigateToNewStory(chatTitle, completion.code);
         }
     }, [activeChatId, activeTitle, conversation.length]);
     const handleSend = async (e) => {
@@ -1631,7 +1646,7 @@ function StoryUIPanel() {
                         saveChats(chats);
                         setRecentChats(chats);
                         // Auto-navigate to the newly created story
-                        navigateToNewStory(chatTitle);
+                        navigateToNewStory(chatTitle, data.code);
                     }
                 }
                 catch (fallbackErr) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tpitre/story-ui",
-  "version": "3.0.0",
+  "version": "3.2.0",
   "description": "AI-powered Storybook story generator with dynamic component discovery",
   "type": "module",
   "main": "dist/index.js",
@@ -36,9 +36,7 @@
     "release:dry-run": "semantic-release --dry-run",
     "commit": "cz",
     "prepare": "husky",
-    "story-ui": "story-ui start --port 4001",
-    "deploy:edge": "./scripts/deploy-edge.sh",
-    "deploy:edge:dry-run": "./scripts/deploy-edge.sh --dry-run"
+    "story-ui": "story-ui start --port 4001"
   },
   "keywords": [
     "storybook",
@@ -78,6 +76,7 @@
     "@emotion/styled": "^11.14.1",
     "@modelcontextprotocol/sdk": "^0.5.0",
     "@mui/material": "^7.2.0",
+    "@types/pg": "^8.15.6",
     "chalk": "^5.3.0",
     "commander": "^11.0.0",
     "cors": "^2.8.5",
@@ -86,6 +85,7 @@
     "glob": "^11.0.3",
     "inquirer": "^9.2.0",
     "node-fetch": "^2.6.7",
+    "pg": "^8.16.3",
     "typescript": "^5.8.3",
     "zod": "^3.22.4"
   },
@@ -105,11 +105,11 @@
     "@types/node": "^20.4.2",
     "@types/node-fetch": "^2.6.12",
     "commitizen": "^4.3.1",
+    "concurrently": "^8.2.0",
     "cz-conventional-changelog": "^3.3.0",
     "husky": "^9.1.7",
     "semantic-release": "^24.2.0",
-    "ts-node": "^10.9.2",
-    "concurrently": "^8.2.0"
+    "ts-node": "^10.9.2"
   },
   "peerDependencies": {
     "@storybook/react": ">=6.0.0",

package/templates/StoryUI/StoryUIPanel.tsx

@@ -291,13 +291,38 @@ const titleToStoryPath = (title: string): string => {
   return `generated-${kebabTitle}--default`;
 };
 
+// Extend window to include our code cache for the Source Code panel
+declare global {
+  interface Window {
+    __STORY_UI_GENERATED_CODE__?: Record<string, string>;
+  }
+}
+
+// Helper to store generated code for the Source Code panel to display
+const storeGeneratedCode = (title: string, code: string) => {
+  const storyPath = titleToStoryPath(title);
+  const topWindow = window.top || window;
+
+  // Store code in the top window so it's accessible from manager frame
+  if (!topWindow.__STORY_UI_GENERATED_CODE__) {
+    topWindow.__STORY_UI_GENERATED_CODE__ = {};
+  }
+  topWindow.__STORY_UI_GENERATED_CODE__[storyPath] = code;
+  console.log(`[Story UI] Stored code for story "${storyPath}" in window cache`);
+};
+
 // Helper to navigate to a newly created story after generation completes
 // In dev mode with HMR, this prevents the "Couldn't find story after HMR" error
 // In all modes, this provides a better UX by auto-navigating to the new story
-const navigateToNewStory = (title: string, delayMs: number = 4000) => {
+const navigateToNewStory = (title: string, code?: string, delayMs: number = 1500) => {
   const storyPath = titleToStoryPath(title);
   console.log(`[Story UI] Will navigate to story "${storyPath}" in ${delayMs}ms...`);
 
+  // Store the code for the Source Code panel if provided
+  if (code) {
+    storeGeneratedCode(title, code);
+  }
+
   setTimeout(() => {
     // Navigate the TOP window (parent Storybook UI), not the iframe
     // The Story UI panel runs inside an iframe, so we need window.top to escape it
@@ -1927,7 +1952,7 @@ function StoryUIPanel() {
       // Auto-navigate to the newly created story after HMR processes the file
       // This prevents the "Couldn't find story after HMR" error by refreshing
       // after the file system has been updated and HMR has processed the change
-      navigateToNewStory(chatTitle);
+      navigateToNewStory(chatTitle, completion.code);
     }
   }, [activeChatId, activeTitle, conversation.length]);
 
@@ -2137,7 +2162,7 @@ function StoryUIPanel() {
             setRecentChats(chats);
 
             // Auto-navigate to the newly created story
-            navigateToNewStory(chatTitle);
+            navigateToNewStory(chatTitle, data.code);
           }
         } catch (fallbackErr: unknown) {
           const errorMessage = fallbackErr instanceof Error ? fallbackErr.message : 'Unknown error';