fix: no hydration when new promise comes in

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
+}

integrations/react-next-15/app/client-component.tsx

@@ -8,10 +8,16 @@ 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()
+
+      console.log('client', count)
+
       return {
         text: 'data from client',
         date: Temporal.PlainDate.from('2023-01-01'),
+        count,
       }
     },
   })
@@ -26,7 +32,7 @@ export function ClientComponent() {
 
   return (
     <div>
-      {query.data.text} - {query.data.date.toJSON()}
+      {query.data.text} - {query.data.date.toJSON()} - {query.data.count}
     </div>
   )
 }

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 })
+}

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,26 @@ 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) => {
+          const shouldDehydrate =
+            defaultShouldDehydrateQuery(query) ||
+            query.state.status === 'pending'
+
+          console.log('shouldDehydrateQuery', query.queryKey, shouldDehydrate)
+
+          return shouldDehydrate
+        },
       },
     },
   })

integrations/react-next-15/app/page.tsx

@@ -2,29 +2,43 @@ 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({
     queryKey: ['data'],
     queryFn: async () => {
-      await sleep(2000)
+      const { count } = await (
+        await fetch('http://localhost:3000/count')
+      ).json()
+
+      console.log('server', count)
+
       return {
         text: 'data from server',
         date: Temporal.PlainDate.from('2024-01-01'),
+        count,
       }
     },
   })
 
+  const state = dehydrate(queryClient)
+
+  console.log('[page] state', state)
+
   return (
     <main>
-      <HydrationBoundary state={dehydrate(queryClient)}>
+      <HydrationBoundary state={state}>
         <ClientComponent />
       </HydrationBoundary>
+      <form action={queryExampleAction}>
+        <button type="submit">Increment</button>
+      </form>
     </main>
   )
 }

integrations/react-next-15/app/providers.tsx

@@ -1,16 +1,24 @@
 'use client'
-import { QueryClientProvider } from '@tanstack/react-query'
-import { ReactQueryDevtools } from '@tanstack/react-query-devtools'
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
+// import { ReactQueryDevtools } from '@tanstack/react-query-devtools'
 import * as React from 'react'
 import { makeQueryClient } from '@/app/make-query-client'
 
+let queryClientSingleton: QueryClient | undefined
+const getQueryClientSingleton = () => {
+  if (typeof window === 'undefined') {
+    return makeQueryClient()
+  }
+  return (queryClientSingleton ??= makeQueryClient())
+}
 export default function Providers({ children }: { children: React.ReactNode }) {
-  const [queryClient] = React.useState(() => makeQueryClient())
+  // const [queryClient] = React.useState(() => makeQueryClient())
+  const queryClient = getQueryClientSingleton()
 
   return (
     <QueryClientProvider client={queryClient}>
       {children}
-      <ReactQueryDevtools />
+      {/* <ReactQueryDevtools /> */}
     </QueryClientProvider>
   )
 }

packages/query-core/package.json

@@ -29,7 +29,8 @@
     "test:lib": "vitest",
     "test:lib:dev": "pnpm run test:lib --watch",
     "test:build": "publint --strict && attw --pack",
-    "build": "tsup"
+    "build": "tsup",
+    "build:dev": "tsup --watch"
   },
   "type": "module",
   "types": "build/legacy/index.d.ts",

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

@@ -1066,4 +1066,87 @@ 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)
+        console.log('queryFn', countRef.current)
+        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),
+    )
+
+    console.log('serialize mock', serializeDataMock.mock.calls)
+
+    expect(serializeDataMock).toHaveBeenCalledTimes(1)
+    expect(serializeDataMock).toHaveBeenCalledWith(0)
+
+    expect(deserializeDataMock).toHaveBeenCalledTimes(1)
+    expect(deserializeDataMock).toHaveBeenCalledWith(0)
+
+    // --- server ---
+    countRef.current++
+    const promise2 = serverQueryClient.prefetchQuery(query)
+
+    dehydrated = dehydrate(serverQueryClient)
+
+    // --- client ---
+
+    hydrate(clientQueryClient, dehydrated)
+
+    await promise2
+    await waitFor(() =>
+      expect(clientQueryClient.getQueryData(query.queryKey)).toBe(1),
+    )
+
+    console.log('serialize mock', serializeDataMock.mock.calls)
+
+    // Why are we getting 3 calls here? Should be 2?
+    // expect(serializeDataMock).toHaveBeenCalledTimes(2)
+    expect(serializeDataMock).toHaveBeenCalledTimes(3)
+    expect(serializeDataMock).toHaveBeenCalledWith(1)
+
+    // expect(deserializeDataMock).toHaveBeenCalledTimes(2)
+    expect(deserializeDataMock).toHaveBeenCalledTimes(3)
+    expect(deserializeDataMock).toHaveBeenCalledWith(1)
+
+    clientQueryClient.clear()
+    serverQueryClient.clear()
+  })
 })

packages/query-core/src/hydration.ts

@@ -71,16 +71,20 @@ function dehydrateQuery(
   query: Query,
   serializeData: TransformerFn,
 ): DehydratedQuery {
+  // console.log('[dehydrateQuery] query', query.queryKey)
   return {
     state: {
       ...query.state,
       ...(query.state.data !== undefined && {
         data: serializeData(query.state.data),
       }),
+      ...(query.state.fetchStatus === 'fetching' && {
+        promiseDehydratedAt: Date.now(),
+      }),
     },
     queryKey: query.queryKey,
     queryHash: query.queryHash,
-    ...(query.state.status === 'pending' && {
+    ...(query.state.fetchStatus === 'fetching' && {
       promise: query.promise?.then(serializeData).catch((error) => {
         if (process.env.NODE_ENV !== 'production') {
           console.error(
@@ -177,6 +181,8 @@ export function hydrate(
     const data =
       state.data === undefined ? state.data : deserializeData(state.data)
 
+    // console.log('[hydrate] query', JSON.stringify(query?.state, null, 2))
+
     // Do not hydrate if an existing query exists with newer data
     if (query) {
       if (query.state.dataUpdatedAt < state.dataUpdatedAt) {

packages/query-core/src/query.ts

@@ -56,6 +56,7 @@ export interface QueryState<TData = unknown, TError = DefaultError> {
   isInvalidated: boolean
   status: QueryStatus
   fetchStatus: FetchStatus
+  promiseDehydratedAt: number
 }
 
 export interface FetchContext<
@@ -679,5 +680,6 @@ function getDefaultState<
     isInvalidated: false,
     status: hasData ? 'success' : 'pending',
     fetchStatus: 'idle',
+    promiseDehydratedAt: 0,
   }
 }

packages/react-query/package.json

@@ -31,6 +31,7 @@
     "test:build": "publint --strict && attw --pack",
     "build": "pnpm build:tsup && pnpm build:codemods",
     "build:tsup": "tsup",
+    "build:tsup:dev": "tsup --watch",
     "build:codemods": "cpy ../query-codemods/* ./build/codemods"
   },
   "type": "module",

packages/react-query/src/HydrationBoundary.tsx

@@ -35,6 +35,11 @@ export const HydrationBoundary = ({
     DehydratedState['queries'] | undefined
   >()
 
+  // console.log(
+  //   '[HydrationBoundary] rendering with state',
+  //   JSON.stringify(state, null, 2),
+  // )
+
   const optionsRef = React.useRef(options)
   optionsRef.current = options
 
@@ -68,12 +73,33 @@ export const HydrationBoundary = ({
       for (const dehydratedQuery of queries) {
         const existingQuery = queryCache.get(dehydratedQuery.queryHash)
 
+        // console.log(
+        //   '[HydrationBoundary] existingQuery',
+        //   JSON.stringify(existingQuery?.state, null, 2),
+        // )
+        // console.log(
+        //   '[HydrationBoundary] dehydratedQuery',
+        //   JSON.stringify(dehydratedQuery.state, null, 2),
+        // )
         if (!existingQuery) {
           newQueries.push(dehydratedQuery)
         } else {
           const hydrationIsNewer =
             dehydratedQuery.state.dataUpdatedAt >
-            existingQuery.state.dataUpdatedAt
+              existingQuery.state.dataUpdatedAt ||
+            // @ts-expect-error
+            dehydratedQuery.promise?.status !== existingQuery.promise?.status
+
+          console.log(
+            '[HydrationBoundary] hydrationIsNewer',
+            dehydratedQuery.queryKey,
+            hydrationIsNewer,
+            dehydratedQuery.state.promiseDehydratedAt,
+            existingQuery.state.promiseDehydratedAt,
+            // @ts-expect-error
+            dehydratedQuery.promise?.status,
+          )
+
           const queryAlreadyQueued = hydrationQueue?.find(
             (query) => query.queryHash === dehydratedQuery.queryHash,
           )
@@ -102,6 +128,13 @@ export const HydrationBoundary = ({
     }
   }, [client, hydrationQueue, state])
 
+  React.useEffect(() => {
+    console.log(
+      'hydrationQueue changed',
+      JSON.stringify(hydrationQueue, null, 2),
+    )
+  }, [hydrationQueue])
+
   React.useEffect(() => {
     if (hydrationQueue) {
       hydrate(client, { queries: hydrationQueue }, optionsRef.current)