Fix misc hydration bugs

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

@@ -473,10 +473,12 @@ describe('dehydration and rehydration', () => {
     const serverAddTodo = vi
       .fn()
       .mockImplementation(() => Promise.reject(new Error('offline')))
-    const serverOnMutate = vi.fn().mockImplementation((variables) => {
-      const optimisticTodo = { id: 1, text: variables.text }
-      return { optimisticTodo }
-    })
+    const serverOnMutate = vi
+      .fn()
+      .mockImplementation((variables: { text: string }) => {
+        const optimisticTodo = { id: 1, text: variables.text }
+        return { optimisticTodo }
+      })
     const serverOnSuccess = vi.fn()
 
     const serverClient = new QueryClient()
@@ -511,13 +513,17 @@ describe('dehydration and rehydration', () => {
     const parsed = JSON.parse(stringified)
     const client = new QueryClient()
 
-    const clientAddTodo = vi.fn().mockImplementation((variables) => {
-      return { id: 2, text: variables.text }
-    })
-    const clientOnMutate = vi.fn().mockImplementation((variables) => {
-      const optimisticTodo = { id: 1, text: variables.text }
-      return { optimisticTodo }
-    })
+    const clientAddTodo = vi
+      .fn()
+      .mockImplementation((variables: { text: string }) => {
+        return { id: 2, text: variables.text }
+      })
+    const clientOnMutate = vi
+      .fn()
+      .mockImplementation((variables: { text: string }) => {
+        const optimisticTodo = { id: 1, text: variables.text }
+        return { optimisticTodo }
+      })
     const clientOnSuccess = vi.fn()
 
     client.setMutationDefaults(['addTodo'], {
@@ -1116,6 +1122,60 @@ describe('dehydration and rehydration', () => {
     serverQueryClient.clear()
   })
 
+  test('should not overwrite query in cache if existing query is newer (with promise)', async () => {
+    // --- server ---
+
+    const serverQueryClient = new QueryClient({
+      defaultOptions: {
+        dehydrate: {
+          shouldDehydrateQuery: () => true,
+        },
+      },
+    })
+
+    const promise = serverQueryClient.prefetchQuery({
+      queryKey: ['data'],
+      queryFn: async () => {
+        await sleep(10)
+        return 'server data'
+      },
+    })
+
+    const dehydrated = dehydrate(serverQueryClient)
+
+    await vi.advanceTimersByTimeAsync(10)
+    await promise
+
+    // Pretend the output of this server part is cached for a long time
+
+    // --- client ---
+
+    await vi.advanceTimersByTimeAsync(10_000) // Arbitrary time in the future
+
+    const clientQueryClient = new QueryClient()
+
+    clientQueryClient.setQueryData(['data'], 'newer data', {
+      updatedAt: Date.now(),
+    })
+
+    hydrate(clientQueryClient, dehydrated)
+
+    // If the query was hydrated in error, it would still take some time for it
+    // to end up in the cache, so for the test to fail properly on regressions,
+    // wait for the fetchStatus to be idle
+    await vi.waitFor(() =>
+      expect(clientQueryClient.getQueryState(['data'])?.fetchStatus).toBe(
+        'idle',
+      ),
+    )
+    await vi.waitFor(() =>
+      expect(clientQueryClient.getQueryData(['data'])).toBe('newer data'),
+    )
+
+    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)
@@ -1291,4 +1351,55 @@ describe('dehydration and rehydration', () => {
     process.env.NODE_ENV = originalNodeEnv
     consoleMock.mockRestore()
   })
+
+  // When React hydrates promises across RSC/client boundaries, it passes
+  // them as special ReactPromise types. There are situations where the
+  // promise might have time to resolve before we end up hydrating it, in
+  // which case React will have made it a special synchronous thenable where
+  // .then() resolves immediately.
+  // In these cases it's important we hydrate the data synchronously, or else
+  // the data in the cache wont match the content that was rendered on the server.
+  // What can end up happening otherwise is that the content is visible from the
+  // server, but the client renders a Suspense fallback, only to immediately show
+  // the data again.
+  test('should rehydrate synchronous thenable immediately', async () => {
+    // --- server ---
+
+    const serverQueryClient = new QueryClient({
+      defaultOptions: {
+        dehydrate: {
+          shouldDehydrateQuery: () => true,
+        },
+      },
+    })
+    const originalPromise = serverQueryClient.prefetchQuery({
+      queryKey: ['data'],
+      queryFn: () => null,
+    })
+
+    const dehydrated = dehydrate(serverQueryClient)
+
+    // --- server end ---
+
+    // Simulate a synchronous thenable
+    // @ts-expect-error
+    dehydrated.queries[0].promise.then = (cb) => {
+      cb?.('server data')
+    }
+
+    // --- client ---
+
+    const clientQueryClient = new QueryClient()
+    hydrate(clientQueryClient, dehydrated)
+
+    // If data is already resolved, it should end up in the cache immediately
+    expect(clientQueryClient.getQueryData(['data'])).toBe('server data')
+
+    // Need to await the original promise or else it will get a cancellation
+    // error and test will fail
+    await originalPromise
+
+    clientQueryClient.clear()
+    serverQueryClient.clear()
+  })
 })

packages/query-core/src/hydration.ts

@@ -1,3 +1,4 @@
+import { tryResolveSync } from './thenable'
 import type {
   DefaultError,
   MutationKey,
@@ -46,6 +47,10 @@
   state: QueryState
   promise?: Promise<unknown>
   meta?: QueryMeta
+  // 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.
+  dehydratedAt?: number
 }
 
 export interface DehydratedState {
@@ -74,6 +79,7 @@
   shouldRedactErrors: (error: unknown) => boolean,
 ): DehydratedQuery {
   return {
+    dehydratedAt: Date.now(),
     state: {
       ...query.state,
       ...(query.state.data !== undefined && {
@@ -189,52 +195,73 @@
     )
   })
 
-  queries.forEach(({ queryKey, state, queryHash, meta, promise }) => {
-    let query = queryCache.get(queryHash)
-
-    const data =
-      state.data === undefined ? state.data : deserializeData(state.data)
-
-    // Do not hydrate if an existing query exists with newer data
-    if (query) {
-      if (query.state.dataUpdatedAt < state.dataUpdatedAt) {
-        // omit fetchStatus from dehydrated state
-        // so that query stays in its current fetchStatus
-        const { fetchStatus: _ignored, ...serializedState } = state
-        query.setState({
-          ...serializedState,
-          data,
+  queries.forEach(
+    ({ queryKey, state, queryHash, meta, promise, dehydratedAt }) => {
+      const syncData = promise ? tryResolveSync(promise) : undefined
+      const rawData = state.data === undefined ? syncData?.data : state.data
+      const data = rawData === undefined ? rawData : deserializeData(rawData)
+
+      let query = queryCache.get(queryHash)
+      const existingQueryIsPending = query?.state.status === 'pending'
+
+      // Do not hydrate if an existing query exists with newer data
+      if (query) {
+        const hasNewerSyncData =
+          syncData &&
+          // We only need this undefined check to handle older dehydration
+          // payloads that might not have dehydratedAt
+          dehydratedAt !== undefined &&
+          dehydratedAt > query.state.dataUpdatedAt
+        if (
+          state.dataUpdatedAt > query.state.dataUpdatedAt ||
+          hasNewerSyncData
+        ) {
+          // omit fetchStatus from dehydrated state
+          // so that query stays in its current fetchStatus
+          const { fetchStatus: _ignored, ...serializedState } = state
+          query.setState({
+            ...serializedState,
+            data,
+          })
+        }
+      } else {
+        // Restore query
+        query = queryCache.build(
+          client,
+          {
+            ...client.getDefaultOptions().hydrate?.queries,
+            ...options?.defaultOptions?.queries,
+            queryKey,
+            queryHash,
+            meta,
+          },
+          // Reset fetch status to idle to avoid
+          // query being stuck in fetching state upon hydration
+          {
+            ...state,
+            data,
+            fetchStatus: 'idle',
+            status: data !== undefined ? 'success' : state.status,
+          },
+        )
+      }
+
+      if (
+        promise &&
+        !existingQueryIsPending &&
+        // Only hydrate if dehydration is newer than any existing data,
+        // this is always true for new queries
+        (dehydratedAt === undefined || dehydratedAt > query.state.dataUpdatedAt)
+      ) {
+        // This doesn't actually fetch - it just creates a retryer
+        // which will re-use the passed `initialPromise`
+        // Note that we need to call these even when data was synchronously
+        // available, as we still need to set up the retryer
+        void query.fetch(undefined, {
+          // RSC transformed promises are not thenable
+          initialPromise: Promise.resolve(promise).then(deserializeData),
         })
       }
-    } else {
-      // Restore query
-      query = queryCache.build(
-        client,
-        {
-          ...client.getDefaultOptions().hydrate?.queries,
-          ...options?.defaultOptions?.queries,
-          queryKey,
-          queryHash,
-          meta,
-        },
-        // Reset fetch status to idle to avoid
-        // query being stuck in fetching state upon hydration
-        {
-          ...state,
-          data,
-          fetchStatus: 'idle',
-        },
-      )
-    }
-
-    if (promise) {
-      // Note: `Promise.resolve` required cause
-      // RSC transformed promises are not thenable
-      const initialPromise = Promise.resolve(promise).then(deserializeData)
-
-      // this doesn't actually fetch - it just creates a retryer
-      // which will re-use the passed `initialPromise`
-      void query.fetch(undefined, { initialPromise })
-    }
-  })
+    },
+  )
 }

packages/query-core/src/thenable.ts

@@ -80,3 +80,30 @@
 
   return thenable
 }
+
+/**
+ * This function takes a Promise-like input and detects whether the data
+ * is synchronously available or not.
+ *
+ * It does not inspect .status, .value or .reason properties of the promise,
+ * as those are not always available, and the .status of React's promises
+ * should not be considered part of the public API.
+ */
+export function tryResolveSync(promise: Promise<unknown> | Thenable<unknown>) {
+  let data: unknown
+
+  promise
+    .then((result) => {
+      data = result
+      return result
+    })
+    // This can be unavailable on certain kinds of thenable's
+    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+    ?.catch(() => {})
+
+  if (data !== undefined) {
+    return { data }
+  }
+
+  return undefined
+}

packages/react-query/src/HydrationBoundary.tsx

@@ -24,13 +24,6 @@ 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 = {},
@@ -45,7 +38,7 @@ export const HydrationBoundary = ({
   const optionsRef = React.useRef(options)
   optionsRef.current = options
 
-  // This useMemo is for performance reasons only, everything inside it _must_
+  // This useMemo is for performance reasons only, everything inside it must
   // be safe to run in every render and code here should be read as "in render".
   //
   // This code needs to happen during the render phase, because after initial
@@ -80,10 +73,11 @@ export const HydrationBoundary = ({
         } else {
           const hydrationIsNewer =
             dehydratedQuery.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)
+              existingQuery.state.dataUpdatedAt ||
+            (dehydratedQuery.promise &&
+              existingQuery.state.status !== 'pending' &&
+              dehydratedQuery.dehydratedAt !== undefined &&
+              dehydratedQuery.dehydratedAt > existingQuery.state.dataUpdatedAt)
 
           const queryAlreadyQueued = hydrationQueue?.find(
             (query) => query.queryHash === dehydratedQuery.queryHash,
@@ -116,6 +110,7 @@ export const HydrationBoundary = ({
   React.useEffect(() => {
     if (hydrationQueue) {
       hydrate(client, { queries: hydrationQueue }, optionsRef.current)
+      // eslint-disable-next-line @eslint-react/hooks-extra/no-direct-set-state-in-use-effect
       setHydrationQueue(undefined)
     }
   }, [client, hydrationQueue])