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

- Move python.analysis.restartLanguageServer command registration from
  individual language server managers to LanguageServerWatcher.register(),
  so it is always available regardless of which language server type is active
- Fix async/forEach bug in restartLanguageServers() — forEach ignores
  returned Promises, so awaits inside were silently dropped; replaced with
  for...of
- Fix resource URI reconstruction in restartLanguageServers() — map keys
  'Pylance'/'None' are not URIs, Jedi keys are fsPath strings needing
  Uri.file() not Uri.parse()
- Remove now-redundant static commandDispose and command registration from
  NodeLanguageServerManager and JediLanguageServerManager
- Update unit tests: fix disposable count assertions, fix ICommandManager
  mocks, add test for restart command handler

Author: arymaheshwari

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,15 @@ 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

@@ -27,6 +27,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 +120,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 +136,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 +169,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 +185,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 +497,77 @@ 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('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());