CSF-Factories: Preserve leading comments when adding imports

code/lib/cli-storybook/src/codemod/helpers/config-to-csf-factory.test.ts

@@ -58,6 +58,28 @@ describe('main/preview codemod: general parsing functionality', () => {
       });
     `);
   });
+
+  it('should preserve leading comments when adding import', async () => {
+    await expect(
+      transform(dedent`
+        // @ts-check
+        /** @license MIT */
+        export default {
+          stories: ['../src/**/*.stories.@(js|jsx|ts|tsx)'],
+          framework: '@storybook/react-vite',
+        };
+      `)
+    ).resolves.toMatchInlineSnapshot(`
+      // @ts-check
+      /** @license MIT */
+      import { defineMain } from '@storybook/react-vite/node';
+
+      export default defineMain({
+        stories: ['../src/**/*.stories.@(js|jsx|ts|tsx)'],
+        framework: '@storybook/react-vite',
+      });
+    `);
+  });
   it('should wrap defineMain call from const declared default export with different type annotations', async () => {
     const typedVariants = [
       'export default config;',

code/lib/cli-storybook/src/codemod/helpers/config-to-csf-factory.ts

@@ -7,6 +7,7 @@ import picocolors from 'picocolors';
 
 import type { FileInfo } from '../../automigrate/codemod';
 import {
+  addImportToTop,
   cleanupTypeImports,
   getConfigProperties,
   removeExportDeclarations,
@@ -197,7 +198,7 @@ export async function configToCsfFactory(
     }
   } else {
     // if not, add import { defineMain } from '@storybook/framework'
-    programNode.body.unshift(configImport);
+    addImportToTop(programNode, configImport);
   }
 
   // Remove type imports – now inferred – from @storybook/* packages

code/lib/cli-storybook/src/codemod/helpers/csf-factories-utils.ts

@@ -119,3 +119,29 @@ export function getConfigProperties(
   // error TS4058: Return type of exported function has or is using name 'ObjectProperty' from external module "/tmp/storybook/code/core/dist/babel/index" but cannot be named.
   return properties as any;
 }
+
+/**
+ * Adds an import declaration to the beginning of the program while preserving any leading comments
+ * (like license headers or @ts-check directives).
+ *
+ * When using `programNode.body.unshift()`, the import would be placed before any leading comments
+ * attached to the first node. This function transfers those comments to the new import so they
+ * remain at the top of the file.
+ *
+ * Note: We use the `comments` property (used by recast for printing) rather than `leadingComments`
+ * (used by babel internally) to ensure proper output formatting.
+ */
+export function addImportToTop(programNode: t.Program, importDecl: t.ImportDeclaration): void {
+  const firstNode = programNode.body[0] as t.Node & { comments?: t.Comment[] };
+
+  if (firstNode && firstNode.leadingComments && firstNode.leadingComments.length > 0) {
+    // Transfer leading comments from the first node to the import using 'comments' property
+    // which is what recast uses for printing (not 'leadingComments')
+    (importDecl as t.Node & { comments?: t.Comment[] }).comments = firstNode.leadingComments;
+    // Clear comments from the original first node to avoid duplication
+    firstNode.leadingComments = [];
+    firstNode.comments = [];
+  }
+
+  programNode.body.unshift(importDecl);
+}

code/lib/cli-storybook/src/codemod/helpers/story-to-csf-factory.test.ts

@@ -45,6 +45,28 @@ describe('stories codemod', () => {
       `);
     });
 
+    it('should preserve leading comments when adding import', async () => {
+      await expect(
+        transform(dedent`
+            // @ts-check
+            /**
+             * @license MIT
+             * Copyright 2024
+             */
+            const meta = { title: 'Component' };
+            export default meta;
+            export const A = {};
+          `)
+      ).resolves.toMatchInlineSnapshot(`
+        // @ts-check
+        /** @license MIT Copyright 2024 */
+        import preview from '#.storybook/preview';
+
+        const meta = preview.meta({ title: 'Component' });
+        export const A = meta.story();
+      `);
+    });
+
     it('should transform and wrap inline default exported meta', async () => {
       await expect(
         transform(dedent`

code/lib/cli-storybook/src/codemod/helpers/story-to-csf-factory.ts

@@ -5,7 +5,7 @@ import { logger } from 'storybook/internal/node-logger';
 import path from 'path';
 
 import type { FileInfo } from '../../automigrate/codemod';
-import { cleanupTypeImports } from './csf-factories-utils';
+import { addImportToTop, cleanupTypeImports } from './csf-factories-utils';
 import { removeUnusedTypes } from './remove-unused-types';
 
 const typesDisallowList = [
@@ -324,7 +324,7 @@ export async function storyToCsfFactory(
       [t.importDefaultSpecifier(t.identifier(sbConfigImportName))],
       t.stringLiteral(previewPath)
     );
-    programNode.body.unshift(configImport);
+    addImportToTop(programNode, configImport);
   }
 
   removeUnusedTypes(programNode, csf._ast);