Handle head.js on client-side navigation

packages/next/client/components/app-router.tsx

@@ -109,6 +109,37 @@ type AppRouterProps = {
   assetPrefix: string
 }
 
+function findHeadInCache(
+  cache: CacheNode,
+  parallelRoutes: FlightRouterState[1]
+): React.ReactNode {
+  const isLastItem = Object.keys(parallelRoutes).length === 0
+  if (isLastItem) {
+    return cache.head
+  }
+  for (const key in parallelRoutes) {
+    const [segment, childParallelRoutes] = parallelRoutes[key]
+    const childSegmentMap = cache.parallelRoutes.get(key)
+    if (!childSegmentMap) {
+      continue
+    }
+
+    const cacheKey = Array.isArray(segment) ? segment[1] : segment
+
+    const cacheNode = childSegmentMap.get(cacheKey)
+    if (!cacheNode) {
+      continue
+    }
+
+    const item = findHeadInCache(cacheNode, childParallelRoutes)
+    if (item) {
+      return item
+    }
+  }
+
+  return undefined
+}
+
 /**
  * The global router that wraps the application components.
  */
@@ -147,6 +178,10 @@ function Router({
     sync,
   ] = useReducerWithReduxDevtools(reducer, initialState)
 
+  const head = useMemo(() => {
+    return findHeadInCache(cache, tree[1])
+  }, [cache, tree])
+
   useEffect(() => {
     // Ensure initialParallelRoutes is cleaned up from memory once it's used.
     initialParallelRoutes = null!
@@ -357,6 +392,13 @@ function Router({
     }
   }, [onPopState])
 
+  const content = (
+    <>
+      {head || initialHead}
+      {cache.subTreeData}
+    </>
+  )
+
   return (
     <PathnameContext.Provider value={pathname}>
       <SearchParamsContext.Provider value={searchParams}>
@@ -378,15 +420,9 @@ function Router({
               }}
             >
               {HotReloader ? (
-                <HotReloader assetPrefix={assetPrefix}>
-                  {initialHead}
-                  {cache.subTreeData}
-                </HotReloader>
+                <HotReloader assetPrefix={assetPrefix}>{content}</HotReloader>
               ) : (
-                <>
-                  {initialHead}
-                  {cache.subTreeData}
-                </>
+                content
               )}
             </LayoutRouterContext.Provider>
           </AppRouterContext.Provider>

packages/next/client/components/layout-router.tsx

@@ -170,7 +170,7 @@ export function InnerLayoutRouter({
   }
 
   // When childNode is not available during rendering client-side we need to fetch it from the server.
-  if (!childNode) {
+  if (!childNode || childNode.status === CacheStates.LAZY_INITIALIZED) {
     /**
      * Router state with refetch marker added
      */
@@ -184,7 +184,14 @@ export function InnerLayoutRouter({
       status: CacheStates.DATA_FETCH,
       data: fetchServerResponse(new URL(url, location.origin), refetchTree),
       subTreeData: null,
-      parallelRoutes: new Map(),
+      head:
+        childNode && childNode.status === CacheStates.LAZY_INITIALIZED
+          ? childNode.head
+          : undefined,
+      parallelRoutes:
+        childNode && childNode.status === CacheStates.LAZY_INITIALIZED
+          ? childNode.parallelRoutes
+          : new Map(),
     })
     // In the above case childNode was set on childNodes, so we have to get it from the cacheNodes again.
     childNode = childNodes.get(path)

packages/next/client/components/reducer.ts

@@ -164,7 +164,6 @@ function fillCacheWithNewSubTreeData(
   const existingChildCacheNode = existingChildSegmentMap.get(segmentForCache)
   let childCacheNode = childSegmentMap.get(segmentForCache)
 
-  // In case of last segment start the fetch at this level and don't copy further down.
   if (isLastEntry) {
     if (
       !childCacheNode ||
@@ -193,7 +192,7 @@ function fillCacheWithNewSubTreeData(
         childCacheNode,
         existingChildCacheNode,
         flightDataPath[2],
-        /* flightDataPath[4] */ undefined
+        flightDataPath[4]
       )
 
       childSegmentMap.set(segmentForCache, childCacheNode)
@@ -309,12 +308,20 @@ function fillCacheWithPrefetchedSubTreeData(
 
   if (isLastEntry) {
     if (!existingChildCacheNode) {
-      existingChildSegmentMap.set(segmentForCache, {
+      const childCacheNode: CacheNode = {
         status: CacheStates.READY,
         data: null,
         subTreeData: flightDataPath[3],
         parallelRoutes: new Map(),
-      })
+      }
+
+      fillLazyItemsTillLeafWithHead(
+        childCacheNode,
+        existingChildCacheNode,
+        flightDataPath[2],
+        flightDataPath[4]
+      )
+      existingChildSegmentMap.set(segmentForCache, childCacheNode)
     }
 
     return
@@ -997,7 +1004,6 @@ function clientReducer(
       const flightDataPath = flightData[0]
 
       // The one before last item is the router state tree patch
-      // eslint-disable-next-line @typescript-eslint/no-unused-vars
       const [treePatch, subTreeData, head] = flightDataPath.slice(-3)
 
       // Path without the last segment, router state, and the subTreeData
@@ -1027,6 +1033,7 @@ function clientReducer(
 
       if (flightDataPath.length === 3) {
         cache.subTreeData = subTreeData
+        fillLazyItemsTillLeafWithHead(cache, state.cache, treePatch, head)
       } else {
         // Copy subTreeData for the root node of the cache.
         cache.subTreeData = state.cache.subTreeData
@@ -1124,7 +1131,6 @@ function clientReducer(
 
       // Slices off the last segment (which is at -4) as it doesn't exist in the tree yet
       const flightSegmentPath = flightDataPath.slice(0, -4)
-      // eslint-disable-next-line @typescript-eslint/no-unused-vars
       const [treePatch, subTreeData, head] = flightDataPath.slice(-3)
 
       const newTree = applyRouterStatePatchToTree(
@@ -1152,6 +1158,7 @@ function clientReducer(
       // Root refresh
       if (flightDataPath.length === 3) {
         cache.subTreeData = subTreeData
+        fillLazyItemsTillLeafWithHead(cache, state.cache, treePatch, head)
       } else {
         // Copy subTreeData for the root node of the cache.
         cache.subTreeData = state.cache.subTreeData
@@ -1276,7 +1283,7 @@ function clientReducer(
       }
 
       // Given the path can only have two items the items are only the router state and subTreeData for the root.
-      const [treePatch, subTreeData] = flightDataPath
+      const [treePatch, subTreeData, head] = flightDataPath
       const newTree = applyRouterStatePatchToTree(
         // TODO-APP: remove ''
         [''],
@@ -1302,6 +1309,7 @@ function clientReducer(
 
       // Set subTreeData for the root node of the cache.
       cache.subTreeData = subTreeData
+      fillLazyItemsTillLeafWithHead(cache, state.cache, treePatch, head)
 
       return {
         // Set href, this doesn't reuse the state.canonicalUrl as because of concurrent rendering the href might change between dispatching and applying.
@@ -1334,8 +1342,7 @@ function clientReducer(
       const flightDataPath = flightData[0]
 
       // The one before last item is the router state tree patch
-      // eslint-disable-next-line @typescript-eslint/no-unused-vars
-      const [treePatch, subTreeData, head] = flightDataPath.slice(-3)
+      const [treePatch, subTreeData] = flightDataPath.slice(-3)
 
       // TODO-APP: Verify if `null` can't be returned from user code.
       // If subTreeData is null the prefetch did not provide a component tree.

packages/next/server/app-render.tsx

@@ -516,10 +516,11 @@ export type FlightDataPath =
   // Looks somewhat like this
   | [
       // Holds full path to the segment.
-      ...FlightSegmentPath,
+      ...FlightSegmentPath[],
       /* segment of the rendered slice: */ Segment,
       /* treePatch */ FlightRouterState,
-      /* subTreeData: */ React.ReactNode | null // Can be null during prefetch if there's no loading component
+      /* subTreeData: */ React.ReactNode | null, // Can be null during prefetch if there's no loading component
+      /* head */ React.ReactNode | null
     ]
 
 /**
@@ -1381,13 +1382,15 @@ export async function renderToHTMLOrFlight(
         isFirst,
         flightRouterState,
         parentRendered,
+        rscPayloadHead,
       }: {
         createSegmentPath: CreateSegmentPath
         loaderTreeToFilter: LoaderTree
         parentParams: { [key: string]: string | string[] }
         isFirst: boolean
         flightRouterState?: FlightRouterState
         parentRendered?: boolean
+        rscPayloadHead: React.ReactNode
       }): Promise<FlightDataPath> => {
         const [segment, parallelRoutes] = loaderTreeToFilter
         const parallelRoutesKeys = Object.keys(parallelRoutes)
@@ -1444,7 +1447,7 @@ export async function renderToHTMLOrFlight(
                   ).Component
                 ),
             isPrefetch && !Boolean(loaderTreeToFilter[2].loading) ? null : (
-              <>{null}</> // TODO: change this to head tags.
+              <>{rscPayloadHead}</>
             ),
           ]
         }
@@ -1467,6 +1470,7 @@ export async function renderToHTMLOrFlight(
               flightRouterState && flightRouterState[1][parallelRouteKey],
             parentRendered: parentRendered || renderComponentsOnThisLevel,
             isFirst: false,
+            rscPayloadHead,
           })
 
           if (typeof path[path.length - 1] !== 'string') {
@@ -1477,6 +1481,7 @@ export async function renderToHTMLOrFlight(
         return [actualSegment]
       }
 
+      const rscPayloadHead = await resolveHead(loaderTree, {})
       // Flight data that is going to be passed to the browser.
       // Currently a single item array but in the future multiple patches might be combined in a single request.
       const flightData: FlightData = [
@@ -1488,6 +1493,7 @@ export async function renderToHTMLOrFlight(
             parentParams: {},
             flightRouterState: providedFlightRouterState,
             isFirst: true,
+            rscPayloadHead,
           })
         ).slice(1),
       ]

test/e2e/app-dir/head.test.ts

@@ -3,6 +3,7 @@ import cheerio from 'cheerio'
 import { createNext, FileRef } from 'e2e-utils'
 import { NextInstance } from 'test/lib/next-modes/base'
 import { renderViaHTTP } from 'next-test-utils'
+import webdriver from 'next-webdriver'
 
 describe('app dir head', () => {
   if ((global as any).isNextDeploy) {
@@ -91,6 +92,27 @@ describe('app dir head', () => {
         headTags.find((el) => el.attribs.src === '/another.js')
       ).toBeTruthy()
     })
+
+    it('should apply head when navigating client-side', async () => {
+      const browser = await webdriver(next.url, '/')
+
+      const getTitle = () => browser.elementByCss('title').text()
+
+      expect(await getTitle()).toBe('hello from index')
+      await browser
+        .elementByCss('#to-blog')
+        .click()
+        .waitForElementByCss('#layout', 2000)
+
+      expect(await getTitle()).toBe('hello from blog layout')
+      await browser.back().waitForElementByCss('#to-blog', 2000)
+      expect(await getTitle()).toBe('hello from index')
+      await browser
+        .elementByCss('#to-blog-slug')
+        .click()
+        .waitForElementByCss('#layout', 2000)
+      expect(await getTitle()).toBe('hello from dynamic blog page post-1')
+    })
   }
 
   runTests()