Implement workspace checkpoints to save and restore file states manually or automatically before applying AI-generated changes

Author: robertpiosik

apps/editor/package.json

@@ -285,6 +285,17 @@
         "category": "Code Web Chat",
         "icon": "$(search)"
       },
+      {
+        "command": "codeWebChat.checkpoints",
+        "title": "%codeWebChat.checkpoints.title%",
+        "category": "Code Web Chat",
+        "icon": "$(history)"
+      },
+      {
+        "command": "codeWebChat.createNewCheckpoint",
+        "title": "%codeWebChat.createNewCheckpoint.title%",
+        "category": "Code Web Chat"
+      },
       {
         "command": "codeWebChat.settings",
         "title": "%codeWebChat.settings.title%",

apps/editor/src/commands/apply-chat-response-command/apply-chat-response-command.ts

@@ -1,6 +1,11 @@
 import * as vscode from 'vscode'
 import * as fs from 'fs'
+import {
+  create_checkpoint,
+  delete_checkpoint
+} from '../checkpoints-command/actions'
 import { FileInPreview } from '@shared/types/file-in-preview'
+import { get_checkpoint_path } from '../checkpoints-command/utils'
 import { PanelProvider } from '@/views/panel/backend/panel-provider'
 import { dictionary } from '@shared/constants/dictionary'
 import { WorkspaceProvider } from '@/context/providers/workspace/workspace-provider'
@@ -15,6 +20,7 @@ import {
   process_chat_response,
   ApplyChatResponseCommandArgs
 } from './response-processor'
+import { Checkpoint } from '../checkpoints-command/types'
 import { CHECKPOINTS_STATE_KEY } from '@/constants/state-keys'
 import { ResponseHistoryItem } from '@shared/types/response-history-item'
 import { ApiManager } from '@/services/api-manager'
@@ -24,12 +30,6 @@ import {
   CwcPreviewProvider
 } from './utils/preview/virtual-document-provider'
 import { parse_response } from './utils/clipboard-parser'
-import { Checkpoint } from '@/features/checkpoints/types'
-import {
-  create_checkpoint,
-  delete_checkpoint
-} from '@/features/checkpoints/actions'
-import { get_checkpoint_path } from '@/features/checkpoints/utils'
 
 let in_progress = false
 

apps/editor/src/commands/checkpoints-command/actions/clear-all-checkpoints.ts

@@ -0,0 +1,64 @@
+import * as vscode from 'vscode'
+import {
+  CHECKPOINTS_STATE_KEY,
+  TEMPORARY_CHECKPOINT_STATE_KEY
+} from '../../../constants/state-keys'
+import { get_checkpoint_path } from '../utils'
+import { get_checkpoints } from './get-checkpoints'
+import { Logger } from '@shared/utils/logger'
+import { PanelProvider } from '@/views/panel/backend/panel-provider'
+import type { Checkpoint } from '../types'
+
+export const clear_all_checkpoints = async (
+  context: vscode.ExtensionContext,
+  panel_provider: PanelProvider
+) => {
+  const clear_task = async () => {
+    const checkpoints = await get_checkpoints(context)
+    for (const checkpoint of checkpoints) {
+      try {
+        const checkpoint_path = get_checkpoint_path(checkpoint.timestamp)
+        await vscode.workspace.fs.delete(vscode.Uri.file(checkpoint_path), {
+          recursive: true
+        })
+      } catch (error) {
+        Logger.warn({
+          function_name: 'clear_all_checkpoints',
+          message: 'Could not delete checkpoint file',
+          data: error
+        })
+      }
+    }
+    const temp_checkpoint = context.workspaceState.get<Checkpoint>(
+      TEMPORARY_CHECKPOINT_STATE_KEY
+    )
+    if (temp_checkpoint) {
+      try {
+        const checkpoint_path = get_checkpoint_path(temp_checkpoint.timestamp)
+        await vscode.workspace.fs.delete(vscode.Uri.file(checkpoint_path), {
+          recursive: true
+        })
+      } catch (error) {
+        Logger.warn({
+          function_name: 'clear_all_checkpoints',
+          message: 'Could not delete temporary checkpoint file',
+          data: error
+        })
+      }
+    }
+    await context.workspaceState.update(
+      TEMPORARY_CHECKPOINT_STATE_KEY,
+      undefined
+    )
+    await context.workspaceState.update(CHECKPOINTS_STATE_KEY, [])
+    await panel_provider.send_checkpoints()
+  }
+  await vscode.window.withProgress(
+    {
+      location: vscode.ProgressLocation.Notification,
+      title: 'Clearing all checkpoints...',
+      cancellable: false
+    },
+    clear_task
+  )
+}

apps/editor/src/commands/checkpoints-command/actions/create-checkpoint.ts

@@ -0,0 +1,267 @@
+import * as vscode from 'vscode'
+import {
+  CHECKPOINTS_STATE_KEY,
+  TEMPORARY_CHECKPOINT_STATE_KEY,
+  CHECKPOINT_OPERATION_IN_PROGRESS_STATE_KEY
+} from '../../../constants/state-keys'
+import { WorkspaceProvider } from '../../../context/providers/workspace/workspace-provider'
+import type { Checkpoint, CheckpointTab, CheckpointTrigger } from '../types'
+import { get_checkpoints } from './get-checkpoints'
+import {
+  copy_optimised_recursively,
+  get_checkpoint_path,
+  get_incremented_description
+} from '../utils'
+import {
+  get_git_diff,
+  get_git_info,
+  is_git_repository
+} from '../utils/git-utils'
+import * as path from 'path'
+import { Logger } from '@shared/utils/logger'
+import { PanelProvider } from '@/views/panel/backend/panel-provider'
+import { dictionary } from '@shared/constants/dictionary'
+import { response_preview_promise_resolve } from '../../apply-chat-response-command/utils/preview'
+import { remove_old_checkpoints } from './get-checkpoints'
+
+export const create_checkpoint = async (params: {
+  workspace_provider: WorkspaceProvider
+  context: vscode.ExtensionContext
+  panel_provider: PanelProvider
+  trigger?: CheckpointTrigger
+  description?: string
+  silent?: boolean
+}): Promise<Checkpoint | undefined> => {
+  try {
+    const trigger = params.trigger ?? 'manual'
+    const are_automatic_checkpoints_disabled = vscode.workspace
+      .getConfiguration('codeWebChat')
+      .get<boolean>('areAutomaticCheckpointsDisabled', false)
+    const is_automatic_checkpoint = trigger != 'manual'
+
+    if (are_automatic_checkpoints_disabled && is_automatic_checkpoint) {
+      Logger.info({
+        function_name: 'create_checkpoint',
+        message: 'Automatic checkpoint creation is disabled. Skipping.',
+        data: { trigger }
+      })
+      return undefined
+    }
+
+    const operation_in_progress = params.context.workspaceState.get<number>(
+      CHECKPOINT_OPERATION_IN_PROGRESS_STATE_KEY
+    )
+    if (
+      operation_in_progress &&
+      Date.now() - operation_in_progress < 60 * 1000
+    ) {
+      vscode.window.showWarningMessage(
+        dictionary.warning_message.CHECKPOINT_OPERATION_IN_PROGRESS
+      )
+      return undefined
+    }
+
+    const old_temp_checkpoint = params.context.workspaceState.get<Checkpoint>(
+      TEMPORARY_CHECKPOINT_STATE_KEY
+    )
+    if (old_temp_checkpoint) {
+      try {
+        const checkpoint_path = get_checkpoint_path(
+          old_temp_checkpoint.timestamp
+        )
+        await vscode.workspace.fs.delete(vscode.Uri.file(checkpoint_path), {
+          recursive: true
+        })
+      } catch (error) {
+        Logger.warn({
+          function_name: 'create_checkpoint',
+          message: 'Could not delete old temporary checkpoint file',
+          data: error
+        })
+      }
+      await params.context.workspaceState.update(
+        TEMPORARY_CHECKPOINT_STATE_KEY,
+        undefined
+      )
+    }
+    const timestamp = Date.now()
+    const workspace_folders = vscode.workspace.workspaceFolders!
+
+    const folder_git_statuses = await Promise.all(
+      workspace_folders.map((folder) => is_git_repository(folder))
+    )
+    let new_checkpoint: Checkpoint | undefined
+
+    const create_checkpoint_task = async () => {
+      await params.context.workspaceState.update(
+        CHECKPOINT_OPERATION_IN_PROGRESS_STATE_KEY,
+        Date.now()
+      )
+      await vscode.workspace.saveAll()
+      const checkpoint_dir_path = get_checkpoint_path(timestamp)
+      const checkpoint_dir_uri = vscode.Uri.file(checkpoint_dir_path)
+      await vscode.workspace.fs.createDirectory(checkpoint_dir_uri)
+
+      const git_data: Record<string, any> = {}
+      let uses_git = false
+
+      for (const [index, folder] of workspace_folders.entries()) {
+        const has_git = folder_git_statuses[index]
+
+        if (has_git) {
+          uses_git = true
+          const git_info = await get_git_info(folder)
+          const diff = await get_git_diff(folder)
+
+          if (git_info && diff !== null) {
+            const folder_name = folder.name
+            git_data[folder_name] = {
+              branch: git_info.branch,
+              commit_hash: git_info.commit_hash,
+              folder_name
+            }
+
+            const diff_file_path = path.join(
+              checkpoint_dir_path,
+              `${folder_name}.diff`
+            )
+            await vscode.workspace.fs.writeFile(
+              vscode.Uri.file(diff_file_path),
+              Buffer.from(diff, 'utf8')
+            )
+          } else {
+            Logger.warn({
+              function_name: 'create_checkpoint',
+              message: `Failed to get git information for repository in ${folder.name}. Falling back to file copy.`
+            })
+            const dest_folder_path =
+              workspace_folders.length > 1
+                ? path.join(checkpoint_dir_path, folder.name)
+                : checkpoint_dir_path
+            const dest_folder_uri = vscode.Uri.file(dest_folder_path)
+            if (workspace_folders.length > 1) {
+              await vscode.workspace.fs.createDirectory(dest_folder_uri)
+            }
+
+            const entries = await vscode.workspace.fs.readDirectory(folder.uri)
+            for (const [name] of entries) {
+              await copy_optimised_recursively(
+                vscode.Uri.joinPath(folder.uri, name),
+                vscode.Uri.joinPath(dest_folder_uri, name),
+                folder.uri.fsPath,
+                params.workspace_provider
+              )
+            }
+          }
+        } else {
+          const dest_folder_path =
+            workspace_folders.length > 1
+              ? path.join(checkpoint_dir_path, folder.name)
+              : checkpoint_dir_path
+          const dest_folder_uri = vscode.Uri.file(dest_folder_path)
+          if (workspace_folders.length > 1) {
+            await vscode.workspace.fs.createDirectory(dest_folder_uri)
+          }
+
+          const entries = await vscode.workspace.fs.readDirectory(folder.uri)
+          for (const [name] of entries) {
+            await copy_optimised_recursively(
+              vscode.Uri.joinPath(folder.uri, name),
+              vscode.Uri.joinPath(dest_folder_uri, name),
+              folder.uri.fsPath,
+              params.workspace_provider
+            )
+          }
+        }
+      }
+
+      let checkpoints = await get_checkpoints(params.context)
+
+      // Remove checkpoints older than 48 hours
+      checkpoints = await remove_old_checkpoints(checkpoints)
+
+      const final_description = get_incremented_description(
+        trigger,
+        params.description,
+        checkpoints
+      )
+
+      const active_tabs: CheckpointTab[] = []
+      const tab_groups = vscode.window.tabGroups.all
+      for (const group of tab_groups) {
+        for (const tab of group.tabs) {
+          if (tab.input instanceof vscode.TabInputText) {
+            active_tabs.push({
+              uri: (tab.input as vscode.TabInputText).uri.toString(),
+              view_column: group.viewColumn,
+              is_active: tab.isActive,
+              is_group_active: group.isActive
+            })
+          }
+        }
+      }
+
+      const checkpoint_object: Checkpoint = {
+        trigger,
+        timestamp,
+        description: final_description,
+        uses_git,
+        git_data: Object.keys(git_data).length > 0 ? git_data : undefined,
+        checked_files: params.workspace_provider.get_all_checked_paths(),
+        active_tabs
+      }
+
+      if (trigger == 'manual' && !response_preview_promise_resolve) {
+        checkpoint_object.response_history = [
+          ...params.panel_provider.response_history
+        ]
+      }
+
+      checkpoints.unshift(checkpoint_object)
+      await params.context.workspaceState.update(
+        CHECKPOINTS_STATE_KEY,
+        checkpoints
+      )
+      await params.panel_provider.send_checkpoints()
+      await params.context.workspaceState.update(
+        CHECKPOINT_OPERATION_IN_PROGRESS_STATE_KEY,
+        undefined
+      )
+      new_checkpoint = checkpoint_object
+    }
+
+    let timer: ReturnType<typeof setTimeout> | undefined
+    let did_show_modal = false
+
+    try {
+      if (!params.silent) {
+        timer = setTimeout(() => {
+          did_show_modal = true
+          params.panel_provider.send_message({
+            command: 'SHOW_PROGRESS',
+            title: 'Creating checkpoint...'
+          })
+        }, 1000)
+      }
+
+      await create_checkpoint_task()
+    } finally {
+      if (timer) clearTimeout(timer)
+      if (did_show_modal) {
+        params.panel_provider.send_message({
+          command: 'HIDE_PROGRESS'
+        })
+      }
+    }
+    return new_checkpoint
+  } catch (err: any) {
+    vscode.window.showErrorMessage(
+      `Failed to create checkpoint: ${err.message}`
+    )
+    await params.context.workspaceState.update(
+      CHECKPOINT_OPERATION_IN_PROGRESS_STATE_KEY,
+      undefined
+    )
+    return undefined
+  }
+}