fix: allow some recursion for middleware subrequests (#60615)

This alters the behavior of the subrequest check to allow for 5
recursive calls to match Vercel production, ref [Slack
thread](https://pinglabsworkspace.slack.com/archives/C052S77L05C/p1694729495655489).

> [!NOTE]
> Currently limited by fetches having to forward the subrequest header
for each request which isn't ideal. Need some assistant on how to access
the request in the module context fetch override.
> No forwarding: 
![CleanShot 2024-02-05 at 22 52
10@2x](https://github.com/vercel/next.js/assets/51714798/8ae79f00-f987-4919-946c-d8363d540cef)
> With forwarding: 
![CleanShot 2024-02-05 at 22 51
31@2x](https://github.com/vercel/next.js/assets/51714798/32bd4072-9373-4cb0-ab05-f862a818e0d7)

---------

Co-authored-by: JJ Kasper <jj@jjsweb.site>

packages/next/src/server/web/sandbox/context.ts

@@ -232,6 +232,10 @@ const NativeModuleMap = (() => {
   return new Map(Object.entries(mods))
 })()
 
+export const requestStore = new AsyncLocalStorage<{
+  headers: Headers
+}>()
+
 /**
  * Create a module cache specific for the provided parameters. It includes
  * a runtime context, require cache and paths cache.
@@ -335,6 +339,19 @@ Learn More: https://nextjs.org/docs/messages/edge-dynamic-code-evaluation`),
         }
 
         init.headers = new Headers(init.headers ?? {})
+
+        // Forward subrequest header from incoming request to outgoing request
+        const store = requestStore.getStore()
+        if (
+          store?.headers.has('x-middleware-subrequest') &&
+          !init.headers.has('x-middleware-subrequest')
+        ) {
+          init.headers.set(
+            'x-middleware-subrequest',
+            store.headers.get('x-middleware-subrequest') ?? ''
+          )
+        }
+
         const prevs =
           init.headers.get(`x-middleware-subrequest`)?.split(':') || []
         const value = prevs.concat(options.moduleName).join(':')

packages/next/src/server/web/sandbox/sandbox.ts

@@ -1,7 +1,7 @@
 import type { NodejsRequestData, FetchEventResult, RequestData } from '../types'
 import type { EdgeFunctionDefinition } from '../../../build/webpack/plugins/middleware-plugin'
 import type { EdgeRuntime } from 'next/dist/compiled/edge-runtime'
-import { getModuleContext } from './context'
+import { getModuleContext, requestStore } from './context'
 import { requestToBodyStream } from '../../body-streams'
 import { NEXT_RSC_UNION_QUERY } from '../../../client/components/app-router-headers'
 
@@ -82,7 +82,14 @@ export const run = withTaggedErrors(async function runWithTaggedErrors(params) {
   const runtime = await getRuntimeContext(params)
   const subreq = params.request.headers[`x-middleware-subrequest`]
   const subrequests = typeof subreq === 'string' ? subreq.split(':') : []
-  if (subrequests.includes(params.name)) {
+
+  const MAX_RECURSION_DEPTH = 5
+  const depth = subrequests.reduce(
+    (acc, curr) => (curr === params.name ? acc + 1 : acc),
+    0
+  )
+
+  if (depth >= MAX_RECURSION_DEPTH) {
     return {
       waitUntil: Promise.resolve(),
       response: new runtime.context.Response(null, {
@@ -108,17 +115,26 @@ export const run = withTaggedErrors(async function runWithTaggedErrors(params) {
 
   params.request.url = urlInstance.toString()
 
+  const headers = new Headers()
+  for (const [key, value] of Object.entries(params.request.headers)) {
+    headers.set(key, value?.toString() ?? '')
+  }
+
   try {
-    const result = await edgeFunction({
-      request: {
-        ...params.request,
-        body:
-          cloned && requestToBodyStream(runtime.context, KUint8Array, cloned),
-      },
+    let result: FetchEventResult | undefined = undefined
+    await requestStore.run({ headers }, async () => {
+      result = await edgeFunction({
+        request: {
+          ...params.request,
+          body:
+            cloned && requestToBodyStream(runtime.context, KUint8Array, cloned),
+        },
+      })
+      for (const headerName of FORBIDDEN_HEADERS) {
+        result.response.headers.delete(headerName)
+      }
     })
-    for (const headerName of FORBIDDEN_HEADERS) {
-      result.response.headers.delete(headerName)
-    }
+    if (!result) throw new Error('Edge function did not return a response')
     return result
   } finally {
     await params.request.body?.finalize()

test/e2e/app-dir/app-routes-subrequests/app-routes-subrequests.test.ts

@@ -0,0 +1,18 @@
+import { createNextDescribe } from 'e2e-utils'
+
+const bathPath = process.env.BASE_PATH ?? ''
+
+createNextDescribe(
+  'app-routes-subrequests',
+  {
+    files: __dirname,
+    skipDeployment: true,
+  },
+  ({ next }) => {
+    it('shortcuts after 5 subrequests', async () => {
+      expect(JSON.parse(await next.render(bathPath + '/'))).toEqual({
+        count: 5,
+      })
+    })
+  }
+)

test/e2e/app-dir/app-routes-subrequests/app/route.ts

@@ -0,0 +1,11 @@
+import { NextRequest, NextResponse } from 'next/server'
+
+export const runtime = 'edge'
+
+let count = 0
+
+export const GET = async (req: NextRequest) => {
+  await fetch(req.nextUrl)
+  count++
+  return NextResponse.json({ count })
+}

test/e2e/app-dir/app-routes-subrequests/next.config.js

@@ -0,0 +1,10 @@
+/**
+ * @type {import('next').NextConfig}
+ */
+const config = {
+  typescript: {
+    ignoreBuildErrors: true,
+  },
+}
+
+module.exports = config