feat(rsc): configure server reference markers

Author: james-elicx

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

@@ -53,4 +53,48 @@ describe('server reference manifest', () => {
     expect(manifest.code.match(/\bcached\b/g)).toHaveLength(2)
     expect(manifest.code.match(/\baction\b/g)).toHaveLength(2)
   })
+
+  it('preserves metadata for configured server reference markers', async () => {
+    const marker = '/* framework-server-reference */'
+    const plugins = vitePluginRscMinimal({
+      enableActionEncryption: false,
+      serverReferenceMarkers: [marker, ''],
+    })
+    const minimalPlugin = plugins.find(
+      (plugin) => plugin.name === 'rsc:minimal',
+    )!
+    const manager = (minimalPlugin.api as any).manager
+    const useServerPlugin = plugins.find(
+      (plugin) => plugin.name === 'rsc:use-server',
+    )!
+    const transformHandler = (useServerPlugin.transform as any).handler
+    const transformContext = {
+      environment: { name: 'rsc', mode: 'build' },
+      error(error: unknown) {
+        throw error
+      },
+    }
+
+    const markedId = '/root/marked.ts'
+    manager.serverReferenceMetaMap[markedId] = {
+      importId: markedId,
+      referenceKey: 'marked',
+      exportNames: ['cached'],
+    }
+    await transformHandler.call(
+      transformContext,
+      `${marker}\nexport {}`,
+      markedId,
+    )
+    expect(manager.serverReferenceMetaMap[markedId]).toBeDefined()
+
+    const unmarkedId = '/root/unmarked.ts'
+    manager.serverReferenceMetaMap[unmarkedId] = {
+      importId: unmarkedId,
+      referenceKey: 'unmarked',
+      exportNames: ['stale'],
+    }
+    await transformHandler.call(transformContext, 'export {}', unmarkedId)
+    expect(manager.serverReferenceMetaMap[unmarkedId]).toBeUndefined()
+  })
 })

packages/plugin-rsc/src/plugin.ts

@@ -171,6 +171,13 @@ class RscPluginManager {
 }
 
 export type RscPluginOptions = {
+  /**
+   * Source markers emitted by preceding transforms that register server
+   * references in the plugin manager. Modules containing a marker keep their
+   * existing server-reference metadata when they do not contain `"use server"`.
+   * @experimental
+   */
+  serverReferenceMarkers?: string[]
   /**
    * shorthand for configuring `environments.(name).build.rollupOptions.input.index`
    */
@@ -1926,7 +1933,7 @@ function vitePluginDefineEncryptionKey(
 function vitePluginUseServer(
   useServerPluginOptions: Pick<
     RscPluginOptions,
-    'enableActionEncryption' | 'environment'
+    'enableActionEncryption' | 'environment' | 'serverReferenceMarkers'
   >,
   manager: RscPluginManager,
 ): Plugin[] {
@@ -1946,7 +1953,13 @@ function vitePluginUseServer(
         // filter: { code: 'use server' },
         async handler(code, id) {
           if (!code.includes('use server')) {
-            delete manager.serverReferenceMetaMap[id]
+            const hasServerReferenceMarker =
+              useServerPluginOptions.serverReferenceMarkers?.some(
+                (marker) => marker.length > 0 && code.includes(marker),
+              ) ?? false
+            if (!hasServerReferenceMarker) {
+              delete manager.serverReferenceMetaMap[id]
+            }
             return
           }
           let ast = await parseAstAsync(code)