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,