fix: no hydration when new promise comes in (#8383)

* add failing repro test * update assertinos * add logg * ehm - maybe fix? * rm -only * make example * upd * ad debug logs * more debugging * push * maybe??? * rm log * revert * fix: ? * fix: check for pending status again otherwise, we risk including promises that happen because of background updates (think persistQueryClient) * fix: clear serverQueryClient between "requests" otherwise, we are re-using the cache and the query won't be in "pending" state the second time around * chore: remove logs * rethrow next build error * kick off ci again * add `shouldRedactError` option * pluralize * docs * chore: more memory * chore: revert more memory * don't compare statuses if they don't exist * ci: apply automated fixes * lint --------- Co-authored-by: Dominik Dorfmeister <[email protected]> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
TanStack · Feb 17, 2025 · 6ca0eb7 · 6ca0eb7
1 parent bb25d06
commit 6ca0eb7
Show file tree

Hide file tree

Showing 11 changed files with 207 additions and 19 deletions.
diff --git a/docs/framework/react/guides/advanced-ssr.md b/docs/framework/react/guides/advanced-ssr.md
@@ -392,6 +392,14 @@ function makeQueryClient() {
         shouldDehydrateQuery: (query) =>
           defaultShouldDehydrateQuery(query) ||
           query.state.status === 'pending',
+        shouldRedactErrors: (error) => {
+          // We should not catch Next.js server errors
+          // as that's how Next.js detects dynamic pages
+          // so we cannot redact them.
+          // Next.js also automatically redacts errors for us
+          // with better digests.
+          return false
+        },
       },
     },
   })

diff --git a/docs/framework/react/reference/hydration.md b/docs/framework/react/reference/hydration.md
@@ -38,6 +38,12 @@ const dehydratedState = dehydrate(queryClient, {
     - Defaults to only including successful queries
     - If you would like to extend the function while retaining the default behavior, import and execute `defaultShouldDehydrateQuery` as part of the return statement
   - `serializeData?: (data: any) => any` A function to transform (serialize) data during dehydration.
+  - `shouldRedactErrors?: (error: unknown) => boolean`
+    - Optional
+    - Whether to redact errors from the server during dehydration.
+    - The function is called for each error in the cache
+      - Return `true` to redact this error, or `false` otherwise
+    - Defaults to redacting all errors
 
 **Returns**
 

diff --git a/integrations/react-next-15/app/_action.ts b/integrations/react-next-15/app/_action.ts
@@ -0,0 +1,11 @@
+'use server'
+
+import { revalidatePath } from 'next/cache'
+import { countRef } from './make-query-client'
+
+export async function queryExampleAction() {
+  await Promise.resolve()
+  countRef.current++
+  revalidatePath('/', 'page')
+  return undefined
+}
diff --git a/integrations/react-next-15/app/client-component.tsx b/integrations/react-next-15/app/client-component.tsx
@@ -8,10 +8,14 @@ export function ClientComponent() {
   const query = useQuery({
     queryKey: ['data'],
     queryFn: async () => {
-      await new Promise((r) => setTimeout(r, 1000))
+      const { count } = await (
+        await fetch('http://localhost:3000/count')
+      ).json()
+
       return {
         text: 'data from client',
         date: Temporal.PlainDate.from('2023-01-01'),
+        count,
       }
     },
   })
@@ -26,7 +30,7 @@ export function ClientComponent() {
 
   return (
     <div>
-      {query.data.text} - {query.data.date.toJSON()}
+      {query.data.text} - {query.data.date.toJSON()} - {query.data.count}
     </div>
   )
 }
diff --git a/integrations/react-next-15/app/count/route.ts b/integrations/react-next-15/app/count/route.ts
@@ -0,0 +1,5 @@
+import { countRef } from '../make-query-client'
+
+export const GET = () => {
+  return Response.json({ count: countRef.current })
+}
diff --git a/integrations/react-next-15/app/make-query-client.ts b/integrations/react-next-15/app/make-query-client.ts
@@ -10,6 +10,10 @@ const plainDate = {
   test: (v) => v instanceof Temporal.PlainDate,
 } satisfies TsonType<Temporal.PlainDate, string>
 
+export const countRef = {
+  current: 0,
+}
+
 export const tson = createTson({
   types: [plainDate],
 })
@@ -22,16 +26,27 @@ export function makeQueryClient() {
          * Called when the query is rebuilt from a prefetched
          * promise, before the query data is put into the cache.
          */
-        deserializeData: tson.deserialize,
+        deserializeData: (data) => {
+          return tson.deserialize(data)
+        },
       },
       queries: {
         staleTime: 60 * 1000,
       },
       dehydrate: {
-        serializeData: tson.serialize,
-        shouldDehydrateQuery: (query) =>
-          defaultShouldDehydrateQuery(query) ||
-          query.state.status === 'pending',
+        serializeData: (data) => {
+          return tson.serialize(data)
+        },
+        shouldDehydrateQuery: (query) => {
+          return (
+            defaultShouldDehydrateQuery(query) ||
+            query.state.status === 'pending'
+          )
+        },
+        shouldRedactErrors: (error) => {
+          // Next.js automatically redacts errors for us
+          return false
+        },
       },
     },
   })

diff --git a/integrations/react-next-15/app/page.tsx b/integrations/react-next-15/app/page.tsx
@@ -1,30 +1,41 @@
+import { headers } from 'next/headers'
 import React from 'react'
 import { HydrationBoundary, dehydrate } from '@tanstack/react-query'
 import { Temporal } from '@js-temporal/polyfill'
 import { ClientComponent } from './client-component'
-import { makeQueryClient, tson } from './make-query-client'
+import { makeQueryClient } from './make-query-client'
+import { queryExampleAction } from './_action'
 
-const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms))
-
-export default async function Home() {
+export default function Home() {
   const queryClient = makeQueryClient()
 
-  void queryClient.prefetchQuery({
+  queryClient.prefetchQuery({
     queryKey: ['data'],
     queryFn: async () => {
-      await sleep(2000)
+      const { count } = await (
+        await fetch('http://localhost:3000/count', {
+          headers: await headers(),
+        })
+      ).json()
+
       return {
         text: 'data from server',
         date: Temporal.PlainDate.from('2024-01-01'),
+        count,
       }
     },
   })
 
+  const state = dehydrate(queryClient)
+
   return (
     <main>
-      <HydrationBoundary state={dehydrate(queryClient)}>
+      <HydrationBoundary state={state}>
         <ClientComponent />
       </HydrationBoundary>
+      <form action={queryExampleAction}>
+        <button type="submit">Increment</button>
+      </form>
     </main>
   )
 }
diff --git a/integrations/react-next-15/app/providers.tsx b/integrations/react-next-15/app/providers.tsx
@@ -1,11 +1,34 @@
+// In Next.js, this file would be called: app/providers.tsx
 'use client'
-import { QueryClientProvider } from '@tanstack/react-query'
+
+// Since QueryClientProvider relies on useContext under the hood, we have to put 'use client' on top
+import { QueryClientProvider, isServer } from '@tanstack/react-query'
 import { ReactQueryDevtools } from '@tanstack/react-query-devtools'
-import * as React from 'react'
+import type { QueryClient } from '@tanstack/react-query'
 import { makeQueryClient } from '@/app/make-query-client'
 
+let browserQueryClient: QueryClient | undefined = undefined
+
+function getQueryClient() {
+  if (isServer) {
+    // Server: always make a new query client
+    return makeQueryClient()
+  } else {
+    // Browser: make a new query client if we don't already have one
+    // This is very important, so we don't re-make a new client if React
+    // suspends during the initial render. This may not be needed if we
+    // have a suspense boundary BELOW the creation of the query client
+    if (!browserQueryClient) browserQueryClient = makeQueryClient()
+    return browserQueryClient
+  }
+}
+
 export default function Providers({ children }: { children: React.ReactNode }) {
-  const [queryClient] = React.useState(() => makeQueryClient())
+  // NOTE: Avoid useState when initializing the query client if you don't
+  //       have a suspense boundary between this and the code that may
+  //       suspend because React will throw away the client on the initial
+  //       render if it suspends and there is no boundary
+  const queryClient = getQueryClient()
 
   return (
     <QueryClientProvider client={queryClient}>

diff --git a/packages/query-core/src/__tests__/hydration.test.tsx b/packages/query-core/src/__tests__/hydration.test.tsx
@@ -1066,4 +1066,80 @@ describe('dehydration and rehydration', () => {
     clientQueryClient.clear()
     serverQueryClient.clear()
   })
+
+  test('should overwrite data when a new promise is streamed in', async () => {
+    const serializeDataMock = vi.fn((data: any) => data)
+    const deserializeDataMock = vi.fn((data: any) => data)
+
+    const countRef = { current: 0 }
+    // --- server ---
+    const serverQueryClient = createQueryClient({
+      defaultOptions: {
+        dehydrate: {
+          shouldDehydrateQuery: () => true,
+          serializeData: serializeDataMock,
+        },
+      },
+    })
+
+    const query = {
+      queryKey: ['data'],
+      queryFn: async () => {
+        await sleep(10)
+        return countRef.current
+      },
+    }
+
+    const promise = serverQueryClient.prefetchQuery(query)
+
+    let dehydrated = dehydrate(serverQueryClient)
+
+    // --- client ---
+
+    const clientQueryClient = createQueryClient({
+      defaultOptions: {
+        hydrate: {
+          deserializeData: deserializeDataMock,
+        },
+      },
+    })
+
+    hydrate(clientQueryClient, dehydrated)
+
+    await promise
+    await waitFor(() =>
+      expect(clientQueryClient.getQueryData(query.queryKey)).toBe(0),
+    )
+
+    expect(serializeDataMock).toHaveBeenCalledTimes(1)
+    expect(serializeDataMock).toHaveBeenCalledWith(0)
+
+    expect(deserializeDataMock).toHaveBeenCalledTimes(1)
+    expect(deserializeDataMock).toHaveBeenCalledWith(0)
+
+    // --- server ---
+    countRef.current++
+    serverQueryClient.clear()
+    const promise2 = serverQueryClient.prefetchQuery(query)
+
+    dehydrated = dehydrate(serverQueryClient)
+
+    // --- client ---
+
+    hydrate(clientQueryClient, dehydrated)
+
+    await promise2
+    await waitFor(() =>
+      expect(clientQueryClient.getQueryData(query.queryKey)).toBe(1),
+    )
+
+    expect(serializeDataMock).toHaveBeenCalledTimes(2)
+    expect(serializeDataMock).toHaveBeenCalledWith(1)
+
+    expect(deserializeDataMock).toHaveBeenCalledTimes(2)
+    expect(deserializeDataMock).toHaveBeenCalledWith(1)
+
+    clientQueryClient.clear()
+    serverQueryClient.clear()
+  })
 })
diff --git a/packages/query-core/src/hydration.ts b/packages/query-core/src/hydration.ts
@@ -22,6 +22,7 @@ export interface DehydrateOptions {
   serializeData?: TransformerFn
   shouldDehydrateMutation?: (mutation: Mutation) => boolean
   shouldDehydrateQuery?: (query: Query) => boolean
+  shouldRedactErrors?: (error: unknown) => boolean
 }
 
 export interface HydrateOptions {
@@ -70,6 +71,7 @@ function dehydrateMutation(mutation: Mutation): DehydratedMutation {
 function dehydrateQuery(
   query: Query,
   serializeData: TransformerFn,
+  shouldRedactErrors: (error: unknown) => boolean,
 ): DehydratedQuery {
   return {
     state: {
@@ -82,6 +84,11 @@ function dehydrateQuery(
     queryHash: query.queryHash,
     ...(query.state.status === 'pending' && {
       promise: query.promise?.then(serializeData).catch((error) => {
+        if (!shouldRedactErrors(error)) {
+          // Reject original error if it should not be redacted
+          return Promise.reject(error)
+        }
+        // If not in production, log original error before rejecting redacted error
         if (process.env.NODE_ENV !== 'production') {
           console.error(
             `A query that was dehydrated as pending ended up rejecting. [${query.queryHash}]: ${error}; The error will be redacted in production builds`,
@@ -102,6 +109,10 @@ export function defaultShouldDehydrateQuery(query: Query) {
   return query.state.status === 'success'
 }
 
+export function defaultshouldRedactErrors(_: unknown) {
+  return true
+}
+
 export function dehydrate(
   client: QueryClient,
   options: DehydrateOptions = {},
@@ -123,6 +134,11 @@ export function dehydrate(
     client.getDefaultOptions().dehydrate?.shouldDehydrateQuery ??
     defaultShouldDehydrateQuery
 
+  const shouldRedactErrors =
+    options.shouldRedactErrors ??
+    client.getDefaultOptions().dehydrate?.shouldRedactErrors ??
+    defaultshouldRedactErrors
+
   const serializeData =
     options.serializeData ??
     client.getDefaultOptions().dehydrate?.serializeData ??
@@ -132,7 +148,9 @@ export function dehydrate(
     .getQueryCache()
     .getAll()
     .flatMap((query) =>
-      filterQuery(query) ? [dehydrateQuery(query, serializeData)] : [],
+      filterQuery(query)
+        ? [dehydrateQuery(query, serializeData, shouldRedactErrors)]
+        : [],
     )
 
   return { mutations, queries }

diff --git a/packages/react-query/src/HydrationBoundary.tsx b/packages/react-query/src/HydrationBoundary.tsx
@@ -24,6 +24,13 @@ export interface HydrationBoundaryProps {
   queryClient?: QueryClient
 }
 
+const hasProperty = <TKey extends string>(
+  obj: unknown,
+  key: TKey,
+): obj is { [k in TKey]: unknown } => {
+  return typeof obj === 'object' && obj !== null && key in obj
+}
+
 export const HydrationBoundary = ({
   children,
   options = {},
@@ -73,7 +80,11 @@ export const HydrationBoundary = ({
         } else {
           const hydrationIsNewer =
             dehydratedQuery.state.dataUpdatedAt >
-            existingQuery.state.dataUpdatedAt
+              existingQuery.state.dataUpdatedAt || // RSC special serialized then-able chunks
+            (hasProperty(dehydratedQuery.promise, 'status') &&
+              hasProperty(existingQuery.promise, 'status') &&
+              dehydratedQuery.promise.status !== existingQuery.promise.status)
+
           const queryAlreadyQueued = hydrationQueue?.find(
             (query) => query.queryHash === dehydratedQuery.queryHash,
           )