From e066db763a06f0a9dcff922fecb35271c0f9a92f Mon Sep 17 00:00:00 2001
From: one <wangan.cs@gmail.com>
Date: Thu, 6 Mar 2025 20:31:21 +0800
Subject: [PATCH] fix: distinguish model mention sources

---
 .../src/pages/home/Inputbar/Inputbar.tsx      | 22 ++++++++++++-------
 .../home/Inputbar/MentionModelsButton.tsx     | 14 +++++++++---
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
index bc440720..82a458ce 100644
--- a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
+++ b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
@@ -102,6 +102,8 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
 
   const [tokenCount, setTokenCount] = useState(0)
 
+  const [mentionFromKeyboard, setMentionFromKeyboard] = useState(false)
+
   const debouncedEstimate = useCallback(
     debounce((newText) => {
       if (showInputEstimatedTokens) {
@@ -192,6 +194,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
         const cursorPosition = textArea.selectionStart
         const textBeforeCursor = text.substring(0, cursorPosition)
         if (cursorPosition === 0 || textBeforeCursor.endsWith(' ')) {
+          setMentionFromKeyboard(true)
           EventEmitter.emit(EVENT_NAMES.SHOW_MODEL_SELECTOR)
           setIsMentionPopupOpen(true)
           return
@@ -557,16 +560,18 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
     setSelectedKnowledgeBases(bases ?? [])
   }
 
-  const onMentionModel = (model: Model) => {
+  const onMentionModel = (model: Model, fromKeyboard: boolean = false) => {
     const textArea = textareaRef.current?.resizableTextArea?.textArea
     if (textArea) {
-      const cursorPosition = textArea.selectionStart
-      const textBeforeCursor = text.substring(0, cursorPosition)
-      const lastAtIndex = textBeforeCursor.lastIndexOf('@')
+      if (fromKeyboard) {
+        const cursorPosition = textArea.selectionStart
+        const textBeforeCursor = text.substring(0, cursorPosition)
+        const lastAtIndex = textBeforeCursor.lastIndexOf('@')
 
-      if (lastAtIndex !== -1) {
-        const newText = text.substring(0, lastAtIndex) + text.substring(cursorPosition)
-        setText(newText)
+        if (lastAtIndex !== -1) {
+          const newText = text.substring(0, lastAtIndex) + text.substring(cursorPosition)
+          setText(newText)
+        }
       }
 
       setMentionModels((prev) => [...prev, model])
@@ -574,6 +579,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
       setTimeout(() => {
         textareaRef.current?.focus()
       }, 0)
+      setMentionFromKeyboard(false)
     }
   }
 
@@ -672,7 +678,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
               </Tooltip>
               <MentionModelsButton
                 mentionModels={mentionModels}
-                onMentionModel={onMentionModel}
+                onMentionModel={(model) => onMentionModel(model, mentionFromKeyboard)}
                 ToolbarButton={ToolbarButton}
               />
               <Tooltip placement="top" title={t('chat.input.web_search')} arrow>
diff --git a/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx b/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
index ba0d5761..75e1267e 100644
--- a/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
+++ b/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
@@ -14,7 +14,7 @@ import styled, { createGlobalStyle } from 'styled-components'
 
 interface Props {
   mentionModels: Model[]
-  onMentionModel: (model: Model) => void
+  onMentionModel: (model: Model, fromKeyboard?: boolean) => void
   ToolbarButton: any
 }
 
@@ -30,6 +30,8 @@ const MentionModelsButton: FC<Props> = ({ mentionModels, onMentionModel: onSelec
   const itemRefs = useRef<Array<HTMLDivElement | null>>([])
   // Add a new state to track if menu was dismissed
   const [menuDismissed, setMenuDismissed] = useState(false)
+  // Add a state to track if the model selector was triggered by keyboard
+  const [fromKeyboard, setFromKeyboard] = useState(false)
 
   const setItemRef = (index: number, el: HTMLDivElement | null) => {
     itemRefs.current[index] = el
@@ -49,7 +51,7 @@ const MentionModelsButton: FC<Props> = ({ mentionModels, onMentionModel: onSelec
     if (mentionModels.some((selected) => getModelUniqId(selected) === getModelUniqId(model))) {
       return
     }
-    onSelect(model)
+    onSelect(model, fromKeyboard)
     setIsOpen(false)
   }
 
@@ -190,6 +192,7 @@ const MentionModelsButton: FC<Props> = ({ mentionModels, onMentionModel: onSelec
       setSelectedIndex(0)
       setSearchText('')
       setMenuDismissed(false) // Reset dismissed flag when manually showing selector
+      setFromKeyboard(true) // Set fromKeyboard to true when triggered by keyboard
     }
 
     const handleKeyDown = (e: KeyboardEvent) => {
@@ -307,7 +310,12 @@ const MentionModelsButton: FC<Props> = ({ mentionModels, onMentionModel: onSelec
         dropdownRender={() => menu}
         trigger={['click']}
         open={isOpen}
-        onOpenChange={setIsOpen}
+        onOpenChange={(open) => {
+          setIsOpen(open)
+          if (open) {
+            setFromKeyboard(false) // Set fromKeyboard to false when opened by button click
+          }
+        }}
         overlayClassName="mention-models-dropdown">
         <Tooltip placement="top" title={t('agents.edit.model.select.title')} arrow>
           <ToolbarButton type="text" ref={dropdownRef}>