From 401331688921f4ef40ba410144afcc9e892e2cff Mon Sep 17 00:00:00 2001 From: Zhiming Ma Date: Sat, 2 Nov 2024 03:30:41 +0800 Subject: [PATCH] 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 | 30 +++----- clients/tabby-agent/src/dataStore/dataFile.ts | 4 + clients/tabby-agent/src/dataStore/index.ts | 75 ++++++++++++------- clients/tabby-agent/src/feature.ts | 3 +- clients/tabby-agent/src/protocol.ts | 72 ++++++++++++++---- clients/tabby-agent/src/server.ts | 20 +++-- clients/tabby-agent/src/status.ts | 70 ++++++++--------- clients/vscode/src/lsp/DataStoreFeature.ts | 49 +++++++----- 8 files changed, 198 insertions(+), 125 deletions(-) diff --git a/clients/tabby-agent/src/config/index.ts b/clients/tabby-agent/src/config/index.ts index 86a1710db39e..2b925b66afb8 100644 --- a/clients/tabby-agent/src/config/index.ts +++ b/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) => { - 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) => { + 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) { diff --git a/clients/tabby-agent/src/dataStore/dataFile.ts b/clients/tabby-agent/src/dataStore/dataFile.ts index f83e4e80aa3d..72479782c09b 100644 --- a/clients/tabby-agent/src/dataStore/dataFile.ts +++ b/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 { diff --git a/clients/tabby-agent/src/dataStore/index.ts b/clients/tabby-agent/src/dataStore/index.ts index 1ee2053ea6af..993570e73e4d 100644 --- a/clients/tabby-agent/src/dataStore/index.ts +++ b/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 = {}; private lspConnection: Connection | undefined = undefined; + private lspInitialized = false; + private fileDataStore: FileDataStore | undefined = undefined; async preInitialize(): Promise { @@ -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 { + async initialize( + connection: Connection, + clientCapabilities: ClientCapabilities, + _clientProvidedConfig: ClientProvidedConfig, + dataStoreRecords: DataStoreRecords | undefined, + ): Promise { 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); } } diff --git a/clients/tabby-agent/src/feature.ts b/clients/tabby-agent/src/feature.ts index 18e331c5c19f..4113e0f64eea 100644 --- a/clients/tabby-agent/src/feature.ts +++ b/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; initialized?(connection: Connection): void | Promise; diff --git a/clients/tabby-agent/src/protocol.ts b/clients/tabby-agent/src/protocol.ts index 147fc85f234f..73dd9f044130 100644 --- a/clients/tabby-agent/src/protocol.ts +++ b/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(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(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(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(method); +} + +export type DataStoreRecords = Record; + /** * [Tabby] Language Support Declaration Request(↪️) * diff --git a/clients/tabby-agent/src/server.ts b/clients/tabby-agent/src/server.ts index 0acc20a01ffb..3536ea2c088c 100644 --- a/clients/tabby-agent/src/server.ts +++ b/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 { 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) { diff --git a/clients/tabby-agent/src/status.ts b/clients/tabby-agent/src/status.ts index e91b30680476..96712ca02bdb 100644 --- a/clients/tabby-agent/src/status.ts +++ b/clients/tabby-agent/src/status.ts @@ -147,45 +147,39 @@ export class StatusProvider extends EventEmitter implements Feature { const issues = Array.isArray(params.issues) ? params.issues : [params.issues]; const dataStore = this.dataStore; switch (params.operation) { - case "add": - if (dataStore) { - const current = dataStore.data.statusIgnoredIssues ?? []; - dataStore.data.statusIgnoredIssues = current.concat(issues).distinct(); - this.logger.debug( - "Adding ignored issues: [" + - current.join(",") + - "] -> [" + - dataStore.data.statusIgnoredIssues.join(",") + - "].", - ); - await dataStore.save(); - return true; - } - break; - case "remove": - if (dataStore) { - const current = dataStore.data.statusIgnoredIssues ?? []; - dataStore.data.statusIgnoredIssues = current.filter((item) => !issues.includes(item)); - this.logger.debug( - "Removing ignored issues: [" + - current.join(",") + - "] -> [" + - dataStore.data.statusIgnoredIssues.join(",") + - "].", - ); + case "add": { + const current = dataStore.data.statusIgnoredIssues ?? []; + dataStore.data.statusIgnoredIssues = current.concat(issues).distinct(); + this.logger.debug( + "Adding ignored issues: [" + + current.join(",") + + "] -> [" + + dataStore.data.statusIgnoredIssues.join(",") + + "].", + ); + await dataStore.save(); + return true; + } + case "remove": { + const current = dataStore.data.statusIgnoredIssues ?? []; + dataStore.data.statusIgnoredIssues = current.filter((item) => !issues.includes(item)); + this.logger.debug( + "Removing ignored issues: [" + + current.join(",") + + "] -> [" + + dataStore.data.statusIgnoredIssues.join(",") + + "].", + ); - await dataStore.save(); - return true; - } - break; - case "removeAll": - if (dataStore) { - dataStore.data.statusIgnoredIssues = []; - this.logger.debug("Removing all ignored issues."); - await dataStore.save(); - return true; - } - break; + await dataStore.save(); + return true; + } + case "removeAll": { + dataStore.data.statusIgnoredIssues = []; + this.logger.debug("Removing all ignored issues."); + await dataStore.save(); + return true; + } default: break; } diff --git a/clients/vscode/src/lsp/DataStoreFeature.ts b/clients/vscode/src/lsp/DataStoreFeature.ts index 0d36718dd0dc..7ca13c1916ba 100644 --- a/clients/vscode/src/lsp/DataStoreFeature.ts +++ b/clients/vscode/src/lsp/DataStoreFeature.ts @@ -1,27 +1,36 @@ +import { EventEmitter } from "events"; import { env, ExtensionContext } from "vscode"; import { BaseLanguageClient, StaticFeature, FeatureState, Disposable } from "vscode-languageclient"; import { + InitializeParams, ClientCapabilities, - DataStoreGetRequest, - DataStoreGetParams, - DataStoreSetRequest, - DataStoreSetParams, + DataStoreRecords, + DataStoreDidUpdateNotification, + DataStoreUpdateRequest, } from "tabby-agent"; -export class DataStoreFeature implements StaticFeature { +export class DataStoreFeature extends EventEmitter implements StaticFeature { private disposables: Disposable[] = []; constructor( private readonly context: ExtensionContext, private readonly client: BaseLanguageClient, - ) {} + ) { + super(); + } getState(): FeatureState { return { kind: "static" }; } - fillInitializeParams() { - // nothing + fillInitializeParams(params: InitializeParams) { + if (env.appHost === "desktop") { + return; + } + params.initializationOptions = { + ...params.initializationOptions, + dataStoreRecords: this.getRecords(), + }; } fillClientCapabilities(capabilities: ClientCapabilities): void { @@ -43,19 +52,23 @@ export class DataStoreFeature implements StaticFeature { return; } this.disposables.push( - this.client.onRequest(DataStoreGetRequest.type, (params: DataStoreGetParams) => { - const data: Record = this.context.globalState.get("data", {}); - return data[params.key]; - }), - ); - this.disposables.push( - this.client.onRequest(DataStoreSetRequest.type, async (params: DataStoreSetParams) => { - const data: Record = this.context.globalState.get("data", {}); - data[params.key] = params.value; - this.context.globalState.update("data", data); + this.client.onRequest(DataStoreUpdateRequest.type, async (params: DataStoreRecords) => { + this.update(params); return true; }), ); + this.on("didUpdate", (records) => { + this.client.sendNotification(DataStoreDidUpdateNotification.type, records); + }); + } + + private getRecords(): DataStoreRecords { + return this.context.globalState.get("data", {}); + } + + private update(records: DataStoreRecords) { + this.context.globalState.update("data", records); + this.emit("didUpdate", records); } clear(): void {