ultisuite-client/lib/api/query-provider.tsx

"use client"

import { useState } from "react"
import { QueryClient } from "@tanstack/react-query"
import {
  PersistQueryClientProvider,
  type PersistedClient,
} from "@tanstack/react-query-persist-client"
import { openDB, type IDBPDatabase } from "idb"
import type { Persister } from "@tanstack/react-query-persist-client"
import {
  isPreviewThumbQueryKey,
  revokePreviewBlobData,
} from "@/lib/api/preview-blob-url"
import { ApiRequestError } from "@/lib/api/client"

const DB_NAME = "ultimail-query-cache"
const STORE_NAME = "query-cache"

let dbPromise: Promise<IDBPDatabase> | null = null

function getDb() {
  if (!dbPromise) {
    dbPromise = openDB(DB_NAME, 1, {
      upgrade(db) {
        db.createObjectStore(STORE_NAME)
      },
    })
  }
  return dbPromise
}

const idbPersister: Persister = {
  persistClient: async (client: PersistedClient) => {
    const db = await getDb()
    await db.put(STORE_NAME, client, "cache")
  },
  restoreClient: async (): Promise<PersistedClient | undefined> => {
    const db = await getDb()
    const restored = await db.get<PersistedClient>(STORE_NAME, "cache")
    if (!restored?.clientState?.queries) return restored
    restored.clientState.queries = restored.clientState.queries.filter(
      (entry) => !isPreviewThumbQueryKey(entry.queryKey)
    )
    return restored
  },
  removeClient: async () => {
    const db = await getDb()
    await db.delete(STORE_NAME, "cache")
  },
}

function attachPreviewBlobGc(queryClient: QueryClient) {
  return queryClient.getQueryCache().subscribe((event) => {
    if (event.type === "removed") {
      revokePreviewBlobData(event.query.state.data)
    }
  })
}

function makeQueryClient() {
  return new QueryClient({
    defaultOptions: {
      queries: {
        gcTime: 1000 * 60 * 60 * 24,
        staleTime: 1000 * 60 * 5,
        networkMode: "offlineFirst",
        retry: (failureCount, error) => {
          if (error instanceof ApiRequestError && error.status === 401) {
            return false
          }
          return failureCount < 3
        },
      },
      mutations: {
        networkMode: "offlineFirst",
      },
    },
  })
}

export function QueryProvider({ children }: { children: React.ReactNode }) {
  const [queryClient] = useState(() => {
    const client = makeQueryClient()
    attachPreviewBlobGc(client)
    return client
  })

  return (
    <PersistQueryClientProvider
      client={queryClient}
      persistOptions={{
        persister: idbPersister,
        dehydrateOptions: {
          shouldDehydrateQuery: (query) => {
            if (query.state.status !== "success") return false
            if (isPreviewThumbQueryKey(query.queryKey)) return false
            const root = query.queryKey[0]
            // Mail list/detail queries change often and pollute SSR hydration when restored from IDB.
            if (
              root === "messages" ||
              root === "message" ||
              root === "mail-search" ||
              root === "thread" ||
              root === "demo"
            ) {
              return false
            }
            return true
          },
        },
      }}
    >
      {children}
    </PersistQueryClientProvider>
  )
}