Implement file rename tracking and handling within the code review process, and introduce a dedicated handler to restore files from review data

Author: robertpiosik

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

@@ -35,6 +35,7 @@ import {
   promote_temporary_checkpoint
 } from '../checkpoints-command/actions'
 import { FileInReview } from '@shared/types/file-in-review'
+import { handle_restore_review } from './handlers/restore-review-handler'
 
 let ongoing_review_cleanup_promise: Promise<void> | null = null
 
@@ -436,6 +437,22 @@ export const apply_chat_response_command = (
       }
 
       const review_data = await (async (): Promise<ReviewData | null> => {
+        if (args?.files_with_content) {
+          const result = await handle_restore_review(args.files_with_content)
+          if (result.success && result.original_states) {
+            update_undo_button_state(
+              result.original_states,
+              chat_response,
+              args?.original_editor_state
+            )
+            return {
+              original_states: result.original_states,
+              chat_response
+            }
+          }
+          return null
+        }
+
         const is_single_root_folder_workspace =
           vscode.workspace.workspaceFolders?.length == 1
 
@@ -444,13 +461,6 @@ export const apply_chat_response_command = (
           is_single_root_folder_workspace
         })
 
-        if (args?.files_with_content) {
-          clipboard_content = {
-            type: 'files',
-            files: args.files_with_content as ClipboardFile[]
-          }
-        }
-
         if (
           clipboard_content.type == 'code-completion' &&
           clipboard_content.code_completion

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

@@ -0,0 +1,217 @@
+import * as vscode from 'vscode'
+import * as path from 'path'
+import * as fs from 'fs'
+import { Logger } from '@shared/utils/logger'
+import { dictionary } from '@shared/constants/dictionary'
+import { create_safe_path } from '@/utils/path-sanitizer'
+import { OriginalFileState } from '../types/original-file-state'
+import { FileInReview } from '@shared/types/file-in-review'
+import { remove_directory_if_empty } from '../utils/file-operations'
+
+export const handle_restore_review = async (
+  files: FileInReview[]
+): Promise<{ success: boolean; original_states?: OriginalFileState[] }> => {
+  Logger.info({
+    function_name: 'handle_restore_review',
+    message: 'start',
+    data: { file_count: files.length }
+  })
+
+  if (
+    !vscode.workspace.workspaceFolders ||
+    vscode.workspace.workspaceFolders.length == 0
+  ) {
+    vscode.window.showErrorMessage(
+      dictionary.error_message.NO_WORKSPACE_FOLDER_OPEN
+    )
+    Logger.warn({
+      function_name: 'handle_restore_review',
+      message: 'No workspace folder open.'
+    })
+    return { success: false }
+  }
+
+  const workspace_map = new Map<string, string>()
+  vscode.workspace.workspaceFolders.forEach((folder) => {
+    workspace_map.set(folder.name, folder.uri.fsPath)
+  })
+  const default_workspace = vscode.workspace.workspaceFolders[0].uri.fsPath
+
+  const original_states: OriginalFileState[] = []
+
+  for (const file of files) {
+    let workspace_root = default_workspace
+    if (file.workspace_name && workspace_map.has(file.workspace_name)) {
+      workspace_root = workspace_map.get(file.workspace_name)!
+    } else if (file.workspace_name) {
+      Logger.warn({
+        function_name: 'handle_restore_review',
+        message: `Workspace '${file.workspace_name}' not found for file '${file.file_path}'. Using default.`
+      })
+    }
+
+    const safe_path = create_safe_path(workspace_root, file.file_path)
+
+    if (!safe_path) {
+      Logger.warn({
+        function_name: 'handle_restore_review',
+        message: 'Unsafe file path detected and skipped',
+        data: file.file_path
+      })
+      continue
+    }
+
+    const file_exists = fs.existsSync(safe_path)
+    let original_content = ''
+    if (file_exists) {
+      try {
+        original_content = fs.readFileSync(safe_path, 'utf8')
+      } catch (e) {
+        Logger.warn({
+          function_name: 'handle_restore_review',
+          message: 'Could not read original content of file',
+          data: { file_path: safe_path, error: e }
+        })
+      }
+    }
+
+    original_states.push({
+      file_path: file.file_path,
+      content: original_content,
+      is_new: !file_exists,
+      workspace_name: file.workspace_name,
+      is_deleted: file.is_deleted
+    })
+
+    if (file.is_deleted) {
+      if (file_exists) {
+        try {
+          const tabs_to_close: vscode.Tab[] = []
+          for (const tab_group of vscode.window.tabGroups.all) {
+            tabs_to_close.push(
+              ...tab_group.tabs.filter((tab) => {
+                const tab_uri = (tab.input as any)?.uri as
+                  | vscode.Uri
+                  | undefined
+                return tab_uri && tab_uri.fsPath === safe_path
+              })
+            )
+          }
+
+          if (tabs_to_close.length > 0) {
+            await vscode.window.tabGroups.close(tabs_to_close)
+          }
+
+          fs.unlinkSync(safe_path)
+          Logger.info({
+            function_name: 'handle_restore_review',
+            message: 'File deleted',
+            data: safe_path
+          })
+          await remove_directory_if_empty({
+            dir_path: path.dirname(safe_path),
+            workspace_root
+          })
+        } catch (error) {
+          Logger.error({
+            function_name: 'handle_restore_review',
+            message: 'Error deleting file',
+            data: { safe_path, error }
+          })
+          vscode.window.showErrorMessage(
+            `Failed to delete file: ${file.file_path}`
+          )
+        }
+      }
+    } else if (file.content !== undefined) {
+      // Create or update
+      if (file_exists) {
+        try {
+          const file_uri = vscode.Uri.file(safe_path)
+          const document = await vscode.workspace.openTextDocument(file_uri)
+          const editor = await vscode.window.showTextDocument(document)
+          let final_content = file.content
+          if (
+            original_content.endsWith('\n') &&
+            !final_content.endsWith('\n')
+          ) {
+            final_content += '\n'
+          }
+          await editor.edit((edit) => {
+            edit.replace(
+              new vscode.Range(
+                document.positionAt(0),
+                document.positionAt(document.getText().length)
+              ),
+              final_content
+            )
+          })
+
+          await document.save()
+          Logger.info({
+            function_name: 'handle_restore_review',
+            message: 'Existing file replaced and saved',
+            data: safe_path
+          })
+        } catch (error) {
+          Logger.error({
+            function_name: 'handle_restore_review',
+            message: 'Error updating file',
+            data: { safe_path, error }
+          })
+          vscode.window.showErrorMessage(
+            `Failed to update file: ${file.file_path}`
+          )
+        }
+      } else {
+        // New file
+        const directory = path.dirname(safe_path)
+        if (!fs.existsSync(directory)) {
+          try {
+            fs.mkdirSync(directory, { recursive: true })
+            Logger.info({
+              function_name: 'handle_restore_review',
+              message: 'Directory created',
+              data: directory
+            })
+          } catch (error) {
+            Logger.error({
+              function_name: 'handle_restore_review',
+              message: 'Failed to create directory',
+              data: { directory, error, file_path: file.file_path }
+            })
+            vscode.window.showErrorMessage(
+              dictionary.error_message.FAILED_TO_CREATE_DIRECTORY(
+                file.file_path
+              )
+            )
+            continue
+          }
+        }
+
+        try {
+          fs.writeFileSync(safe_path, file.content)
+          Logger.info({
+            function_name: 'handle_restore_review',
+            message: 'New file created',
+            data: safe_path
+          })
+          const document = await vscode.workspace.openTextDocument(safe_path)
+          await vscode.window.showTextDocument(document, { preview: false })
+        } catch (error) {
+          Logger.error({
+            function_name: 'handle_restore_review',
+            message: 'Failed to write new file',
+            data: { safe_path, error, file_path: file.file_path }
+          })
+          vscode.window.showErrorMessage(
+            dictionary.error_message.FAILED_TO_WRITE_FILE(file.file_path)
+          )
+          continue
+        }
+      }
+    }
+  }
+
+  return { success: true, original_states }
+}

packages/vscode/src/commands/apply-chat-response-command/types/original-file-state.ts

@@ -1,12 +1,13 @@
 export type OriginalFileState = {
   file_path: string
   content: string
-  is_new: boolean
+  is_new?: boolean
   workspace_name?: string
   cursor_offset?: number
   new_file_path?: string
   file_path_to_restore?: string
   is_fallback?: boolean
   diff_fallback_method?: 'recount' | 'search_and_replace'
   is_replaced?: boolean
+  is_deleted?: boolean
 }

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

@@ -326,7 +326,7 @@ export const undo_files = async (params: {
         continue // Skip this file
       }
 
-      if (state.is_new) {
+      if (state.is_new && !state.file_path_to_restore) {
         if (fs.existsSync(safe_path)) {
           // Close any open editors for this file before deleting
           const uri = vscode.Uri.file(safe_path)