diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 645b59010..7773a438a 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -24,7 +24,6 @@ import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom, - sidebarSelectedDocumentsAtom, } from "@/atoms/chat/mentioned-documents.atom"; import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom"; import { @@ -221,12 +220,10 @@ export default function NewChatPage() { // Get disabled tools from the tool toggle UI const disabledTools = useAtomValue(disabledToolsAtom); - // Get mentioned document IDs from the composer (derived from @ mentions + sidebar selections) + // Get mentioned document IDs from the composer. const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom); const mentionedDocuments = useAtomValue(mentionedDocumentsAtom); - const sidebarDocuments = useAtomValue(sidebarSelectedDocumentsAtom); const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom); - const setSidebarDocuments = useSetAtom(sidebarSelectedDocumentsAtom); const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom); const setCurrentThreadState = useSetAtom(currentThreadAtom); const setTargetCommentId = useSetAtom(setTargetCommentIdAtom); @@ -326,7 +323,6 @@ export default function NewChatPage() { setCurrentThread(null); setMentionedDocuments([]); tokenUsageStore.clear(); - setSidebarDocuments([]); setMessageDocumentsMap({}); clearPlanOwnerRegistry(); closeReportPanel(); @@ -394,7 +390,6 @@ export default function NewChatPage() { urlChatId, setMessageDocumentsMap, setMentionedDocuments, - setSidebarDocuments, closeReportPanel, closeEditorPanel, removeChatTab, @@ -600,15 +595,14 @@ export default function NewChatPage() { messageLength: userQuery.length, }); - // Combine @-mention chips + sidebar selections for display & persistence + // Collect unique mentioned docs for display & persistence const allMentionedDocs: MentionedDocumentInfo[] = []; const seenDocKeys = new Set(); - for (const doc of [...mentionedDocuments, ...sidebarDocuments]) { + for (const doc of mentionedDocuments) { const key = `${doc.document_type}:${doc.id}`; - if (!seenDocKeys.has(key)) { - seenDocKeys.add(key); - allMentionedDocs.push({ id: doc.id, title: doc.title, document_type: doc.document_type }); - } + if (seenDocKeys.has(key)) continue; + seenDocKeys.add(key); + allMentionedDocs.push({ id: doc.id, title: doc.title, document_type: doc.document_type }); } if (allMentionedDocs.length > 0) { @@ -710,7 +704,6 @@ export default function NewChatPage() { // Clear mentioned documents after capturing them if (hasDocumentIds || hasSurfsenseDocIds) { setMentionedDocuments([]); - setSidebarDocuments([]); } const response = await fetch(`${backendUrl}/api/v1/new_chat`, { @@ -994,9 +987,7 @@ export default function NewChatPage() { messages, mentionedDocumentIds, mentionedDocuments, - sidebarDocuments, setMentionedDocuments, - setSidebarDocuments, setMessageDocumentsMap, setAgentCreatedDocuments, queryClient, diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index ee93a409a..9c4546237 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -10,21 +10,11 @@ import type { Document } from "@/contracts/types/document.types"; export const mentionedDocumentsAtom = atom[]>([]); /** - * Atom to store documents selected via the sidebar checkboxes / row clicks. - * These are NOT inserted as chips – the composer shows a count badge instead. - */ -export const sidebarSelectedDocumentsAtom = atom< - Pick[] ->([]); - -/** - * Derived read-only atom that merges @-mention chips and sidebar selections - * into a single deduplicated set of document IDs for the backend. + * Derived read-only atom that maps deduplicated mentioned docs + * into backend payload fields. */ export const mentionedDocumentIdsAtom = atom((get) => { - const chipDocs = get(mentionedDocumentsAtom); - const sidebarDocs = get(sidebarSelectedDocumentsAtom); - const allDocs = [...chipDocs, ...sidebarDocs]; + const allDocs = get(mentionedDocumentsAtom); const seen = new Set(); const deduped = allDocs.filter((d) => { const key = `${d.document_type}:${d.id}`; diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index 45ad219dd..05277f508 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -11,25 +11,14 @@ import { useRef, useState, } from "react"; -import { flushSync } from "react-dom"; -import { createRoot } from "react-dom/client"; +import { renderToStaticMarkup } from "react-dom/server"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document } from "@/contracts/types/document.types"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { cn } from "@/lib/utils"; -// Render a React element to an HTML string on the client without pulling -// `react-dom/server` into the bundle. `createRoot` + `flushSync` use the -// same `react-dom` package React itself imports, so this adds zero new -// runtime weight. function renderElementToHTML(element: ReactElement): string { - const container = document.createElement("div"); - const root = createRoot(container); - flushSync(() => { - root.render(element); - }); - const html = container.innerHTML; - root.unmount(); - return html; + return renderToStaticMarkup(element); } export interface MentionedDocument { @@ -44,7 +33,10 @@ export interface InlineMentionEditorRef { setText: (text: string) => void; getText: () => string; getMentionedDocuments: () => MentionedDocument[]; - insertDocumentChip: (doc: Pick) => void; + insertDocumentChip: ( + doc: Pick, + options?: { removeTriggerText?: boolean } + ) => void; removeDocumentChip: (docId: number, docType?: string) => void; setDocumentChipStatus: ( docId: number, @@ -66,7 +58,6 @@ interface InlineMentionEditorProps { onKeyDown?: (e: React.KeyboardEvent) => void; disabled?: boolean; className?: string; - initialDocuments?: MentionedDocument[]; initialText?: string; } @@ -118,7 +109,6 @@ export const InlineMentionEditor = forwardRef(null); const [isEmpty, setIsEmpty] = useState(true); const [mentionedDocs, setMentionedDocs] = useState>( - () => new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d])) + () => new Map() ); const isComposingRef = useRef(false); + const lastSelectionRangeRef = useRef(null); + const isRangeInsideEditor = useCallback((range: Range | null): range is Range => { + if (!range || !editorRef.current) return false; + return ( + editorRef.current.contains(range.startContainer) && + editorRef.current.contains(range.endContainer) + ); + }, []); + const isSelectionInsideEditor = useCallback( + (selection: Selection | null): selection is Selection => { + if (!selection || selection.rangeCount === 0 || !editorRef.current) return false; + const range = selection.getRangeAt(0); + return isRangeInsideEditor(range); + }, + [isRangeInsideEditor] + ); + + const rememberSelection = useCallback(() => { + const selection = window.getSelection(); + if (!isSelectionInsideEditor(selection)) return; + lastSelectionRangeRef.current = selection.getRangeAt(0).cloneRange(); + }, [isSelectionInsideEditor]); + + const restoreRememberedSelection = useCallback((): Selection | null => { + const selection = window.getSelection(); + if (!selection) return null; + if (!isRangeInsideEditor(lastSelectionRangeRef.current)) return null; + selection.removeAllRanges(); + selection.addRange(lastSelectionRangeRef.current.cloneRange()); + return selection; + }, [isRangeInsideEditor]); - // Sync initial documents useEffect(() => { - if (initialDocuments.length > 0) { - setMentionedDocs( - new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d])) - ); - } - }, [initialDocuments]); + const handleSelectionChange = () => { + if (document.activeElement !== editorRef.current) return; + rememberSelection(); + }; + document.addEventListener("selectionchange", handleSelectionChange); + return () => document.removeEventListener("selectionchange", handleSelectionChange); + }, [rememberSelection]); useEffect(() => { if (!initialText || !editorRef.current) return; @@ -145,7 +166,7 @@ export const InlineMentionEditor = forwardRef { @@ -211,6 +232,19 @@ export const InlineMentionEditor = forwardRef) => { + const docs = docsOverride + ? Array.from(docsOverride.values()) + : Array.from(mentionedDocs.values()); + const text = getText(); + const empty = text.length === 0 && docs.length === 0; + setIsEmpty(empty); + onChange?.(text, docs); + }, + [getText, mentionedDocs, onChange] + ); + // Create a chip element for a document const createChipElement = useCallback( (doc: MentionedDocument): HTMLSpanElement => { @@ -246,10 +280,11 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(docKey); + syncEditorState(next); return next; }); onDocumentRemove?.(doc.id, doc.document_type); @@ -294,13 +329,17 @@ export const InlineMentionEditor = forwardRef) => { + ( + doc: Pick, + options?: { removeTriggerText?: boolean } + ) => { if (!editorRef.current) return; + const removeTriggerText = options?.removeTriggerText ?? true; // Validate required fields for type safety if (typeof doc.id !== "number" || typeof doc.title !== "string") { @@ -315,25 +354,51 @@ export const InlineMentionEditor = forwardRef new Map(prev).set(docKey, mentionDoc)); + const nextDocs = new Map(mentionedDocs); + nextDocs.set(docKey, mentionDoc); // Find and remove the @query text const selection = window.getSelection(); - if (!selection || selection.rangeCount === 0) { - // No selection, just append + const hasActiveSelection = isSelectionInsideEditor(selection); + const resolvedSelection = hasActiveSelection ? selection : restoreRememberedSelection(); + if ( + !resolvedSelection || + resolvedSelection.rangeCount === 0 || + !isSelectionInsideEditor(resolvedSelection) + ) { + // No valid in-editor selection: deterministically insert at end. + editorRef.current.focus(); + const endSelection = window.getSelection(); + if (!endSelection) return; + const endRange = document.createRange(); + endRange.selectNodeContents(editorRef.current); + endRange.collapse(false); + endSelection.removeAllRanges(); + endSelection.addRange(endRange); + const chip = createChipElement(mentionDoc); - editorRef.current.appendChild(chip); - editorRef.current.appendChild(document.createTextNode(" ")); - focusAtEnd(); + endRange.insertNode(chip); + endRange.setStartAfter(chip); + endRange.collapse(true); + const space = document.createTextNode(" "); + endRange.insertNode(space); + endRange.setStartAfter(space); + endRange.collapse(true); + endSelection.removeAllRanges(); + endSelection.addRange(endRange); + + syncEditorState(nextDocs); + rememberSelection(); return; } // Find the @ symbol before the cursor and remove it along with any query text - const range = selection.getRangeAt(0); + const range = resolvedSelection.getRangeAt(0); const textNode = range.startContainer; - if (textNode.nodeType === Node.TEXT_NODE) { + if (textNode.nodeType === Node.TEXT_NODE && removeTriggerText) { const text = textNode.textContent || ""; const cursorPos = range.startOffset; @@ -369,8 +434,9 @@ export const InlineMentionEditor = forwardRef { - onChange(getText(), getMentionedDocuments()); - }, 0); - } + syncEditorState(nextDocs); }, - [createChipElement, focusAtEnd, getText, getMentionedDocuments, onChange] + [ + createChipElement, + isSelectionInsideEditor, + mentionedDocs, + rememberSelection, + restoreRememberedSelection, + syncEditorState, + ] ); // Clear the editor const clear = useCallback(() => { if (editorRef.current) { editorRef.current.innerHTML = ""; - setIsEmpty(true); - setMentionedDocs(new Map()); + const emptyDocs = new Map(); + setMentionedDocs(emptyDocs); + syncEditorState(emptyDocs); } - }, []); + }, [syncEditorState]); // Replace editor content with plain text and place cursor at end const setText = useCallback( (text: string) => { if (!editorRef.current) return; editorRef.current.innerText = text; - const empty = text.length === 0; - setIsEmpty(empty); - onChange?.(text, Array.from(mentionedDocs.values())); + syncEditorState(); focusAtEnd(); }, - [focusAtEnd, onChange, mentionedDocs] + [focusAtEnd, syncEditorState] ); const setDocumentChipStatus = useCallback( @@ -473,7 +547,7 @@ export const InlineMentionEditor = forwardRef { if (!editorRef.current) return; - const chipKey = `${docType ?? "UNKNOWN"}:${docId}`; + const chipKey = getMentionDocKey({ id: docId, document_type: docType }); const chips = editorRef.current.querySelectorAll( `span[${CHIP_DATA_ATTR}="true"]` ); @@ -486,14 +560,11 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(chipKey); + syncEditorState(next); return next; }); - - const text = getText(); - const empty = text.length === 0 && mentionedDocs.size <= 1; - setIsEmpty(empty); }, - [getText, mentionedDocs.size] + [syncEditorState] ); // Expose methods via ref @@ -594,6 +665,7 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(chipKey); + syncEditorState(next); return next; }); // Notify parent that a document was removed @@ -676,10 +753,14 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(chipKey); + syncEditorState(next); return next; }); // Notify parent that a document was removed @@ -691,7 +772,7 @@ export const InlineMentionEditor = forwardRef - {/** biome-ignore lint/a11y/useSemanticElements: */} + {/* biome-ignore lint/a11y/noStaticElementInteractions: contenteditable mention editor requires a div for inline chips */}
{/* Placeholder with fade animation on change */} {isEmpty && ( diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index e7ae2f471..cf99598f1 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -39,12 +39,10 @@ import { import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom"; import { mentionedDocumentsAtom, - sidebarSelectedDocumentsAtom, } from "@/atoms/chat/mentioned-documents.atom"; import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom"; import { connectorDialogOpenAtom } from "@/atoms/connector-dialog/connector-dialog.atoms"; import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms"; -import { documentsSidebarOpenAtom } from "@/atoms/documents/ui.atoms"; import { membersAtom } from "@/atoms/members/members-query.atoms"; import { globalNewLLMConfigsAtom, @@ -91,6 +89,7 @@ import { useBatchCommentsPreload } from "@/hooks/use-comments"; import { useCommentsSync } from "@/hooks/use-comments-sync"; import { useMediaQuery } from "@/hooks/use-media-query"; import { useElectronAPI } from "@/hooks/use-platform"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { captureDisplayToPngDataUrl } from "@/lib/chat/display-media-capture"; import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/lib/layout-events"; import { cn } from "@/lib/utils"; @@ -364,12 +363,14 @@ const ClipboardChip: FC<{ text: string; onDismiss: () => void }> = ({ text, onDi const Composer: FC = () => { // Document mention state (atoms persist across component remounts) const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom); - const setSidebarDocs = useSetAtom(sidebarSelectedDocumentsAtom); const [showDocumentPopover, setShowDocumentPopover] = useState(false); const [showPromptPicker, setShowPromptPicker] = useState(false); const [mentionQuery, setMentionQuery] = useState(""); const [actionQuery, setActionQuery] = useState(""); const editorRef = useRef(null); + const prevMentionedDocsRef = useRef< + Map> + >(new Map()); const documentPickerRef = useRef(null); const promptPickerRef = useRef(null); const viewportRef = useRef(null); @@ -605,7 +606,6 @@ const Composer: FC = () => { aui.composer().send(); editorRef.current?.clear(); setMentionedDocuments([]); - setSidebarDocs([]); // With turnAnchor="top", ViewportSlack adds min-height to the last // assistant message so that scrolling-to-bottom actually positions the @@ -652,43 +652,71 @@ const Composer: FC = () => { clipboardInitialText, aui, setMentionedDocuments, - setSidebarDocs, threadViewportStore, ]); const handleDocumentRemove = useCallback( (docId: number, docType?: string) => { - setMentionedDocuments((prev) => - prev.filter((doc) => !(doc.id === docId && doc.document_type === docType)) - ); + setMentionedDocuments((prev) => { + if (!docType) { + // Defensive fallback: keep UI in sync even when chip type is unavailable. + return prev.filter((doc) => doc.id !== docId); + } + const removedKey = getMentionDocKey({ id: docId, document_type: docType }); + return prev.filter((doc) => getMentionDocKey(doc) !== removedKey); + }); }, [setMentionedDocuments] ); const handleDocumentsMention = useCallback( (documents: Pick[]) => { - const existingKeys = new Set(mentionedDocuments.map((d) => `${d.document_type}:${d.id}`)); - const newDocs = documents.filter( - (doc) => !existingKeys.has(`${doc.document_type}:${doc.id}`) - ); + const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? []; + const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc))); - for (const doc of newDocs) { + for (const doc of documents) { + const key = getMentionDocKey(doc); + if (editorDocKeys.has(key)) continue; editorRef.current?.insertDocumentChip(doc); } setMentionedDocuments((prev) => { - const existingKeySet = new Set(prev.map((d) => `${d.document_type}:${d.id}`)); - const uniqueNewDocs = documents.filter( - (doc) => !existingKeySet.has(`${doc.document_type}:${doc.id}`) - ); + const existingKeySet = new Set(prev.map((d) => getMentionDocKey(d))); + const uniqueNewDocs = documents.filter((doc) => !existingKeySet.has(getMentionDocKey(doc))); return [...prev, ...uniqueNewDocs]; }); setMentionQuery(""); }, - [mentionedDocuments, setMentionedDocuments] + [setMentionedDocuments] ); + useEffect(() => { + const editor = editorRef.current; + const nextDocsMap = new Map(mentionedDocuments.map((doc) => [getMentionDocKey(doc), doc])); + const prevDocsMap = prevMentionedDocsRef.current; + + if (!editor) { + prevMentionedDocsRef.current = nextDocsMap; + return; + } + + const editorKeys = new Set(editor.getMentionedDocuments().map(getMentionDocKey)); + + for (const [key, doc] of nextDocsMap) { + if (prevDocsMap.has(key) || editorKeys.has(key)) continue; + editor.insertDocumentChip(doc, { removeTriggerText: false }); + } + + for (const [key, doc] of prevDocsMap) { + if (!nextDocsMap.has(key)) { + editor.removeDocumentChip(doc.id, doc.document_type); + } + } + + prevMentionedDocsRef.current = nextDocsMap; + }, [mentionedDocuments]); + return ( = ({ isBlockedByOtherUser = false }) => { const mentionedDocuments = useAtomValue(mentionedDocumentsAtom); - const sidebarDocs = useAtomValue(sidebarSelectedDocumentsAtom); - const setDocumentsSidebarOpen = useSetAtom(documentsSidebarOpenAtom); const setConnectorDialogOpen = useSetAtom(connectorDialogOpenAtom); const [toolsPopoverOpen, setToolsPopoverOpen] = useState(false); const isDesktop = useMediaQuery("(min-width: 640px)"); @@ -1226,15 +1252,6 @@ const ComposerAction: FC = ({ isBlockedByOtherUser = false )} - {sidebarDocs.length > 0 && ( - - )}
{!hasModelConfigured && (
diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx index 86863a501..fb7212119 100644 --- a/surfsense_web/components/assistant-ui/user-message.tsx +++ b/surfsense_web/components/assistant-ui/user-message.tsx @@ -1,11 +1,12 @@ import { ActionBarPrimitive, AuiIf, MessagePrimitive, useAuiState } from "@assistant-ui/react"; import { useAtomValue } from "jotai"; -import { CheckIcon, CopyIcon, FileText, Pencil } from "lucide-react"; +import { CheckIcon, CopyIcon, Pencil } from "lucide-react"; import Image from "next/image"; import { type FC, useState } from "react"; import { currentThreadAtom } from "@/atoms/chat/current-thread.atom"; import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; interface AuthorMetadata { displayName: string | null; @@ -48,6 +49,19 @@ const UserAvatar: FC = ({ displayName, avatarUrl }) => { export const UserMessage: FC = () => { const messageId = useAuiState(({ message }) => message?.id); + const messageText = useAuiState(({ message }) => + (message?.content ?? []) + .map((part) => + typeof part === "object" && + part !== null && + "type" in part && + (part as { type?: string }).type === "text" && + "text" in part + ? String((part as { text?: string }).text ?? "") + : "" + ) + .join("") + ); const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom); const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined; const metadata = useAuiState(({ message }) => message?.metadata); @@ -63,22 +77,12 @@ export const UserMessage: FC = () => {
- {mentionedDocs && mentionedDocs.length > 0 && ( -
- {mentionedDocs?.map((doc) => ( - - - {doc.title} - - ))} -
- )}
- + {mentionedDocs && mentionedDocs.length > 0 ? ( + + ) : ( + + )}
@@ -95,6 +99,64 @@ export const UserMessage: FC = () => { ); }; +const UserMessageWithMentionChips: FC<{ + text: string; + mentionedDocs: { id: number; title: string; document_type: string }[]; +}> = ({ text, mentionedDocs }) => { + type Segment = + | { type: "text"; value: string; start: number } + | { type: "mention"; doc: { id: number; title: string; document_type: string }; start: number }; + + const tokens = mentionedDocs + .map((doc) => ({ doc, token: `@${doc.title}` })) + .sort((a, b) => b.token.length - a.token.length); + + const segments: Segment[] = []; + let i = 0; + let buffer = ""; + let bufferStart = 0; + while (i < text.length) { + const tokenMatch = tokens.find(({ token }) => text.startsWith(token, i)); + if (tokenMatch) { + if (buffer) { + segments.push({ type: "text", value: buffer, start: bufferStart }); + buffer = ""; + } + segments.push({ type: "mention", doc: tokenMatch.doc, start: i }); + i += tokenMatch.token.length; + bufferStart = i; + continue; + } + if (!buffer) bufferStart = i; + buffer += text[i]; + i += 1; + } + if (buffer) { + segments.push({ type: "text", value: buffer, start: bufferStart }); + } + + return ( + + {segments.map((segment) => + segment.type === "text" ? ( + {segment.value} + ) : ( + + + {getConnectorIcon(segment.doc.document_type ?? "UNKNOWN", "h-3 w-3")} + + {segment.doc.title} + + ) + )} + + ); +}; + const UserActionBar: FC = () => { const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); diff --git a/surfsense_web/components/documents/FolderTreeView.tsx b/surfsense_web/components/documents/FolderTreeView.tsx index 9b7a393d8..2063fbee5 100644 --- a/surfsense_web/components/documents/FolderTreeView.tsx +++ b/surfsense_web/components/documents/FolderTreeView.tsx @@ -7,6 +7,7 @@ import { DndProvider } from "react-dnd"; import { HTML5Backend } from "react-dnd-html5-backend"; import { renamingFolderIdAtom } from "@/atoms/documents/folder.atoms"; import type { DocumentTypeEnum } from "@/contracts/types/document.types"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { DocumentNode, type DocumentNodeDoc } from "./DocumentNode"; import { type FolderDisplay, FolderNode } from "./FolderNode"; @@ -17,7 +18,7 @@ interface FolderTreeViewProps { documents: DocumentNodeDoc[]; expandedIds: Set; onToggleExpand: (folderId: number) => void; - mentionedDocIds: Set; + mentionedDocKeys: Set; onToggleChatMention: ( doc: { id: number; title: string; document_type: string }, isMentioned: boolean @@ -62,7 +63,7 @@ export function FolderTreeView({ documents, expandedIds, onToggleExpand, - mentionedDocIds, + mentionedDocKeys, onToggleChatMention, onToggleFolderSelect, onRenameFolder, @@ -181,7 +182,7 @@ export function FolderTreeView({ function compute(folderId: number): { selected: number; total: number } { const directDocs = (docsByFolder[folderId] ?? []).filter(isSelectable); - let selected = directDocs.filter((d) => mentionedDocIds.has(d.id)).length; + let selected = directDocs.filter((d) => mentionedDocKeys.has(getMentionDocKey(d))).length; let total = directDocs.length; for (const child of foldersByParent[folderId] ?? []) { @@ -202,7 +203,7 @@ export function FolderTreeView({ if (states[f.id] === undefined) compute(f.id); } return states; - }, [folders, docsByFolder, foldersByParent, mentionedDocIds]); + }, [folders, docsByFolder, foldersByParent, mentionedDocKeys]); const folderMap = useMemo(() => { const map: Record = {}; @@ -276,7 +277,7 @@ export function FolderTreeView({ key={`doc-${d.id}`} doc={d} depth={depth} - isMentioned={mentionedDocIds.has(d.id)} + isMentioned={mentionedDocKeys.has(getMentionDocKey(d))} onToggleChatMention={onToggleChatMention} onPreview={onPreviewDocument} onEdit={onEditDocument} @@ -356,7 +357,7 @@ export function FolderTreeView({ key={`doc-${d.id}`} doc={d} depth={depth} - isMentioned={mentionedDocIds.has(d.id)} + isMentioned={mentionedDocKeys.has(getMentionDocKey(d))} onToggleChatMention={onToggleChatMention} onPreview={onPreviewDocument} onEdit={onEditDocument} diff --git a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx index 0a147f7b7..d20aea2cd 100644 --- a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx +++ b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx @@ -23,7 +23,9 @@ import { useTranslations } from "next-intl"; import type React from "react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; -import { sidebarSelectedDocumentsAtom } from "@/atoms/chat/mentioned-documents.atom"; +import { + mentionedDocumentsAtom, +} from "@/atoms/chat/mentioned-documents.atom"; import { connectorDialogOpenAtom } from "@/atoms/connector-dialog/connector-dialog.atoms"; import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms"; import { deleteDocumentMutationAtom } from "@/atoms/documents/document-mutation.atoms"; @@ -72,6 +74,7 @@ import type { DocumentTypeEnum } from "@/contracts/types/document.types"; import { useDebouncedValue } from "@/hooks/use-debounced-value"; import { useMediaQuery } from "@/hooks/use-media-query"; import { useElectronAPI, usePlatform } from "@/hooks/use-platform"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { anonymousChatApiService } from "@/lib/apis/anonymous-chat-api.service"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { foldersApiService } from "@/lib/apis/folders-api.service"; @@ -425,8 +428,11 @@ function AuthenticatedDocumentsSidebarBase({ }, [refreshWatchedIds]); const { mutateAsync: deleteDocumentMutation } = useAtomValue(deleteDocumentMutationAtom); - const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); - const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); + const [sidebarDocs, setSidebarDocs] = useAtom(mentionedDocumentsAtom); + const mentionedDocKeys = useMemo( + () => new Set(sidebarDocs.map((d) => getMentionDocKey(d))), + [sidebarDocs] + ); // Folder state const [expandedFolderMap, setExpandedFolderMap] = useAtom(expandedFolderIdsAtom); @@ -874,11 +880,12 @@ function AuthenticatedDocumentsSidebarBase({ const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { + const key = getMentionDocKey(doc); if (isMentioned) { - setSidebarDocs((prev) => prev.filter((d) => d.id !== doc.id)); + setSidebarDocs((prev) => prev.filter((d) => getMentionDocKey(d) !== key)); } else { setSidebarDocs((prev) => { - if (prev.some((d) => d.id === doc.id)) return prev; + if (prev.some((d) => getMentionDocKey(d) === key)) return prev; return [ ...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, @@ -909,9 +916,9 @@ function AuthenticatedDocumentsSidebarBase({ if (selectAll) { setSidebarDocs((prev) => { - const existingIds = new Set(prev.map((d) => d.id)); + const existingDocKeys = new Set(prev.map((d) => getMentionDocKey(d))); const newDocs = subtreeDocs - .filter((d) => !existingIds.has(d.id)) + .filter((d) => !existingDocKeys.has(getMentionDocKey(d))) .map((d) => ({ id: d.id, title: d.title, @@ -920,8 +927,8 @@ function AuthenticatedDocumentsSidebarBase({ return newDocs.length > 0 ? [...prev, ...newDocs] : prev; }); } else { - const idsToRemove = new Set(subtreeDocs.map((d) => d.id)); - setSidebarDocs((prev) => prev.filter((d) => !idsToRemove.has(d.id))); + const keysToRemove = new Set(subtreeDocs.map((d) => getMentionDocKey(d))); + setSidebarDocs((prev) => prev.filter((d) => !keysToRemove.has(getMentionDocKey(d)))); } }, [treeDocuments, foldersByParent, setSidebarDocs] @@ -1157,7 +1164,7 @@ function AuthenticatedDocumentsSidebarBase({ documents={searchFilteredDocuments} expandedIds={expandedIds} onToggleExpand={toggleFolderExpand} - mentionedDocIds={mentionedDocIds} + mentionedDocKeys={mentionedDocKeys} onToggleChatMention={handleToggleChatMention} onToggleFolderSelect={handleToggleFolderSelect} onRenameFolder={handleRenameFolder} @@ -1585,16 +1592,20 @@ function AnonymousDocumentsSidebar({ const [isUploading, setIsUploading] = useState(false); const [search, setSearch] = useState(""); - const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); - const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); + const [sidebarDocs, setSidebarDocs] = useAtom(mentionedDocumentsAtom); + const mentionedDocKeys = useMemo( + () => new Set(sidebarDocs.map((d) => getMentionDocKey(d))), + [sidebarDocs] + ); const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { + const key = getMentionDocKey(doc); if (isMentioned) { - setSidebarDocs((prev) => prev.filter((d) => d.id !== doc.id)); + setSidebarDocs((prev) => prev.filter((d) => getMentionDocKey(d) !== key)); } else { setSidebarDocs((prev) => { - if (prev.some((d) => d.id === doc.id)) return prev; + if (prev.some((d) => getMentionDocKey(d) === key)) return prev; return [ ...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, @@ -1814,7 +1825,7 @@ function AnonymousDocumentsSidebar({ documents={searchFilteredDocs} expandedIds={new Set()} onToggleExpand={() => {}} - mentionedDocIds={mentionedDocIds} + mentionedDocKeys={mentionedDocKeys} onToggleChatMention={handleToggleChatMention} onToggleFolderSelect={() => {}} onRenameFolder={() => gate("rename folders")} diff --git a/surfsense_web/lib/chat/mention-doc-key.ts b/surfsense_web/lib/chat/mention-doc-key.ts new file mode 100644 index 000000000..5dfa11ea3 --- /dev/null +++ b/surfsense_web/lib/chat/mention-doc-key.ts @@ -0,0 +1,8 @@ +type MentionKeyInput = { + id: number; + document_type?: string | null; +}; + +export function getMentionDocKey(doc: MentionKeyInput): string { + return `${doc.document_type ?? "UNKNOWN"}:${doc.id}`; +}