From bd794af7b5a385331c9d808bb27242a624899c04 Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Wed, 4 Dec 2024 13:57:01 -0500
Subject: [PATCH] [Segment Cache] Add isHeadPartial

Similar to the previous PR, but for the head, which is delivered
separately from the segments. We can only skip the dynamic request if
this value is `false`.
---
 .../apply-router-state-patch-to-tree.test.tsx |  4 ++-
 .../fill-cache-with-new-subtree-data.test.tsx |  1 +
 ...ll-lazy-items-till-leaf-with-head.test.tsx |  3 +-
 ...te-cache-below-flight-segmentpath.test.tsx |  1 +
 .../should-hard-navigate.test.tsx             |  7 ++--
 .../client/components/segment-cache/cache.ts  |  7 ++++
 .../next/src/client/flight-data-helpers.ts    |  8 +++--
 .../next/src/server/app-render/app-render.tsx | 33 +++++++++++++++++--
 .../app-render/collect-segment-data.tsx       |  4 +++
 packages/next/src/server/app-render/types.ts  |  1 +
 .../walk-tree-with-flight-router-state.tsx    |  5 +++
 11 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/packages/next/src/client/components/router-reducer/apply-router-state-patch-to-tree.test.tsx b/packages/next/src/client/components/router-reducer/apply-router-state-patch-to-tree.test.tsx
index 86c16eac07e8b3..78999d2c80f91a 100644
--- a/packages/next/src/client/components/router-reducer/apply-router-state-patch-to-tree.test.tsx
+++ b/packages/next/src/client/components/router-reducer/apply-router-state-patch-to-tree.test.tsx
@@ -37,6 +37,7 @@ const getFlightData = (): FlightData => {
       <>
         <title>About page!</title>
       </>,
+      false,
     ],
   ]
 }
@@ -52,7 +53,8 @@ describe('applyRouterStatePatchToTree', () => {
 
     // Mirrors the way router-reducer values are passed in.
     const flightDataPath = flightData[0]
-    const [treePatch /*, cacheNodeSeedData, head*/] = flightDataPath.slice(-3)
+    const [treePatch /*, cacheNodeSeedData, head, isHeadPartial*/] =
+      flightDataPath.slice(-4)
     const flightSegmentPath = flightDataPath.slice(0, -4)
 
     const newRouterStateTree = applyRouterStatePatchToTree(
diff --git a/packages/next/src/client/components/router-reducer/fill-cache-with-new-subtree-data.test.tsx b/packages/next/src/client/components/router-reducer/fill-cache-with-new-subtree-data.test.tsx
index 39f1654977e3c8..94285f2d910bcd 100644
--- a/packages/next/src/client/components/router-reducer/fill-cache-with-new-subtree-data.test.tsx
+++ b/packages/next/src/client/components/router-reducer/fill-cache-with-new-subtree-data.test.tsx
@@ -12,6 +12,7 @@ const getFlightData = (): NormalizedFlightData[] => {
       tree: ['about', { children: ['', {}] }],
       seedData: ['about', <h1>SubTreeData Injected!</h1>, {}, null, false],
       head: '<title>Head Injected!</title>',
+      isHeadPartial: false,
       isRootRender: false,
     },
   ]
diff --git a/packages/next/src/client/components/router-reducer/fill-lazy-items-till-leaf-with-head.test.tsx b/packages/next/src/client/components/router-reducer/fill-lazy-items-till-leaf-with-head.test.tsx
index 147b9bc56044d9..d19f5648974a8f 100644
--- a/packages/next/src/client/components/router-reducer/fill-lazy-items-till-leaf-with-head.test.tsx
+++ b/packages/next/src/client/components/router-reducer/fill-lazy-items-till-leaf-with-head.test.tsx
@@ -98,7 +98,8 @@ describe('fillLazyItemsTillLeafWithHead', () => {
 
     // Mirrors the way router-reducer values are passed in.
     const flightDataPath = flightData[0]
-    const [treePatch, cacheNodeSeedData, head] = flightDataPath.slice(-3)
+    const [treePatch, cacheNodeSeedData, head /*, isHeadPartial */] =
+      flightDataPath.slice(-4)
     fillLazyItemsTillLeafWithHead(
       cache,
       existingCache,
diff --git a/packages/next/src/client/components/router-reducer/invalidate-cache-below-flight-segmentpath.test.tsx b/packages/next/src/client/components/router-reducer/invalidate-cache-below-flight-segmentpath.test.tsx
index e7be100bb5c844..eb3efeabf7260d 100644
--- a/packages/next/src/client/components/router-reducer/invalidate-cache-below-flight-segmentpath.test.tsx
+++ b/packages/next/src/client/components/router-reducer/invalidate-cache-below-flight-segmentpath.test.tsx
@@ -13,6 +13,7 @@ const getFlightData = (): NormalizedFlightData[] => {
       tree: ['about', { children: ['', {}] }],
       seedData: ['about', <h1>About Page!</h1>, {}, null, false],
       head: '<title>About page!</title>',
+      isHeadPartial: false,
       isRootRender: false,
     },
   ]
diff --git a/packages/next/src/client/components/router-reducer/should-hard-navigate.test.tsx b/packages/next/src/client/components/router-reducer/should-hard-navigate.test.tsx
index 2c5075dc8d5a36..5783ad74a7024b 100644
--- a/packages/next/src/client/components/router-reducer/should-hard-navigate.test.tsx
+++ b/packages/next/src/client/components/router-reducer/should-hard-navigate.test.tsx
@@ -50,7 +50,7 @@ describe('shouldHardNavigate', () => {
 
     // Mirrors the way router-reducer values are passed in.
     const flightDataPath = flightData[0]
-    const flightSegmentPath = flightDataPath.slice(0, -3)
+    const flightSegmentPath = flightDataPath.slice(0, -4)
 
     const result = shouldHardNavigate(
       ['', ...flightSegmentPath],
@@ -107,7 +107,7 @@ describe('shouldHardNavigate', () => {
 
     // Mirrors the way router-reducer values are passed in.
     const flightDataPath = flightData[0]
-    const flightSegmentPath = flightDataPath.slice(0, -3)
+    const flightSegmentPath = flightDataPath.slice(0, -4)
 
     const result = shouldHardNavigate(
       ['', ...flightSegmentPath],
@@ -153,6 +153,7 @@ describe('shouldHardNavigate', () => {
           ],
           [['id', '123', 'd'], {}, null],
           null,
+          false,
         ],
       ]
     }
@@ -164,7 +165,7 @@ describe('shouldHardNavigate', () => {
 
     // Mirrors the way router-reducer values are passed in.
     const flightDataPath = flightData[0]
-    const flightSegmentPath = flightDataPath.slice(0, -3)
+    const flightSegmentPath = flightDataPath.slice(0, -4)
 
     const result = shouldHardNavigate(
       ['', ...flightSegmentPath],
diff --git a/packages/next/src/client/components/segment-cache/cache.ts b/packages/next/src/client/components/segment-cache/cache.ts
index 8be3535929fb1a..cfdd93bfd9b7ad 100644
--- a/packages/next/src/client/components/segment-cache/cache.ts
+++ b/packages/next/src/client/components/segment-cache/cache.ts
@@ -80,6 +80,7 @@ type PendingRouteCacheEntry = RouteCacheEntryShared & {
   canonicalUrl: null
   tree: null
   head: null
+  isHeadPartial: true
 }
 
 type RejectedRouteCacheEntry = RouteCacheEntryShared & {
@@ -88,6 +89,7 @@ type RejectedRouteCacheEntry = RouteCacheEntryShared & {
   canonicalUrl: null
   tree: null
   head: null
+  isHeadPartial: true
 }
 
 export type FulfilledRouteCacheEntry = RouteCacheEntryShared & {
@@ -96,6 +98,7 @@ export type FulfilledRouteCacheEntry = RouteCacheEntryShared & {
   canonicalUrl: string
   tree: TreePrefetch
   head: React.ReactNode | null
+  isHeadPartial: boolean
 }
 
 export type RouteCacheEntry =
@@ -281,6 +284,7 @@ export function requestRouteCacheEntryFromCache(
     blockedTasks: null,
     tree: null,
     head: null,
+    isHeadPartial: true,
     // If the request takes longer than a minute, a subsequent request should
     // retry instead of waiting for this one.
     //
@@ -420,6 +424,7 @@ function fulfillRouteCacheEntry(
   entry: PendingRouteCacheEntry,
   tree: TreePrefetch,
   head: React.ReactNode,
+  isHeadPartial: boolean,
   staleAt: number,
   couldBeIntercepted: boolean,
   canonicalUrl: string
@@ -428,6 +433,7 @@ function fulfillRouteCacheEntry(
   fulfilledEntry.status = EntryStatus.Fulfilled
   fulfilledEntry.tree = tree
   fulfilledEntry.head = head
+  fulfilledEntry.isHeadPartial = isHeadPartial
   fulfilledEntry.staleAt = staleAt
   fulfilledEntry.couldBeIntercepted = couldBeIntercepted
   fulfilledEntry.canonicalUrl = canonicalUrl
@@ -532,6 +538,7 @@ async function fetchRouteOnCacheMiss(
       entry,
       serverData.tree,
       serverData.head,
+      serverData.isHeadPartial,
       Date.now() + serverData.staleTime,
       couldBeIntercepted,
       canonicalUrl
diff --git a/packages/next/src/client/flight-data-helpers.ts b/packages/next/src/client/flight-data-helpers.ts
index 01fa57faff42b5..1f9930220e6ed6 100644
--- a/packages/next/src/client/flight-data-helpers.ts
+++ b/packages/next/src/client/flight-data-helpers.ts
@@ -20,6 +20,7 @@ export type NormalizedFlightData = {
   tree: FlightRouterState
   seedData: CacheNodeSeedData | null
   head: React.ReactNode | null
+  isHeadPartial: boolean
   isRootRender: boolean
 }
 
@@ -31,9 +32,9 @@ export function getFlightDataPartsFromPath(
   flightDataPath: FlightDataPath
 ): NormalizedFlightData {
   // tree, seedData, and head are *always* the last three items in the `FlightDataPath`.
-  const [tree, seedData, head] = flightDataPath.slice(-3)
+  const [tree, seedData, head, isHeadPartial] = flightDataPath.slice(-4)
   // The `FlightSegmentPath` is everything except the last three items. For a root render, it won't be present.
-  const segmentPath = flightDataPath.slice(0, -3)
+  const segmentPath = flightDataPath.slice(0, -4)
 
   return {
     // TODO: Unify these two segment path helpers. We are inconsistently pushing an empty segment ("")
@@ -47,7 +48,8 @@ export function getFlightDataPartsFromPath(
     tree,
     seedData,
     head,
-    isRootRender: flightDataPath.length === 3,
+    isHeadPartial,
+    isRootRender: flightDataPath.length === 4,
   }
 }
 
diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx
index e829eb6efabcf9..16243888ee2262 100644
--- a/packages/next/src/server/app-render/app-render.tsx
+++ b/packages/next/src/server/app-render/app-render.tsx
@@ -10,6 +10,7 @@ import type {
   RSCPayload,
   FlightData,
   InitialRSCPayload,
+  FlightDataPath,
 } from './types'
 import {
   workAsyncStorage,
@@ -780,6 +781,16 @@ async function getRSCPayload(
 
   const globalErrorStyles = await getGlobalErrorStyles(tree, ctx)
 
+  // Assume the head we're rendering contains only partial data if PPR is
+  // enabled and this is a statically generated response. This is used by the
+  // client Segment Cache after a prefetch to determine if it can skip the
+  // second request to fill in the dynamic data.
+  //
+  // See similar comment in create-component-tree.tsx for more context.
+  const isPossiblyPartialHead =
+    workStore.isStaticGeneration &&
+    ctx.renderOpts.experimental.isRoutePPREnabled === true
+
   return {
     // See the comment above the `Preloads` component (below) for why this is part of the payload
     P: <Preloads preloadCallbacks={preloadCallbacks} />,
@@ -787,7 +798,14 @@ async function getRSCPayload(
     p: ctx.assetPrefix,
     c: prepareInitialCanonicalUrl(url),
     i: !!couldBeIntercepted,
-    f: [[initialTree, seedData, initialHead]],
+    f: [
+      [
+        initialTree,
+        seedData,
+        initialHead,
+        isPossiblyPartialHead,
+      ] as FlightDataPath,
+    ],
     m: missingSlots,
     G: [GlobalError, globalErrorStyles],
     s: typeof ctx.renderOpts.postponed === 'string',
@@ -877,13 +895,24 @@ async function getErrorRSCPayload(
 
   const globalErrorStyles = await getGlobalErrorStyles(tree, ctx)
 
+  const isPossiblyPartialHead =
+    workStore.isStaticGeneration &&
+    ctx.renderOpts.experimental.isRoutePPREnabled === true
+
   return {
     b: ctx.renderOpts.buildId,
     p: ctx.assetPrefix,
     c: prepareInitialCanonicalUrl(url),
     m: undefined,
     i: false,
-    f: [[initialTree, initialSeedData, initialHead]],
+    f: [
+      [
+        initialTree,
+        initialSeedData,
+        initialHead,
+        isPossiblyPartialHead,
+      ] as FlightDataPath,
+    ],
     G: [GlobalError, globalErrorStyles],
     s: typeof ctx.renderOpts.postponed === 'string',
     S: workStore.isStaticGeneration,
diff --git a/packages/next/src/server/app-render/collect-segment-data.tsx b/packages/next/src/server/app-render/collect-segment-data.tsx
index a942588f5b1b9c..9c5591dd631950 100644
--- a/packages/next/src/server/app-render/collect-segment-data.tsx
+++ b/packages/next/src/server/app-render/collect-segment-data.tsx
@@ -25,6 +25,7 @@ export type RootTreePrefetch = {
   buildId: string
   tree: TreePrefetch
   head: React.ReactNode | null
+  isHeadPartial: boolean
   staleTime: number
 }
 
@@ -194,6 +195,8 @@ async function PrefetchTreeData({
     segmentTasks
   )
 
+  const isHeadPartial = await isPartialRSCData(head, clientModules)
+
   // Notify the abort controller that we're done processing the route tree.
   // Anything async that happens after this point must be due to hanging
   // promises in the original stream.
@@ -204,6 +207,7 @@ async function PrefetchTreeData({
     buildId,
     tree,
     head,
+    isHeadPartial,
     staleTime,
   }
   return treePrefetch
diff --git a/packages/next/src/server/app-render/types.ts b/packages/next/src/server/app-render/types.ts
index 750b52568b03a0..8a81b581cdbe87 100644
--- a/packages/next/src/server/app-render/types.ts
+++ b/packages/next/src/server/app-render/types.ts
@@ -105,6 +105,7 @@ export type FlightDataSegment = [
   /* treePatch */ FlightRouterState,
   /* cacheNodeSeedData */ CacheNodeSeedData | null, // Can be null during prefetch if there's no loading component
   /* head */ React.ReactNode | null,
+  /* isHeadPartial */ boolean,
 ]
 
 export type FlightDataPath =
diff --git a/packages/next/src/server/app-render/walk-tree-with-flight-router-state.tsx b/packages/next/src/server/app-render/walk-tree-with-flight-router-state.tsx
index 100ef3ceeb5daf..66f361aee96337 100644
--- a/packages/next/src/server/app-render/walk-tree-with-flight-router-state.tsx
+++ b/packages/next/src/server/app-render/walk-tree-with-flight-router-state.tsx
@@ -140,7 +140,11 @@ export async function walkTreeWithFlightRouterState({
           overriddenSegment,
           routerState,
           null,
+          // TODO: It's possible that all the segment data was prefetched during
+          // a navigation, but the head was not. Should we send it down
+          // here anyway?
           null,
+          false,
         ] satisfies FlightDataSegment,
       ]
     } else {
@@ -170,6 +174,7 @@ export async function walkTreeWithFlightRouterState({
           routerState,
           seedData,
           rscPayloadHead,
+          false,
         ] satisfies FlightDataSegment,
       ]
     }