feat(rsc): improve directive transform primitives

packages/plugin-rsc/src/plugin.ts

@@ -39,6 +39,12 @@ import {
   withResolvedIdProxy,
 } from './plugins/resolved-id-proxy'
 import { scanBuildStripPlugin } from './plugins/scan'
+import {
+  SERVER_FUNCTION_DIRECTIVE_MARKER,
+  type ServerFunctionDirective,
+  type ServerFunctionDirectiveContext,
+  vitePluginServerFunctionDirectives,
+} from './plugins/server-function-directives'
 import {
   parseCssVirtual,
   toCssVirtual,
@@ -171,6 +177,8 @@ class RscPluginManager {
 }
 
 export type RscPluginOptions = {
+  /** @experimental */
+  serverFunctionDirectives?: ServerFunctionDirective[]
   /**
    * shorthand for configuring `environments.(name).build.rollupOptions.input.index`
    */
@@ -283,6 +291,8 @@ export type RscPluginOptions = {
   customClientEntry?: boolean
 }
 
+export type { ServerFunctionDirective, ServerFunctionDirectiveContext }
+
 export type PluginApi = {
   manager: RscPluginManager
 }
@@ -339,6 +349,30 @@ export function vitePluginRscMinimal(
     },
     ...vitePluginRscCore(),
     ...vitePluginUseClient(rscPluginOptions, manager),
+    ...(rscPluginOptions.serverFunctionDirectives?.length
+      ? [
+          vitePluginServerFunctionDirectives({
+            definitions: rscPluginOptions.serverFunctionDirectives,
+            manager,
+            serverEnvironmentName: rscPluginOptions.environment?.rsc ?? 'rsc',
+            browserEnvironmentName:
+              rscPluginOptions.environment?.browser ?? 'client',
+            rscRuntime: resolvePackage(`${PKG_NAME}/react/rsc`),
+            ssrRuntime: resolvePackage(`${PKG_NAME}/react/ssr`),
+            browserRuntime: resolvePackage(`${PKG_NAME}/react/browser`),
+            encryptionRuntime: resolvePackage(
+              `${PKG_NAME}/utils/encryption-runtime`,
+            ),
+            expandExportAll: (context, code, ast, id) =>
+              transformExpandExportAll({
+                code,
+                ast,
+                importer: id,
+                ...createTransformExpandExportAllContext(context),
+              }),
+          }),
+        ]
+      : []),
     ...vitePluginUseServer(rscPluginOptions, manager),
     ...vitePluginDefineEncryptionKey(rscPluginOptions),
     {
@@ -1946,7 +1980,9 @@ function vitePluginUseServer(
         // filter: { code: 'use server' },
         async handler(code, id) {
           if (!code.includes('use server')) {
-            delete manager.serverReferenceMetaMap[id]
+            if (!code.includes(SERVER_FUNCTION_DIRECTIVE_MARKER)) {
+              delete manager.serverReferenceMetaMap[id]
+            }
             return
           }
           let ast = await parseAstAsync(code)
@@ -2017,11 +2053,15 @@ function vitePluginUseServer(
               delete manager.serverReferenceMetaMap[id]
               return
             }
+            const customExportNames =
+              manager.serverReferenceMetaMap[id]?.exportNames ?? []
             manager.serverReferenceMetaMap[id] = {
               importId: id,
               referenceKey: getNormalizedId(),
-              exportNames:
-                'names' in result ? result.names : result.exportNames,
+              exportNames: [
+                ...customExportNames,
+                ...('names' in result ? result.names : result.exportNames),
+              ],
             }
             const importSource = resolvePackage(`${PKG_NAME}/react/rsc`)
             output.prepend(