[Segment Cache] Fix: Key by rewritten search (vercel#81986)

When a request URL is rewritten, the resulting data must be keyed by its
rewritten search params — the ones that were used by the server to
render the page — not the original search params. This works in most
places by encoding the rewritten search params into page's segment key
in the server response. However, the Segment Cache implementation did
not handle this correctly.

The fix is to read the rewritten search params from the
x-nextjs-rewritten-query header in the server response.

In upcoming PRs, we will use a similar approach for regular route
params, too, so that we can lift the params out of the body of the
server response.

---------

Co-authored-by: Janka Uryga <[email protected]>

Author: acdlite

packages/next/src/client/components/segment-cache-impl/cache-key.ts

@@ -1,3 +1,6 @@
+import { NEXT_REWRITTEN_QUERY_HEADER } from '../app-router-headers'
+import type { RSCResponse } from '../router-reducer/fetch-server-response'
+
 // TypeScript trick to simulate opaque types, like in Flow.
 type Opaque<K, T> = T & { __brand: K }
 
@@ -29,3 +32,18 @@ export function createCacheKey(
   } as RouteCacheKey
   return cacheKey
 }
+
+export function getRenderedSearch(response: RSCResponse): NormalizedSearch {
+  // If the server performed a rewrite, the search params used to render the
+  // page will be different from the params in the request URL. In this case,
+  // the response will include a header that gives the rewritten search query.
+  const rewrittenQuery = response.headers.get(NEXT_REWRITTEN_QUERY_HEADER)
+  if (rewrittenQuery !== null) {
+    return (
+      rewrittenQuery === '' ? '' : '?' + rewrittenQuery
+    ) as NormalizedSearch
+  }
+  // If the header is not present, there was no rewrite, so we use the search
+  // query of the response URL.
+  return new URL(response.url).search as NormalizedSearch
+}

packages/next/src/client/components/segment-cache-impl/cache.ts

@@ -42,6 +42,7 @@ import type {
   NormalizedSearch,
   RouteCacheKey,
 } from './cache-key'
+import { getRenderedSearch } from './cache-key'
 import { createTupleMap, type TupleMap, type Prefix } from './tuple-map'
 import { createLRU } from './lru'
 import {
@@ -130,6 +131,7 @@ type PendingRouteCacheEntry = RouteCacheEntryShared & {
   status: EntryStatus.Empty | EntryStatus.Pending
   blockedTasks: Set<PrefetchTask> | null
   canonicalUrl: null
+  renderedSearch: null
   tree: null
   head: HeadData | null
   isHeadPartial: true
@@ -140,6 +142,7 @@ type RejectedRouteCacheEntry = RouteCacheEntryShared & {
   status: EntryStatus.Rejected
   blockedTasks: Set<PrefetchTask> | null
   canonicalUrl: null
+  renderedSearch: null
   tree: null
   head: null
   isHeadPartial: true
@@ -150,6 +153,7 @@ export type FulfilledRouteCacheEntry = RouteCacheEntryShared & {
   status: EntryStatus.Fulfilled
   blockedTasks: null
   canonicalUrl: string
+  renderedSearch: NormalizedSearch
   tree: RouteTree
   head: HeadData
   isHeadPartial: boolean
@@ -413,29 +417,32 @@ export function getSegmentKeypathForTask(
   // cache entry is valid for all possible search param values.
   const isDynamicTask = task.includeDynamicData || !route.isPPREnabled
   return isDynamicTask && path.endsWith('/' + PAGE_SEGMENT_KEY)
-    ? [path, task.key.search]
+    ? [path, route.renderedSearch]
     : [path]
 }
 
 export function readSegmentCacheEntry(
   now: number,
-  routeCacheKey: RouteCacheKey,
+  route: FulfilledRouteCacheEntry,
   path: string
 ): SegmentCacheEntry | null {
   if (!path.endsWith('/' + PAGE_SEGMENT_KEY)) {
     // Fast path. Search params only exist on page segments.
     return readExactSegmentCacheEntry(now, [path])
   }
 
-  // Page segments may or may not contain search params. If they were prefetched
-  // using a dynamic request, then we will have an entry with search params.
-  // Check for that case first.
-  const entryWithSearchParams = readExactSegmentCacheEntry(now, [
-    path,
-    routeCacheKey.search,
-  ])
-  if (entryWithSearchParams !== null) {
-    return entryWithSearchParams
+  const renderedSearch = route.renderedSearch
+  if (renderedSearch !== null) {
+    // Page segments may or may not contain search params. If they were prefetched
+    // using a dynamic request, then we will have an entry with search params.
+    // Check for that case first.
+    const entryWithSearchParams = readExactSegmentCacheEntry(now, [
+      path,
+      renderedSearch,
+    ])
+    if (entryWithSearchParams !== null) {
+      return entryWithSearchParams
+    }
   }
 
   // If we did not find an entry with the given search params, check for a
@@ -550,6 +557,7 @@ export function readOrCreateRouteCacheEntry(
     couldBeIntercepted: true,
     // Similarly, we don't yet know if the route supports PPR.
     isPPREnabled: false,
+    renderedSearch: null,
 
     // LRU-related fields
     keypath: null,
@@ -783,6 +791,7 @@ function fulfillRouteCacheEntry(
   staleAt: number,
   couldBeIntercepted: boolean,
   canonicalUrl: string,
+  renderedSearch: NormalizedSearch,
   isPPREnabled: boolean
 ): FulfilledRouteCacheEntry {
   const fulfilledEntry: FulfilledRouteCacheEntry = entry as any
@@ -793,6 +802,7 @@ function fulfillRouteCacheEntry(
   fulfilledEntry.staleAt = staleAt
   fulfilledEntry.couldBeIntercepted = couldBeIntercepted
   fulfilledEntry.canonicalUrl = canonicalUrl
+  fulfilledEntry.renderedSearch = renderedSearch
   fulfilledEntry.isPPREnabled = isPPREnabled
   pingBlockedTasks(entry)
   return fulfilledEntry
@@ -1133,6 +1143,11 @@ export async function fetchRouteOnCacheMiss(
         return null
       }
 
+      // Get the search params that were used to render the target page. This may
+      // be different from the search params in the request URL, if the page
+      // was rewritten.
+      const renderedSearch = getRenderedSearch(response)
+
       const staleTimeMs = serverData.staleTime * 1000
       fulfillRouteCacheEntry(
         entry,
@@ -1142,6 +1157,7 @@ export async function fetchRouteOnCacheMiss(
         Date.now() + staleTimeMs,
         couldBeIntercepted,
         canonicalUrl,
+        renderedSearch,
         routeIsPPREnabled
       )
     } else {
@@ -1366,6 +1382,19 @@ export async function fetchSegmentPrefetchesUsingDynamicRequest(
       return null
     }
 
+    const renderedSearch = getRenderedSearch(response)
+    if (renderedSearch !== route.renderedSearch) {
+      // The search params that were used to render the target page are
+      // different from the search params in the request URL. This only happens
+      // when there's a dynamic rewrite in between the tree prefetch and the
+      // data prefetch.
+      // TODO: For now, since this is an edge case, we reject the prefetch, but
+      // the proper way to handle this is to evict the stale route tree entry
+      // then fill the cache with the new response.
+      rejectSegmentEntriesIfStillPending(spawnedEntries, Date.now() + 10 * 1000)
+      return null
+    }
+
     // Track when the network connection closes.
     const closed = createPromiseWithResolvers<void>()
 
@@ -1462,6 +1491,11 @@ function writeDynamicTreeResponseIntoCache(
   const isResponsePartial =
     response.headers.get(NEXT_DID_POSTPONE_HEADER) === '1'
 
+  // Get the search params that were used to render the target page. This may
+  // be different from the search params in the request URL, if the page
+  // was rewritten.
+  const renderedSearch = getRenderedSearch(response)
+
   const fulfilledEntry = fulfillRouteCacheEntry(
     entry,
     convertRootFlightRouterStateToRouteTree(flightRouterState),
@@ -1470,6 +1504,7 @@ function writeDynamicTreeResponseIntoCache(
     now + staleTimeMs,
     couldBeIntercepted,
     canonicalUrl,
+    renderedSearch,
     routeIsPPREnabled
   )
 

packages/next/src/client/components/segment-cache-impl/navigation.ts

@@ -22,12 +22,10 @@ import {
   readSegmentCacheEntry,
   waitForSegmentCacheEntry,
   type RouteTree,
+  type FulfilledRouteCacheEntry,
 } from './cache'
-import { createCacheKey, type RouteCacheKey } from './cache-key'
-import {
-  addSearchParamsIfPageSegment,
-  PAGE_SEGMENT_KEY,
-} from '../../../shared/lib/segment'
+import { createCacheKey } from './cache-key'
+import { addSearchParamsIfPageSegment } from '../../../shared/lib/segment'
 import { NavigationResultTag } from '../segment-cache'
 
 type MPANavigationResult = {
@@ -116,7 +114,7 @@ export function navigate(
   const route = readRouteCacheEntry(now, cacheKey)
   if (route !== null && route.status === EntryStatus.Fulfilled) {
     // We have a matching prefetch.
-    const snapshot = readRenderSnapshotFromCache(now, cacheKey, route.tree)
+    const snapshot = readRenderSnapshotFromCache(now, route, route.tree)
     const prefetchFlightRouterState = snapshot.flightRouterState
     const prefetchSeedData = snapshot.seedData
     const prefetchHead = route.head
@@ -255,7 +253,7 @@ function navigationTaskToResult(
 
 function readRenderSnapshotFromCache(
   now: number,
-  routeCacheKey: RouteCacheKey,
+  route: FulfilledRouteCacheEntry,
   tree: RouteTree
 ): { flightRouterState: FlightRouterState; seedData: CacheNodeSeedData } {
   let childRouterStates: { [parallelRouteKey: string]: FlightRouterState } = {}
@@ -266,11 +264,7 @@ function readRenderSnapshotFromCache(
   if (slots !== null) {
     for (const parallelRouteKey in slots) {
       const childTree = slots[parallelRouteKey]
-      const childResult = readRenderSnapshotFromCache(
-        now,
-        routeCacheKey,
-        childTree
-      )
+      const childResult = readRenderSnapshotFromCache(now, route, childTree)
       childRouterStates[parallelRouteKey] = childResult.flightRouterState
       childSeedDatas[parallelRouteKey] = childResult.seedData
     }
@@ -280,7 +274,7 @@ function readRenderSnapshotFromCache(
   let loading: LoadingModuleData | Promise<LoadingModuleData> = null
   let isPartial: boolean = true
 
-  const segmentEntry = readSegmentCacheEntry(now, routeCacheKey, tree.key)
+  const segmentEntry = readSegmentCacheEntry(now, route, tree.key)
   if (segmentEntry !== null) {
     switch (segmentEntry.status) {
       case EntryStatus.Fulfilled: {
@@ -315,23 +309,20 @@ function readRenderSnapshotFromCache(
     }
   }
 
-  const segment =
-    tree.segment === PAGE_SEGMENT_KEY && routeCacheKey.search
-      ? // The navigation implementation expects the search params to be
-        // included in the segment. However, the Segment Cache tracks search
-        // params separately from the rest of the segment key. So we need to
-        // add them back here.
-        //
-        // See corresponding comment in convertFlightRouterStateToTree.
-        //
-        // TODO: What we should do instead is update the navigation diffing
-        // logic to compare search params explicitly. This is a temporary
-        // solution until more of the Segment Cache implementation has settled.
-        addSearchParamsIfPageSegment(
-          tree.segment,
-          Object.fromEntries(new URLSearchParams(routeCacheKey.search))
-        )
-      : tree.segment
+  // The navigation implementation expects the search params to be
+  // included in the segment. However, the Segment Cache tracks search
+  // params separately from the rest of the segment key. So we need to
+  // add them back here.
+  //
+  // See corresponding comment in convertFlightRouterStateToTree.
+  //
+  // TODO: What we should do instead is update the navigation diffing
+  // logic to compare search params explicitly. This is a temporary
+  // solution until more of the Segment Cache implementation has settled.
+  const segment = addSearchParamsIfPageSegment(
+    tree.segment,
+    Object.fromEntries(new URLSearchParams(route.renderedSearch))
+  )
 
   return {
     flightRouterState: [

packages/next/src/client/components/segment-cache-impl/scheduler.ts

@@ -1,6 +1,7 @@
 import type {
   FlightRouterState,
   Segment as FlightRouterStateSegment,
+  Segment,
 } from '../../../server/app-render/types'
 import { HasLoadingBoundary } from '../../../server/app-render/types'
 import { matchSegment } from '../match-segments'
@@ -28,6 +29,10 @@ import {
 } from './cache'
 import type { RouteCacheKey } from './cache-key'
 import { getCurrentCacheVersion, PrefetchPriority } from '../segment-cache'
+import {
+  addSearchParamsIfPageSegment,
+  PAGE_SEGMENT_KEY,
+} from '../../../shared/lib/segment'
 
 const scheduleMicrotask =
   typeof queueMicrotask === 'function'
@@ -657,7 +662,11 @@ function diffRouteTreeAgainstCurrent(
         oldTreeChild?.[0]
       if (
         oldTreeChildSegment !== undefined &&
-        matchSegment(newTreeChildSegment, oldTreeChildSegment)
+        doesCurrentSegmentMatchCachedSegment(
+          route,
+          newTreeChildSegment,
+          oldTreeChildSegment
+        )
       ) {
         // This segment is already part of the current route. Keep traversing.
         const requestTreeChild = diffRouteTreeAgainstCurrent(
@@ -1177,6 +1186,34 @@ function upsertSegmentOnCompletion(
   }, noop)
 }
 
+function doesCurrentSegmentMatchCachedSegment(
+  route: FulfilledRouteCacheEntry,
+  currentSegment: Segment,
+  cachedSegment: Segment
+): boolean {
+  if (cachedSegment === PAGE_SEGMENT_KEY) {
+    // In the FlightRouterState stored by the router, the page segment has the
+    // rendered search params appended to the name of the segment. In the
+    // prefetch cache, however, this is stored separately. So, when comparing
+    // the router's current FlightRouterState to the cached FlightRouterState,
+    // we need to make sure we compare both parts of the segment.
+    // TODO: This is not modeled clearly. We use the same type,
+    // FlightRouterState, for both the CacheNode tree _and_ the prefetch cache
+    // _and_ the server response format, when conceptually those are three
+    // different things and treated in different ways. We should encode more of
+    // this information into the type design so mistakes are less likely.
+    return (
+      currentSegment ===
+      addSearchParamsIfPageSegment(
+        PAGE_SEGMENT_KEY,
+        Object.fromEntries(new URLSearchParams(route.renderedSearch))
+      )
+    )
+  }
+  // Non-page segments are compared using the same function as the server
+  return matchSegment(cachedSegment, currentSegment)
+}
+
 // -----------------------------------------------------------------------------
 // The remainder of the module is a MinHeap implementation. Try not to put any
 // logic below here unless it's related to the heap algorithm. We can extract