From 2e0023fb3c397ea776c1e6b8b247e75ac6af0163 Mon Sep 17 00:00:00 2001
From: Andrew Eisenberg <aeisenberg@github.com>
Date: Fri, 24 Mar 2023 17:15:49 +0000
Subject: [PATCH] Restart language server on command

Ensure that the language server is restarted when the "Restart
Query Server" command is invoked.
---
 extensions/ql-vscode/CHANGELOG.md     |  2 ++
 extensions/ql-vscode/src/extension.ts | 10 +++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/extensions/ql-vscode/CHANGELOG.md b/extensions/ql-vscode/CHANGELOG.md
index 56dcf5fc9da..910758c116b 100644
--- a/extensions/ql-vscode/CHANGELOG.md
+++ b/extensions/ql-vscode/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## [UNRELEASED]
 
+- Restart the CodeQL language server whenever the _CodeQL: Restart Query Server_ command is invoked. This avoids bugs where the CLI version changes to support new language features, but the language server is not updated. [#2238](https://github.com/github/vscode-codeql/pull/2238)
+
 ## 1.8.1 - 23 March 2023
 
 - Show data flow paths of a variant analysis in a new tab. [#2172](https://github.com/github/vscode-codeql/pull/2172) & [#2182](https://github.com/github/vscode-codeql/pull/2182)
diff --git a/extensions/ql-vscode/src/extension.ts b/extensions/ql-vscode/src/extension.ts
index 18d3c1a66c4..ebfdf76d201 100644
--- a/extensions/ql-vscode/src/extension.ts
+++ b/extensions/ql-vscode/src/extension.ts
@@ -161,6 +161,7 @@ function getCommands(
   app: App,
   cliServer: CodeQLCliServer,
   queryRunner: QueryRunner,
+  ideServer: LanguageClient,
 ): BaseCommands {
   const getCliVersion = async () => {
     try {
@@ -177,9 +178,12 @@ function getCommands(
     "codeQL.restartQueryServer": async () =>
       withProgress(
         async (progress: ProgressCallback, token: CancellationToken) => {
-          // We restart the CLI server too, to ensure they are the same version
+          // Restart all of the spawned servers: cli, query, and language.
           cliServer.restartCliServer();
-          await queryRunner.restartQueryServer(progress, token);
+          await Promise.all([
+            queryRunner.restartQueryServer(progress, token),
+            ideServer.restart(),
+          ]);
           void showAndLogInformationMessage("CodeQL Query Server restarted.", {
             outputLogger: queryServerLogger,
           });
@@ -889,7 +893,7 @@ async function activateWithInstalledDistribution(
   void extLogger.log("Registering top-level command palette commands.");
 
   const allCommands: AllExtensionCommands = {
-    ...getCommands(app, cliServer, qs),
+    ...getCommands(app, cliServer, qs, client),
     ...getQueryEditorCommands({
       commandManager: app.commands,
       queryRunner: qs,