Create server module map proxy only once

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

@@ -49,7 +49,10 @@ import { warn } from '../../build/output/log'
 import { RequestCookies, ResponseCookies } from '../web/spec-extension/cookies'
 import { HeadersAdapter } from '../web/spec-extension/adapters/headers'
 import { fromNodeOutgoingHttpHeaders } from '../web/utils'
-import { selectWorkerForForwarding, type ServerModuleMap } from './action-utils'
+import {
+  selectWorkerForForwarding,
+  type ServerModuleMap,
+} from './manifests-singleton'
 import { isNodeNextRequest, isWebNextRequest } from '../base-http/helpers'
 import { RedirectStatusCode } from '../../client/components/redirect-status-code'
 import { synchronizeMutableCookies } from '../async-storage/request-store'

packages/next/src/server/app-render/manifests-singleton.ts

@@ -3,9 +3,19 @@ import type { ClientReferenceManifest } from '../../build/webpack/plugins/flight
 import type { DeepReadonly } from '../../shared/lib/deep-readonly'
 import { InvariantError } from '../../shared/lib/invariant-error'
 import { normalizeAppPath } from '../../shared/lib/router/utils/app-paths'
-import { createServerModuleMap, type ServerModuleMap } from './action-utils'
+import { pathHasPrefix } from '../../shared/lib/router/utils/path-has-prefix'
+import { removePathPrefix } from '../../shared/lib/router/utils/remove-path-prefix'
 import { workAsyncStorage } from './work-async-storage.external'
 
+export interface ServerModuleMap {
+  readonly [name: string]: {
+    readonly id: string | number
+    readonly name: string
+    readonly chunks: Readonly<Array<string>> // currently not used
+    readonly async?: boolean
+  }
+}
+
 // This is a global singleton that is, among other things, also used to
 // encode/decode bound args of server function closures. This can't be using a
 // AsyncLocalStorage as it might happen at the module level.
@@ -162,6 +172,105 @@ function createProxiedClientReferenceManifest(
   ) as DeepReadonly<ClientReferenceManifest>
 }
 
+/**
+ * This function creates a Flight-acceptable server module map proxy from our
+ * Server Reference Manifest similar to our client module map. This is because
+ * our manifest contains a lot of internal Next.js data that are relevant to the
+ * runtime, workers, etc. that React doesn't need to know.
+ */
+function createServerModuleMap(): ServerModuleMap {
+  return new Proxy(
+    {},
+    {
+      get: (_, id: string) => {
+        const workers =
+          getServerActionsManifest()[
+            process.env.NEXT_RUNTIME === 'edge' ? 'edge' : 'node'
+          ]?.[id]?.workers
+
+        if (!workers) {
+          return undefined
+        }
+
+        const workStore = workAsyncStorage.getStore()
+
+        let workerEntry:
+          | { moduleId: string | number; async: boolean }
+          | undefined
+
+        if (workStore) {
+          workerEntry = workers[normalizeWorkerPageName(workStore.page)]
+        } else {
+          // If there's no work store defined, we can assume that a server
+          // module map is needed during module evaluation, e.g. to create a
+          // server action using a higher-order function. Therefore it should be
+          // safe to return any entry from the manifest that matches the action
+          // ID. They all refer to the same module ID, which must also exist in
+          // the current page bundle. TODO: This is currently not guaranteed in
+          // Turbopack, and needs to be fixed.
+          workerEntry = Object.values(workers).at(0)
+        }
+
+        if (!workerEntry) {
+          return undefined
+        }
+
+        const { moduleId, async } = workerEntry
+
+        return { id: moduleId, name: id, chunks: [], async }
+      },
+    }
+  )
+}
+
+/**
+ * The flight entry loader keys actions by bundlePath. bundlePath corresponds
+ * with the relative path (including 'app') to the page entrypoint.
+ */
+function normalizeWorkerPageName(pageName: string) {
+  if (pathHasPrefix(pageName, 'app')) {
+    return pageName
+  }
+
+  return 'app' + pageName
+}
+
+/**
+ * Converts a bundlePath (relative path to the entrypoint) to a routable page
+ * name.
+ */
+function denormalizeWorkerPageName(bundlePath: string) {
+  return normalizeAppPath(removePathPrefix(bundlePath, 'app'))
+}
+
+/**
+ * Checks if the requested action has a worker for the current page.
+ * If not, it returns the first worker that has a handler for the action.
+ */
+export function selectWorkerForForwarding(
+  actionId: string,
+  pageName: string
+): string | undefined {
+  const serverActionsManifest = getServerActionsManifest()
+  const workers =
+    serverActionsManifest[
+      process.env.NEXT_RUNTIME === 'edge' ? 'edge' : 'node'
+    ][actionId]?.workers
+
+  // There are no workers to handle this action, nothing to forward to.
+  if (!workers) {
+    return
+  }
+
+  // If there is an entry for the current page, we don't need to forward.
+  if (workers[normalizeWorkerPageName(pageName)]) {
+    return
+  }
+
+  // Otherwise, grab the first worker that has a handler for this action id.
+  return denormalizeWorkerPageName(Object.keys(workers)[0])
+}
+
 export function setManifestsSingleton({
   page,
   clientReferenceManifest,
@@ -180,10 +289,6 @@ export function setManifestsSingleton({
     )
 
     existingSingleton.serverActionsManifest = serverActionsManifest
-
-    existingSingleton.serverModuleMap = createServerModuleMap({
-      serverActionsManifest,
-    })
   } else {
     const clientReferenceManifestsPerRoute = new Map<
       string,
@@ -194,15 +299,11 @@ export function setManifestsSingleton({
       clientReferenceManifestsPerRoute
     )
 
-    const serverModuleMap = createServerModuleMap({
-      serverActionsManifest,
-    })
-
     globalThisWithManifests[MANIFESTS_SINGLETON] = {
       clientReferenceManifestsPerRoute,
       proxiedClientReferenceManifest,
       serverActionsManifest,
-      serverModuleMap,
+      serverModuleMap: createServerModuleMap(),
     }
   }
 }