feat: add InlineMentionEditor component for document mentions

- Introduced InlineMentionEditor to allow users to mention documents inline using '@'.
- Integrated the new editor into the Composer component, replacing the previous textarea input.
- Implemented functionality for handling document chips, mention triggers, and document removal.
- Enhanced user experience with real-time updates and improved keyboard navigation for mentions.

surfsense_web/components/assistant-ui/inline-mention-editor.tsx

@@ -0,0 +1,487 @@
+"use client";
+
+import { X } from "lucide-react";
+import {
+	forwardRef,
+	useCallback,
+	useEffect,
+	useImperativeHandle,
+	useRef,
+	useState,
+} from "react";
+import type { Document } from "@/contracts/types/document.types";
+import { cn } from "@/lib/utils";
+
+export interface MentionedDocument {
+	id: number;
+	title: string;
+	document_type?: string;
+}
+
+export interface InlineMentionEditorRef {
+	focus: () => void;
+	clear: () => void;
+	getText: () => string;
+	getMentionedDocuments: () => MentionedDocument[];
+	insertDocumentChip: (doc: Document) => void;
+}
+
+interface InlineMentionEditorProps {
+	placeholder?: string;
+	onMentionTrigger?: (query: string) => void;
+	onMentionClose?: () => void;
+	onSubmit?: () => void;
+	onChange?: (text: string, docs: MentionedDocument[]) => void;
+	onDocumentRemove?: (docId: number) => void;
+	onKeyDown?: (e: React.KeyboardEvent) => void;
+	disabled?: boolean;
+	className?: string;
+	initialDocuments?: MentionedDocument[];
+}
+
+// Unique data attribute to identify chip elements
+const CHIP_DATA_ATTR = "data-mention-chip";
+const CHIP_ID_ATTR = "data-mention-id";
+
+export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMentionEditorProps>(
+	(
+		{
+			placeholder = "Type @ to mention documents...",
+			onMentionTrigger,
+			onMentionClose,
+			onSubmit,
+			onChange,
+			onDocumentRemove,
+			onKeyDown,
+			disabled = false,
+			className,
+			initialDocuments = [],
+		},
+		ref
+	) => {
+		const editorRef = useRef<HTMLDivElement>(null);
+		const [isEmpty, setIsEmpty] = useState(true);
+		const [mentionedDocs, setMentionedDocs] = useState<Map<number, MentionedDocument>>(
+			() => new Map(initialDocuments.map((d) => [d.id, d]))
+		);
+		const isComposingRef = useRef(false);
+		const lastCaretPositionRef = useRef<{ node: Node; offset: number } | null>(null);
+
+		// Sync initial documents
+		useEffect(() => {
+			if (initialDocuments.length > 0) {
+				setMentionedDocs(new Map(initialDocuments.map((d) => [d.id, d])));
+			}
+		}, [initialDocuments]);
+
+		// Save caret position before any operations that might lose it
+		const saveCaretPosition = useCallback(() => {
+			const selection = window.getSelection();
+			if (selection && selection.rangeCount > 0) {
+				const range = selection.getRangeAt(0);
+				lastCaretPositionRef.current = {
+					node: range.startContainer,
+					offset: range.startOffset,
+				};
+			}
+		}, []);
+
+		// Restore caret position
+		const restoreCaretPosition = useCallback(() => {
+			if (lastCaretPositionRef.current && editorRef.current) {
+				const { node, offset } = lastCaretPositionRef.current;
+				try {
+					const selection = window.getSelection();
+					const range = document.createRange();
+					range.setStart(node, offset);
+					range.collapse(true);
+					selection?.removeAllRanges();
+					selection?.addRange(range);
+				} catch {
+					// Node might not exist anymore, focus at end
+					focusAtEnd();
+				}
+			}
+		}, []);
+
+		// Focus at the end of the editor
+		const focusAtEnd = useCallback(() => {
+			if (!editorRef.current) return;
+			editorRef.current.focus();
+			const selection = window.getSelection();
+			const range = document.createRange();
+			range.selectNodeContents(editorRef.current);
+			range.collapse(false);
+			selection?.removeAllRanges();
+			selection?.addRange(range);
+		}, []);
+
+		// Get plain text content (excluding chips)
+		const getText = useCallback((): string => {
+			if (!editorRef.current) return "";
+			
+			let text = "";
+			const walker = document.createTreeWalker(
+				editorRef.current,
+				NodeFilter.SHOW_TEXT | NodeFilter.SHOW_ELEMENT,
+				{
+					acceptNode: (node) => {
+						// Skip chip elements entirely
+						if (node.nodeType === Node.ELEMENT_NODE) {
+							const el = node as Element;
+							if (el.hasAttribute(CHIP_DATA_ATTR)) {
+								return NodeFilter.FILTER_REJECT; // Skip this subtree
+							}
+							return NodeFilter.FILTER_SKIP; // Continue into children
+						}
+						return NodeFilter.FILTER_ACCEPT;
+					},
+				}
+			);
+
+			let node: Node | null;
+			while ((node = walker.nextNode())) {
+				if (node.nodeType === Node.TEXT_NODE) {
+					text += node.textContent;
+				}
+			}
+
+			return text.trim();
+		}, []);
+
+		// Get all mentioned documents
+		const getMentionedDocuments = useCallback((): MentionedDocument[] => {
+			return Array.from(mentionedDocs.values());
+		}, [mentionedDocs]);
+
+		// Create a chip element for a document
+		const createChipElement = useCallback((doc: MentionedDocument): HTMLSpanElement => {
+			const chip = document.createElement("span");
+			chip.setAttribute(CHIP_DATA_ATTR, "true");
+			chip.setAttribute(CHIP_ID_ATTR, String(doc.id));
+			chip.contentEditable = "false";
+			chip.className =
+				"inline-flex items-center gap-0.5 mx-0.5 px-1 rounded bg-primary/10 text-xs font-medium text-primary border border-primary/20 select-none";
+			chip.style.userSelect = "none";
+			chip.style.verticalAlign = "baseline";
+
+			const titleSpan = document.createElement("span");
+			titleSpan.className = "max-w-[80px] truncate";
+			titleSpan.textContent = doc.title;
+			titleSpan.title = doc.title;
+
+			const removeBtn = document.createElement("button");
+			removeBtn.type = "button";
+			removeBtn.className =
+				"size-3 flex items-center justify-center rounded-full hover:bg-primary/20 transition-colors ml-0.5";
+			removeBtn.innerHTML = `<svg xmlns="http://www.w3.org/2000/svg" width="8" height="8" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2.5" stroke-linecap="round" stroke-linejoin="round"><path d="M18 6 6 18"/><path d="m6 6 12 12"/></svg>`;
+			removeBtn.onclick = (e) => {
+				e.preventDefault();
+				e.stopPropagation();
+				chip.remove();
+				setMentionedDocs((prev) => {
+					const next = new Map(prev);
+					next.delete(doc.id);
+					return next;
+				});
+				// Notify parent that a document was removed
+				onDocumentRemove?.(doc.id);
+				focusAtEnd();
+			};
+
+			chip.appendChild(titleSpan);
+			chip.appendChild(removeBtn);
+
+			return chip;
+		}, [focusAtEnd, onDocumentRemove]);
+
+		// Insert a document chip at the current cursor position
+		const insertDocumentChip = useCallback(
+			(doc: Document) => {
+				if (!editorRef.current) return;
+
+				const mentionDoc: MentionedDocument = {
+					id: doc.id,
+					title: doc.title,
+					document_type: doc.document_type,
+				};
+
+				// Add to mentioned docs map
+				setMentionedDocs((prev) => new Map(prev).set(doc.id, mentionDoc));
+
+				// Find and remove the @query text
+				const selection = window.getSelection();
+				if (!selection || selection.rangeCount === 0) {
+					// No selection, just append
+					const chip = createChipElement(mentionDoc);
+					editorRef.current.appendChild(chip);
+					editorRef.current.appendChild(document.createTextNode(" "));
+					focusAtEnd();
+					return;
+				}
+
+				// Find the @ symbol before the cursor and remove it along with any query text
+				const range = selection.getRangeAt(0);
+				const textNode = range.startContainer;
+
+				if (textNode.nodeType === Node.TEXT_NODE) {
+					const text = textNode.textContent || "";
+					const cursorPos = range.startOffset;
+
+					// Find the @ symbol before cursor
+					let atIndex = -1;
+					for (let i = cursorPos - 1; i >= 0; i--) {
+						if (text[i] === "@") {
+							atIndex = i;
+							break;
+						}
+					}
+
+					if (atIndex !== -1) {
+						// Remove @query and insert chip
+						const beforeAt = text.slice(0, atIndex);
+						const afterCursor = text.slice(cursorPos);
+
+						// Create chip
+						const chip = createChipElement(mentionDoc);
+
+						// Replace text node content
+						const parent = textNode.parentNode;
+						if (parent) {
+							const beforeNode = document.createTextNode(beforeAt);
+							const afterNode = document.createTextNode(" " + afterCursor);
+
+							parent.insertBefore(beforeNode, textNode);
+							parent.insertBefore(chip, textNode);
+							parent.insertBefore(afterNode, textNode);
+							parent.removeChild(textNode);
+
+							// Set cursor after the chip
+							const newRange = document.createRange();
+							newRange.setStart(afterNode, 1);
+							newRange.collapse(true);
+							selection.removeAllRanges();
+							selection.addRange(newRange);
+						}
+					} else {
+						// No @ found, just insert at cursor
+						const chip = createChipElement(mentionDoc);
+						range.insertNode(chip);
+						range.setStartAfter(chip);
+						range.collapse(true);
+						
+						// Add space after chip
+						const space = document.createTextNode(" ");
+						range.insertNode(space);
+						range.setStartAfter(space);
+						range.collapse(true);
+					}
+				} else {
+					// Not in a text node, append to editor
+					const chip = createChipElement(mentionDoc);
+					editorRef.current.appendChild(chip);
+					editorRef.current.appendChild(document.createTextNode(" "));
+					focusAtEnd();
+				}
+
+				// Update empty state
+				setIsEmpty(false);
+
+				// Trigger onChange
+				if (onChange) {
+					setTimeout(() => {
+						onChange(getText(), getMentionedDocuments());
+					}, 0);
+				}
+			},
+			[createChipElement, focusAtEnd, getText, getMentionedDocuments, onChange]
+		);
+
+		// Clear the editor
+		const clear = useCallback(() => {
+			if (editorRef.current) {
+				editorRef.current.innerHTML = "";
+				setIsEmpty(true);
+				setMentionedDocs(new Map());
+			}
+		}, []);
+
+		// Expose methods via ref
+		useImperativeHandle(ref, () => ({
+			focus: () => editorRef.current?.focus(),
+			clear,
+			getText,
+			getMentionedDocuments,
+			insertDocumentChip,
+		}));
+
+		// Handle input changes
+		const handleInput = useCallback(() => {
+			if (!editorRef.current) return;
+
+			const text = getText();
+			const empty = text.length === 0 && mentionedDocs.size === 0;
+			setIsEmpty(empty);
+
+			// Check for @ mentions
+			const selection = window.getSelection();
+			if (selection && selection.rangeCount > 0) {
+				const range = selection.getRangeAt(0);
+				const textNode = range.startContainer;
+
+				if (textNode.nodeType === Node.TEXT_NODE) {
+					const textContent = textNode.textContent || "";
+					const cursorPos = range.startOffset;
+
+					// Look for @ before cursor
+					let atIndex = -1;
+					for (let i = cursorPos - 1; i >= 0; i--) {
+						if (textContent[i] === "@") {
+							atIndex = i;
+							break;
+						}
+						// Stop if we hit a space (@ must be at word boundary)
+						if (textContent[i] === " " || textContent[i] === "\n") {
+							break;
+						}
+					}
+
+					if (atIndex !== -1) {
+						const query = textContent.slice(atIndex + 1, cursorPos);
+						// Only trigger if query doesn't start with space
+						if (!query.startsWith(" ")) {
+							onMentionTrigger?.(query);
+						} else {
+							onMentionClose?.();
+						}
+					} else {
+						onMentionClose?.();
+					}
+				}
+			}
+
+			// Notify parent of change
+			onChange?.(text, Array.from(mentionedDocs.values()));
+		}, [getText, mentionedDocs, onChange, onMentionTrigger, onMentionClose]);
+
+		// Handle keydown
+		const handleKeyDown = useCallback(
+			(e: React.KeyboardEvent<HTMLDivElement>) => {
+				// Let parent handle navigation keys when mention popover is open
+				if (onKeyDown) {
+					onKeyDown(e);
+					if (e.defaultPrevented) return;
+				}
+
+				// Handle Enter for submit (without shift)
+				if (e.key === "Enter" && !e.shiftKey) {
+					e.preventDefault();
+					onSubmit?.();
+					return;
+				}
+
+				// Handle backspace on chips
+				if (e.key === "Backspace") {
+					const selection = window.getSelection();
+					if (selection && selection.rangeCount > 0) {
+						const range = selection.getRangeAt(0);
+						if (range.collapsed) {
+							// Check if cursor is right after a chip
+							const node = range.startContainer;
+							const offset = range.startOffset;
+
+							if (node.nodeType === Node.TEXT_NODE && offset === 0) {
+								// Check previous sibling
+								const prevSibling = node.previousSibling;
+								if (prevSibling && (prevSibling as Element).hasAttribute?.(CHIP_DATA_ATTR)) {
+									e.preventDefault();
+									const chipId = Number((prevSibling as Element).getAttribute(CHIP_ID_ATTR));
+									prevSibling.parentNode?.removeChild(prevSibling);
+									setMentionedDocs((prev) => {
+										const next = new Map(prev);
+										next.delete(chipId);
+										return next;
+									});
+									// Notify parent that a document was removed
+									onDocumentRemove?.(chipId);
+								}
+							} else if (node.nodeType === Node.ELEMENT_NODE && offset > 0) {
+								// Check if previous child is a chip
+								const prevChild = (node as Element).childNodes[offset - 1];
+								if (prevChild && (prevChild as Element).hasAttribute?.(CHIP_DATA_ATTR)) {
+									e.preventDefault();
+									const chipId = Number((prevChild as Element).getAttribute(CHIP_ID_ATTR));
+									prevChild.parentNode?.removeChild(prevChild);
+									setMentionedDocs((prev) => {
+										const next = new Map(prev);
+										next.delete(chipId);
+										return next;
+									});
+									// Notify parent that a document was removed
+									onDocumentRemove?.(chipId);
+								}
+							}
+						}
+					}
+				}
+			},
+			[onKeyDown, onSubmit, onDocumentRemove]
+		);
+
+		// Handle paste - strip formatting
+		const handlePaste = useCallback((e: React.ClipboardEvent) => {
+			e.preventDefault();
+			const text = e.clipboardData.getData("text/plain");
+			document.execCommand("insertText", false, text);
+		}, []);
+
+		// Handle composition (for IME input)
+		const handleCompositionStart = useCallback(() => {
+			isComposingRef.current = true;
+		}, []);
+
+		const handleCompositionEnd = useCallback(() => {
+			isComposingRef.current = false;
+			handleInput();
+		}, [handleInput]);
+
+		return (
+			<div className="relative w-full">
+				<div
+					ref={editorRef}
+					contentEditable={!disabled}
+					suppressContentEditableWarning
+					onInput={handleInput}
+					onKeyDown={handleKeyDown}
+					onPaste={handlePaste}
+					onCompositionStart={handleCompositionStart}
+					onCompositionEnd={handleCompositionEnd}
+					className={cn(
+						"min-h-[24px] max-h-32 overflow-y-auto",
+						"text-sm outline-none",
+						"whitespace-pre-wrap break-words",
+						disabled && "opacity-50 cursor-not-allowed",
+						className
+					)}
+					style={{ wordBreak: "break-word" }}
+					data-placeholder={placeholder}
+					aria-label="Message input with inline mentions"
+					role="textbox"
+					aria-multiline="true"
+				/>
+				{/* Placeholder */}
+				{isEmpty && (
+					<div
+						className="absolute top-0 left-0 pointer-events-none text-muted-foreground text-sm"
+						aria-hidden="true"
+					>
+						{placeholder}
+					</div>
+				)}
+			</div>
+		);
+	}
+);
+
+InlineMentionEditor.displayName = "InlineMentionEditor";
+

surfsense_web/components/assistant-ui/thread.tsx

@@ -7,7 +7,7 @@ import {
 	MessagePrimitive,
 	ThreadPrimitive,
 	useAssistantState,
-	useMessage,
+	useComposerRuntime,
 	useThreadViewport,
 } from "@assistant-ui/react";
 import { useAtom, useAtomValue, useSetAtom } from "jotai";
@@ -31,7 +31,6 @@ import {
 	Search,
 	Sparkles,
 	SquareIcon,
-	X,
 } from "lucide-react";
 import Link from "next/link";
 import { useParams } from "next/navigation";
@@ -65,6 +64,10 @@ import {
 	ComposerAttachments,
 	UserMessageAttachments,
 } from "@/components/assistant-ui/attachment";
+import {
+	InlineMentionEditor,
+	type InlineMentionEditorRef,
+} from "@/components/assistant-ui/inline-mention-editor";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
@@ -240,7 +243,7 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
  * Uses useThreadViewport to scroll to bottom when thinking steps change,
  * ensuring the user always sees the latest content during streaming.
  */
-const ThinkingStepsScrollHandler: FC = () => {
+const _ThinkingStepsScrollHandler: FC = () => {
 	const thinkingStepsMap = useContext(ThinkingStepsContext);
 	const viewport = useThreadViewport();
 	const isRunning = useAssistantState(({ thread }) => thread.isRunning);
@@ -412,177 +415,140 @@ const Composer: FC = () => {
 	const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
 	const [showDocumentPopover, setShowDocumentPopover] = useState(false);
 	const [mentionQuery, setMentionQuery] = useState("");
-	const inputRef = useRef<HTMLTextAreaElement | null>(null);
+	const editorRef = useRef<InlineMentionEditorRef>(null);
+	const editorContainerRef = useRef<HTMLDivElement>(null);
 	const documentPickerRef = useRef<DocumentsDataTableRef>(null);
 	const { search_space_id } = useParams();
 	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
+	const composerRuntime = useComposerRuntime();
 
 	// Sync mentioned document IDs to atom for use in chat request
 	useEffect(() => {
 		setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id));
 	}, [mentionedDocuments, setMentionedDocumentIds]);
 
-	// Extract mention query (text after @)
-	const extractMentionQuery = useCallback((value: string): string => {
-		const atIndex = value.lastIndexOf("@");
-		if (atIndex === -1) return "";
-		return value.slice(atIndex + 1);
+	// Handle text change from inline editor - sync with assistant-ui composer
+	const handleEditorChange = useCallback(
+		(text: string) => {
+			composerRuntime.setText(text);
+		},
+		[composerRuntime]
+	);
+
+	// Handle @ mention trigger from inline editor
+	const handleMentionTrigger = useCallback((query: string) => {
+		setShowDocumentPopover(true);
+		setMentionQuery(query);
 	}, []);
 
-	const handleKeyUp = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-		const textarea = e.currentTarget;
-		const value = textarea.value;
-
-		// Open document picker when user types '@'
-		if (e.key === "@" || (e.key === "2" && e.shiftKey)) {
-			setShowDocumentPopover(true);
-			setMentionQuery("");
-			return;
-		}
-
-		// Check if value contains @ and extract query
-		if (value.includes("@")) {
-			const query = extractMentionQuery(value);
-
-			// Close popup if query starts with space (user typed "@ ")
-			if (query.startsWith(" ")) {
-				setShowDocumentPopover(false);
-				setMentionQuery("");
-				return;
-			}
-
-			// Reopen popup if @ is present and query doesn't start with space
-			// (handles case where user deleted the space after @)
-			if (!showDocumentPopover) {
-				setShowDocumentPopover(true);
-			}
-			setMentionQuery(query);
-		} else {
-			// Close popover if '@' is no longer in the input (user deleted it)
-			if (showDocumentPopover) {
-				setShowDocumentPopover(false);
-				setMentionQuery("");
-			}
-		}
-	};
-
-	const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-		// When popup is open, handle navigation keys
+	// Handle mention close
+	const handleMentionClose = useCallback(() => {
 		if (showDocumentPopover) {
-			if (e.key === "ArrowDown") {
-				e.preventDefault();
-				documentPickerRef.current?.moveDown();
-				return;
-			}
-			if (e.key === "ArrowUp") {
-				e.preventDefault();
-				documentPickerRef.current?.moveUp();
-				return;
-			}
-			if (e.key === "Enter") {
-				e.preventDefault();
-				documentPickerRef.current?.selectHighlighted();
-				return;
-			}
-			if (e.key === "Escape") {
-				e.preventDefault();
-				setShowDocumentPopover(false);
-				setMentionQuery("");
-				return;
-			}
+			setShowDocumentPopover(false);
+			setMentionQuery("");
 		}
+	}, [showDocumentPopover]);
 
-		// Remove last document chip when pressing backspace at the beginning of input
-		if (e.key === "Backspace" && mentionedDocuments.length > 0) {
-			const textarea = e.currentTarget;
-			const selectionStart = textarea.selectionStart;
-			const selectionEnd = textarea.selectionEnd;
-
-			// Only remove chip if cursor is at position 0 and nothing is selected
-			if (selectionStart === 0 && selectionEnd === 0) {
-				e.preventDefault();
-				// Remove the last document chip
-				setMentionedDocuments((prev) => prev.slice(0, -1));
-			}
-		}
-	};
-
-	const handleDocumentsMention = (documents: Document[]) => {
-		// Update mentioned documents (merge with existing, avoid duplicates)
-		setMentionedDocuments((prev) => {
-			const existingIds = new Set(prev.map((d) => d.id));
-			const newDocs = documents.filter((doc) => !existingIds.has(doc.id));
-			return [...prev, ...newDocs];
-		});
-
-		// Clean up the '@...' mention text from input
-		if (inputRef.current) {
-			const input = inputRef.current;
-			const currentValue = input.value;
-			const atIndex = currentValue.lastIndexOf("@");
-
-			if (atIndex !== -1) {
-				// Remove @ and everything after it
-				const newValue = currentValue.slice(0, atIndex);
-				const nativeInputValueSetter = Object.getOwnPropertyDescriptor(
-					window.HTMLTextAreaElement.prototype,
-					"value"
-				)?.set;
-				if (nativeInputValueSetter) {
-					nativeInputValueSetter.call(input, newValue);
-					input.dispatchEvent(new Event("input", { bubbles: true }));
+	// Handle keyboard navigation when popover is open
+	const handleKeyDown = useCallback(
+		(e: React.KeyboardEvent) => {
+			if (showDocumentPopover) {
+				if (e.key === "ArrowDown") {
+					e.preventDefault();
+					documentPickerRef.current?.moveDown();
+					return;
+				}
+				if (e.key === "ArrowUp") {
+					e.preventDefault();
+					documentPickerRef.current?.moveUp();
+					return;
+				}
+				if (e.key === "Enter") {
+					e.preventDefault();
+					documentPickerRef.current?.selectHighlighted();
+					return;
+				}
+				if (e.key === "Escape") {
+					e.preventDefault();
+					setShowDocumentPopover(false);
+					setMentionQuery("");
+					return;
 				}
 			}
-			// Focus the input so user can continue typing
-			input.focus();
+		},
+		[showDocumentPopover]
+	);
+
+	// Handle submit from inline editor (Enter key)
+	const handleSubmit = useCallback(() => {
+		if (!showDocumentPopover) {
+			composerRuntime.send();
+			// Clear the editor after sending
+			editorRef.current?.clear();
+			setMentionedDocuments([]);
+			setMentionedDocumentIds([]);
 		}
+	}, [showDocumentPopover, composerRuntime, setMentionedDocuments, setMentionedDocumentIds]);
 
-		// Reset mention query
-		setMentionQuery("");
-	};
+	// Handle document removal from inline editor
+	const handleDocumentRemove = useCallback(
+		(docId: number) => {
+			setMentionedDocuments((prev) => {
+				const updated = prev.filter((doc) => doc.id !== docId);
+				// Immediately sync document IDs to avoid race conditions
+				setMentionedDocumentIds(updated.map((doc) => doc.id));
+				return updated;
+			});
+		},
+		[setMentionedDocuments, setMentionedDocumentIds]
+	);
 
-	const handleRemoveDocument = (docId: number) => {
-		setMentionedDocuments((prev) => prev.filter((doc) => doc.id !== docId));
-	};
+	// Handle document selection from picker
+	const handleDocumentsMention = useCallback(
+		(documents: Document[]) => {
+			// Insert chips into the inline editor for each new document
+			const existingIds = new Set(mentionedDocuments.map((d) => d.id));
+			const newDocs = documents.filter((doc) => !existingIds.has(doc.id));
+
+			for (const doc of newDocs) {
+				editorRef.current?.insertDocumentChip(doc);
+			}
+
+			// Update mentioned documents state
+			setMentionedDocuments((prev) => {
+				const existingIdSet = new Set(prev.map((d) => d.id));
+				const uniqueNewDocs = documents.filter((doc) => !existingIdSet.has(doc.id));
+				const updated = [...prev, ...uniqueNewDocs];
+				// Immediately sync document IDs to avoid race conditions
+				setMentionedDocumentIds(updated.map((doc) => doc.id));
+				return updated;
+			});
+
+			// Reset mention query but keep popover open for more selections
+			setMentionQuery("");
+		},
+		[mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds]
+	);
 
 	return (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
 			<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone flex w-full flex-col rounded-2xl border-input bg-muted px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:border-dashed data-[dragging=true]:bg-accent/50">
 				<ComposerAttachments />
-				{/* -------- Input field with inline document chips -------- */}
-				<div className="aui-composer-input-wrapper flex flex-wrap items-center gap-1.5 px-3 pt-2 pb-6">
-					{/* Inline document chips */}
-					{mentionedDocuments.map((doc) => (
-						<span
-							key={doc.id}
-							className="inline-flex items-center gap-1 pl-2 pr-1 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20 shrink-0"
-							title={doc.title}
-						>
-							<span className="max-w-[120px] truncate">{doc.title}</span>
-							<button
-								type="button"
-								onClick={() => handleRemoveDocument(doc.id)}
-								className="size-4 flex items-center justify-center rounded-full hover:bg-primary/20 transition-colors"
-								aria-label={`Remove ${doc.title}`}
-							>
-								<X className="size-3" />
-							</button>
-						</span>
-					))}
-					{/* Text input */}
-					<ComposerPrimitive.Input
-						ref={inputRef}
-						onKeyUp={handleKeyUp}
+				{/* -------- Inline Mention Editor -------- */}
+				<div
+					ref={editorContainerRef}
+					className="aui-composer-input-wrapper px-3 pt-3 pb-6"
+				>
+					<InlineMentionEditor
+						ref={editorRef}
+						placeholder="Ask SurfSense (type @ to mention docs)"
+						onMentionTrigger={handleMentionTrigger}
+						onMentionClose={handleMentionClose}
+						onChange={handleEditorChange}
+						onDocumentRemove={handleDocumentRemove}
+						onSubmit={handleSubmit}
 						onKeyDown={handleKeyDown}
-						placeholder={
-							mentionedDocuments.length > 0
-								? "Ask about these documents..."
-								: "Ask SurfSense (type @ to mention docs)"
-						}
-						className="aui-composer-input flex-1 min-w-[120px] max-h-32 resize-none bg-transparent text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0 py-1"
-						rows={1}
-						autoFocus
-						aria-label="Message input"
+						className="min-h-[24px]"
 					/>
 				</div>
 
@@ -605,11 +571,11 @@ const Composer: FC = () => {
 								style={{
 									zIndex: 9999,
 									backgroundColor: "#18181b",
-									bottom: inputRef.current
-										? `${window.innerHeight - inputRef.current.getBoundingClientRect().top + 8}px`
+									bottom: editorContainerRef.current
+										? `${window.innerHeight - editorContainerRef.current.getBoundingClientRect().top + 8}px`
 										: "200px",
-									left: inputRef.current
-										? `${inputRef.current.getBoundingClientRect().left}px`
+									left: editorContainerRef.current
+										? `${editorContainerRef.current.getBoundingClientRect().left}px`
 										: "50%",
 								}}
 							>