feat(rsc): ability to not load server reference module during createFromReadableStream on rsc environment

packages/plugin-rsc/examples/rsc-server-reference-replay/.gitignore

@@ -0,0 +1,2 @@
+.flight-cache
+dist

packages/plugin-rsc/examples/rsc-server-reference-replay/README.md

@@ -0,0 +1,46 @@
+# Persisted Flight server reference replay
+
+This example persists a Flight payload containing a server reference, restarts the server, and replays the payload without importing the server action in the RSC environment. The action is imported only when the replayed form invokes it.
+
+The example intentionally uses a native form without JavaScript so the final step exercises `decodeAction`.
+
+## Development manual test
+
+Start the first process:
+
+```bash
+pnpm dev
+```
+
+1. Visit `http://localhost:5173/cache` and confirm the page displays `true`.
+2. Stop the development server without changing the source graph.
+3. Run `pnpm dev` again.
+4. Visit `http://localhost:5173/` and confirm the page displays `false`.
+5. Select **Invoke replayed action** and confirm the response displays `true`.
+
+The source graph must remain unchanged across the restart so its development server-reference IDs remain stable.
+
+## Production manual test
+
+Build the example once:
+
+```bash
+pnpm build
+```
+
+Start the first process:
+
+```bash
+pnpm preview
+```
+
+1. Visit `http://localhost:4173/cache`.
+2. Confirm the page displays `Action imported in the RSC environment: true`.
+3. Stop the preview server without rebuilding.
+4. Run `pnpm preview` again.
+5. Visit `http://localhost:4173/` and confirm the page displays `false`.
+6. Select **Invoke replayed action** and confirm the response displays `true`.
+
+The same production build must be used for both processes because the persisted Flight payload contains build-specific server-reference IDs.
+
+Delete `.flight-cache` to reset the example.

packages/plugin-rsc/examples/rsc-server-reference-replay/package.json

@@ -0,0 +1,22 @@
+{
+  "name": "@vitejs/plugin-rsc-examples-rsc-server-reference-replay",
+  "version": "0.0.0",
+  "private": true,
+  "license": "MIT",
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "vite build",
+    "preview": "vite preview"
+  },
+  "dependencies": {
+    "react": "^19.2.7",
+    "react-dom": "^19.2.7"
+  },
+  "devDependencies": {
+    "@types/react": "^19.2.17",
+    "@types/react-dom": "^19.2.3",
+    "@vitejs/plugin-rsc": "latest",
+    "vite": "^8.1.4"
+  }
+}

packages/plugin-rsc/examples/rsc-server-reference-replay/src/action-import-state.ts

@@ -0,0 +1,9 @@
+let actionImported = false
+
+export function markActionImported() {
+  actionImported = true
+}
+
+export function wasActionImported() {
+  return actionImported
+}

packages/plugin-rsc/examples/rsc-server-reference-replay/src/action.ts

@@ -0,0 +1,7 @@
+'use server'
+
+import { markActionImported } from './action-import-state'
+
+markActionImported()
+
+export async function replayedAction() {}

packages/plugin-rsc/examples/rsc-server-reference-replay/src/cached-content.tsx

@@ -0,0 +1,9 @@
+import { replayedAction } from './action'
+
+export function CachedContent() {
+  return (
+    <form action={replayedAction}>
+      <button type="submit">Invoke replayed action</button>
+    </form>
+  )
+}

packages/plugin-rsc/examples/rsc-server-reference-replay/src/framework/entry.browser.ts

@@ -0,0 +1 @@
+// This example only tests progressive enhancement without JavaScript.

packages/plugin-rsc/examples/rsc-server-reference-replay/src/framework/entry.rsc.tsx

@@ -0,0 +1,71 @@
+import { readFile, writeFile } from 'node:fs/promises'
+import { resolve } from 'node:path'
+import {
+  createFromReadableStream,
+  decodeAction,
+  renderToReadableStream,
+} from '@vitejs/plugin-rsc/rsc'
+import { wasActionImported } from '../action-import-state'
+import { Root } from '../root'
+
+const cacheFile = resolve('.flight-cache')
+
+async function handler(request: Request): Promise<Response> {
+  const url = new URL(request.url)
+
+  if (request.method === 'POST') {
+    const action = await decodeAction(await request.formData())
+    await action()
+  }
+
+  let bytes: Uint8Array
+  if (url.pathname === '/cache') {
+    const { CachedContent } = await import('../cached-content')
+    const stream = renderToReadableStream(<CachedContent />)
+    bytes = new Uint8Array(await new Response(stream).arrayBuffer())
+    await writeFile(cacheFile, bytes)
+  } else {
+    try {
+      bytes = await readFile(cacheFile)
+    } catch (error) {
+      if ((error as NodeJS.ErrnoException).code === 'ENOENT') {
+        return new Response('Visit /cache before replaying the saved Flight.', {
+          status: 404,
+        })
+      }
+      throw error
+    }
+  }
+
+  const cachedContent = await createFromReadableStream<React.ReactNode>(
+    toReadableStream(bytes),
+    {},
+    { preserveServerReferences: true },
+  )
+  const rscStream = renderToReadableStream({
+    root: (
+      <Root
+        actionImported={wasActionImported()}
+        cachedContent={cachedContent}
+      />
+    ),
+  })
+  const ssr = await import.meta.viteRsc.loadModule<
+    typeof import('./entry.ssr')
+  >('ssr', 'index')
+  const htmlStream = await ssr.renderHtml(rscStream)
+  return new Response(htmlStream, {
+    headers: { 'content-type': 'text/html;charset=utf-8' },
+  })
+}
+
+function toReadableStream(bytes: Uint8Array): ReadableStream<Uint8Array> {
+  return new ReadableStream({
+    start(controller) {
+      controller.enqueue(bytes)
+      controller.close()
+    },
+  })
+}
+
+export default { fetch: handler }

packages/plugin-rsc/examples/rsc-server-reference-replay/src/framework/entry.ssr.tsx

@@ -0,0 +1,17 @@
+import { createFromReadableStream } from '@vitejs/plugin-rsc/ssr'
+import React from 'react'
+import { renderToReadableStream } from 'react-dom/server.edge'
+
+type RscPayload = { root: React.ReactNode }
+
+export async function renderHtml(
+  rscStream: ReadableStream<Uint8Array>,
+): Promise<ReadableStream<Uint8Array>> {
+  let payload: Promise<RscPayload> | undefined
+  function SsrRoot() {
+    payload ??= createFromReadableStream<RscPayload>(rscStream)
+    return React.use(payload).root
+  }
+
+  return renderToReadableStream(<SsrRoot />)
+}

packages/plugin-rsc/examples/rsc-server-reference-replay/src/root.tsx

@@ -0,0 +1,19 @@
+export function Root(props: {
+  actionImported: boolean
+  cachedContent: React.ReactNode
+}) {
+  return (
+    <html>
+      <body>
+        <h1>Persisted Flight server reference</h1>
+        <p>
+          Action imported in the RSC environment:{' '}
+          <output data-testid="action-imported">
+            {String(props.actionImported)}
+          </output>
+        </p>
+        {props.cachedContent}
+      </body>
+    </html>
+  )
+}

packages/plugin-rsc/examples/rsc-server-reference-replay/tsconfig.json

@@ -0,0 +1,17 @@
+{
+  "compilerOptions": {
+    "erasableSyntaxOnly": true,
+    "allowImportingTsExtensions": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "skipLibCheck": true,
+    "verbatimModuleSyntax": true,
+    "noEmit": true,
+    "moduleResolution": "Bundler",
+    "module": "ESNext",
+    "target": "ESNext",
+    "lib": ["ESNext", "DOM"],
+    "types": ["vite/client", "@vitejs/plugin-rsc/types"],
+    "jsx": "react-jsx"
+  }
+}

packages/plugin-rsc/examples/rsc-server-reference-replay/vite.config.ts

@@ -0,0 +1,29 @@
+import rsc from '@vitejs/plugin-rsc'
+import { defineConfig } from 'vite'
+
+export default defineConfig({
+  plugins: [rsc()],
+  environments: {
+    rsc: {
+      build: {
+        rollupOptions: {
+          input: { index: './src/framework/entry.rsc.tsx' },
+        },
+      },
+    },
+    ssr: {
+      build: {
+        rollupOptions: {
+          input: { index: './src/framework/entry.ssr.tsx' },
+        },
+      },
+    },
+    client: {
+      build: {
+        rollupOptions: {
+          input: { index: './src/framework/entry.browser.ts' },
+        },
+      },
+    },
+  },
+})

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_REFERENCE_PROXY_PREFIX,
   SERVER_REFERENCE_PREFIX,
   createReferenceCacheTag,
   removeReferenceCacheTag,
@@ -27,6 +28,31 @@ 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_REFERENCE_PROXY_PREFIX)) {
+        id = id.slice(SERVER_REFERENCE_PROXY_PREFIX.length)
+        id = removeReferenceCacheTag(id)
+        const target = {} as any
+        const getOrCreateServerReference = (name: string) => {
+          return (target[name] ??= ReactServer.registerServerReference(
+            () => {
+              throw new Error(
+                `Unexpectedly decoded server reference '${id}#${name}' is called on server`,
+              )
+            },
+            id,
+            name,
+          ))
+        }
+        return new Proxy(target, {
+          getOwnPropertyDescriptor(_target, name) {
+            if (typeof name !== 'string' || name === 'then') {
+              return Reflect.getOwnPropertyDescriptor(target, name)
+            }
+            getOrCreateServerReference(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 +97,12 @@ export async function loadServerAction(id: string): Promise<Function> {
   return mod[name]
 }
 
-export function createServerManifest(): BundlerConfig {
+export function createServerManifest(options?: {
+  preserveServerReferences?: boolean
+}): BundlerConfig {
+  const prefix = options?.preserveServerReferences
+    ? SERVER_REFERENCE_PROXY_PREFIX
+    : ''
   const cacheTag = import.meta.env.DEV ? createReferenceCacheTag() : ''
 
   return new Proxy(
@@ -83,7 +114,7 @@ export function createServerManifest(): BundlerConfig {
         tinyassert(id)
         tinyassert(name)
         return {
-          id: SERVER_REFERENCE_PREFIX + id + cacheTag,
+          id: SERVER_REFERENCE_PREFIX + 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_REFERENCE_PROXY_PREFIX = '$$proxy:'
+
 // 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

@@ -346,7 +346,7 @@ export function vitePluginRscMinimal(
       apply: 'serve',
       load: {
         filter: { id: prefixRegex('\0virtual:vite-rsc/reference-validation?') },
-        handler(id, _options) {
+        async handler(id, _options) {
           if (id.startsWith('\0virtual:vite-rsc/reference-validation?')) {
             const parsed = parseReferenceValidationVirtual(id)
             assert(parsed)
@@ -359,9 +359,19 @@ export function vitePluginRscMinimal(
               }
             }
             if (parsed.type === 'server') {
-              const meta = Object.values(manager.serverReferenceMetaMap).find(
+              let meta = Object.values(manager.serverReferenceMetaMap).find(
                 (meta) => meta.referenceKey === parsed.id,
               )
+              if (!meta) {
+                try {
+                  await (this.environment as DevEnvironment).transformRequest(
+                    parsed.id,
+                  )
+                } catch {}
+                meta = Object.values(manager.serverReferenceMetaMap).find(
+                  (meta) => meta.referenceKey === parsed.id,
+                )
+              }
               if (meta) {
                 return `export {}`
               }

packages/plugin-rsc/src/react/rsc/client.ts

@@ -13,12 +13,20 @@ import type {
 export function createFromReadableStream<T>(
   stream: ReadableStream<Uint8Array>,
   options: CreateFromReadableStreamEdgeOptions = {},
+  extraOptions?: {
+    /**
+     * @experimental
+     */
+    preserveServerReferences?: boolean
+  },
 ): Promise<T> {
   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: extraOptions?.preserveServerReferences,
+      }),
       moduleMap: createServerDecodeClientManifest(),
     },
     ...options,