Fix/infinite query hydration 8825

.changeset/stupid-seals-live.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/query-core': patch
+---
+
+fix: preserve infinite query behavior during SSR hydration (#8825)

packages/query-core/src/__tests__/hydration.test.tsx

@@ -1385,4 +1385,235 @@ describe('dehydration and rehydration', () => {
     // error and test will fail
     await originalPromise
   })
+
+  test('should preserve queryType for infinite queries during hydration', async () => {
+    const queryCache = new QueryCache()
+    const queryClient = new QueryClient({ queryCache })
+
+    await vi.waitFor(() =>
+      queryClient.prefetchInfiniteQuery({
+        queryKey: ['infinite'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({
+            items: [`page-${pageParam}`],
+            nextCursor: pageParam + 1,
+          })),
+        initialPageParam: 0,
+        getNextPageParam: (
+          lastPage: { items: Array<string>; nextCursor: number },
+        ) => lastPage.nextCursor,
+      }),
+    )
+
+    const dehydrated = dehydrate(queryClient)
+
+    const infiniteQueryState = dehydrated.queries.find(
+      (q) => q.queryKey[0] === 'infinite',
+    )
+    expect(infiniteQueryState?.queryType).toBe('infiniteQuery')
+
+    const hydrationCache = new QueryCache()
+    const hydrationClient = new QueryClient({ queryCache: hydrationCache })
+    hydrate(hydrationClient, dehydrated)
+
+    const hydratedQuery = hydrationCache.find({ queryKey: ['infinite'] })
+    expect(hydratedQuery?.state.data).toBeDefined()
+    expect(hydratedQuery?.state.data).toHaveProperty('pages')
+    expect(hydratedQuery?.state.data).toHaveProperty('pageParams')
+    expect((hydratedQuery?.state.data as any).pages).toHaveLength(1)
+  })
+
+  test('should attach infiniteQueryBehavior during hydration', async () => {
+    const queryCache = new QueryCache()
+    const queryClient = new QueryClient({ queryCache })
+
+    await vi.waitFor(() =>
+      queryClient.prefetchInfiniteQuery({
+        queryKey: ['infinite-with-behavior'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({ data: `page-${pageParam}`, next: pageParam + 1 })),
+        initialPageParam: 0,
+        getNextPageParam: (lastPage: { data: string; next: number }) =>
+          lastPage.next,
+      }),
+    )
+
+    const dehydrated = dehydrate(queryClient)
+
+    const hydrationCache = new QueryCache()
+    const hydrationClient = new QueryClient({ queryCache: hydrationCache })
+    hydrate(hydrationClient, dehydrated)
+
+    const result = await vi.waitFor(() =>
+      hydrationClient.fetchInfiniteQuery({
+        queryKey: ['infinite-with-behavior'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({ data: `page-${pageParam}`, next: pageParam + 1 })),
+        initialPageParam: 0,
+        getNextPageParam: (lastPage: { data: string; next: number }) =>
+          lastPage.next,
+      }),
+    )
+
+    expect(result.pages).toHaveLength(1)
+    expect(result.pageParams).toHaveLength(1)
+  })
+
+  test('should restore infinite query type through dehydrate and hydrate cycle', async () => {
+    const serverClient = new QueryClient({ queryCache: new QueryCache() })
+
+    await vi.waitFor(() =>
+      serverClient.prefetchInfiniteQuery({
+        queryKey: ['infinite-type-restore'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({ items: [`item-${pageParam}`], next: (pageParam as number) + 1 })),
+        initialPageParam: 0,
+        getNextPageParam: (lastPage: { items: Array<string>; next: number }) =>
+          lastPage.next,
+      }),
+    )
+
+    const dehydrated = dehydrate(serverClient)
+
+    const dehydratedQuery = dehydrated.queries.find(
+      (q) => q.queryKey[0] === 'infinite-type-restore',
+    )
+    expect(dehydratedQuery?.queryType).toBe('infiniteQuery')
+
+    const clientCache = new QueryCache()
+    const clientClient = new QueryClient({ queryCache: clientCache })
+    hydrate(clientClient, dehydrated)
+
+    const hydratedQuery = clientCache.find({ queryKey: ['infinite-type-restore'] })
+    expect(hydratedQuery?.type).toBe('infiniteQuery')
+  })
+
+  test('should preserve pages structure when refetching infinite query after hydration', async () => {
+    const serverClient = new QueryClient({ queryCache: new QueryCache() })
+
+    await vi.waitFor(() =>
+      serverClient.prefetchInfiniteQuery({
+        queryKey: ['refetch'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({
+            items: [`page-${pageParam}`],
+            next: (pageParam as number) + 1,
+          })),
+        initialPageParam: 0,
+        getNextPageParam: (lastPage: { items: Array<string>; next: number }) =>
+          lastPage.next,
+      }),
+    )
+
+    const dehydrated = dehydrate(serverClient)
+
+    const clientCache = new QueryCache()
+    const clientClient = new QueryClient({ queryCache: clientCache })
+    hydrate(clientClient, dehydrated)
+
+    const beforeRefetch = clientClient.getQueryData<{
+      pages: Array<{ items: Array<string>; next: number }>
+      pageParams: Array<unknown>
+    }>(['refetch'])
+    expect(beforeRefetch?.pages).toHaveLength(1)
+    expect(beforeRefetch?.pageParams).toHaveLength(1)
+
+    const result = await vi.waitFor(() =>
+      clientClient.fetchInfiniteQuery({
+        queryKey: ['refetch'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({
+            items: [`page-${pageParam}`],
+            next: (pageParam as number) + 1,
+          })),
+        initialPageParam: 0,
+        getNextPageParam: (lastPage: { items: Array<string>; next: number }) =>
+          lastPage.next,
+      }),
+    )
+
+    expect(result).toHaveProperty('pages')
+    expect(result).toHaveProperty('pageParams')
+    expect(Array.isArray(result.pages)).toBe(true)
+    expect(result.pages).toHaveLength(1)
+    expect(result.pages[0]).toHaveProperty('items')
+  })
+
+  test('should retain infinite query type after subsequent setOptions calls', async () => {
+    const serverClient = new QueryClient({ queryCache: new QueryCache() })
+
+    await vi.waitFor(() =>
+      serverClient.prefetchInfiniteQuery({
+        queryKey: ['infinite-setoptions-guard'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({ data: `p${pageParam}`, next: (pageParam as number) + 1 })),
+        initialPageParam: 0,
+        getNextPageParam: (lastPage: { data: string; next: number }) =>
+          lastPage.next,
+      }),
+    )
+
+    const dehydrated = dehydrate(serverClient)
+
+    const clientCache = new QueryCache()
+    const clientClient = new QueryClient({ queryCache: clientCache })
+    hydrate(clientClient, dehydrated)
+
+    const query = clientCache.find({ queryKey: ['infinite-setoptions-guard'] })!
+    expect(query.type).toBe('infiniteQuery')
+
+    query.setOptions({ queryKey: ['infinite-setoptions-guard'] })
+    expect(query.type).toBe('infiniteQuery')
+  })
+
+  test('should restore all pages when refetching multi-page infinite query after hydration', async () => {
+    const serverClient = new QueryClient({ queryCache: new QueryCache() })
+
+    await vi.waitFor(() =>
+      serverClient.prefetchInfiniteQuery({
+        queryKey: ['infinite-multipage-restore'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({
+            items: [`item-${pageParam}`],
+            next: (pageParam as number) + 1,
+          })),
+        initialPageParam: 0,
+        pages: 2,
+        getNextPageParam: (lastPage: { items: Array<string>; next: number }) =>
+          lastPage.next,
+      }),
+    )
+
+    const dehydrated = dehydrate(serverClient)
+
+    const clientCache = new QueryCache()
+    const clientClient = new QueryClient({ queryCache: clientCache })
+    hydrate(clientClient, dehydrated)
+
+    const beforeRefetch = clientClient.getQueryData<{
+      pages: Array<unknown>
+      pageParams: Array<unknown>
+    }>(['infinite-multipage-restore'])
+    expect(beforeRefetch?.pages).toHaveLength(2)
+
+    const result = await vi.waitFor(() =>
+      clientClient.fetchInfiniteQuery({
+        queryKey: ['infinite-multipage-restore'],
+        queryFn: async ({ pageParam }) =>
+          sleep(0).then(() => ({
+            items: [`item-${pageParam}`],
+            next: (pageParam as number) + 1,
+          })),
+        initialPageParam: 0,
+        pages: 2,
+        getNextPageParam: (lastPage: { items: Array<string>; next: number }) =>
+          lastPage.next,
+      }),
+    )
+
+    expect(result.pages).toHaveLength(2)
+    expect(result.pageParams).toHaveLength(2)
+    expect(result.pages[0]).toHaveProperty('items')
+    expect(result.pages[1]).toHaveProperty('items')
+  })
 })

packages/query-core/src/__tests__/infiniteQueryObserver.test.tsx

@@ -235,8 +235,7 @@ describe('InfiniteQueryObserver', () => {
 
     const result = observer.getOptimisticResult(options)
 
-    expect(options.behavior).toBeDefined()
-    expect(options.behavior?.onFetch).toBeDefined()
+    expect(options._type).toBe('infiniteQuery')
 
     expect(result).toMatchObject({
       data: undefined,

packages/query-core/src/hydration.ts

@@ -52,6 +52,7 @@ interface DehydratedQuery {
   // without it which we need to handle for backwards compatibility.
   // This should be changed to required in the future.
   dehydratedAt?: number
+  queryType?: 'query' | 'infiniteQuery'
 }
 
 export interface DehydratedState {
@@ -113,6 +114,7 @@ function dehydrateQuery(
     },
     queryKey: query.queryKey,
     queryHash: query.queryHash,
+    queryType: query.type,
     ...(query.state.status === 'pending' && {
       promise: dehydratePromise(),
     }),
@@ -209,7 +211,15 @@ export function hydrate(
   })
 
   queries.forEach(
-    ({ queryKey, state, queryHash, meta, promise, dehydratedAt }) => {
+    ({
+      queryKey,
+      state,
+      queryHash,
+      meta,
+      promise,
+      dehydratedAt,
+      queryType,
+    }) => {
       const syncData = promise ? tryResolveSync(promise) : undefined
       const rawData = state.data === undefined ? syncData?.data : state.data
       const data = rawData === undefined ? rawData : deserializeData(rawData)
@@ -239,16 +249,19 @@ export function hydrate(
           })
         }
       } else {
+        const queryOptions = {
+          ...client.getDefaultOptions().hydrate?.queries,
+          ...options?.defaultOptions?.queries,
+          queryKey,
+          queryHash,
+          meta,
+          _type: queryType === 'infiniteQuery' ? ('infiniteQuery' as const) : undefined,
+        }
+
         // Restore query
         query = queryCache.build(
           client,
-          {
-            ...client.getDefaultOptions().hydrate?.queries,
-            ...options?.defaultOptions?.queries,
-            queryKey,
-            queryHash,
-            meta,
-          },
+          queryOptions,
           // Reset fetch status to idle to avoid
           // query being stuck in fetching state upon hydration
           {

packages/query-core/src/infiniteQueryObserver.ts

@@ -1,9 +1,5 @@
 import { QueryObserver } from './queryObserver'
-import {
-  hasNextPage,
-  hasPreviousPage,
-  infiniteQueryBehavior,
-} from './infiniteQueryBehavior'
+import { hasNextPage, hasPreviousPage } from './infiniteQueryBehavior'
 import type { Subscribable } from './subscribable'
 import type {
   DefaultError,
@@ -95,7 +91,7 @@ export class InfiniteQueryObserver<
   ): void {
     super.setOptions({
       ...options,
-      behavior: infiniteQueryBehavior(),
+      _type: 'infiniteQuery',
     })
   }
 
@@ -108,7 +104,7 @@ export class InfiniteQueryObserver<
       TPageParam
     >,
   ): InfiniteQueryObserverResult<TData, TError> {
-    options.behavior = infiniteQueryBehavior()
+    options._type = 'infiniteQuery'
     return super.getOptimisticResult(options) as InfiniteQueryObserverResult<
       TData,
       TError

packages/query-core/src/query.ts

@@ -10,6 +10,7 @@ import {
 import { notifyManager } from './notifyManager'
 import { CancelledError, canFetch, createRetryer } from './retryer'
 import { Removable } from './removable'
+import { infiniteQueryBehavior } from './infiniteQueryBehavior'
 import type { QueryCache } from './queryCache'
 import type { QueryClient } from './queryClient'
 import type {
@@ -166,6 +167,7 @@ export class Query<
   queryHash: string
   options!: QueryOptions<TQueryFnData, TError, TData, TQueryKey>
   state: QueryState<TData, TError>
+  type: 'query' | 'infiniteQuery'
 
   #initialState: QueryState<TData, TError>
   #revertState?: QueryState<TData, TError>
@@ -181,6 +183,7 @@ export class Query<
 
     this.#abortSignalConsumed = false
     this.#defaultOptions = config.defaultOptions
+    this.type = 'query'
     this.setOptions(config.options)
     this.observers = []
     this.#client = config.client
@@ -204,6 +207,10 @@ export class Query<
   ): void {
     this.options = { ...this.#defaultOptions, ...options }
 
+    if (options?._type) {
+      this.type = options._type
+    }
+
     this.updateGcTime(this.options.gcTime)
 
     // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
@@ -510,7 +517,13 @@ export class Query<
 
     const context = createFetchContext()
 
-    this.options.behavior?.onFetch(context, this as unknown as Query)
+    const behavior =
+      this.type === 'infiniteQuery'
+        ? (infiniteQueryBehavior(
+            (this.options as { pages?: number }).pages,
+          ) as QueryBehavior<TQueryFnData, TError, TData, TQueryKey>)
+        : this.options.behavior
+    behavior?.onFetch(context, this as unknown as Query)
 
     // Store state in case the current fetch needs to be reverted
     this.#revertState = this.state