"use client"; import { useQuery as useZeroQuery } from "@rocicorp/zero/react"; import { keepPreviousData, useQuery } from "@tanstack/react-query"; import { Folder as FolderIcon } from "lucide-react"; import { forwardRef, useCallback, useDeferredValue, useEffect, useImperativeHandle, useMemo, useRef, useState, } from "react"; import { FOLDER_MENTION_DOCUMENT_TYPE, type MentionedDocumentInfo, } from "@/atoms/chat/mentioned-documents.atom"; import { Skeleton } from "@/components/ui/skeleton"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document, SearchDocumentTitlesResponse } from "@/contracts/types/document.types"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { cn } from "@/lib/utils"; import { queries } from "@/zero/queries"; export interface DocumentMentionPickerRef { selectHighlighted: () => void; moveUp: () => void; moveDown: () => void; } interface DocumentMentionPickerProps { searchSpaceId: number; onSelectionChange: (mentions: MentionedDocumentInfo[]) => void; onDone: () => void; initialSelectedDocuments?: MentionedDocumentInfo[]; externalSearch?: string; } const PAGE_SIZE = 20; const MIN_SEARCH_LENGTH = 2; const DEBOUNCE_MS = 100; /** * Custom debounce hook that delays value updates until user input stabilizes. * Preferred over throttling for search inputs as it reduces API request frequency * and prevents race conditions from stale responses overtaking recent ones. */ function useDebounced(value: T, delay = DEBOUNCE_MS) { const [debounced, setDebounced] = useState(value); const timeoutRef = useRef | undefined>(undefined); useEffect(() => { if (timeoutRef.current) { clearTimeout(timeoutRef.current); } timeoutRef.current = setTimeout(() => { setDebounced(value); }, delay); return () => { if (timeoutRef.current) { clearTimeout(timeoutRef.current); } }; }, [value, delay]); return debounced; } export const DocumentMentionPicker = forwardRef< DocumentMentionPickerRef, DocumentMentionPickerProps >(function DocumentMentionPicker( { searchSpaceId, onSelectionChange, onDone, initialSelectedDocuments = [], externalSearch = "" }, ref ) { // Debounced search value to minimize API calls and prevent race conditions const search = externalSearch; const debouncedSearch = useDebounced(search, DEBOUNCE_MS); // Deferred snapshot of debouncedSearch — client-side filtering uses this so it // is treated as a non-urgent update, keeping the input responsive. const deferredSearch = useDeferredValue(debouncedSearch); const [highlightedIndex, setHighlightedIndex] = useState(0); const itemRefs = useRef>(new Map()); const scrollContainerRef = useRef(null); const shouldScrollRef = useRef(false); // Keyboard navigation scroll flag // Pagination state for infinite scroll const [accumulatedDocuments, setAccumulatedDocuments] = useState< Pick[] >([]); const [currentPage, setCurrentPage] = useState(0); const [hasMore, setHasMore] = useState(false); const [isLoadingMore, setIsLoadingMore] = useState(false); // Folders for this search space — pulled from Zero so the picker // stays consistent with the documents sidebar (same source of // truth, automatic updates on rename/delete). const [zeroFolders] = useZeroQuery(queries.folders.bySpace({ searchSpaceId })); /** * Search Strategy: * - Single character (length === 1): Client-side filtering for instant results * - Two or more characters (length >= 2): Server-side search with pg_trgm index * This hybrid approach optimizes UX by providing immediate feedback for short queries * while leveraging efficient database indexing for longer, more specific searches. */ const isSearchValid = debouncedSearch.trim().length >= MIN_SEARCH_LENGTH; const shouldSearch = debouncedSearch.trim().length > 0; const isSingleCharSearch = debouncedSearch.trim().length === 1; // Reset pagination state when search query or search space changes. // Documents are not cleared to maintain visual continuity during fetches. // biome-ignore lint/correctness/useExhaustiveDependencies: Intentional reset on search/space change useEffect(() => { setCurrentPage(0); setHasMore(false); setHighlightedIndex(0); }, [debouncedSearch, searchSpaceId]); // Query parameters for lightweight title search endpoint const titleSearchParams = useMemo( () => ({ search_space_id: searchSpaceId, page: 0, page_size: PAGE_SIZE, ...(isSearchValid ? { title: debouncedSearch.trim() } : {}), }), [searchSpaceId, debouncedSearch, isSearchValid] ); const surfsenseDocsQueryParams = useMemo(() => { const params: { page: number; page_size: number; title?: string } = { page: 0, page_size: PAGE_SIZE, }; if (isSearchValid) { params.title = debouncedSearch.trim(); } return params; }, [debouncedSearch, isSearchValid]); /** * TanStack Query for document title search. * - Uses AbortSignal for automatic request cancellation on query key changes * - placeholderData: keepPreviousData maintains UI stability during fetches * - Only triggers server-side search when isSearchValid (2+ characters) */ const { data: titleSearchResults, isLoading: isTitleSearchLoading } = useQuery({ queryKey: ["document-titles", titleSearchParams], queryFn: ({ signal }) => documentsApiService.searchDocumentTitles({ queryParams: titleSearchParams }, signal), staleTime: 60 * 1000, enabled: !!searchSpaceId && currentPage === 0 && (!shouldSearch || isSearchValid), placeholderData: keepPreviousData, }); /** * TanStack Query for SurfSense documentation. * - Uses AbortSignal for automatic request cancellation * - placeholderData: keepPreviousData prevents UI flicker during refetches */ const { data: surfsenseDocs, isLoading: isSurfsenseDocsLoading } = useQuery({ queryKey: ["surfsense-docs-mention", debouncedSearch, isSearchValid], queryFn: ({ signal }) => documentsApiService.getSurfsenseDocs({ queryParams: surfsenseDocsQueryParams }, signal), staleTime: 3 * 60 * 1000, enabled: !shouldSearch || isSearchValid, placeholderData: keepPreviousData, }); // Post-fetch filter to eliminate false positives from backend fuzzy matching const filterBySearchTerm = useCallback( (docs: Pick[]) => { if (!isSearchValid) return docs; // No filtering when not searching const searchLower = debouncedSearch.trim().toLowerCase(); return docs.filter((doc) => doc.title.toLowerCase().includes(searchLower)); }, [debouncedSearch, isSearchValid] ); // Combine and update document list when first page data arrives useEffect(() => { if (currentPage === 0) { const combinedDocs: Pick[] = []; // SurfSense docs displayed first in the list if (surfsenseDocs?.items) { for (const doc of surfsenseDocs.items) { combinedDocs.push({ id: doc.id, title: doc.title, document_type: "SURFSENSE_DOCS", }); } } if (titleSearchResults?.items) { combinedDocs.push(...titleSearchResults.items); setHasMore(titleSearchResults.has_more); } setAccumulatedDocuments(filterBySearchTerm(combinedDocs)); } }, [titleSearchResults, surfsenseDocs, currentPage, filterBySearchTerm]); // Load next page for infinite scroll pagination const loadNextPage = useCallback(async () => { if (isLoadingMore || !hasMore) return; const nextPage = currentPage + 1; setIsLoadingMore(true); try { const queryParams = { search_space_id: searchSpaceId, page: nextPage, page_size: PAGE_SIZE, ...(isSearchValid ? { title: debouncedSearch.trim() } : {}), }; const response: SearchDocumentTitlesResponse = await documentsApiService.searchDocumentTitles( { queryParams } ); setAccumulatedDocuments((prev) => [...prev, ...response.items]); setHasMore(response.has_more); setCurrentPage(nextPage); } catch (error) { console.error("Failed to load next page:", error); } finally { setIsLoadingMore(false); } }, [currentPage, hasMore, isLoadingMore, debouncedSearch, searchSpaceId, isSearchValid]); // Trigger pagination when user scrolls near the bottom (50px threshold) const handleScroll = useCallback( (e: React.UIEvent) => { const target = e.currentTarget; const scrollBottom = target.scrollHeight - target.scrollTop - target.clientHeight; if (scrollBottom < 50 && hasMore && !isLoadingMore) { loadNextPage(); } }, [hasMore, isLoadingMore, loadNextPage] ); /** * Client-side filtering for single character searches. * Filters cached documents locally for instant feedback without additional API calls. * Server-side search is reserved for 2+ character queries to leverage database indexing. * Uses deferredSearch (a deferred snapshot of debouncedSearch) so this memo is treated * as non-urgent — React can interrupt it to keep the input responsive. */ const clientFilteredDocs = useMemo(() => { if (!isSingleCharSearch) return null; const searchLower = deferredSearch.trim().toLowerCase(); return accumulatedDocuments.filter((doc) => doc.title.toLowerCase().includes(searchLower)); }, [isSingleCharSearch, deferredSearch, accumulatedDocuments]); // Select data source based on search length: client-filtered for single char, server results for 2+ const actualDocuments = isSingleCharSearch ? (clientFilteredDocs ?? []) : accumulatedDocuments; // Only show loading spinner on initial load (no documents yet), not during subsequent searches const actualLoading = (isTitleSearchLoading || isSurfsenseDocsLoading) && currentPage === 0 && !isSingleCharSearch && accumulatedDocuments.length === 0; // Partition documents by type for grouped UI rendering const surfsenseDocsList = useMemo( () => actualDocuments.filter((doc) => doc.document_type === "SURFSENSE_DOCS"), [actualDocuments] ); const userDocsList = useMemo( () => actualDocuments.filter((doc) => doc.document_type !== "SURFSENSE_DOCS"), [actualDocuments] ); // Folder mention candidates filtered by the current search term. // Single-char and server-search both use the same client filter // — folder counts in a workspace are tiny compared to docs, so we // don't need a paged endpoint. Empty search shows all folders. const folderMentions: MentionedDocumentInfo[] = useMemo(() => { const all = (zeroFolders ?? []).map((f) => ({ id: f.id, title: f.name, document_type: FOLDER_MENTION_DOCUMENT_TYPE, kind: "folder" as const, })); if (!shouldSearch) return all; const needle = (isSingleCharSearch ? deferredSearch : debouncedSearch).trim().toLowerCase(); if (!needle) return all; return all.filter((f) => f.title.toLowerCase().includes(needle)); }, [zeroFolders, debouncedSearch, deferredSearch, isSingleCharSearch, shouldSearch]); // Doc-shape entries reuse their ``document_type`` discriminator; // folder entries lift the existing kind-aware key so the same // matchers used by the chip atom apply unchanged. const selectedKeys = useMemo( () => new Set(initialSelectedDocuments.map((d) => getMentionDocKey(d))), [initialSelectedDocuments] ); // Combined navigation order: SurfSense docs -> User docs -> Folders. // Mirrors the on-screen ordering so keyboard arrows match what the // user sees. const selectableMentions = useMemo(() => { const docs: MentionedDocumentInfo[] = actualDocuments.map((doc) => ({ id: doc.id, title: doc.title, document_type: doc.document_type, kind: "doc" as const, })); const ordered = [...docs, ...folderMentions]; return ordered.filter((m) => !selectedKeys.has(getMentionDocKey(m))); }, [actualDocuments, folderMentions, selectedKeys]); const handleSelectMention = useCallback( (mention: MentionedDocumentInfo) => { onSelectionChange([...initialSelectedDocuments, mention]); onDone(); }, [initialSelectedDocuments, onSelectionChange, onDone] ); // Auto-scroll highlighted item into view (keyboard navigation only, not mouse hover) useEffect(() => { if (!shouldScrollRef.current) { return; } shouldScrollRef.current = false; const rafId = requestAnimationFrame(() => { const item = itemRefs.current.get(highlightedIndex); const container = scrollContainerRef.current; if (item && container) { const itemRect = item.getBoundingClientRect(); const containerRect = container.getBoundingClientRect(); const padding = 8; const isAboveViewport = itemRect.top < containerRect.top + padding; const isBelowViewport = itemRect.bottom > containerRect.bottom - padding; if (isAboveViewport || isBelowViewport) { const itemOffsetTop = item.offsetTop; const containerHeight = container.clientHeight; const itemHeight = item.offsetHeight; const targetScrollTop = itemOffsetTop - containerHeight / 2 + itemHeight / 2; const maxScrollTop = container.scrollHeight - containerHeight; const clampedScrollTop = Math.max(0, Math.min(targetScrollTop, maxScrollTop)); container.scrollTo({ top: clampedScrollTop, behavior: "smooth", }); } } }); return () => cancelAnimationFrame(rafId); }, [highlightedIndex]); // Reset highlight position when search query changes const prevSearchRef = useRef(search); if (prevSearchRef.current !== search) { prevSearchRef.current = search; if (highlightedIndex !== 0) { setHighlightedIndex(0); } } // Expose navigation and selection methods to parent component via ref useImperativeHandle( ref, () => ({ selectHighlighted: () => { if (selectableMentions[highlightedIndex]) { handleSelectMention(selectableMentions[highlightedIndex]); } }, moveUp: () => { shouldScrollRef.current = true; setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableMentions.length - 1)); }, moveDown: () => { shouldScrollRef.current = true; setHighlightedIndex((prev) => (prev < selectableMentions.length - 1 ? prev + 1 : 0)); }, }), [selectableMentions, highlightedIndex, handleSelectMention] ); // Keyboard navigation handler for arrow keys, Enter, and Escape const handleKeyDown = useCallback( (e: React.KeyboardEvent) => { if (selectableMentions.length === 0) return; switch (e.key) { case "ArrowDown": e.preventDefault(); shouldScrollRef.current = true; setHighlightedIndex((prev) => (prev < selectableMentions.length - 1 ? prev + 1 : 0)); break; case "ArrowUp": e.preventDefault(); shouldScrollRef.current = true; setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableMentions.length - 1)); break; case "Enter": e.preventDefault(); if (selectableMentions[highlightedIndex]) { handleSelectMention(selectableMentions[highlightedIndex]); } break; case "Escape": e.preventDefault(); onDone(); break; } }, [selectableMentions, highlightedIndex, handleSelectMention, onDone] ); return (
{/* Scrollable document list with responsive height */}
{actualLoading ? (
{["a", "b", "c", "d", "e"].map((id, i) => (
= 3 && "hidden sm:flex" )} >
))}
) : actualDocuments.length > 0 || folderMentions.length > 0 ? (
{/* SurfSense Documentation */} {surfsenseDocsList.length > 0 && ( <>
SurfSense Docs
{surfsenseDocsList.map((doc) => { const mention: MentionedDocumentInfo = { id: doc.id, title: doc.title, document_type: doc.document_type, kind: "doc", }; const docKey = getMentionDocKey(mention); const isAlreadySelected = selectedKeys.has(docKey); const selectableIndex = selectableMentions.findIndex( (m) => getMentionDocKey(m) === docKey ); const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; return ( ); })} )} {/* User Documents */} {userDocsList.length > 0 && ( <> {surfsenseDocsList.length > 0 && (
)}
Your Documents
{userDocsList.map((doc) => { const mention: MentionedDocumentInfo = { id: doc.id, title: doc.title, document_type: doc.document_type, kind: "doc", }; const docKey = getMentionDocKey(mention); const isAlreadySelected = selectedKeys.has(docKey); const selectableIndex = selectableMentions.findIndex( (m) => getMentionDocKey(m) === docKey ); const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; return ( ); })} )} {/* Folders — single source of truth is Zero (same store that powers the documents sidebar). Selecting a folder inserts a folder chip whose path the agent can walk with ``ls`` / ``find_documents``. */} {folderMentions.length > 0 && ( <> {(surfsenseDocsList.length > 0 || userDocsList.length > 0) && (
)}
Folders
{folderMentions.map((folder) => { const folderKey = getMentionDocKey(folder); const isAlreadySelected = selectedKeys.has(folderKey); const selectableIndex = selectableMentions.findIndex( (m) => getMentionDocKey(m) === folderKey ); const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex; return ( ); })} )} {/* Pagination loading indicator */} {isLoadingMore && (
)}
) : (

No matching documents

)}
); });