diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index ee93a409a..27a205af2 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -11,12 +11,24 @@ export const mentionedDocumentsAtom = atom[] >([]); +export interface SidebarMentionEvent { + kind: "add" | "remove"; + docs: Pick[]; + nonce: number; +} + +/** + * Event atom used to tell the composer that documents were selected/unselected + * from sidebar checkboxes, so chips can be inserted/removed in-editor. + */ +export const sidebarMentionEventAtom = atom(null); + /** * Derived read-only atom that merges @-mention chips and sidebar selections * into a single deduplicated set of document IDs for the backend. diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index 45ad219dd..31b2a0566 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -44,7 +44,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, @@ -129,6 +132,40 @@ export const InlineMentionEditor = forwardRef new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d])) ); const isComposingRef = useRef(false); + const lastSelectionRangeRef = useRef(null); + const isSelectionInsideEditor = useCallback( + (selection: Selection | null): selection is Selection => { + if (!selection || selection.rangeCount === 0 || !editorRef.current) return false; + const range = selection.getRangeAt(0); + return editorRef.current.contains(range.startContainer); + }, + [] + ); + + 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 (!lastSelectionRangeRef.current) return selection; + selection.removeAllRanges(); + selection.addRange(lastSelectionRangeRef.current.cloneRange()); + return selection; + }, []); + + useEffect(() => { + const handleSelectionChange = () => { + if (document.activeElement !== editorRef.current) return; + rememberSelection(); + }; + document.addEventListener("selectionchange", handleSelectionChange); + return () => document.removeEventListener("selectionchange", handleSelectionChange); + }, [rememberSelection]); + // Sync initial documents useEffect(() => { @@ -157,7 +194,7 @@ export const InlineMentionEditor = forwardRef { @@ -299,8 +336,12 @@ 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") { @@ -320,20 +361,23 @@ export const InlineMentionEditor = forwardRef - {/** biome-ignore lint/a11y/useSemanticElements: */}
{ // Document mention state (atoms persist across component remounts) const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom); const setSidebarDocs = useSetAtom(sidebarSelectedDocumentsAtom); + const [sidebarMentionEvent, setSidebarMentionEvent] = useAtom(sidebarMentionEventAtom); const [showDocumentPopover, setShowDocumentPopover] = useState(false); const [showPromptPicker, setShowPromptPicker] = useState(false); const [mentionQuery, setMentionQuery] = useState(""); @@ -660,6 +662,42 @@ const Composer: FC = () => { [mentionedDocuments, setMentionedDocuments] ); + useEffect(() => { + if (!sidebarMentionEvent) return; + + const eventDocs = sidebarMentionEvent.docs; + if (eventDocs.length === 0) { + setSidebarMentionEvent(null); + return; + } + + const docKey = (doc: Pick) => + `${doc.document_type}:${doc.id}`; + const mentionedKeys = new Set(mentionedDocuments.map(docKey)); + + if (sidebarMentionEvent.kind === "add") { + const docsToAdd = eventDocs.filter((doc) => !mentionedKeys.has(docKey(doc))); + for (const doc of docsToAdd) { + editorRef.current?.insertDocumentChip(doc, { removeTriggerText: false }); + } + if (docsToAdd.length > 0) { + setMentionedDocuments((prev) => { + const existing = new Set(prev.map(docKey)); + const uniqueAdds = docsToAdd.filter((doc) => !existing.has(docKey(doc))); + return uniqueAdds.length > 0 ? [...prev, ...uniqueAdds] : prev; + }); + } + } else { + const removeKeys = new Set(eventDocs.map(docKey)); + for (const doc of eventDocs) { + editorRef.current?.removeDocumentChip(doc.id, doc.document_type); + } + setMentionedDocuments((prev) => prev.filter((doc) => !removeKeys.has(docKey(doc)))); + } + + setSidebarMentionEvent(null); + }, [sidebarMentionEvent, mentionedDocuments, setMentionedDocuments, setSidebarMentionEvent]); + return ( new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); // Folder state @@ -857,19 +861,42 @@ function AuthenticatedDocumentsSidebarBase({ const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { + const key = `${doc.document_type}:${doc.id}`; if (isMentioned) { - setSidebarDocs((prev) => prev.filter((d) => d.id !== doc.id)); + setSidebarDocs((prev) => prev.filter((d) => `${d.document_type}:${d.id}` !== key)); + setSidebarMentionEvent({ + kind: "remove", + docs: [ + { + id: doc.id, + title: doc.title, + document_type: doc.document_type as DocumentTypeEnum, + }, + ], + nonce: Date.now(), + }); } else { setSidebarDocs((prev) => { - if (prev.some((d) => d.id === doc.id)) return prev; + if (prev.some((d) => `${d.document_type}:${d.id}` === key)) return prev; return [ ...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, ]; }); + setSidebarMentionEvent({ + kind: "add", + docs: [ + { + id: doc.id, + title: doc.title, + document_type: doc.document_type as DocumentTypeEnum, + }, + ], + nonce: Date.now(), + }); } }, - [setSidebarDocs] + [setSidebarDocs, setSidebarMentionEvent] ); const handleToggleFolderSelect = useCallback( @@ -891,10 +918,18 @@ function AuthenticatedDocumentsSidebarBase({ if (subtreeDocs.length === 0) return; if (selectAll) { + const existingKeys = new Set(sidebarDocs.map((d) => `${d.document_type}:${d.id}`)); + const docsToAdd = subtreeDocs + .filter((d) => !existingKeys.has(`${d.document_type}:${d.id}`)) + .map((d) => ({ + id: d.id, + title: d.title, + document_type: d.document_type as DocumentTypeEnum, + })); setSidebarDocs((prev) => { - const existingIds = new Set(prev.map((d) => d.id)); + const existingDocKeys = new Set(prev.map((d) => `${d.document_type}:${d.id}`)); const newDocs = subtreeDocs - .filter((d) => !existingIds.has(d.id)) + .filter((d) => !existingDocKeys.has(`${d.document_type}:${d.id}`)) .map((d) => ({ id: d.id, title: d.title, @@ -902,12 +937,35 @@ function AuthenticatedDocumentsSidebarBase({ })); return newDocs.length > 0 ? [...prev, ...newDocs] : prev; }); + if (docsToAdd.length > 0) { + setSidebarMentionEvent({ + kind: "add", + docs: docsToAdd, + nonce: Date.now(), + }); + } } 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) => `${d.document_type}:${d.id}`)); + const docsToRemove = sidebarDocs + .filter((d) => keysToRemove.has(`${d.document_type}:${d.id}`)) + .map((d) => ({ + id: d.id, + title: d.title, + document_type: d.document_type as DocumentTypeEnum, + })); + setSidebarDocs((prev) => + prev.filter((d) => !keysToRemove.has(`${d.document_type}:${d.id}`)) + ); + if (docsToRemove.length > 0) { + setSidebarMentionEvent({ + kind: "remove", + docs: docsToRemove, + nonce: Date.now(), + }); + } } }, - [treeDocuments, foldersByParent, setSidebarDocs] + [treeDocuments, foldersByParent, sidebarDocs, setSidebarDocs, setSidebarMentionEvent] ); const searchFilteredDocuments = useMemo(() => { @@ -1568,23 +1626,47 @@ function AnonymousDocumentsSidebar({ const [search, setSearch] = useState(""); const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); + const setSidebarMentionEvent = useSetAtom(sidebarMentionEventAtom); const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { + const key = `${doc.document_type}:${doc.id}`; if (isMentioned) { - setSidebarDocs((prev) => prev.filter((d) => d.id !== doc.id)); + setSidebarDocs((prev) => prev.filter((d) => `${d.document_type}:${d.id}` !== key)); + setSidebarMentionEvent({ + kind: "remove", + docs: [ + { + id: doc.id, + title: doc.title, + document_type: doc.document_type as DocumentTypeEnum, + }, + ], + nonce: Date.now(), + }); } else { setSidebarDocs((prev) => { - if (prev.some((d) => d.id === doc.id)) return prev; + if (prev.some((d) => `${d.document_type}:${d.id}` === key)) return prev; return [ ...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, ]; }); + setSidebarMentionEvent({ + kind: "add", + docs: [ + { + id: doc.id, + title: doc.title, + document_type: doc.document_type as DocumentTypeEnum, + }, + ], + nonce: Date.now(), + }); } }, - [setSidebarDocs] + [setSidebarDocs, setSidebarMentionEvent] ); const uploadedDoc = anonMode.isAnonymous ? anonMode.uploadedDoc : null;