From 8a4492d5abe46d9f2e208a971c2b155cabed9772 Mon Sep 17 00:00:00 2001
From: Ian Clanton-Thuon <iclanton@users.noreply.github.com>
Date: Thu, 16 Jan 2025 14:18:57 -0800
Subject: [PATCH] Add support for the --docs parameter to the Heft storybook
 plugin. (#5080)

---
 ...dd-docs-to-storybook_2025-01-16-21-48.json |  10 ++
 .../heft-storybook-plugin/heft-plugin.json    |   5 +
 .../src/StorybookPlugin.ts                    | 131 +++++++++++-------
 3 files changed, 96 insertions(+), 50 deletions(-)
 create mode 100644 common/changes/@rushstack/heft-storybook-plugin/add-docs-to-storybook_2025-01-16-21-48.json

diff --git a/common/changes/@rushstack/heft-storybook-plugin/add-docs-to-storybook_2025-01-16-21-48.json b/common/changes/@rushstack/heft-storybook-plugin/add-docs-to-storybook_2025-01-16-21-48.json
new file mode 100644
index 00000000000..6da1ab942b1
--- /dev/null
+++ b/common/changes/@rushstack/heft-storybook-plugin/add-docs-to-storybook_2025-01-16-21-48.json
@@ -0,0 +1,10 @@
+{
+  "changes": [
+    {
+      "packageName": "@rushstack/heft-storybook-plugin",
+      "comment": "Add support for the `--docs` parameter.",
+      "type": "minor"
+    }
+  ],
+  "packageName": "@rushstack/heft-storybook-plugin"
+}
\ No newline at end of file
diff --git a/heft-plugins/heft-storybook-plugin/heft-plugin.json b/heft-plugins/heft-storybook-plugin/heft-plugin.json
index 2b552d232db..c961ec466f9 100644
--- a/heft-plugins/heft-storybook-plugin/heft-plugin.json
+++ b/heft-plugins/heft-storybook-plugin/heft-plugin.json
@@ -18,6 +18,11 @@
           "longName": "--storybook-test",
           "description": "Executes a stripped down build-storybook for testing purposes.",
           "parameterKind": "flag"
+        },
+        {
+          "longName": "--docs",
+          "description": "Execute storybook in docs mode.",
+          "parameterKind": "flag"
         }
       ]
     }
diff --git a/heft-plugins/heft-storybook-plugin/src/StorybookPlugin.ts b/heft-plugins/heft-storybook-plugin/src/StorybookPlugin.ts
index 9d396767630..40b16a949e3 100644
--- a/heft-plugins/heft-storybook-plugin/src/StorybookPlugin.ts
+++ b/heft-plugins/heft-storybook-plugin/src/StorybookPlugin.ts
@@ -157,7 +157,9 @@ export interface IStorybookPluginOptions {
   captureWebpackStats?: boolean;
 }
 
-interface IRunStorybookOptions {
+interface IRunStorybookOptions extends IPrepareStorybookOptions {
+  logger: IScopedLogger;
+  isServeMode: boolean;
   workingDirectory: string;
   resolvedModulePath: string;
   outputFolder: string | undefined;
@@ -165,6 +167,15 @@ interface IRunStorybookOptions {
   verbose: boolean;
 }
 
+interface IPrepareStorybookOptions extends IStorybookPluginOptions {
+  logger: IScopedLogger;
+  taskSession: IHeftTaskSession;
+  heftConfiguration: HeftConfiguration;
+  isServeMode: boolean;
+  isTestMode: boolean;
+  isDocsMode: boolean;
+}
+
 const DEFAULT_STORYBOOK_VERSION: StorybookCliVersion = StorybookCliVersion.STORYBOOK7;
 const DEFAULT_STORYBOOK_CLI_CONFIG: Record<StorybookCliVersion, IStorybookCliCallingConfig> = {
   [StorybookCliVersion.STORYBOOK6]: {
@@ -190,14 +201,12 @@ const DEFAULT_STORYBOOK_CLI_CONFIG: Record<StorybookCliVersion, IStorybookCliCal
   }
 };
 
+const STORYBOOK_FLAG_NAME: '--storybook' = '--storybook';
 const STORYBOOK_TEST_FLAG_NAME: '--storybook-test' = '--storybook-test';
+const DOCS_FLAG_NAME: '--docs' = '--docs';
 
 /** @public */
 export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPluginOptions> {
-  private _logger!: IScopedLogger;
-  private _isServeMode: boolean = false;
-  private _isTestMode: boolean = false;
-
   /**
    * Generate typings for Sass files before TypeScript compilation.
    */
@@ -206,11 +215,12 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
     heftConfiguration: HeftConfiguration,
     options: IStorybookPluginOptions
   ): void {
-    this._logger = taskSession.logger;
+    const logger: IScopedLogger = taskSession.logger;
     const storybookParameter: CommandLineFlagParameter =
-      taskSession.parameters.getFlagParameter('--storybook');
+      taskSession.parameters.getFlagParameter(STORYBOOK_FLAG_NAME);
     const storybookTestParameter: CommandLineFlagParameter =
       taskSession.parameters.getFlagParameter(STORYBOOK_TEST_FLAG_NAME);
+    const docsParameter: CommandLineFlagParameter = taskSession.parameters.getFlagParameter(DOCS_FLAG_NAME);
 
     const parseResult: IParsedPackageNameOrError = PackageName.tryParse(options.storykitPackageName);
     if (parseResult.error) {
@@ -221,27 +231,23 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
       );
     }
 
-    if (storybookTestParameter.value) {
-      this._isTestMode = true;
-    }
-
     // Only tap if the --storybook flag is present.
     if (storybookParameter.value) {
       const configureWebpackTap: () => Promise<false> = async () => {
         // Discard Webpack's configuration to prevent Webpack from running
-        this._logger.terminal.writeLine(
+        logger.terminal.writeLine(
           'The command line includes "--storybook", redirecting Webpack to Storybook'
         );
         return false;
       };
 
+      let isServeMode: boolean = false;
       taskSession.requestAccessToPluginByName(
         '@rushstack/heft-webpack4-plugin',
         WEBPACK4_PLUGIN_NAME,
         (accessor: IWebpack4PluginAccessor) => {
-          if (accessor.parameters.isServeMode) {
-            this._isServeMode = true;
-          }
+          isServeMode = accessor.parameters.isServeMode;
+
           // Discard Webpack's configuration to prevent Webpack from running only when performing Storybook build
           accessor.hooks.onLoadConfiguration.tapPromise(PLUGIN_NAME, configureWebpackTap);
         }
@@ -251,31 +257,37 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
         '@rushstack/heft-webpack5-plugin',
         WEBPACK5_PLUGIN_NAME,
         (accessor: IWebpack5PluginAccessor) => {
-          if (accessor.parameters.isServeMode) {
-            this._isServeMode = true;
-          }
+          isServeMode = accessor.parameters.isServeMode;
+
           // Discard Webpack's configuration to prevent Webpack from running only when performing Storybook build
           accessor.hooks.onLoadConfiguration.tapPromise(PLUGIN_NAME, configureWebpackTap);
         }
       );
 
       taskSession.hooks.run.tapPromise(PLUGIN_NAME, async (runOptions: IHeftTaskRunHookOptions) => {
-        const runStorybookOptions: IRunStorybookOptions = await this._prepareStorybookAsync(
+        const runStorybookOptions: IRunStorybookOptions = await this._prepareStorybookAsync({
+          logger,
           taskSession,
           heftConfiguration,
-          options
-        );
+          isServeMode,
+          isTestMode: storybookTestParameter.value,
+          isDocsMode: docsParameter.value,
+          ...options
+        });
         await this._runStorybookAsync(runStorybookOptions, options);
       });
     }
   }
 
-  private async _prepareStorybookAsync(
-    taskSession: IHeftTaskSession,
-    heftConfiguration: HeftConfiguration,
-    options: IStorybookPluginOptions
-  ): Promise<IRunStorybookOptions> {
-    const { storykitPackageName, staticBuildOutputFolder } = options;
+  private async _prepareStorybookAsync(options: IPrepareStorybookOptions): Promise<IRunStorybookOptions> {
+    const {
+      logger,
+      taskSession,
+      heftConfiguration,
+      storykitPackageName,
+      staticBuildOutputFolder,
+      isTestMode
+    } = options;
     const storybookCliVersion: `${StorybookCliVersion}` = this._getStorybookVersion(options);
     const storyBookCliConfig: IStorybookCliCallingConfig = DEFAULT_STORYBOOK_CLI_CONFIG[storybookCliVersion];
     const cliPackageName: string = options.cliPackageName ?? storyBookCliConfig.packageName;
@@ -283,11 +295,11 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
       ? StorybookBuildMode.WATCH
       : StorybookBuildMode.BUILD;
 
-    if (buildMode === StorybookBuildMode.WATCH && this._isTestMode) {
+    if (buildMode === StorybookBuildMode.WATCH && isTestMode) {
       throw new Error(`The ${STORYBOOK_TEST_FLAG_NAME} flag is not supported in watch mode`);
     }
     if (
-      this._isTestMode &&
+      isTestMode &&
       (storybookCliVersion === StorybookCliVersion.STORYBOOK6 ||
         storybookCliVersion === StorybookCliVersion.STORYBOOK7)
     ) {
@@ -296,7 +308,7 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
       );
     }
 
-    this._logger.terminal.writeVerboseLine(`Probing for "${storykitPackageName}"`);
+    logger.terminal.writeVerboseLine(`Probing for "${storykitPackageName}"`);
     // Example: "/path/to/my-project/node_modules/my-storykit"
     let storykitFolderPath: string;
     try {
@@ -308,9 +320,9 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
       throw new Error(`The ${taskSession.taskName} task cannot start: ` + (ex as Error).message);
     }
 
-    this._logger.terminal.writeVerboseLine(`Found "${storykitPackageName}" in ` + storykitFolderPath);
+    logger.terminal.writeVerboseLine(`Found "${storykitPackageName}" in ` + storykitFolderPath);
 
-    this._logger.terminal.writeVerboseLine(`Probing for "${cliPackageName}" in "${storykitPackageName}"`);
+    logger.terminal.writeVerboseLine(`Probing for "${cliPackageName}" in "${storykitPackageName}"`);
     // Example: "/path/to/my-project/node_modules/my-storykit/node_modules/@storybook/cli"
     let storyBookCliPackage: string;
     try {
@@ -322,7 +334,7 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
       throw new Error(`The ${taskSession.taskName} task cannot start: ` + (ex as Error).message);
     }
 
-    this._logger.terminal.writeVerboseLine(`Found "${cliPackageName}" in ` + storyBookCliPackage);
+    logger.terminal.writeVerboseLine(`Found "${cliPackageName}" in ` + storyBookCliPackage);
 
     const storyBookPackagePackageJsonFile: string = path.join(storyBookCliPackage, FileConstants.PackageJson);
     const packageJson: IPackageJson = await JsonFile.loadAsync(storyBookPackagePackageJsonFile);
@@ -333,7 +345,7 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
     }
     const [moduleExecutableName, ...moduleDefaultArgs] = storyBookCliConfig.command[buildMode];
     const modulePath: string | undefined = packageJson.bin[moduleExecutableName];
-    this._logger.terminal.writeVerboseLine(
+    logger.terminal.writeVerboseLine(
       `Found storybook "${modulePath}" for "${buildMode}" mode in "${cliPackageName}"`
     );
 
@@ -353,12 +365,12 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
       buildMode === StorybookBuildMode.WATCH ? undefined : staticBuildOutputFolder;
 
     if (!modulePath) {
-      this._logger.terminal.writeVerboseLine(
+      logger.terminal.writeVerboseLine(
         'No matching module path option specified in heft.json, so bundling will proceed without Storybook'
       );
     }
 
-    this._logger.terminal.writeVerboseLine(`Resolving modulePath "${modulePath}"`);
+    logger.terminal.writeVerboseLine(`Resolving modulePath "${modulePath}"`);
     let resolvedModulePath: string;
     try {
       resolvedModulePath = Import.resolveModule({
@@ -368,7 +380,7 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
     } catch (ex) {
       throw new Error(`The ${taskSession.taskName} task cannot start: ` + (ex as Error).message);
     }
-    this._logger.terminal.writeVerboseLine(`Resolved modulePath is "${resolvedModulePath}"`);
+    logger.terminal.writeVerboseLine(`Resolved modulePath is "${resolvedModulePath}"`);
 
     // Example: "/path/to/my-project/.storybook"
     const dotStorybookFolderPath: string = `${heftConfiguration.buildFolderPath}/.storybook`;
@@ -390,6 +402,7 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
     });
 
     return {
+      ...options,
       workingDirectory: heftConfiguration.buildFolderPath,
       resolvedModulePath,
       moduleDefaultArgs,
@@ -402,10 +415,10 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
     runStorybookOptions: IRunStorybookOptions,
     options: IStorybookPluginOptions
   ): Promise<void> {
-    const { resolvedModulePath, verbose } = runStorybookOptions;
+    const { logger, resolvedModulePath, verbose, isServeMode, isTestMode, isDocsMode } = runStorybookOptions;
     let { workingDirectory, outputFolder } = runStorybookOptions;
-    this._logger.terminal.writeLine('Running Storybook compilation');
-    this._logger.terminal.writeVerboseLine(`Loading Storybook module "${resolvedModulePath}"`);
+    logger.terminal.writeLine('Running Storybook compilation');
+    logger.terminal.writeVerboseLine(`Loading Storybook module "${resolvedModulePath}"`);
     const storybookCliVersion: `${StorybookCliVersion}` = this._getStorybookVersion(options);
 
     /**
@@ -424,7 +437,7 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
         baseFolderPath: workingDirectory
       });
 
-      this._logger.terminal.writeVerboseLine(`Changing Storybook working directory to "${workingDirectory}"`);
+      logger.terminal.writeVerboseLine(`Changing Storybook working directory to "${workingDirectory}"`);
     }
 
     const storybookArgs: string[] = runStorybookOptions.moduleDefaultArgs ?? [];
@@ -432,26 +445,34 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
     if (outputFolder) {
       storybookArgs.push('--output-dir', outputFolder);
     }
+
     if (options.captureWebpackStats) {
       storybookArgs.push('--webpack-stats-json');
     }
+
     if (!verbose) {
       storybookArgs.push('--quiet');
     }
-    if (this._isTestMode) {
+
+    if (isTestMode) {
       storybookArgs.push('--test');
     }
 
-    if (this._isServeMode) {
+    if (isDocsMode) {
+      storybookArgs.push('--docs');
+    }
+
+    if (isServeMode) {
       // Instantiate storybook runner synchronously for incremental builds
       // this ensure that the process is not killed when heft watcher detects file changes
       this._invokeSync(
+        logger,
         resolvedModulePath,
         storybookArgs,
         storybookCliVersion === StorybookCliVersion.STORYBOOK8
       );
     } else {
-      await this._invokeAsSubprocessAsync(resolvedModulePath, storybookArgs, workingDirectory);
+      await this._invokeAsSubprocessAsync(logger, resolvedModulePath, storybookArgs, workingDirectory);
     }
   }
 
@@ -462,7 +483,12 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
    * @param cwd - working directory
    * @returns
    */
-  private async _invokeAsSubprocessAsync(command: string, args: string[], cwd: string): Promise<void> {
+  private async _invokeAsSubprocessAsync(
+    logger: IScopedLogger,
+    command: string,
+    args: string[],
+    cwd: string
+  ): Promise<void> {
     return await new Promise<void>((resolve, reject) => {
       const storybookEnv: NodeJS.ProcessEnv = { ...process.env };
       const forkedProcess: child_process.ChildProcess = child_process.fork(command, args, {
@@ -479,12 +505,12 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
       if (childPid === undefined) {
         throw new InternalError(`Failed to spawn child process`);
       }
-      this._logger.terminal.writeVerboseLine(`Started storybook process #${childPid}`);
+      logger.terminal.writeVerboseLine(`Started storybook process #${childPid}`);
 
       // Apply the pipe here instead of doing it in the forked process args due to a bug in Node
       // We will output stderr to the normal stdout stream since all output is piped through
       // stdout. We have to rely on the exit code to determine if there was an error.
-      const terminal: ITerminal = this._logger.terminal;
+      const terminal: ITerminal = logger.terminal;
       const terminalOutStream: TerminalStreamWritable = new TerminalStreamWritable({
         terminal,
         severity: TerminalProviderSeverity.log
@@ -520,8 +546,13 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
    * @param args - storybook args
    * @param cwd - working directory
    */
-  private _invokeSync(command: string, args: string[], patchNpmConfigUserAgent: boolean): void {
-    this._logger.terminal.writeLine('Launching ' + command);
+  private _invokeSync(
+    logger: IScopedLogger,
+    command: string,
+    args: string[],
+    patchNpmConfigUserAgent: boolean
+  ): void {
+    logger.terminal.writeLine('Launching ' + command);
 
     // simulate storybook cli command
     const originalArgv: string[] = process.argv;
@@ -547,7 +578,7 @@ export default class StorybookPlugin implements IHeftTaskPlugin<IStorybookPlugin
     // restore original heft process argv
     process.argv = originalArgv;
 
-    this._logger.terminal.writeVerboseLine('Completed synchronous portion of launching startupModulePath');
+    logger.terminal.writeVerboseLine('Completed synchronous portion of launching startupModulePath');
   }
 
   private _getStorybookVersion(options: IStorybookPluginOptions): `${StorybookCliVersion}` {