feat: optimize mention search

.vscode/settings.json

@@ -57,6 +57,7 @@
     "deepseek",
     "esno",
     "execa",
+    "flexsearch",
     "Flytek",
     "fullpath",
     "ianvs",

package.json

@@ -417,6 +417,7 @@
     "eslint-plugin-unused-imports": "^3.2.0",
     "esno": "^4.7.0",
     "execa": "^9.3.1",
+    "flexsearch": "^0.7.43",
     "form-data": "^4.0.0",
     "framer-motion": "^11.5.4",
     "fs-extra": "^11.2.0",

src/extension/polyfill.ts

@@ -17,21 +17,21 @@ const enableFetchPolyfill = async () => {
 
   // fuck, vscode fetch not working on v1.92.0
   // we add a polyfill here
-  const {
-    default: fetch,
-    Headers,
-    Request,
-    Response
-  } = await import('node-fetch')
-  const { default: FormData } = await import('form-data')
+  // const {
+  //   default: fetch,
+  //   Headers,
+  //   Request,
+  //   Response
+  // } = await import('node-fetch')
+  // const { default: FormData } = await import('form-data')
 
-  Object.assign(globalThis, {
-    fetch,
-    FormData,
-    Headers,
-    Request,
-    Response
-  })
+  // Object.assign(globalThis, {
+  //   fetch,
+  //   FormData,
+  //   Headers,
+  //   Request,
+  //   Response
+  // })
 }
 
 export const enablePolyfill = async () => {

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

@@ -1,4 +1,4 @@
-import React, { useEffect, useMemo, useRef, useState } from 'react'
+import React, { useEffect, useRef, useState } from 'react'
 import {
   Command,
   CommandEmpty,
@@ -11,6 +11,7 @@ import {
   PopoverContent,
   PopoverTrigger
 } from '@webview/components/ui/popover'
+import { useFilteredMentionOptions } from '@webview/hooks/chat/use-filtered-mention-options'
 import { useControllableState } from '@webview/hooks/use-controllable-state'
 import { useKeyboardNavigation } from '@webview/hooks/use-keyboard-navigation'
 import { IMentionStrategy, MentionOption } from '@webview/types/chat'
@@ -43,30 +44,24 @@ export const MentionSelector: React.FC<MentionSelectorProps> = ({
   children
 }) => {
   const commandRef = useRef<HTMLDivElement>(null)
-  const [currentOptions, setCurrentOptions] =
-    useState<MentionOption[]>(mentionOptions)
   const maxItemLength = mentionOptions.length > 8 ? mentionOptions.length : 8
+  const [optionsStack, setOptionsStack] = useState<MentionOption[][]>([
+    mentionOptions
+  ])
 
   const [isOpen = false, setIsOpen] = useControllableState({
     prop: open,
     defaultProp: false,
     onChange: onOpenChange
   })
 
-  useEffect(() => {
-    if (!isOpen) {
-      setCurrentOptions(mentionOptions)
-    }
-  }, [isOpen, mentionOptions])
+  const currentOptions = optionsStack[optionsStack.length - 1] || []
 
-  const filteredOptions = useMemo(() => {
-    if (!searchQuery) return currentOptions.slice(0, maxItemLength)
-    return currentOptions
-      .filter(option =>
-        option.label.toLowerCase().includes(searchQuery.toLowerCase())
-      )
-      .slice(0, maxItemLength)
-  }, [currentOptions, searchQuery, maxItemLength])
+  const { filteredOptions, isFlattened } = useFilteredMentionOptions({
+    currentOptions,
+    searchQuery,
+    maxItemLength
+  })
 
   const itemRefs = useRef<(HTMLDivElement | null)[]>([])
   const { focusedIndex, setFocusedIndex, handleKeyDown } =
@@ -82,9 +77,28 @@ export const MentionSelector: React.FC<MentionSelectorProps> = ({
     setFocusedIndex(0)
   }, [filteredOptions])
 
+  useEffect(() => {
+    if (!isOpen) {
+      setOptionsStack([mentionOptions])
+    }
+  }, [isOpen, mentionOptions])
+
   const handleSelect = (option: MentionOption) => {
+    if (isFlattened) {
+      if (option.mentionStrategy) {
+        onSelect({
+          name: option.label,
+          strategy: option.mentionStrategy,
+          strategyAddData: option.data || { label: option.label }
+        })
+      }
+      setIsOpen(false)
+      return
+    }
+
+    // not flattened
     if (option.children) {
-      setCurrentOptions(option.children)
+      setOptionsStack(prevStack => [...prevStack, option.children || []])
       onCloseWithoutSelect?.()
     } else {
       if (option.mentionStrategy) {

src/webview/hooks/chat/use-filtered-mention-options.ts

@@ -0,0 +1,159 @@
+import { useEffect, useMemo, useRef, useState } from 'react'
+import { MentionOption, SearchSortStrategy } from '@webview/types/chat'
+import Index from 'flexsearch/dist/module'
+
+const flattenCurrentLevelOptions = (
+  options: MentionOption[]
+): MentionOption[] =>
+  options.reduce((acc: MentionOption[], option) => {
+    if (option.children) {
+      return [...acc, ...option.children]
+    }
+    return [...acc, option]
+  }, [])
+
+export interface UseFilteredMentionOptions {
+  currentOptions: MentionOption[]
+  searchQuery: string
+  maxItemLength: number
+}
+
+export const useFilteredMentionOptions = (props: UseFilteredMentionOptions) => {
+  const { currentOptions, searchQuery, maxItemLength } = props
+  const [isFlattened, setIsFlattened] = useState(false)
+  const currentOptionsSearchServiceRef = useRef<SearchService>(
+    new SearchService()
+  )
+  const flattenedOptionsSearchServiceRef = useRef<SearchService>(
+    new SearchService()
+  )
+
+  useEffect(() => {
+    currentOptionsSearchServiceRef.current.indexOptions(currentOptions)
+  }, [currentOptions])
+
+  const filteredOptions = useMemo(() => {
+    if (!searchQuery) return currentOptions.slice(0, maxItemLength)
+
+    let matches = currentOptionsSearchServiceRef.current.search(searchQuery)
+
+    if (matches.length > 0) {
+      setIsFlattened(false)
+      return matches.slice(0, maxItemLength)
+    }
+
+    // If no matches, try flattened options
+    const flattenedOptions = flattenCurrentLevelOptions(currentOptions)
+    flattenedOptionsSearchServiceRef.current.indexOptions(flattenedOptions)
+    matches = flattenedOptionsSearchServiceRef.current.search(searchQuery)
+
+    if (matches.length > 0) {
+      setIsFlattened(true)
+      return matches.slice(0, maxItemLength)
+    }
+
+    setIsFlattened(false)
+    return []
+  }, [searchQuery, currentOptions, maxItemLength])
+
+  return { filteredOptions, isFlattened }
+}
+
+class SearchService {
+  private index!: Index
+
+  private optionsMap: Map<string, MentionOption> = new Map()
+
+  constructor() {
+    this.init()
+  }
+
+  init() {
+    this.index = new Index({
+      tokenize: 'full',
+      cache: true,
+      optimize: true,
+      // 中文 https://github.com/nextapps-de/flexsearch?tab=readme-ov-file#cjk-word-break-chinese-japanese-korean
+      // 同时支持中文和英文搜索 https://github.com/nextapps-de/flexsearch/issues/202
+      encode(str: string) {
+        // eslint-disable-next-line no-control-regex
+        const cjkItems = str.replace(/[\x00-\x7F]/g, '').split('')
+        const asciiItems = str.split(/\W+/)
+        return cjkItems.concat(asciiItems)
+      }
+    })
+    this.optionsMap.clear()
+  }
+
+  indexOptions(options: MentionOption[]) {
+    this.init()
+    options.forEach(option => {
+      const { id } = option // Use a unique identifier, preferably option.id if available
+      this.optionsMap.set(id, option)
+      this.index.add(id, option.label)
+      option.searchKeywords?.forEach(keyword => {
+        this.index.add(id, keyword)
+      })
+    })
+  }
+
+  search(query: string): MentionOption[] {
+    const results = this.index.search(query) as string[]
+    const matchedOptions = results
+      .map(id => this.optionsMap.get(id))
+      .filter(Boolean) as MentionOption[]
+
+    return this.sortOptionsByStrategy(query, matchedOptions)
+  }
+
+  private sortOptionsByStrategy(
+    query: string,
+    options: MentionOption[]
+  ): MentionOption[] {
+    return options.sort((a, b) => {
+      const scoreA = this.getMatchScore(query, a)
+      const scoreB = this.getMatchScore(query, b)
+
+      // Higher scores come first
+      return scoreB - scoreA
+    })
+  }
+
+  private getMatchScore(query: string, option: MentionOption): number {
+    const label = option.label.toLowerCase()
+    const q = query.toLowerCase()
+
+    // Exact match gets the highest score
+    if (label === q) return 1000
+
+    // Prefix match is second best
+    if (label.startsWith(q)) return 500 + q.length / label.length
+
+    // EndMatch strategy
+    if (option.searchSortStrategy === SearchSortStrategy.EndMatch) {
+      // Calculate the longest matching length from the end
+      let matchLength = 0
+      for (let i = 1; i <= Math.min(label.length, q.length); i++) {
+        if (label.slice(-i) === q.slice(-i)) {
+          matchLength = i
+        } else {
+          break
+        }
+      }
+
+      // If the query is a suffix of the label, give a higher score
+      if (matchLength === q.length) {
+        return 200 + matchLength
+      }
+
+      // Partial end match
+      return matchLength
+    }
+
+    // Contains match
+    if (label.includes(q)) return 50
+
+    // No match
+    return 0
+  }
+}