Implement Prune Context mode with dedicated prompt instructions, history management, tool configurations, and compact token count display

Author: robertpiosik

packages/vscode/src/context/providers/open-editors/open-editors-provider.ts

@@ -29,6 +29,14 @@ export class OpenEditorsProvider
   private _shared_state: SharedFileState
   private _config_change_handler: vscode.Disposable
   private _workspace_provider: WorkspaceProvider
+  private _use_compact_token_count: boolean = false
+
+  public set_use_compact_token_count(use_compact: boolean): void {
+    if (this._use_compact_token_count != use_compact) {
+      this._use_compact_token_count = use_compact
+      this.refresh()
+    }
+  }
 
   constructor(
     workspace_folders: vscode.WorkspaceFolder[],
@@ -184,7 +192,9 @@ export class OpenEditorsProvider
 
     element.checkboxState = checkbox_state
 
-    const token_count = element.tokenCount
+    const token_count = this._use_compact_token_count
+      ? element.compactTokenCount
+      : element.tokenCount
 
     let final_description = element.description || ''
 

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

@@ -66,6 +66,15 @@ export class WorkspaceProvider
   public gitignore_initialization: Promise<void>
   public ranges_initialization: Promise<void>
 
+  private _use_compact_token_count: boolean = false
+
+  public set_use_compact_token_count(use_compact: boolean): void {
+    if (this._use_compact_token_count != use_compact) {
+      this._use_compact_token_count = use_compact
+      this.refresh()
+    }
+  }
+
   constructor(
     workspace_folders: vscode.WorkspaceFolder[],
     private _context: vscode.ExtensionContext
@@ -523,8 +532,12 @@ export class WorkspaceProvider
 
     element.checkboxState = checkbox_state
 
-    const total_token_count = element.tokenCount
-    const selected_token_count = element.selectedTokenCount
+    const total_token_count = this._use_compact_token_count
+      ? element.compactTokenCount
+      : element.tokenCount
+    const selected_token_count = this._use_compact_token_count
+      ? element.selectedCompactTokenCount
+      : element.selectedTokenCount
 
     const formatted_total =
       total_token_count !== undefined && total_token_count > 0
@@ -700,11 +713,12 @@ export class WorkspaceProvider
             vscode.TreeItemCheckboxState.Unchecked,
           false, // Is not symbolic link
           false, // Is not open file
-          total_tokens,
-          selected_tokens,
+          total_tokens.total,
+          total_tokens.compact,
+          selected_tokens.total,
+          selected_tokens.compact,
           undefined,
-          true, // Is workspace root
-          undefined
+          true // Is workspace root
         )
       )
       if (context_view) {

packages/vscode/src/services/model-providers-manager.ts

@@ -45,6 +45,7 @@ export type CodeCompletionsConfigs = ToolConfig[]
 export type EditContextConfigs = ToolConfig[]
 export type IntelligentUpdateConfigs = ToolConfig[]
 export type CommitMessagesConfigs = ToolConfig[]
+export type PruneContextConfigs = ToolConfig[]
 
 export class ModelProvidersManager {
   private _providers: Provider[] = []
@@ -421,6 +422,18 @@ export class ModelProvidersManager {
     )
   }
 
+  public async get_prune_context_tool_configs(): Promise<PruneContextConfigs> {
+    await this._load_promise
+    return this._get_tool_configs_from_settings('configurationsForPruneContext')
+  }
+
+  public async save_prune_context_tool_configs(configs: PruneContextConfigs) {
+    await this._save_tool_configs_to_settings(
+      'configurationsForPruneContext',
+      configs
+    )
+  }
+
   public async update_provider_name_in_configs(params: {
     old_name: string
     new_name: string
@@ -432,7 +445,8 @@ export class ModelProvidersManager {
       'configurationsForCodeCompletions',
       'configurationsForEditContext',
       'configurationsForIntelligentUpdate',
-      'configurationsForCommitMessages'
+      'configurationsForCommitMessages',
+      'configurationsForPruneContext'
     ]
 
     for (const key of settings_keys) {

packages/vscode/src/views/panel/backend/message-handlers/handle-get-api-tool-configurations.ts

@@ -15,11 +15,13 @@ export const handle_get_api_tool_configurations = async (
   const [
     edit_context_configs,
     code_completions_configs,
-    default_code_completions_config
+    default_code_completions_config,
+    prune_context_configs
   ] = await Promise.all([
     providers_manager.get_edit_context_tool_configs(),
     providers_manager.get_code_completions_tool_configs(),
-    providers_manager.get_default_code_completions_config()
+    providers_manager.get_default_code_completions_config(),
+    providers_manager.get_prune_context_tool_configs()
   ])
 
   const is_config_default = (
@@ -46,6 +48,10 @@ export const handle_get_api_tool_configurations = async (
       ...config,
       id: get_tool_config_id(config),
       is_default: is_config_default(config, default_code_completions_config)
+    })),
+    'prune-context': prune_context_configs.map((config) => ({
+      ...config,
+      id: get_tool_config_id(config)
     }))
   }
 

packages/vscode/src/views/panel/backend/message-handlers/handle-get-history.ts

@@ -3,6 +3,7 @@ import {
   HISTORY_CODE_COMPLETIONS_STATE_KEY,
   HISTORY_EDIT_STATE_KEY,
   HISTORY_NO_CONTEXT_STATE_KEY,
+  HISTORY_PRUNE_CONTEXT_STATE_KEY,
   HistoryEntry
 } from '@/constants/state-keys'
 import { PanelProvider } from '@/views/panel/backend/panel-provider'
@@ -21,12 +22,16 @@ export const handle_get_history = (panel_provider: PanelProvider): void => {
   const no_context_history = panel_provider.context.workspaceState.get<
     HistoryEntry[]
   >(HISTORY_NO_CONTEXT_STATE_KEY, [])
+  const prune_context_history = panel_provider.context.workspaceState.get<
+    HistoryEntry[]
+  >(HISTORY_PRUNE_CONTEXT_STATE_KEY, [])
 
   panel_provider.send_message({
     command: 'CHAT_HISTORY',
     ask: ask_history.map((h) => h.text),
     edit_context: edit_history.map((h) => h.text),
     no_context: no_context_history.map((h) => h.text),
-    code_completions: code_completions_history.map((h) => h.text)
+    code_completions: code_completions_history.map((h) => h.text),
+    prune_context: prune_context_history.map((h) => h.text)
   })
 }

packages/vscode/src/views/panel/backend/message-handlers/handle-get-instructions.ts

@@ -8,6 +8,7 @@ export const handle_get_instructions = (
     ask: panel_provider.ask_instructions,
     edit_context: panel_provider.edit_instructions,
     no_context: panel_provider.no_context_instructions,
-    code_completions: panel_provider.code_completion_instructions
+    code_completions: panel_provider.code_completion_instructions,
+    prune_context: panel_provider.prune_context_instructions
   })
 }

packages/vscode/src/views/panel/backend/message-handlers/handle-reorder-api-tool-configurations.ts

@@ -49,5 +49,22 @@ export const handle_reorder_api_tool_configurations = async (
         reordered_configs
       )
     }
+  } else if (message.mode == 'prune-context') {
+    const current_configs =
+      await providers_manager.get_prune_context_tool_configs()
+    const reordered_configs = reordered_ids
+      .map((id) => {
+        const found = current_configs.find((p) => get_tool_config_id(p) === id)
+        if (!found) {
+          console.error(`Config with id ${id} not found during reorder.`)
+          return null
+        }
+        return found
+      })
+      .filter((p): p is ToolConfig => p !== null)
+
+    if (reordered_configs.length === current_configs.length) {
+      await providers_manager.save_prune_context_tool_configs(reordered_configs)
+    }
   }
 }

packages/vscode/src/views/panel/backend/message-handlers/handle-save-history.ts

@@ -7,7 +7,9 @@ import {
   PINNED_HISTORY_ASK_STATE_KEY,
   PINNED_HISTORY_CODE_COMPLETIONS_STATE_KEY,
   PINNED_HISTORY_EDIT_STATE_KEY,
-  PINNED_HISTORY_NO_CONTEXT_STATE_KEY
+  PINNED_HISTORY_NO_CONTEXT_STATE_KEY,
+  HISTORY_PRUNE_CONTEXT_STATE_KEY,
+  PINNED_HISTORY_PRUNE_CONTEXT_STATE_KEY
 } from '@/constants/state-keys'
 import { PanelProvider } from '@/views/panel/backend/panel-provider'
 import { SaveHistoryMessage } from '@/views/panel/types/messages'
@@ -35,6 +37,10 @@ export const handle_save_history = async (
       key = HISTORY_CODE_COMPLETIONS_STATE_KEY
       pinned_key = PINNED_HISTORY_CODE_COMPLETIONS_STATE_KEY
       break
+    case 'prune-context':
+      key = HISTORY_PRUNE_CONTEXT_STATE_KEY
+      pinned_key = PINNED_HISTORY_PRUNE_CONTEXT_STATE_KEY
+      break
   }
   if (key) {
     const text_history: string[] = message.messages

packages/vscode/src/views/panel/backend/message-handlers/handle-show-history-quick-pick.ts

@@ -10,6 +10,8 @@ import {
   PINNED_HISTORY_CODE_COMPLETIONS_STATE_KEY,
   PINNED_HISTORY_EDIT_STATE_KEY,
   PINNED_HISTORY_NO_CONTEXT_STATE_KEY,
+  HISTORY_PRUNE_CONTEXT_STATE_KEY,
+  PINNED_HISTORY_PRUNE_CONTEXT_STATE_KEY,
   HistoryEntry
 } from '@/constants/state-keys'
 import { PanelProvider } from '@/views/panel/backend/panel-provider'
@@ -50,6 +52,10 @@ export const handle_show_history_quick_pick = async (
       history_key = HISTORY_NO_CONTEXT_STATE_KEY
       pinned_history_key = PINNED_HISTORY_NO_CONTEXT_STATE_KEY
       break
+    case 'prune-context':
+      history_key = HISTORY_PRUNE_CONTEXT_STATE_KEY
+      pinned_history_key = PINNED_HISTORY_PRUNE_CONTEXT_STATE_KEY
+      break
   }
 
   if (!history_key || !pinned_history_key) return
@@ -169,6 +175,9 @@ export const handle_show_history_quick_pick = async (
       case 'code-completions':
         panel_provider.code_completion_instructions = text
         break
+      case 'prune-context':
+        panel_provider.prune_context_instructions = text
+        break
       default:
         return
     }
@@ -178,7 +187,8 @@ export const handle_show_history_quick_pick = async (
       ask: panel_provider.ask_instructions,
       edit_context: panel_provider.edit_instructions,
       no_context: panel_provider.no_context_instructions,
-      code_completions: panel_provider.code_completion_instructions
+      code_completions: panel_provider.code_completion_instructions,
+      prune_context: panel_provider.prune_context_instructions
     })
   }
 

packages/vscode/src/views/panel/backend/message-handlers/handle-toggle-pinned-api-tool-configuration.ts

@@ -17,6 +17,8 @@ export const handle_toggle_pinned_api_tool_configuration = async (
     configs = await providers_manager.get_edit_context_tool_configs()
   } else if (mode == 'code-completions') {
     configs = await providers_manager.get_code_completions_tool_configs()
+  } else if (mode == 'prune-context') {
+    configs = await providers_manager.get_prune_context_tool_configs()
   } else {
     return
   }
@@ -30,6 +32,8 @@ export const handle_toggle_pinned_api_tool_configuration = async (
     await providers_manager.save_edit_context_tool_configs(configs)
   } else if (mode == 'code-completions') {
     await providers_manager.save_code_completions_tool_configs(configs)
+  } else if (mode == 'prune-context') {
+    await providers_manager.save_prune_context_tool_configs(configs)
   }
   await handle_get_api_tool_configurations(panel_provider)
 }