feat(rsc): support custom server function directives

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

@@ -0,0 +1,37 @@
+import { beforeAll, describe, expect, it, vi } from 'vitest'
+
+vi.mock('@vitejs/plugin-rsc/vendor/react-server-dom/server.edge', () => ({
+  registerClientReference: vi.fn(),
+  registerServerReference(reference: Function, id: string, name: string) {
+    return Object.defineProperties(reference, {
+      $$typeof: { value: Symbol.for('react.server.reference') },
+      $$id: { value: `${id}#${name}` },
+      $$bound: { value: null, writable: true },
+    })
+  },
+}))
+
+const { createServerManifest, setRequireModule } = await import('./rsc')
+
+beforeAll(() => {
+  setRequireModule({
+    load() {
+      throw new Error('preserved references must not load their implementation')
+    },
+  })
+})
+
+describe('createServerManifest', () => {
+  it('preserves server references without loading their implementation', async () => {
+    const manifest = createServerManifest({ preserveServerReferences: true })
+    const entry = manifest['module-id#action']!
+    expect(entry.id).toContain('$$decode-server-reference:module-id')
+
+    const module = await (globalThis as any).__vite_rsc_require__(entry.id)
+    expect(Object.prototype.hasOwnProperty.call(module, 'action')).toBe(true)
+    const reference = module.action
+    expect(reference.$$typeof).toBe(Symbol.for('react.server.reference'))
+    expect(reference.$$id).toBe('module-id#action')
+    expect(reference.$$bound).toBeNull()
+  })
+})

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

@@ -2,6 +2,7 @@ import { memoize, tinyassert } from '@hiogawa/utils'
 import type { BundlerConfig, ImportManifestEntry, ModuleMap } from '../types'
 import {
   SERVER_DECODE_CLIENT_PREFIX,
+  SERVER_DECODE_REFERENCE_PREFIX,
   SERVER_REFERENCE_PREFIX,
   createReferenceCacheTag,
   removeReferenceCacheTag,
@@ -27,6 +28,28 @@ export function setRequireModule(options: {
   // need memoize to return stable promise from __webpack_require__
   ;(globalThis as any).__vite_rsc_server_require__ = memoize(
     async (id: string) => {
+      if (id.startsWith(SERVER_DECODE_REFERENCE_PREFIX)) {
+        id = id.slice(SERVER_DECODE_REFERENCE_PREFIX.length)
+        id = removeReferenceCacheTag(id)
+        const target = {} as Record<string, unknown>
+        return new Proxy(target, {
+          getOwnPropertyDescriptor(_target, name) {
+            if (typeof name !== 'string' || name === 'then') {
+              return Reflect.getOwnPropertyDescriptor(target, name)
+            }
+            target[name] ??= ReactServer.registerServerReference(
+              () => {
+                throw new Error(
+                  `Unexpectedly preserved server reference '${id}#${name}' is called on server`,
+                )
+              },
+              id,
+              name,
+            )
+            return Reflect.getOwnPropertyDescriptor(target, name)
+          },
+        })
+      }
       if (id.startsWith(SERVER_DECODE_CLIENT_PREFIX)) {
         // decode client reference on the server
         id = id.slice(SERVER_DECODE_CLIENT_PREFIX.length)
@@ -71,7 +94,9 @@ export async function loadServerAction(id: string): Promise<Function> {
   return mod[name]
 }
 
-export function createServerManifest(): BundlerConfig {
+export function createServerManifest(options?: {
+  preserveServerReferences?: boolean
+}): BundlerConfig {
   const cacheTag = import.meta.env.DEV ? createReferenceCacheTag() : ''
 
   return new Proxy(
@@ -83,7 +108,13 @@ export function createServerManifest(): BundlerConfig {
         tinyassert(id)
         tinyassert(name)
         return {
-          id: SERVER_REFERENCE_PREFIX + id + cacheTag,
+          id:
+            SERVER_REFERENCE_PREFIX +
+            (options?.preserveServerReferences
+              ? SERVER_DECODE_REFERENCE_PREFIX
+              : '') +
+            id +
+            cacheTag,
           name,
           chunks: [],
           async: true,

packages/plugin-rsc/src/core/shared.ts

@@ -3,6 +3,8 @@ export const SERVER_REFERENCE_PREFIX = '$$server:'
 
 export const SERVER_DECODE_CLIENT_PREFIX = '$$decode-client:'
 
+export const SERVER_DECODE_REFERENCE_PREFIX = '$$decode-server-reference:'
+
 // cache bust memoized require promise during dev
 export function createReferenceCacheTag(): string {
   const cache = Math.random().toString(36).slice(2)

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,14 @@ function vitePluginUseServer(
               delete manager.serverReferenceMetaMap[id]
               return
             }
+            const customExportNames =
+              manager.serverReferenceMetaMap[id]?.exportNames ?? []
+            const exportNames =
+              'names' in result ? result.names : result.exportNames
             manager.serverReferenceMetaMap[id] = {
               importId: id,
               referenceKey: getNormalizedId(),
-              exportNames:
-                'names' in result ? result.names : result.exportNames,
+              exportNames: [...new Set([...customExportNames, ...exportNames])],
             }
             const importSource = resolvePackage(`${PKG_NAME}/react/rsc`)
             output.prepend(
@@ -2094,7 +2133,7 @@ function vitePluginUseServer(
       for (const meta of Object.values(manager.serverReferenceMetaMap)) {
         const key = JSON.stringify(meta.referenceKey)
         const id = JSON.stringify(meta.importId)
-        const exports = meta.exportNames
+        const exports = [...new Set(meta.exportNames)]
           .map((name) => (name === 'default' ? 'default: _default' : name))
           .sort()
         code += `