feat(chat): init steps for supporting chat history and small refactoring

package.json

@@ -36,7 +36,8 @@
       "view/title": [
         {
           "command": "firecoder.startNewChat",
-          "group": "navigation@1"
+          "group": "navigation",
+          "when": "view === firecoder.chat-gui"
         }
       ]
     },

src/common/panel/chat.ts

@@ -3,7 +3,8 @@ import * as vscode from "vscode";
 import { getUri } from "../utils/getUri";
 import { getNonce } from "../utils/getNonce";
 import { chat } from "../chat";
-import { ChatMessage } from "../prompt/promptChat";
+import { Chat, ChatMessage } from "../prompt/promptChat";
+import { state } from "../utils/state";
 
 export type MessageType =
   | {
@@ -15,11 +16,43 @@ export type MessageType =
   | {
       type: "e2w-response";
       id: string;
-      command: string;
       done: boolean;
       data: any;
     };
 
+type MessageToExtention =
+  | {
+      type: "send-message";
+      data: ChatMessage[];
+    }
+  | {
+      type: "abort-generate";
+      id: string;
+    }
+  | {
+      type: "get-chat";
+      chatId: string;
+    }
+  | {
+      type: "save-chat";
+      chatId: string;
+      data: Chat;
+    }
+  | {
+      type: "get-chats";
+    }
+  | {
+      type: "delete-chat";
+      chatId: string;
+    }
+  | {
+      type: "delete-chats";
+    };
+
+type MessageFromWebview = MessageToExtention & {
+  id: string;
+};
+
 export class ChatPanel implements vscode.WebviewViewProvider {
   private disposables: Disposable[] = [];
   private webview: Webview | undefined;
@@ -50,6 +83,7 @@ export class ChatPanel implements vscode.WebviewViewProvider {
       "css",
       "main.css",
     ]);
+
     const scriptUri = getUri(webview, extensionUri, [
       "webviews",
       "build",
@@ -95,49 +129,70 @@ export class ChatPanel implements vscode.WebviewViewProvider {
 
   private setWebviewMessageListener(webview: Webview) {
     webview.onDidReceiveMessage(
-      async (message: any) => {
+      async (message: MessageFromWebview) => {
         if (message.type in this.messageCallback) {
           this.messageCallback[message.type]();
           return;
         }
+
         const type = message.type;
 
         switch (type) {
-          case "sendMessage":
+          case "send-message":
             await this.handleStartGeneration({
+              id: message.id,
               chatMessage: message.data,
-              messageId: message.messageId,
-              messageType: message.type,
             });
-            return;
+            break;
+          case "get-chat":
+            await this.handleGetChat({
+              id: message.id,
+              chatId: message.chatId,
+            });
+            break;
+          case "save-chat":
+            await this.handleSaveChat({
+              id: message.id,
+              chatId: message.chatId,
+              history: message.data,
+            });
+            break;
+          case "delete-chat":
+            await this.handleDeleteChat({
+              id: message.id,
+              chatId: message.chatId,
+            });
+            break;
+          case "delete-chats":
+            await this.handleDeleteChats({
+              id: message.id,
+            });
+            break;
+          case "get-chats":
+            await this.handleGetChats({
+              id: message.id,
+            });
+            break;
+          default:
+            break;
         }
       },
       undefined,
       this.disposables
     );
   }
 
-  private addMessageListener(
-    commandOrMessageId: string,
-    callback: (message: any) => void
-  ) {
-    this.messageCallback[commandOrMessageId] = callback;
-  }
-
   private async handleStartGeneration({
-    messageId,
-    messageType,
+    id,
     chatMessage,
   }: {
-    messageId: string;
-    messageType: string;
+    id: string;
     chatMessage: ChatMessage[];
   }) {
-    const sendResponse = (messageToResponse: any, done: boolean) => {
+    const sendResponse = (messageToResponse: string, done: boolean) => {
       this.postMessage({
         type: "e2w-response",
-        id: messageId,
-        command: messageType,
+        id: id,
         data: messageToResponse,
         done: done,
       });
@@ -158,8 +213,91 @@ export class ChatPanel implements vscode.WebviewViewProvider {
     sendResponse("", true);
   }
 
+  private async handleGetChat({ chatId, id }: { chatId: string; id: string }) {
+    const sendResponse = (messageToResponse: Chat | null, done: boolean) => {
+      this.postMessage({
+        type: "e2w-response",
+        id: id,
+        data: messageToResponse,
+        done: done,
+      });
+    };
+
+    const history = state.global.get(`chat-${chatId}`);
+    if (history) {
+      sendResponse(history, true);
+    } else {
+      sendResponse(null, true);
+    }
+  }
+
+  private async handleSaveChat({
+    chatId,
+    history,
+    id,
+  }: {
+    chatId: string;
+    history: Chat;
+    id: string;
+  }) {
+    await state.global.update(`chat-${chatId}`, history);
+    await this.postMessage({
+      type: "e2w-response",
+      id: id,
+      data: "",
+      done: true,
+    });
+  }
+
+  private async handleDeleteChat({
+    chatId,
+    id,
+  }: {
+    chatId: string;
+    id: string;
+  }) {
+    await state.global.delete(`chat-${chatId}`);
+    await this.postMessage({
+      type: "e2w-response",
+      id: id,
+      data: "",
+      done: true,
+    });
+  }
+
+  private async handleDeleteChats({ id }: { id: string }) {
+    await state.global.deleteChats();
+    await this.postMessage({
+      type: "e2w-response",
+      id: id,
+      data: "",
+      done: true,
+    });
+  }
+
+  private async handleGetChats({ id }: { id: string }) {
+    const chats = state.global.getChats();
+    await this.postMessage({
+      type: "e2w-response",
+      id: id,
+      data: chats,
+      done: true,
+    });
+  }
+
+  private addMessageListener(
+    commandOrMessageId: string,
+    callback: (message: any) => void
+  ) {
+    this.messageCallback[commandOrMessageId] = callback;
+  }
+
+  private async postMessage(message: MessageType) {
+    await this.webview?.postMessage(message);
+  }
+
   public async sendMessageToWebview(
-    command: MessageType["command"],
+    command: string,
     data: MessageType["data"]
   ) {
     const message: MessageType = {
@@ -170,8 +308,4 @@ export class ChatPanel implements vscode.WebviewViewProvider {
     };
     await this.postMessage(message);
   }
-
-  private async postMessage(message: MessageType) {
-    await this.webview?.postMessage(message);
-  }
 }

src/common/prompt/promptChat.ts

@@ -1,6 +1,14 @@
 export type ChatMessage = {
   role: string;
   content: string;
+  // chatMessageId: string;
+};
+
+export type Chat = {
+  messages: ChatMessage[];
+  chatId: string;
+  date: number;
+  title: string;
 };
 
 const promptBaseDefault = `You are an AI programming assistant, utilizing the DeepSeek Coder model, developed by DeepSeek Company, and you only answer questions related to computer science. For politically sensitive questions, security and privacy issues, and other non-computer science questions, you will refuse to answer.

src/common/utils/state.ts

@@ -1,23 +1,17 @@
 import * as vscode from "vscode";
 import type { Spec } from "../download";
+import { Chat } from "../prompt/promptChat";
 
 const StateValues = {
-  inlineSuggestModeAuto: {
-    default: true,
-  },
-  serverSpec: {
-    default: null,
-  },
+  inlineSuggestModeAuto: true,
+  serverSpec: null,
+};
+type StateValuesType = {
+  inlineSuggestModeAuto: boolean;
+  serverSpec: Spec | null;
+  [key: `chat-${string}`]: Chat | undefined;
 };
 
-interface StateValuesType extends Record<keyof typeof StateValues, any> {
-  inlineSuggestModeAuto: {
-    possibleValues: boolean;
-  };
-  serverSpec: {
-    possibleValues: Spec | null;
-  };
-}
 class State {
   state?: vscode.Memento;
   constructor() {}
@@ -26,18 +20,44 @@ class State {
     this.state = state;
   }
 
-  public get<T extends keyof StateValuesType>(
+  public get<T extends keyof StateValuesType & string>(
     key: T
-  ): StateValuesType[T]["possibleValues"] {
-    return this.state?.get(key) ?? StateValues[key]["default"];
+  ): StateValuesType[T] {
+    // @ts-ignore
+    return this.state?.get(key) ?? StateValues[key];
   }
 
   public async update<T extends keyof StateValuesType>(
     key: T,
-    value: StateValuesType[T]["possibleValues"]
+    value: StateValuesType[T]
   ) {
     await this.state?.update(key, value);
   }
+
+  public getChats(): Chat[] {
+    const allKeys = (this.state?.keys() ||
+      []) as unknown as (keyof StateValuesType)[];
+
+    return allKeys
+      .filter((key) => key.startsWith("chat-"))
+      .map((key) => {
+        return this.get(key as `chat-${string}`) as Chat;
+      });
+  }
+  public async delete<T extends keyof StateValuesType>(key: T) {
+    await this.state?.update(key, undefined);
+  }
+
+  public async deleteChats() {
+    const allKeys = (this.state?.keys() ||
+      []) as unknown as (keyof StateValuesType)[];
+
+    await Promise.all(
+      allKeys
+        .filter((key) => key.startsWith("chat-"))
+        .map((key) => this.delete(key as `chat-${string}`))
+    );
+  }
 }
 
 export const state = {

src/extension.ts

@@ -29,7 +29,7 @@ export async function activate(context: vscode.ExtensionContext) {
   );
   context.subscriptions.push(
     vscode.commands.registerCommand("firecoder.startNewChat", async () => {
-      await provider.sendMessageToWebview("startNewChat", {});
+      await provider.sendMessageToWebview("start-new-chat", {});
     })
   );