From 1bb9f479e1dcdaa85b5cce8aef4e750c900a623a Mon Sep 17 00:00:00 2001 From: "DESKTOP-RTLN3BA\\$punk" Date: Sat, 28 Feb 2026 01:54:54 -0800 Subject: [PATCH] feat: refactor document fetching and improve comment batching - Replaced the useDocuments hook with React Query for better caching and deduplication of document requests. - Updated the ConnectorIndicator component to fetch document type counts using a new atom for real-time updates. - Enhanced the useComments hook to manage batch requests more effectively, reducing race conditions and improving performance. - Set default query options in the query client to optimize stale time and refetch behavior. --- .../assistant-ui/connector-popup.tsx | 9 +- surfsense_web/hooks/use-comments.ts | 31 ++++++- surfsense_web/hooks/use-documents.ts | 88 +++++++++---------- surfsense_web/lib/query-client/client.ts | 9 +- 4 files changed, 80 insertions(+), 57 deletions(-) diff --git a/surfsense_web/components/assistant-ui/connector-popup.tsx b/surfsense_web/components/assistant-ui/connector-popup.tsx index 98964013d..332694676 100644 --- a/surfsense_web/components/assistant-ui/connector-popup.tsx +++ b/surfsense_web/components/assistant-ui/connector-popup.tsx @@ -5,6 +5,7 @@ import { AlertTriangle, Cable, Settings } from "lucide-react"; import Link from "next/link"; import { useSearchParams } from "next/navigation"; import type { FC } from "react"; +import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; import { globalNewLLMConfigsAtom, llmPreferencesAtom, @@ -19,7 +20,6 @@ import { Spinner } from "@/components/ui/spinner"; import { Tabs, TabsContent } from "@/components/ui/tabs"; import type { SearchSourceConnector } from "@/contracts/types/connector.types"; import { useConnectorsElectric } from "@/hooks/use-connectors-electric"; -import { useDocuments } from "@/hooks/use-documents"; import { useInbox } from "@/hooks/use-inbox"; import { cn } from "@/lib/utils"; import { ConnectorDialogHeader } from "./connector-popup/components/connector-dialog-header"; @@ -62,10 +62,9 @@ export const ConnectorIndicator: FC<{ hideTrigger?: boolean }> = ({ hideTrigger const llmConfigLoading = preferencesLoading || globalConfigsLoading; - // Fetch document type counts using Electric SQL + PGlite for real-time updates - const { typeCounts: documentTypeCounts, loading: documentTypesLoading } = useDocuments( - searchSpaceId ? Number(searchSpaceId) : null - ); + // Fetch document type counts via the lightweight /type-counts endpoint (cached 10 min) + const { data: documentTypeCounts, isFetching: documentTypesLoading } = + useAtomValue(documentTypeCountsAtom); // Fetch notifications to detect indexing failures const { inboxItems = [] } = useInbox( diff --git a/surfsense_web/hooks/use-comments.ts b/surfsense_web/hooks/use-comments.ts index 562f7ae02..c02f9fe16 100644 --- a/surfsense_web/hooks/use-comments.ts +++ b/surfsense_web/hooks/use-comments.ts @@ -11,9 +11,26 @@ interface UseCommentsOptions { // --------------------------------------------------------------------------- // Module-level coordination: when a batch request is in-flight, individual // useComments queryFns piggy-back on it instead of making their own requests. +// +// _batchReady is a promise that resolves once the batch useEffect has had a +// chance to set _batchInflight. Individual queryFns await this gate before +// deciding whether to piggy-back or fetch on their own, eliminating the +// previous race where setTimeout(0) was not enough. // --------------------------------------------------------------------------- let _batchInflight: Promise | null = null; let _batchTargetIds = new Set(); +let _batchReady: Promise | null = null; +let _resolveBatchReady: (() => void) | null = null; + +function resetBatchGate() { + _batchReady = new Promise((r) => { + _resolveBatchReady = r; + }); +} + +// Open the initial gate immediately (no batch pending yet) +resetBatchGate(); +_resolveBatchReady?.(); export function useComments({ messageId, enabled = true }: UseCommentsOptions) { const queryClient = useQueryClient(); @@ -21,9 +38,11 @@ export function useComments({ messageId, enabled = true }: UseCommentsOptions) { return useQuery({ queryKey: cacheKeys.comments.byMessage(messageId), queryFn: async () => { - // Yield one macro-task so the batch prefetch useEffect (which sets - // _batchInflight) has a chance to fire before we decide to fetch. - await new Promise((r) => setTimeout(r, 0)); + // Wait for the batch gate so the useEffect in useBatchCommentsPreload + // has a chance to set _batchInflight before we decide. + if (_batchReady) { + await _batchReady; + } if (_batchInflight && _batchTargetIds.has(messageId)) { await _batchInflight; @@ -57,6 +76,9 @@ export function useBatchCommentsPreload(messageIds: number[]) { if (key === prevKeyRef.current) return; prevKeyRef.current = key; + // Open a new gate so individual queryFns wait for us + resetBatchGate(); + _batchTargetIds = new Set(messageIds); let cancelled = false; @@ -80,6 +102,9 @@ export function useBatchCommentsPreload(messageIds: number[]) { _batchInflight = promise; + // Release the gate — individual queryFns can now check _batchInflight + _resolveBatchReady?.(); + return () => { cancelled = true; if (_batchInflight === promise) { diff --git a/surfsense_web/hooks/use-documents.ts b/surfsense_web/hooks/use-documents.ts index 55d48c4f1..36a359696 100644 --- a/surfsense_web/hooks/use-documents.ts +++ b/surfsense_web/hooks/use-documents.ts @@ -1,5 +1,6 @@ "use client"; +import { useQuery } from "@tanstack/react-query"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import type { DocumentTypeEnum } from "@/contracts/types/document.types"; import { documentsApiService } from "@/lib/apis/documents-api.service"; @@ -183,56 +184,47 @@ export function useDocuments( [] ); - // EFFECT 1: Load ALL documents from API (PRIMARY source of truth) - // No type filter — always fetches everything so typeCounts stay complete + // STEP 1: Load ALL documents from API (PRIMARY source of truth). + // Uses React Query for automatic deduplication, caching, and staleTime so + // multiple components mounting useDocuments(sameId) share a single request. + const { + data: apiResponse, + isLoading: apiLoading, + error: apiError, + } = useQuery({ + queryKey: ["documents", "all", searchSpaceId], + queryFn: () => + documentsApiService.getDocuments({ + queryParams: { + search_space_id: searchSpaceId!, + page: 0, + page_size: -1, + }, + }), + enabled: !!searchSpaceId, + staleTime: 30_000, + }); + + // Seed local state from API response (runs once per fresh fetch) useEffect(() => { - if (!searchSpaceId) { - setLoading(false); - return; + if (!apiResponse) return; + populateUserCache(apiResponse.items); + const docs = apiResponse.items.map(apiToDisplayDoc); + setAllDocuments(docs); + apiLoadedRef.current = true; + setError(null); + }, [apiResponse, populateUserCache, apiToDisplayDoc]); + + // Propagate loading / error from React Query + useEffect(() => { + setLoading(apiLoading); + }, [apiLoading]); + + useEffect(() => { + if (apiError) { + setError(apiError instanceof Error ? apiError : new Error("Failed to load documents")); } - - // Capture validated value for async closure - const spaceId = searchSpaceId; - - let mounted = true; - apiLoadedRef.current = false; - - async function loadFromApi() { - try { - setLoading(true); - console.log("[useDocuments] Loading from API (source of truth):", spaceId); - - const response = await documentsApiService.getDocuments({ - queryParams: { - search_space_id: spaceId, - page: 0, - page_size: -1, // Fetch all documents (unfiltered) - }, - }); - - if (!mounted) return; - - populateUserCache(response.items); - const docs = response.items.map(apiToDisplayDoc); - setAllDocuments(docs); - apiLoadedRef.current = true; - setError(null); - console.log("[useDocuments] API loaded", docs.length, "documents"); - } catch (err) { - if (!mounted) return; - console.error("[useDocuments] API load failed:", err); - setError(err instanceof Error ? err : new Error("Failed to load documents")); - } finally { - if (mounted) setLoading(false); - } - } - - loadFromApi(); - - return () => { - mounted = false; - }; - }, [searchSpaceId, populateUserCache, apiToDisplayDoc]); + }, [apiError]); // EFFECT 2: Start Electric sync + live query for real-time updates // No type filter — syncs and queries ALL documents; filtering is client-side diff --git a/surfsense_web/lib/query-client/client.ts b/surfsense_web/lib/query-client/client.ts index 6c7b9ded3..0dcc2ef03 100644 --- a/surfsense_web/lib/query-client/client.ts +++ b/surfsense_web/lib/query-client/client.ts @@ -1,3 +1,10 @@ import { QueryClient } from "@tanstack/react-query"; -export const queryClient = new QueryClient(); +export const queryClient = new QueryClient({ + defaultOptions: { + queries: { + staleTime: 30_000, + refetchOnWindowFocus: false, + }, + }, +});