refactor: replace document type counts atom with real-time hook

- Removed the `documentTypeCountsAtom` and its associated logic from the document query atoms.
- Introduced `useZeroDocumentTypeCounts` hook to provide real-time document type counts, enhancing responsiveness as documents are indexed.
- Updated components to utilize the new hook for fetching document type counts, ensuring instant updates in the UI.

surfsense_web/atoms/documents/document-query.atoms.ts

@@ -1,38 +0,0 @@
-import { atomWithQuery } from "jotai-tanstack-query";
-import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
-import type { SearchDocumentsRequest } from "@/contracts/types/document.types";
-import { documentsApiService } from "@/lib/apis/documents-api.service";
-import { cacheKeys } from "@/lib/query-client/cache-keys";
-import { globalDocumentsQueryParamsAtom } from "./ui.atoms";
-
-export const documentsAtom = atomWithQuery((get) => {
-	const searchSpaceId = get(activeSearchSpaceIdAtom);
-	const queryParams = get(globalDocumentsQueryParamsAtom);
-
-	return {
-		queryKey: cacheKeys.documents.globalQueryParams(queryParams),
-		enabled: !!searchSpaceId,
-		queryFn: async () => {
-			return documentsApiService.getDocuments({
-				queryParams: queryParams,
-			});
-		},
-	};
-});
-
-export const documentTypeCountsAtom = atomWithQuery((get) => {
-	const searchSpaceId = get(activeSearchSpaceIdAtom);
-
-	return {
-		queryKey: cacheKeys.documents.typeCounts(searchSpaceId ?? undefined),
-		enabled: !!searchSpaceId,
-		staleTime: 10 * 60 * 1000, // 10 minutes
-		queryFn: async () => {
-			return documentsApiService.getDocumentTypeCounts({
-				queryParams: {
-					search_space_id: searchSpaceId ?? undefined,
-				},
-			});
-		},
-	};
-});

surfsense_web/components/assistant-ui/connector-popup.tsx

@@ -4,7 +4,7 @@ import { useAtomValue, useSetAtom } from "jotai";
 import { AlertTriangle, Cable, Settings } from "lucide-react";
 import { forwardRef, useEffect, useImperativeHandle, useMemo, useState } from "react";
 import { createPortal } from "react-dom";
-import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
+import { useZeroDocumentTypeCounts } from "@/hooks/use-zero-document-type-counts";
 import { statusInboxItemsAtom } from "@/atoms/inbox/status-inbox.atom";
 import {
 	globalNewLLMConfigsAtom,
@@ -72,9 +72,9 @@ export const ConnectorIndicator = forwardRef<ConnectorIndicatorHandle, Connector
 
 		const llmConfigLoading = preferencesLoading || globalConfigsLoading;
 
-		// Fetch document type counts via the lightweight /type-counts endpoint (cached 10 min)
+		// Real-time document type counts via Zero (updates instantly as docs are indexed)
-		const { data: documentTypeCounts, isFetching: documentTypesLoading } =
+		const documentTypeCounts = useZeroDocumentTypeCounts(searchSpaceId);
-			useAtomValue(documentTypeCountsAtom);
+		const documentTypesLoading = documentTypeCounts === undefined;
 
 		// Read status inbox items from shared atom (populated by LayoutDataProvider)
 		// instead of creating a duplicate useInbox("status") hook.

surfsense_web/components/onboarding-tour.tsx

@@ -7,7 +7,7 @@ import { useTheme } from "next-themes";
 import { useCallback, useEffect, useRef, useState } from "react";
 import { createPortal } from "react-dom";
 import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms";
-import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
+import { useZeroDocumentTypeCounts } from "@/hooks/use-zero-document-type-counts";
 import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
 import { currentUserAtom } from "@/atoms/user/user-query.atoms";
 import { useIsMobile } from "@/hooks/use-mobile";
@@ -451,8 +451,8 @@ export function OnboardingTour() {
 		enabled: !!searchSpaceId,
 	});
 
-	// Get document type counts
+	// Real-time document type counts via Zero
-	const { data: documentTypeCounts } = useAtomValue(documentTypeCountsAtom);
+	const documentTypeCounts = useZeroDocumentTypeCounts(searchSpaceId);
 
 	// Get connectors
 	const { data: connectors = [] } = useAtomValue(connectorsAtom);

surfsense_web/hooks/use-zero-document-type-counts.ts

@@ -0,0 +1,31 @@
+"use client";
+
+import { useQuery } from "@rocicorp/zero/react";
+import { useMemo } from "react";
+import { queries } from "@/zero/queries";
+
+/**
+ * Real-time document type counts derived from Zero's live document sync.
+ * Updates instantly as documents are created, deleted, or change type.
+ */
+export function useZeroDocumentTypeCounts(
+	searchSpaceId: number | string | null
+): Record<string, number> | undefined {
+	const numericId = searchSpaceId != null ? Number(searchSpaceId) : null;
+
+	const [zeroDocuments] = useQuery(
+		queries.documents.bySpace({ searchSpaceId: numericId ?? -1 })
+	);
+
+	return useMemo(() => {
+		if (!zeroDocuments || numericId == null) return undefined;
+
+		const counts: Record<string, number> = {};
+		for (const doc of zeroDocuments) {
+			if (doc.id != null && doc.title != null && doc.title !== "") {
+				counts[doc.documentType] = (counts[doc.documentType] || 0) + 1;
+			}
+		}
+		return counts;
+	}, [zeroDocuments, numericId]);
+}