Add panel mode state management and optimize token count calculation by moving it to context initialization

Author: robertpiosik

packages/vscode/src/constants/state-keys.ts

@@ -44,6 +44,7 @@ export const CHAT_EDIT_FORMAT_STATE_KEY = 'chat-edit-format'
 export const API_EDIT_FORMAT_STATE_KEY = 'api-edit-format'
 export const WEB_MODE_STATE_KEY = 'web-mode'
 export const API_MODE_STATE_KEY = 'api-mode'
+export const PANEL_MODE_STATE_KEY = 'panel-mode'
 
 export const INSTRUCTIONS_EDIT_CONTEXT_STATE_KEY = 'instructions-edit-context'
 export const INSTRUCTIONS_ASK_STATE_KEY = 'instructions-ask'

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

@@ -16,7 +16,6 @@ import {
   CONTEXT_CHECKED_URLS_STATE_KEY,
   CONTEXT_CHECKED_TIMESTAMPS_STATE_KEY,
   DUPLICATE_WORKSPACE_CONTEXT_STATE_KEY,
-  SAVED_CONTEXTS_STATE_KEY,
   RANGES_STATE_KEY,
   type DuplicateWorkspaceContext
 } from '../constants/state-keys'
@@ -82,12 +81,6 @@ const restore_duplicated_workspace_context = async (
             duplicated_context.ranges
           )
         }
-        if (duplicated_context.saved_contexts) {
-          await context.workspaceState.update(
-            SAVED_CONTEXTS_STATE_KEY,
-            duplicated_context.saved_contexts
-          )
-        }
         if (duplicated_context.open_editors) {
           for (const editor_info of duplicated_context.open_editors) {
             try {
@@ -150,8 +143,12 @@ export const context_initialization = async (
   const update_view_badges = async () => {
     let context_token_count = 0
     if (context_provider && context_view) {
-      context_token_count =
+      const files_count =
         await workspace_provider.get_checked_files_token_count()
+      const websites_count =
+        websites_provider.get_checked_websites_token_count()
+      context_token_count = files_count + websites_count
+
       workspace_view.badge = {
         value: round_token_count_for_badge(context_token_count),
         tooltip: context_token_count
@@ -160,7 +157,7 @@ export const context_initialization = async (
       }
       context_view.badge = undefined
     }
-    token_count_emitter.emit('token-count-updated')
+    token_count_emitter.emit('token-count-updated', context_token_count)
   }
 
   const shared_state = SharedFileState.get_instance()

…e-handlers/handle-save-home-view-type.ts …end/message-handlers/handle-save-mode.tspackages/vscode/src/views/panel/backend/message-handlers/handle-save-home-view-type.ts renamed to packages/vscode/src/views/panel/backend/message-handlers/handle-save-mode.ts packages/vscode/src/views/panel/backend/message-handlers/handle-save-home-view-type.ts renamed to packages/vscode/src/views/panel/backend/message-handlers/handle-save-mode.ts

@@ -1,13 +1,18 @@
 import { PanelProvider } from '@/views/panel/backend/panel-provider'
 import { SaveModeMessage } from '@/views/panel/types/messages'
+import { PANEL_MODE_STATE_KEY } from '@/constants/state-keys'
 
 export const handle_save_mode = async (
   panel_provider: PanelProvider,
   message: SaveModeMessage
 ): Promise<void> => {
   panel_provider.mode = message.mode
-  panel_provider.send_message({
-    command: 'MODE',
-    mode: message.mode
-  })
+  await panel_provider.context.workspaceState.update(
+    PANEL_MODE_STATE_KEY,
+    message.mode
+  )
+  await panel_provider.context.globalState.update(
+    PANEL_MODE_STATE_KEY,
+    message.mode
+  )
 }

packages/vscode/src/views/panel/backend/message-handlers/index.ts

@@ -21,7 +21,7 @@ export * from './handle-show-prompt-template-quick-pick'
 export * from './handle-update-preset'
 export * from './handle-pick-chatbot'
 export * from './handle-pick-open-router-model'
-export * from './handle-save-home-view-type'
+export * from './handle-save-mode'
 export * from './handle-get-mode'
 export * from './handle-edit-context'
 export * from './handle-code-completion'

packages/vscode/src/views/panel/backend/panel-provider.ts

@@ -43,6 +43,7 @@ import {
   handle_at_sign_quick_pick,
   handle_get_mode_web,
   handle_save_mode_web,
+  handle_save_mode,
   handle_get_mode_api,
   handle_save_mode_api,
   handle_get_mode,
@@ -89,6 +90,7 @@ import {
   INSTRUCTIONS_NO_CONTEXT_STATE_KEY,
   LAST_SELECTED_CODE_COMPLETION_CONFIG_ID_STATE_KEY,
   LAST_SELECTED_EDIT_CONTEXT_CONFIG_ID_STATE_KEY,
+  PANEL_MODE_STATE_KEY,
   WEB_MODE_STATE_KEY,
   get_recently_used_presets_or_groups_key
 } from '@/constants/state-keys'
@@ -210,6 +212,10 @@ export class PanelProvider implements vscode.WebviewViewProvider {
       this.context.workspaceState.get<EditFormat>(API_EDIT_FORMAT_STATE_KEY) ??
       this.context.globalState.get<EditFormat>(API_EDIT_FORMAT_STATE_KEY) ??
       'whole'
+    this.mode =
+      this.context.workspaceState.get<Mode>(PANEL_MODE_STATE_KEY) ??
+      this.context.globalState.get<Mode>(PANEL_MODE_STATE_KEY) ??
+      MODE.WEB
 
     this.web_prompt_type = this.context.workspaceState.get<WebPromptType>(
       WEB_MODE_STATE_KEY,
@@ -282,9 +288,13 @@ export class PanelProvider implements vscode.WebviewViewProvider {
       }
     )
 
-    token_count_emitter.on('token-count-updated', () => {
+    token_count_emitter.on('token-count-updated', (token_count: number) => {
       if (this._webview_view) {
-        this.calculate_token_count()
+        this.send_message({
+          command: 'TOKEN_COUNT_UPDATED',
+          token_count
+        })
+
         this.send_context_files()
       }
     })
@@ -347,13 +357,6 @@ export class PanelProvider implements vscode.WebviewViewProvider {
         vscode.window.activeTextEditor &&
         event.document === vscode.window.activeTextEditor.document
       ) {
-        if (
-          (this.mode == MODE.WEB &&
-            this.web_prompt_type == 'code-completions') ||
-          (this.mode == MODE.API && this.api_prompt_type == 'code-completions')
-        ) {
-          this.calculate_token_count()
-        }
         this.send_currently_open_file_text()
       }
     })
@@ -507,8 +510,6 @@ export class PanelProvider implements vscode.WebviewViewProvider {
             handle_request_editor_state(this)
           } else if (message.command == 'REQUEST_EDITOR_SELECTION_STATE') {
             handle_request_editor_selection_state(this)
-          } else if (message.command == 'GET_CURRENT_TOKEN_COUNT') {
-            this.calculate_token_count()
           } else if (message.command == 'GET_CONTEXT_SIZE_WARNING_THRESHOLD') {
             this._send_context_size_warning_threshold()
           } else if (message.command == 'REQUEST_CURRENTLY_OPEN_FILE_TEXT') {
@@ -579,8 +580,7 @@ export class PanelProvider implements vscode.WebviewViewProvider {
           } else if (message.command == 'CARET_POSITION_CHANGED') {
             this.caret_position = message.caret_position
           } else if (message.command == 'SAVE_MODE') {
-            this.mode = message.mode
-            this.calculate_token_count()
+            await handle_save_mode(this, message)
           } else if (message.command == 'GET_MODE') {
             handle_get_mode(this)
           } else if (message.command == 'GET_VERSION') {
@@ -671,78 +671,6 @@ export class PanelProvider implements vscode.WebviewViewProvider {
     )
   }
 
-  public calculate_token_count() {
-    if (this.mode == MODE.WEB && this.web_prompt_type == 'no-context') {
-      this.send_message({
-        command: 'TOKEN_COUNT_UPDATED',
-        token_count: 0
-      })
-      return
-    }
-
-    const active_editor = vscode.window.activeTextEditor
-
-    const is_code_completions_mode =
-      (this.mode == MODE.WEB && this.web_prompt_type == 'code-completions') ||
-      (this.mode == MODE.API && this.api_prompt_type == 'code-completions')
-
-    Promise.all([
-      this.workspace_provider.get_checked_files_token_count({
-        exclude_file_path:
-          is_code_completions_mode && active_editor
-            ? active_editor.document.uri.fsPath
-            : undefined
-      }),
-      this.websites_provider.get_checked_websites_token_count()
-    ])
-      .then(([workspace_tokens, websites_tokens]) => {
-        let current_token_count = workspace_tokens + websites_tokens
-
-        if (active_editor && is_code_completions_mode) {
-          const document = active_editor.document
-          const text = document.getText()
-          const file_path = document.uri.fsPath
-          const workspace_root =
-            this.workspace_provider.get_workspace_root_for_file(file_path)
-          let content_xml = ''
-
-          if (!workspace_root) {
-            content_xml = `<file path="${file_path}">\n<![CDATA[\n${text}\n]]>\n</file>\n`
-          } else {
-            const relative_path = path.relative(workspace_root, file_path)
-            if (this.workspace_provider.get_workspace_roots().length > 1) {
-              const workspace_name =
-                this.workspace_provider.get_workspace_name(workspace_root)
-              content_xml = `<file path="${workspace_name}/${relative_path}">\n<![CDATA[\n${text}\n]]>\n</file>\n`
-            } else {
-              content_xml = `<file path="${relative_path}">\n<![CDATA[\n${text}\n]]>\n</file>\n`
-            }
-          }
-          const file_token_count = Math.floor(content_xml.length / 4)
-          current_token_count += file_token_count
-        }
-
-        this.send_message({
-          command: 'TOKEN_COUNT_UPDATED',
-          token_count: current_token_count
-        })
-      })
-      .catch((error) => {
-        Logger.error({
-          function_name: 'calculate_token_count',
-          message: 'Error calculating token count',
-          data: error
-        })
-        vscode.window.showErrorMessage(
-          dictionary.error_message.ERROR_CALCULATING_TOKEN_COUNT(error.message)
-        )
-        this.send_message({
-          command: 'TOKEN_COUNT_UPDATED',
-          token_count: 0
-        })
-      })
-  }
-
   public send_presets_to_webview(_: vscode.Webview) {
     const config = vscode.workspace.getConfiguration('codeWebChat')
     const web_modes: WebPromptType[] = [

packages/vscode/src/views/panel/frontend/Main/Main.tsx

@@ -160,7 +160,6 @@ export const Main: React.FC<Props> = (props) => {
     const initial_messages: FrontendMessage[] = [
       { command: 'GET_PRESETS' },
       { command: 'GET_HISTORY' },
-      { command: 'GET_CURRENT_TOKEN_COUNT' },
       { command: 'GET_INSTRUCTIONS' },
       { command: 'GET_EDIT_FORMAT' },
       { command: 'GET_EDIT_FORMAT_INSTRUCTIONS' },

packages/vscode/src/views/panel/frontend/hooks/use-panel.ts

@@ -433,9 +433,6 @@ export const use_panel = (vscode: any) => {
       command: 'SAVE_WEB_MODE',
       mode: new_mode
     })
-    post_message(vscode, {
-      command: 'GET_CURRENT_TOKEN_COUNT'
-    })
   }
 
   const handle_api_prompt_type_change = (new_mode: ApiPromptType) => {
@@ -446,9 +443,6 @@ export const use_panel = (vscode: any) => {
       command: 'SAVE_API_MODE',
       mode: new_mode
     })
-    post_message(vscode, {
-      command: 'GET_CURRENT_TOKEN_COUNT'
-    })
   }
 
   const handle_discard_user_changes_in_preview = (file: {

packages/vscode/src/views/panel/types/messages.ts

@@ -124,10 +124,6 @@ export interface SaveHistoryMessage extends BaseMessage {
   mode: 'ask' | 'edit-context' | 'no-context' | 'code-completions'
 }
 
-export interface GetCurrentTokenCountMessage extends BaseMessage {
-  command: 'GET_CURRENT_TOKEN_COUNT'
-}
-
 export interface GetContextSizeWarningThresholdMessage extends BaseMessage {
   command: 'GET_CONTEXT_SIZE_WARNING_THRESHOLD'
 }
@@ -469,7 +465,6 @@ export type FrontendMessage =
   | GetHistoryMessage
   | GetResponseHistoryMessage
   | SaveHistoryMessage
-  | GetCurrentTokenCountMessage
   | GetContextSizeWarningThresholdMessage
   | UpdatePresetMessage
   | DeletePresetGroupOrSeparatorMessage