From acdb4feee67a4b0ffb6d4dfdd8e48f12df631d36 Mon Sep 17 00:00:00 2001
From: Wolmir Nemitz <wolmir.nemitz@gmail.com>
Date: Mon, 3 Oct 2022 09:25:37 -0300
Subject: [PATCH] Add config option to disable the language server

---
 extension/package.json                    |  6 ++++++
 extension/package.nls.json                |  2 ++
 extension/src/extension.ts                |  3 +++
 extension/src/lspClient/languageClient.ts | 26 +++++++++++++++++++++--
 extension/src/vscode/config.ts            |  1 +
 5 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/extension/package.json b/extension/package.json
index 4aab7a8dfb..eece9f6842 100644
--- a/extension/package.json
+++ b/extension/package.json
@@ -548,6 +548,12 @@
           "type": "boolean",
           "default": null
         },
+        "dvc.disableLanguageServer": {
+          "title": "%config.disableLanguageServer.title%",
+          "description": "%config.disableLanguageServer.description%",
+          "type": "boolean",
+          "default": null
+        },
         "dvc.dvcPath": {
           "title": "%config.dvcPath.title%",
           "description": "%config.dvcPath.description%",
diff --git a/extension/package.nls.json b/extension/package.nls.json
index 2d51650d3a..09c6b69e06 100644
--- a/extension/package.nls.json
+++ b/extension/package.nls.json
@@ -82,6 +82,8 @@
   "config.doNotShowWalkthroughAfterInstall.title": "Do not show the Get Started page.",
   "config.doNotShowUnableToFilter.description": "Do not warn before disabling auto-apply filters when these would result in too many experiments being selected.",
   "config.doNotShowUnableToFilter.title": "Do not warn before disabling auto-apply filters to experiment selection.",
+  "config.disableLanguageServer.description": "Disable the language server for dvc.yaml files.",
+  "config.disableLanguageServer.title": "Disable language server",
   "config.dvcPath.description": "Path or shell command to the DVC binary. Required unless Microsoft's Python extension is installed and the `dvc` package found in its environment.",
   "config.dvcPath.title": "DVC Path (or shell command)",
   "config.expTableHeadMaxLayers.title": "Maximum depth of Experiment table head rows",
diff --git a/extension/src/extension.ts b/extension/src/extension.ts
index bdb0968528..5dd1a4e58f 100644
--- a/extension/src/extension.ts
+++ b/extension/src/extension.ts
@@ -51,6 +51,7 @@ import { collectWorkspaceScale } from './telemetry/collect'
 import { createFileSystemWatcher } from './fileSystem/watcher'
 import { GitExecutor } from './cli/git/executor'
 import { GitReader } from './cli/git/reader'
+import { LanguageClientWrapper } from './lspClient/languageClient'
 
 export class Extension extends Disposable implements IExtension {
   protected readonly internalCommands: InternalCommands
@@ -92,6 +93,8 @@ export class Extension extends Disposable implements IExtension {
     this.setCommandsAvailability(false)
     this.setProjectAvailability()
 
+    this.dispose.track(new LanguageClientWrapper())
+
     this.resourceLocator = this.dispose.track(
       new ResourceLocator(context.extensionUri)
     )
diff --git a/extension/src/lspClient/languageClient.ts b/extension/src/lspClient/languageClient.ts
index 29b8afd6b6..6259761aac 100644
--- a/extension/src/lspClient/languageClient.ts
+++ b/extension/src/lspClient/languageClient.ts
@@ -5,7 +5,9 @@ import {
   TransportKind
 } from 'vscode-languageclient/node'
 import { documentSelector, serverModule } from 'dvc-vscode-lsp'
+import { ConfigurationChangeEvent, workspace } from 'vscode'
 import { Disposable } from '../class/dispose'
+import { ConfigKey, getConfigValue } from '../vscode/config'
 
 export class LanguageClientWrapper extends Disposable {
   private client: LanguageClient
@@ -26,12 +28,24 @@ export class LanguageClientWrapper extends Disposable {
       )
     )
 
-    // Start the client. This will also launch the server
     this.start()
   }
 
   async start() {
-    await this.client.start()
+    if (!this.languageServerDisabled()) {
+      // This will also start the server process
+      await this.client.start()
+
+      workspace.onDidChangeConfiguration((event: ConfigurationChangeEvent) => {
+        if (
+          event.affectsConfiguration(ConfigKey.DISABLE_LANGUAGE_SERVER) &&
+          this.isRunning() &&
+          this.languageServerDisabled()
+        ) {
+          this.stop()
+        }
+      })
+    }
 
     return this
   }
@@ -40,6 +54,14 @@ export class LanguageClientWrapper extends Disposable {
     this.client.stop()
   }
 
+  private isRunning() {
+    return this.client.isRunning()
+  }
+
+  private languageServerDisabled() {
+    return getConfigValue<boolean>(ConfigKey.DISABLE_LANGUAGE_SERVER)
+  }
+
   private getServerOptions(): ServerOptions {
     const debugOptions = { execArgv: ['--nolazy', '--inspect=6009'] }
 
diff --git a/extension/src/vscode/config.ts b/extension/src/vscode/config.ts
index ad3e91f9f9..441f57dce0 100644
--- a/extension/src/vscode/config.ts
+++ b/extension/src/vscode/config.ts
@@ -5,6 +5,7 @@ export enum ConfigKey {
   DO_NOT_SHOW_CLI_UNAVAILABLE = 'dvc.doNotShowCliUnavailable',
   DO_NOT_SHOW_WALKTHROUGH_AFTER_INSTALL = 'dvc.doNotShowWalkthroughAfterInstall',
   DO_NOT_SHOW_UNABLE_TO_FILTER = 'dvc.doNotShowUnableToFilter',
+  DISABLE_LANGUAGE_SERVER = 'dvc.disableLanguageServer',
   EXP_TABLE_HEAD_MAX_LAYERS = 'dvc.expTableHeadMaxLayers',
   DVC_PATH = 'dvc.dvcPath',
   PYTHON_PATH = 'dvc.pythonPath'