Repository: vercel/next.js
Pull request: Always apply render preparations after running an action

This PR centralizes some logic that we always want to run after executing a server action into `executeActionAndPrepareForRender`. This fixes some issues where we were only executing revalidates, setting draft mode, and synchronizing mutated cookies for fetch actions, but not MPA actions.

I've also moved the logic for `addRevalidationHeader` to the one place where it makes sense -- after executing a fetch action. This header is only used by the client router and only makes sense to add if we actually executed an action, so e.g. calling it for each possible error type was unnecessary.

Changed files:
filename: packages/next/src/server/app-render/action-handler.ts
patch: @@ -501,20 +501,6 @@ 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`.
-    // We want the render to see any cookie writes that we performed during the action,
-    // so we need to update the immutable cookies to reflect the changes.
-    synchronizeMutableCookies(requestStore)
-
-    // The server action might have toggled draft mode, so we need to reflect
-    // that in the work store to be up-to-date for subsequent rendering.
-    workStore.isDraftMode = requestStore.draftMode.isEnabled
-
-    return generateFlight(...args)
-  }
-
   // When running actions the default is no-store, you can still `cache: 'force-cache'`
   workStore.fetchCache = 'default-no-store'
 
@@ -566,7 +552,6 @@ export async function handleAction({
 
       if (isFetchAction) {
         res.statusCode = 500
-        await executeRevalidates(workStore)
 
         const promise = Promise.reject(error)
         try {
@@ -581,10 +566,10 @@ export async function handleAction({
 
         return {
           type: 'done',
-          result: await finalizeAndGenerateFlight(req, ctx, requestStore, {
+          result: await generateFlight(req, ctx, requestStore, {
             actionResult: promise,
-            // if the page was not revalidated, we can skip the rendering the flight tree
-            skipFlight: !workStore.pathWasRevalidated,
+            // We didn't execute an action, so no revalidations could have occurred. We can skip rendering the page.
+            skipFlight: true,
             temporaryReferences,
           }),
         }
@@ -672,16 +657,13 @@ export async function handleAction({
               // Only warn if it's a server action, otherwise skip for other post requests
               warnBadServerActionRequest()
 
-              let actionReturnedState: unknown
-              requestStore.phase = 'action'
-              try {
-                actionReturnedState = await workUnitAsyncStorage.run(
-                  requestStore,
-                  action
+              const actionReturnedState =
+                await executeActionAndPrepareForRender(
+                  action as () => Promise<unknown>,
+                  [],
+                  workStore,
+                  requestStore
                 )
-              } finally {
-                requestStore.phase = 'render'
-              }
 
               formState = await decodeFormState(
                 actionReturnedState,
@@ -846,16 +828,13 @@ export async function handleAction({
               // Only warn if it's a server action, otherwise skip for other post requests
               warnBadServerActionRequest()
 
-              let actionReturnedState: unknown
-              requestStore.phase = 'action'
-              try {
-                actionReturnedState = await workUnitAsyncStorage.run(
-                  requestStore,
-                  action
+              const actionReturnedState =
+                await executeActionAndPrepareForRender(
+                  action as () => Promise<unknown>,
+                  [],
+                  workStore,
+                  requestStore
                 )
-              } finally {
-                requestStore.phase = 'render'
-              }
 
               formState = await decodeFormState(
                 actionReturnedState,
@@ -938,22 +917,18 @@ export async function handleAction({
           actionId!
         ]
 
-      let returnVal: unknown
-      requestStore.phase = 'action'
-      try {
-        returnVal = await workUnitAsyncStorage.run(requestStore, () =>
-          actionHandler.apply(null, boundActionArguments)
-        )
-      } finally {
-        requestStore.phase = 'render'
-      }
+      const returnVal = await executeActionAndPrepareForRender(
+        actionHandler,
+        boundActionArguments,
+        workStore,
+        requestStore
+      ).finally(() => {
+        addRevalidationHeader(res, { workStore, requestStore })
+      })
 
       // For form actions, we need to continue rendering the page.
       if (isFetchAction) {
-        await executeRevalidates(workStore)
-        addRevalidationHeader(res, { workStore, requestStore })
-
-        actionResult = await finalizeAndGenerateFlight(req, ctx, requestStore, {
+        actionResult = await generateFlight(req, ctx, requestStore, {
           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,
@@ -972,9 +947,6 @@ export async function handleAction({
       const redirectUrl = getURLFromRedirectError(err)
       const redirectType = getRedirectTypeFromError(err)
 
-      await executeRevalidates(workStore)
-      addRevalidationHeader(res, { workStore, requestStore })
-
       // if it's a fetch action, we'll set the status code for logging/debugging purposes
       // but we won't set a Location header, as the redirect will be handled by the client router
       res.statusCode = RedirectStatusCode.SeeOther
@@ -1002,9 +974,6 @@ export async function handleAction({
     } else if (isHTTPAccessFallbackError(err)) {
       res.statusCode = getAccessFallbackHTTPStatus(err)
 
-      await executeRevalidates(workStore)
-      addRevalidationHeader(res, { workStore, requestStore })
-
       if (isFetchAction) {
         const promise = Promise.reject(err)
         try {
@@ -1018,7 +987,7 @@ export async function handleAction({
         }
         return {
           type: 'done',
-          result: await finalizeAndGenerateFlight(req, ctx, requestStore, {
+          result: await generateFlight(req, ctx, requestStore, {
             skipFlight: false,
             actionResult: promise,
             temporaryReferences,
@@ -1035,7 +1004,6 @@ export async function handleAction({
       // so that we can respond with a 413 to requests that break the body size limit
       // (but if we do that, we also need to make sure that whatever handles the non-fetch error path below does the same)
       res.statusCode = 500
-      await executeRevalidates(workStore)
       const promise = Promise.reject(err)
       try {
         // we need to await the promise to trigger the rejection early
@@ -1062,6 +1030,38 @@ export async function handleAction({
   }
 }
 
+async function executeActionAndPrepareForRender<
+  TFn extends (...args: any[]) => Promise<any>,
+>(
+  action: TFn,
+  args: Parameters<TFn>,
+  workStore: WorkStore,
+  requestStore: RequestStore
+): Promise<Awaited<ReturnType<TFn>>> {
+  requestStore.phase = 'action'
+  try {
+    return await workUnitAsyncStorage.run(requestStore, () =>
+      action.apply(null, args)
+    )
+  } finally {
+    requestStore.phase = 'render'
+
+    // When we switch to the render phase, cookies() will return
+    // `workUnitStore.cookies` instead of `workUnitStore.userspaceMutableCookies`.
+    // We want the render to see any cookie writes that we performed during the action,
+    // so we need to update the immutable cookies to reflect the changes.
+    synchronizeMutableCookies(requestStore)
+
+    // The server action might have toggled draft mode, so we need to reflect
+    // that in the work store to be up-to-date for subsequent rendering.
+    workStore.isDraftMode = requestStore.draftMode.isEnabled
+
+    // If the action called revalidateTag/revalidatePath, then that might affect data used by the subsequent render,
+    // so we need to make sure all revalidations are applied before that
+    await executeRevalidates(workStore)
+  }
+}
+
 /**
  * Attempts to find the module ID for the action from the module map. When this fails, it could be a deployment skew where
  * the action came from a different deployment. It could also simply be an invalid POST request that is not a server action.

filename: test/e2e/app-dir/use-cache/use-cache.test.ts
patch: @@ -740,96 +740,121 @@ describe('use-cache', () => {
     expect(text).toBe('This page could not be found.')
   })
 
-  it('should not read nor write cached data when draft mode is enabled', async () => {
-    const browser = await next.browser('/draft-mode')
+  it.each([
+    { description: 'js enabled', disableJavaScript: false },
+    { description: 'js disabled', disableJavaScript: true },
+  ])(
+    'should not read nor write cached data when draft mode is enabled: $description',
+    async ({ disableJavaScript }) => {
+      const browser = await next.browser('/draft-mode', {
+        // This test relies on a server action to set draft mode.
+        // To ensure that it works for both fetch actions and MPA actions,
+        // we test it with javascript disabled too.
+        // (this is because of a bug where draft mode status was not correctly propagated to the workStore for MPA actions)
+        disableJavaScript,
+      })
 
-    expect(await browser.elementByCss('button#toggle').text()).toBe(
-      'Enable Draft Mode'
-    )
+      const refreshAfterServerAction = async () => {
+        if (disableJavaScript) {
+          // browser.refresh() seems to automatically resubmit POST requests,
+          // so if we submitted an MPA action, it'll trigger the action again,
+          // which in this case will toggle draftMode again.
+          await browser.get(new URL('/draft-mode', next.url).href)
+        } else {
+          await browser.refresh()
+        }
+      }
 
-    let initialTopLevelValue = await browser.elementById('top-level').text()
+      expect(await browser.elementByCss('button#toggle').text()).toBe(
+        'Enable Draft Mode'
+      )
 
-    if (isNextDeploy) {
-      await retry(async () => {
-        // Wait for the background revalidation after the deployment to settle.
-        await browser.refresh()
+      let initialTopLevelValue = await browser.elementById('top-level').text()
 
-        expect(await browser.elementById('top-level').text()).not.toBe(
-          initialTopLevelValue
-        )
-      })
+      if (isNextDeploy) {
+        await retry(async () => {
+          // Wait for the background revalidation after the deployment to settle.
+          await browser.refresh()
 
-      initialTopLevelValue = await browser.elementById('top-level').text()
-    }
+          expect(await browser.elementById('top-level').text()).not.toBe(
+            initialTopLevelValue
+          )
+        })
 
-    // Draft mode is disabled, cached data should be returned on refresh.
+        initialTopLevelValue = await browser.elementById('top-level').text()
+      }
 
-    const initialClosureValue = await browser.elementById('closure').text()
+      // Draft mode is disabled, cached data should be returned on refresh.
 
-    await browser.refresh()
+      const initialClosureValue = await browser.elementById('closure').text()
 
-    expect(await browser.elementById('top-level').text()).toBe(
-      initialTopLevelValue
-    )
-    expect(await browser.elementById('closure').text()).toBe(
-      initialClosureValue
-    )
+      await browser.refresh()
 
-    await browser.elementByCss('button#toggle').click()
-    await browser.waitForElementByCss('button#toggle:enabled')
+      expect(await browser.elementById('top-level').text()).toBe(
+        initialTopLevelValue
+      )
+      expect(await browser.elementById('closure').text()).toBe(
+        initialClosureValue
+      )
 
-    expect(await browser.elementByCss('button#toggle').text()).toBe(
-      'Disable Draft Mode'
-    )
+      await browser.elementByCss('button#toggle').click()
+      await browser.waitForElementByCss('button#toggle:enabled')
 
-    // Draft mode is now enabled, no cached data should be returned on refresh.
+      expect(await browser.elementByCss('button#toggle').text()).toBe(
+        'Disable Draft Mode'
+      )
 
-    const newTopLevelValue = await browser.elementById('top-level').text()
-    const newClosureValue = await browser.elementById('closure').text()
-    console.log(await browser.elementById('top-level').text())
+      // Draft mode is now enabled, no cached data should be returned on refresh.
 
-    expect(newTopLevelValue).not.toBe(initialTopLevelValue)
-    expect(newClosureValue).not.toBe(initialClosureValue)
+      const newTopLevelValue = await browser.elementById('top-level').text()
+      const newClosureValue = await browser.elementById('closure').text()
+      console.log(await browser.elementById('top-level').text())
 
-    await browser.refresh()
+      expect(newTopLevelValue).not.toBe(initialTopLevelValue)
+      expect(newClosureValue).not.toBe(initialClosureValue)
 
-    expect(await browser.elementById('top-level').text()).not.toBe(
-      newTopLevelValue
-    )
-    console.log(await browser.elementById('top-level').text())
+      await refreshAfterServerAction()
 
-    expect(await browser.elementById('closure').text()).not.toBe(
-      newClosureValue
-    )
+      expect(await browser.elementById('top-level').text()).not.toBe(
+        newTopLevelValue
+      )
+      console.log(await browser.elementById('top-level').text())
 
-    // Accessing request-scoped data should still not be allowed.
-    expect(
-      await browser
-        .elementById('is-accessing-request-scoped-data-allowed-in-use-cache')
-        .text()
-    ).toBe('false')
+      expect(await browser.elementById('closure').text()).not.toBe(
+        newClosureValue
+      )
 
-    await browser.elementByCss('button#toggle').click()
-    await browser.waitForElementByCss('button#toggle:enabled')
+      // Accessing request-scoped data should still not be allowed.
+      expect(
+        await browser
+          .elementById('is-accessing-request-scoped-data-allowed-in-use-cache')
+          .text()
+      ).toBe('false')
 
-    expect(await browser.elementByCss('button#toggle').text()).toBe(
-      'Enable Draft Mode'
-    )
+      await browser.elementByCss('button#toggle').click()
+      await browser.waitForElementByCss('button#toggle:enabled')
 
-    // Draft mode is disabled again, the initially cached data should be
-    // returned again.
+      expect(await browser.elementByCss('button#toggle').text()).toBe(
+        'Enable Draft Mode'
+      )
 
-    console.log(await browser.elementById('top-level').text())
-    await browser.refresh()
-    console.log(await browser.elementById('top-level').text())
+      // Draft mode is disabled again, the initially cached data should be
+      // returned again.
 
-    expect(await browser.elementById('top-level').text()).toBe(
-      initialTopLevelValue
-    )
-    expect(await browser.elementById('closure').text()).toBe(
-      initialClosureValue
-    )
-  })
+      console.log(await browser.elementById('top-level').text())
+
+      await refreshAfterServerAction()
+
+      console.log(await browser.elementById('top-level').text())
+
+      expect(await browser.elementById('top-level').text()).toBe(
+        initialTopLevelValue
+      )
+      expect(await browser.elementById('closure').text()).toBe(
+        initialClosureValue
+      )
+    }
+  )
 
   if (isNextDev) {
     it('should not have unhandled rejection of Request data promises when use cache is enabled without dynamicIO', async () => {