feat(rsc): preserve decoded server references

Author: james-elicx

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/react/rsc.ts

@@ -40,16 +40,22 @@ export function renderToReadableStream<T>(
 
 export function createFromReadableStream<T>(
   stream: ReadableStream<Uint8Array>,
-  options: CreateFromReadableStreamEdgeOptions = {},
+  options: CreateFromReadableStreamEdgeOptions & {
+    /** Preserve decoded server references so the value can be rendered again. */
+    serverReferences?: 'resolve' | 'preserve'
+  } = {},
 ): Promise<T> {
+  const { serverReferences = 'resolve', ...reactOptions } = options
   return ReactClient.createFromReadableStream(stream, {
     serverConsumerManifest: {
       // https://github.com/facebook/react/pull/31300
       // https://github.com/vercel/next.js/pull/71527
-      serverModuleMap: createServerManifest(),
+      serverModuleMap: createServerManifest({
+        preserveServerReferences: serverReferences === 'preserve',
+      }),
       moduleMap: createServerDecodeClientManifest(),
     },
-    ...options,
+    ...reactOptions,
   })
 }