forward missing server actions to valid worker if one exists

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

@@ -45,6 +45,7 @@ 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 } from './action-utils'
 
 function formDataFromSearchQueryString(query: string) {
   const searchParams = new URLSearchParams(query)
@@ -149,6 +150,96 @@ async function addRevalidationHeader(
   )
 }
 
+/**
+ * Forwards a server action request to a separate worker. Used when the requested action is not available in the current worker.
+ */
+async function createForwardedActionResponse(
+  req: IncomingMessage,
+  res: ServerResponse,
+  host: Host,
+  workerPathname: string,
+  basePath: string,
+  staticGenerationStore: StaticGenerationStore
+) {
+  if (!host) {
+    throw new Error(
+      'Invariant: Missing `host` header from a forwarded Server Actions request.'
+    )
+  }
+
+  const forwardedHeaders = getForwardedHeaders(req, res)
+
+  // indicate that this action request was forwarded from another worker
+  // we use this to skip rendering the flight tree so that we don't update the UI
+  // with the response from the forwarded worker
+  forwardedHeaders.set('x-action-forwarded', '1')
+
+  const proto =
+    staticGenerationStore.incrementalCache?.requestProtocol || 'https'
+
+  // For standalone or the serverful mode, use the internal origin directly
+  // other than the host headers from the request.
+  const origin = process.env.__NEXT_PRIVATE_ORIGIN || `${proto}://${host.value}`
+
+  const fetchUrl = new URL(`${origin}${basePath}${workerPathname}`)
+
+  try {
+    let readableStream: ReadableStream<Uint8Array> | undefined
+    if (process.env.NEXT_RUNTIME === 'edge') {
+      const webRequest = req as unknown as WebNextRequest
+      if (!webRequest.body) {
+        throw new Error('invariant: Missing request body.')
+      }
+
+      readableStream = webRequest.body
+    } else {
+      // Convert the Node.js readable stream to a Web Stream.
+      readableStream = new ReadableStream({
+        start(controller) {
+          req.on('data', (chunk) => {
+            controller.enqueue(new Uint8Array(chunk))
+          })
+          req.on('end', () => {
+            controller.close()
+          })
+          req.on('error', (err) => {
+            controller.error(err)
+          })
+        },
+      })
+    }
+
+    // Forward the request to the new worker
+    const response = await fetch(fetchUrl, {
+      method: 'POST',
+      body: readableStream,
+      duplex: 'half',
+      headers: forwardedHeaders,
+      next: {
+        // @ts-ignore
+        internal: 1,
+      },
+    })
+
+    if (response.headers.get('content-type') === RSC_CONTENT_TYPE_HEADER) {
+      // copy the headers from the redirect response to the response we're sending
+      for (const [key, value] of response.headers) {
+        if (!actionsForbiddenHeaders.includes(key)) {
+          res.setHeader(key, value)
+        }
+      }
+
+      return new FlightRenderResult(response.body!)
+    } else {
+      // Since we aren't consuming the response body, we cancel it to avoid memory leaks
+      response.body?.cancel()
+    }
+  } catch (err) {
+    // we couldn't stream the forwarded response, so we'll just do a normal redirect
+    console.error(`failed to forward action response`, err)
+  }
+}
+
 async function createRedirectRenderResult(
   req: IncomingMessage,
   res: ServerResponse,
@@ -204,9 +295,7 @@ async function createRedirectRenderResult(
     }
 
     // Ensures that when the path was revalidated we don't return a partial response on redirects
-    // if (staticGenerationStore.pathWasRevalidated) {
     forwardedHeaders.delete('next-router-state-tree')
-    // }
 
     try {
       const response = await fetch(fetchUrl, {
@@ -308,6 +397,7 @@ export async function handleAction({
     }
 > {
   const contentType = req.headers['content-type']
+  const { serverActionsManifest, page } = ctx.renderOpts
 
   const { actionId, isURLEncodedAction, isMultipartAction, isFetchAction } =
     getServerActionRequestMetadata(req)
@@ -430,6 +520,31 @@ export async function handleAction({
   let actionResult: RenderResult | undefined
   let formState: any | undefined
   let actionModId: string | undefined
+  const actionWasForwarded = Boolean(req.headers['x-action-forwarded'])
+
+  if (actionId) {
+    const forwardedWorker = selectWorkerForForwarding(
+      actionId,
+      page,
+      serverActionsManifest
+    )
+
+    // If forwardedWorker is truthy, it means there isn't a worker for the action
+    // in the current handler, so we forward the request to a worker that has the action.
+    if (forwardedWorker) {
+      return {
+        type: 'done',
+        result: await createForwardedActionResponse(
+          req,
+          res,
+          host,
+          forwardedWorker,
+          ctx.renderOpts.basePath,
+          staticGenerationStore
+        ),
+      }
+    }
+  }
 
   try {
     await actionAsyncStorage.run({ isAction: true }, async () => {
@@ -627,8 +742,9 @@ To configure the body size limit for Server Actions, see: https://nextjs.org/doc
 
         actionResult = await generateFlight(ctx, {
           actionResult: Promise.resolve(returnVal),
-          // if the page was not revalidated, we can skip the rendering the flight tree
-          skipFlight: !staticGenerationStore.pathWasRevalidated,
+          // if the page was not revalidated, or if the action was forwarded from another worker, we can skip the rendering the flight tree
+          skipFlight:
+            !staticGenerationStore.pathWasRevalidated || actionWasForwarded,
         })
       }
     })
@@ -734,8 +850,9 @@ To configure the body size limit for Server Actions, see: https://nextjs.org/doc
         type: 'done',
         result: await generateFlight(ctx, {
           actionResult: promise,
-          // if the page was not revalidated, we can skip the rendering the flight tree
-          skipFlight: !staticGenerationStore.pathWasRevalidated,
+          // if the page was not revalidated, or if the action was forwarded from another worker, we can skip the rendering the flight tree
+          skipFlight:
+            !staticGenerationStore.pathWasRevalidated || actionWasForwarded,
         }),
       }
     }

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

@@ -1,4 +1,7 @@
 import type { ActionManifest } from '../../build/webpack/plugins/flight-client-entry-plugin'
+import { normalizeAppPath } from '../../shared/lib/router/utils/app-paths'
+import { pathHasPrefix } from '../../shared/lib/router/utils/path-has-prefix'
+import { removePathPrefix } from '../../shared/lib/router/utils/remove-path-prefix'
 
 // This function creates a Flight-acceptable server module map proxy from our
 // Server Reference Manifest similar to our client module map.
@@ -18,11 +21,57 @@ export function createServerModuleMap({
         return {
           id: serverActionsManifest[
             process.env.NEXT_RUNTIME === 'edge' ? 'edge' : 'node'
-          ][id].workers['app' + pageName],
+          ][id].workers[normalizeWorkerPageName(pageName)],
           name: id,
           chunks: [],
         }
       },
     }
   )
 }
+
+/**
+ * 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,
+  serverActionsManifest: ActionManifest
+) {
+  const workers =
+    serverActionsManifest[
+      process.env.NEXT_RUNTIME === 'edge' ? 'edge' : 'node'
+    ][actionId]?.workers
+  const workerName = normalizeWorkerPageName(pageName)
+
+  // no workers, nothing to forward to
+  if (!workers) return
+
+  // if there is a worker for this page, no need to forward it.
+  if (workers[workerName]) {
+    return
+  }
+
+  // otherwise, grab the first worker that has a handler for this action id
+  return denormalizeWorkerPageName(Object.keys(workers)[0])
+}
+
+/**
+ * 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'))
+}

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

@@ -512,6 +512,45 @@ createNextDescribe(
       ).toBe(true)
     })
 
+    it.each(['node', 'edge'])(
+      'should forward action request to a worker that contains the action handler (%s)',
+      async (runtime) => {
+        const cliOutputIndex = next.cliOutput.length
+        const browser = await next.browser(`/delayed-action/${runtime}`)
+
+        // confirm there's no data yet
+        expect(await browser.elementById('delayed-action-result').text()).toBe(
+          ''
+        )
+
+        // Trigger the delayed action. This will sleep for a few seconds before dispatching the server action handler
+        await browser.elementById('run-action').click()
+
+        // navigate away from the page
+        await browser
+          .elementByCss(`[href='/delayed-action/${runtime}/other']`)
+          .click()
+          .waitForElementByCss('#other-page')
+
+        await retry(async () => {
+          expect(
+            await browser.elementById('delayed-action-result').text()
+          ).toMatch(
+            // matches a Math.random() string
+            /0\.\d+/
+          )
+        })
+
+        // make sure that we still are rendering other-page content
+        expect(await browser.hasElementByCssSelector('#other-page')).toBe(true)
+
+        // make sure we didn't get any errors in the console
+        expect(next.cliOutput.slice(cliOutputIndex)).not.toContain(
+          'Failed to find Server Action'
+        )
+      }
+    )
+
     if (isNextStart) {
       it('should not expose action content in sourcemaps', async () => {
         const sourcemap = (

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

@@ -0,0 +1,8 @@
+'use server'
+import { revalidatePath } from 'next/cache'
+
+export const action = async () => {
+  console.log('revalidating')
+  revalidatePath('/delayed-action', 'page')
+  return Math.random()
+}

test/e2e/app-dir/actions/app/delayed-action/button.tsx

@@ -0,0 +1,21 @@
+'use client'
+
+import { useContext } from 'react'
+import { action } from './actions'
+import { DataContext } from './context'
+
+export function Button() {
+  const { setData } = useContext(DataContext)
+  const handleClick = async () => {
+    await new Promise((res) => setTimeout(res, 1000))
+
+    const result = await action()
+
+    setData(result)
+  }
+  return (
+    <button onClick={handleClick} id="run-action">
+      Run Action
+    </button>
+  )
+}

test/e2e/app-dir/actions/app/delayed-action/context.tsx

@@ -0,0 +1,6 @@
+import React from 'react'
+
+export const DataContext = React.createContext<{
+  data: number | null
+  setData: (number: number) => void
+}>({ data: null, setData: () => {} })

test/e2e/app-dir/actions/app/delayed-action/edge/layout.tsx

@@ -0,0 +1 @@
+export { default } from '../layout-edge'

test/e2e/app-dir/actions/app/delayed-action/edge/other/page.tsx

@@ -0,0 +1 @@
+export { default } from '../../other-page'

test/e2e/app-dir/actions/app/delayed-action/edge/page.tsx

@@ -0,0 +1,15 @@
+import Link from 'next/link'
+import { Button } from '../button'
+
+export default function Page() {
+  return (
+    <>
+      <div>
+        <Link href="/delayed-action/edge/other">Navigate to Other Page</Link>
+      </div>
+      <div>
+        <Button />
+      </div>
+    </>
+  )
+}

test/e2e/app-dir/actions/app/delayed-action/layout-edge.tsx

@@ -0,0 +1,17 @@
+'use client'
+
+export const runtime = 'edge'
+
+import { useState } from 'react'
+import { DataContext } from './context'
+
+export default function Layout({ children }) {
+  const [data, setData] = useState<number | null>(null)
+
+  return (
+    <DataContext.Provider value={{ data, setData }}>
+      <div>{children}</div>
+      <div id="delayed-action-result">{data}</div>
+    </DataContext.Provider>
+  )
+}

test/e2e/app-dir/actions/app/delayed-action/layout-node.tsx

@@ -0,0 +1,15 @@
+'use client'
+
+import { useState } from 'react'
+import { DataContext } from './context'
+
+export default function Layout({ children }) {
+  const [data, setData] = useState<number | null>(null)
+
+  return (
+    <DataContext.Provider value={{ data, setData }}>
+      <div>{children}</div>
+      <div id="delayed-action-result">{data}</div>
+    </DataContext.Provider>
+  )
+}

test/e2e/app-dir/actions/app/delayed-action/node/layout.tsx

@@ -0,0 +1 @@
+export { default } from '../layout-node'

test/e2e/app-dir/actions/app/delayed-action/node/other/page.tsx

@@ -0,0 +1 @@
+export { default } from '../../other-page'