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(
-      this.requestAsyncStorage,
-      requestContext,
+    const response: unknown = await this.actionAsyncStorage.run(
+      {
+        isAppRoute: true,
+      },
       () =>
-        StaticGenerationAsyncStorageWrapper.wrap(
-          this.staticGenerationAsyncStorage,
-          staticGenerationContext,
-          (staticGenerationStore) => {
-            // 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.
-            switch (this.dynamic) {
-              case 'force-dynamic':
-                // The dynamic property is set to force-dynamic, so we should
-                // force the page to be dynamic.
-                staticGenerationStore.forceDynamic = true
-                this.staticGenerationBailout(`force-dynamic`, {
-                  dynamic: this.dynamic,
-                })
-                break
-              case 'force-static':
-                // The dynamic property is set to force-static, so we should
-                // force the page to be static.
-                staticGenerationStore.forceStatic = true
-                break
-              case 'error':
-                // The dynamic property is set to error, so we should throw an
-                // error if the page is being statically generated.
-                staticGenerationStore.dynamicShouldError = true
-                break
-              default:
-                break
-            }
-
-            // If the static generation store does not have a revalidate value
-            // set, then we should set it the revalidate value from the userland
-            // module or default to false.
-            staticGenerationStore.revalidate ??=
-              this.userland.revalidate ?? false
-
-            // Wrap the request so we can add additional functionality to cases
-            // that might change it's output or affect the rendering.
-            const wrappedRequest = proxyRequest(
-              request,
-              { dynamic: this.dynamic },
-              {
-                headerHooks: this.headerHooks,
-                serverHooks: this.serverHooks,
-                staticGenerationBailout: this.staticGenerationBailout,
-              }
-            )
+        RequestAsyncStorageWrapper.wrap(
+          this.requestAsyncStorage,
+          requestContext,
+          () =>
+            StaticGenerationAsyncStorageWrapper.wrap(
+              this.staticGenerationAsyncStorage,
+              staticGenerationContext,
+              (staticGenerationStore) => {
+                // 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.
+                switch (this.dynamic) {
+                  case 'force-dynamic':
+                    // The dynamic property is set to force-dynamic, so we should
+                    // force the page to be dynamic.
+                    staticGenerationStore.forceDynamic = true
+                    this.staticGenerationBailout(`force-dynamic`, {
+                      dynamic: this.dynamic,
+                    })
+                    break
+                  case 'force-static':
+                    // The dynamic property is set to force-static, so we should
+                    // force the page to be static.
+                    staticGenerationStore.forceStatic = true
+                    break
+                  case 'error':
+                    // The dynamic property is set to error, so we should throw an
+                    // error if the page is being statically generated.
+                    staticGenerationStore.dynamicShouldError = true
+                    break
+                  default:
+                    break
+                }
+
+                // If the static generation store does not have a revalidate value
+                // set, then we should set it the revalidate value from the userland
+                // module or default to false.
+                staticGenerationStore.revalidate ??=
+                  this.userland.revalidate ?? false
+
+                // Wrap the request so we can add additional functionality to cases
+                // that might change it's output or affect the rendering.
+                const wrappedRequest = proxyRequest(
+                  request,
+                  { dynamic: this.dynamic },
+                  {
+                    headerHooks: this.headerHooks,
+                    serverHooks: this.serverHooks,
+                    staticGenerationBailout: this.staticGenerationBailout,
+                  }
+                )
 
-            // TODO: propagate this pathname from route matcher
-            const route = getPathnameFromAbsolutePath(this.resolvedPagePath)
-            getTracer().getRootSpanAttributes()?.set('next.route', route)
-            return getTracer().trace(
-              AppRouteRouteHandlersSpan.runHandler,
-              {
-                spanName: `executing api route (app) ${route}`,
-                attributes: {
-                  'next.route': route,
-                },
-              },
-              async () => {
-                // Patch the global fetch.
-                patchFetch({
-                  serverHooks: this.serverHooks,
-                  staticGenerationAsyncStorage:
-                    this.staticGenerationAsyncStorage,
-                })
-                const res = await handler(wrappedRequest, {
-                  params: context.params,
-                })
-
-                await Promise.all(
-                  staticGenerationStore.pendingRevalidates || []
+                // TODO: propagate this pathname from route matcher
+                const route = getPathnameFromAbsolutePath(this.resolvedPagePath)
+                getTracer().getRootSpanAttributes()?.set('next.route', route)
+                return getTracer().trace(
+                  AppRouteRouteHandlersSpan.runHandler,
+                  {
+                    spanName: `executing api route (app) ${route}`,
+                    attributes: {
+                      'next.route': route,
+                    },
+                  },
+                  async () => {
+                    // Patch the global fetch.
+                    patchFetch({
+                      serverHooks: this.serverHooks,
+                      staticGenerationAsyncStorage:
+                        this.staticGenerationAsyncStorage,
+                    })
+                    const res = await handler(wrappedRequest, {
+                      params: context.params,
+                    })
+
+                    await Promise.all(
+                      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