[pull] canary from vercel:canary

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