From 82a1873c74b5b2185b9992ea1d7a5d2ae2f7c82d Mon Sep 17 00:00:00 2001
From: Zhiming Ma <codes.icy@gmail.com>
Date: Thu, 16 Jan 2025 12:39:27 +0800
Subject: [PATCH] feat(vscode): support store chat session state.

---
 clients/vscode/src/chat/chatPanel.ts    |  4 +--
 clients/vscode/src/chat/createClient.ts |  2 ++
 clients/vscode/src/chat/html/main.html  | 13 +++-----
 clients/vscode/src/chat/webview.ts      | 42 ++++++++++++++++++++-----
 4 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/clients/vscode/src/chat/chatPanel.ts b/clients/vscode/src/chat/chatPanel.ts
index 575a916da63f..ebcee50a20c1 100644
--- a/clients/vscode/src/chat/chatPanel.ts
+++ b/clients/vscode/src/chat/chatPanel.ts
@@ -1,12 +1,10 @@
 import { window, ExtensionContext, ViewColumn } from "vscode";
-import { v4 as uuid } from "uuid";
 import { ChatWebview } from "./webview";
 import type { Client } from "../lsp/Client";
 import type { GitProvider } from "../git/GitProvider";
 
 export async function createChatPanel(context: ExtensionContext, client: Client, gitProvider: GitProvider) {
-  const id = uuid();
-  const panel = window.createWebviewPanel(`tabby.chat.panel-${id}`, "Tabby", ViewColumn.One, {
+  const panel = window.createWebviewPanel(`tabby.chat.panel`, "Tabby", ViewColumn.One, {
     retainContextWhenHidden: true,
   });
 
diff --git a/clients/vscode/src/chat/createClient.ts b/clients/vscode/src/chat/createClient.ts
index 924ee0f546d0..084d4f546035 100644
--- a/clients/vscode/src/chat/createClient.ts
+++ b/clients/vscode/src/chat/createClient.ts
@@ -36,6 +36,8 @@ export function createClient(webview: Webview, api: ClientApiMethods): ServerApi
       openExternal: api.openExternal,
       readWorkspaceGitRepositories: api.readWorkspaceGitRepositories,
       getActiveEditorSelection: api.getActiveEditorSelection,
+      fetchSessionState: api.fetchSessionState,
+      storeSessionState: api.storeSessionState,
     },
   });
 }
diff --git a/clients/vscode/src/chat/html/main.html b/clients/vscode/src/chat/html/main.html
index 964329d537a1..f259cb3b7c4e 100644
--- a/clients/vscode/src/chat/html/main.html
+++ b/clients/vscode/src/chat/html/main.html
@@ -57,13 +57,6 @@ <h4 class="title">Welcome to Tabby Chat!</h4>
         chatIframe.style.display = "block";
       };
 
-      const checkFocused = () => {
-        vscode.postMessage({
-          action: "checkFocusedResult",
-          focused: document.hasFocus(),
-        });
-      };
-
       window.addEventListener("message", (event) => {
         if (event.data) {
           if (event.data.action === "dispatchKeyboardEvent") {
@@ -71,7 +64,11 @@ <h4 class="title">Welcome to Tabby Chat!</h4>
           } else if (event.data.action === "showChatPanel") {
             showChatIframe();
           } else if (event.data.action === "checkFocused") {
-            checkFocused();
+            vscode.postMessage({
+              id: event.data.id,
+              action: "jsCallback",
+              args: [document.hasFocus()],
+            });
           } else if (event.data.action === "postMessageToChatPanel") {
             chatIframe.contentWindow.postMessage(event.data.message, "*");
           } else {
diff --git a/clients/vscode/src/chat/webview.ts b/clients/vscode/src/chat/webview.ts
index 60ee353d276c..b0d36d0e3f66 100644
--- a/clients/vscode/src/chat/webview.ts
+++ b/clients/vscode/src/chat/webview.ts
@@ -28,6 +28,7 @@ import type {
   EditorFileContext,
 } from "tabby-chat-panel";
 import * as semver from "semver";
+import { v4 as uuid } from "uuid";
 import type { StatusInfo, Config } from "tabby-agent";
 import type { GitProvider } from "../git/GitProvider";
 import type { Client as LspClient } from "../lsp/Client";
@@ -70,8 +71,11 @@ export class ChatWebview {
   // A number to ensure the html is reloaded when assigned a new value
   private reloadCount = 0;
 
-  // A callback list for `isFocused` method
-  private pendingFocusCheckCallbacks: ((focused: boolean) => void)[] = [];
+  // A callback list for invoke javascript function by postMessage
+  private pendingCallbacks = new Map<string, (...arg: unknown[]) => void>();
+
+  // Store the chat state to be reload when webview is reloaded
+  private sessionState: Record<string, unknown> = {};
 
   constructor(
     private readonly context: ExtensionContext,
@@ -135,9 +139,9 @@ export class ChatWebview {
             this.client?.updateTheme(event.style, this.getColorThemeString());
             return;
           }
-          case "checkFocusedResult": {
-            this.pendingFocusCheckCallbacks.forEach((cb) => cb(event.focused));
-            this.pendingFocusCheckCallbacks = [];
+          case "jsCallback": {
+            this.pendingCallbacks.get(event.id)?.(...event.args);
+            this.pendingCallbacks.delete(event.id);
             return;
           }
         }
@@ -164,8 +168,11 @@ export class ChatWebview {
       return false;
     }
     return new Promise((resolve) => {
-      webview.postMessage({ action: "checkFocused" });
-      this.pendingFocusCheckCallbacks.push(resolve);
+      const id = uuid();
+      this.pendingCallbacks.set(id, (...args) => {
+        resolve(args[0] as boolean);
+      });
+      webview.postMessage({ id, action: "checkFocused" });
     });
   }
 
@@ -464,6 +471,27 @@ export class ChatWebview {
         const fileContext = await getFileContextFromSelection(editor, this.gitProvider);
         return fileContext;
       },
+
+      fetchSessionState: async (keys?: string[] | undefined): Promise<Record<string, unknown> | null> => {
+        if (!keys) {
+          return { ...this.sessionState };
+        }
+
+        const filtered: Record<string, unknown> = {};
+        for (const key of keys) {
+          if (key in this.sessionState) {
+            filtered[key] = this.sessionState[key];
+          }
+        }
+        return filtered;
+      },
+
+      storeSessionState: async (state: Record<string, unknown>) => {
+        this.sessionState = {
+          ...this.sessionState,
+          ...state,
+        };
+      },
     });
   }