Fix race condition when setting client reference manifests (#71741)

Similar to how #71669 fixed a race condition when setting the manifests
singleton regarding the server action manifests, this PR fixes the same
race condition for the client reference manifests.

Author: unstubbable

packages/next/src/build/templates/edge-app-route.ts

@@ -17,6 +17,7 @@ const rscServerManifest = maybeJSONParse(self.__RSC_SERVER_MANIFEST)
 
 if (rscManifest && rscServerManifest) {
   setReferenceManifestsSingleton({
+    page: 'VAR_PAGE',
     clientReferenceManifest: rscManifest,
     serverActionsManifest: rscServerManifest,
     serverModuleMap: createServerModuleMap({

packages/next/src/build/templates/edge-ssr-app.ts

@@ -56,6 +56,7 @@ const interceptionRouteRewrites =
 
 if (rscManifest && rscServerManifest) {
   setReferenceManifestsSingleton({
+    page: 'VAR_PAGE',
     clientReferenceManifest: rscManifest,
     serverActionsManifest: rscServerManifest,
     serverModuleMap: createServerModuleMap({

packages/next/src/build/webpack/plugins/flight-manifest-plugin.ts

@@ -72,12 +72,21 @@ export interface ManifestNode {
   }
 }
 
-export type ClientReferenceManifest = {
+export interface ClientReferenceManifestForRsc {
+  clientModules: ManifestNode
+  rscModuleMapping: {
+    [moduleId: string]: ManifestNode
+  }
+  edgeRscModuleMapping: {
+    [moduleId: string]: ManifestNode
+  }
+}
+
+export interface ClientReferenceManifest extends ClientReferenceManifestForRsc {
   readonly moduleLoading: {
     prefix: string
     crossOrigin: string | null
   }
-  clientModules: ManifestNode
   ssrModuleMapping: {
     [moduleId: string]: ManifestNode
   }
@@ -90,12 +99,6 @@ export type ClientReferenceManifest = {
   entryJSFiles?: {
     [entry: string]: string[]
   }
-  rscModuleMapping: {
-    [moduleId: string]: ManifestNode
-  }
-  edgeRscModuleMapping: {
-    [moduleId: string]: ManifestNode
-  }
 }
 
 function getAppPathRequiredChunks(

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

@@ -1060,6 +1060,7 @@ async function renderToHTMLOrFlightImpl(
   const serverModuleMap = createServerModuleMap({ serverActionsManifest })
 
   setReferenceManifestsSingleton({
+    page: workStore.page,
     clientReferenceManifest,
     serverActionsManifest,
     serverModuleMap,

packages/next/src/server/app-render/encryption-utils.ts

@@ -1,6 +1,11 @@
 import type { ActionManifest } from '../../build/webpack/plugins/flight-client-entry-plugin'
-import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight-manifest-plugin'
+import type {
+  ClientReferenceManifest,
+  ClientReferenceManifestForRsc,
+} from '../../build/webpack/plugins/flight-manifest-plugin'
 import type { DeepReadonly } from '../../shared/lib/deep-readonly'
+import { InvariantError } from '../../shared/lib/invariant-error'
+import { workAsyncStorage } from './work-async-storage.external'
 
 let __next_loaded_action_key: CryptoKey
 
@@ -64,10 +69,12 @@ const SERVER_ACTION_MANIFESTS_SINGLETON = Symbol.for(
 )
 
 export function setReferenceManifestsSingleton({
+  page,
   clientReferenceManifest,
   serverActionsManifest,
   serverModuleMap,
 }: {
+  page: string
   clientReferenceManifest: DeepReadonly<ClientReferenceManifest>
   serverActionsManifest: DeepReadonly<ActionManifest>
   serverModuleMap: {
@@ -78,9 +85,19 @@ export function setReferenceManifestsSingleton({
     }
   }
 }) {
-  // @ts-ignore
+  // @ts-expect-error
+  const clientReferenceManifestsPerPage = globalThis[
+    SERVER_ACTION_MANIFESTS_SINGLETON
+  ]?.clientReferenceManifestsPerPage as
+    | undefined
+    | DeepReadonly<Record<string, ClientReferenceManifest>>
+
+  // @ts-expect-error
   globalThis[SERVER_ACTION_MANIFESTS_SINGLETON] = {
-    clientReferenceManifest,
+    clientReferenceManifestsPerPage: {
+      ...clientReferenceManifestsPerPage,
+      [normalizePage(page)]: clientReferenceManifest,
+    },
     serverActionsManifest,
     serverModuleMap,
   }
@@ -100,29 +117,49 @@ export function getServerModuleMap() {
   }
 
   if (!serverActionsManifestSingleton) {
-    throw new Error(
-      'Missing manifest for Server Actions. This is a bug in Next.js'
-    )
+    throw new InvariantError('Missing manifest for Server Actions.')
   }
 
   return serverActionsManifestSingleton.serverModuleMap
 }
 
-export function getClientReferenceManifestSingleton() {
+export function getClientReferenceManifestForRsc(): DeepReadonly<ClientReferenceManifestForRsc> {
   const serverActionsManifestSingleton = (globalThis as any)[
     SERVER_ACTION_MANIFESTS_SINGLETON
   ] as {
-    clientReferenceManifest: DeepReadonly<ClientReferenceManifest>
-    serverActionsManifest: DeepReadonly<ActionManifest>
+    clientReferenceManifestsPerPage: DeepReadonly<
+      Record<string, ClientReferenceManifest>
+    >
   }
 
   if (!serverActionsManifestSingleton) {
-    throw new Error(
-      'Missing manifest for Server Actions. This is a bug in Next.js'
-    )
+    throw new InvariantError('Missing manifest for Server Actions.')
+  }
+
+  const { clientReferenceManifestsPerPage } = serverActionsManifestSingleton
+  const workStore = workAsyncStorage.getStore()
+
+  if (!workStore) {
+    // If there's no work store defined, we can assume that a client reference
+    // manifest is needed during module evaluation, e.g. to create a server
+    // action using a higher-order function. This might also use client
+    // components which need to be serialized by Flight, and therefore client
+    // references need to be resolvable. To make this work, we're returning a
+    // merged manifest across all pages. This is fine as long as the module IDs
+    // are not page specific, which they are not for Webpack. TODO: Fix this in
+    // Turbopack.
+    return mergeClientReferenceManifests(clientReferenceManifestsPerPage)
   }
 
-  return serverActionsManifestSingleton.clientReferenceManifest
+  const page = normalizePage(workStore.page)
+
+  const clientReferenceManifest = clientReferenceManifestsPerPage[page]
+
+  if (!clientReferenceManifest) {
+    throw new InvariantError(`Missing Client Reference Manifest for ${page}.`)
+  }
+
+  return clientReferenceManifest
 }
 
 export async function getActionEncryptionKey() {
@@ -133,22 +170,19 @@ export async function getActionEncryptionKey() {
   const serverActionsManifestSingleton = (globalThis as any)[
     SERVER_ACTION_MANIFESTS_SINGLETON
   ] as {
-    clientReferenceManifest: DeepReadonly<ClientReferenceManifest>
     serverActionsManifest: DeepReadonly<ActionManifest>
   }
 
   if (!serverActionsManifestSingleton) {
-    throw new Error(
-      'Missing manifest for Server Actions. This is a bug in Next.js'
-    )
+    throw new InvariantError('Missing manifest for Server Actions.')
   }
 
   const rawKey =
     process.env.NEXT_SERVER_ACTIONS_ENCRYPTION_KEY ||
     serverActionsManifestSingleton.serverActionsManifest.encryptionKey
 
   if (rawKey === undefined) {
-    throw new Error('Missing encryption key for Server Actions')
+    throw new InvariantError('Missing encryption key for Server Actions')
   }
 
   __next_loaded_action_key = await crypto.subtle.importKey(
@@ -161,3 +195,40 @@ export async function getActionEncryptionKey() {
 
   return __next_loaded_action_key
 }
+
+function normalizePage(page: string): string {
+  return page.replace(/\/(page|route)$/, '')
+}
+
+function mergeClientReferenceManifests(
+  clientReferenceManifestsPerPage: DeepReadonly<
+    Record<string, ClientReferenceManifest>
+  >
+): ClientReferenceManifestForRsc {
+  const clientReferenceManifests = Object.values(
+    clientReferenceManifestsPerPage as Record<string, ClientReferenceManifest>
+  )
+
+  const mergedClientReferenceManifest: ClientReferenceManifestForRsc = {
+    clientModules: {},
+    edgeRscModuleMapping: {},
+    rscModuleMapping: {},
+  }
+
+  for (const clientReferenceManifest of clientReferenceManifests) {
+    mergedClientReferenceManifest.clientModules = {
+      ...mergedClientReferenceManifest.clientModules,
+      ...clientReferenceManifest.clientModules,
+    }
+    mergedClientReferenceManifest.edgeRscModuleMapping = {
+      ...mergedClientReferenceManifest.edgeRscModuleMapping,
+      ...clientReferenceManifest.edgeRscModuleMapping,
+    }
+    mergedClientReferenceManifest.rscModuleMapping = {
+      ...mergedClientReferenceManifest.rscModuleMapping,
+      ...clientReferenceManifest.rscModuleMapping,
+    }
+  }
+
+  return mergedClientReferenceManifest
+}

packages/next/src/server/app-render/encryption.ts

@@ -12,7 +12,7 @@ import {
   decrypt,
   encrypt,
   getActionEncryptionKey,
-  getClientReferenceManifestSingleton,
+  getClientReferenceManifestForRsc,
   getServerModuleMap,
   stringToUint8Array,
 } from './encryption-utils'
@@ -70,11 +70,11 @@ async function encodeActionBoundArg(actionId: string, arg: string) {
 
 // Encrypts the action's bound args into a string.
 export async function encryptActionBoundArgs(actionId: string, args: any[]) {
-  const clientReferenceManifestSingleton = getClientReferenceManifestSingleton()
+  const { clientModules } = getClientReferenceManifestForRsc()
 
   // Using Flight to serialize the args into a string.
   const serialized = await streamToString(
-    renderToReadableStream(args, clientReferenceManifestSingleton.clientModules)
+    renderToReadableStream(args, clientModules)
   )
 
   // Encrypt the serialized string with the action id as the salt.
@@ -90,16 +90,17 @@ export async function decryptActionBoundArgs(
   actionId: string,
   encrypted: Promise<string>
 ) {
-  const clientReferenceManifestSingleton = getClientReferenceManifestSingleton()
+  const { edgeRscModuleMapping, rscModuleMapping } =
+    getClientReferenceManifestForRsc()
 
   // Decrypt the serialized string with the action id as the salt.
-  const decryped = await decodeActionBoundArg(actionId, await encrypted)
+  const decrypted = await decodeActionBoundArg(actionId, await encrypted)
 
   // Using Flight to deserialize the args from the string.
   const deserialized = await createFromReadableStream(
     new ReadableStream({
       start(controller) {
-        controller.enqueue(textEncoder.encode(decryped))
+        controller.enqueue(textEncoder.encode(decrypted))
         controller.close()
       },
     }),
@@ -109,9 +110,7 @@ export async function decryptActionBoundArgs(
         // to be added to the current execution. Instead, we'll wait for any ClientReference
         // to be emitted which themselves will handle the preloading.
         moduleLoading: null,
-        moduleMap: isEdgeRuntime
-          ? clientReferenceManifestSingleton.edgeRscModuleMapping
-          : clientReferenceManifestSingleton.rscModuleMapping,
+        moduleMap: isEdgeRuntime ? edgeRscModuleMapping : rscModuleMapping,
         serverModuleMap: getServerModuleMap(),
       },
     }

packages/next/src/server/load-components.ts

@@ -177,6 +177,7 @@ async function loadComponentsImpl<N = any>({
   // to them at the top level of the page module.
   if (serverActionsManifest && clientReferenceManifest) {
     setReferenceManifestsSingleton({
+      page,
       clientReferenceManifest,
       serverActionsManifest,
       serverModuleMap: createServerModuleMap({

packages/next/src/server/use-cache/use-cache-wrapper.ts

@@ -25,10 +25,10 @@ import { makeHangingPromise } from '../dynamic-rendering-utils'
 
 import { cacheScopeAsyncLocalStorage } from '../async-storage/cache-scope.external'
 
-import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight-manifest-plugin'
+import type { ClientReferenceManifestForRsc } from '../../build/webpack/plugins/flight-manifest-plugin'
 
 import {
-  getClientReferenceManifestSingleton,
+  getClientReferenceManifestForRsc,
   getServerModuleMap,
 } from '../app-render/encryption-utils'
 import type { CacheScopeStore } from '../async-storage/cache-scope.external'
@@ -66,7 +66,7 @@ function generateCacheEntry(
   workStore: WorkStore,
   outerWorkUnitStore: WorkUnitStore | undefined,
   cacheScope: undefined | CacheScopeStore,
-  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
   fn: any
 ): Promise<[ReadableStream, Promise<CacheEntry>]> {
@@ -90,7 +90,7 @@ function generateCacheEntryWithRestoredWorkStore(
   workStore: WorkStore,
   outerWorkUnitStore: WorkUnitStore | undefined,
   cacheScope: undefined | CacheScopeStore,
-  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
   fn: any
 ) {
@@ -128,7 +128,7 @@ function generateCacheEntryWithRestoredWorkStore(
 function generateCacheEntryWithCacheContext(
   workStore: WorkStore,
   outerWorkUnitStore: WorkUnitStore | undefined,
-  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
   fn: any
 ) {
@@ -298,7 +298,7 @@ async function generateCacheEntryImpl(
   workStore: WorkStore,
   outerWorkUnitStore: WorkUnitStore | undefined,
   innerCacheStore: UseCacheStore,
-  clientReferenceManifest: DeepReadonly<ClientReferenceManifest>,
+  clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
   fn: any
 ): Promise<[ReadableStream, Promise<CacheEntry>]> {
@@ -455,7 +455,7 @@ export function cache(kind: string, id: string, fn: any) {
 
       // Get the clientReferenceManifest while we're still in the outer Context.
       // In case getClientReferenceManifestSingleton is implemented using AsyncLocalStorage.
-      const clientReferenceManifest = getClientReferenceManifestSingleton()
+      const clientReferenceManifest = getClientReferenceManifestForRsc()
 
       // Because the Action ID is not yet unique per implementation of that Action we can't
       // safely reuse the results across builds yet. In the meantime we add the buildId to the

test/e2e/app-dir/actions/app-action.test.ts

@@ -1587,6 +1587,16 @@ describe('app-dir action handling', () => {
       expect(html).not.toContain('qwerty123')
       expect(html).not.toContain('some-module-level-encryption-value')
     })
+
+    it('should be able to resolve other server actions and client components', async () => {
+      const browser = await next.browser('/encryption')
+      expect(await browser.elementByCss('p').text()).toBe('initial')
+      await browser.elementByCss('button').click()
+
+      await retry(async () => {
+        expect(await browser.elementByCss('p').text()).toBe('hello from client')
+      })
+    })
   })
 
   describe('redirects', () => {

test/e2e/app-dir/actions/app/encryption/client.js

@@ -0,0 +1,5 @@
+'use client'
+
+export function Client() {
+  return 'hello from client'
+}