Update new prompt input to support recent changes to current prompt i…

…nput

lib/prompt-editor/src/nodes/ContextItemMentionNode.tsx

@@ -1,14 +1,12 @@
 import {
     CONTEXT_ITEM_MENTION_NODE_TYPE,
     type ContextItem,
-    ContextItemSource,
     FILE_CONTEXT_MENTION_PROVIDER,
     REMOTE_REPOSITORY_PROVIDER_URI,
     SYMBOL_CONTEXT_MENTION_PROVIDER,
     type SerializedContextItem,
     type SerializedContextItemMentionNode,
     contextItemMentionNodeDisplayText,
-    displayPath,
     serializeContextItem,
 } from '@sourcegraph/cody-shared'
 import {
@@ -24,11 +22,10 @@ import {
     TextNode,
 } from 'lexical'
 import { AtSignIcon } from 'lucide-react'
-import * as v from 'valibot'
-import { URI } from 'vscode-uri'
 import { iconForProvider } from '../mentions/mentionMenu/MentionMenuItem'
 import styles from './ContextItemMentionNode.module.css'
 import { MentionComponent } from './MentionComponent'
+import { tooltipForContextItem } from './tooltip'
 
 export const MENTION_CLASS_NAME = styles.contextItemMentionNode
 
@@ -130,26 +127,7 @@ export class ContextItemMentionNode extends DecoratorNode<JSX.Element> {
     }
 
     private getTooltip(): string | undefined {
-        if (this.contextItem.type === 'repository') {
-            return `Repository: ${this.contextItem.repoName ?? this.contextItem.title ?? 'unknown'}`
-        }
-        if (this.contextItem.type === 'tree') {
-            return this.contextItem.title || 'Local workspace'
-        }
-        if (this.contextItem.type === 'file') {
-            return this.contextItem.isTooLarge
-                ? this.contextItem.source === ContextItemSource.Initial
-                    ? 'File is too large. Select a smaller range of lines from the file.'
-                    : 'File is too large. Try adding the file again with a smaller range of lines.'
-                : displayPath(URI.parse(this.contextItem.uri))
-        }
-        if (v.is(OpenCtxItemWithTooltipSchema, this.contextItem)) {
-            return this.contextItem.mention.data.tooltip
-        }
-        if (this.contextItem.type === 'openctx') {
-            return this.contextItem.uri
-        }
-        return undefined
+        return tooltipForContextItem(this.contextItem)
     }
 
     decorate(_editor: LexicalEditor, _config: EditorConfig): JSX.Element {
@@ -213,15 +191,3 @@ export function $createContextItemTextNode(contextItem: ContextItem): TextNode {
     const textNode = new TextNode(contextItemMentionNodeDisplayText(serializeContextItem(contextItem)))
     return $applyNodeReplacement(textNode)
 }
-
-/**
- * The structure of an openctx context item with a tooltip.
- */
-const OpenCtxItemWithTooltipSchema = v.object({
-    type: v.literal('openctx'),
-    mention: v.object({
-        data: v.object({
-            tooltip: v.string(),
-        }),
-    }),
-})

lib/prompt-editor/src/nodes/tooltip.ts

@@ -0,0 +1,38 @@
+import { ContextItemSource, type SerializedContextItem, displayPath } from '@sourcegraph/cody-shared'
+import * as v from 'valibot'
+import { URI } from 'vscode-uri'
+
+/**
+ * The structure of an openctx context item with a tooltip.
+ */
+const OpenCtxItemWithTooltipSchema = v.object({
+    type: v.literal('openctx'),
+    mention: v.object({
+        data: v.object({
+            tooltip: v.string(),
+        }),
+    }),
+})
+
+export function tooltipForContextItem(item: SerializedContextItem): string | undefined {
+    if (item.type === 'repository') {
+        return `Repository: ${item.repoName ?? item.title ?? 'unknown'}`
+    }
+    if (item.type === 'tree') {
+        return item.title || 'Local workspace'
+    }
+    if (item.type === 'file') {
+        return item.isTooLarge
+            ? item.source === ContextItemSource.Initial
+                ? 'File is too large. Select a smaller range of lines from the file.'
+                : 'File is too large. Try adding the file again with a smaller range of lines.'
+            : displayPath(URI.parse(item.uri))
+    }
+    if (v.is(OpenCtxItemWithTooltipSchema, item)) {
+        return item.mention.data.tooltip
+    }
+    if (item.type === 'openctx') {
+        return item.uri
+    }
+    return undefined
+}

lib/prompt-editor/src/v2/MentionView.tsx

@@ -1,20 +1,18 @@
 import { shift, size, useDismiss, useFloating, useHover, useInteractions } from '@floating-ui/react'
 import {
-    ContextItemSource,
     FILE_CONTEXT_MENTION_PROVIDER,
     REMOTE_REPOSITORY_PROVIDER_URI,
     SYMBOL_CONTEXT_MENTION_PROVIDER,
     type SerializedContextItem,
-    displayPath,
 } from '@sourcegraph/cody-shared'
 import clsx from 'clsx'
 import { AtSignIcon } from 'lucide-react'
 import type { Node } from 'prosemirror-model'
 import type { NodeView } from 'prosemirror-view'
 import { useState } from 'react'
 import { type Root, createRoot } from 'react-dom/client'
-import { URI } from 'vscode-uri'
 import { iconForProvider } from '../mentions/mentionMenu/MentionMenuItem'
+import { tooltipForContextItem } from '../nodes/tooltip'
 import styles from './MentionView.module.css'
 
 export class MentionView implements NodeView {
@@ -106,26 +104,6 @@ const MentionChip: React.FC<MentionChipProps> = props => {
     )
 }
 
-function tooltipForContextItem(item: SerializedContextItem): string | undefined {
-    if (item.type === 'repository') {
-        return `Repository: ${item.repoName ?? item.title ?? 'unknown'}`
-    }
-    if (item.type === 'tree') {
-        return item.title || 'Local workspace'
-    }
-    if (item.type === 'file') {
-        return item.isTooLarge
-            ? item.source === ContextItemSource.Initial
-                ? 'File is too large. Select a smaller range of lines from the file.'
-                : 'File is too large. Try adding the file again with a smaller range of lines.'
-            : displayPath(URI.parse(item.uri))
-    }
-    if (item.type === 'openctx') {
-        return item.uri
-    }
-    return undefined
-}
-
 function iconForContextItem(item: SerializedContextItem): React.ComponentType {
     let providerURI = 'unknown'
     switch (item.type) {

lib/prompt-editor/src/v2/PromptEditor.tsx

@@ -64,6 +64,21 @@ export interface PromptEditorRefAPI {
     setFocus(focus: boolean, options?: { moveCursorToEnd?: boolean }): Promise<void>
     appendText(text: string): Promise<void>
     addMentions(items: ContextItem[], position?: 'before' | 'after', sep?: string): Promise<void>
+    /**
+     * 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,
+        focusEditor?: boolean
+    ): Promise<void>
     filterMentions(filter: (item: SerializedContextItem) => boolean): Promise<void>
     setInitialContextMentions(items: ContextItem[]): Promise<void>
     setEditorState(state: SerializedPromptEditorState): void
@@ -189,6 +204,14 @@ export const PromptEditor: FunctionComponent<Props> = ({
             ): Promise<void> {
                 api.addMentions(items, position, sep)
             },
+            async upsertMentions(
+                items: ContextItem[],
+                position: 'before' | 'after' = 'after',
+                sep = ' ',
+                focusEditor = true
+            ): Promise<void> {
+                api.upsertMentions(items, position, sep, focusEditor)
+            },
             async setInitialContextMentions(items: ContextItem[]): Promise<void> {
                 api.setInitialContextMentions(items)
             },

lib/prompt-editor/src/v2/promptInput-react.ts

@@ -66,6 +66,12 @@ interface PromptEditorAPI {
     setFocus(focus: boolean, options?: { moveCursorToEnd?: boolean }): void
     appendText(text: string): void
     addMentions(items: ContextItem[], position?: 'before' | 'after', sep?: string): void
+    upsertMentions(
+        items: ContextItem[],
+        position?: 'before' | 'after',
+        sep?: string,
+        focusEditor?: boolean
+    ): void
     filterMentions(filter: (item: SerializedContextItem) => boolean): void
     setInitialContextMentions(items: ContextItem[]): void
     setDocument(doc: Node): void
@@ -184,6 +190,22 @@ export const usePromptInput = (options: PromptEditorOptions): [PromptInputActor,
                     separator: sep,
                 })
             },
+            upsertMentions(
+                items: ContextItem[],
+                position: 'before' | 'after' = 'after',
+                sep = ' ',
+                focusEditor = true
+            ) {
+                editor.send({
+                    type: 'document.mentions.upsert',
+                    items: items.map(serializeContextItem),
+                    position,
+                    separator: sep,
+                })
+                if (focusEditor) {
+                    editor.send({ type: 'focus' })
+                }
+            },
             filterMentions(filter: (item: SerializedContextItem) => boolean) {
                 editor.send({ type: 'document.mentions.filter', filter })
             },

lib/prompt-editor/src/v2/promptInput.test.ts

@@ -16,6 +16,7 @@ import {
     type DataLoaderInput,
     type MenuItem,
     type PromptInputOptions,
+    getMentions,
     promptInput,
     schema,
 } from './promptInput'
@@ -184,6 +185,69 @@ describe('mentions', () => {
         })
     })
 
+    describe('document.mentions.upsert', () => {
+        test('append', () => {
+            const input = createInput(['before ', cm('file1'), ' after'])
+            input.send({
+                type: 'document.mentions.upsert',
+                items: [cm('file2'), cm('file3')],
+                position: 'after',
+                separator: ' ! ',
+            })
+
+            expect(getText(input)).toMatchInlineSnapshot(`"before file1 after file2 ! file3 ! "`)
+        })
+
+        test('prepend', () => {
+            const input = createInput(['before ', cm('file1'), ' after'])
+            input.send({
+                type: 'document.mentions.upsert',
+                items: [cm('file2'), cm('file3')],
+                position: 'before',
+                separator: ' ! ',
+            })
+
+            expect(getText(input)).toMatchInlineSnapshot(`"file2 ! file3 ! before file1 after"`)
+        })
+
+        test('update mention', () => {
+            const input = createInput([
+                'before ',
+                {
+                    type: 'openctx',
+                    uri: 'file:///file1.txt',
+                    title: 'test',
+                    provider: 'openctx',
+                    providerUri: REMOTE_FILE_PROVIDER_URI,
+                },
+                ' after',
+            ])
+
+            const newMentionData = { uri: 'uri1', data: 1 }
+            input.send({
+                type: 'document.mentions.upsert',
+                items: [
+                    {
+                        type: 'openctx',
+                        uri: 'file:///file1.txt',
+                        title: '|test updated|',
+                        provider: 'openctx',
+                        providerUri: REMOTE_FILE_PROVIDER_URI,
+                        mention: newMentionData,
+                    },
+                ],
+                position: 'after',
+                separator: ' ! ',
+            })
+
+            expect(getText(input)).toMatchInlineSnapshot(`"before |test updated| after"`)
+            const state = getEditorState(input)
+            expect(getMentions(state.doc)).toEqual([
+                expect.objectContaining({ mention: newMentionData }),
+            ])
+        })
+    })
+
     test('document.mentions.filter', () => {
         const editor = createInput(['1 ', cm('file1'), ' 2 ', cm('file2'), ' 3 ', cm('file3')])
         editor.send({ type: 'document.mentions.filter', filter: item => item.uri === 'file2' })