Support temporary references for server actions (#71230)

Implements https://x.com/sebmarkbage/status/1842240773867126939

supersedes and closes #66054 (with permission)

Author: unstubbable

packages/next/src/client/app-index.tsx

@@ -139,10 +139,10 @@ const readable = new ReadableStream({
   },
 })
 
-const initialServerResponse = createFromReadableStream(readable, {
-  callServer,
-  findSourceMapURL,
-})
+const initialServerResponse = createFromReadableStream<InitialRSCPayload>(
+  readable,
+  { callServer, findSourceMapURL }
+)
 
 // React overrides `.then` and doesn't return a new promise chain,
 // so we wrap the action queue in a promise to ensure that its value
@@ -151,7 +151,7 @@ const initialServerResponse = createFromReadableStream(readable, {
 const pendingActionQueue: Promise<AppRouterActionQueue> = new Promise(
   (resolve, reject) => {
     initialServerResponse.then(
-      (initialRSCPayload: InitialRSCPayload) => {
+      (initialRSCPayload) => {
         resolve(
           createMutableActionQueue(
             createInitialRouterState({
@@ -173,7 +173,7 @@ const pendingActionQueue: Promise<AppRouterActionQueue> = new Promise(
 )
 
 function ServerRoot(): React.ReactNode {
-  const initialRSCPayload = use<InitialRSCPayload>(initialServerResponse)
+  const initialRSCPayload = use(initialServerResponse)
   const actionQueue = use<AppRouterActionQueue>(pendingActionQueue)
 
   const router = (

packages/next/src/client/components/router-reducer/reducers/server-action-reducer.ts

@@ -16,7 +16,7 @@ import {
 // import { createFromFetch } from 'react-server-dom-webpack/client'
 // // eslint-disable-next-line import/no-extraneous-dependencies
 // import { encodeReply } from 'react-server-dom-webpack/client'
-const { createFromFetch, encodeReply } = (
+const { createFromFetch, createTemporaryReferenceSet, encodeReply } = (
   !!process.env.NEXT_RUNTIME
     ? // eslint-disable-next-line import/no-extraneous-dependencies
       require('react-server-dom-webpack/client.edge')
@@ -70,7 +70,8 @@ async function fetchServerAction(
   nextUrl: ReadonlyReducerState['nextUrl'],
   { actionId, actionArgs }: ServerActionAction
 ): Promise<FetchServerActionResult> {
-  const body = await encodeReply(actionArgs)
+  const temporaryReferences = createTemporaryReferenceSet()
+  const body = await encodeReply(actionArgs, { temporaryReferences })
 
   const res = await fetch('', {
     method: 'POST',
@@ -140,7 +141,7 @@ async function fetchServerAction(
   if (contentType?.startsWith(RSC_CONTENT_TYPE_HEADER)) {
     const response: ActionFlightResponse = await createFromFetch(
       Promise.resolve(res),
-      { callServer, findSourceMapURL }
+      { callServer, findSourceMapURL, temporaryReferences }
     )
 
     if (location) {

packages/next/src/server/app-render/action-handler.ts

@@ -44,6 +44,7 @@ import { selectWorkerForForwarding } from './action-utils'
 import { isNodeNextRequest, isWebNextRequest } from '../base-http/helpers'
 import { RedirectStatusCode } from '../../client/components/redirect-status-code'
 import { synchronizeMutableCookies } from '../async-storage/request-store'
+import type { TemporaryReferenceSet } from 'react-server-dom-webpack/server.edge'
 
 function formDataFromSearchQueryString(query: string) {
   const searchParams = new URLSearchParams(query)
@@ -394,12 +395,11 @@ function limitUntrustedHeaderValueForLogs(value: string) {
 
 type ServerModuleMap = Record<
   string,
-  | {
-      id: string
-      chunks: string[]
-      name: string
-    }
-  | undefined
+  {
+    id: string
+    chunks: string[]
+    name: string
+  }
 >
 
 type ServerActionsConfig = {
@@ -460,6 +460,8 @@ export async function handleAction({
     )
   }
 
+  let temporaryReferences: TemporaryReferenceSet | undefined
+
   const finalizeAndGenerateFlight: GenerateFlight = (...args) => {
     // When we switch to the render phase, cookies() will return
     // `workUnitStore.cookies` instead of `workUnitStore.userspaceMutableCookies`.
@@ -562,6 +564,7 @@ export async function handleAction({
             actionResult: promise,
             // if the page was not revalidated, we can skip the rendering the flight tree
             skipFlight: !workStore.pathWasRevalidated,
+            temporaryReferences,
           }),
         }
       }
@@ -617,19 +620,31 @@ export async function handleAction({
         process.env.NEXT_RUNTIME === 'edge' &&
         isWebNextRequest(req)
       ) {
-        // Use react-server-dom-webpack/server.edge
-        const { decodeReply, decodeAction, decodeFormState } = ComponentMod
         if (!req.body) {
           throw new Error('invariant: Missing request body.')
         }
 
         // TODO: add body limit
 
+        // Use react-server-dom-webpack/server.edge
+        const {
+          createTemporaryReferenceSet,
+          decodeReply,
+          decodeAction,
+          decodeFormState,
+        } = ComponentMod
+
+        temporaryReferences = createTemporaryReferenceSet()
+
         if (isMultipartAction) {
           // TODO-APP: Add streaming support
           const formData = await req.request.formData()
           if (isFetchAction) {
-            boundActionArguments = await decodeReply(formData, serverModuleMap)
+            boundActionArguments = await decodeReply(
+              formData,
+              serverModuleMap,
+              { temporaryReferences }
+            )
           } else {
             const action = await decodeAction(formData, serverModuleMap)
             if (typeof action === 'function') {
@@ -672,11 +687,16 @@ export async function handleAction({
 
           if (isURLEncodedAction) {
             const formData = formDataFromSearchQueryString(actionData)
-            boundActionArguments = await decodeReply(formData, serverModuleMap)
+            boundActionArguments = await decodeReply(
+              formData,
+              serverModuleMap,
+              { temporaryReferences }
+            )
           } else {
             boundActionArguments = await decodeReply(
               actionData,
-              serverModuleMap
+              serverModuleMap,
+              { temporaryReferences }
             )
           }
         }
@@ -688,11 +708,16 @@ export async function handleAction({
       ) {
         // Use react-server-dom-webpack/server.node which supports streaming
         const {
+          createTemporaryReferenceSet,
           decodeReply,
           decodeReplyFromBusboy,
           decodeAction,
           decodeFormState,
-        } = require(`./react-server.node`)
+        } = require(
+          `./react-server.node`
+        ) as typeof import('./react-server.node')
+
+        temporaryReferences = createTemporaryReferenceSet()
 
         const { Transform } =
           require('node:stream') as typeof import('node:stream')
@@ -742,7 +767,8 @@ export async function handleAction({
 
             boundActionArguments = await decodeReplyFromBusboy(
               busboy,
-              serverModuleMap
+              serverModuleMap,
+              { temporaryReferences }
             )
           } else {
             // React doesn't yet publish a busboy version of decodeAction
@@ -772,7 +798,11 @@ export async function handleAction({
               // Only warn if it's a server action, otherwise skip for other post requests
               warnBadServerActionRequest()
               const actionReturnedState = await action()
-              formState = await decodeFormState(actionReturnedState, formData)
+              formState = await decodeFormState(
+                actionReturnedState,
+                formData,
+                serverModuleMap
+              )
             }
 
             // Skip the fetch path
@@ -799,11 +829,16 @@ export async function handleAction({
 
           if (isURLEncodedAction) {
             const formData = formDataFromSearchQueryString(actionData)
-            boundActionArguments = await decodeReply(formData, serverModuleMap)
+            boundActionArguments = await decodeReply(
+              formData,
+              serverModuleMap,
+              { temporaryReferences }
+            )
           } else {
             boundActionArguments = await decodeReply(
               actionData,
-              serverModuleMap
+              serverModuleMap,
+              { temporaryReferences }
             )
           }
         }
@@ -855,6 +890,7 @@ export async function handleAction({
           actionResult: Promise.resolve(returnVal),
           // if the page was not revalidated, or if the action was forwarded from another worker, we can skip the rendering the flight tree
           skipFlight: !workStore.pathWasRevalidated || actionWasForwarded,
+          temporaryReferences,
         })
       }
     })
@@ -929,6 +965,7 @@ export async function handleAction({
           result: await finalizeAndGenerateFlight(req, ctx, {
             skipFlight: false,
             actionResult: promise,
+            temporaryReferences,
           }),
         }
       }
@@ -963,6 +1000,7 @@ export async function handleAction({
           actionResult: promise,
           // if the page was not revalidated, or if the action was forwarded from another worker, we can skip the rendering the flight tree
           skipFlight: !workStore.pathWasRevalidated || actionWasForwarded,
+          temporaryReferences,
         }),
       }
     }

packages/next/src/server/app-render/app-render.tsx

@@ -492,6 +492,7 @@ async function generateDynamicFlightRenderResult(
     skipFlight: boolean
     componentTree?: CacheNodeSeedData
     preloadCallbacks?: PreloadCallbacks
+    temporaryReferences?: WeakMap<any, string>
   }
 ): Promise<RenderResult> {
   const renderOpts = ctx.renderOpts
@@ -517,6 +518,7 @@ async function generateDynamicFlightRenderResult(
     ctx.clientReferenceManifest.clientModules,
     {
       onError,
+      temporaryReferences: options?.temporaryReferences,
     }
   )
 

packages/next/src/server/app-render/entry-base.ts

@@ -1,5 +1,6 @@
 // eslint-disable-next-line import/no-extraneous-dependencies
 export {
+  createTemporaryReferenceSet,
   renderToReadableStream,
   decodeReply,
   decodeAction,

packages/next/src/server/app-render/react-server.node.ts

@@ -2,6 +2,7 @@
 
 // eslint-disable-next-line import/no-extraneous-dependencies
 export {
+  createTemporaryReferenceSet,
   decodeReply,
   decodeReplyFromBusboy,
   decodeAction,