Implement MutableRequestCookies in server entries (#48847)

Similar to #47922 but based off the latest server implementation and #48626:

> This PR implements the MutableRequestCookies instance for cookies() based on the current async context, so we can allow setting cookies in certain places such as Server Functions and Route handlers. Note that to support Route Handlers, we need to also implement the logic of merging Response's Set-Cookie header and the cookies() mutations, hence it's not included in this PR.
>
> fix [NEXT-942](https://linear.app/vercel/issue/NEXT-942)

This PR also adds the same support for Custom Routes.

cc @styfle.

fix NEXT-942, fix NEXT-941.

packages/next/src/client/components/action-async-storage.ts

@@ -1,47 +1,11 @@
 import type { AsyncLocalStorage } from 'async_hooks'
+import { createAsyncLocalStorage } from './async-local-storage'
 
 export interface ActionStore {
-  readonly isAction: boolean
+  readonly isAction?: boolean
+  readonly isAppRoute?: boolean
 }
 
 export type ActionAsyncStorage = AsyncLocalStorage<ActionStore>
 
-let createAsyncLocalStorage: () => ActionAsyncStorage
-if (process.env.NEXT_RUNTIME === 'edge') {
-  createAsyncLocalStorage = () => {
-    let store: ActionStore | undefined
-    return {
-      disable() {
-        throw new Error(
-          'Invariant: AsyncLocalStorage accessed in runtime where it is not available'
-        )
-      },
-      getStore() {
-        return store
-      },
-      async run<R>(s: ActionStore, fn: () => R): Promise<R> {
-        store = s
-        try {
-          return await fn()
-        } finally {
-          store = undefined
-        }
-      },
-      exit() {
-        throw new Error(
-          'Invariant: AsyncLocalStorage accessed in runtime where it is not available'
-        )
-      },
-      enterWith() {
-        throw new Error(
-          'Invariant: AsyncLocalStorage accessed in runtime where it is not available'
-        )
-      },
-    } as ActionAsyncStorage
-  }
-} else {
-  createAsyncLocalStorage =
-    require('./async-local-storage').createAsyncLocalStorage
-}
-
 export const actionAsyncStorage: ActionAsyncStorage = createAsyncLocalStorage()

packages/next/src/client/components/headers.ts

@@ -2,6 +2,7 @@ import { RequestCookiesAdapter } from '../../server/web/spec-extension/adapters/
 import { HeadersAdapter } from '../../server/web/spec-extension/adapters/headers'
 import { RequestCookies } from '../../server/web/spec-extension/cookies'
 import { requestAsyncStorage } from './request-async-storage'
+import { actionAsyncStorage } from './action-async-storage'
 import { staticGenerationBailout } from './static-generation-bailout'
 
 export function headers() {
@@ -42,5 +43,13 @@ export function cookies() {
     )
   }
 
+  const asyncActionStore = actionAsyncStorage.getStore()
+  if (
+    asyncActionStore &&
+    (asyncActionStore.isAction || asyncActionStore.isAppRoute)
+  ) {
+    return requestStore.mutableCookies
+  }
+
   return requestStore.cookies
 }

packages/next/src/client/components/request-async-storage.ts

@@ -1,4 +1,5 @@
 import type { AsyncLocalStorage } from 'async_hooks'
+import type { RequestCookies } from '../../server/web/spec-extension/cookies'
 import type { PreviewData } from '../../../types'
 import type { ReadonlyHeaders } from '../../server/web/spec-extension/adapters/headers'
 import type { ReadonlyRequestCookies } from '../../server/web/spec-extension/adapters/request-cookies'
@@ -8,6 +9,7 @@ import { createAsyncLocalStorage } from './async-local-storage'
 export interface RequestStore {
   readonly headers: ReadonlyHeaders
   readonly cookies: ReadonlyRequestCookies
+  readonly mutableCookies: RequestCookies
   readonly previewData: PreviewData
 }
 

packages/next/src/server/async-storage/request-async-storage-wrapper.ts

@@ -14,6 +14,7 @@ import {
   type ReadonlyHeaders,
 } from '../web/spec-extension/adapters/headers'
 import {
+  MutableRequestCookiesAdapter,
   RequestCookiesAdapter,
   type ReadonlyRequestCookies,
 } from '../web/spec-extension/adapters/request-cookies'
@@ -35,6 +36,14 @@ function getCookies(
   return RequestCookiesAdapter.seal(cookies)
 }
 
+function getMutableCookies(
+  headers: Headers | IncomingHttpHeaders,
+  res: ServerResponse | BaseNextResponse | undefined
+): RequestCookies {
+  const cookies = new RequestCookies(HeadersAdapter.from(headers))
+  return MutableRequestCookiesAdapter.seal(cookies, res)
+}
+
 /**
  * Tries to get the preview data on the request for the given route. This
  * isn't enabled in the edge runtime yet.
@@ -80,6 +89,7 @@ export const RequestAsyncStorageWrapper: AsyncStorageWrapper<
     const cache: {
       headers?: ReadonlyHeaders
       cookies?: ReadonlyRequestCookies
+      mutableCookies?: RequestCookies
     } = {}
 
     const store: RequestStore = {
@@ -101,6 +111,12 @@ export const RequestAsyncStorageWrapper: AsyncStorageWrapper<
 
         return cache.cookies
       },
+      get mutableCookies() {
+        if (!cache.mutableCookies) {
+          cache.mutableCookies = getMutableCookies(req.headers, res)
+        }
+        return cache.mutableCookies
+      },
       previewData,
     }
 

packages/next/src/server/future/route-modules/app-route/module.ts

@@ -31,6 +31,9 @@ import { RouteKind } from '../../route-kind'
 import * as Log from '../../../../build/output/log'
 import { autoImplementMethods } from './helpers/auto-implement-methods'
 import { getNonStaticMethods } from './helpers/get-non-static-methods'
+import { SYMBOL_MODIFY_COOKIE_VALUES } from '../../../web/spec-extension/adapters/request-cookies'
+import { ResponseCookies } from '../../../web/spec-extension/cookies'
+import { HeadersAdapter } from '../../../web/spec-extension/adapters/headers'
 
 /**
  * AppRouteRouteHandlerContext is the context that is passed to the route
@@ -249,93 +252,138 @@ export class AppRouteRouteModule extends RouteModule<
     // Run the handler with the request AsyncLocalStorage to inject the helper
     // support. We set this to `unknown` because the type is not known until
     // runtime when we do a instanceof check below.
-    const response: unknown = await RequestAsyncStorageWrapper.wrap(
+    const response: unknown = await this.actionAsyncStorage.run(
-      this.requestAsyncStorage,
+      {
-      requestContext,
+        isAppRoute: true,
+      },
       () =>
-        StaticGenerationAsyncStorageWrapper.wrap(
+        RequestAsyncStorageWrapper.wrap(
-          this.staticGenerationAsyncStorage,
+          this.requestAsyncStorage,
-          staticGenerationContext,
+          requestContext,
-          (staticGenerationStore) => {
+          () =>
-            // Check to see if we should bail out of static generation based on
+            StaticGenerationAsyncStorageWrapper.wrap(
-            // having non-static methods.
+              this.staticGenerationAsyncStorage,
-            if (this.nonStaticMethods) {
+              staticGenerationContext,
-              this.staticGenerationBailout(
+              (staticGenerationStore) => {
-                `non-static methods used ${this.nonStaticMethods.join(', ')}`
+                // Check to see if we should bail out of static generation based on
-              )
+                // having non-static methods.
-            }
+                if (this.nonStaticMethods) {
+                  this.staticGenerationBailout(
+                    `non-static methods used ${this.nonStaticMethods.join(
+                      ', '
+                    )}`
+                  )
+                }
 
-            // Update the static generation store based on the dynamic property.
+                // Update the static generation store based on the dynamic property.
-            switch (this.dynamic) {
+                switch (this.dynamic) {
-              case 'force-dynamic':
+                  case 'force-dynamic':
-                // The dynamic property is set to force-dynamic, so we should
+                    // The dynamic property is set to force-dynamic, so we should
-                // force the page to be dynamic.
+                    // force the page to be dynamic.
-                staticGenerationStore.forceDynamic = true
+                    staticGenerationStore.forceDynamic = true
-                this.staticGenerationBailout(`force-dynamic`, {
+                    this.staticGenerationBailout(`force-dynamic`, {
-                  dynamic: this.dynamic,
+                      dynamic: this.dynamic,
-                })
+                    })
-                break
+                    break
-              case 'force-static':
+                  case 'force-static':
-                // The dynamic property is set to force-static, so we should
+                    // The dynamic property is set to force-static, so we should
-                // force the page to be static.
+                    // force the page to be static.
-                staticGenerationStore.forceStatic = true
+                    staticGenerationStore.forceStatic = true
-                break
+                    break
-              case 'error':
+                  case 'error':
-                // The dynamic property is set to error, so we should throw an
+                    // The dynamic property is set to error, so we should throw an
-                // error if the page is being statically generated.
+                    // error if the page is being statically generated.
-                staticGenerationStore.dynamicShouldError = true
+                    staticGenerationStore.dynamicShouldError = true
-                break
+                    break
-              default:
+                  default:
-                break
+                    break
-            }
+                }
 
-            // If the static generation store does not have a revalidate value
+                // If the static generation store does not have a revalidate value
-            // set, then we should set it the revalidate value from the userland
+                // set, then we should set it the revalidate value from the userland
-            // module or default to false.
+                // module or default to false.
-            staticGenerationStore.revalidate ??=
+                staticGenerationStore.revalidate ??=
-              this.userland.revalidate ?? false
+                  this.userland.revalidate ?? false
 
-            // Wrap the request so we can add additional functionality to cases
+                // Wrap the request so we can add additional functionality to cases
-            // that might change it's output or affect the rendering.
+                // that might change it's output or affect the rendering.
-            const wrappedRequest = proxyRequest(
+                const wrappedRequest = proxyRequest(
-              request,
+                  request,
-              { dynamic: this.dynamic },
+                  { dynamic: this.dynamic },
-              {
+                  {
-                headerHooks: this.headerHooks,
+                    headerHooks: this.headerHooks,
-                serverHooks: this.serverHooks,
+                    serverHooks: this.serverHooks,
-                staticGenerationBailout: this.staticGenerationBailout,
+                    staticGenerationBailout: this.staticGenerationBailout,
-              }
+                  }
-            )
+                )
-
+
-            // TODO: propagate this pathname from route matcher
+                // TODO: propagate this pathname from route matcher
-            const route = getPathnameFromAbsolutePath(this.resolvedPagePath)
+                const route = getPathnameFromAbsolutePath(this.resolvedPagePath)
-            getTracer().getRootSpanAttributes()?.set('next.route', route)
+                getTracer().getRootSpanAttributes()?.set('next.route', route)
-            return getTracer().trace(
+                return getTracer().trace(
-              AppRouteRouteHandlersSpan.runHandler,
+                  AppRouteRouteHandlersSpan.runHandler,
-              {
+                  {
-                spanName: `executing api route (app) ${route}`,
+                    spanName: `executing api route (app) ${route}`,
-                attributes: {
+                    attributes: {
-                  'next.route': route,
+                      'next.route': route,
-                },
+                    },
-              },
+                  },
-              async () => {
+                  async () => {
-                // Patch the global fetch.
+                    // Patch the global fetch.
-                patchFetch({
+                    patchFetch({
-                  serverHooks: this.serverHooks,
+                      serverHooks: this.serverHooks,
-                  staticGenerationAsyncStorage:
+                      staticGenerationAsyncStorage:
-                    this.staticGenerationAsyncStorage,
+                        this.staticGenerationAsyncStorage,
-                })
+                    })
-                const res = await handler(wrappedRequest, {
+                    const res = await handler(wrappedRequest, {
-                  params: context.params,
+                      params: context.params,
-                })
+                    })
-
+
-                await Promise.all(
+                    await Promise.all(
-                  staticGenerationStore.pendingRevalidates || []
+                      staticGenerationStore.pendingRevalidates || []
+                    )
+
+                    // It's possible cookies were set in the handler, so we need
+                    // to merge the modified cookies and the returned response
+                    // here.
+                    // TODO: Move this into a helper function.
+                    const requestStore = this.requestAsyncStorage.getStore()
+                    if (requestStore && requestStore.mutableCookies) {
+                      const modifiedCookieValues = (
+                        requestStore.mutableCookies as any
+                      )[SYMBOL_MODIFY_COOKIE_VALUES] as [string, string][]
+                      if (modifiedCookieValues.length) {
+                        // Return a new response that extends the response with
+                        // the modified cookies as fallbacks. `res`' cookies
+                        // will still take precedence.
+                        const resCookies = new ResponseCookies(
+                          HeadersAdapter.from(res.headers)
+                        )
+                        const finalCookies = resCookies.getAll()
+
+                        // Set the modified cookies as fallbacks.
+                        modifiedCookieValues.forEach((cookie) =>
+                          resCookies.set(cookie[0], cookie[1])
+                        )
+                        // Set the original cookies as the final values.
+                        finalCookies.forEach((cookie) => resCookies.set(cookie))
+
+                        return new Response(res.body, {
+                          status: res.status,
+                          statusText: res.statusText,
+                          headers: {
+                            ...res.headers,
+                            'Set-Cookie': resCookies.toString(),
+                          },
+                        })
+                      }
+                    }
+
+                    return res
+                  }
                 )
-                return res
               }
             )
-          }
         )
     )
 

packages/next/src/server/future/route-modules/route-module.ts

@@ -15,6 +15,8 @@ const headerHooks =
   require('next/dist/client/components/headers') as typeof import('../../../client/components/headers')
 const { staticGenerationBailout } =
   require('next/dist/client/components/static-generation-bailout') as typeof import('../../../client/components/static-generation-bailout')
+const { actionAsyncStorage } =
+  require('next/dist/client/components/action-async-storage') as typeof import('../../../client/components/action-async-storage')
 
 /**
  * RouteModuleOptions is the options that are passed to the route module, other
@@ -72,6 +74,12 @@ export abstract class RouteModule<
    */
   public readonly staticGenerationBailout = staticGenerationBailout
 
+  /**
+   * A reference to the mutation related async storage, such as mutations of
+   * cookies.
+   */
+  public readonly actionAsyncStorage = actionAsyncStorage
+
   /**
    * The userland module. This is the module that is exported from the user's
    * code. This is marked as readonly to ensure that the module is not mutated

packages/next/src/server/web/spec-extension/adapters/request-cookies.ts

@@ -1,4 +1,7 @@
 import type { RequestCookies } from '../cookies'
+import type { BaseNextResponse } from '../../../base-http'
+import type { ServerResponse } from 'http'
+
 import { ReflectAdapter } from './reflect'
 
 /**
@@ -37,3 +40,92 @@ export class RequestCookiesAdapter {
     })
   }
 }
+
+export const SYMBOL_MODIFY_COOKIE_VALUES = Symbol.for('next.mutated.cookies')
+
+export class MutableRequestCookiesAdapter {
+  public static seal(
+    cookies: RequestCookies,
+    res: ServerResponse | BaseNextResponse | undefined
+  ): RequestCookies {
+    let modifiedValues: [string, string][] = []
+    const modifiedCookies = new Set<string>()
+    const updateResponseCookies = () => {
+      const allCookies = cookies.getAll()
+      modifiedValues = allCookies
+        .filter((c) => modifiedCookies.has(c.name))
+        .map((c) => [c.name, c.value])
+      if (res) {
+        res.setHeader(
+          'Set-Cookie',
+          modifiedValues.map((c) => `${c[0]}=${c[1]}`)
+        )
+      }
+    }
+
+    return new Proxy(cookies, {
+      get(target, prop, receiver) {
+        switch (prop) {
+          // A special symbol to get the modified cookie values
+          case SYMBOL_MODIFY_COOKIE_VALUES:
+            return modifiedValues
+
+          // TODO: Throw error if trying to set a cookie after the response
+          // headers have been set.
+          case 'clear':
+            return function () {
+              for (const c of cookies.getAll()) {
+                modifiedCookies.add(c.name)
+              }
+              try {
+                return cookies.clear()
+              } finally {
+                updateResponseCookies()
+              }
+            }
+          case 'delete':
+            return function (names: string | string[]) {
+              if (Array.isArray(names)) {
+                names.forEach((name) => modifiedCookies.add(name))
+              } else {
+                modifiedCookies.add(names)
+              }
+              try {
+                return cookies.delete(names)
+              } finally {
+                updateResponseCookies()
+              }
+            }
+          case 'set':
+            return function (
+              ...args:
+                | [string, string]
+                | [
+                    options: NonNullable<
+                      ReturnType<InstanceType<typeof RequestCookies>['get']>
+                    >
+                  ]
+            ) {
+              const [key, value] = args
+              if (typeof key === 'string') {
+                modifiedCookies.add(key)
+                try {
+                  return cookies.set(key, value!)
+                } finally {
+                  updateResponseCookies()
+                }
+              }
+              modifiedCookies.add(key.name)
+              try {
+                return cookies.set(key)
+              } finally {
+                updateResponseCookies()
+              }
+            }
+          default:
+            return ReflectAdapter.get(target, prop, receiver)
+        }
+      },
+    })
+  }
+}

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

@@ -42,6 +42,24 @@ createNextDescribe(
         const res = (await browser.elementByCss('h1').text()) || ''
         return res.includes('Mozilla') ? 'UA' : ''
       }, 'UA')
+
+      // Set cookies
+      await browser.elementByCss('#setCookie').click()
+      await check(async () => {
+        const res = (await browser.elementByCss('h1').text()) || ''
+        const id = res.split(':')
+        return id[0] === id[1] && id[0] === id[2] && id[0]
+          ? 'same'
+          : 'different'
+      }, 'same')
+    })
+
+    it('should support setting cookies in route handlers with the correct overrides', async () => {
+      const res = await next.fetch('/handler')
+      const setCookieHeader = res.headers.get('set-cookie')
+      expect(setCookieHeader).toInclude('bar=bar2;')
+      expect(setCookieHeader).toInclude('baz=baz2;')
+      expect(setCookieHeader).toInclude('foo=foo1;')
     })
 
     it('should support formData and redirect', async () => {

test/e2e/app-dir/actions/app/handler/route.js

@@ -0,0 +1,12 @@
+import { cookies } from 'next/headers'
+
+export const GET = async () => {
+  cookies().set('foo', 'foo1')
+  cookies().set('bar', 'bar1')
+  return new Response('Hello, world!', {
+    headers: [
+      ['Set-Cookie', 'bar=bar2'],
+      ['Set-Cookie', 'baz=baz2'],
+    ],
+  })
+}

test/e2e/app-dir/actions/app/header/actions.js

@@ -9,3 +9,8 @@ export async function getCookie(name) {
 export async function getHeader(name) {
   return headers().get(name)
 }
+
+export async function setCookie(name, value) {
+  cookies().set(name, value)
+  return cookies().get(name)
+}

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

@@ -1,6 +1,6 @@
 import UI from './ui'
 
-import { getCookie, getHeader } from './actions'
+import { getCookie, getHeader, setCookie } from './actions'
 import { validator } from './validator'
 
 export default function Page() {
@@ -9,6 +9,7 @@ export default function Page() {
     <UI
       getCookie={getCookie}
       getHeader={getHeader}
+      setCookie={setCookie}
       getAuthedUppercase={validator(async (str) => {
         'use server'
         return prefix + ' ' + str.toUpperCase()

test/e2e/app-dir/actions/app/header/ui.js

@@ -2,7 +2,12 @@
 
 import { useState } from 'react'
 
-export default function UI({ getCookie, getHeader, getAuthedUppercase }) {
+export default function UI({
+  getCookie,
+  getHeader,
+  setCookie,
+  getAuthedUppercase,
+}) {
   const [result, setResult] = useState('')
 
   return (
@@ -20,6 +25,23 @@ export default function UI({ getCookie, getHeader, getAuthedUppercase }) {
       >
         getCookie
       </button>
+      <button
+        id="setCookie"
+        onClick={async () => {
+          // set cookie on server side
+          const random = Math.random()
+          const res = await setCookie('random-server', random)
+          setResult(
+            random +
+              ':' +
+              res.value +
+              ':' +
+              document.cookie.match(/random-server=([^;]+)/)?.[1]
+          )
+        }}
+      >
+        setCookie
+      </button>
       <button
         id="header"
         onClick={async () => {