Avoid server action endpoint function indirection (#71572)

This PR removes the extra wrapper functions and `.bind()` calls from the
action entrypoint loader. Instead, the loader module now only re-exports
the actions from the original module using the IDs as export names. This
ensures that the same action instance will always remain identical. It
also unblocks the combined usage of `"use cache"` and `"use server"`
functions.

---------

Co-authored-by: Shu Ding <[email protected]>

Author: unstubbable

packages/next/src/build/webpack/loaders/next-flight-action-entry-loader.ts

@@ -20,29 +20,12 @@ function nextFlightActionEntryLoader(this: any) {
     .flat()
 
   return `
-const actions = {
 ${individualActions
   .map(([id, path, name]) => {
-    return `'${id}': () => import(/* webpackMode: "eager" */ ${JSON.stringify(
-      path
-    )}).then(mod => mod[${JSON.stringify(name)}]),`
+    // Re-export the same functions from the original module path as action IDs.
+    return `export { ${name} as "${id}" } from ${JSON.stringify(path)}`
   })
   .join('\n')}
-}
-
-async function endpoint(id, ...args) {
-  const action = await actions[id]()
-  return action.apply(null, args)
-}
-
-// Using CJS to avoid this to be tree-shaken away due to unused exports.
-module.exports = {
-${individualActions
-  .map(([id]) => {
-    return `  '${id}': endpoint.bind(null, '${id}'),`
-  })
-  .join('\n')}
-}
 `
 }
 

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

@@ -55,7 +55,11 @@ const PLUGIN_NAME = 'FlightClientEntryPlugin'
 type Actions = {
   [actionId: string]: {
     workers: {
-      [name: string]: string | number
+      [name: string]:
+        | { moduleId: string | number; async: boolean }
+        // TODO: This is legacy for Turbopack, and needs to be changed to the
+        // object above.
+        | string
     }
     // Record which layer the action is in (rsc or sc_action), in the specific entry.
     layer: {
@@ -79,15 +83,15 @@ const pluginState = getProxiedPluginState({
   actionModServerId: {} as Record<
     string,
     {
-      server?: string | number
-      client?: string | number
+      server?: { moduleId: string | number; async: boolean }
+      client?: { moduleId: string | number; async: boolean }
     }
   >,
   actionModEdgeServerId: {} as Record<
     string,
     {
-      server?: string | number
-      client?: string | number
+      server?: { moduleId: string | number; async: boolean }
+      client?: { moduleId: string | number; async: boolean }
     }
   >,
 
@@ -879,7 +883,11 @@ export class FlightClientEntryPlugin {
             layer: {},
           }
         }
-        currentCompilerServerActions[id].workers[bundlePath] = ''
+        currentCompilerServerActions[id].workers[bundlePath] = {
+          moduleId: '', // TODO: What's the meaning of this?
+          async: false,
+        }
+
         currentCompilerServerActions[id].layer[bundlePath] = fromClient
           ? WEBPACK_LAYERS.actionBrowser
           : WEBPACK_LAYERS.reactServerComponents
@@ -928,6 +936,13 @@ export class FlightClientEntryPlugin {
           }
 
           compilation.hooks.succeedEntry.call(dependency, options, module)
+
+          compilation.moduleGraph
+            .getExportsInfo(module)
+            .setUsedInUnknownWay(
+              this.isEdgeServer ? EDGE_RUNTIME_WEBPACK : DEFAULT_RUNTIME_WEBPACK
+            )
+
           return resolve(module)
         }
       )
@@ -958,7 +973,10 @@ export class FlightClientEntryPlugin {
         if (!mapping[chunkGroup.name]) {
           mapping[chunkGroup.name] = {}
         }
-        mapping[chunkGroup.name][fromClient ? 'client' : 'server'] = modId
+        mapping[chunkGroup.name][fromClient ? 'client' : 'server'] = {
+          moduleId: modId,
+          async: compilation.moduleGraph.isAsync(mod),
+        }
       }
     })
 

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

@@ -18,13 +18,18 @@ export function createServerModuleMap({
     {},
     {
       get: (_, id: string) => {
-        return {
-          id: serverActionsManifest[
+        const workerEntry =
+          serverActionsManifest[
             process.env.NEXT_RUNTIME === 'edge' ? 'edge' : 'node'
-          ][id].workers[normalizeWorkerPageName(pageName)],
-          name: id,
-          chunks: [],
+          ][id].workers[normalizeWorkerPageName(pageName)]
+
+        if (typeof workerEntry === 'string') {
+          return { id: workerEntry, name: id, chunks: [] }
         }
+
+        const { moduleId, async } = workerEntry
+
+        return { id: moduleId, name: id, chunks: [], async }
       },
     }
   )

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

@@ -828,6 +828,24 @@ describe('app-dir action handling', () => {
     ).toBe(true)
   })
 
+  // we don't have access to runtime logs on deploy
+  if (!isNextDeploy) {
+    it('should keep action instances identical', async () => {
+      const logs: string[] = []
+      next.on('stdout', (log) => {
+        logs.push(log)
+      })
+
+      const browser = await next.browser('/identity')
+
+      await browser.elementByCss('button').click()
+
+      await retry(() => {
+        expect(logs.join('')).toContain('result: true')
+      })
+    })
+  }
+
   it.each(['node', 'edge'])(
     'should forward action request to a worker that contains the action handler (%s)',
     async (runtime) => {

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

@@ -0,0 +1,13 @@
+'use client'
+
+export function Client({ foo, b }) {
+  return (
+    <button
+      onClick={() => {
+        foo(b)
+      }}
+    >
+      Trigger
+    </button>
+  )
+}

test/e2e/app-dir/actions/app/identity/page.js

@@ -0,0 +1,14 @@
+import { Client } from './client'
+
+export default async function Page() {
+  async function b() {
+    'use server'
+  }
+
+  async function foo(a) {
+    'use server'
+    console.log('result:', a === b)
+  }
+
+  return <Client foo={foo} b={b} />
+}