From 2437716752866bd588f49b27068cf3aa023a2501 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 12 May 2026 23:18:45 +0530 Subject: [PATCH] refactor(assistant-ui): enhance mention chip handling and editor focus behavior --- .../assistant-ui/inline-mention-editor.tsx | 241 ++++++++++-------- .../components/assistant-ui/thread.tsx | 108 +++++--- 2 files changed, 215 insertions(+), 134 deletions(-) diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index c2b896794..da5fe878f 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -1,6 +1,7 @@ "use client"; import { Folder as FolderIcon, X as XIcon } from "lucide-react"; +import type { NodeEntry, TElement } from "platejs"; import type { PlateElementProps } from "platejs/react"; import { createPlatePlugin, @@ -56,10 +57,7 @@ export interface InlineMentionEditorRef { setText: (text: string) => void; getText: () => string; getMentionedDocuments: () => MentionedDocument[]; - insertMentionChip: ( - mention: MentionChipInput, - options?: { removeTriggerText?: boolean } - ) => void; + insertMentionChip: (mention: MentionChipInput, options?: { removeTriggerText?: boolean }) => void; /** * @deprecated Use ``insertMentionChip``. Kept for one transition * cycle so we don't break ad-hoc callers; prefer the new name. @@ -338,17 +336,21 @@ export const InlineMentionEditor = forwardRef { - const el = editableRef.current; - if (!el) return; - el.focus(); - const selection = window.getSelection(); - const range = document.createRange(); - range.selectNodeContents(el); - range.collapse(false); - selection?.removeAllRanges(); - selection?.addRange(range); - }, []); + try { + editor.tf.select(editor.api.end([])); + editor.tf.focus(); + } catch { + editableRef.current?.focus(); + } + }, [editor]); const getCurrentValue = useCallback( () => (editor.children as ComposerValue) ?? EMPTY_VALUE, @@ -396,17 +398,30 @@ export const InlineMentionEditor = forwardRef { if (typeof mention.id !== "number" || typeof mention.title !== "string") return; const removeTriggerText = options?.removeTriggerText ?? true; - const current = getCurrentValue(); - const selection = editor.selection; const kind: MentionKind = mention.kind ?? "doc"; const document_type = - mention.document_type ?? - (kind === "folder" ? FOLDER_MENTION_DOCUMENT_TYPE : undefined); + mention.document_type ?? (kind === "folder" ? FOLDER_MENTION_DOCUMENT_TYPE : undefined); const mentionNode: MentionElementNode = { type: MENTION_TYPE, id: mention.id, @@ -416,60 +431,49 @@ export const InlineMentionEditor = forwardRef { + const selection = editor.selection; - const block = current[cursorCtx.blockIndex]; - const currentChild = getTextNode(block.children[cursorCtx.childIndex]); - if (!currentChild) { - const children = [...block.children]; - children.splice(cursorCtx.childIndex + 1, 0, mentionNode, { text: " " }); - const next = [...current]; - next[cursorCtx.blockIndex] = { ...block, children }; - setValue(next as ComposerValue); - requestAnimationFrame(focusAtEnd); - return; - } - - const text = currentChild.text; - let removeStart = cursorCtx.cursor; - if (removeTriggerText) { - for (let i = cursorCtx.cursor - 1; i >= 0; i--) { - if (text[i] === "@") { - removeStart = i; - break; + // No active editor selection — typically because focus + // moved to a picker/dropdown. Snap the caret to the end + // of the document so the chip appends cleanly instead + // of disappearing into a dead range. + if (!selection) { + editor.tf.select(editor.api.end([])); + } else if (removeTriggerText) { + // Delete the in-progress "@query" text so the chip + // stands in for it. Mirrors the old splice but lets + // Slate keep selection sane through the edit. + const cursorCtx = getCursorTextContext(getCurrentValue(), selection); + if (cursorCtx) { + const text = cursorCtx.text; + let triggerIndex = -1; + for (let i = cursorCtx.cursor - 1; i >= 0; i--) { + if (text[i] === "@") { + triggerIndex = i; + break; + } + if (text[i] === " " || text[i] === "\n") break; + } + if (triggerIndex >= 0 && triggerIndex < cursorCtx.cursor) { + const path = [cursorCtx.blockIndex, cursorCtx.childIndex]; + editor.tf.delete({ + at: { + anchor: { path, offset: triggerIndex }, + focus: { path, offset: cursorCtx.cursor }, + }, + }); + } } - if (text[i] === " " || text[i] === "\n") break; } - } - const before = text.slice(0, removeStart); - const after = text.slice(cursorCtx.cursor); - const replacement: ComposerNode[] = []; - if (before.length > 0) replacement.push({ text: before }); - replacement.push(mentionNode); - replacement.push({ text: ` ${after}` }); - - const children = [...block.children]; - children.splice(cursorCtx.childIndex, 1, ...replacement); - const next = [...current]; - next[cursorCtx.blockIndex] = { ...block, children }; - setValue(next as ComposerValue); - requestAnimationFrame(focusAtEnd); + editor.tf.insertNodes([mentionNode, { text: " " }] as unknown as TElement[], { + select: true, + }); + }); + editor.tf.focus(); }, - [editor.selection, focusAtEnd, getCurrentValue, setValue] + [editor, getCurrentValue] ); // Backwards-compatible shim — pre-folder callers pass a doc-only @@ -485,24 +489,34 @@ export const InlineMentionEditor = forwardRef { - const current = getCurrentValue(); - let changed = false; - const next = current.map((block) => { - const children = block.children.filter((node) => { - if (!isMentionNode(node)) return true; - const match = - node.id === docId && (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN"); - if (match) changed = true; - return !match; - }); - return { ...block, children: children.length ? children : [{ text: "" }] }; + const match = (n: unknown) => { + if (!n || typeof n !== "object" || !("type" in n)) return false; + const node = n as MentionElementNode; + if (node.type !== MENTION_TYPE) return false; + if (node.id !== docId) return false; + return (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN"); + }; + + const entries = Array.from(editor.api.nodes({ at: [], match })) as NodeEntry[]; + if (entries.length === 0) return; + editor.tf.withoutNormalizing(() => { + for (const [, path] of entries.reverse()) { + editor.tf.removeNodes({ at: path }); + } }); - if (!changed) return; - setValue(next as ComposerValue); }, - [getCurrentValue, setValue] + [editor] ); // Combined "remove chip end-to-end" used by both the Backspace @@ -517,6 +531,11 @@ export const InlineMentionEditor = forwardRef { - const current = getCurrentValue(); - let changed = false; - const next = current.map((block) => ({ - ...block, - children: block.children.map((node) => { - if (!isMentionNode(node)) return node; - const sameType = (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN"); - if (node.id !== docId || !sameType) return node; - changed = true; - return { - ...node, - statusLabel, - statusKind: statusLabel ? statusKind : undefined, - }; - }), - })); - if (!changed) return; - setValue(next as ComposerValue); + const match = (n: unknown) => { + if (!n || typeof n !== "object" || !("type" in n)) return false; + const node = n as MentionElementNode; + if (node.type !== MENTION_TYPE) return false; + if (node.id !== docId) return false; + return (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN"); + }; + + editor.tf.setNodes( + { + statusLabel, + statusKind: statusLabel ? statusKind : undefined, + } as Partial, + { at: [], match } + ); }, - [getCurrentValue, setValue] + [editor] ); const clear = useCallback(() => { setValue(EMPTY_VALUE); - }, [setValue]); + // ``tf.setValue`` (inside ``setValue``) wipes the editor's + // selection — without this, after the user presses Enter to + // submit, the composer is left with no caret and they would + // have to click before typing again. + requestAnimationFrame(focusAtEnd); + }, [focusAtEnd, setValue]); const setText = useCallback( (text: string) => { @@ -567,7 +588,22 @@ export const InlineMentionEditor = forwardRef ({ - focus: () => editableRef.current?.focus(), + // If we already have a Plate selection (user was typing + // before focus left), preserve it — just refocus. If we + // don't (first mount, or focus was lost without a + // surviving selection), seed a selection at end-of-doc + // so the contentEditable shows a caret instead of an + // invisible focus ring. + focus: () => { + try { + if (!editor.selection) { + editor.tf.select(editor.api.end([])); + } + editor.tf.focus(); + } catch { + editableRef.current?.focus(); + } + }, clear, setText, getText, @@ -579,6 +615,7 @@ export const InlineMentionEditor = forwardRef { const promptPickerRef = useRef(null); const { search_space_id, chat_id } = useParams(); const aui = useAui(); - const hasAutoFocusedRef = useRef(false); + // Gate the always-focused composer behaviour to desktop. On mobile, + // programmatic focus pops the soft keyboard, which would be jarring + // whenever a picker closes or the user navigates between threads. + const isDesktop = useMediaQuery("(min-width: 640px)"); const electronAPI = useElectronAPI(); const [clipboardInitialText, setClipboardInitialText] = useState(); @@ -437,16 +441,24 @@ const Composer: FC = () => { ); useBatchCommentsPreload(assistantDbMessageIds); - // Auto-focus editor on new chat page after mount + // Always-focused composer (Claude-style). Runs as a reactive + // invariant: whenever the composer is mounted and no transient + // picker has taken over keyboard input, the editor should be the + // focused element. This naturally restores focus after pickers + // close, after the user switches threads, and on first mount — + // replacing the previous one-shot ``hasAutoFocusedRef`` gate that + // only worked on the welcome screen. + // + // Gated on ``isDesktop`` so we don't repeatedly summon the mobile + // soft keyboard whenever any of the deps change. ``threadId`` is + // read so the effect re-fires when the user switches between two + // non-empty threads (where the Composer instance is reused). useEffect(() => { - if (isThreadEmpty && !hasAutoFocusedRef.current && editorRef.current) { - const timeoutId = setTimeout(() => { - editorRef.current?.focus(); - hasAutoFocusedRef.current = true; - }, 100); - return () => clearTimeout(timeoutId); - } - }, [isThreadEmpty]); + if (!isDesktop) return; + if (showDocumentPopover || showPromptPicker) return; + void threadId; + editorRef.current?.focus(); + }, [isDesktop, showDocumentPopover, showPromptPicker, threadId]); // Close document picker when a slide-out panel (inbox, shared/private chats) opens useEffect(() => { @@ -458,12 +470,45 @@ const Composer: FC = () => { return () => window.removeEventListener(SLIDEOUT_PANEL_OPENED_EVENT, handler); }, []); - // Sync editor text with assistant-ui composer runtime + // Sync editor text with the assistant-ui composer runtime and + // reconcile the chip atom from the editor's reported docs. + // + // The editor is the source of truth for which chips exist on + // screen. Reconciling here covers every deletion path Plate can + // produce (the explicit Backspace handler, the X-button, + // Cmd+Backspace, range-select+Delete, cut, paste-over) without + // needing per-keybinding plumbing. Without this, paths that bypass + // ``onDocumentRemove`` left the atom carrying stale entries that + // the picker would re-emit via ``initialSelectedDocuments`` and + // resurface as chips on the next selection. + // + // The setter returns ``prev`` when the chip set is unchanged so + // pure-text keystrokes don't churn the atom (Jotai compares by + // reference for store change notifications). const handleEditorChange = useCallback( - (text: string) => { + (text: string, docs: MentionedDocument[]) => { aui.composer().setText(text); + setMentionedDocuments((prev) => { + if (prev.length === docs.length) { + const editorKeys = new Set(docs.map((d) => getMentionDocKey(d))); + if (prev.every((d) => editorKeys.has(getMentionDocKey(d)))) { + return prev; + } + } + return docs.map((d) => ({ + id: d.id, + title: d.title, + // ``MentionedDocument.document_type`` is optional but + // the atom shape requires a string. ``"UNKNOWN"`` is + // the same sentinel ``getMentionDocKey`` and the + // editor's match predicates already use, so the key + // is stable across the round trip. + document_type: d.document_type ?? "UNKNOWN", + kind: d.kind, + })); + }); }, - [aui] + [aui, setMentionedDocuments] ); // Open document picker when @ mention is triggered @@ -621,27 +666,26 @@ const Composer: FC = () => { [setMentionedDocuments] ); - const handleDocumentsMention = useCallback( - (mentions: MentionedDocumentInfo[]) => { - const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? []; - const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc))); + const handleDocumentsMention = useCallback((mentions: MentionedDocumentInfo[]) => { + const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? []; + const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc))); - for (const mention of mentions) { - const key = getMentionDocKey(mention); - if (editorDocKeys.has(key)) continue; - editorRef.current?.insertMentionChip(mention); - } + for (const mention of mentions) { + const key = getMentionDocKey(mention); + if (editorDocKeys.has(key)) continue; + editorRef.current?.insertMentionChip(mention); + // Track within the loop so duplicates in the same batch + // (defensive — the picker shouldn't produce them today) + // can't slip through as double-inserted chips. + editorDocKeys.add(key); + } - setMentionedDocuments((prev) => { - const existingKeySet = new Set(prev.map((d) => getMentionDocKey(d))); - const uniqueNew = mentions.filter((m) => !existingKeySet.has(getMentionDocKey(m))); - return [...prev, ...uniqueNew]; - }); - - setMentionQuery(""); - }, - [setMentionedDocuments] - ); + // Atom is reconciled by the editor's ``onChange`` after each + // ``insertMentionChip`` (see ``handleEditorChange``); writing + // here would be a second, divergent write path — exactly the + // shape that let stale entries resurface in the past. + setMentionQuery(""); + }, []); useEffect(() => { const editor = editorRef.current;