Rename context merge/replace state keys and standardize merge/replace behavior across all file selection flows

Author: robertpiosik

apps/editor/src/commands/apply-context-command/sources/commit-files-source.ts

@@ -3,7 +3,7 @@ import * as path from 'path'
 import * as fs from 'fs'
 import { execSync } from 'child_process'
 import { WorkspaceProvider } from '../../../context/providers/workspace/workspace-provider'
-import { LAST_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY } from '../../../constants/state-keys'
+import { LAST_APPLY_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY } from '../../../constants/state-keys'
 import { GIT_LOG_SINCE_DURATION } from '../../../constants/values'
 import { get_git_repository } from '@/utils/git-repository-utils'
 import { Logger } from '@shared/utils/logger'
@@ -274,7 +274,7 @@ export const handle_commit_files_source = async (
 
             const last_choice_label =
               extension_context.workspaceState.get<string>(
-                LAST_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY
+                LAST_APPLY_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY
               )
 
             const quick_pick = vscode.window.createQuickPick()
@@ -321,7 +321,7 @@ export const handle_commit_files_source = async (
               return
             } else {
               await extension_context.workspaceState.update(
-                LAST_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY,
+                LAST_APPLY_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY,
                 choice.label
               )
 

apps/editor/src/commands/apply-context-command/sources/unstaged-files-source.ts

@@ -5,7 +5,7 @@ import { WorkspaceProvider } from '../../../context/providers/workspace/workspac
 import { dictionary } from '@shared/constants/dictionary'
 import { Logger } from '@shared/utils/logger'
 import { display_token_count } from '../../../utils/display-token-count'
-import { LAST_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY } from '../../../constants/state-keys'
+import { LAST_APPLY_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY } from '../../../constants/state-keys'
 import { t } from '@/i18n'
 
 export const handle_unstaged_files_source = async (
@@ -200,7 +200,7 @@ export const handle_unstaged_files_source = async (
 
           const last_choice_label =
             extension_context.workspaceState.get<string>(
-              LAST_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY
+              LAST_APPLY_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY
             )
 
           const quick_pick_merge = vscode.window.createQuickPick()
@@ -250,7 +250,7 @@ export const handle_unstaged_files_source = async (
             return
           } else {
             await extension_context.workspaceState.update(
-              LAST_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY,
+              LAST_APPLY_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY,
               choice.label
             )
 

apps/editor/src/commands/find-relevant-files-command/find-relevant-files-command.ts

@@ -131,7 +131,8 @@ export const find_relevant_files_command = (
             const apply_result = await show_results_and_apply({
               extracted_files: api_result,
               analysis,
-              workspace_provider
+              workspace_provider,
+              extension_context
             })
 
             if (apply_result === 'back') {

apps/editor/src/commands/find-relevant-files-command/utils/show-results-and-apply.ts

@@ -5,20 +5,26 @@ import { WorkspaceProvider } from '../../../context/providers/workspace/workspac
 import { display_token_count } from '../../../utils/display-token-count'
 import { FileAnalysisResult } from './analyze-workspace-files'
 import { t } from '@/i18n'
+import { dictionary } from '@shared/constants/dictionary'
+import { LAST_FIND_RELEVANT_FILES_MERGE_REPLACE_OPTION_STATE_KEY } from '../../../constants/state-keys'
 
 export const show_results_and_apply = async (params: {
   extracted_files: string[]
   analysis: FileAnalysisResult
   workspace_provider: WorkspaceProvider
+  extension_context: vscode.ExtensionContext
 }): Promise<'success' | 'back' | 'cancel'> => {
-  const currently_checked = params.workspace_provider.get_checked_files()
   const absolute_paths: string[] = []
 
   for (const rel_path of params.extracted_files) {
     const potential_abs = path.join(params.analysis.workspace_root, rel_path)
     if (fs.existsSync(potential_abs)) absolute_paths.push(potential_abs)
   }
 
+  const files_in_searched_folder = params.analysis.files_data.map(
+    (f) => f.file_path
+  )
+
   const close_button = {
     iconPath: new vscode.ThemeIcon('close'),
     tooltip: t('common.close')
@@ -51,88 +57,203 @@ export const show_results_and_apply = async (params: {
     })
   )
 
-  const quick_pick = vscode.window.createQuickPick<
-    vscode.QuickPickItem & { file_path: string }
-  >()
-  quick_pick.items = quick_pick_items
-  quick_pick.selectedItems = quick_pick_items.filter((item) =>
-    currently_checked.includes(item.file_path)
-  )
-  quick_pick.canSelectMany = true
-  quick_pick.title = t('command.find-relevant-files.quick-pick.title')
-  quick_pick.placeholder = t(
-    'command.find-relevant-files.quick-pick.placeholder'
-  )
-  quick_pick.ignoreFocusOut = true
-  quick_pick.buttons = [vscode.QuickInputButtons.Back, close_button]
+  while (true) {
+    const currently_checked = params.workspace_provider.get_checked_files()
 
-  return new Promise<'success' | 'back' | 'cancel'>((resolve) => {
-    let is_resolved = false
+    const quick_pick = vscode.window.createQuickPick<
+      vscode.QuickPickItem & { file_path: string }
+    >()
+    quick_pick.items = quick_pick_items
+    quick_pick.selectedItems = quick_pick_items.filter((item) =>
+      currently_checked.includes(item.file_path)
+    )
+    quick_pick.canSelectMany = true
+    quick_pick.title = t('command.find-relevant-files.quick-pick.title')
+    quick_pick.placeholder = t(
+      'command.find-relevant-files.quick-pick.placeholder'
+    )
+    quick_pick.ignoreFocusOut = true
+    quick_pick.buttons = [vscode.QuickInputButtons.Back, close_button]
 
-    quick_pick.onDidTriggerButton((button) => {
-      if (button === vscode.QuickInputButtons.Back) {
-        is_resolved = true
-        resolve('back')
-        quick_pick.hide()
-      } else if (button === close_button) {
+    const list_selection = await new Promise<
+      | readonly (vscode.QuickPickItem & { file_path: string })[]
+      | 'back'
+      | 'cancel'
+    >((resolve) => {
+      let is_resolved = false
+
+      quick_pick.onDidTriggerButton((button) => {
+        if (button === vscode.QuickInputButtons.Back) {
+          is_resolved = true
+          resolve('back')
+          quick_pick.hide()
+        } else if (button === close_button) {
+          is_resolved = true
+          resolve('cancel')
+          quick_pick.hide()
+        }
+      })
+
+      quick_pick.onDidAccept(() => {
         is_resolved = true
-        resolve('cancel')
+        resolve(quick_pick.selectedItems)
         quick_pick.hide()
-      }
+      })
+
+      quick_pick.onDidTriggerItemButton(async (e) => {
+        if (e.button === open_file_button) {
+          try {
+            const doc = await vscode.workspace.openTextDocument(
+              e.item.file_path
+            )
+            await vscode.window.showTextDocument(doc, { preview: true })
+          } catch (error) {
+            vscode.window.showErrorMessage(
+              t('command.find-relevant-files.error.opening-file', {
+                error: String(error)
+              })
+            )
+          }
+        }
+      })
+
+      quick_pick.onDidHide(() => {
+        if (!is_resolved) {
+          resolve('back')
+        }
+        quick_pick.dispose()
+      })
+
+      quick_pick.show()
     })
 
-    quick_pick.onDidAccept(async () => {
-      is_resolved = true
-      const selected_paths = quick_pick.selectedItems.map(
-        (item) => item.file_path
-      )
-      const unchecked_paths = absolute_paths.filter(
-        (file_path) => !selected_paths.includes(file_path)
-      )
+    if (list_selection === 'back' || list_selection === 'cancel') {
+      return list_selection
+    }
+
+    const selected_paths = list_selection.map((item) => item.file_path)
+    const unchecked_paths = absolute_paths.filter(
+      (file_path) => !selected_paths.includes(file_path)
+    )
+
+    let paths_to_apply: string[] = []
+    let should_continue_loop = false
+
+    const currently_checked_in_folder = currently_checked.filter((f) =>
+      files_in_searched_folder.includes(f)
+    )
 
-      const paths_to_apply = [
-        ...new Set([
-          ...currently_checked.filter((p) => !unchecked_paths.includes(p)),
-          ...selected_paths
-        ])
+    if (currently_checked_in_folder.length > 0) {
+      const selected_paths_set = new Set(selected_paths)
+      const is_identical =
+        currently_checked_in_folder.length == selected_paths_set.size &&
+        currently_checked_in_folder.every((file) =>
+          selected_paths_set.has(file)
+        )
+
+      if (is_identical) {
+        vscode.window.showInformationMessage(
+          dictionary.information_message.CONTEXT_ALREADY_SET
+        )
+        return 'success'
+      }
+
+      const quick_pick_options = [
+        {
+          label: t('command.apply-context.action.replace.label'),
+          description: t('command.apply-context.action.replace.description')
+        },
+        {
+          label: t('command.apply-context.action.merge.label'),
+          description: t('command.apply-context.action.merge.description')
+        }
       ]
-      await params.workspace_provider.set_checked_files(paths_to_apply)
-
-      const newly_selected_count = selected_paths.filter(
-        (p) => !currently_checked.includes(p)
-      ).length
-      vscode.window.showInformationMessage(
-        t('command.find-relevant-files.success.added', {
-          count: newly_selected_count
-        })
-      )
 
-      resolve('success')
-      quick_pick.hide()
-    })
+      const last_choice_label =
+        params.extension_context.workspaceState.get<string>(
+          LAST_FIND_RELEVANT_FILES_MERGE_REPLACE_OPTION_STATE_KEY
+        )
+
+      const quick_pick_merge = vscode.window.createQuickPick()
+      quick_pick_merge.items = quick_pick_options
+      quick_pick_merge.placeholder = t('command.apply-context.unstaged.apply', {
+        count: selected_paths.length
+      })
+      quick_pick_merge.buttons = [vscode.QuickInputButtons.Back]
 
-    quick_pick.onDidTriggerItemButton(async (e) => {
-      if (e.button === open_file_button) {
-        try {
-          const doc = await vscode.workspace.openTextDocument(e.item.file_path)
-          await vscode.window.showTextDocument(doc, { preview: true })
-        } catch (error) {
-          vscode.window.showErrorMessage(
-            t('command.find-relevant-files.error.opening-file', {
-              error: String(error)
-            })
-          )
+      if (last_choice_label) {
+        const active_item = quick_pick_options.find(
+          (opt) => opt.label === last_choice_label
+        )
+        if (active_item) {
+          quick_pick_merge.activeItems = [active_item]
         }
       }
-    })
 
-    quick_pick.onDidHide(() => {
-      if (!is_resolved) {
-        resolve('back')
+      const choice = await new Promise<
+        vscode.QuickPickItem | 'back' | undefined
+      >((resolve_choice) => {
+        let is_accepted = false
+        quick_pick_merge.onDidTriggerButton((button) => {
+          if (button === vscode.QuickInputButtons.Back) {
+            resolve_choice('back')
+            quick_pick_merge.hide()
+          }
+        })
+        quick_pick_merge.onDidAccept(() => {
+          is_accepted = true
+          resolve_choice(quick_pick_merge.selectedItems[0])
+          quick_pick_merge.hide()
+        })
+        quick_pick_merge.onDidHide(() => {
+          if (!is_accepted) resolve_choice('back')
+          quick_pick_merge.dispose()
+        })
+        quick_pick_merge.show()
+      })
+
+      if (choice === 'back') {
+        should_continue_loop = true
+      } else if (!choice) {
+        return 'cancel'
+      } else {
+        await params.extension_context.workspaceState.update(
+          LAST_FIND_RELEVANT_FILES_MERGE_REPLACE_OPTION_STATE_KEY,
+          choice.label
+        )
+
+        if (choice.label == t('command.apply-context.action.merge.label')) {
+          paths_to_apply = [
+            ...new Set([
+              ...currently_checked.filter((p) => !unchecked_paths.includes(p)),
+              ...selected_paths
+            ])
+          ]
+        } else {
+          paths_to_apply = [
+            ...new Set([
+              ...currently_checked.filter(
+                (p) => !files_in_searched_folder.includes(p)
+              ),
+              ...selected_paths
+            ])
+          ]
+        }
       }
-      quick_pick.dispose()
-    })
+    } else {
+      paths_to_apply = [...new Set([...currently_checked, ...selected_paths])]
+    }
+
+    if (should_continue_loop) {
+      continue
+    }
+
+    await params.workspace_provider.set_checked_files(paths_to_apply)
 
-    quick_pick.show()
-  })
+    vscode.window.showInformationMessage(
+      t('command.find-relevant-files.success.added')
+    )
+
+    return 'success'
+  }
 }

apps/editor/src/constants/state-keys.ts

@@ -10,8 +10,10 @@ export const LAST_RANGES_SAVE_LOCATION_STATE_KEY = 'last-ranges-save-location'
 export const LAST_APPLY_CONTEXT_OPTION_STATE_KEY = 'last-apply-context-option'
 export const CONTEXT_CHECKED_PATHS_STATE_KEY = 'context-checked-paths'
 export const CONTEXT_CHECKED_TIMESTAMPS_STATE_KEY = 'context-checked-timestamps'
-export const LAST_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY =
-  'last-context-merge-replace-option'
+export const LAST_APPLY_CONTEXT_MERGE_REPLACE_OPTION_STATE_KEY =
+  'last-apply-context-merge-replace-option'
+export const LAST_FIND_RELEVANT_FILES_MERGE_REPLACE_OPTION_STATE_KEY =
+  'last-find-relevant-files-merge-replace-option'
 export const LAST_REFACTOR_INSTRUCTION_SOURCE_STATE_KEY =
   'last-refactor-instruction-source'
 export const LAST_REFACTOR_INSTRUCTION_STATE_KEY = 'last-refactor-instruction'