diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx
index cf6d469eeb90a..a35963b3c3351 100644
--- a/packages/next/src/server/app-render/app-render.tsx
+++ b/packages/next/src/server/app-render/app-render.tsx
@@ -1194,6 +1194,7 @@ async function renderToHTMLOrFlightImpl(
const notFoundLoaderTree = createNotFoundLoaderTree(loaderTree)
res.statusCode = 404
const stream = await renderToStreamWithTracing(
+ requestStore,
req,
res,
ctx,
@@ -1219,6 +1220,7 @@ async function renderToHTMLOrFlightImpl(
}
const stream = await renderToStreamWithTracing(
+ requestStore,
req,
res,
ctx,
@@ -1345,6 +1347,7 @@ export const renderToHTMLOrFlight: AppPageRender = (
}
async function renderToStream(
+ requestStore: RequestStore,
req: BaseNextRequest,
res: BaseNextResponse,
ctx: AppRenderContext,
@@ -1435,9 +1438,17 @@ async function renderToStream(
try {
// This is a dynamic render. We don't do dynamic tracking because we're not prerendering
- const RSCPayload = await getRSCPayload(tree, ctx, res.statusCode === 404)
+ const RSCPayload = await workUnitAsyncStorage.run(
+ requestStore,
+ getRSCPayload,
+ tree,
+ ctx,
+ res.statusCode === 404
+ )
reactServerResult = new ReactServerResult(
- ComponentMod.renderToReadableStream(
+ workUnitAsyncStorage.run(
+ requestStore,
+ ComponentMod.renderToReadableStream,
RSCPayload,
clientReferenceManifest.clientModules,
{
@@ -1475,7 +1486,9 @@ async function renderToStream(
const resume = require('react-dom/server.edge')
.resume as (typeof import('react-dom/server.edge'))['resume']
- const htmlStream = await resume(
+ const htmlStream = await workUnitAsyncStorage.run(
+ requestStore,
+ resume,
- ),
- streamOptions: {
- nonce: ctx.nonce,
- // Include hydration scripts in the HTML
- bootstrapScripts: [errorBootstrapScript],
- formState,
- },
- })
+ const fizzStream = await workUnitAsyncStorage.run(
+ requestStore,
+ renderToInitialFizzStream,
+ {
+ ReactDOMServer: require('react-dom/server.edge'),
+ element: (
+
+ ),
+ streamOptions: {
+ nonce: ctx.nonce,
+ // Include hydration scripts in the HTML
+ bootstrapScripts: [errorBootstrapScript],
+ formState,
+ },
+ }
+ )
/**
* Rules of Static & Dynamic HTML:
diff --git a/packages/next/src/server/route-modules/app-route/module.ts b/packages/next/src/server/route-modules/app-route/module.ts
index 61e6afa2394c5..fb69109dda946 100644
--- a/packages/next/src/server/route-modules/app-route/module.ts
+++ b/packages/next/src/server/route-modules/app-route/module.ts
@@ -41,7 +41,7 @@ import {
} from '../../app-render/work-async-storage.external'
import {
workUnitAsyncStorage,
- type WorkUnitStore,
+ type RequestStore,
type PrerenderStore,
} from '../../app-render/work-unit-async-storage.external'
import {
@@ -276,7 +276,10 @@ export class AppRouteRouteModule extends RouteModule<
handler: AppRouteHandlerFn,
actionStore: ActionStore,
workStore: WorkStore,
- workUnitStore: WorkUnitStore,
+ // @TODO refactor to not take this argument but instead construct the RequestStore
+ // inside this function. Right now we get passed a RequestStore even when
+ // we're going to do a prerender. We should probably just split do up into prexecute and execute
+ requestStore: RequestStore,
implicitTags: string[],
request: NextRequest,
context: AppRouteRouteHandlerContext
@@ -503,7 +506,12 @@ export class AppRouteRouteModule extends RouteModule<
)
}
} else {
- res = await handler(request, handlerContext)
+ res = await workUnitAsyncStorage.run(
+ requestStore,
+ handler,
+ request,
+ handlerContext
+ )
}
} catch (err) {
if (isRedirectError(err)) {
@@ -518,8 +526,10 @@ export class AppRouteRouteModule extends RouteModule<
// Let's append any cookies that were added by the
// cookie API.
- if (workUnitStore.type === 'request') {
- appendMutableCookies(headers, workUnitStore.mutableCookies)
+ // TODO leaving the gate here b/c it indicates that we we might not actually want to do this
+ // on every `do` call. During prerender there should be no mutableCookies because
+ if (requestStore.type === 'request') {
+ appendMutableCookies(headers, requestStore.mutableCookies)
}
// Return the redirect response.
@@ -567,8 +577,8 @@ export class AppRouteRouteModule extends RouteModule<
// here.
const headers = new Headers(res.headers)
if (
- workUnitStore.type === 'request' &&
- appendMutableCookies(headers, workUnitStore.mutableCookies)
+ requestStore.type === 'request' &&
+ appendMutableCookies(headers, requestStore.mutableCookies)
) {
return new Response(res.body, {
status: res.status,