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