Add active editor state restoration after applying chat response changes

Author: robertpiosik

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

@@ -113,6 +113,37 @@ export const preview_handler = async (params: {
         original_editor_state: params.original_editor_state
       })
 
+      if (preview_result.active_editor_state) {
+        try {
+          const uri = vscode.Uri.file(
+            preview_result.active_editor_state.file_path
+          )
+          const file_exists = await vscode.workspace.fs.stat(uri).then(
+            () => true,
+            () => false
+          )
+
+          if (file_exists) {
+            const document = await vscode.workspace.openTextDocument(uri)
+            const editor = await vscode.window.showTextDocument(document, {
+              preview: false
+            })
+            const position = preview_result.active_editor_state.position
+            editor.selection = new vscode.Selection(position, position)
+            editor.revealRange(
+              new vscode.Range(position, position),
+              vscode.TextEditorRevealType.InCenter
+            )
+          }
+        } catch (error) {
+          Logger.error({
+            function_name: 'preview_handler',
+            message: 'Failed to restore active editor state',
+            data: error
+          })
+        }
+      }
+
       return true
     } else {
       update_undo_button_state({

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

@@ -27,6 +27,10 @@ export const preview = async (params: {
   accepted_files: ReviewableFile[]
   rejected_states: OriginalFileState[]
   created_at?: number
+  active_editor_state?: {
+    file_path: string
+    position: vscode.Position
+  }
 } | null> => {
   if (!vscode.workspace.workspaceFolders?.length) {
     vscode.window.showErrorMessage(
@@ -206,10 +210,29 @@ export const preview = async (params: {
           })
           .filter((state): state is OriginalFileState => state !== undefined)
 
+        let active_editor_state:
+          | { file_path: string; position: vscode.Position }
+          | undefined
+
+        if (result.active_file_path && result.active_position) {
+          const active_path = result.active_file_path
+          const matching_prepared = prepared_files.find(
+            (pf) => pf.temp_file_path === active_path
+          )
+
+          active_editor_state = {
+            file_path: matching_prepared
+              ? matching_prepared.sanitized_path
+              : active_path,
+            position: result.active_position
+          }
+        }
+
         return {
           accepted_files,
           rejected_states,
-          created_at: decision.created_at
+          created_at: decision.created_at,
+          active_editor_state
         }
       }
 

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

@@ -1,4 +1,5 @@
 import { FileInPreview } from '@shared/types/file-in-preview'
+import * as vscode from 'vscode'
 
 export type CodeReviewDecision =
   | { jump_to: { file_path: string; workspace_name?: string } }
@@ -8,6 +9,8 @@ export type CodeReviewResult = {
   decision: CodeReviewDecision
   new_content: string
   temp_file_path: string
+  active_file_path?: string
+  active_position?: vscode.Position
 }
 
 export type ReviewableFile = FileInPreview & {

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

@@ -54,10 +54,21 @@ export const show_diff_with_actions = async (
         final_content = right_doc.getText()
       } catch (error) {}
 
+      const active_editor = vscode.window.activeTextEditor
+      let active_file_path: string | undefined
+      let active_position: vscode.Position | undefined
+
+      if (active_editor) {
+        active_file_path = active_editor.document.uri.fsPath
+        active_position = active_editor.selection.active
+      }
+
       resolve({
         decision,
         new_content: final_content,
-        temp_file_path: prepared_file.temp_file_path
+        temp_file_path: prepared_file.temp_file_path,
+        active_file_path,
+        active_position
       })
     }
   }).finally(() => {