feat(rsc): ability to preserve decoded server references

packages/plugin-rsc/e2e/basic.test.ts

@@ -927,7 +927,7 @@
           `/* color: rgb(0, 165, 255); */`,
         ),
       )
       await expect(page.locator('.test-style-server')).toHaveCSS(
         'color',
         'rgb(0, 0, 0)',
       )
@@ -1781,6 +1781,26 @@
     )
   })
 
+  test('use cache replays Flight with framework server action', async ({
+    page,
+  }) => {
+    await page.goto(f.url())
+    await waitForHydration(page)
+    const locator = page.getByTestId(
+      'test-use-cache-flight-replay-server-action',
+    )
+    await expect(
+      locator.getByTestId('test-use-cache-flight-replay-server-action-cache'),
+    ).toHaveText('cached product card render count: 1')
+    await locator.getByRole('button', { name: 'add cached product' }).click()
+    await expect(
+      locator.getByTestId('test-use-cache-flight-replay-server-action-result'),
+    ).toHaveText(/^added rsc-product-1 with framework action \([1-9]\d*\)$/)
+    await expect(
+      locator.getByTestId('test-use-cache-flight-replay-server-action-cache'),
+    ).toHaveText('cached product card render count: 1')
+  })
+
   test('hydration mismatch', async ({ page }) => {
     const errors: Error[] = []
     page.on('pageerror', (error) => {

packages/plugin-rsc/examples/basic/src/framework/entry.rsc.tsx

@@ -9,6 +9,7 @@ import {
 import type React from 'react'
 import type { ReactFormState } from 'react-dom/client'
 import { parseRenderRequest } from './request.tsx'
+import { loadFrameworkServerReference } from './server-reference-runtime.ts'
 
 // The schema of payload which is serialized into RSC stream on rsc environment
 // and deserialized on ssr/client environments.
@@ -50,7 +51,9 @@ async function handleRequest({
         : await request.text()
       temporaryReferences = createTemporaryReferenceSet()
       const args = await decodeReply(body, { temporaryReferences })
-      const action = await loadServerAction(renderRequest.actionId)
+      const action =
+        loadFrameworkServerReference(renderRequest.actionId) ??
+        (await loadServerAction(renderRequest.actionId))
       try {
         const data = await action.apply(null, args)
         returnValue = { ok: true, data }

packages/plugin-rsc/examples/basic/src/framework/server-reference-runtime.ts

@@ -0,0 +1,14 @@
+import { registerServerReference } from '@vitejs/plugin-rsc/rsc'
+
+const frameworkServerReferences = new Map<string, Function>()
+
+export function registerFrameworkServerReference<
+  T extends (...args: any[]) => unknown,
+>(reference: T, id: string, name: string): T {
+  frameworkServerReferences.set(`${id}#${name}`, reference)
+  return registerServerReference(reference, id, name) as T
+}
+
+export function loadFrameworkServerReference(id: string): Function | undefined {
+  return frameworkServerReferences.get(id)
+}

packages/plugin-rsc/examples/basic/src/framework/use-cache-runtime.tsx

@@ -65,6 +65,9 @@ export default function cacheWrapper(fn: (...args: any[]) => Promise<unknown>) {
     const result = createFromReadableStream(stream, {
       environmentName: 'Cache',
       replayConsoleLogs: true,
+      // Cached RSC streams can contain framework-owned server references whose
+      // implementation modules are not resolvable by the app bundler runtime.
+      serverReferences: 'preserve',
       temporaryReferences: clientTemporaryReferences,
     })
     return result

packages/plugin-rsc/examples/basic/src/routes/use-cache/client.tsx

@@ -0,0 +1,29 @@
+'use client'
+
+import React from 'react'
+
+export function TestUseCacheFlightReplayServerActionClient(props: {
+  addToCart: (productId: string) => Promise<string>
+  productId: string
+  renderCount: number
+}) {
+  const [result, setResult] = React.useState('idle')
+
+  return (
+    <div data-testid="test-use-cache-flight-replay-server-action">
+      <span data-testid="test-use-cache-flight-replay-server-action-cache">
+        cached product card render count: {props.renderCount}
+      </span>
+      <button
+        onClick={async () => {
+          setResult(await props.addToCart(props.productId))
+        }}
+      >
+        add cached product
+      </button>
+      <span data-testid="test-use-cache-flight-replay-server-action-result">
+        {result}
+      </span>
+    </div>
+  )
+}

packages/plugin-rsc/examples/basic/src/routes/use-cache/server.tsx

@@ -1,11 +1,14 @@
+import { registerFrameworkServerReference } from '../../framework/server-reference-runtime'
 import { revalidateCache } from '../../framework/use-cache-runtime'
+import { TestUseCacheFlightReplayServerActionClient } from './client'
 
 export function TestUseCache() {
   return (
     <>
       <TestUseCacheFn />
       <TestUseCacheComponent />
       <TestUseCacheClosure />
+      <TestUseCacheFlightReplayServerAction />
     </>
   )
 }
@@ -103,3 +106,32 @@ let outerFnArg = ''
 let innerFnArg = ''
 let innerFnCount = 0
 let actionCount2 = 0
+
+async function TestUseCacheFlightReplayServerAction() {
+  return <CachedProductCard productId="rsc-product-1" />
+}
+
+async function CachedProductCard(props: { productId: string }) {
+  'use cache'
+  cachedProductCardRenderCount++
+  return (
+    <TestUseCacheFlightReplayServerActionClient
+      key={props.productId}
+      addToCart={addCachedProductToCart}
+      productId={props.productId}
+      renderCount={cachedProductCardRenderCount}
+    />
+  )
+}
+
+let cachedProductCardRenderCount = 0
+let cartActionCount = 0
+
+const addCachedProductToCart = registerFrameworkServerReference(
+  async (productId: string) => {
+    cartActionCount++
+    return `added ${productId} with framework action (${cartActionCount})`
+  },
+  'framework:cached-flight-product-card',
+  'addToCart',
+)

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.test.ts

@@ -0,0 +1,100 @@
+import { describe, expect, it } from 'vitest'
+import { vitePluginRscMinimal } from './plugin'
+
+describe('server reference manifest', () => {
+  it('preserves and deduplicates existing export names', async () => {
+    const plugins = vitePluginRscMinimal({ enableActionEncryption: false })
+    const minimalPlugin = plugins.find(
+      (plugin) => plugin.name === 'rsc:minimal',
+    )!
+    const manager = (minimalPlugin.api as any).manager
+    manager.config = {
+      command: 'build',
+      root: '/root',
+    }
+
+    const id = '/root/actions.ts'
+    manager.serverReferenceMetaMap[id] = {
+      importId: id,
+      referenceKey: 'existing',
+      exportNames: ['cached', 'action', 'cached'],
+    }
+
+    const useServerPlugin = plugins.find(
+      (plugin) => plugin.name === 'rsc:use-server',
+    )!
+    const transformHandler = (useServerPlugin.transform as any).handler
+    await transformHandler.call(
+      {
+        environment: { name: 'rsc', mode: 'build' },
+        error(error: unknown) {
+          throw error
+        },
+      },
+      `"use server"; export async function action() {}`,
+      id,
+    )
+
+    expect(manager.serverReferenceMetaMap[id].exportNames).toEqual([
+      'cached',
+      'action',
+    ])
+
+    const manifestPlugin = plugins.find(
+      (plugin) => plugin.name === 'rsc:virtual-vite-rsc/server-references',
+    )!
+    const loadHandler = (manifestPlugin.load as any).handler
+    const manifest = await loadHandler.call(
+      { environment: { mode: 'build' } },
+      '\0virtual:vite-rsc/server-references',
+      {},
+    )
+
+    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()
+  })
+})