Refactor the Chat Input component to correctly handle and display #SavedContext keywords as interactive, deletable tokens, mirroring the behavior of other context keywords

Author: robertpiosik

packages/ui/src/components/editor/panel/ChatInput/ChatInput.module.scss

@@ -175,6 +175,10 @@
     content: '\ebdb';
   }
 
+  &--saved-context &__icon::before {
+    content: '\eab3';
+  }
+
   &--selection &__icon::before {
     content: '\eb84';
   }

packages/ui/src/components/editor/panel/ChatInput/ChatInput.tsx

@@ -300,6 +300,21 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
         apply_deletion(start_index, start_index + search_pattern.length)
       }
     }
+    // Check if it's a saved context keyword
+    else if (
+      keyword_element.classList.contains(styles['keyword--saved-context'])
+    ) {
+      const context_type = keyword_element.dataset.contextType
+      const context_name = keyword_element.dataset.contextName
+      if (!context_type || !context_name) return
+
+      const search_pattern = `#SavedContext:${context_type} "${context_name}"`
+      const start_index = props.value.indexOf(search_pattern)
+
+      if (start_index !== -1) {
+        apply_deletion(start_index, start_index + search_pattern.length)
+      }
+    }
     // Check if it's a selection keyword
     else if (keyword_element.classList.contains(styles['keyword--selection'])) {
       const search_pattern = '#Selection'

packages/ui/src/components/editor/panel/ChatInput/hooks/use-handlers.ts

@@ -27,6 +27,13 @@ const reconstruct_raw_value_from_node = (node: Node): string => {
         const branchName = el.dataset.branchName
         return branchName ? `#Changes:${branchName}` : ''
       }
+      case 'saved-context-keyword': {
+        const contextType = el.dataset.contextType
+        const contextName = el.dataset.contextName
+        return contextType && contextName
+          ? `#SavedContext:${contextType} "${contextName}"`
+          : ''
+      }
       case 'selection-keyword':
         return '#Selection'
     }
@@ -194,6 +201,34 @@ export const use_handlers = (
     return false
   }
 
+  const handle_saved_context_keyword_deletion = (
+    raw_pos: number,
+    context_file_paths: string[]
+  ): boolean => {
+    const text_before_cursor = props.value.substring(0, raw_pos)
+    const match = text_before_cursor.match(
+      /#SavedContext:(?:WorkspaceState|JSON)\s+"[^"]+"$/
+    )
+
+    if (match) {
+      const start_of_match = raw_pos - match[0].length
+      const new_value =
+        props.value.substring(0, start_of_match) +
+        props.value.substring(raw_pos)
+      const new_raw_cursor_pos = start_of_match
+      props.on_change(new_value)
+
+      set_caret_position_after_change(
+        input_ref,
+        new_raw_cursor_pos,
+        new_value,
+        context_file_paths
+      )
+      return true
+    }
+    return false
+  }
+
   const handle_selection_keyword_deletion = (
     raw_pos: number,
     context_file_paths: string[]
@@ -239,6 +274,10 @@ export const use_handlers = (
       return handle_changes_keyword_deletion(raw_pos, context_file_paths)
     }
 
+    if (el.dataset.type === 'saved-context-keyword') {
+      return handle_saved_context_keyword_deletion(raw_pos, context_file_paths)
+    }
+
     if (el.dataset.type === 'selection-keyword') {
       return handle_selection_keyword_deletion(raw_pos, context_file_paths)
     }
@@ -301,7 +340,8 @@ export const use_handlers = (
       if (
         el.dataset.type === 'file-keyword' ||
         el.dataset.type === 'changes-keyword' ||
-        el.dataset.type === 'selection-keyword'
+        el.dataset.type === 'selection-keyword' ||
+        el.dataset.type === 'saved-context-keyword'
       ) {
         const display_pos = get_caret_position_from_div(input_ref.current)
         if (handle_keyword_deletion_by_backspace(el, display_pos)) {

packages/ui/src/components/editor/panel/ChatInput/utils/get-display-text.ts

@@ -2,18 +2,31 @@ export const get_display_text = (
   text: string,
   context_file_paths: string[]
 ): string => {
-  const regex = /`([^\s`]*\.[^\s`]+)`|(#Changes:[^\s,;:!?]+)/g
-  return text.replace(regex, (match, file_path, changes_keyword) => {
-    if (file_path) {
-      if (context_file_paths.includes(file_path)) {
-        const filename = file_path.split('/').pop() || file_path
-        return filename
+  const regex =
+    /`([^\s`]*\.[^\s`]+)`|(#Changes:[^\s,;:!?]+)|(#SavedContext:(?:WorkspaceState|JSON)\s+"([^"]+)")/g
+  return text.replace(
+    regex,
+    (
+      match,
+      file_path,
+      changes_keyword,
+      saved_context_keyword,
+      context_name
+    ) => {
+      if (file_path) {
+        if (context_file_paths.includes(file_path)) {
+          const filename = file_path.split('/').pop() || file_path
+          return filename
+        }
+        return match
+      }
+      if (changes_keyword) {
+        return 'Changes'
+      }
+      if (saved_context_keyword) {
+        return `Saved context ${context_name}`
       }
       return match
     }
-    if (changes_keyword) {
-      return 'Changes'
-    }
-    return match
-  })
+  )
 }

packages/ui/src/components/editor/panel/ChatInput/utils/get-highlighted-text.tsx

@@ -47,28 +47,41 @@ export const get_highlighted_text = (params: {
   has_active_selection: boolean
   context_file_paths: string[]
 }): string => {
+  const saved_context_regex_part =
+    '#SavedContext:(?:WorkspaceState|JSON)\\s+"[^"]+"'
   if (params.is_in_code_completions_mode) {
-    const regex = /(#SavedContext:(?:WorkspaceState|JSON)\s+"[^"]+")/g
+    const regex = new RegExp(`(${saved_context_regex_part})`, 'g')
     const parts = params.text.split(regex)
     return parts
       .map((part) => {
-        if (
-          part &&
-          /^#SavedContext:(?:WorkspaceState|JSON)\s+"[^"]+"$/.test(part)
-        ) {
-          return `<span class="${
-            styles['selection-keyword']
-          }" contentEditable="false" data-type="selection-keyword">${escape_html(
-            part
-          )}</span>`
+        const saved_context_match = part.match(
+          /^#SavedContext:(WorkspaceState|JSON)\s+"([^"]+)"$/
+        )
+        if (part && saved_context_match) {
+          const context_type = saved_context_match[1]
+          const context_name = saved_context_match[2]
+          return `<span class="${cn(
+            styles['keyword'],
+            styles['keyword--saved-context']
+          )}" contentEditable="false" data-type="saved-context-keyword" data-context-type="${escape_html(
+            context_type
+          )}" data-context-name="${escape_html(
+            context_name
+          )}" title="Saved Context: ${escape_html(
+            context_name
+          )}"><span class="${styles['keyword__icon']}"></span><span class="${
+            styles['keyword__text']
+          }">Saved context ${escape_html(context_name)}</span></span>`
         }
         return process_text_part_for_files(part, params.context_file_paths)
       })
       .join('')
   }
 
-  const regex =
-    /(#Selection|#Changes:[^\s,;:!?]+|#SavedContext:(?:WorkspaceState|JSON)\s+"[^"]+")/g
+  const regex = new RegExp(
+    `(#Selection|#Changes:[^\\s,;:!?]+|${saved_context_regex_part})`,
+    'g'
+  )
   const parts = params.text.split(regex)
 
   return parts
@@ -92,15 +105,24 @@ export const get_highlighted_text = (params: {
           styles['keyword__text']
         }">Changes</span></span>`
       }
-      if (
-        part &&
-        /^#SavedContext:(?:WorkspaceState|JSON)\s+"[^"]+"$/.test(part)
-      ) {
-        return `<span class="${
-          styles['selection-keyword']
-        }" contentEditable="false" data-type="selection-keyword">${escape_html(
-          part
-        )}</span>`
+      const saved_context_match = part.match(
+        /^#SavedContext:(WorkspaceState|JSON)\s+"([^"]+)"$/
+      )
+      if (part && saved_context_match) {
+        const context_type = saved_context_match[1]
+        const context_name = saved_context_match[2]
+        return `<span class="${cn(
+          styles['keyword'],
+          styles['keyword--saved-context']
+        )}" contentEditable="false" data-type="saved-context-keyword" data-context-type="${escape_html(
+          context_type
+        )}" data-context-name="${escape_html(
+          context_name
+        )}" title="Saved Context: ${escape_html(context_name)}"><span class="${
+          styles['keyword__icon']
+        }"></span><span class="${
+          styles['keyword__text']
+        }">Saved context ${escape_html(context_name)}</span></span>`
       }
 
       return process_text_part_for_files(part, params.context_file_paths)

packages/ui/src/components/editor/panel/ChatInput/utils/position-mapping.ts

@@ -7,13 +7,16 @@ export const map_display_pos_to_raw_pos = (
   let current_display_pos = 0
   let last_raw_index = 0
 
-  const regex = /`([^\s`]*\.[^\s`]+)`|(#Changes:[^\s,;:!?]+)|(#Selection)/g
+  const regex =
+    /`([^\s`]*\.[^\s`]+)`|(#Changes:[^\s,;:!?]+)|(#Selection)|(#SavedContext:(?:WorkspaceState|JSON)\s+"([^"]+)")/g
   let match
 
   while ((match = regex.exec(raw_text)) !== null) {
     const file_path = match[1]
     const changes_keyword = match[2]
     const selection_keyword = match[3]
+    const saved_context_keyword = match[4]
+    const context_name = match[5]
 
     let is_replacement_match = false
     let display_match_length = 0
@@ -28,6 +31,9 @@ export const map_display_pos_to_raw_pos = (
     } else if (selection_keyword) {
       display_match_length = 'Selection'.length
       is_replacement_match = true
+    } else if (saved_context_keyword) {
+      display_match_length = `Saved context ${context_name}`.length
+      is_replacement_match = true
     }
 
     if (!is_replacement_match) {
@@ -71,13 +77,16 @@ export const map_raw_pos_to_display_pos = (
   let current_raw_pos = 0
   let last_raw_index = 0
 
-  const regex = /`([^\s`]*\.[^\s`]+)`|(#Changes:[^\s,;:!?]+)|(#Selection)/g
+  const regex =
+    /`([^\s`]*\.[^\s`]+)`|(#Changes:[^\s,;:!?]+)|(#Selection)|(#SavedContext:(?:WorkspaceState|JSON)\s+"([^"]+)")/g
   let match
 
   while ((match = regex.exec(raw_text)) !== null) {
     const file_path = match[1]
     const changes_keyword = match[2]
     const selection_keyword = match[3]
+    const saved_context_keyword = match[4]
+    const context_name = match[5]
 
     let is_replacement_match = false
     let display_match_length = 0
@@ -92,6 +101,9 @@ export const map_raw_pos_to_display_pos = (
     } else if (selection_keyword) {
       display_match_length = 'Selection'.length
       is_replacement_match = true
+    } else if (saved_context_keyword) {
+      display_match_length = `Saved context ${context_name}`.length
+      is_replacement_match = true
     }
 
     if (!is_replacement_match) {

packages/vscode/src/views/panel/backend/utils/hash-sign-quick-pick.ts

@@ -12,19 +12,16 @@ export async function hash_sign_quick_pick(params: {
 }): Promise<string | undefined> {
   let items = [
     {
-      label: '#Selection',
-      description:
-        'Places symbol that is then replaced by text selection from the active editor'
+      label: '$(list-flat) Selection',
+      description: 'Text selection from the active editor'
     },
     {
-      label: '#Changes',
-      description:
-        'Places symbol that is then replaced by a diff of changes between the selected branch'
+      label: '$(git-pull-request-draft) Changes',
+      description: 'Diff of changes between the selected branch'
     },
     {
-      label: '#SavedContext',
-      description:
-        'Places symbol that is then replaced by files from a saved context'
+      label: '$(checklist) Saved context',
+      description: 'Include files from the saved context'
     }
   ]
 
@@ -43,11 +40,11 @@ export async function hash_sign_quick_pick(params: {
       return
     }
 
-    if (selected.label == '#Selection') {
+    if (selected.label == '$(list-flat) Selection') {
       return '#Selection '
     }
 
-    if (selected.label == '#Changes') {
+    if (selected.label == '$(git-pull-request-draft) Changes') {
       try {
         const workspace_folders = vscode.workspace.workspaceFolders
         if (!workspace_folders || workspace_folders.length == 0) {
@@ -134,7 +131,7 @@ export async function hash_sign_quick_pick(params: {
       continue
     }
 
-    if (selected.label == '#SavedContext') {
+    if (selected.label == '$(checklist) Saved context') {
       const workspace_root = vscode.workspace.workspaceFolders?.[0]?.uri.fsPath
       if (!workspace_root) {
         vscode.window.showErrorMessage(