From 62d53fe4d590e2eda27d02e9f6e602e601ff108f Mon Sep 17 00:00:00 2001
From: liangfung <1098486429@qq.com>
Date: Tue, 7 Nov 2023 23:00:12 +0800
Subject: [PATCH] fix: select completion and replace with doc name

---
 ee/tabby-ui/components/prompt-form.tsx | 44 ++++++++++++++++++--------
 ee/tabby-ui/components/ui/combobox.tsx |  9 ++----
 2 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/ee/tabby-ui/components/prompt-form.tsx b/ee/tabby-ui/components/prompt-form.tsx
index dd6b5e1..165c4c9 100644
--- a/ee/tabby-ui/components/prompt-form.tsx
+++ b/ee/tabby-ui/components/prompt-form.tsx
@@ -25,7 +25,6 @@ import { cn } from '@/lib/utils'
 import fetcher from '@/lib/tabby-fetcher'
 import { debounce, has } from 'lodash-es'
 import type { ISearchHit, SearchReponse } from '@/lib/types'
-import { lightfair } from 'react-syntax-highlighter/dist/esm/styles/hljs'
 
 export interface PromptProps
   extends Pick<UseChatHelpers, 'input' | 'setInput'> {
@@ -45,6 +44,9 @@ export function PromptForm({
   >(null)
   const latestFetchKey = React.useRef('')
   const inputRef = React.useRef<HTMLTextAreaElement>(null)
+  // store the input selection for replacing inputValue
+  const prevInputSelectionEnd = React.useRef<number>()
+  // for updating the input selection after replacing
   const nextInputSelectionRange = React.useRef<[number, number]>()
   const [options, setOptions] = React.useState<SearchReponse['hits']>([])
   const [selectedCompletionsMap, setSelectedCompletionsMap] = React.useState<
@@ -53,10 +55,9 @@ export function PromptForm({
 
   useSWR<SearchReponse>(queryCompletionUrl, fetcher, {
     revalidateOnFocus: false,
-    dedupingInterval: 500,
+    dedupingInterval: 0,
     onSuccess: (data, key) => {
       if (key !== latestFetchKey.current) return
-
       setOptions(data?.hits ?? [])
     }
   })
@@ -88,21 +89,23 @@ export function PromptForm({
     }, 200)
   }, [])
 
-  const handleCompletionSelect = (
-    inputRef: React.RefObject<HTMLTextAreaElement | HTMLInputElement>,
-    item: ISearchHit
-  ) => {
-    const selectionEnd = inputRef.current?.selectionEnd ?? 0
+  const handleCompletionSelect = (item: ISearchHit) => {
+    const selectionEnd = prevInputSelectionEnd.current ?? 0
     const queryNameMatches = getSearchCompletionQueryName(input, selectionEnd)
     if (queryNameMatches) {
       setSelectedCompletionsMap({
         ...selectedCompletionsMap,
-        [queryNameMatches[0]]: item
+        [`@${item.doc?.name}`]: item
       })
-      // insert a space to break the search query
-      setInput(input.slice(0, selectionEnd) + ' ' + input.slice(selectionEnd))
+      const replaceString = `@${item?.doc?.name} `
+      const prevInput = input
+        .substring(0, selectionEnd)
+        .replace(new RegExp(queryNameMatches[0]), '')
+      const nextSelectionEnd = prevInput.length + replaceString.length
       // store the selection range and update it when layout
-      nextInputSelectionRange.current = [selectionEnd + 1, selectionEnd + 1]
+      nextInputSelectionRange.current = [nextSelectionEnd, nextSelectionEnd]
+      // insert a space to break the search query
+      setInput(prevInput + replaceString + input.slice(selectionEnd))
     }
     setOptions([])
   }
@@ -132,6 +135,17 @@ export function PromptForm({
     await onSubmit(finalInput)
   }
 
+  const handleTextareaKeyDown = (
+    e: React.KeyboardEvent<HTMLTextAreaElement>,
+    isOpen: boolean
+  ) => {
+    if (isOpen && ['ArrowRight', 'ArrowLeft', 'Home', 'End'].includes(e.key)) {
+      setOptions([])
+    } else {
+      onKeyDown(e)
+    }
+  }
+
   return (
     <form onSubmit={handlePromptSubmit} ref={formRef}>
       <Combobox
@@ -164,11 +178,14 @@ export function PromptForm({
                     ref={inputRef}
                     onChange={e => {
                       if (has(e, 'target.value')) {
+                        prevInputSelectionEnd.current = e.target.selectionEnd
                         setInput(e.target.value)
                         handleSearchCompletion(e)
+                      } else {
+                        prevInputSelectionEnd.current = undefined
                       }
                     }}
-                    onKeyDown={onKeyDown}
+                    onKeyDown={e => handleTextareaKeyDown(e, open)}
                   />
                   <div className="absolute right-0 top-4 sm:right-4">
                     <Tooltip>
@@ -189,6 +206,7 @@ export function PromptForm({
               </ComboboxAnchor>
               <ComboboxContent
                 align="start"
+                side="top"
                 onOpenAutoFocus={e => e.preventDefault()}
                 className="w-[60vw] md:w-[430px]"
               >
diff --git a/ee/tabby-ui/components/ui/combobox.tsx b/ee/tabby-ui/components/ui/combobox.tsx
index fda256c..a227236 100644
--- a/ee/tabby-ui/components/ui/combobox.tsx
+++ b/ee/tabby-ui/components/ui/combobox.tsx
@@ -48,7 +48,7 @@ export const ComboboxTextarea = React.forwardRef<
             onKeyDown: (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
               if (e.key === 'Enter' && open) {
                 e.preventDefault()
-              } else if (!open) {
+              } else {
                 onKeyDown?.(e)
               }
             },
@@ -154,10 +154,7 @@ ComboboxOption.displayName = 'ComboboxOption'
 
 interface ComboboxProps<T> {
   options: T[] | undefined
-  onSelect?: (
-    ref: React.RefObject<HTMLTextAreaElement | HTMLInputElement>,
-    data: T
-  ) => void
+  onSelect?: (data: T) => void
   inputRef?: React.RefObject<HTMLTextAreaElement | HTMLInputElement>
   children?:
     | React.ReactNode
@@ -202,7 +199,7 @@ export function Combobox<T extends { id: number }>({
     isOpen: manualOpen,
     onSelectedItemChange({ selectedItem }) {
       if (selectedItem) {
-        onSelect?.(inputRef, selectedItem)
+        onSelect?.(selectedItem)
         setManualOpen(false)
       }
     },