fix(agent&vscode): fix data store for vscode web extension (#3327)

* fix(agent&vscode): fix the dataStore in web extension.

* fix(agent): register once lisenter instead of on.

clients/tabby-agent/src/config/index.ts

@@ -82,7 +82,7 @@ export class Configurations extends EventEmitter implements Feature {
 
   private clientCapabilities: ClientCapabilities | undefined = undefined;
 
-  constructor(private readonly dataStore?: DataStore) {
+  constructor(private readonly dataStore: DataStore) {
     super();
   }
 
@@ -115,24 +115,20 @@ export class Configurations extends EventEmitter implements Feature {
       const configFile = this.configFile;
       await configFile.load();
       configFile.on("updated", async () => {
-        if (this.dataStore) {
-          this.serverProvided = this.pickStoredServerProvidedConfig(this.dataStore.data);
-        }
+        this.serverProvided = this.pickStoredServerProvidedConfig(this.dataStore.data);
         this.update();
       });
       configFile.watch();
     }
 
-    if (this.dataStore) {
-      this.serverProvided = this.pickStoredServerProvidedConfig(this.dataStore.data);
-      this.dataStore.on("updated", async (data: Partial<StoredData>) => {
-        const serverProvidedConfig = this.pickStoredServerProvidedConfig(data);
-        if (!deepEqual(serverProvidedConfig, this.serverProvided)) {
-          this.serverProvided = serverProvidedConfig;
-          this.update();
-        }
-      });
-    }
+    this.serverProvided = this.pickStoredServerProvidedConfig(this.dataStore.data);
+    this.dataStore.on("updated", async (data: Partial<StoredData>) => {
+      const serverProvidedConfig = this.pickStoredServerProvidedConfig(data);
+      if (!deepEqual(serverProvidedConfig, this.serverProvided)) {
+        this.serverProvided = serverProvidedConfig;
+        this.update();
+      }
+    });
 
     this.update();
   }
@@ -185,9 +181,7 @@ export class Configurations extends EventEmitter implements Feature {
       const old = this.clientProvided;
       this.clientProvided = config;
       this.emit("clientProvidedConfigUpdated", config, old);
-      if (this.dataStore) {
-        this.serverProvided = this.pickStoredServerProvidedConfig(this.dataStore.data);
-      }
+      this.serverProvided = this.pickStoredServerProvidedConfig(this.dataStore.data);
       this.update();
     }
   }
@@ -197,7 +191,7 @@ export class Configurations extends EventEmitter implements Feature {
       this.serverProvided = config;
       this.update();
     }
-    if (save && this.dataStore) {
+    if (save) {
       const mergedLocalConfig = mergeConfig(this.defaultConfig, this.configFile, this.clientProvided);
       const serverEndpoint = mergedLocalConfig.server.endpoint;
       if (!this.dataStore.data.serverConfig) {

clients/tabby-agent/src/dataStore/dataFile.ts

@@ -38,6 +38,10 @@ export class FileDataStore extends EventEmitter {
     this.watcher.on("add", onUpdated);
     this.watcher.on("change", onUpdated);
   }
+
+  stopWatch() {
+    this.watcher?.close();
+  }
 }
 
 export function getFileDataStore(): FileDataStore | undefined {

clients/tabby-agent/src/dataStore/index.ts

@@ -1,16 +1,16 @@
 import type { Connection } from "vscode-languageserver";
 import type {
   ClientCapabilities,
+  ClientProvidedConfig,
+  DataStoreRecords,
   ServerCapabilities,
   StatusIssuesName,
-  DataStoreGetParams,
-  DataStoreSetParams,
 } from "../protocol";
 import type { TabbyServerProvidedConfig } from "../http/tabbyApiClient";
 import type { Feature } from "../feature";
 import type { FileDataStore } from "./dataFile";
 import { EventEmitter } from "events";
-import { DataStoreGetRequest, DataStoreSetRequest } from "../protocol";
+import { DataStoreDidUpdateNotification, DataStoreUpdateRequest } from "../protocol";
 import { getFileDataStore } from "./dataFile";
 import deepEqual from "deep-equal";
 
@@ -24,6 +24,8 @@ export class DataStore extends EventEmitter implements Feature {
   public data: Partial<StoredData> = {};
 
   private lspConnection: Connection | undefined = undefined;
+  private lspInitialized = false;
+
   private fileDataStore: FileDataStore | undefined = undefined;
 
   async preInitialize(): Promise<void> {
@@ -37,48 +39,67 @@ export class DataStore extends EventEmitter implements Feature {
         if (!deepEqual(data, this.data)) {
           const old = this.data;
           this.data = data;
-          this.emit("updated", data, old);
+          this.emit("updated", this.data, old);
         }
       });
       dataStore.watch();
     }
   }
 
-  async initialize(connection: Connection, clientCapabilities: ClientCapabilities): Promise<ServerCapabilities> {
+  async initialize(
+    connection: Connection,
+    clientCapabilities: ClientCapabilities,
+    _clientProvidedConfig: ClientProvidedConfig,
+    dataStoreRecords: DataStoreRecords | undefined,
+  ): Promise<ServerCapabilities> {
     if (clientCapabilities.tabby?.dataStore) {
       this.lspConnection = connection;
 
-      try {
-        // FIXME(@icycodes): This try-catch block avoids the initialization error in current version.
-        // As the lsp client has not be initialized, the request will always throws errors,
-        // the dataStore data should be initialized by initializationOptions, and be synced by notifications.
-        const params: DataStoreGetParams = { key: "data" };
-        const data = await connection.sendRequest(DataStoreGetRequest.type, params);
-        if (!deepEqual(data, this.data)) {
+      // When dataStore is provided by the LSP connection, do not use the file data store anymore.
+      const dataStore = this.fileDataStore;
+      if (dataStore) {
+        dataStore.stopWatch();
+        this.fileDataStore = undefined;
+        const old = this.data;
+        this.data = dataStoreRecords ?? {};
+        this.emit("updated", this.data, old);
+      } else {
+        this.data = dataStoreRecords ?? {};
+      }
+
+      connection.onNotification(DataStoreDidUpdateNotification.type, async (params) => {
+        const records = params ?? {};
+        if (!deepEqual(records, this.data)) {
           const old = this.data;
-          this.data = data;
-          this.emit("updated", data, old);
+          this.data = records;
+          this.emit("updated", this.data, old);
         }
-      } catch (error) {
-        // ignore
-      }
+      });
     }
     return {};
   }
 
-  async save() {
+  async initialized() {
     if (this.lspConnection) {
-      const params: DataStoreGetParams = { key: "data" };
-      const old = await this.lspConnection.sendRequest(DataStoreGetRequest.type, params);
-
-      if (!deepEqual(old, this.data)) {
-        const params: DataStoreSetParams = { key: "data", value: this.data };
-        await this.lspConnection.sendRequest(DataStoreSetRequest.type, params);
-        this.emit("updated", this.data, old);
-      }
+      this.lspInitialized = true;
+      this.emit("initialized");
     }
+  }
 
-    if (this.fileDataStore) {
+  async save() {
+    if (this.lspConnection) {
+      const connection = this.lspConnection;
+      const sendUpdateRequest = async () => {
+        await connection.sendRequest(DataStoreUpdateRequest.type, this.data);
+      };
+      if (this.lspInitialized) {
+        await sendUpdateRequest();
+      } else {
+        this.once("initialized", async () => {
+          await sendUpdateRequest();
+        });
+      }
+    } else if (this.fileDataStore) {
       await this.fileDataStore.write(this.data);
     }
   }

clients/tabby-agent/src/feature.ts

@@ -1,11 +1,12 @@
 import type { Connection } from "vscode-languageserver";
-import type { ClientCapabilities, ClientProvidedConfig, ServerCapabilities } from "./protocol";
+import type { ClientCapabilities, ClientProvidedConfig, ServerCapabilities, DataStoreRecords } from "./protocol";
 
 export interface Feature {
   initialize(
     connection: Connection,
     clientCapabilities: ClientCapabilities,
     clientProvidedConfig: ClientProvidedConfig,
+    dataStoreRecords: DataStoreRecords | undefined,
   ): ServerCapabilities | Promise<ServerCapabilities>;
 
   initialized?(connection: Connection): void | Promise<void>;

clients/tabby-agent/src/protocol.ts

@@ -59,19 +59,26 @@ export namespace InitializeRequest {
 export type InitializeParams = LspInitializeParams & {
   clientInfo?: ClientInfo;
   capabilities: ClientCapabilities;
-  initializationOptions?: {
-    config?: ClientProvidedConfig;
-    /**
-     * ClientInfo also can be provided in InitializationOptions, will be merged with the one in InitializeParams.
-     * This is useful for the clients that don't support changing the ClientInfo in InitializeParams.
-     */
-    clientInfo?: ClientInfo;
-    /**
-     * ClientCapabilities also can be provided in InitializationOptions, will be merged with the one in InitializeParams.
-     * This is useful for the clients that don't support changing the ClientCapabilities in InitializeParams.
-     */
-    clientCapabilities?: ClientCapabilities;
-  };
+  initializationOptions?: InitializationOptions;
+};
+
+export type InitializationOptions = {
+  config?: ClientProvidedConfig;
+  /**
+   * ClientInfo also can be provided in InitializationOptions, will be merged with the one in InitializeParams.
+   * This is useful for the clients that don't support changing the ClientInfo in InitializeParams.
+   */
+  clientInfo?: ClientInfo;
+  /**
+   * ClientCapabilities also can be provided in InitializationOptions, will be merged with the one in InitializeParams.
+   * This is useful for the clients that don't support changing the ClientCapabilities in InitializeParams.
+   */
+  clientCapabilities?: ClientCapabilities;
+  /**
+   * The data store records that should be initialized when the server starts. This is useful for the clients that
+   * provides the dataStore capability.
+   */
+  dataStoreRecords?: DataStoreRecords;
 };
 
 export type InitializeResult = LspInitializeResult & {
@@ -126,9 +133,9 @@ export type ClientCapabilities = LspClientCapabilities & {
      */
     workspaceFileSystem?: boolean;
     /**
-     * The client supports:
-     * - `tabby/dataStore/get`
-     * - `tabby/dataStore/set`
+     * The client provides a initial data store records for initialization and supports methods:
+     * - `tabby/dataStore/didUpdate`
+     * - `tabby/dataStore/update`
      * When not provided, the server will try to fallback to the default data store,
      *  a file-based data store (~/.tabby-client/agent/data.json), which is not available in the browser.
      */
@@ -939,6 +946,7 @@ export type ReadFileResult = {
 };
 
 /**
+ * @deprecated see {@link InitializationOptions} and {@link DataStoreDidUpdateNotification}
  * [Tabby] DataStore Get Request(↪️)
  *
  * This method is sent from the server to the client to get the value of the given key.
@@ -952,11 +960,13 @@ export namespace DataStoreGetRequest {
   export const type = new ProtocolRequestType<DataStoreGetParams, any, never, void, void>(method);
 }
 
+/** @deprecated */
 export type DataStoreGetParams = {
   key: string;
 };
 
 /**
+ * @deprecated see {@link DataStoreUpdateRequest}
  * [Tabby] DataStore Set Request(↪️)
  *
  * This method is sent from the server to the client to set the value of the given key.
@@ -970,11 +980,41 @@ export namespace DataStoreSetRequest {
   export const type = new ProtocolRequestType<DataStoreSetParams, boolean, never, void, void>(method);
 }
 
+/** @deprecated */
 export type DataStoreSetParams = {
   key: string;
   value: any;
 };
 
+/**
+ * [Tabby] DataStore DidUpdate Notification(➡️)
+ *
+ * This method is sent from the client to the server to notify that the data store records has been updated.
+ * - method: `tabby/dataStore/didUpdate`
+ * - params: {@link DataStoreDidChangeParams}
+ */
+export namespace DataStoreDidUpdateNotification {
+  export const method = "tabby/dataStore/didUpdate";
+  export const messageDirection = MessageDirection.clientToServer;
+  export const type = new ProtocolNotificationType<DataStoreRecords, void>(method);
+}
+
+/**
+ * [Tabby] DataStore Update Request(↪️)
+ *
+ * This method is sent from the server to the client to update the data store records.
+ * - method: `tabby/dataStore/update`
+ * - params: {@link DataStoreUpdateParams}
+ * - result: boolean
+ */
+export namespace DataStoreUpdateRequest {
+  export const method = "tabby/dataStore/update";
+  export const messageDirection = MessageDirection.serverToClient;
+  export const type = new ProtocolRequestType<DataStoreRecords, boolean, never, void, void>(method);
+}
+
+export type DataStoreRecords = Record<string, any>;
+
 /**
  * [Tabby] Language Support Declaration Request(↪️)
  *

clients/tabby-agent/src/server.ts

@@ -13,6 +13,7 @@ import {
   ClientCapabilities,
   ServerCapabilities,
   ClientProvidedConfig,
+  DataStoreRecords,
   AgentServerInfoRequest,
   AgentServerInfoSync,
   ServerInfo,
@@ -154,6 +155,7 @@ export class Server {
     this.clientCapabilities = clientCapabilities;
 
     const clientProvidedConfig: ClientProvidedConfig = params.initializationOptions?.config ?? {};
+    const dataStoreRecords: DataStoreRecords | undefined = params.initializationOptions?.dataStoreRecords;
 
     const baseCapabilities: ServerCapabilities = {
       textDocumentSync: {
@@ -176,7 +178,7 @@ export class Server {
     };
 
     this.logger.debug("Initializing internal components...");
-    await this.dataStore.initialize(this.connection, clientCapabilities);
+    await this.dataStore.initialize(this.connection, clientCapabilities, clientProvidedConfig, dataStoreRecords);
     await this.configurations.initialize(this.connection, clientCapabilities, clientProvidedConfig);
     await this.anonymousUsageLogger.initialize(clientInfo);
     await this.tabbyApiClient.initialize(clientInfo);
@@ -196,7 +198,7 @@ export class Server {
       this.commandProvider,
       this.fileTracker,
     ].mapAsync((feature: Feature) => {
-      return feature.initialize(this.connection, clientCapabilities, clientProvidedConfig);
+      return feature.initialize(this.connection, clientCapabilities, clientProvidedConfig, dataStoreRecords);
     });
     this.logger.debug("Feature components initialized.");
 
@@ -217,11 +219,15 @@ export class Server {
 
   private async initialized(): Promise<void> {
     this.logger.info("Received initialized notification.");
-    await [this.configurations, this.statusProvider, this.completionProvider, this.chatFeature].mapAsync(
-      (feature: Feature) => {
-        return feature.initialized?.(this.connection);
-      },
-    );
+    await [
+      this.dataStore,
+      this.configurations,
+      this.statusProvider,
+      this.completionProvider,
+      this.chatFeature,
+    ].mapAsync((feature: Feature) => {
+      return feature.initialized?.(this.connection);
+    });
 
     // FIXME(@icycodes): remove deprecated methods
     if (this.clientCapabilities?.tabby?.agent) {