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.

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
+	// Fetch document type counts via the lightweight /type-counts endpoint (cached 10 min)
-	const { typeCounts: documentTypeCounts, loading: documentTypesLoading } = useDocuments(
+	const { data: documentTypeCounts, isFetching: documentTypesLoading } =
-		searchSpaceId ? Number(searchSpaceId) : null
+		useAtomValue(documentTypeCountsAtom);
-	);
 
 	// Fetch notifications to detect indexing failures
 	const { inboxItems = [] } = useInbox(

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<void> | null = null;
 let _batchTargetIds = new Set<number>();
+let _batchReady: Promise<void> | null = null;
+let _resolveBatchReady: (() => void) | null = null;
+
+function resetBatchGate() {
+	_batchReady = new Promise<void>((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
+			// Wait for the batch gate so the useEffect in useBatchCommentsPreload
-			// _batchInflight) has a chance to fire before we decide to fetch.
+			// has a chance to set _batchInflight before we decide.
-			await new Promise<void>((r) => setTimeout(r, 0));
+			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) {

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)
+	// STEP 1: Load ALL documents from API (PRIMARY source of truth).
-	// No type filter — always fetches everything so typeCounts stay complete
+	// Uses React Query for automatic deduplication, caching, and staleTime so
-	useEffect(() => {
+	// multiple components mounting useDocuments(sameId) share a single request.
-		if (!searchSpaceId) {
+	const {
-			setLoading(false);
+		data: apiResponse,
-			return;
+		isLoading: apiLoading,
-		}
+		error: apiError,
-
+	} = useQuery({
-		// Capture validated value for async closure
+		queryKey: ["documents", "all", searchSpaceId],
-		const spaceId = searchSpaceId;
+		queryFn: () =>
-
+			documentsApiService.getDocuments({
-		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,
+					search_space_id: searchSpaceId!,
 					page: 0,
-						page_size: -1, // Fetch all documents (unfiltered)
+					page_size: -1,
 				},
+			}),
+		enabled: !!searchSpaceId,
+		staleTime: 30_000,
 	});
 
-				if (!mounted) return;
+	// Seed local state from API response (runs once per fresh fetch)
-
+	useEffect(() => {
-				populateUserCache(response.items);
+		if (!apiResponse) return;
-				const docs = response.items.map(apiToDisplayDoc);
+		populateUserCache(apiResponse.items);
+		const docs = apiResponse.items.map(apiToDisplayDoc);
 		setAllDocuments(docs);
 		apiLoadedRef.current = true;
 		setError(null);
-				console.log("[useDocuments] API loaded", docs.length, "documents");
+	}, [apiResponse, populateUserCache, apiToDisplayDoc]);
-			} 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();
+	// Propagate loading / error from React Query
+	useEffect(() => {
+		setLoading(apiLoading);
+	}, [apiLoading]);
 
-		return () => {
+	useEffect(() => {
-			mounted = false;
+		if (apiError) {
-		};
+			setError(apiError instanceof Error ? apiError : new Error("Failed to load documents"));
-	}, [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

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