refactor(assistant-ui): enhance mention chip handling and editor focus behavior

This commit is contained in:
Anish Sarkar 2026-05-12 23:18:45 +05:30
parent 0c2beb7ce8
commit 2437716752
2 changed files with 215 additions and 134 deletions

View file

@ -1,6 +1,7 @@
"use client"; "use client";
import { Folder as FolderIcon, X as XIcon } from "lucide-react"; import { Folder as FolderIcon, X as XIcon } from "lucide-react";
import type { NodeEntry, TElement } from "platejs";
import type { PlateElementProps } from "platejs/react"; import type { PlateElementProps } from "platejs/react";
import { import {
createPlatePlugin, createPlatePlugin,
@ -56,10 +57,7 @@ export interface InlineMentionEditorRef {
setText: (text: string) => void; setText: (text: string) => void;
getText: () => string; getText: () => string;
getMentionedDocuments: () => MentionedDocument[]; getMentionedDocuments: () => MentionedDocument[];
insertMentionChip: ( insertMentionChip: (mention: MentionChipInput, options?: { removeTriggerText?: boolean }) => void;
mention: MentionChipInput,
options?: { removeTriggerText?: boolean }
) => void;
/** /**
* @deprecated Use ``insertMentionChip``. Kept for one transition * @deprecated Use ``insertMentionChip``. Kept for one transition
* cycle so we don't break ad-hoc callers; prefer the new name. * cycle so we don't break ad-hoc callers; prefer the new name.
@ -338,17 +336,21 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
value: initialText ? toValueFromText(initialText) : EMPTY_VALUE, value: initialText ? toValueFromText(initialText) : EMPTY_VALUE,
}); });
// Move the caret to the end of the document and focus the editor.
// Routes through Plate's transforms so ``editor.selection`` and
// the DOM selection stay in sync — bypassing Plate (via raw
// ``window.getSelection``) was the prior implementation and is
// what made the caret disappear after every ``setValue``-based
// mutation. Falls back to DOM focus if Plate's API throws (e.g.
// during a transient unmount race).
const focusAtEnd = useCallback(() => { const focusAtEnd = useCallback(() => {
const el = editableRef.current; try {
if (!el) return; editor.tf.select(editor.api.end([]));
el.focus(); editor.tf.focus();
const selection = window.getSelection(); } catch {
const range = document.createRange(); editableRef.current?.focus();
range.selectNodeContents(el); }
range.collapse(false); }, [editor]);
selection?.removeAllRanges();
selection?.addRange(range);
}, []);
const getCurrentValue = useCallback( const getCurrentValue = useCallback(
() => (editor.children as ComposerValue) ?? EMPTY_VALUE, () => (editor.children as ComposerValue) ?? EMPTY_VALUE,
@ -396,17 +398,30 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
[editor, emitState] [editor, emitState]
); );
// Insert a mention chip at the current caret. Uses Plate
// transforms so Slate keeps the editor selection valid through
// the edit.
//
// Critical detail: the chip is a void inline element. Inserting
// it on its own with ``{ select: true }`` would land the caret
// inside the void's empty ``children: [{ text: "" }]`` — a point
// the browser can't render a caret on, which is what made the
// cursor disappear or jump to the wrong side of the chip after
// insertion. Inserting ``[mentionNode, { text: " " }]`` as a
// single array means the *last* inserted node is a text node, so
// ``{ select: true }`` resolves to that text node's end (offset
// 1, after the trailing space) — a real, renderable text point.
// The whole sequence stays inside ``withoutNormalizing`` so the
// optional trigger-text delete and the chip insert show up as a
// single undo step.
const insertMentionChip = useCallback( const insertMentionChip = useCallback(
(mention: MentionChipInput, options?: { removeTriggerText?: boolean }) => { (mention: MentionChipInput, options?: { removeTriggerText?: boolean }) => {
if (typeof mention.id !== "number" || typeof mention.title !== "string") return; if (typeof mention.id !== "number" || typeof mention.title !== "string") return;
const removeTriggerText = options?.removeTriggerText ?? true; const removeTriggerText = options?.removeTriggerText ?? true;
const current = getCurrentValue();
const selection = editor.selection;
const kind: MentionKind = mention.kind ?? "doc"; const kind: MentionKind = mention.kind ?? "doc";
const document_type = const document_type =
mention.document_type ?? mention.document_type ?? (kind === "folder" ? FOLDER_MENTION_DOCUMENT_TYPE : undefined);
(kind === "folder" ? FOLDER_MENTION_DOCUMENT_TYPE : undefined);
const mentionNode: MentionElementNode = { const mentionNode: MentionElementNode = {
type: MENTION_TYPE, type: MENTION_TYPE,
id: mention.id, id: mention.id,
@ -416,60 +431,49 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
children: [{ text: "" }], children: [{ text: "" }],
}; };
const cursorCtx = getCursorTextContext(current, selection); editor.tf.withoutNormalizing(() => {
if (!cursorCtx) { const selection = editor.selection;
const lastBlock = current[current.length - 1] ?? { type: "p", children: [{ text: "" }] };
const appended: ComposerValue = [
...current.slice(0, -1),
{
...lastBlock,
children: [...lastBlock.children, mentionNode, { text: " " }],
},
];
setValue(appended);
requestAnimationFrame(focusAtEnd);
return;
}
const block = current[cursorCtx.blockIndex]; // No active editor selection — typically because focus
const currentChild = getTextNode(block.children[cursorCtx.childIndex]); // moved to a picker/dropdown. Snap the caret to the end
if (!currentChild) { // of the document so the chip appends cleanly instead
const children = [...block.children]; // of disappearing into a dead range.
children.splice(cursorCtx.childIndex + 1, 0, mentionNode, { text: " " }); if (!selection) {
const next = [...current]; editor.tf.select(editor.api.end([]));
next[cursorCtx.blockIndex] = { ...block, children }; } else if (removeTriggerText) {
setValue(next as ComposerValue); // Delete the in-progress "@query" text so the chip
requestAnimationFrame(focusAtEnd); // stands in for it. Mirrors the old splice but lets
return; // Slate keep selection sane through the edit.
} const cursorCtx = getCursorTextContext(getCurrentValue(), selection);
if (cursorCtx) {
const text = currentChild.text; const text = cursorCtx.text;
let removeStart = cursorCtx.cursor; let triggerIndex = -1;
if (removeTriggerText) { for (let i = cursorCtx.cursor - 1; i >= 0; i--) {
for (let i = cursorCtx.cursor - 1; i >= 0; i--) { if (text[i] === "@") {
if (text[i] === "@") { triggerIndex = i;
removeStart = i; break;
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); editor.tf.insertNodes([mentionNode, { text: " " }] as unknown as TElement[], {
const after = text.slice(cursorCtx.cursor); select: true,
const replacement: ComposerNode[] = []; });
if (before.length > 0) replacement.push({ text: before }); });
replacement.push(mentionNode); editor.tf.focus();
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.selection, focusAtEnd, getCurrentValue, setValue] [editor, getCurrentValue]
); );
// Backwards-compatible shim — pre-folder callers pass a doc-only // Backwards-compatible shim — pre-folder callers pass a doc-only
@ -485,24 +489,34 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
[insertMentionChip] [insertMentionChip]
); );
// Remove the chip(s) matching the given (id, document_type) pair.
// Goes through ``tf.removeNodes`` so Slate keeps the surrounding
// selection valid — the previous ``setValue``-based filter wiped
// selection on every removal, which is why the caret vanished
// when the X button was clicked. Iterates descending so removing
// one entry doesn't invalidate the path of subsequent matches.
// In practice chips are deduped by ``getMentionDocKey`` so this
// loop runs at most once; the descending iteration is defense
// against any future divergence.
const removeDocumentChip = useCallback( const removeDocumentChip = useCallback(
(docId: number, docType?: string) => { (docId: number, docType?: string) => {
const current = getCurrentValue(); const match = (n: unknown) => {
let changed = false; if (!n || typeof n !== "object" || !("type" in n)) return false;
const next = current.map((block) => { const node = n as MentionElementNode;
const children = block.children.filter((node) => { if (node.type !== MENTION_TYPE) return false;
if (!isMentionNode(node)) return true; if (node.id !== docId) return false;
const match = return (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN");
node.id === docId && (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN"); };
if (match) changed = true;
return !match; const entries = Array.from(editor.api.nodes({ at: [], match })) as NodeEntry[];
}); if (entries.length === 0) return;
return { ...block, children: children.length ? children : [{ text: "" }] }; 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 // Combined "remove chip end-to-end" used by both the Backspace
@ -517,6 +531,11 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
[onDocumentRemove, removeDocumentChip] [onDocumentRemove, removeDocumentChip]
); );
// Update the streaming status on a chip in place. ``tf.setNodes``
// merges the partial props onto every node matching the
// predicate without rebuilding the document, so the user's
// selection stays put — important because status transitions
// arrive as backend events while the user may be mid-typing.
const setDocumentChipStatus = useCallback( const setDocumentChipStatus = useCallback(
( (
docId: number, docId: number,
@ -524,31 +543,33 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
statusLabel: string | null, statusLabel: string | null,
statusKind: MentionStatusKind = "pending" statusKind: MentionStatusKind = "pending"
) => { ) => {
const current = getCurrentValue(); const match = (n: unknown) => {
let changed = false; if (!n || typeof n !== "object" || !("type" in n)) return false;
const next = current.map((block) => ({ const node = n as MentionElementNode;
...block, if (node.type !== MENTION_TYPE) return false;
children: block.children.map((node) => { if (node.id !== docId) return false;
if (!isMentionNode(node)) return node; return (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN");
const sameType = (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN"); };
if (node.id !== docId || !sameType) return node;
changed = true; editor.tf.setNodes(
return { {
...node, statusLabel,
statusLabel, statusKind: statusLabel ? statusKind : undefined,
statusKind: statusLabel ? statusKind : undefined, } as Partial<TElement>,
}; { at: [], match }
}), );
}));
if (!changed) return;
setValue(next as ComposerValue);
}, },
[getCurrentValue, setValue] [editor]
); );
const clear = useCallback(() => { const clear = useCallback(() => {
setValue(EMPTY_VALUE); 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( const setText = useCallback(
(text: string) => { (text: string) => {
@ -567,7 +588,22 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
useImperativeHandle( useImperativeHandle(
ref, ref,
() => ({ () => ({
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, clear,
setText, setText,
getText, getText,
@ -579,6 +615,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
}), }),
[ [
clear, clear,
editor,
getMentionedDocs, getMentionedDocs,
getText, getText,
insertMentionChip, insertMentionChip,

View file

@ -62,6 +62,7 @@ import { useDocumentUploadDialog } from "@/components/assistant-ui/document-uplo
import { import {
InlineMentionEditor, InlineMentionEditor,
type InlineMentionEditorRef, type InlineMentionEditorRef,
type MentionedDocument,
} from "@/components/assistant-ui/inline-mention-editor"; } from "@/components/assistant-ui/inline-mention-editor";
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import { UserMessage } from "@/components/assistant-ui/user-message"; import { UserMessage } from "@/components/assistant-ui/user-message";
@ -384,7 +385,10 @@ const Composer: FC = () => {
const promptPickerRef = useRef<PromptPickerRef>(null); const promptPickerRef = useRef<PromptPickerRef>(null);
const { search_space_id, chat_id } = useParams(); const { search_space_id, chat_id } = useParams();
const aui = useAui(); 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 electronAPI = useElectronAPI();
const [clipboardInitialText, setClipboardInitialText] = useState<string | undefined>(); const [clipboardInitialText, setClipboardInitialText] = useState<string | undefined>();
@ -437,16 +441,24 @@ const Composer: FC = () => {
); );
useBatchCommentsPreload(assistantDbMessageIds); 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(() => { useEffect(() => {
if (isThreadEmpty && !hasAutoFocusedRef.current && editorRef.current) { if (!isDesktop) return;
const timeoutId = setTimeout(() => { if (showDocumentPopover || showPromptPicker) return;
editorRef.current?.focus(); void threadId;
hasAutoFocusedRef.current = true; editorRef.current?.focus();
}, 100); }, [isDesktop, showDocumentPopover, showPromptPicker, threadId]);
return () => clearTimeout(timeoutId);
}
}, [isThreadEmpty]);
// Close document picker when a slide-out panel (inbox, shared/private chats) opens // Close document picker when a slide-out panel (inbox, shared/private chats) opens
useEffect(() => { useEffect(() => {
@ -458,12 +470,45 @@ const Composer: FC = () => {
return () => window.removeEventListener(SLIDEOUT_PANEL_OPENED_EVENT, handler); 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( const handleEditorChange = useCallback(
(text: string) => { (text: string, docs: MentionedDocument[]) => {
aui.composer().setText(text); 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<MentionedDocumentInfo>((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 // Open document picker when @ mention is triggered
@ -621,27 +666,26 @@ const Composer: FC = () => {
[setMentionedDocuments] [setMentionedDocuments]
); );
const handleDocumentsMention = useCallback( const handleDocumentsMention = useCallback((mentions: MentionedDocumentInfo[]) => {
(mentions: MentionedDocumentInfo[]) => { const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? [];
const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? []; const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc)));
const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc)));
for (const mention of mentions) { for (const mention of mentions) {
const key = getMentionDocKey(mention); const key = getMentionDocKey(mention);
if (editorDocKeys.has(key)) continue; if (editorDocKeys.has(key)) continue;
editorRef.current?.insertMentionChip(mention); 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) => { // Atom is reconciled by the editor's ``onChange`` after each
const existingKeySet = new Set(prev.map((d) => getMentionDocKey(d))); // ``insertMentionChip`` (see ``handleEditorChange``); writing
const uniqueNew = mentions.filter((m) => !existingKeySet.has(getMentionDocKey(m))); // here would be a second, divergent write path — exactly the
return [...prev, ...uniqueNew]; // shape that let stale entries resurface in the past.
}); setMentionQuery("");
}, []);
setMentionQuery("");
},
[setMentionedDocuments]
);
useEffect(() => { useEffect(() => {
const editor = editorRef.current; const editor = editorRef.current;