Restore the user's original editor tab groups and active editor focus after the chat response preview process concludes

Author: robertpiosik

packages/vscode/src/commands/apply-chat-response-command/apply-chat-response-command.ts

@@ -23,13 +23,22 @@ import { ResponseHistoryItem } from '@shared/types/response-history-item'
 let in_progress = false
 let placeholder_created_at_for_update: number | undefined
 
-export const apply_chat_response_command = (
-  params: {
-    context: vscode.ExtensionContext
-    panel_provider: PanelProvider
-    workspace_provider: WorkspaceProvider
-  }
-) => {
+interface SavedEditorState {
+  uri: string
+  view_column: vscode.ViewColumn
+  is_active: boolean
+}
+
+interface SavedTabGroups {
+  editors: SavedEditorState[]
+  active_editor_uri?: string
+}
+
+export const apply_chat_response_command = (params: {
+  context: vscode.ExtensionContext
+  panel_provider: PanelProvider
+  workspace_provider: WorkspaceProvider
+}) => {
   return vscode.commands.registerCommand(
     'codeWebChat.applyChatResponse',
     async (args?: CommandArgs) => {
@@ -57,12 +66,10 @@ export const apply_chat_response_command = (
       }
 
       if (in_progress && !response_preview_promise_resolve) {
-        // Running intelligent update before response preview...
         return
       }
 
       if (response_preview_promise_resolve) {
-        // Previewing response...
         const history = params.panel_provider.response_history
 
         const new_item: ResponseHistoryItem = {
@@ -95,7 +102,28 @@ export const apply_chat_response_command = (
 
       in_progress = true
       let before_checkpoint: Checkpoint | undefined
+      let saved_tab_groups: SavedTabGroups | undefined
+
       try {
+        // Save current tab groups before entering preview
+        saved_tab_groups = {
+          editors: [],
+          active_editor_uri:
+            vscode.window.activeTextEditor?.document.uri.toString()
+        }
+
+        for (const tab_group of vscode.window.tabGroups.all) {
+          for (const tab of tab_group.tabs) {
+            if (tab.input instanceof vscode.TabInputText) {
+              saved_tab_groups.editors.push({
+                uri: tab.input.uri.toString(),
+                view_column: tab_group.viewColumn,
+                is_active: tab.isActive
+              })
+            }
+          }
+        }
+
         before_checkpoint = await create_checkpoint(
           params.workspace_provider,
           params.context,
@@ -224,7 +252,9 @@ export const apply_chat_response_command = (
             })
           }
 
-          const history_for_checkpoint = [...params.panel_provider.response_history]
+          const history_for_checkpoint = [
+            ...params.panel_provider.response_history
+          ]
           const changes_accepted = await preview_handler({
             original_states: review_data.original_states,
             chat_response: review_data.chat_response,
@@ -260,6 +290,11 @@ export const apply_chat_response_command = (
 
             before_checkpoint = undefined
           }
+
+          // Restore tab groups after preview ends
+          if (saved_tab_groups) {
+            await restore_tab_groups(saved_tab_groups)
+          }
         }
       } catch (err: any) {
         vscode.window.showErrorMessage(
@@ -278,3 +313,39 @@ export const apply_chat_response_command = (
     }
   )
 }
+
+async function restore_tab_groups(saved_state: SavedTabGroups): Promise<void> {
+  try {
+    // Close all current tabs first
+    await vscode.commands.executeCommand('workbench.action.closeAllEditors')
+
+    // Reopen saved editors in their original view columns
+    for (const editor of saved_state.editors) {
+      try {
+        const uri = vscode.Uri.parse(editor.uri)
+        await vscode.window.showTextDocument(uri, {
+          viewColumn: editor.view_column,
+          preview: false,
+          preserveFocus: !editor.is_active
+        })
+      } catch (error) {
+        // File might have been deleted or is no longer accessible
+        console.error(`Failed to restore editor for ${editor.uri}:`, error)
+      }
+    }
+
+    // Restore active editor focus if it was saved
+    if (saved_state.active_editor_uri) {
+      try {
+        const active_uri = vscode.Uri.parse(saved_state.active_editor_uri)
+        await vscode.window.showTextDocument(active_uri, {
+          preserveFocus: false
+        })
+      } catch (error) {
+        console.error('Failed to restore active editor:', error)
+      }
+    }
+  } catch (error) {
+    console.error('Failed to restore tab groups:', error)
+  }
+}

packages/vscode/src/commands/apply-chat-response-command/handlers/diff-handler.ts

@@ -476,8 +476,7 @@ export const apply_git_patch = async (
       undefined
 
     if (patch_info.is_renaming) {
-      // Skip git apply and fallbacks, renaming here is a base case
-      // where no content is changing, just paths. File is already copied.
+      // Skip git apply and fallbacks, renaming here is a base case where no content is changing, just paths. File is already copied.
       success = true
     }
 

packages/vscode/src/commands/apply-chat-response-command/handlers/fast-replace-handler.ts

@@ -87,11 +87,8 @@ export const handle_fast_replace = async (
       }
     }
 
-    // --- Application Phase ---
-    // Process all files that were deemed safe.
     const original_states: OriginalFileState[] = []
     for (const file of safe_files) {
-      // Changed from accepted_files to safe_files
       let workspace_root = default_workspace
       if (file.workspace_name && workspace_map.has(file.workspace_name)) {
         workspace_root = workspace_map.get(file.workspace_name)!
@@ -208,8 +205,8 @@ export const handle_fast_replace = async (
 
     Logger.info({
       function_name: 'handle_fast_replace',
-      message: 'Files replaced successfully', // Updated log message
-      data: { file_count: safe_files.length } // Changed from accepted_files to safe_files
+      message: 'Files replaced successfully',
+      data: { file_count: safe_files.length }
     })
     return { success: true, original_states }
   } catch (error: any) {

packages/vscode/src/commands/apply-chat-response-command/handlers/restore-preview-handler.ts

@@ -126,7 +126,6 @@ export const handle_restore_preview = async (
           }
         }
       } else if (file.content !== undefined) {
-        // Create or update
         if (file_exists) {
           try {
             const file_uri = vscode.Uri.file(safe_path)
@@ -166,7 +165,6 @@ export const handle_restore_preview = async (
             )
           }
         } else {
-          // New file
           const directory = path.dirname(safe_path)
           if (!fs.existsSync(directory)) {
             try {

packages/vscode/src/commands/apply-chat-response-command/response-processor.ts

@@ -452,8 +452,8 @@ export const process_chat_response = async (
     } else {
       let has_truncated_fragments = false
       if (
-        args?.edit_format === undefined || // Is undefined when invoked manually
-        args?.edit_format == 'truncated'
+        args?.edit_format === undefined ||
+        args?.edit_format == 'truncated' // Is undefined when invoked manually
       ) {
         has_truncated_fragments = check_for_truncated_fragments(files)
       }

packages/vscode/src/commands/apply-chat-response-command/utils/diff-processor.ts

@@ -47,11 +47,6 @@ export const process_diff = async (params: {
       message: 'File saved successfully'
     })
   } catch (error) {
-    // Logger.error({
-    //   function_name: 'process_diff',
-    //   message: 'Error saving file',
-    //   data: error
-    // })
     throw new Error('Failed to save file after applying diff patch')
   }
 }
@@ -259,17 +254,6 @@ export const apply_diff = (params: {
       if (!found) {
         search_replace_block.search_block_start_index = -2 // Not found
 
-        // Logger.error({
-        //   function_name: 'apply_diff',
-        //   message: `Search block not found: ${search_string}`
-        // })
-        // Logger.error({
-        //   function_name: 'apply_diff',
-        //   message: `search_replace_block: ${JSON.stringify(
-        //     search_replace_block
-        //   )}`
-        // })
-
         throw new Error(`Search block not found: ${search_string}`)
       }
     }
@@ -279,7 +263,7 @@ export const apply_diff = (params: {
     )
     valid_blocks.sort((a, b) => b.get_start_index() - a.get_start_index())
 
-    const result_lines = [...original_code_lines] // Operate on a copy
+    const result_lines = [...original_code_lines]
 
     for (const block of valid_blocks) {
       const start_index =
@@ -288,17 +272,9 @@ export const apply_diff = (params: {
       const replacement_content = block.replace_lines
 
       if (start_index < 0 || start_index > result_lines.length) {
-        // start_index can be == result_lines.length for appending
-        // Logger.error({
-        //   function_name: 'apply_diff',
-        //   message: `Invalid start index ${start_index} for block application. Max index: ${
-        //     result_lines.length - 1
-        //   }`
-        // })
         continue
       }
 
-      // Ensure search_count does not exceed available lines from start_index
       const actual_search_count = Math.min(
         search_count,
         result_lines.length - start_index

packages/vscode/src/commands/apply-chat-response-command/utils/file-operations.ts

@@ -204,10 +204,8 @@ const relocate_file = async (params: {
       await vscode.commands.executeCommand('workbench.action.closeActiveEditor')
     }
 
-    // Move the file
     fs.renameSync(old_safe_path, new_safe_path)
 
-    // Clean up empty directory if needed
     await remove_directory_if_empty({
       dir_path: path.dirname(old_safe_path),
       workspace_root: params.workspace_root

packages/vscode/src/commands/apply-chat-response-command/utils/preview/workspace-listener.ts

@@ -109,7 +109,6 @@ export const setup_workspace_listeners = (
           file: changed_file_in_review.reviewable_file
         })
       } else {
-        // A file not in the review has been changed. Add it to the review.
         const doc = event.document
         if (doc.uri.scheme != 'file') return
         const workspace_folder = vscode.workspace.getWorkspaceFolder(doc.uri)
@@ -616,7 +615,6 @@ export const setup_workspace_listeners = (
     }
 
     if (!is_checked) {
-      // This is the unchecking logic
       // Before reverting, capture current content if file exists
       if (fs.existsSync(file_to_toggle.sanitized_path)) {
         const document = await vscode.workspace.openTextDocument(
@@ -644,7 +642,6 @@ export const setup_workspace_listeners = (
           Buffer.from(file_to_toggle.original_content, 'utf8')
         )
       }
-      // This is the checking logic
     } else {
       if (file_to_toggle.reviewable_file.is_deleted) {
         try {