ensure cookies set in middleware can be read in a server action (#67924)

### What
If middleware targets a server action handler and sets or updates a
cookie, the newly updated cookie would not be reflected in the
`cookies()` response of the action handler

### Why
In #65008 we fixed a bug where cookies set in middleware were not
reflected in the `cookies()` call in a server component from the same
request. We did this by introducing a `x-middleware-set-cookie` header,
that signaled to downstream handlers that middleware had run on the
request & set a cookie. However this handling was only applied to the
sealed/read-only cookies. Cookies accessed from a server action use
`mutableCookies`, since those aren't frozen as a server action is
allowed to modify cookies.

### How
This pulls the cookie merge handling into a function and applies the
merge to `mutableCookies`.

Fixes #67814
Closes NDX-95

packages/next/src/server/async-storage/with-request-store.ts

@@ -80,6 +80,35 @@ export type RequestContext = {
   serverComponentsHmrCache?: ServerComponentsHmrCache
 }
 
+/**
+ * If middleware set cookies in this request (indicated by `x-middleware-set-cookie`),
+ * then merge those into the existing cookie object, so that when `cookies()` is accessed
+ * it's able to read the newly set cookies.
+ */
+function mergeMiddlewareCookies(
+  req: RequestContext['req'],
+  existingCookies: RequestCookies | ResponseCookies
+) {
+  if (
+    'x-middleware-set-cookie' in req.headers &&
+    typeof req.headers['x-middleware-set-cookie'] === 'string'
+  ) {
+    const setCookieValue = req.headers['x-middleware-set-cookie']
+    const responseHeaders = new Headers()
+
+    for (const cookie of splitCookiesString(setCookieValue)) {
+      responseHeaders.append('set-cookie', cookie)
+    }
+
+    const responseCookies = new ResponseCookies(responseHeaders)
+
+    // Transfer cookies from ResponseCookies to RequestCookies
+    for (const cookie of responseCookies.getAll()) {
+      existingCookies.set(cookie.name, cookie.value ?? '')
+    }
+  }
+}
+
 export const withRequestStore: WithStore<RequestStore, RequestContext> = <
   Result,
 >(
@@ -129,24 +158,7 @@ export const withRequestStore: WithStore<RequestStore, RequestContext> = <
           HeadersAdapter.from(req.headers)
         )
 
-        if (
-          'x-middleware-set-cookie' in req.headers &&
-          typeof req.headers['x-middleware-set-cookie'] === 'string'
-        ) {
-          const setCookieValue = req.headers['x-middleware-set-cookie']
-          const responseHeaders = new Headers()
-
-          for (const cookie of splitCookiesString(setCookieValue)) {
-            responseHeaders.append('set-cookie', cookie)
-          }
-
-          const responseCookies = new ResponseCookies(responseHeaders)
-
-          // Transfer cookies from ResponseCookies to RequestCookies
-          for (const cookie of responseCookies.getAll()) {
-            requestCookies.set(cookie.name, cookie.value ?? '')
-          }
-        }
+        mergeMiddlewareCookies(req, requestCookies)
 
         // Seal the cookies object that'll freeze out any methods that could
         // mutate the underlying data.
@@ -157,11 +169,15 @@ export const withRequestStore: WithStore<RequestStore, RequestContext> = <
     },
     get mutableCookies() {
       if (!cache.mutableCookies) {
-        cache.mutableCookies = getMutableCookies(
+        const mutableCookies = getMutableCookies(
           req.headers,
           renderOpts?.onUpdateCookies ||
             (res ? defaultOnUpdateCookies : undefined)
         )
+
+        mergeMiddlewareCookies(req, mutableCookies)
+
+        cache.mutableCookies = mutableCookies
       }
       return cache.mutableCookies
     },

test/e2e/app-dir/app-middleware/app-middleware.test.ts

@@ -162,6 +162,31 @@ describe('app-dir with middleware', () => {
       )
     })
   })
+
+  it('should be possible to read cookies that are set during the middleware handling of a server action', async () => {
+    const browser = await next.browser('/rsc-cookies')
+    const initialRandom1 = await browser.elementById('rsc-cookie-1').text()
+    const initialRandom2 = await browser.elementById('rsc-cookie-2').text()
+    const totalCookies = await browser.elementById('total-cookies').text()
+
+    // cookies were set in middleware, assert they are present and match the Math.random() pattern
+    expect(initialRandom1).toMatch(/Cookie 1: \d+\.\d+/)
+    expect(initialRandom2).toMatch(/Cookie 2: \d+\.\d+/)
+    expect(totalCookies).toBe('Total Cookie Length: 2')
+
+    expect(await browser.eval('document.cookie')).toBeTruthy()
+
+    await browser.deleteCookies()
+
+    // assert that document.cookie is empty
+    expect(await browser.eval('document.cookie')).toBeFalsy()
+
+    await browser.elementById('submit-server-action').click()
+
+    await retry(() => {
+      expect(next.cliOutput).toMatch(/\[Cookie From Action\]: \d+\.\d+/)
+    })
+  })
 })
 
 describe('app dir - middleware without pages dir', () => {

test/e2e/app-dir/app-middleware/app/rsc-cookies/page.js

@@ -11,6 +11,20 @@ export default function Page() {
       <p id="rsc-cookie-2">Cookie 2: {rscCookie2}</p>
       <p id="total-cookies">Total Cookie Length: {cookies().size}</p>
       <Link href="/rsc-cookies-delete">To Delete Cookies Route</Link>
+
+      <form
+        action={async () => {
+          'use server'
+          console.log(
+            '[Cookie From Action]:',
+            cookies().get('rsc-cookie-value-1').value
+          )
+        }}
+      >
+        <button type="submit" id="submit-server-action">
+          server action
+        </button>
+      </form>
     </div>
   )
 }