From 533084b433c81a8a96c6e37e2a6c8ae08fb2baa0 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Thu, 25 Dec 2025 13:44:18 +0530 Subject: [PATCH] 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. --- .../assistant-ui/inline-mention-editor.tsx | 487 ++++++++++++++++++ .../components/assistant-ui/thread.tsx | 272 +++++----- 2 files changed, 606 insertions(+), 153 deletions(-) create mode 100644 surfsense_web/components/assistant-ui/inline-mention-editor.tsx diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx new file mode 100644 index 000000000..2e5d1938d --- /dev/null +++ b/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( + ( + { + placeholder = "Type @ to mention documents...", + onMentionTrigger, + onMentionClose, + onSubmit, + onChange, + onDocumentRemove, + onKeyDown, + disabled = false, + className, + initialDocuments = [], + }, + ref + ) => { + const editorRef = useRef(null); + const [isEmpty, setIsEmpty] = useState(true); + const [mentionedDocs, setMentionedDocs] = useState>( + () => 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 = ``; + 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) => { + // 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 ( +
+
+ {/* Placeholder */} + {isEmpty && ( + + )} +
+ ); + } +); + +InlineMentionEditor.displayName = "InlineMentionEditor"; + diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 7f6cedabc..9d69ffc77 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/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(null); + const editorRef = useRef(null); + const editorContainerRef = useRef(null); const documentPickerRef = useRef(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) => { - 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) => { - // 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 ( - {/* -------- Input field with inline document chips -------- */} -
- {/* Inline document chips */} - {mentionedDocuments.map((doc) => ( - - {doc.title} - - - ))} - {/* Text input */} - + 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]" />
@@ -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%", }} >