feat(react-query): support server snapshots during hydration

- Add optional `serverSnapshot` support to `QueryClientProvider`
- Prevent hydration mismatches when streamed queries resolve early

Author: DogPawHat

.changeset/tidy-poems-guess.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/react-query': minor
+---
+
+Add an optional `serverSnapshot` prop to `QueryClientProvider`. When supplied with the same `DehydratedState` used to hydrate the client, `useQuery` and the other hooks built on `useBaseQuery` replay the frozen server-rendered result during hydration (through `useSyncExternalStore`'s server snapshot) before switching to the live cache. This prevents hydration mismatches when a streamed query promise resolves before the browser hydrates, which previously made the first client render differ from the server markup.

packages/react-query/src/QueryClientProvider.tsx

@@ -1,7 +1,8 @@
 'use client'
 import * as React from 'react'
 
-import type { QueryClient } from '@tanstack/query-core'
+import { QueryCache, QueryClient } from '@tanstack/query-core'
+import type { DehydratedState } from '@tanstack/query-core'
 
 /**
  * The context that `useQueryClient` reads from. `QueryClientProvider` is the normal way to set it.
@@ -10,6 +11,16 @@ export const QueryClientContext = React.createContext<QueryClient | undefined>(
   undefined,
 )
 
+/**
+ * Internal context that carries the frozen server snapshot used during hydration. The value is a
+ * throwaway `QueryClient` whose cache holds the exact query state that produced the server markup,
+ * so hooks can replay it instead of reading newer live cache data. `undefined` outside of SSR
+ * hydration (i.e. on the server without a provided snapshot, or after hydration).
+ */
+export const QueryServerSnapshotContext = React.createContext<
+  QueryClient | undefined
+>(undefined)
+
 /**
  * The `useQueryClient` hook returns the current `QueryClient` instance.
  *
@@ -42,6 +53,17 @@ export type QueryClientProviderProps = {
    * The `QueryClient` instance to provide.
    */
   client: QueryClient
+  /**
+   * Optional frozen snapshot of the query state that produced the server-rendered markup. When
+   * provided, hooks replay this state during hydration (via `useSyncExternalStore`'s server
+   * snapshot) before switching to the live cache, so the first client render matches the server
+   * output even if the live cache already advanced (e.g. a streamed promise resolved before
+   * hydration).
+   *
+   * Typically this is the same `DehydratedState` that was passed to `hydrate`/`HydrationBoundary`.
+   * Passing the server state here mirrors React Redux's `Provider serverState` API.
+   */
+  serverSnapshot?: DehydratedState | null
   /**
    * The components that get access to the provided `QueryClient`.
    */
@@ -70,6 +92,7 @@ export type QueryClientProviderProps = {
 export const QueryClientProvider = ({
   client,
   children,
+  serverSnapshot,
 }: QueryClientProviderProps): React.JSX.Element => {
   React.useEffect(() => {
     client.mount()
@@ -78,9 +101,34 @@ export const QueryClientProvider = ({
     }
   }, [client])
 
+  // A throwaway client whose cache holds the server-rendered query state. We build queries directly
+  // from the dehydrated state (rather than calling `hydrate`) so the frozen `fetchStatus` is
+  // preserved and hooks replay exactly what the server rendered.
+  const snapshotClient = React.useMemo(() => {
+    if (!serverSnapshot) {
+      return undefined
+    }
+
+    const queryCache = new QueryCache()
+    const frozenClient = new QueryClient({ queryCache })
+
+    serverSnapshot.queries.forEach(({ queryKey, queryHash, state, meta }) => {
+      queryCache.build(
+        frozenClient,
+        { queryKey, queryHash, meta },
+        // Build from a copy so the caller's dehydrated state is never mutated.
+        { ...state },
+      )
+    })
+
+    return frozenClient
+  }, [serverSnapshot])
+
   return (
     <QueryClientContext.Provider value={client}>
-      {children}
+      <QueryServerSnapshotContext.Provider value={snapshotClient}>
+        {children}
+      </QueryServerSnapshotContext.Provider>
     </QueryClientContext.Provider>
   )
 }

packages/react-query/src/__tests__/ssr-hydration.test.tsx

@@ -14,10 +14,14 @@ import {
 } from '..'
 import { setIsServer } from './utils'
 
-const ReactHydrate = (element: React.ReactElement, container: Element) => {
+const ReactHydrate = (
+  element: React.ReactElement,
+  container: Element,
+  options?: { onRecoverableError?: (error: unknown) => void },
+) => {
   let root: any
   act(() => {
-    root = hydrateRoot(container, element)
+    root = hydrateRoot(container, element, options)
   })
   return () => {
     root.unmount()
@@ -275,4 +279,198 @@ describe('Server side rendering with de/rehydration', () => {
     queryClient.clear()
     consoleMock.mockRestore()
   })
+
+  it('should not mismatch when a pending prefetched query resolves before client hydration', async () => {
+    const key = queryKey()
+
+    let resolveQuery!: (value: string) => void
+    const pendingQuery = new Promise<string>((resolve) => {
+      resolveQuery = resolve
+    })
+    const queryFn = vi.fn(() => pendingQuery)
+    const renderedStates: Array<string> = []
+
+    // -- Shared part --
+    function SuccessComponent() {
+      const result = useQuery({
+        queryKey: key,
+        queryFn,
+      })
+      const rendered = `SuccessComponent - status:${result.status} fetching:${result.isFetching} data:${result.data}`
+      renderedStates.push(rendered)
+      return rendered
+    }
+
+    // -- Server part --
+    setIsServer(true)
+
+    const prefetchClient = new QueryClient()
+    // Prefetch without awaiting: the query is dehydrated while still pending.
+    prefetchClient.prefetchQuery({ queryKey: key, queryFn }).catch(noop)
+    // Let the retryer start so the pending promise is part of the dehydrated state.
+    await vi.advanceTimersByTimeAsync(1)
+    const dehydratedStateServer = dehydrate(prefetchClient, {
+      shouldDehydrateQuery: () => true,
+    })
+    expect(dehydratedStateServer.queries[0]?.promise).toBeDefined()
+
+    const renderCache = new QueryCache()
+    const renderClient = new QueryClient({ queryCache: renderCache })
+    hydrate(renderClient, dehydratedStateServer)
+    const markup = ReactDOMServer.renderToString(
+      <QueryClientProvider client={renderClient}>
+        <SuccessComponent />
+      </QueryClientProvider>,
+    )
+    renderClient.clear()
+    setIsServer(false)
+
+    const expectedMarkup =
+      'SuccessComponent - status:pending fetching:true data:undefined'
+
+    expect(markup).toBe(expectedMarkup)
+
+    // -- Client part --
+    renderedStates.length = 0
+    const el = document.createElement('div')
+    el.innerHTML = markup
+
+    const queryCache = new QueryCache()
+    const queryClient = new QueryClient({ queryCache })
+    // Pass the dehydrated promise through (not JSON-serializable) to mimic a
+    // framework streaming the pending query's promise to the browser.
+    hydrate(queryClient, dehydratedStateServer)
+
+    // The streamed promise resolves before React hydrates, so the live cache is
+    // already successful while the server markup shows the pending state.
+    resolveQuery('success!')
+    await vi.advanceTimersByTimeAsync(1)
+
+    expect(queryClient.getQueryData(key)).toBe('success!')
+
+    const onRecoverableError = vi.fn()
+    const unmount = ReactHydrate(
+      <QueryClientProvider
+        client={queryClient}
+        serverSnapshot={dehydratedStateServer}
+      >
+        <SuccessComponent />
+      </QueryClientProvider>,
+      el,
+      { onRecoverableError },
+    )
+
+    // Hydration must not report a mismatch, and the first client render must
+    // replay the frozen server snapshot (pending) even though the live cache is
+    // already successful.
+    expect(onRecoverableError).toHaveBeenCalledTimes(0)
+    expect(renderedStates[0]).toBe(expectedMarkup)
+
+    await vi.advanceTimersByTimeAsync(50)
+    expect(renderedStates.at(-1)).toBe(
+      'SuccessComponent - status:success fetching:false data:success!',
+    )
+    expect(el.innerHTML).toBe(
+      'SuccessComponent - status:success fetching:false data:success!',
+    )
+
+    unmount()
+    queryClient.clear()
+  })
+
+  // Adapted from the reproduction in
+  // https://github.com/TanStack/query/issues/9399#issuecomment-4323008704 — the streamed promise
+  // is simulated with a synchronously-resolvable thenable, so the client's `hydrate` resolves it
+  // via `tryResolveSync` during the first render.
+  it('should not mismatch on a query whose streamed promise is synchronously resolved by hydrate', async () => {
+    const key = queryKey()
+    const renderedStates: Array<string> = []
+
+    function SuccessComponent() {
+      const result = useQuery({
+        queryKey: key,
+        queryFn: () => Promise.resolve('success!'),
+      })
+      const rendered = `SuccessComponent - status:${result.status} fetching:${result.isFetching} data:${result.data}`
+      renderedStates.push(rendered)
+      return rendered
+    }
+
+    // -- Server --
+    setIsServer(true)
+    const prefetchClient = new QueryClient({
+      defaultOptions: { dehydrate: { shouldDehydrateQuery: () => true } },
+    })
+    let resolvePrefetch: ((value: string) => void) | undefined
+    const prefetchPromise = new Promise<string>((resolve) => {
+      resolvePrefetch = resolve
+    })
+    void prefetchClient.prefetchQuery({
+      queryKey: key,
+      queryFn: () => prefetchPromise,
+    })
+
+    const dehydrated = dehydrate(prefetchClient)
+    expect(dehydrated.queries[0]?.state.status).toBe('pending')
+
+    const renderClient = new QueryClient()
+    hydrate(renderClient, dehydrated)
+    const markup = ReactDOMServer.renderToString(
+      <QueryClientProvider client={renderClient}>
+        <SuccessComponent />
+      </QueryClientProvider>,
+    )
+    renderClient.clear()
+    setIsServer(false)
+
+    const expectedMarkup =
+      'SuccessComponent - status:pending fetching:true data:undefined'
+    expect(markup).toBe(expectedMarkup)
+
+    // The promise resolves *between* SSR and client hydration (streamed value arrives).
+    resolvePrefetch?.('success!')
+    const promiseRef = dehydrated.queries[0]?.promise
+    if (promiseRef) {
+      // Synchronously-resolvable thenable, mirroring a streamed React promise.
+      // @ts-expect-error deliberately replacing the native `then` so it resolves synchronously
+      promiseRef.then = (cb?: (value: unknown) => unknown) => {
+        cb?.('success!')
+        return promiseRef
+      }
+    }
+
+    // -- Client --
+    renderedStates.length = 0
+    const el = document.createElement('div')
+    el.innerHTML = markup
+    const queryClient = new QueryClient()
+    hydrate(queryClient, dehydrated)
+
+    expect(queryClient.getQueryData(key)).toBe('success!')
+
+    const onRecoverableError = vi.fn()
+    const unmount = ReactHydrate(
+      <QueryClientProvider client={queryClient} serverSnapshot={dehydrated}>
+        <SuccessComponent />
+      </QueryClientProvider>,
+      el,
+      { onRecoverableError },
+    )
+
+    // No mismatch, and the first client render replays the pending server snapshot even though
+    // `hydrate` resolved the streamed promise synchronously.
+    expect(onRecoverableError).toHaveBeenCalledTimes(0)
+    expect(renderedStates[0]).toBe(expectedMarkup)
+
+    await vi.advanceTimersByTimeAsync(50)
+    expect(renderedStates.at(-1)).toBe(
+      'SuccessComponent - status:success fetching:false data:success!',
+    )
+    expect(el.innerHTML).toBe(
+      'SuccessComponent - status:success fetching:false data:success!',
+    )
+
+    unmount()
+    queryClient.clear()
+  })
 })