Fix accessing headers in progressively enhanced form actions (vercel#74196)

When scoping the `requestStore` to dynamic renders in vercel#72312, we missed
the case when a server action is invoked before hydration is complete.

This leads to an error when a server action tries to read headers,
cookies, or anything else that requires the presence of the
`requestStore`.

This fixes that for both the Node.js and Edge runtimes. It also appears
that progressively enhanced form actions did not work at all before in
the Edge runtime due to a missing `await` of `decodeFormState`.

fixes vercel#73992
closes NAR-55

Author: unstubbable

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

@@ -657,12 +657,19 @@ export async function handleAction({
             if (typeof action === 'function') {
               // Only warn if it's a server action, otherwise skip for other post requests
               warnBadServerActionRequest()
-              const actionReturnedState = await action()
-              formState = decodeFormState(
+
+              const actionReturnedState = await workUnitAsyncStorage.run(
+                requestStore,
+                action
+              )
+
+              formState = await decodeFormState(
                 actionReturnedState,
                 formData,
                 serverModuleMap
               )
+
+              requestStore.phase = 'render'
             }
 
             // Skip the fetch path
@@ -804,12 +811,19 @@ export async function handleAction({
             if (typeof action === 'function') {
               // Only warn if it's a server action, otherwise skip for other post requests
               warnBadServerActionRequest()
-              const actionReturnedState = await action()
+
+              const actionReturnedState = await workUnitAsyncStorage.run(
+                requestStore,
+                action
+              )
+
               formState = await decodeFormState(
                 actionReturnedState,
                 formData,
                 serverModuleMap
               )
+
+              requestStore.phase = 'render'
             }
 
             // Skip the fetch path

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

@@ -1,6 +1,6 @@
 /* eslint-disable jest/no-standalone-expect */
 import { nextTestSetup } from 'e2e-utils'
-import { check } from 'next-test-utils'
+import { retry } from 'next-test-utils'
 import type { Response } from 'playwright'
 
 describe('app-dir action progressive enhancement', () => {
@@ -13,7 +13,7 @@ describe('app-dir action progressive enhancement', () => {
   })
 
   it('should support formData and redirect without JS', async () => {
-    let responseCode
+    let responseCode: number | undefined
     const browser = await next.browser('/server', {
       disableJavaScript: true,
       beforePageLoad(page) {
@@ -27,12 +27,14 @@ describe('app-dir action progressive enhancement', () => {
       },
     })
 
-    await browser.eval(`document.getElementById('name').value = 'test'`)
-    await browser.elementByCss('#submit').click()
+    await browser.elementById('name').type('test')
+    await browser.elementById('submit').click()
 
-    await check(() => {
-      return browser.eval('window.location.pathname + window.location.search')
-    }, '/header?name=test&hidden-info=hi')
+    await retry(async () => {
+      expect(await browser.url()).toBe(
+        `${next.url}/header?name=test&hidden-info=hi`
+      )
+    })
 
     expect(responseCode).toBe(303)
   })
@@ -42,11 +44,42 @@ describe('app-dir action progressive enhancement', () => {
       disableJavaScript: true,
     })
 
-    await browser.eval(`document.getElementById('client-name').value = 'test'`)
-    await browser.elementByCss('#there').click()
+    await browser.elementById('client-name').type('test')
+    await browser.elementById('there').click()
 
-    await check(() => {
-      return browser.eval('window.location.pathname + window.location.search')
-    }, '/header?name=test&hidden-info=hi')
+    await retry(async () => {
+      expect(await browser.url()).toBe(
+        `${next.url}/header?name=test&hidden-info=hi`
+      )
+    })
   })
+
+  it.each(['edge', 'node'])(
+    'should support headers and cookies without JS (runtime: %s)',
+    async (runtime) => {
+      const browser = await next.browser(`/header/${runtime}/form`, {
+        disableJavaScript: true,
+      })
+
+      await browser.elementById('get-referer').click()
+
+      await retry(async () => {
+        expect(await browser.elementById('referer').text()).toBe(
+          `${next.url}/header/${runtime}/form`
+        )
+      })
+
+      await browser.elementById('set-cookie').click()
+
+      await retry(async () => {
+        expect(await browser.elementById('referer').text()).toBe('')
+      })
+
+      await browser.elementById('get-cookie').click()
+
+      await retry(async () => {
+        expect(await browser.elementById('cookie').text()).toBe('42')
+      })
+    }
+  )
 })

test/e2e/app-dir/actions/app/header/edge/form/page.tsx

@@ -0,0 +1,3 @@
+export { default } from '../../node/form/page'
+
+export const runtime = 'edge'

test/e2e/app-dir/actions/app/header/node/form/page.tsx

@@ -0,0 +1,32 @@
+'use client'
+
+import { useActionState } from 'react'
+import { getCookie, getHeader, setCookie } from '../../actions'
+
+const getReferrer = getHeader.bind(null, 'referer')
+const setTestCookie = setCookie.bind(null, 'test', '42')
+const getTestCookie = getCookie.bind(null, 'test')
+
+export default function Page() {
+  const [referer, getReferrerAction] = useActionState(getReferrer, null)
+
+  const [cookie, getTestCookieAction] = useActionState<ReturnType<
+    typeof getCookie
+  > | null>(getTestCookie, null)
+
+  return (
+    <>
+      <form action={getReferrerAction}>
+        <p id="referer">{referer}</p>
+        <button id="get-referer">Get Referer</button>
+      </form>
+      <form action={getTestCookieAction}>
+        <p id="cookie">{cookie?.value}</p>
+        <button id="set-cookie" formAction={setTestCookie}>
+          Set Cookie
+        </button>
+        <button id="get-cookie">Get Cookie</button>
+      </form>
+    </>
+  )
+}