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

@@ -1386,6 +1386,252 @@ describe('dehydration and rehydration', () => {
     await originalPromise
   })
 
+  it('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('infinite')
+
+    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)
+  })
+
+  it('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)
+  })
+
+  it('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 + 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('infinite')
+
+    const clientCache = new QueryCache()
+    const clientClient = new QueryClient({ queryCache: clientCache })
+    hydrate(clientClient, dehydrated)
+
+    const hydratedQuery = clientCache.find({
+      queryKey: ['infinite-type-restore'],
+    })
+    expect(hydratedQuery?.queryType).toBe('infinite')
+  })
+
+  it('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 + 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 + 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')
+  })
+
+  it('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 + 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.queryType).toBe('infinite')
+
+    query.setOptions({ queryKey: ['infinite-setoptions-guard'] })
+    expect(query.queryType).toBe('infinite')
+  })
+
+  it('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 + 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 + 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')
+  })
+
   // Companion to the test above: when the query already exists in the cache
   // (e.g. after an initial render or a first hydration pass), the same
   // synchronous thenable resolution must also produce status: 'success'.

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('infinite')
 
     expect(result).toMatchObject({
       data: undefined,

packages/query-core/src/hydration.ts

@@ -48,6 +48,7 @@ interface DehydratedQuery {
   state: QueryState
   promise?: Promise<unknown>
   meta?: QueryMeta
+  queryType?: 'infinite'
   // This is only optional because older versions of Query might have dehydrated
   // without it which we need to handle for backwards compatibility.
   // This should be changed to required in the future.
@@ -117,6 +118,7 @@ function dehydrateQuery(
       promise: dehydratePromise(),
     }),
     ...(query.meta && { meta: query.meta }),
+    ...(query.queryType && { queryType: query.queryType }),
   }
 }
 
@@ -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)
@@ -260,6 +270,7 @@ export function hydrate(
             queryKey,
             queryHash,
             meta,
+            _type: queryType,
           },
           // 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,
@@ -93,10 +89,8 @@ export class InfiniteQueryObserver<
       TPageParam
     >,
   ): void {
-    super.setOptions({
-      ...options,
-      behavior: infiniteQueryBehavior(),
-    })
+    options._type = 'infinite'
+    super.setOptions(options)
   }
 
   getOptimisticResult(
@@ -108,7 +102,7 @@ export class InfiniteQueryObserver<
       TPageParam
     >,
   ): InfiniteQueryObserverResult<TData, TError> {
-    options.behavior = infiniteQueryBehavior()
+    options._type = 'infinite'
     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>
+  #queryType?: 'infinite'
 
   #initialState: QueryState<TData, TError>
   #revertState?: QueryState<TData, TError>
@@ -195,6 +197,10 @@ export class Query<
     return this.options.meta
   }
 
+  get queryType() {
+    return this.#queryType
+  }
+
   get promise(): Promise<TData> | undefined {
     return this.#retryer?.promise
   }
@@ -204,6 +210,10 @@ export class Query<
   ): void {
     this.options = { ...this.#defaultOptions, ...options }
 
+    if (options?._type) {
+      this.#queryType = options._type
+    }
+
     this.updateGcTime(this.options.gcTime)
 
     // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
@@ -511,7 +521,13 @@ export class Query<
 
     const context = createFetchContext()
 
-    this.options.behavior?.onFetch(context, this as unknown as Query)
+    const behavior =
+      this.#queryType === 'infinite'
+        ? (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