Fix restartLanguageServer command not found in command palette (#22045)

src/client/activation/jedi/manager.ts

@@ -13,7 +13,6 @@ import { IServiceContainer } from '../../ioc/types';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { captureTelemetry } from '../../telemetry';
 import { EventName } from '../../telemetry/constants';
-import { Commands } from '../commands';
 import { JediLanguageClientMiddleware } from './languageClientMiddleware';
 import { ILanguageServerAnalysisOptions, ILanguageServerManager, ILanguageServerProxy } from '../types';
 import { traceDecoratorError, traceDecoratorVerbose, traceVerbose } from '../../logging';
@@ -27,8 +26,6 @@ export class JediLanguageServerManager implements ILanguageServerManager {
 
     private disposables: IDisposable[] = [];
 
-    private static commandDispose: IDisposable;
-
     private connected = false;
 
     private lsVersion: string | undefined;
@@ -37,15 +34,8 @@ export class JediLanguageServerManager implements ILanguageServerManager {
         private readonly serviceContainer: IServiceContainer,
         private readonly analysisOptions: ILanguageServerAnalysisOptions,
         private readonly languageServerProxy: ILanguageServerProxy,
-        commandManager: ICommandManager,
-    ) {
-        if (JediLanguageServerManager.commandDispose) {
-            JediLanguageServerManager.commandDispose.dispose();
-        }
-        JediLanguageServerManager.commandDispose = commandManager.registerCommand(Commands.RestartLS, () => {
-            this.restartLanguageServer().ignoreErrors();
-        });
-    }
+        _commandManager: ICommandManager,
+    ) {}
 
     private static versionTelemetryProps(instance: JediLanguageServerManager) {
         return {
@@ -55,7 +45,6 @@ export class JediLanguageServerManager implements ILanguageServerManager {
 
     public dispose(): void {
         this.stopLanguageServer().ignoreErrors();
-        JediLanguageServerManager.commandDispose.dispose();
         this.disposables.forEach((d) => d.dispose());
     }
 

src/client/activation/node/manager.ts

@@ -9,7 +9,6 @@ import { IServiceContainer } from '../../ioc/types';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { captureTelemetry, sendTelemetryEvent } from '../../telemetry';
 import { EventName } from '../../telemetry/constants';
-import { Commands } from '../commands';
 import { NodeLanguageClientMiddleware } from './languageClientMiddleware';
 import { ILanguageServerAnalysisOptions, ILanguageServerManager } from '../types';
 import { traceDecoratorError, traceDecoratorVerbose } from '../../logging';
@@ -31,23 +30,13 @@ export class NodeLanguageServerManager implements ILanguageServerManager {
 
     private started = false;
 
-    private static commandDispose: IDisposable;
-
     constructor(
         private readonly serviceContainer: IServiceContainer,
         private readonly analysisOptions: ILanguageServerAnalysisOptions,
         private readonly languageServerProxy: NodeLanguageServerProxy,
-        commandManager: ICommandManager,
+        _commandManager: ICommandManager,
         private readonly extensions: IExtensions,
-    ) {
-        if (NodeLanguageServerManager.commandDispose) {
-            NodeLanguageServerManager.commandDispose.dispose();
-        }
-        NodeLanguageServerManager.commandDispose = commandManager.registerCommand(Commands.RestartLS, () => {
-            sendTelemetryEvent(EventName.LANGUAGE_SERVER_RESTART, undefined, { reason: 'command' });
-            this.restartLanguageServer().ignoreErrors();
-        });
-    }
+    ) {}
 
     private static versionTelemetryProps(instance: NodeLanguageServerManager) {
         return {
@@ -57,7 +46,6 @@ export class NodeLanguageServerManager implements ILanguageServerManager {
 
     public dispose(): void {
         this.stopLanguageServer().ignoreErrors();
-        NodeLanguageServerManager.commandDispose.dispose();
         this.disposables.forEach((d) => d.dispose());
     }
 

src/client/languageServer/watcher.ts

@@ -30,6 +30,7 @@ import { ILanguageServerExtensionManager, ILanguageServerWatcher } from './types
 import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { StopWatch } from '../common/utils/stopWatch';
+import { Commands } from '../activation/commands';
 
 @injectable()
 /**
@@ -113,6 +114,13 @@ export class LanguageServerWatcher implements IExtensionActivationService, ILang
                 }),
             );
 
+            this.disposables.push(
+                this.commandManager.registerCommand(Commands.RestartLS, async () => {
+                    sendTelemetryEvent(EventName.LANGUAGE_SERVER_RESTART, undefined, { reason: 'command' });
+                    await this.restartLanguageServers();
+                }),
+            );
+
             this.disposables.push(
                 new LanguageServerChangeHandler(
                     this.languageServerType,
@@ -195,12 +203,13 @@ export class LanguageServerWatcher implements IExtensionActivationService, ILang
     }
 
     public async restartLanguageServers(): Promise<void> {
-        this.workspaceLanguageServers.forEach(async (_, resourceString) => {
-            sendTelemetryEvent(EventName.LANGUAGE_SERVER_RESTART, undefined, { reason: 'notebooksExperiment' });
-            const resource = Uri.parse(resourceString);
+        const resourceStrings = [...this.workspaceLanguageServers.keys()];
+        for (const resourceString of resourceStrings) {
+            const resource =
+                resourceString === 'Pylance' || resourceString === 'None' ? undefined : Uri.file(resourceString);
             await this.stopLanguageServer(resource);
             await this.startLanguageServer(this.languageServerType, resource);
-        });
+        }
     }
 
     public async get(resource?: Resource): Promise<ILanguageServerExtensionManager> {

src/test/languageServer/watcher.unit.test.ts

@@ -3,12 +3,13 @@
 
 import * as assert from 'assert';
 import * as sinon from 'sinon';
-import { ConfigurationChangeEvent, Uri, WorkspaceFolder, WorkspaceFoldersChangeEvent } from 'vscode';
+import { ConfigurationChangeEvent, Disposable, Uri, WorkspaceFolder, WorkspaceFoldersChangeEvent } from 'vscode';
 import { JediLanguageServerManager } from '../../client/activation/jedi/manager';
 import { NodeLanguageServerManager } from '../../client/activation/node/manager';
 import { ILanguageServerOutputChannel, LanguageServerType } from '../../client/activation/types';
 import { IApplicationShell, ICommandManager, IWorkspaceService } from '../../client/common/application/types';
 import { IFileSystem } from '../../client/common/platform/types';
+import { PYLANCE_EXTENSION_ID } from '../../client/common/constants';
 import {
     IConfigurationService,
     IDisposable,
@@ -17,6 +18,7 @@ import {
     IInterpreterPathService,
 } from '../../client/common/types';
 import { LanguageService } from '../../client/common/utils/localize';
+import { createDeferred } from '../../client/common/utils/async';
 import { IEnvironmentVariablesProvider } from '../../client/common/variables/types';
 import { IInterpreterHelper, IInterpreterService } from '../../client/interpreter/contracts';
 import { IServiceContainer } from '../../client/ioc/types';
@@ -27,6 +29,7 @@ import { ILanguageServerExtensionManager } from '../../client/languageServer/typ
 import { LanguageServerWatcher } from '../../client/languageServer/watcher';
 import * as Logging from '../../client/logging';
 import { PythonEnvironment } from '../../client/pythonEnvironments/info';
+import { Commands } from '../../client/activation/commands';
 
 suite('Language server watcher', () => {
     let watcher: LanguageServerWatcher;
@@ -119,7 +122,11 @@ suite('Language server watcher', () => {
                     /* do nothing */
                 },
             } as unknown) as IWorkspaceService,
-            {} as ICommandManager,
+            ({
+                registerCommand: () => {
+                    /* do nothing */
+                },
+            } as unknown) as ICommandManager,
             {} as IFileSystem,
             ({
                 getExtension: () => undefined,
@@ -131,7 +138,7 @@ suite('Language server watcher', () => {
             disposables,
         );
         watcher.register();
-        assert.strictEqual(disposables.length, 11);
+        assert.strictEqual(disposables.length, 13);
     });
 
     test('The constructor should not add a listener to onDidChange to the list of disposables if it is not a trusted workspace', () => {
@@ -164,7 +171,11 @@ suite('Language server watcher', () => {
                     /* do nothing */
                 },
             } as unknown) as IWorkspaceService,
-            {} as ICommandManager,
+            ({
+                registerCommand: () => {
+                    /* do nothing */
+                },
+            } as unknown) as ICommandManager,
             {} as IFileSystem,
             ({
                 getExtension: () => undefined,
@@ -176,7 +187,7 @@ suite('Language server watcher', () => {
             disposables,
         );
         watcher.register();
-        assert.strictEqual(disposables.length, 10);
+        assert.strictEqual(disposables.length, 12);
     });
 
     test(`When starting the language server, the language server extension manager should not be undefined`, async () => {
@@ -488,6 +499,169 @@ suite('Language server watcher', () => {
         assert.ok(startLanguageServerFake.calledTwice);
     });
 
+    test('When the "Restart Language Server" command is executed, all running language servers should be stopped and restarted', async () => {
+        let restartCommandHandler: (() => Promise<void>) | undefined;
+
+        watcher = new LanguageServerWatcher(
+            ({
+                get: () => {
+                    /* do nothing */
+                },
+            } as unknown) as IServiceContainer,
+            {} as ILanguageServerOutputChannel,
+            {
+                getSettings: () => ({ languageServer: LanguageServerType.None }),
+            } as IConfigurationService,
+            {} as IExperimentService,
+            ({
+                getActiveWorkspaceUri: () => undefined,
+            } as unknown) as IInterpreterHelper,
+            ({
+                onDidChange: () => {
+                    /* do nothing */
+                },
+            } as unknown) as IInterpreterPathService,
+            ({
+                getActiveInterpreter: () => 'python',
+                onDidChangeInterpreterInformation: () => {
+                    /* do nothing */
+                },
+            } as unknown) as IInterpreterService,
+            {} as IEnvironmentVariablesProvider,
+            ({
+                getWorkspaceFolder: (uri: Uri) => ({ uri }),
+                onDidChangeConfiguration: () => {
+                    /* do nothing */
+                },
+                onDidChangeWorkspaceFolders: () => {
+                    /* do nothing */
+                },
+            } as unknown) as IWorkspaceService,
+            ({
+                registerCommand: (command: string, handler: () => Promise<void>) => {
+                    if (command === Commands.RestartLS) {
+                        restartCommandHandler = handler;
+                    }
+                },
+            } as unknown) as ICommandManager,
+            {} as IFileSystem,
+            ({
+                getExtension: () => undefined,
+                onDidChange: () => {
+                    /* do nothing */
+                },
+            } as unknown) as IExtensions,
+            {} as IApplicationShell,
+            disposables,
+        );
+        watcher.register();
+
+        await watcher.startLanguageServer(LanguageServerType.None);
+
+        const stopLanguageServerStub = sandbox.stub(NoneLSExtensionManager.prototype, 'stopLanguageServer');
+        stopLanguageServerStub.returns(Promise.resolve());
+        const startLanguageServerStub = sandbox.stub(NoneLSExtensionManager.prototype, 'startLanguageServer');
+        startLanguageServerStub.returns(Promise.resolve());
+
+        assert.ok(restartCommandHandler, 'Restart command handler should be registered');
+        await restartCommandHandler!();
+
+        assert.ok(stopLanguageServerStub.calledOnce, 'stopLanguageServer should be called once');
+        assert.ok(startLanguageServerStub.calledOnce, 'startLanguageServer should be called once');
+    });
+
+    test('The restart command waits for the Pylance client to stop before starting it again', async () => {
+        let restartCommandHandler: (() => Promise<void>) | undefined;
+        const stopStarted = createDeferred<void>();
+        const stopCompleted = createDeferred<void>();
+        const startClient = sandbox.stub().resolves();
+        const stopClient = sandbox.stub().callsFake(async () => {
+            stopStarted.resolve();
+            await stopCompleted.promise;
+        });
+        const pylanceExtension = {
+            isActive: true,
+            packageJSON: { version: '1.0.0' },
+            exports: {
+                client: {
+                    isEnabled: () => true,
+                    start: startClient,
+                    stop: stopClient,
+                },
+            },
+        };
+        const resource = Uri.file('workspace');
+        const interpreter = { path: 'python', version: { major: 3, minor: 12 } };
+        const registerEvent = () => new Disposable(() => {});
+        const testDisposables: IDisposable[] = [];
+
+        watcher = new LanguageServerWatcher(
+            {} as IServiceContainer,
+            {} as ILanguageServerOutputChannel,
+            {
+                getSettings: () => ({ languageServer: LanguageServerType.Node }),
+            } as IConfigurationService,
+            {} as IExperimentService,
+            ({
+                getActiveWorkspaceUri: () => ({ folderUri: resource }),
+            } as unknown) as IInterpreterHelper,
+            ({ onDidChange: registerEvent } as unknown) as IInterpreterPathService,
+            ({
+                getActiveInterpreter: () => interpreter,
+                onDidChangeInterpreterInformation: registerEvent,
+            } as unknown) as IInterpreterService,
+            {} as IEnvironmentVariablesProvider,
+            ({
+                isTrusted: true,
+                getWorkspaceFolder: () => ({ uri: resource }),
+                onDidChangeConfiguration: registerEvent,
+                onDidChangeWorkspaceFolders: registerEvent,
+            } as unknown) as IWorkspaceService,
+            ({
+                registerCommand: (command: string, handler: () => Promise<void>) => {
+                    if (command === Commands.RestartLS) {
+                        restartCommandHandler = handler;
+                    }
+                    return registerEvent();
+                },
+            } as unknown) as ICommandManager,
+            {} as IFileSystem,
+            ({
+                getExtension: (id: string) => (id === PYLANCE_EXTENSION_ID ? pylanceExtension : undefined),
+                onDidChange: registerEvent,
+            } as unknown) as IExtensions,
+            {} as IApplicationShell,
+            testDisposables,
+        );
+
+        try {
+            await watcher.activate(resource);
+            const originalManager = await watcher.get(resource);
+            const disposeManager = sandbox.spy(originalManager, 'dispose');
+            assert.ok(originalManager instanceof PylanceLSExtensionManager);
+            assert.ok(startClient.calledOnce, 'The dummy Pylance client should initially be started');
+            assert.ok(restartCommandHandler, 'The restart command should be registered');
+
+            const restarting = restartCommandHandler!();
+            await stopStarted.promise;
+            assert.ok(startClient.calledOnce, 'Pylance must not start again while its previous client is stopping');
+            stopCompleted.resolve();
+            await restarting;
+
+            assert.ok(stopClient.calledOnce, 'The previous Pylance client should be stopped');
+            assert.ok(disposeManager.calledOnce, 'The previous Pylance manager should be disposed');
+            assert.ok(startClient.calledTwice, 'The Pylance client should be started again');
+            assert.notStrictEqual(await watcher.get(resource), originalManager);
+        } finally {
+            stopCompleted.resolve();
+            for (const disposable of testDisposables) {
+                await disposable.dispose();
+            }
+        }
+
+        assert.ok(stopClient.calledTwice, 'The restarted Pylance client should also be stopped during cleanup');
+    });
+
     test('When starting a language server with a Python 2.7 interpreter and the python.languageServer setting is Jedi, do not instantiate a language server', async () => {
         const startLanguageServerStub = sandbox.stub(NoneLSExtensionManager.prototype, 'startLanguageServer');
         startLanguageServerStub.returns(Promise.resolve());