feat(rsc): support observing imports/exports during build scan

packages/plugin-rsc/src/plugin.ts

@@ -97,8 +97,25 @@ type ServerReferenceMeta = {
   referenceKey: string
   // TODO: tree shake unused server functions
   exportNames: string[]
+  inlineExportNames?: string[]
 }
 
+type ScanBuildObserver = (
+  event:
+    | {
+        type: 'reset'
+        environmentName: string
+      }
+    | {
+        type: 'module'
+        environmentName: string
+        code: string
+        imports: readonly esModuleLexer.ImportSpecifier[]
+        exports: readonly esModuleLexer.ExportSpecifier[]
+        info: Rollup.ModuleInfo
+      },
+) => void
+
 const PKG_NAME = '@vitejs/plugin-rsc'
 const REACT_SERVER_DOM_NAME = `${PKG_NAME}/vendor/react-server-dom`
 
@@ -128,6 +145,7 @@ class RscPluginManager {
   clientReferenceGroups: Record</* group name*/ string, ClientReferenceMeta[]> =
     {}
   serverReferenceMetaMap: Record<string, ServerReferenceMeta> = {}
+  scanBuildObservers: Set<ScanBuildObserver> = new Set()
   serverResourcesMetaMap: Record<string, { key: string }> = {}
   environmentImportMetaMap: Record<
     string, // sourceEnv
@@ -2039,6 +2057,7 @@ function vitePluginUseServer(
               referenceKey: getNormalizedId(),
               exportNames:
                 'names' in result ? result.names : result.exportNames,
+              inlineExportNames: 'names' in result ? result.names : [],
             }
             const importSource = resolvePackage(`${PKG_NAME}/react/rsc`)
             output.prepend(

packages/plugin-rsc/src/plugins/fixtures/scan-observer/actions.js

@@ -0,0 +1 @@
+export const action = 1

packages/plugin-rsc/src/plugins/fixtures/scan-observer/entry.js

@@ -0,0 +1,2 @@
+import { action } from './actions.js'
+export { action as submit } from './actions.js'

packages/plugin-rsc/src/plugins/fixtures/server-reference-meta/entry.js

@@ -0,0 +1,2 @@
+import './inline.js'
+import './module.js'

packages/plugin-rsc/src/plugins/fixtures/server-reference-meta/inline.js

@@ -0,0 +1,3 @@
+export async function inlineAction() {
+  'use server'
+}

packages/plugin-rsc/src/plugins/fixtures/server-reference-meta/module.js

@@ -0,0 +1,3 @@
+'use server'
+
+export async function moduleAction() {}

packages/plugin-rsc/src/plugins/scan.integration.test.ts

@@ -0,0 +1,123 @@
+import path from 'node:path'
+import type { ExportSpecifier, ImportSpecifier } from 'es-module-lexer'
+import { build } from 'vite'
+import { describe, expect, it, vi } from 'vitest'
+import type { RscPluginManager } from '../plugin'
+import { scanBuildStripPlugin } from './scan'
+
+const root = path.join(import.meta.dirname, 'fixtures/scan-observer')
+
+describe(scanBuildStripPlugin, () => {
+  it('does not notify observers outside scan builds', async () => {
+    const observer = vi.fn()
+    const manager = {
+      isScanBuild: false,
+      scanBuildObservers: new Set([observer]),
+    } as unknown as RscPluginManager
+
+    await build({
+      root,
+      logLevel: 'silent',
+      build: {
+        write: false,
+        rollupOptions: { input: path.join(root, 'entry.js') },
+      },
+      plugins: [scanBuildStripPlugin({ manager })],
+    })
+
+    expect(observer).not.toHaveBeenCalled()
+  })
+
+  it('emits ordered reset and module events with raw scan metadata', async () => {
+    const observer = vi.fn()
+    const secondObserver = vi.fn()
+    const manager = {
+      isScanBuild: true,
+      scanBuildObservers: new Set([observer, secondObserver]),
+    } as unknown as RscPluginManager
+
+    await build({
+      root,
+      logLevel: 'silent',
+      build: {
+        write: false,
+        rollupOptions: { input: path.join(root, 'entry.js') },
+      },
+      plugins: [scanBuildStripPlugin({ manager })],
+    })
+
+    expect(secondObserver.mock.calls).toEqual(observer.mock.calls)
+    expect(observer.mock.calls[0]![0]).toEqual({
+      type: 'reset',
+      environmentName: 'client',
+    })
+
+    const moduleEvents = observer.mock.calls
+      .map(([event]) => event)
+      .filter((event) => event.type === 'module')
+    expect(moduleEvents).toHaveLength(2)
+
+    const entryEvent = moduleEvents.find((event) =>
+      event.info.id.endsWith('/entry.js'),
+    )
+    expect(entryEvent).toMatchObject({
+      type: 'module',
+      environmentName: 'client',
+      code: expect.stringContaining("import { action } from './actions.js'"),
+      info: {
+        id: expect.stringMatching(/\/entry\.js$/),
+        importedIds: [expect.stringMatching(/\/actions\.js$/)],
+      },
+    })
+    expect(entryEvent.imports.map((item: ImportSpecifier) => item.n)).toEqual([
+      './actions.js',
+      './actions.js',
+    ])
+    expect(entryEvent.exports.map((item: ExportSpecifier) => item.n)).toEqual([
+      'submit',
+    ])
+
+    const actionsEvent = moduleEvents.find((event) =>
+      event.info.id.endsWith('/actions.js'),
+    )
+    expect(actionsEvent).toMatchObject({
+      type: 'module',
+      environmentName: 'client',
+      code: 'export const action = 1\n',
+      imports: [],
+      exports: [expect.objectContaining({ n: 'action' })],
+      info: {
+        id: expect.stringMatching(/\/actions\.js$/),
+        importedIds: [],
+      },
+    })
+  })
+
+  it('deduplicates module events across shared scan plugins', async () => {
+    const observer = vi.fn()
+    const manager = {
+      isScanBuild: true,
+      scanBuildObservers: new Set([observer]),
+    } as unknown as RscPluginManager
+
+    await build({
+      root,
+      logLevel: 'silent',
+      build: {
+        write: false,
+        rollupOptions: { input: path.join(root, 'entry.js') },
+      },
+      plugins: [
+        scanBuildStripPlugin({ manager }),
+        scanBuildStripPlugin({ manager }),
+      ],
+    })
+
+    const moduleIds = observer.mock.calls
+      .map(([event]) => event)
+      .filter((event) => event.type === 'module')
+      .map((event) => path.basename(event.info.id))
+      .sort()
+    expect(moduleIds).toEqual(['actions.js', 'entry.js'])
+  })
+})

packages/plugin-rsc/src/plugins/scan.test.ts

@@ -1,13 +1,14 @@
 import * as esModuleLexer from 'es-module-lexer'
-import { beforeAll, describe, expect, it } from 'vitest'
+import { beforeAll, describe, expect, it, vi } from 'vitest'
+import vitePluginRsc from '../plugin'
 import { transformScanBuildStrip } from './scan'
 
 describe(transformScanBuildStrip, () => {
   beforeAll(async () => {
     await esModuleLexer.init
   })
 
-  it('basic', async () => {
+  it('strips modules to their imports', async () => {
     const input = `\
 import { a } from "a";
 import "b";
@@ -27,4 +28,26 @@ export default "foo";
       "
     `)
   })
+
+  it('reports the existing lexer result when observed', async () => {
+    const input = `import { action } from './actions'; export { action as submit }`
+    const onLexed = vi.fn()
+
+    await transformScanBuildStrip(input, onLexed)
+
+    expect(onLexed).toHaveBeenCalledOnce()
+    const [imports, exports] = onLexed.mock.calls[0]!
+    expect(
+      imports.map((item: esModuleLexer.ImportSpecifier) => item.n),
+    ).toEqual(['./actions'])
+    expect(
+      exports.map((item: esModuleLexer.ExportSpecifier) => item.n),
+    ).toEqual(['submit'])
+  })
+
+  it('keeps the existing scan plugin composition', () => {
+    expect(
+      vitePluginRsc().filter((plugin) => plugin.name === 'rsc:scan-strip'),
+    ).toHaveLength(2)
+  })
 })

packages/plugin-rsc/src/plugins/scan.ts

@@ -4,35 +4,98 @@ import { walk } from 'estree-walker'
 import { parseAstAsync, type Plugin } from 'vite'
 import type { RscPluginManager } from '../plugin'
 
+type ScanBuildModule = {
+  code: string
+  imports: readonly esModuleLexer.ImportSpecifier[]
+  exports: readonly esModuleLexer.ExportSpecifier[]
+}
+
+const scanBuildModulesMap = new WeakMap<
+  RscPluginManager,
+  Map<string, ScanBuildModule>
+>()
+
 // During scan build, we strip all code but imports to
 // traverse module graph faster and just discover client/server references.
 export function scanBuildStripPlugin({
   manager,
 }: {
   manager: RscPluginManager
 }): Plugin {
+  let scanBuildModules = scanBuildModulesMap.get(manager)
+  if (!scanBuildModules) {
+    scanBuildModules = new Map()
+    scanBuildModulesMap.set(manager, scanBuildModules)
+  }
   return {
     name: 'rsc:scan-strip',
     apply: 'build',
     enforce: 'post',
+    buildStart() {
+      if (manager.isScanBuild && manager.scanBuildObservers.size > 0) {
+        scanBuildModules.clear()
+        for (const observer of manager.scanBuildObservers) {
+          observer({
+            type: 'reset',
+            environmentName: this.environment.name,
+          })
+        }
+      }
+    },
     transform: {
       filter: {
         id: { exclude: exactRegex('\0rolldown/runtime.js') },
       },
-      async handler(code, _id, _options) {
+      async handler(code, id, _options) {
         if (!manager.isScanBuild) return
-        const output = await transformScanBuildStrip(code)
-        return { code: output, map: { mappings: '' } }
+        return {
+          code: await transformScanBuildStrip(
+            code,
+            manager.scanBuildObservers.size > 0
+              ? (imports, exports) => {
+                  if (!scanBuildModules.has(id)) {
+                    scanBuildModules.set(id, {
+                      code,
+                      imports,
+                      exports,
+                    })
+                  }
+                }
+              : undefined,
+          ),
+          map: { mappings: '' },
+        }
       },
     },
+    moduleParsed(info) {
+      if (!manager.isScanBuild || manager.scanBuildObservers.size === 0) return
+      const lexed = scanBuildModules.get(info.id)
+      if (!lexed) return
+      scanBuildModules.delete(info.id)
+      for (const observer of manager.scanBuildObservers) {
+        observer({
+          type: 'module',
+          environmentName: this.environment.name,
+          info,
+          ...lexed,
+        })
+      }
+    },
   }
 }
 
 // https://github.com/vitejs/vite/blob/86d2e8be50be535494734f9f5f5236c61626b308/packages/vite/src/node/plugins/importMetaGlob.ts#L113
 const importGlobRE = /\bimport\.meta\.glob(?:<\w+>)?\s*\(/g
 
-export async function transformScanBuildStrip(code: string): Promise<string> {
-  const [imports] = esModuleLexer.parse(code)
+export async function transformScanBuildStrip(
+  code: string,
+  onLexed?: (
+    imports: readonly esModuleLexer.ImportSpecifier[],
+    exports: readonly esModuleLexer.ExportSpecifier[],
+  ) => void | Promise<void>,
+): Promise<string> {
+  const [imports, exports] = esModuleLexer.parse(code)
+  await onLexed?.(imports, exports)
   let output = imports
     .map((e) => e.n && `import ${JSON.stringify(e.n)};\n`)
     .filter(Boolean)

packages/plugin-rsc/src/plugins/server-reference-meta.integration.test.ts

@@ -0,0 +1,44 @@
+import path from 'node:path'
+import { build } from 'vite'
+import { describe, expect, it } from 'vitest'
+import { type PluginApi, vitePluginRscMinimal } from '../plugin'
+
+const root = path.join(import.meta.dirname, 'fixtures/server-reference-meta')
+
+describe('server reference metadata', () => {
+  it('distinguishes inline actions from module-level actions', async () => {
+    const plugins = vitePluginRscMinimal({
+      enableActionEncryption: false,
+      environment: { rsc: 'client' },
+    })
+    const manager = (
+      plugins.find((plugin) => plugin.name === 'rsc:minimal')!.api as PluginApi
+    ).manager
+
+    await build({
+      root,
+      logLevel: 'silent',
+      build: {
+        write: false,
+        rollupOptions: { input: path.join(root, 'entry.js') },
+      },
+      plugins,
+    })
+
+    const inlineMeta = Object.values(manager.serverReferenceMetaMap).find(
+      (meta) => meta.importId.endsWith('/inline.js'),
+    )
+    expect(inlineMeta?.inlineExportNames).toEqual(inlineMeta?.exportNames)
+    expect(inlineMeta?.inlineExportNames).toEqual([
+      expect.stringMatching(/inlineAction$/),
+    ])
+
+    const moduleMeta = Object.values(manager.serverReferenceMetaMap).find(
+      (meta) => meta.importId.endsWith('/module.js'),
+    )
+    expect(moduleMeta).toMatchObject({
+      exportNames: ['moduleAction'],
+      inlineExportNames: [],
+    })
+  })
+})