From b0eac5fdf6fe0f4607295bc858ea73628208a0f1 Mon Sep 17 00:00:00 2001
From: Sma1lboy <541898146chen@gmail.com>
Date: Sun, 12 Jan 2025 23:19:26 -0500
Subject: [PATCH] refactor(chat): enhance at-mention handling and improve file
 item processing

---
 ee/tabby-ui/components/chat/chat.tsx          |  51 +++++++--
 .../components/chat/form-editor/types.ts      |  10 +-
 .../components/chat/form-editor/utils.tsx     | 107 +++++++++++++++++-
 ee/tabby-ui/components/chat/prompt-form.css   |  17 +++
 ee/tabby-ui/components/chat/prompt-form.tsx   |   7 +-
 .../components/message-markdown/index.tsx     |  27 +++++
 6 files changed, 198 insertions(+), 21 deletions(-)
 create mode 100644 ee/tabby-ui/components/chat/prompt-form.css

diff --git a/ee/tabby-ui/components/chat/chat.tsx b/ee/tabby-ui/components/chat/chat.tsx
index 0a28ce5374c5..2daba42a901c 100644
--- a/ee/tabby-ui/components/chat/chat.tsx
+++ b/ee/tabby-ui/components/chat/chat.tsx
@@ -51,7 +51,11 @@ import {
 import { ChatPanel, ChatPanelRef } from './chat-panel'
 import { ChatScrollAnchor } from './chat-scroll-anchor'
 import { EmptyScreen } from './empty-screen'
-import { FILEITEM_REGEX } from './form-editor/utils'
+import { FileItem } from './form-editor/types'
+import {
+  FILEITEM_REGEX,
+  replaceAtMentionPlaceHolderWithAt
+} from './form-editor/utils'
 import { QuestionAnswerList } from './question-answer'
 
 type ChatContextValue = {
@@ -278,7 +282,8 @@ function ChatRenderer(
     // delete message pair
     const nextQaPairs = qaPairs.filter(o => o.user.id !== userMessageId)
     setQaPairs(nextQaPairs)
-    setInput(userMessage.message)
+    // FIXME: put this transformer to somewhere else, both case in message markdown and edit could be cover by same method
+    setInput(replaceAtMentionPlaceHolderWithAt(userMessage.message))
     if (userMessage.activeContext) {
       openInEditor(getFileLocationFromContext(userMessage.activeContext))
     }
@@ -487,7 +492,13 @@ function ChatRenderer(
 
       setQaPairs(nextQaPairs)
 
-      sendUserMessage(...generateRequestPayload(newUserMessage))
+      // FIXME: we don't need to passing placeholder to backend
+      sendUserMessage(
+        ...generateRequestPayload({
+          ...newUserMessage,
+          message: replaceAtMentionPlaceHolderWithAt(userMessage.message)
+        })
+      )
     }
   )
 
@@ -508,26 +519,46 @@ function ChatRenderer(
   }
 
   const handleSubmit = async (value: string) => {
-    const fileItems: any[] = []
+    const fileItems: FileItem[] = []
     let newValue = value
-
     let match
     while ((match = FILEITEM_REGEX.exec(value)) !== null) {
       try {
         const parsedItem = JSON.parse(match[1])
         fileItems.push(parsedItem)
-
-        const replacement = `@${
+        const labelName =
           parsedItem.label.split('/').pop() || parsedItem.label || 'unknown'
-        }`
-        newValue = newValue.replace(match[0], replacement)
+        newValue = newValue.replace(match[0], `@${labelName}`)
       } catch (error) {
         continue
       }
     }
+
+    // read all at file and push to relevant context, which will request to backend server later
+    let fileContents: Context[] = []
+    if (readFileContent && fileItems.length > 0) {
+      fileContents = await Promise.all(
+        fileItems.map(async item => {
+          const content = await readFileContent({ filepath: item.filepath })
+          return {
+            filepath:
+              'filepath' in item.filepath
+                ? item.filepath.filepath
+                : item.filepath.uri,
+            content: content ?? '',
+            git_url:
+              'git_url' in item.filepath
+                ? (item.filepath.git_url as string)
+                : '',
+            kind: 'file'
+          }
+        })
+      )
+    }
+
     sendUserChat({
       message: value,
-      relevantContext: relevantContext
+      relevantContext: [...fileContents, ...relevantContext]
     })
 
     setRelevantContext([])
diff --git a/ee/tabby-ui/components/chat/form-editor/types.ts b/ee/tabby-ui/components/chat/form-editor/types.ts
index 975aff24e472..1d16251d2fd2 100644
--- a/ee/tabby-ui/components/chat/form-editor/types.ts
+++ b/ee/tabby-ui/components/chat/form-editor/types.ts
@@ -1,3 +1,5 @@
+import { ListFileItem } from 'tabby-chat-panel/index'
+
 /**
  * PromptProps defines the props for the PromptForm component.
  */
@@ -31,15 +33,11 @@ export interface PromptFormRef {
   input: string
 }
 
+// TODO: move this into chat-panel in next iterate
 /**
  * Represents a file item inside the workspace.
- * (You can add more properties if needed)
  */
-export interface FileItem {
-  label: string
-  id?: string
-  // ... any other fields that you might have
-}
+export type FileItem = ListFileItem
 
 /**
  * Represents a file source item for mention suggestions.
diff --git a/ee/tabby-ui/components/chat/form-editor/utils.tsx b/ee/tabby-ui/components/chat/form-editor/utils.tsx
index 55ad9e5405a8..e11e68755bfa 100644
--- a/ee/tabby-ui/components/chat/form-editor/utils.tsx
+++ b/ee/tabby-ui/components/chat/form-editor/utils.tsx
@@ -1,5 +1,6 @@
 import Mention from '@tiptap/extension-mention'
 import { NodeViewWrapper, ReactNodeViewRenderer } from '@tiptap/react'
+import { Filepath } from 'tabby-chat-panel/index'
 
 import { cn } from '@/lib/utils'
 
@@ -53,16 +54,17 @@ export function shortenLabel(label: string, suffixLength = 15): string {
  */
 export const MentionComponent = ({ node }: { node: any }) => {
   return (
-    <NodeViewWrapper className="inline">
+    <NodeViewWrapper className="inline-block align-middle -my-1">
       <span
         className={cn(
-          'inline-flex items-center rounded bg-muted px-1.5 py-0.5 text-sm font-medium text-white',
-          'ring-1 ring-inset ring-muted'
+          'bg-muted prose inline-flex items-center rounded px-1.5 py-0.5 text-sm font-medium text-white',
+          'ring-muted ring-1 ring-inset',
+          'relative top-[0.1em]'
         )}
         data-category={node.attrs.category}
       >
         <FileItemIcon />
-        <span>{node.attrs.name}</span>
+        <span className="relative top-[-0.5px]">{node.attrs.name}</span>
       </span>
     </NodeViewWrapper>
   )
@@ -193,3 +195,100 @@ export const MentionList = ({
     </div>
   )
 }
+
+// Some utils function help to extract place holder
+export function replaceAtMentionPlaceHolderWithAt(value: string) {
+  // eslint-disable-next-line no-console
+  console.log('before value: ', value)
+
+  let newValue = value
+
+  let match
+  while ((match = FILEITEM_REGEX.exec(value)) !== null) {
+    try {
+      const parsedItem = JSON.parse(match[1])
+      const labelName =
+        parsedItem.label.split('/').pop() || parsedItem.label || 'unknown'
+      newValue = newValue.replace(match[0], `@${labelName}`)
+    } catch (error) {
+      continue
+    }
+  }
+
+  // eslint-disable-next-line no-console
+  console.log('new value:', newValue)
+  return newValue
+}
+
+interface ReplaceResult {
+  newValue: string
+  fileItems: FileItem[]
+}
+
+export const FILEITEM_AT_REGEX = /\[\[fileItemAt: (\d+)\]\]/g
+
+// Some utils function help to extract place holder
+// replace at mention JSON placeholder to something like [[fileItemAt: id]] which ad is string
+// return a list of FileItem with unique id
+// also return string already replaced
+// example:
+//  [[fileItem:{"label":"src/CodeActions.ts","filepath":{"kind":"git","filepath":"clients/vscode/src/CodeActions.ts",
+// "gitUrl":"git@github.com:Sma1lboy/tabby.git"}}]] explain this  [[fileItem:{"label":"src/CodeActions.ts",
+// "filepath":{"kind":"git","filepath":"clients/vscode/src/CodeActions.ts","gitUrl":"git@github.com:Sma1lboy/tabby.git"}}]]
+// will replaced as [[fileItemAt: idx0]] explain this [[fileItemAt: idx1]]
+// and with [fileItem1, fileItem2]
+export function replaceAtMentionPlaceHolderWithAtPlaceHolder(
+  value: string
+): ReplaceResult {
+  // eslint-disable-next-line no-console
+  console.log('before value: ', value)
+
+  let newValue = value
+  let match
+  const fileItems: FileItem[] = []
+  let idx = 0
+
+  while ((match = FILEITEM_REGEX.exec(value)) !== null) {
+    try {
+      const parsedItem = JSON.parse(match[1])
+
+      fileItems.push({
+        ...parsedItem
+      })
+
+      newValue = newValue.replace(match[0], `[[fileItemAt: ${idx}]]`)
+
+      idx++
+    } catch (error) {
+      continue
+    }
+  }
+
+  // eslint-disable-next-line no-console
+  console.log('new value:', newValue)
+  return {
+    newValue,
+    fileItems
+  }
+}
+
+// utils function
+export function getFilepathStringByChatPanelFilePath(
+  filepath: Filepath
+): string {
+  return 'filepath' in filepath ? filepath.filepath : filepath.uri
+}
+
+export function getLastSegmentFromPath(filepath: string): string {
+  if (!filepath) {
+    return 'unknown'
+  }
+
+  const normalizedPath = filepath.replace(/\\/g, '/')
+
+  const cleanPath = normalizedPath.replace(/\/+$/, '')
+  const segments = cleanPath.split('/')
+  const lastSegment = segments[segments.length - 1]
+
+  return lastSegment || 'unknown'
+}
diff --git a/ee/tabby-ui/components/chat/prompt-form.css b/ee/tabby-ui/components/chat/prompt-form.css
new file mode 100644
index 000000000000..ab3b269a6e19
--- /dev/null
+++ b/ee/tabby-ui/components/chat/prompt-form.css
@@ -0,0 +1,17 @@
+.ProseMirror {
+  > * + * {
+    margin-top: 0.75em;
+  }
+
+  &:focus {
+    outline: none !important;
+  }
+
+  .ProseMirror-selectednode {
+    outline: none !important;
+  }
+
+  ::selection {
+    background: transparent;
+  }
+}
diff --git a/ee/tabby-ui/components/chat/prompt-form.tsx b/ee/tabby-ui/components/chat/prompt-form.tsx
index 7844834e2a03..31f453570c7d 100644
--- a/ee/tabby-ui/components/chat/prompt-form.tsx
+++ b/ee/tabby-ui/components/chat/prompt-form.tsx
@@ -11,6 +11,9 @@ import Placeholder from '@tiptap/extension-placeholder'
 import Text from '@tiptap/extension-text'
 import { EditorContent, useEditor } from '@tiptap/react'
 
+import './prompt-form.css'
+
+import { cn } from '@/lib/utils'
 import { Button } from '@/components/ui/button'
 import { IconArrowElbow, IconEdit } from '@/components/ui/icons'
 import { Popover, PopoverAnchor, PopoverContent } from '@/components/ui/popover'
@@ -283,7 +286,9 @@ function PromptFormRenderer(
             {/* TipTap editor content */}
             <EditorContent
               editor={editor}
-              className="prose overflow-hidden break-words text-white focus:outline-none"
+              className={cn(
+                'prose overflow-hidden break-words text-white focus:outline-none'
+              )}
             />
           </div>
           {anchorElement}
diff --git a/ee/tabby-ui/components/message-markdown/index.tsx b/ee/tabby-ui/components/message-markdown/index.tsx
index 1ba2c7aa7f85..6bc610ece9b7 100644
--- a/ee/tabby-ui/components/message-markdown/index.tsx
+++ b/ee/tabby-ui/components/message-markdown/index.tsx
@@ -31,6 +31,12 @@ import {
   MARKDOWN_SOURCE_REGEX
 } from '@/lib/constants/regex'
 
+import {
+  FILEITEM_AT_REGEX,
+  getFilepathStringByChatPanelFilePath,
+  getLastSegmentFromPath,
+  replaceAtMentionPlaceHolderWithAtPlaceHolder
+} from '../chat/form-editor/utils'
 import { Mention } from '../mention-tag'
 import { Skeleton } from '../ui/skeleton'
 import { CodeElement } from './code'
@@ -96,6 +102,11 @@ export function MessageMarkdown({
   activeSelection,
   ...rest
 }: MessageMarkdownProps) {
+  // deal with some unresolved file items
+  const msgRes = replaceAtMentionPlaceHolderWithAtPlaceHolder(message)
+  message = msgRes.newValue
+  const fileItems = msgRes.fileItems
+
   const [symbolPositionMap, setSymbolLocationMap] = useState<
     Map<string, SymbolInfo | undefined>
   >(new Map())
@@ -163,6 +174,14 @@ export function MessageMarkdown({
       return { sourceId, className }
     })
 
+    processMatches(FILEITEM_AT_REGEX, AtMentionTag, (match: string) => {
+      return {
+        label: getLastSegmentFromPath(
+          getFilepathStringByChatPanelFilePath(fileItems[+match[1]].filepath)
+        )
+      }
+    })
+
     addTextNode(text.slice(lastIndex))
 
     return elements
@@ -377,6 +396,14 @@ function SourceTag({
   )
 }
 
+function AtMentionTag({ label }: { label: string | undefined }) {
+  return (
+    <span className="bg-muted/50 hover:bg-muted text-muted-foreground prose inline-flex items-center rounded px-1.5 py-0.5 text-sm font-medium text-white">
+      @{label}
+    </span>
+  )
+}
+
 function RelevantDocumentBadge({
   relevantDocument,
   citationIndex