@tanstack/react-db 0.4.1 → 0.5.1

package/src/useLiveQuery.ts

@@ -3,6 +3,7 @@
 import { useRef, useSyncExternalStore } from 'react'
 import {
   BaseQueryBuilder,
+  IR,
   UnhashableQueryIRError,
   createLiveQueryCollection,
   createLiveQueryObserver,
@@ -25,12 +26,16 @@ import type {
   InitialQueryBuilder,
   LiveQueryCollectionConfig,
   LiveQueryObserver,
+  LiveQueryPersistedStatus,
   NonSingleResult,
   QueryBuilder,
   SingleResult,
 } from '@tanstack/db'
 
 const DEFAULT_GC_TIME_MS = 1 // Live queries created by useLiveQuery are cleaned up immediately (0 disables GC)
+// Suspense renders can be abandoned before React subscribes. Keep their
+// collection briefly so a nearby retry can commit without starting over.
+const DEFAULT_SUSPENSE_GC_TIME_MS = 5000
 const DERIVED_IDENTITY_SINGLE_RENDER_WARN_MS = 16
 const DERIVED_IDENTITY_RENDER_COUNT_WARN_THRESHOLD = 10
 const DERIVED_IDENTITY_TOTAL_WARN_MS = 50
@@ -39,6 +44,72 @@ const warnedDerivedIdentityCallsites = new Set<string>()
 const warnedUnhashableIdentityCallsites = new Set<string>()
 const unpreparedQueryValue = Symbol(`unpreparedQueryValue`)
 
+type SuspenseCollection = Collection<object, string | number, {}>
+type SuspenseCollectionEntry = {
+  collection: SuspenseCollection
+  removeStatusListeners: () => void
+}
+
+const unscopedSuspenseCollections = new Map<string, SuspenseCollectionEntry>()
+const suspenseCollectionsByClient = new WeakMap<
+  DbClient,
+  Map<string, SuspenseCollectionEntry>
+>()
+const suspenseSourceIds = new WeakMap<object, number>()
+let nextSuspenseSourceId = 0
+
+function getSuspenseSourceId(source: object): number {
+  let id = suspenseSourceIds.get(source)
+  if (id === undefined) {
+    id = ++nextSuspenseSourceId
+    suspenseSourceIds.set(source, id)
+  }
+  return id
+}
+
+function getUnscopedSuspenseKey(
+  preparedValue: unknown,
+  queryHash: string,
+): string {
+  const query =
+    preparedValue instanceof BaseQueryBuilder
+      ? preparedValue
+      : preparedValue &&
+          typeof preparedValue === `object` &&
+          `query` in preparedValue &&
+          preparedValue.query instanceof BaseQueryBuilder
+        ? preparedValue.query
+        : undefined
+  if (!query) return queryHash
+  const sourceIds = IR.collectCollectionSources(query._getQuery()).map(
+    ({ collection }) => getSuspenseSourceId(collection),
+  )
+  return `${sourceIds.join(`,`)}:${queryHash}`
+}
+
+function getSuspenseCollections(
+  client: DbClient | undefined,
+): Map<string, SuspenseCollectionEntry> {
+  if (!client) return unscopedSuspenseCollections
+  let collections = suspenseCollectionsByClient.get(client)
+  if (!collections) {
+    collections = new Map()
+    suspenseCollectionsByClient.set(client, collections)
+  }
+  return collections
+}
+
+function releaseSuspenseCollection(
+  collections: Map<string, SuspenseCollectionEntry>,
+  queryHash: string,
+  collection: SuspenseCollection,
+): void {
+  const entry = collections.get(queryHash)
+  if (entry?.collection !== collection) return
+  collections.delete(queryHash)
+  entry.removeStatusListeners()
+}
+
 export type DerivedIdentityProfiler = {
   renderCount: number
   totalMs: number
@@ -120,7 +191,16 @@ function getCurrentTime(): number {
 
 function getWarningCallsite(stackIndex: number): string {
   const stack = new Error().stack ?? `unknown`
-  return stack.split(`\n`)[stackIndex]?.trim() ?? stack
+  const lines = stack.split(`\n`)
+  const userFrame = lines
+    .slice(1)
+    .find(
+      (line) =>
+        !line.includes(`useLiveQuery.ts`) &&
+        !line.includes(`useLiveSuspenseQuery.ts`) &&
+        !line.includes(`useLiveInfiniteQuery.ts`),
+    )
+  return userFrame?.trim() ?? lines[stackIndex]?.trim() ?? stack
 }
 
 function warnDerivedIdentityHotPath(
@@ -256,7 +336,10 @@ export function warnUnhashableDerivedIdentity(
   )
 }
 
-function createCollectionFromPreparedQuery(value: unknown) {
+function createCollectionFromPreparedQuery(
+  value: unknown,
+  defaultGcTime = DEFAULT_GC_TIME_MS,
+) {
   if (value === undefined || value === null) {
     return null
   }
@@ -270,14 +353,14 @@ function createCollectionFromPreparedQuery(value: unknown) {
     return createLiveQueryCollection({
       query: value,
       startSync: true,
-      gcTime: DEFAULT_GC_TIME_MS,
+      gcTime: defaultGcTime,
     })
   }
 
   if (typeof value === `object`) {
     return createLiveQueryCollection({
       startSync: true,
-      gcTime: DEFAULT_GC_TIME_MS,
+      gcTime: defaultGcTime,
       ...(value as LiveQueryCollectionConfig<any>),
     })
   }
@@ -367,6 +450,9 @@ export function useLiveQuery<TContext extends Context>(
   status: CollectionStatus // Can't be disabled if always returns QueryBuilder
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -386,6 +472,9 @@ export function useLiveQuery<TContext extends Context>(
   status: UseLiveQueryStatus
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -405,6 +494,9 @@ export function useLiveQuery<TContext extends Context>(
   status: UseLiveQueryStatus
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -428,6 +520,9 @@ export function useLiveQuery<
   status: UseLiveQueryStatus
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -463,6 +558,9 @@ export function useLiveQuery<
   status: UseLiveQueryStatus
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -514,6 +612,9 @@ export function useLiveQuery<TContext extends Context>(
   status: CollectionStatus // Can't be disabled when query always returns a builder
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -530,6 +631,9 @@ export function useLiveQuery<TContext extends Context>(
   status: UseLiveQueryStatus
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -547,6 +651,9 @@ export function useLiveQuery<TContext extends Context>(
   status: CollectionStatus // Can't be disabled when query always returns a builder
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -564,6 +671,9 @@ export function useLiveQuery<TContext extends Context>(
   status: UseLiveQueryStatus
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -613,6 +723,9 @@ export function useLiveQuery<
   status: CollectionStatus // Can't be disabled for pre-created live query collections
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -633,6 +746,9 @@ export function useLiveQuery<
   status: CollectionStatus // Can't be disabled for pre-created live query collections
   isLoading: boolean
   isReady: boolean
+  persistedStatus: LiveQueryPersistedStatus
+  isPersistedReady: boolean
+  persistedError: unknown | undefined
   isIdle: boolean
   isError: boolean
   isCleanedUp: boolean
@@ -643,6 +759,22 @@ export function useLiveQuery<
 export function useLiveQuery(
   configOrQueryOrCollection: any,
   deps?: Array<unknown>,
+) {
+  return useLiveQueryImpl(configOrQueryOrCollection, deps)
+}
+
+/** @internal Shared implementation for the Suspense wrapper. */
+export function useLiveQueryForSuspense(
+  configOrQueryOrCollection: any,
+  deps: Array<unknown> | undefined,
+) {
+  return useLiveQueryImpl(configOrQueryOrCollection, deps, true)
+}
+
+function useLiveQueryImpl(
+  configOrQueryOrCollection: any,
+  deps: Array<unknown> | undefined,
+  forSuspense = false,
 ) {
   const contextDbClient = useOptionalDbClient()
   // Check if it's already a collection
@@ -676,6 +808,7 @@ export function useLiveQuery(
     null,
   )
   const queryHashRef = useRef<string | undefined>(undefined)
+  const suspenseKeyRef = useRef<string | undefined>(undefined)
   const identityErrorRef = useRef<UnhashableQueryIRError | undefined>(undefined)
 
   const queryKey = !inputIsCollection
@@ -747,6 +880,47 @@ export function useLiveQuery(
     warnDeprecatedDepsArray()
   }
 
+  const canReuseSuspenseKey =
+    forSuspense &&
+    !inputIsCollection &&
+    queryHash !== undefined &&
+    !dbClient &&
+    collectionRef.current !== null &&
+    clientRef.current === dbClient &&
+    queryHashRef.current === queryHash &&
+    suspenseKeyRef.current !== undefined
+
+  if (
+    forSuspense &&
+    !inputIsCollection &&
+    queryHash &&
+    !dbClient &&
+    !canReuseSuspenseKey &&
+    preparedQueryValue === unpreparedQueryValue
+  ) {
+    preparedQueryValue = prepareQueryValue(
+      configOrQueryOrCollection,
+      dbClient,
+      deferredCollectionsRef.current,
+    )
+  }
+
+  const suspenseKey =
+    queryHash && !dbClient
+      ? canReuseSuspenseKey
+        ? suspenseKeyRef.current
+        : getUnscopedSuspenseKey(preparedQueryValue, queryHash)
+      : queryHash
+
+  const suspenseCollections =
+    forSuspense && !inputIsCollection && suspenseKey
+      ? getSuspenseCollections(dbClient)
+      : undefined
+  const suspenseEntry = suspenseKey
+    ? suspenseCollections?.get(suspenseKey)
+    : undefined
+  const suspenseCollection = suspenseEntry?.collection
+
   const identityChanged =
     depsRef.current === null ||
     (deps !== undefined
@@ -792,21 +966,59 @@ export function useLiveQuery(
       collectionRef.current = configOrQueryOrCollection
       configRef.current = configOrQueryOrCollection
     } else {
-      if (preparedQueryValue === unpreparedQueryValue) {
-        preparedQueryValue = prepareQueryValue(
-          configOrQueryOrCollection,
-          dbClient,
-          deferredCollectionsRef.current,
-        )
+      if (suspenseCollection) {
+        collectionRef.current = suspenseCollection
+      } else {
+        if (preparedQueryValue === unpreparedQueryValue) {
+          preparedQueryValue = prepareQueryValue(
+            configOrQueryOrCollection,
+            dbClient,
+            deferredCollectionsRef.current,
+          )
+        }
+        collectionRef.current = createCollectionFromPreparedQuery(
+          preparedQueryValue,
+          forSuspense ? DEFAULT_SUSPENSE_GC_TIME_MS : DEFAULT_GC_TIME_MS,
+        ) as SuspenseCollection | null
+        if (suspenseCollections && suspenseKey && collectionRef.current) {
+          const collection = collectionRef.current
+          const removeCleanupListener = collection.on(`status:cleaned-up`, () =>
+            releaseSuspenseCollection(
+              suspenseCollections,
+              suspenseKey,
+              collection,
+            ),
+          )
+          const removeErrorListener = collection.on(`status:error`, () => {
+            // Keep the failed collection through React's immediate retry so
+            // the hook can throw its actual load error to an ErrorBoundary.
+            // Retire it afterward so a later boundary reset starts fresh.
+            setTimeout(
+              () =>
+                releaseSuspenseCollection(
+                  suspenseCollections,
+                  suspenseKey,
+                  collection,
+                ),
+              0,
+            )
+          })
+          const entry: SuspenseCollectionEntry = {
+            collection,
+            removeStatusListeners: () => {
+              removeCleanupListener()
+              removeErrorListener()
+            },
+          }
+          suspenseCollections.set(suspenseKey, entry)
+        }
       }
-      collectionRef.current = createCollectionFromPreparedQuery(
-        preparedQueryValue,
-      ) as Collection<object, string | number, {}>
       configRef.current = configOrQueryOrCollection
       depsRef.current = [...identityDeps]
     }
     clientRef.current = dbClient
     queryHashRef.current = queryHash
+    suspenseKeyRef.current = suspenseKey
     identityErrorRef.current = identityError
   }
 

package/src/useLiveSuspenseQuery.ts

@@ -1,12 +1,13 @@
 'use client'
 
 import { useRef } from 'react'
-import { useLiveQuery } from './useLiveQuery'
+import { useLiveQueryForSuspense } from './useLiveQuery'
 import { getLiveQueryResultInfo } from './live-query-internals'
 import type { UseLiveQueryConfig } from './useLiveQuery'
 import type {
   Collection,
   Context,
+  DbClient,
   GetResult,
   InferResultType,
   InitialQueryBuilder,
@@ -16,6 +17,46 @@ import type {
   SingleResult,
 } from '@tanstack/db'
 
+// React can discard a render that suspends, including its refs. Keep the
+// initial-render failure across retries, scoped by client because streamed
+// preloads of the same collection can fail independently.
+const initialRenderErrors = new WeakMap<
+  Collection<any, any, any>,
+  {
+    byClient: WeakMap<DbClient, { error: unknown; clientQuery?: object }>
+    unscoped?: { error: unknown; clientQuery?: object }
+  }
+>()
+
+function clearInitialRenderError(
+  collection: Collection<any, any, any>,
+  client: DbClient | undefined,
+) {
+  const entry = initialRenderErrors.get(collection)
+  if (!entry) return
+  if (client) entry.byClient.delete(client)
+  else delete entry.unscoped
+}
+
+function rememberInitialRenderError(
+  collection: Collection<any, any, any>,
+  client: DbClient | undefined,
+  error: unknown,
+  clientQuery?: object,
+) {
+  if (collection.status === `cleaned-up`) return
+  let entry = initialRenderErrors.get(collection)
+  if (!entry) {
+    entry = { byClient: new WeakMap() }
+    collection.once(`status:cleaned-up`, () => {
+      initialRenderErrors.delete(collection)
+    })
+    initialRenderErrors.set(collection, entry)
+  }
+  if (client) entry.byClient.set(client, { error, clientQuery })
+  else entry.unscoped = { error }
+}
+
 /**
  * Create a live query with React Suspense support
  * @param queryFn - Query function that defines what data to fetch
@@ -173,19 +214,9 @@ export function useLiveSuspenseQuery(
   // Use useLiveQuery to handle collection management and reactivity
   const result =
     deps === undefined
-      ? useLiveQuery(configOrQueryOrCollection)
-      : useLiveQuery(configOrQueryOrCollection, deps)
-  const queryInfo = getLiveQueryResultInfo(result)
-
-  // Reset promise and ready state when query identity changes
-  if (collectionRef.current !== result.collection) {
-    promiseRef.current = null
-    collectionRef.current = result.collection
-    hasBeenReadyRef.current = false
-  }
+      ? useLiveQueryForSuspense(configOrQueryOrCollection, undefined)
+      : useLiveQueryForSuspense(configOrQueryOrCollection, deps)
 
-  // SUSPENSE LOGIC: Throw promise or error based on collection status
-  // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
   if (!result.isEnabled) {
     // Suspense queries cannot be disabled - this matches TanStack Query's useSuspenseQuery behavior
     throw new Error(
@@ -197,30 +228,82 @@ export function useLiveSuspenseQuery(
     )
   }
 
+  const queryInfo = getLiveQueryResultInfo(result)
+
+  // Reset promise and ready state when query identity changes
+  if (collectionRef.current !== result.collection) {
+    promiseRef.current = null
+    collectionRef.current = result.collection
+    hasBeenReadyRef.current = false
+  }
+
+  // SUSPENSE LOGIC: Throw promise or error based on collection status
+
   const collectionStatus = result.collection.status
 
   // Track when we reach ready state
-  if (result.isReady) {
+  if (result.isReady || queryInfo.observer.isInitialRenderReady()) {
     hasBeenReadyRef.current = true
     promiseRef.current = null
+    clearInitialRenderError(result.collection, queryInfo.client)
   }
 
   const observerError = queryInfo.observer.getError()
-  if (observerError !== undefined && !hasBeenReadyRef.current) {
+  // A client request can reject with undefined, which getError() cannot
+  // distinguish from no error. The request status preserves that distinction.
+  const clientQuery =
+    queryInfo.client && queryInfo.queryHash
+      ? queryInfo.client._getLiveQuery(queryInfo.queryHash)
+      : undefined
+  // A configured persisted restore may still satisfy this render after the
+  // client query fails. Let the observer classify that failure below.
+  if (
+    !hasBeenReadyRef.current &&
+    (result.persistedStatus === `unavailable` ||
+      result.persistedStatus === `error`) &&
+    (observerError !== undefined || clientQuery?.status === `error`)
+  ) {
     promiseRef.current = null
-    throw observerError
+    throw observerError === undefined ? clientQuery?.error : observerError
+  }
+
+  const errorEntry = initialRenderErrors.get(result.collection)
+  const initialRenderError = queryInfo.client
+    ? errorEntry?.byClient.get(queryInfo.client)
+    : errorEntry?.unscoped
+  // A replacement client query must not inherit the prior request's error.
+  if (
+    initialRenderError?.clientQuery !== undefined &&
+    initialRenderError.clientQuery !== clientQuery
+  ) {
+    clearInitialRenderError(result.collection, queryInfo.client)
+  } else if (initialRenderError && !hasBeenReadyRef.current) {
+    promiseRef.current = null
+    throw initialRenderError.error
   }
 
   // Only throw errors during initial load (before first ready)
   // After success, errors surface as stale data (matches TanStack Query behavior)
-  if (collectionStatus === `error` && !hasBeenReadyRef.current) {
+  if (
+    collectionStatus === `error` &&
+    !hasBeenReadyRef.current &&
+    (result.persistedStatus === `unavailable` ||
+      result.persistedStatus === `error`)
+  ) {
     promiseRef.current = null
     // TODO: Once collections hold a reference to their last error object (#671),
     // we should rethrow that actual error instead of creating a generic message
     throw new Error(`Collection "${result.collection.id}" failed to load`)
   }
 
-  if (!hasBeenReadyRef.current && (result.isLoading || result.isIdle)) {
+  if (
+    !hasBeenReadyRef.current &&
+    (result.isLoading ||
+      result.isIdle ||
+      result.isError ||
+      (collectionStatus === `error` &&
+        result.persistedStatus !== `unavailable`))
+  ) {
     if (queryInfo.client?._isSsrStreamingEnabled() && !queryInfo.queryHash) {
       const reason = queryInfo.identityError
         ? `${queryInfo.identityError.reason} at ${queryInfo.identityError.path}`
@@ -229,14 +312,40 @@ export function useLiveSuspenseQuery(
         `Cannot stream this live query during SSR because ${reason}. Provide an explicit serializable queryKey.`,
       )
     }
-
     // Create or reuse promise for current collection
     if (!promiseRef.current) {
-      promiseRef.current = queryInfo.observer.preload()
+      const collection = result.collection
+      const client = queryInfo.client
+      const queryHash = queryInfo.queryHash
+      let active = true
+      const stopWatchingCleanup = collection.once(`status:cleaned-up`, () => {
+        active = false
+      })
+      let preload: Promise<void>
+      try {
+        preload = queryInfo.observer.preloadForInitialRender()
+      } catch (error) {
+        stopWatchingCleanup()
+        throw error
+      }
+      const preloadClientQuery =
+        client && queryHash ? client._getLiveQuery(queryHash) : undefined
+      promiseRef.current = preload
+        .catch((error: unknown) => {
+          const latestClientQuery =
+            client && queryHash ? client._getLiveQuery(queryHash) : undefined
+          if (active && latestClientQuery === preloadClientQuery)
+            rememberInitialRenderError(
+              collection,
+              client,
+              error,
+              preloadClientQuery,
+            )
+          throw error
+        })
+        .finally(stopWatchingCleanup)
     }
-    // THROW PROMISE - React Suspense catches this (React 18+ required)
-    // Note: We don't check React version here. In React <18, this will be caught
-    // by an Error Boundary, which provides a reasonable failure mode.
+    // React Suspense catches this promise and retries after preload settles.
     throw promiseRef.current
   }