feat(mentions): add sidebar mention event atom and enhance document mention handling in the editor

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<Document, "id" | "title" | "document_type">) => void;
+	insertDocumentChip: (
+		doc: Pick<Document, "id" | "title" | "document_type">,
+		options?: { removeTriggerText?: boolean }
+	) => void;
 	removeDocumentChip: (docId: number, docType?: string) => void;
 	setDocumentChipStatus: (
 		docId: number,
@@ -129,6 +132,40 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			() => new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d]))
 		);
 		const isComposingRef = useRef(false);
+		const lastSelectionRangeRef = useRef<Range | null>(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<InlineMentionEditorRef, InlineMent
 			range.insertNode(anchor);
 			anchor.scrollIntoView({ block: "end" });
 			anchor.remove();
-		}, [initialText]); // eslint-disable-line react-hooks/exhaustive-deps
+		}, [initialText]);
 
 		// Focus at the end of the editor
 		const focusAtEnd = useCallback(() => {
@@ -299,8 +336,12 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 
 		// Insert a document chip at the current cursor position
 		const insertDocumentChip = useCallback(
-			(doc: Pick<Document, "id" | "title" | "document_type">) => {
+			(
+				doc: Pick<Document, "id" | "title" | "document_type">,
+				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<InlineMentionEditorRef, InlineMent
 
 				// Find and remove the @query text
 				const selection = window.getSelection();
-				if (!selection || selection.rangeCount === 0) {
+				const hasActiveSelection = isSelectionInsideEditor(selection);
+				const resolvedSelection = hasActiveSelection ? selection : restoreRememberedSelection();
+				if (!resolvedSelection || resolvedSelection.rangeCount === 0) {
 					// No selection, just append
 					const chip = createChipElement(mentionDoc);
 					editorRef.current.appendChild(chip);
 					editorRef.current.appendChild(document.createTextNode(" "));
 					focusAtEnd();
+					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 +413,9 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 							const newRange = document.createRange();
 							newRange.setStart(afterNode, 1);
 							newRange.collapse(true);
-							selection.removeAllRanges();
-							selection.addRange(newRange);
+							resolvedSelection.removeAllRanges();
+							resolvedSelection.addRange(newRange);
+							rememberSelection();
 						}
 					} else {
 						// No @ found, just insert at cursor
@@ -384,13 +429,23 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 						range.insertNode(space);
 						range.setStartAfter(space);
 						range.collapse(true);
+						resolvedSelection.removeAllRanges();
+						resolvedSelection.addRange(range);
+						rememberSelection();
 					}
 				} else {
-					// Not in a text node, append to editor
+					// Either explicit non-trigger insertion or no @query present.
 					const chip = createChipElement(mentionDoc);
-					editorRef.current.appendChild(chip);
-					editorRef.current.appendChild(document.createTextNode(" "));
-					focusAtEnd();
+					range.insertNode(chip);
+					range.setStartAfter(chip);
+					range.collapse(true);
+					const space = document.createTextNode(" ");
+					range.insertNode(space);
+					range.setStartAfter(space);
+					range.collapse(true);
+					resolvedSelection.removeAllRanges();
+					resolvedSelection.addRange(range);
+					rememberSelection();
 				}
 
 				// Update empty state
@@ -403,7 +458,16 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 					}, 0);
 				}
 			},
-			[createChipElement, focusAtEnd, getText, getMentionedDocuments, onChange]
+			[
+				createChipElement,
+				focusAtEnd,
+				getText,
+				getMentionedDocuments,
+				isSelectionInsideEditor,
+				onChange,
+				rememberSelection,
+				restoreRememberedSelection,
+			]
 		);
 
 		// Clear the editor
@@ -594,6 +658,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 
 			// Notify parent of change
 			onChange?.(text, Array.from(mentionedDocs.values()));
+			rememberSelection();
 		}, [
 			getText,
 			mentionedDocs,
@@ -602,6 +667,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			onMentionClose,
 			onActionTrigger,
 			onActionClose,
+			rememberSelection,
 		]);
 
 		// Handle keydown
@@ -713,7 +779,6 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 
 		return (
 			<div className="relative w-full">
-				{/** biome-ignore lint/a11y/useSemanticElements: <not important> */}
 				<div
 					ref={editorRef}
 					contentEditable={!disabled}
@@ -724,6 +789,9 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 					onPaste={handlePaste}
 					onCompositionStart={handleCompositionStart}
 					onCompositionEnd={handleCompositionEnd}
+					onKeyUp={rememberSelection}
+					onMouseUp={rememberSelection}
+					onBlur={rememberSelection}
 					className={cn(
 						"min-h-[24px] max-h-32 overflow-y-auto",
 						"text-sm outline-none",

surfsense_web/components/assistant-ui/thread.tsx

@@ -38,6 +38,7 @@ import {
 import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom";
 import {
 	mentionedDocumentsAtom,
+	sidebarMentionEventAtom,
 	sidebarSelectedDocumentsAtom,
 } from "@/atoms/chat/mentioned-documents.atom";
 import { connectorDialogOpenAtom } from "@/atoms/connector-dialog/connector-dialog.atoms";
@@ -336,6 +337,7 @@ const Composer: FC = () => {
 	// 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<Document, "id" | "title" | "document_type">) =>
+			`${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 (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col gap-2">
 			<ChatSessionStatus

surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx

@@ -23,7 +23,10 @@ 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 {
+	sidebarMentionEventAtom,
+	sidebarSelectedDocumentsAtom,
+} 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";
@@ -413,6 +416,7 @@ function AuthenticatedDocumentsSidebarBase({
 	const { mutateAsync: deleteDocumentMutation } = useAtomValue(deleteDocumentMutationAtom);
 
 	const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom);
+	const setSidebarMentionEvent = useSetAtom(sidebarMentionEventAtom);
 	const mentionedDocIds = useMemo(() => 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;