Add undo/redo functionality for context selection and reorder prompt types

Author: robertpiosik

README.md

@@ -32,7 +32,7 @@ Code Web Chat is a free and open-source AI coding toolkit. You can use CWC in **
 
 ## Introduction
 
-With its focus on context engineering, CWC proposes a non-agentic workflow, where all the necessary for the task files, includnig useful examples, are picked with engineer's intent. Such guided model is easy to prompt and as accurate as it can get. It's a perfect companion for all kinds of coding tasks; from mass renamings to features.
+With its focus on context engineering, CWC proposes a non-agentic workflow, where all the necessary for the task files, includnig useful examples, are picked with engineer's intent. Such guided model is easy to prompt and as accurate as it can get. The tool a perfect companion for all kinds of coding tasks; from mass renamings to complex features.
 
 > [!TIP]
 > Feeling lost in the codebase? Pick context roughly, switch prompt type to _Prune context_, include your task and AI will help narrowing it down to reasonable size.
@@ -118,7 +118,8 @@ Meaningful summaries of changes adhering to your style.
 - `Code Web Chat: Remove File from Context` - Search and remove file (or parent folder via file action) from the context.
 - `Code Web Chat: Copy Context` - Copy XML-formatted checked files from the Workspace view to the clipboard.
 - `Code Web Chat: Copy Context of Open Editors` - Copy XML-formatted checked files from the Open Editors view to the clipboard.
-- `Code Web Chat: Find Paths in Clipboard` - Select files based on paths found in the clipboard text.
+- `Code Web Chat: Undo Context Selection` - Undo the last context change.
+- `Code Web Chat: Redo Context Selection` - Redo the last undone context change.
 
 ## Contributing
 

packages/shared/src/types/prompt-types.ts

@@ -5,4 +5,4 @@ export type WebPromptType =
   | 'prune-context'
   | 'no-context'
 
-export type ApiPromptType = 'edit-context' | 'code-at-cursor' | 'prune-context'
+export type ApiPromptType = 'edit-context' | 'prune-context' | 'code-at-cursor'

packages/vscode/package.json

@@ -106,6 +106,16 @@
         "category": "Code Web Chat",
         "icon": "$(copy)"
       },
+      {
+        "command": "codeWebChat.undoContextSelection",
+        "title": "Undo Context Selection",
+        "category": "Code Web Chat"
+      },
+      {
+        "command": "codeWebChat.redoContextSelection",
+        "title": "Redo Context Selection",
+        "category": "Code Web Chat"
+      },
       {
         "command": "codeWebChat.duplicateWorkspace",
         "title": "Duplicate Workspace in New Window",
@@ -535,22 +545,22 @@
         {
           "command": "codeWebChat.checkParentFolder",
           "when": "view == codeWebChatViewOpenEditors && viewItem == openEditor",
-          "group": "1_modification@3"
+          "group": "2_modification@1"
         },
         {
           "command": "codeWebChat.referenceInPrompt",
           "when": "view == codeWebChatViewOpenEditors && viewItem == openEditor",
-          "group": "1_modification@4"
+          "group": "2_modification@2"
         },
         {
           "command": "codeWebChat.revealInExplorer",
           "when": "view == codeWebChatViewOpenEditors && viewItem == openEditor",
-          "group": "1_modification@5"
+          "group": "2_modification@3"
         },
         {
           "command": "codeWebChat.setRange",
           "when": "view == codeWebChatViewOpenEditors && viewItem == openEditor",
-          "group": "1_modification@6"
+          "group": "2_modification@4"
         },
         {
           "command": "codeWebChat.openInIntegratedTerminal",

packages/vscode/src/context/context-initialization.ts

@@ -164,7 +164,11 @@ export const context_initialization = async (
   }
 
   const shared_state = SharedFileState.get_instance()
-  shared_state.set_providers(workspace_provider, open_editors_provider)
+  shared_state.set_providers(
+    workspace_provider,
+    open_editors_provider,
+    websites_provider
+  )
 
   workspace_provider.load_checked_files_state()
   websites_provider.load_checked_websites_state()

packages/vscode/src/context/providers/websites/websites-provider.ts

@@ -152,6 +152,22 @@ export class WebsitesProvider
     }
   }
 
+  set_checked_websites(urls: string[]): void {
+    this._checked_websites.clear()
+
+    // Only check websites that are currently in the list
+    const valid_urls = new Set(this._websites.map((w) => w.url))
+
+    for (const url of urls) {
+      if (valid_urls.has(url)) {
+        this._checked_websites.set(url, vscode.TreeItemCheckboxState.Checked)
+      }
+    }
+
+    this._onDidChangeCheckedWebsites.fire()
+    this._onDidChangeTreeData.fire()
+  }
+
   clear_checks(): void {
     if (this._checked_websites.size == 0) {
       return

packages/vscode/src/context/shared-file-state.ts

@@ -6,21 +6,32 @@ import {
   FileItem
 } from './providers/workspace/workspace-provider'
 import { OpenEditorsProvider } from './providers/open-editors/open-editors-provider'
+import { WebsitesProvider } from './providers/websites/websites-provider'
+
+type ContextState = {
+  files: string[]
+  websites: string[]
+}
 
 export class SharedFileState {
   private static _instance: SharedFileState
   private _on_did_change_checked_files = new vscode.EventEmitter<void>()
   readonly onDidChangeCheckedFiles = this._on_did_change_checked_files.event
 
   private _workspace_provider?: WorkspaceProvider
+  private _websites_provider?: WebsitesProvider
   private _open_editors_provider?: OpenEditorsProvider
   private _checked_files: Set<string> = new Set()
+  private _checked_websites: Set<string> = new Set()
   // Track which view last unchecked a file
   private _unchecked_in_open_editors: Set<string> = new Set()
   private _unchecked_in_workspace: Set<string> = new Set()
   private _synchronizing_provider: 'workspace' | 'openEditors' | null = null
   private _is_synchronizing: boolean = false
   private _is_initialized: boolean = false
+  private _is_undoing_redoing: boolean = false
+  private _undo_stack: ContextState[] = []
+  private _redo_stack: ContextState[] = []
 
   static get_instance(): SharedFileState {
     if (!SharedFileState._instance) {
@@ -31,40 +42,136 @@ export class SharedFileState {
 
   set_providers(
     workspace_provider: WorkspaceProvider,
-    open_editors_provider: OpenEditorsProvider
+    open_editors_provider: OpenEditorsProvider,
+    websites_provider?: WebsitesProvider
   ) {
     this._workspace_provider = workspace_provider
     this._open_editors_provider = open_editors_provider
+    this._websites_provider = websites_provider
 
     workspace_provider.onDidChangeCheckedFiles(() => {
-      if (!this._is_synchronizing) {
+      if (!this._is_synchronizing && !this._is_undoing_redoing) {
         this._synchronizing_provider = 'workspace'
         this.synchronize_state()
         this._synchronizing_provider = null
       }
     })
 
     open_editors_provider.onDidChangeCheckedFiles(() => {
-      if (!this._is_synchronizing) {
+      if (!this._is_synchronizing && !this._is_undoing_redoing) {
         this._synchronizing_provider = 'openEditors'
         this.synchronize_state()
         this._synchronizing_provider = null
       }
     })
 
+    if (websites_provider) {
+      websites_provider.onDidChangeCheckedWebsites(() => {
+        if (!this._is_undoing_redoing) {
+          this.handle_websites_change()
+        }
+      })
+    }
+
     // Initialize with a synchronization after a small delay to ensure both providers are ready
     setTimeout(() => {
       if (!this._is_initialized) {
         this.synchronize_state()
+        if (this._websites_provider) {
+          const websites = this._websites_provider
+            .get_checked_websites()
+            .map((w) => w.url)
+          this._checked_websites = new Set(websites)
+        }
         this._is_initialized = true
       }
     }, 1000)
   }
 
+  private push_state(files: string[], websites: string[]) {
+    if (!this._is_initialized) return
+    this._undo_stack.push({ files, websites })
+    if (this._undo_stack.length > 50) this._undo_stack.shift()
+    this._redo_stack = []
+  }
+
+  private handle_websites_change() {
+    if (this._is_undoing_redoing) return
+
+    const current_websites =
+      this._websites_provider
+        ?.get_checked_websites()
+        .map((w) => w.url)
+        .sort() || []
+    const prev_websites = Array.from(this._checked_websites).sort()
+
+    if (JSON.stringify(prev_websites) !== JSON.stringify(current_websites)) {
+      this.push_state(
+        Array.from(this._checked_files).sort(),
+        Array.from(this._checked_websites).sort()
+      )
+      this._checked_websites = new Set(current_websites)
+    }
+  }
+
+  async undo() {
+    if (this._undo_stack.length === 0) return
+
+    const current_state = {
+      files: Array.from(this._checked_files).sort(),
+      websites: Array.from(this._checked_websites).sort()
+    }
+    this._redo_stack.push(current_state)
+
+    const previous_state = this._undo_stack.pop()
+    if (previous_state) {
+      this._is_undoing_redoing = true
+      await this.apply_state(previous_state)
+      this._is_undoing_redoing = false
+    }
+  }
+
+  async redo() {
+    if (this._redo_stack.length === 0) return
+
+    const current_state = {
+      files: Array.from(this._checked_files).sort(),
+      websites: Array.from(this._checked_websites).sort()
+    }
+    this._undo_stack.push(current_state)
+
+    const next_state = this._redo_stack.pop()
+    if (next_state) {
+      this._is_undoing_redoing = true
+      await this.apply_state(next_state)
+      this._is_undoing_redoing = false
+    }
+  }
+
+  async apply_state(state: ContextState) {
+    if (this._workspace_provider) {
+      await this._workspace_provider.set_checked_files(state.files)
+    }
+    if (this._open_editors_provider) {
+      await this._open_editors_provider.set_checked_files(state.files)
+    }
+    if (this._websites_provider) {
+      this._websites_provider.set_checked_websites(state.websites)
+    }
+
+    this.update_checked_files_set()
+    this._checked_websites = new Set(state.websites)
+
+    this._on_did_change_checked_files.fire()
+  }
+
   async synchronize_state() {
     if (!this._workspace_provider || !this._open_editors_provider) return
     if (this._is_synchronizing) return
 
+    const prev_files = Array.from(this._checked_files).sort()
+    const prev_websites = Array.from(this._checked_websites).sort()
+
     this._is_synchronizing = true
 
     try {
@@ -144,6 +251,14 @@ export class SharedFileState {
       this.update_checked_files_set()
 
       this._on_did_change_checked_files.fire()
+
+      const curr_files = Array.from(this._checked_files).sort()
+
+      if (!this._is_undoing_redoing && this._is_initialized) {
+        if (JSON.stringify(prev_files) !== JSON.stringify(curr_files)) {
+          this.push_state(prev_files, prev_websites)
+        }
+      }
     } finally {
       this._is_synchronizing = false
     }

packages/vscode/src/extension.ts

@@ -10,6 +10,7 @@ import {
   migrate_saved_contexts_to_global_storage,
   migrate_token_cache_cleanup
 } from './migrations'
+import { SharedFileState } from './context/shared-file-state'
 import {
   apply_chat_response_command,
   add_file_to_context_command,
@@ -170,6 +171,14 @@ export async function activate(context: vscode.ExtensionContext) {
         }
       }
     ),
+    vscode.commands.registerCommand(
+      'codeWebChat.undoContextSelection',
+      async () => SharedFileState.get_instance().undo()
+    ),
+    vscode.commands.registerCommand(
+      'codeWebChat.redoContextSelection',
+      async () => SharedFileState.get_instance().redo()
+    ),
     vscode.commands.registerCommand(
       'codeWebChat.openInIntegratedTerminal',
       async (item: any) => {

packages/vscode/src/views/panel/frontend/Main/MainView/prompt-type-labels.ts

@@ -13,8 +13,8 @@ export const WEB_PROMPT_TYPES = Object.keys(
 
 export const api_prompt_type_labels: Record<ApiPromptType, string> = {
   'edit-context': 'Edit context',
-  'code-at-cursor': 'Code at cursor',
-  'prune-context': 'Prune context'
+  'prune-context': 'Prune context',
+  'code-at-cursor': 'Code at cursor'
 }
 export const API_PROMPT_TYPES = Object.keys(
   api_prompt_type_labels