From 9df7b053d1bbf38dd85738d94fe24c816e9350a2 Mon Sep 17 00:00:00 2001
From: fi3ework <fi3ework@gmail.com>
Date: Tue, 28 May 2024 22:48:51 +0800
Subject: [PATCH] feat(builder): support emit stats

---
 packages/builder-rsbuild/src/index.ts | 42 ++++++++++++++++++++-------
 packages/builder-rsbuild/src/types.ts |  6 ++++
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/packages/builder-rsbuild/src/index.ts b/packages/builder-rsbuild/src/index.ts
index 36b9ae9..1f3ed9e 100644
--- a/packages/builder-rsbuild/src/index.ts
+++ b/packages/builder-rsbuild/src/index.ts
@@ -4,10 +4,7 @@ import type { Options } from '@storybook/types'
 import { dirname, join, parse } from 'path'
 import express from 'express'
 import fs from 'fs-extra'
-import {
-  NoStatsForViteDevError,
-  WebpackInvocationError,
-} from '@storybook/core-events/server-errors'
+import { WebpackInvocationError } from '@storybook/core-events/server-errors'
 import type { RsbuildBuilder } from './types'
 import rsbuildConfig, {
   type RsbuildBuilderOptions,
@@ -19,6 +16,10 @@ export * from './types'
 export * from './preview/virtual-module-mapping'
 
 type RsbuildDevServer = Awaited<ReturnType<typeof createDevServer>>
+type StatsOrMultiStats = Parameters<rsbuildReal.OnAfterBuildFn>[0]['stats']
+export type Stats = NonNullable<
+  Exclude<StatsOrMultiStats, { stats: unknown[] }>
+>
 
 export const printDuration = (startTime: [number, number]) =>
   prettyTime(process.hrtime(startTime))
@@ -98,6 +99,16 @@ export const start: RsbuildBuilder['start'] = async ({
   })
 
   const rsbuildServer = await rsbuildBuild.createDevServer()
+
+  const waitFirstCompileDone = new Promise<StatsOrMultiStats>((resolve) => {
+    rsbuildBuild.onDevCompileDone(({ stats, isFirstCompile }) => {
+      if (!isFirstCompile) {
+        return
+      }
+      resolve(stats)
+    })
+  })
+
   server = rsbuildServer
 
   if (!rsbuildBuild) {
@@ -117,19 +128,20 @@ export const start: RsbuildBuilder['start'] = async ({
 
   router.use(rsbuildServer.middlewares)
   storybookServer.on('upgrade', rsbuildServer.onHTTPUpgrade)
+  const stats = await waitFirstCompileDone
 
   return {
     bail,
-    stats: {
-      toJson: () => {
-        throw new NoStatsForViteDevError()
-      },
-    },
+    stats,
     totalTime: process.hrtime(startTime),
   }
 }
 
-export const build = async ({ options }: BuilderStartOptions) => {
+// explicit type annotation to bypass TypeScript check
+// see: https://github.com/microsoft/TypeScript/issues/47663#issuecomment-1519138189
+export const build: ({
+  options,
+}: BuilderStartOptions) => Promise<Stats> = async ({ options }) => {
   const { createRsbuild } = await executor.get(options)
   const config = await getConfig(options)
   const rsbuildBuild = await createRsbuild({
@@ -140,6 +152,11 @@ export const build = async ({ options }: BuilderStartOptions) => {
   const previewResolvedDir = getAbsolutePath('@storybook/preview')
   const previewDirOrigin = join(previewResolvedDir, 'dist')
   const previewDirTarget = join(options.outputDir || '', `sb-preview`)
+  let stats: Stats
+
+  rsbuildBuild.onAfterBuild((params) => {
+    stats = params.stats as Stats
+  })
 
   const previewFiles = fs.copy(previewDirOrigin, previewDirTarget, {
     filter: (src) => {
@@ -151,7 +168,12 @@ export const build = async ({ options }: BuilderStartOptions) => {
     },
   })
 
+  rsbuildBuild.onAfterBuild((params) => {
+    stats = params.stats as Stats
+  })
+
   await Promise.all([rsbuildBuild.build(), previewFiles])
+  return stats!
 }
 
 export const corePresets = [join(__dirname, './preview-preset.js')]
diff --git a/packages/builder-rsbuild/src/types.ts b/packages/builder-rsbuild/src/types.ts
index 54fbe64..3fa5922 100644
--- a/packages/builder-rsbuild/src/types.ts
+++ b/packages/builder-rsbuild/src/types.ts
@@ -1,8 +1,10 @@
+import type { Stats } from './index'
 import type { RsbuildConfig } from '@rsbuild/core'
 import type {
   Builder,
   Options,
   TypescriptOptions as TypeScriptOptionsBase,
+  BuilderResult as BuilderResultBase,
 } from '@storybook/types'
 import type { PluginTypeCheckerOptions } from '@rsbuild/plugin-type-check'
 
@@ -38,3 +40,7 @@ export type BuilderOptions = {
    */
   rsbuildConfigPath?: string
 }
+
+export interface BuilderResult extends BuilderResultBase {
+  stats?: Stats
+}