From 9db807b149a4c95da6df1437128f1eb7c6212e0c Mon Sep 17 00:00:00 2001 From: Felix Kling Date: Tue, 17 Dec 2024 12:06:20 +0100 Subject: [PATCH] fix(onebox): Do not focus editor when inserting/updating search results context Changing the focus behavior was more difficult than expected. Apparently when Lexical has a Selection then it will 'steal' the focus whenever changes to its content is made. There have been made changes to Lexical just recently that would allow us to avoid this (https://github.com/facebook/lexical/pull/6894) but we haven't updated to this version yet and updating this closely before the EAP seems risky. I found a workaround in one of the discussions that suggests to set the selection to `null` when making a change. In order to do this I refactored some of the function to make them more composable. Instead of calling `editor.update` directly these functions are now supposed to be called from `editor.update`. Since I've added a helper for creating a promise-version of `editor.update` anyway, I also addressed one of the issues I've encountered in the last PR, namely that `onUpdate` is not called when the editor state doesn't change. This will btw also be fixed on the Lexical side by the above mentioned PR. --- lib/prompt-editor/src/PromptEditor.tsx | 257 +++++++++--------- lib/prompt-editor/src/lexicalUtils.ts | 84 +++++- .../messageCell/assistant/SearchResults.tsx | 2 +- 3 files changed, 206 insertions(+), 137 deletions(-) diff --git a/lib/prompt-editor/src/PromptEditor.tsx b/lib/prompt-editor/src/PromptEditor.tsx index ecd280099f91..a576de2c2ab2 100644 --- a/lib/prompt-editor/src/PromptEditor.tsx +++ b/lib/prompt-editor/src/PromptEditor.tsx @@ -1,4 +1,3 @@ -import { $insertFirst } from '@lexical/utils' import { type ContextItem, type SerializedContextItem, @@ -10,7 +9,6 @@ import { } from '@sourcegraph/cody-shared' import { clsx } from 'clsx' import { - $createParagraphNode, $createTextNode, $getRoot, $getSelection, @@ -18,14 +16,21 @@ import { $setSelection, type LexicalEditor, } from 'lexical' -import type { EditorState, RootNode, SerializedEditorState, SerializedLexicalNode } from 'lexical' +import type { EditorState, SerializedEditorState, SerializedLexicalNode } from 'lexical' import isEqual from 'lodash/isEqual' import { type FunctionComponent, useCallback, useEffect, useImperativeHandle, useRef } from 'react' import { BaseEditor } from './BaseEditor' import styles from './PromptEditor.module.css' import { useSetGlobalPromptEditorConfig } from './config' import { isEditorContentOnlyInitialContext, lexicalNodesForContextItems } from './initialContext' -import { $selectEnd, getContextItemsForEditor, visitContextItemsForEditor } from './lexicalUtils' +import { + $insertMentions, + $selectEnd, + getContextItemsForEditor, + getWhitespace, + update, + walkContextItemMentionNodes, +} from './lexicalUtils' import { $createContextItemMentionNode } from './nodes/ContextItemMentionNode' import type { KeyboardEventPluginProps } from './plugins/keyboardEvent' @@ -55,8 +60,18 @@ export interface PromptEditorRefAPI { /** * Similar to `addMentions`, but unlike `addMentions` it doesn't merge mentions with overlapping * ranges. Instead it updates the meta data of existing mentions with the same uri. + * + * @param items The context items to add or update. + * @param position Where to insert the mentions, before or after the current input. Defaults to 'after'. + * @param sep The separator to use between mentions. Defaults to a space. + * @param focusEditor Whether to focus the editor after updating the mentions. Defaults to true. */ - upsertMentions(items: ContextItem[], position?: 'before' | 'after', sep?: string): Promise + upsertMentions( + items: ContextItem[], + position?: 'before' | 'after', + sep?: string, + focusEditor?: boolean + ): Promise filterMentions(filter: (item: SerializedContextItem) => boolean): Promise setInitialContextMentions(items: ContextItem[]): Promise setEditorState(state: SerializedPromptEditorState): void @@ -141,30 +156,31 @@ export const PromptEditor: FunctionComponent = ({ }) }, appendText(text: string): Promise { - return new Promise(resolve => - editorRef.current?.update( - () => { - const root = $getRoot() - root.selectEnd() - $insertNodes([$createTextNode(`${getWhitespace(root)}${text}`)]) - root.selectEnd() - }, - { onUpdate: resolve } - ) - ) + if (!editorRef.current) { + return Promise.resolve() + } + return update(editorRef.current, () => { + const root = $getRoot() + root.selectEnd() + $insertNodes([$createTextNode(`${getWhitespace(root)}${text}`)]) + root.selectEnd() + return true + }) }, filterMentions(filter: (item: SerializedContextItem) => boolean): Promise { - return new Promise(resolve => { - if (!editorRef.current) { - resolve() - return - } + if (!editorRef.current) { + return Promise.resolve() + } - visitContextItemsForEditor(editorRef.current, node => { + return update(editorRef.current, () => { + let updated = false + walkContextItemMentionNodes($getRoot(), node => { if (!filter(node.contextItem)) { node.remove() + updated = true } - }).then(resolve) + }) + return updated }) }, async addMentions( @@ -181,8 +197,12 @@ export const PromptEditor: FunctionComponent = ({ const existingMentions = getContextItemsForEditor(editor) const ops = getMentionOperations(existingMentions, newContextItems) - if (ops.modify.size + ops.delete.size > 0) { - await visitContextItemsForEditor(editor, existing => { + await update(editor, () => { + if (ops.modify.size + ops.delete.size === 0) { + return false + } + + walkContextItemMentionNodes($getRoot(), existing => { const update = ops.modify.get(existing.contextItem) if (update) { // replace the existing mention inline with the new one @@ -192,15 +212,25 @@ export const PromptEditor: FunctionComponent = ({ existing.remove() } }) - } + return true + }) - if (ops.create.length === 0) { - return - } + return update(editor, () => { + if (ops.create.length === 0) { + return false + } - return insertMentions(editor, ops.create, position, sep) + $insertMentions(ops.create, position, sep) + $selectEnd() + return true + }) }, - async upsertMentions(items, position = 'after', sep = ' '): Promise { + async upsertMentions( + items, + position = 'after', + sep = ' ', + focusEditor = true + ): Promise { const editor = editorRef.current if (!editor) { return @@ -217,63 +247,85 @@ export const PromptEditor: FunctionComponent = ({ } } - // ! visitContextItemsForEditor must only be called when we have changes to make. Otherwise - // the promise it returns will never resolve. - if (toUpdate.size > 0) { - await visitContextItemsForEditor(editor, existing => { + await update(editor, () => { + if (toUpdate.size === 0) { + return false + } + + walkContextItemMentionNodes($getRoot(), existing => { const update = toUpdate.get(getKeyForContextItem(existing.contextItem)) if (update) { // replace the existing mention inline with the new one existing.replace($createContextItemMentionNode(update)) } }) - } - if (items.length === toUpdate.size) { - return - } - - return insertMentions( - editor, - items.filter(item => !toUpdate.has(getKeyForContextItem(item))), - position, - sep - ) + if (focusEditor) { + $selectEnd() + } else { + // Workaround for preventing the editor from stealing focus + // (https://github.com/facebook/lexical/issues/2636#issuecomment-1184418601) + // We need this until we can use the new 'skip-dom-selection' tag as + // explained in https://lexical.dev/docs/concepts/selection#focus, introduced + // by https://github.com/facebook/lexical/pull/6894 + $setSelection(null) + } + return true + }) + return update(editor, () => { + if (items.length === toUpdate.size) { + return false + } + $insertMentions( + items.filter(item => !toUpdate.has(getKeyForContextItem(item))), + position, + sep + ) + if (focusEditor) { + $selectEnd() + } else { + // Workaround for preventing the editor from stealing focus + // (https://github.com/facebook/lexical/issues/2636#issuecomment-1184418601) + // We need this until we can use the new 'skip-dom-selection' tag as + // explained in https://lexical.dev/docs/concepts/selection#focus, introduced + // by https://github.com/facebook/lexical/pull/6894 + $setSelection(null) + } + return true + }) }, setInitialContextMentions(items: ContextItem[]): Promise { - return new Promise(resolve => { - const editor = editorRef.current - if (!editor) { - return resolve() + const editor = editorRef.current + if (!editor) { + return Promise.resolve() + } + + return update(editor, () => { + let updated = false + + if (!hasSetInitialContext.current || isEditorContentOnlyInitialContext(editor)) { + if (isEditorContentOnlyInitialContext(editor)) { + // Only clear in this case so that we don't clobber any text that was + // inserted before initial context was received. + $getRoot().clear() + updated = true + } + const nodesToInsert = lexicalNodesForContextItems(items, { + isFromInitialContext: true, + }) + + // Add whitespace after initial context items chips + if (items.length > 0) { + nodesToInsert.push($createTextNode(' ')) + updated = true + } + + $setSelection($getRoot().selectStart()) // insert at start + $insertNodes(nodesToInsert) + $selectEnd() + hasSetInitialContext.current = true } - editor.update( - () => { - if ( - !hasSetInitialContext.current || - isEditorContentOnlyInitialContext(editor) - ) { - if (isEditorContentOnlyInitialContext(editor)) { - // Only clear in this case so that we don't clobber any text that was - // inserted before initial context was received. - $getRoot().clear() - } - const nodesToInsert = lexicalNodesForContextItems(items, { - isFromInitialContext: true, - }) - - // Add whitespace after initial context items chips - if (items.length > 0) { - nodesToInsert.push($createTextNode(' ')) - } - - $setSelection($getRoot().selectStart()) // insert at start - $insertNodes(nodesToInsert) - $selectEnd() - hasSetInitialContext.current = true - } - }, - { onUpdate: resolve } - ) + return updated }) }, }), @@ -283,7 +335,7 @@ export const PromptEditor: FunctionComponent = ({ useSetGlobalPromptEditorConfig() const onBaseEditorChange = useCallback( - (_editorState: EditorState, editor: LexicalEditor, tags: Set): void => { + (_editorState: EditorState, editor: LexicalEditor): void => { if (onChange) { onChange(toSerializedPromptEditorValue(editor)) } @@ -333,53 +385,6 @@ function normalizeEditorStateJSON( return JSON.parse(JSON.stringify(value)) } -function getWhitespace(root: RootNode): string { - const needsWhitespaceBefore = !/(^|\s)$/.test(root.getTextContent()) - return needsWhitespaceBefore ? ' ' : '' -} - -function insertMentions( - editor: LexicalEditor, - items: (SerializedContextItem | ContextItem)[], - position: 'before' | 'after', - sep?: string -): Promise { - return new Promise(resolve => - editor.update( - () => { - const nodesToInsert = lexicalNodesForContextItems( - items, - { - isFromInitialContext: false, - }, - sep - ) - const pNode = $createParagraphNode() - - switch (position) { - case 'before': { - pNode.append(...nodesToInsert) - $insertFirst($getRoot(), pNode) - break - } - case 'after': { - pNode.append( - $createTextNode(getWhitespace($getRoot())), - ...nodesToInsert, - $createTextNode(sep) - ) - $insertNodes([pNode]) - break - } - } - - $selectEnd() - }, - { onUpdate: resolve } - ) - ) -} - /** * Computes a unique key for a context item that can be used in e.g. a Map. * diff --git a/lib/prompt-editor/src/lexicalUtils.ts b/lib/prompt-editor/src/lexicalUtils.ts index 367f2e8bcb62..47ea996b51ba 100644 --- a/lib/prompt-editor/src/lexicalUtils.ts +++ b/lib/prompt-editor/src/lexicalUtils.ts @@ -1,7 +1,11 @@ -import type { SerializedContextItem } from '@sourcegraph/cody-shared' +import { $insertFirst } from '@lexical/utils' +import type { ContextItem, SerializedContextItem } from '@sourcegraph/cody-shared' import { + $createParagraphNode, $createRangeSelection, + $createTextNode, $getRoot, + $insertNodes, $isDecoratorNode, $isElementNode, $isTextNode, @@ -13,6 +17,7 @@ import { type RootNode, type TextNode, } from 'lexical' +import { lexicalNodesForContextItems } from './initialContext' import { ContextItemMentionNode } from './nodes/ContextItemMentionNode' function getLastNode(root: RootNode): ElementNode | TextNode | null { @@ -73,19 +78,78 @@ export function getContextItemsForEditor(editor: LexicalEditor): SerializedConte }) } -export function visitContextItemsForEditor( - editor: LexicalEditor, +export function walkContextItemMentionNodes( + node: LexicalNode, visit: (mention: ContextItemMentionNode) => void -): Promise { +): void { + walkLexicalNodes(node, node => { + if (node instanceof ContextItemMentionNode) { + visit(node) + } + return true + }) +} + +/** + * Inserts the given context items before or after the current value of the editor. + * + * @param items The context items to insert. + * @param position Where to insert the context items, relative to the editor value. Defaults to 'after'. + * @param sep The separator to insert between the current value and the context items. + */ +export function $insertMentions( + items: (SerializedContextItem | ContextItem)[], + position: 'before' | 'after', + sep?: string +): void { + const nodesToInsert = lexicalNodesForContextItems( + items, + { + isFromInitialContext: false, + }, + sep + ) + const pNode = $createParagraphNode() + + switch (position) { + case 'before': { + pNode.append(...nodesToInsert) + $insertFirst($getRoot(), pNode) + break + } + case 'after': { + pNode.append( + $createTextNode(getWhitespace($getRoot())), + ...nodesToInsert, + $createTextNode(sep) + ) + $insertNodes([pNode]) + break + } + } +} + +export function getWhitespace(root: RootNode): string { + const needsWhitespaceBefore = !/(^|\s)$/.test(root.getTextContent()) + return needsWhitespaceBefore ? ' ' : '' +} + +/** + * Helper function to update the editor state and get a promise that resolves when the update is done. + * + * IMPORTANT: The promise will never resolve when the update function does not cause any changes to the editor state. + * (not until we update to a version that includes https://github.com/facebook/lexical/pull/6894). + * To mitigate this, the update function should return a boolean, where `false` indicates that it did not cause any changes, + * in which case the promise will resolve immediately. + */ +export function update(editor: LexicalEditor, updateFn: () => boolean): Promise { return new Promise(resolve => { editor.update( () => { - walkLexicalNodes($getRoot(), node => { - if (node instanceof ContextItemMentionNode) { - visit(node) - } - return true - }) + const result = updateFn() + if (result === false) { + resolve() + } }, { onUpdate: resolve } ) diff --git a/vscode/webviews/chat/cells/messageCell/assistant/SearchResults.tsx b/vscode/webviews/chat/cells/messageCell/assistant/SearchResults.tsx index dc293f1c457a..b808f15fb17c 100644 --- a/vscode/webviews/chat/cells/messageCell/assistant/SearchResults.tsx +++ b/vscode/webviews/chat/cells/messageCell/assistant/SearchResults.tsx @@ -97,7 +97,7 @@ export const SearchResults = ({ }) .filter(isDefined) ) - lastEditorRef.current?.upsertMentions([contextItem]) + lastEditorRef.current?.upsertMentions([contextItem], 'before', ' ', false) } else { lastEditorRef.current?.filterMentions( mention => mention.type !== 'openctx' || mention.providerUri !== CODE_SEARCH_PROVIDER_URI