fix: fix cursor position bug

src/webview/components/chat/editor/chat-editor.tsx

@@ -1,4 +1,4 @@
-import { forwardRef, useImperativeHandle } from 'react'
+import { forwardRef, useCallback, useImperativeHandle } from 'react'
 import { AutoFocusPlugin } from '@lexical/react/LexicalAutoFocusPlugin'
 import {
   LexicalComposer,
@@ -18,7 +18,12 @@ import {
   type MentionPluginProps
 } from '@webview/lexical/plugins/mention-plugin'
 import { cn } from '@webview/utils/common'
-import type { EditorState, LexicalEditor } from 'lexical'
+import {
+  $getSelection,
+  $isRangeSelection,
+  type EditorState,
+  type LexicalEditor
+} from 'lexical'
 
 const onError = (error: unknown) => {
   // eslint-disable-next-line no-console
@@ -46,6 +51,7 @@ export interface ChatEditorProps
 
 export interface ChatEditorRef {
   editor: LexicalEditor
+  insertSpaceAndAt: () => void
 }
 
 export const ChatEditor = forwardRef<ChatEditorRef, ChatEditorProps>(
@@ -103,7 +109,17 @@ const ChatEditorInner = forwardRef<ChatEditorRef, ChatEditorProps>(
   ) => {
     const [editor] = useLexicalComposerContext()
 
-    useImperativeHandle(ref, () => ({ editor }), [editor])
+    const insertSpaceAndAt = useCallback(() => {
+      editor.focus()
+      editor.update(() => {
+        const selection = $getSelection()
+        if ($isRangeSelection(selection)) {
+          selection.insertText(' @')
+        }
+      })
+    }, [editor])
+
+    useImperativeHandle(ref, () => ({ editor, insertSpaceAndAt }), [editor])
 
     const handleKeyDown = (e: React.KeyboardEvent<HTMLDivElement>) => {
       if (e.key === 'Enter' && e.ctrlKey) {

src/webview/components/chat/editor/chat-input.tsx

@@ -65,6 +65,9 @@ export const ChatInput: React.FC<ChatInputProps> = ({
             setContext={setContext}
             conversation={conversation}
             setConversation={setConversation}
+            onClickMentionSelector={() => {
+              editorRef.current?.insertSpaceAndAt()
+            }}
             onClose={focusOnEditor}
           />
           <Button

src/webview/components/chat/selectors/context-selector.tsx

@@ -2,18 +2,13 @@
 import React, { useState } from 'react'
 import { ImageIcon } from '@radix-ui/react-icons'
 import { Button } from '@webview/components/ui/button'
-import { createMentionOptions } from '@webview/lexical/mentions'
 import type {
   ChatContext,
   Conversation,
   ModelOption
 } from '@webview/types/chat'
 import type { Updater } from 'use-immer'
 
-import {
-  MentionSelector,
-  type SelectedMentionStrategy
-} from './mention-selector'
 import { ModelSelector } from './model-selector'
 
 interface ContextSelectorProps {
@@ -22,14 +17,16 @@ interface ContextSelectorProps {
   conversation: Conversation
   setConversation: Updater<Conversation>
   onClose?: () => void
+  onClickMentionSelector?: () => void
 }
 
 export const ContextSelector: React.FC<ContextSelectorProps> = ({
   context,
   setContext,
   conversation,
   setConversation,
-  onClose
+  onClose,
+  onClickMentionSelector
 }) => {
   const [modelOptions] = useState<ModelOption[]>([
     { value: 'gpt-4', label: 'GPT-4' },
@@ -45,11 +42,6 @@ export const ContextSelector: React.FC<ContextSelectorProps> = ({
     })
   }
 
-  const handleSelectMention = (option: SelectedMentionStrategy) => {
-    // Handle mention selection
-    console.log('Selected mention:', option)
-  }
-
   const handleSelectImage = () => {
     const input = document.createElement('input')
     input.type = 'file'
@@ -79,15 +71,9 @@ export const ContextSelector: React.FC<ContextSelectorProps> = ({
           {selectedModel?.label}
         </Button>
       </ModelSelector>
-      <MentionSelector
-        mentionOptions={createMentionOptions()}
-        onSelect={handleSelectMention}
-        onOpenChange={isOpen => !isOpen && onClose?.()}
-      >
-        <Button variant="ghost" size="xs">
-          @ Mention
-        </Button>
-      </MentionSelector>
+      <Button variant="ghost" size="xs" onClick={onClickMentionSelector}>
+        @ Mention
+      </Button>
       <Button variant="ghost" size="xs" onClick={handleSelectImage}>
         <ImageIcon className="h-3 w-3 mr-1" />
         Image

src/webview/components/chat/selectors/mention-selector.tsx

@@ -1,9 +1,8 @@
-import { useCallback, useEffect, useRef, type FC, type ReactNode } from 'react'
+import React, { useEffect, useMemo, useRef, useState } from 'react'
 import {
   Command,
   CommandEmpty,
   CommandGroup,
-  CommandInput,
   CommandItem,
   CommandList
 } from '@webview/components/ui/command'
@@ -12,157 +11,120 @@ import {
   PopoverContent,
   PopoverTrigger
 } from '@webview/components/ui/popover'
-import { useCallbackRef } from '@webview/hooks/use-callback-ref'
 import { useControllableState } from '@webview/hooks/use-controllable-state'
-import type { IMentionStrategy, MentionOption } from '@webview/types/chat'
+import { useKeyboardNavigation } from '@webview/hooks/use-keyboard-navigation'
+import { IMentionStrategy, MentionOption } from '@webview/types/chat'
+import { cn } from '@webview/utils/common'
+import { useEvent } from 'react-use'
 
 export interface SelectedMentionStrategy {
   strategy: IMentionStrategy
   strategyAddData: any
 }
 
 interface MentionSelectorProps {
+  searchQuery?: string
   mentionOptions: MentionOption[]
   onSelect: (option: SelectedMentionStrategy) => void
   open?: boolean
   onOpenChange?: (open: boolean) => void
-  lexicalMode?: boolean
-  searchQuery?: string
-  onSearchQueryChange?: (searchQuery: string) => void
-  children: ReactNode
+  onCloseWithoutSelect?: () => void
+  children: React.ReactNode
 }
 
-export const MentionSelector: FC<MentionSelectorProps> = ({
+export const MentionSelector: React.FC<MentionSelectorProps> = ({
+  searchQuery = '',
   mentionOptions,
   onSelect,
   open,
   onOpenChange,
-  lexicalMode,
-  searchQuery,
-  onSearchQueryChange,
+  onCloseWithoutSelect,
   children
 }) => {
   const commandRef = useRef<HTMLDivElement>(null)
+  const [currentOptions, setCurrentOptions] =
+    useState<MentionOption[]>(mentionOptions)
 
   const [isOpen = false, setIsOpen] = useControllableState({
     prop: open,
     defaultProp: false,
     onChange: onOpenChange
   })
 
-  const [internalSearchQuery, setInternalSearchQuery] = useControllableState({
-    prop: searchQuery,
-    defaultProp: '',
-    onChange: onSearchQueryChange
-  })
-
   useEffect(() => {
-    if (!lexicalMode || mentionOptions.length > 0) return
-
-    setIsOpen(false)
-  }, [mentionOptions.length, setIsOpen, lexicalMode])
-
-  const handleSelect = useCallback(
-    (currentValue: string) => {
-      const selectedOption = mentionOptions.find(
-        option => option.category === currentValue
-      )
-      if (selectedOption) {
-        onSelect({
-          strategy: selectedOption.mentionStrategies[0]!,
-          strategyAddData: { label: selectedOption.label }
-        })
-      }
-      setIsOpen(false)
-    },
-    [mentionOptions, onSelect]
-  )
-
-  const filterMentions = useCallback(
-    (value: string, search: string) => {
-      const option = mentionOptions.find(opt => opt.category === value)
-      if (!option) return 0
-
-      const label = option.label.toLowerCase()
-      search = search.toLowerCase()
-
-      if (label === search) return 1
-      if (label.startsWith(search)) return 0.8
-      if (label.includes(search)) return 0.6
-
-      // Calculate a basic fuzzy match score
-      let score = 0
-      let searchIndex = 0
-      for (let i = 0; i < label.length && searchIndex < search.length; i++) {
-        if (label[i] === search[searchIndex]) {
-          score += 1
-          searchIndex++
-        }
-      }
-      return score / Math.max(label.length, search.length)
-    },
-    [mentionOptions]
-  )
-
-  const handleKeyDown = useCallbackRef((event: KeyboardEvent) => {
-    if (!lexicalMode || !isOpen) return
+    if (!isOpen) {
+      setCurrentOptions(mentionOptions)
+    }
+  }, [isOpen, mentionOptions])
 
-    const commandEl = commandRef.current
-    if (!commandEl) return
+  const filteredOptions = useMemo(() => {
+    if (!searchQuery) return currentOptions
+    return currentOptions.filter(option =>
+      option.label.toLowerCase().includes(searchQuery.toLowerCase())
+    )
+  }, [currentOptions, searchQuery])
 
-    if (['ArrowUp', 'ArrowDown', 'Enter'].includes(event.key)) {
-      event.preventDefault()
+  const itemRefs = useRef<(HTMLDivElement | null)[]>([])
+  const { focusedIndex, setFocusedIndex, handleKeyDown } =
+    useKeyboardNavigation({
+      itemCount: filteredOptions.length,
+      itemRefs,
+      onEnter: el => el?.click()
+    })
 
-      const syntheticEvent = new KeyboardEvent(event.type, {
-        key: event.key,
-        code: event.code,
-        isComposing: event.isComposing,
-        location: event.location,
-        repeat: event.repeat,
-        bubbles: true
-      })
-      commandEl.dispatchEvent(syntheticEvent)
-    }
-  })
+  useEvent('keydown', handleKeyDown)
 
   useEffect(() => {
-    if (!lexicalMode) return
-
-    document.addEventListener('keydown', handleKeyDown)
-    return () => {
-      document.removeEventListener('keydown', handleKeyDown)
+    setFocusedIndex(0)
+  }, [filteredOptions])
+
+  const handleSelect = (option: MentionOption) => {
+    if (option.children) {
+      setCurrentOptions(option.children)
+      onCloseWithoutSelect?.()
+    } else {
+      onSelect({
+        strategy: option.mentionStrategies[0]!,
+        strategyAddData: option.data || { label: option.label }
+      })
+      setIsOpen(false)
     }
-  }, [lexicalMode, handleKeyDown, isOpen])
+  }
 
   return (
     <Popover open={isOpen} onOpenChange={setIsOpen}>
       <PopoverTrigger asChild>{children}</PopoverTrigger>
       <PopoverContent
-        className="w-[200px] p-0"
+        className={cn('w-[200px] p-0', !isOpen && 'hidden')}
         updatePositionStrategy="always"
         side="top"
-        onOpenAutoFocus={e => lexicalMode && e.preventDefault()}
-        onCloseAutoFocus={e => lexicalMode && e.preventDefault()}
+        align="start"
+        onOpenAutoFocus={e => e.preventDefault()}
+        onCloseAutoFocus={e => e.preventDefault()}
         onKeyDown={e => e.stopPropagation()}
       >
-        <Command ref={commandRef} filter={filterMentions}>
-          <CommandInput
-            hidden={lexicalMode}
-            showSearchIcon={false}
-            placeholder="Search mention..."
-            className="h-9"
-            value={internalSearchQuery}
-            onValueChange={setInternalSearchQuery}
-          />
+        <Command ref={commandRef} shouldFilter={false}>
           <CommandList>
-            <CommandEmpty>No mention type found.</CommandEmpty>
-            <CommandGroup>
-              {mentionOptions.map(option => (
+            <CommandEmpty>No results found.</CommandEmpty>
+            <CommandGroup
+              className={cn(filteredOptions.length === 0 ? 'p-0' : 'p-1')}
+            >
+              {filteredOptions.map((option, index) => (
                 <CommandItem
-                  key={option.category}
-                  value={option.category}
-                  onSelect={handleSelect}
-                  className="px-1.5 py-1"
+                  key={option.label}
+                  defaultValue=""
+                  value=""
+                  onSelect={() => handleSelect(option)}
+                  className={cn(
+                    'px-1.5 py-1',
+                    focusedIndex === index &&
+                      'bg-primary text-primary-foreground'
+                  )}
+                  ref={el => {
+                    if (itemRefs.current) {
+                      itemRefs.current[index] = el
+                    }
+                  }}
                 >
                   {option.label}
                 </CommandItem>

src/webview/components/ui/command.tsx

@@ -78,7 +78,7 @@ const CommandEmpty = React.forwardRef<
 >((props, ref) => (
   <CommandPrimitive.Empty
     ref={ref}
-    className="py-6 text-center text-sm"
+    className="relative flex cursor-default select-none items-center rounded-sm px-2 py-1.5 text-sm outline-none"
     {...props}
   />
 ))