Refactor LanguageClientConsumer to wait properly

This makes a lot more sense and fixes a long-standing bug. Also cleaned up a lot
while at it.

Author: andyleejordan

src/features/CodeActions.ts

@@ -5,25 +5,25 @@ import vscode = require("vscode");
 import { ILogger } from "../logging";
 
 export class CodeActionsFeature implements vscode.Disposable {
-    private showDocumentationCommand: vscode.Disposable;
+    private command: vscode.Disposable;
 
     constructor(private log: ILogger) {
         // NOTE: While not exposed to the user via package.json, this is
         // required as the server's code action sends across a command name.
         //
         // TODO: In the far future with LSP 3.19 the server can just set a URL
         // and this can go away. See https://github.com/microsoft/language-server-protocol/issues/1548
-        this.showDocumentationCommand =
+        this.command =
             vscode.commands.registerCommand("PowerShell.ShowCodeActionDocumentation", async (ruleName: string) => {
                 await this.showRuleDocumentation(ruleName);
             });
     }
 
     public dispose(): void {
-        this.showDocumentationCommand.dispose();
+        this.command.dispose();
     }
 
-    public async showRuleDocumentation(ruleId: string): Promise<void> {
+    private async showRuleDocumentation(ruleId: string): Promise<void> {
         const pssaDocBaseURL = "https://docs.microsoft.com/powershell/utility-modules/psscriptanalyzer/rules/";
 
         if (!ruleId) {

src/features/Console.ts

@@ -200,8 +200,8 @@ export class ConsoleFeature extends LanguageClientConsumer {
                 } else {
                     selectionRange = editor.document.lineAt(editor.selection.start.line).range;
                 }
-
-                await this.languageClient?.sendRequest(EvaluateRequestType, {
+                const client = await LanguageClientConsumer.getLanguageClient();
+                await client.sendRequest(EvaluateRequestType, {
                     expression: editor.document.getText(selectionRange),
                 });
 
@@ -217,19 +217,19 @@ export class ConsoleFeature extends LanguageClientConsumer {
         for (const command of this.commands) {
             command.dispose();
         }
+
         for (const handler of this.handlers) {
             handler.dispose();
         }
     }
 
-    public override setLanguageClient(languageClient: LanguageClient): void {
-        this.languageClient = languageClient;
+    public override onLanguageClientSet(languageClient: LanguageClient): void {
         this.handlers = [
-            this.languageClient.onRequest(
+            languageClient.onRequest(
                 ShowChoicePromptRequestType,
                 (promptDetails) => showChoicePrompt(promptDetails)),
 
-            this.languageClient.onRequest(
+            languageClient.onRequest(
                 ShowInputPromptRequestType,
                 (promptDetails) => showInputPrompt(promptDetails)),
         ];

src/features/DebugSession.ts

@@ -126,6 +126,7 @@ export class DebugSessionFeature extends LanguageClientConsumer
     private sessionCount = 1;
     private tempDebugProcess: PowerShellProcess | undefined;
     private tempSessionDetails: IEditorServicesSessionDetails | undefined;
+    private commands: Disposable[];
     private handlers: Disposable[] = [];
 
     constructor(context: ExtensionContext, private sessionManager: SessionManager, private logger: ILogger) {
@@ -138,15 +139,34 @@ export class DebugSessionFeature extends LanguageClientConsumer
             context.subscriptions.push(debug.registerDebugConfigurationProvider("PowerShell", this, triggerKind));
         });
         context.subscriptions.push(debug.registerDebugAdapterDescriptorFactory("PowerShell", this));
+
+        // NOTE: While process and runspace IDs are numbers, command
+        // substitutions in VS Code debug configurations are required to return
+        // strings. Hence to the `toString()` on these.
+        this.commands = [
+            commands.registerCommand("PowerShell.PickPSHostProcess", async () => {
+
+                const processId = await this.pickPSHostProcess();
+                return processId?.toString();
+            }),
+            commands.registerCommand("PowerShell.PickRunspace", async (processId) => {
+                const runspace = await this.pickRunspace(processId);
+                return runspace?.toString();
+            }, this),
+        ];
     }
 
     public dispose(): void {
+        for (const command of this.commands) {
+            command.dispose();
+        }
+
         for (const handler of this.handlers) {
             handler.dispose();
         }
     }
 
-    public override setLanguageClient(languageClient: LanguageClient): void {
+    public override onLanguageClientSet(languageClient: LanguageClient): void {
         this.handlers = [
             languageClient.onNotification(
                 StartDebuggerNotificationType,
@@ -477,7 +497,7 @@ export class DebugSessionFeature extends LanguageClientConsumer
 
         // If nothing is set, prompt for the processId.
         if (!config.customPipeName && !config.processId) {
-            config.processId = await commands.executeCommand("PowerShell.PickPSHostProcess");
+            config.processId = await this.pickPSHostProcess();
             // No process selected. Cancel attach.
             if (!config.processId) {
                 return PREVENT_DEBUG_START;
@@ -511,10 +531,66 @@ export class DebugSessionFeature extends LanguageClientConsumer
 
         return config;
     }
+
+    private async pickPSHostProcess(): Promise<number | undefined> {
+        // Start with the current PowerShell process in the list.
+        const items: IProcessItem[] = [];
+
+        const client = await LanguageClientConsumer.getLanguageClient();
+        const response = await client.sendRequest(GetPSHostProcessesRequestType, {});
+        for (const process of response) {
+            let windowTitle = "";
+            if (process.mainWindowTitle) {
+                windowTitle = `, ${process.mainWindowTitle}`;
+            }
+
+            items.push({
+                label: process.processName,
+                description: `PID: ${process.processId.toString()}${windowTitle}`,
+                processId: process.processId,
+            });
+        }
+
+        if (items.length === 0) {
+            return Promise.reject("There are no PowerShell host processes to attach.");
+        }
+
+        const options: QuickPickOptions = {
+            placeHolder: "Select a PowerShell host process to attach.",
+            matchOnDescription: true,
+            matchOnDetail: true,
+        };
+
+        const item = await window.showQuickPick(items, options);
+
+        return item?.processId ?? undefined;
+    }
+
+    private async pickRunspace(processId: number): Promise<number | undefined> {
+        const client = await LanguageClientConsumer.getLanguageClient();
+        const response = await client.sendRequest(GetRunspaceRequestType, { processId });
+        const items: IRunspaceItem[] = [];
+        for (const runspace of response) {
+            items.push({
+                label: runspace.name,
+                description: `ID: ${runspace.id} - ${runspace.availability}`,
+                id: runspace.id,
+            });
+        }
+
+        const options: QuickPickOptions = {
+            placeHolder: "Select PowerShell runspace to debug",
+            matchOnDescription: true,
+            matchOnDetail: true,
+        };
+
+        const item = await window.showQuickPick(items, options);
+
+        return item?.id ?? undefined;
+    }
 }
 
 export class SpecifyScriptArgsFeature implements Disposable {
-
     private command: Disposable;
     private context: ExtensionContext;
 
@@ -572,111 +648,6 @@ interface IPSHostProcessInfo {
 export const GetPSHostProcessesRequestType =
     new RequestType<IGetPSHostProcessesArguments, IPSHostProcessInfo[], string>("powerShell/getPSHostProcesses");
 
-export class PickPSHostProcessFeature extends LanguageClientConsumer {
-    private command: Disposable;
-    private waitingForClientToken?: CancellationTokenSource;
-    private getLanguageClientResolve?: (value: LanguageClient) => void;
-
-    constructor(private logger: ILogger) {
-        super();
-
-        this.command = commands.registerCommand("PowerShell.PickPSHostProcess", async () => {
-            return this.pickPSHostProcess();
-        });
-    }
-
-    public override setLanguageClient(languageClient: LanguageClient): void {
-        this.languageClient = languageClient;
-
-        if (this.waitingForClientToken && this.getLanguageClientResolve) {
-            this.getLanguageClientResolve(this.languageClient);
-            this.clearWaitingToken();
-        }
-    }
-
-    public dispose(): void {
-        this.command.dispose();
-    }
-
-    private getLanguageClient(): Promise<LanguageClient> {
-        if (this.languageClient !== undefined) {
-            return Promise.resolve(this.languageClient);
-        } else {
-            // If PowerShell isn't finished loading yet, show a loading message
-            // until the LanguageClient is passed on to us
-            this.waitingForClientToken = new CancellationTokenSource();
-
-            return new Promise<LanguageClient>(
-                (resolve, reject) => {
-                    this.getLanguageClientResolve = resolve;
-
-                    void window
-                        .showQuickPick(
-                            ["Cancel"],
-                            { placeHolder: "Attach to PowerShell host process: Please wait, starting PowerShell..." },
-                            this.waitingForClientToken?.token)
-                        .then((response) => {
-                            if (response === "Cancel") {
-                                this.clearWaitingToken();
-                                reject();
-                            }
-                        }, undefined);
-
-                    // Cancel the loading prompt after 60 seconds
-                    setTimeout(() => {
-                        if (this.waitingForClientToken) {
-                            this.clearWaitingToken();
-                            reject();
-
-                            void this.logger.writeAndShowError("Attach to PowerShell host process: PowerShell session took too long to start.");
-                        }
-                    }, 60000);
-                },
-            );
-        }
-    }
-
-    private async pickPSHostProcess(): Promise<number | undefined> {
-        // We need the language client in order to send the request.
-        await this.getLanguageClient();
-
-        // Start with the current PowerShell process in the list.
-        const items: IProcessItem[] = [];
-
-        const response = await this.languageClient?.sendRequest(GetPSHostProcessesRequestType, {});
-        for (const process of response ?? []) {
-            let windowTitle = "";
-            if (process.mainWindowTitle) {
-                windowTitle = `, ${process.mainWindowTitle}`;
-            }
-
-            items.push({
-                label: process.processName,
-                description: `PID: ${process.processId.toString()}${windowTitle}`,
-                processId: process.processId,
-            });
-        }
-
-        if (items.length === 0) {
-            return Promise.reject("There are no PowerShell host processes to attach to.");
-        }
-
-        const options: QuickPickOptions = {
-            placeHolder: "Select a PowerShell host process to attach to",
-            matchOnDescription: true,
-            matchOnDetail: true,
-        };
-
-        const item = await window.showQuickPick(items, options);
-
-        return item?.processId ?? undefined;
-    }
-
-    private clearWaitingToken(): void {
-        this.waitingForClientToken?.dispose();
-        this.waitingForClientToken = undefined;
-    }
-}
 
 interface IRunspaceItem extends QuickPickItem {
     id: number; // payload for the QuickPick UI
@@ -694,99 +665,3 @@ interface IRunspace {
 
 export const GetRunspaceRequestType =
     new RequestType<IGetRunspaceRequestArguments, IRunspace[], string>("powerShell/getRunspace");
-
-export class PickRunspaceFeature extends LanguageClientConsumer {
-    private command: Disposable;
-    private waitingForClientToken?: CancellationTokenSource;
-    private getLanguageClientResolve?: (value: LanguageClient) => void;
-
-    constructor(private logger: ILogger) {
-        super();
-
-        this.command = commands.registerCommand("PowerShell.PickRunspace",
-            async (processId) => { return this.pickRunspace(processId); },
-            this);
-    }
-
-    public override setLanguageClient(languageClient: LanguageClient): void {
-        this.languageClient = languageClient;
-
-        if (this.waitingForClientToken && this.getLanguageClientResolve) {
-            this.getLanguageClientResolve(this.languageClient);
-            this.clearWaitingToken();
-        }
-    }
-
-    public dispose(): void {
-        this.command.dispose();
-    }
-
-    private getLanguageClient(): Promise<LanguageClient> {
-        if (this.languageClient) {
-            return Promise.resolve(this.languageClient);
-        } else {
-            // If PowerShell isn't finished loading yet, show a loading message
-            // until the LanguageClient is passed on to us
-            this.waitingForClientToken = new CancellationTokenSource();
-
-            return new Promise<LanguageClient>(
-                (resolve, reject) => {
-                    this.getLanguageClientResolve = resolve;
-
-                    void window
-                        .showQuickPick(
-                            ["Cancel"],
-                            { placeHolder: "Attach to PowerShell host process: Please wait, starting PowerShell..." },
-                            this.waitingForClientToken?.token)
-                        .then((response) => {
-                            if (response === "Cancel") {
-                                this.clearWaitingToken();
-                                reject();
-                            }
-                        }, undefined);
-
-                    // Cancel the loading prompt after 60 seconds
-                    setTimeout(() => {
-                        if (this.waitingForClientToken) {
-                            this.clearWaitingToken();
-                            reject();
-
-                            void this.logger.writeAndShowError(
-                                "Attach to PowerShell host process: PowerShell session took too long to start.");
-                        }
-                    }, 60000);
-                },
-            );
-        }
-    }
-
-    private async pickRunspace(processId: number): Promise<number | undefined> {
-        // We need the language client in order to send the request.
-        await this.getLanguageClient();
-
-        const response = await this.languageClient?.sendRequest(GetRunspaceRequestType, { processId });
-        const items: IRunspaceItem[] = [];
-        for (const runspace of response ?? []) {
-            items.push({
-                label: runspace.name,
-                description: `ID: ${runspace.id} - ${runspace.availability}`,
-                id: runspace.id,
-            });
-        }
-
-        const options: QuickPickOptions = {
-            placeHolder: "Select PowerShell runspace to debug",
-            matchOnDescription: true,
-            matchOnDetail: true,
-        };
-
-        const item = await window.showQuickPick(items, options);
-
-        return item?.id ?? undefined;
-    }
-
-    private clearWaitingToken(): void {
-        this.waitingForClientToken?.dispose();
-        this.waitingForClientToken = undefined;
-    }
-}