Honor rewrites in fallback hydration

Author: feedthejim

packages/next/src/client/app-index.tsx

@@ -337,8 +337,9 @@ if (instantTestStaticFetch) {
 
     return createInitialRSCPayloadFromFallbackPrerender(
       await processedResponse,
-      fallbackInitialRSCPayload,
-      renderedUrl
+      // Let rewritten-path/query headers on the fetched fallback response win
+      // over the current document URL when the host remaps the request.
+      fallbackInitialRSCPayload
     )
   })()
 } else if (

packages/next/src/client/flight-data-helpers.test.ts

@@ -405,4 +405,110 @@ describe('createInitialRSCPayloadFromFallbackPrerender', () => {
     expect(payload.f[0][2]).toBe('head-node')
     expect(payload.f[0][3]).toBe(true)
   })
+
+  it('prefers rewritten headers over the rendered URL override', () => {
+    const fallbackInitialRSCPayload: InitialRSCPayload = {
+      b: 'build-id',
+      c: ['', 'docs', '%%drp:slug:abc123%%'],
+      q: '',
+      i: false,
+      f: [
+        [
+          [
+            '',
+            {
+              children: [
+                'docs',
+                {
+                  children: [['slug', '%%drp:slug:abc123%%', 'd', null], {}],
+                },
+              ],
+            },
+          ],
+          null,
+          'head-node',
+          false,
+        ],
+      ],
+      m: new Set(),
+      G: [(() => null) as any, undefined],
+      S: false,
+      h: null,
+    }
+
+    const response = new Response(null, {
+      headers: {
+        'x-nextjs-rewritten-path': '/docs/rewritten',
+        'x-nextjs-rewritten-query': 'from=fetch',
+      },
+    })
+
+    const payload = createInitialRSCPayloadFromFallbackPrerender(
+      response,
+      fallbackInitialRSCPayload,
+      new URL('https://example.com/docs/original?from=document')
+    )
+
+    const patchedTree = payload.f[0][0]
+    expect(patchedTree[1].children[0]).toBe('docs')
+    expect(patchedTree[1].children[1].children[0]).toEqual([
+      'slug',
+      'rewritten',
+      'd',
+      null,
+    ])
+    expect(payload.q).toBe('?from=fetch')
+    expect(payload.c.join('/')).toBe('/docs/original?from=document')
+  })
+
+  it('falls back to the rendered URL override when no rewritten headers are present', () => {
+    const fallbackInitialRSCPayload: InitialRSCPayload = {
+      b: 'build-id',
+      c: ['', 'docs', '%%drp:slug:abc123%%'],
+      q: '',
+      i: false,
+      f: [
+        [
+          [
+            '',
+            {
+              children: [
+                'docs',
+                {
+                  children: [['slug', '%%drp:slug:abc123%%', 'd', null], {}],
+                },
+              ],
+            },
+          ],
+          null,
+          'head-node',
+          false,
+        ],
+      ],
+      m: new Set(),
+      G: [(() => null) as any, undefined],
+      S: false,
+      h: null,
+    }
+
+    const response = new Response(null, {
+      headers: {},
+    })
+
+    const payload = createInitialRSCPayloadFromFallbackPrerender(
+      response,
+      fallbackInitialRSCPayload,
+      new URL('https://example.com/docs/original?from=document')
+    )
+
+    const patchedTree = payload.f[0][0]
+    expect(patchedTree[1].children[0]).toBe('docs')
+    expect(patchedTree[1].children[1].children[0]).toEqual([
+      'slug',
+      'original',
+      'd',
+      null,
+    ])
+    expect(payload.q).toBe('?from=document')
+  })
 })

packages/next/src/client/flight-data-helpers.ts

@@ -10,6 +10,10 @@ import type {
 } from '../shared/lib/app-router-types'
 import { PAGE_SEGMENT_KEY } from '../shared/lib/segment'
 import type { NormalizedSearch } from './components/segment-cache/cache-key'
+import {
+  NEXT_REWRITTEN_PATH_HEADER,
+  NEXT_REWRITTEN_QUERY_HEADER,
+} from './components/app-router-headers'
 import {
   getCacheKeyForDynamicParam,
   parseDynamicParamFromURLPart,
@@ -95,10 +99,14 @@ export function createInitialRSCPayloadFromFallbackPrerender(
 
   // Patch the Flight data sent by the server with the correct params parsed
   // from the URL + response object.
-  const renderedPathname =
-    renderedUrlOverride?.pathname ?? getRenderedPathname(response)
-  const renderedSearch =
-    renderedUrlOverride?.search ?? getRenderedSearch(response)
+  const renderedPathname = response.headers.has(NEXT_REWRITTEN_PATH_HEADER)
+    ? getRenderedPathname(response)
+    : ((renderedUrlOverride?.pathname ??
+        getRenderedPathname(response)) as string)
+  const renderedSearch = response.headers.has(NEXT_REWRITTEN_QUERY_HEADER)
+    ? getRenderedSearch(response)
+    : ((renderedUrlOverride?.search ??
+        getRenderedSearch(response)) as NormalizedSearch)
   const canonicalUrl = createHrefFromUrl(
     renderedUrlOverride ?? new URL(location.href)
   )