From 6844f88925a6381ccaf64164480f4ac50601a672 Mon Sep 17 00:00:00 2001
From: Andy Jordan <andy.jordan@microsoft.com>
Date: Fri, 30 Sep 2022 09:22:36 -0700
Subject: [PATCH] Handle busy notification for all PowerShell tasks

---
 src/features/Console.ts | 29 +-------------------------
 src/main.ts             |  2 +-
 src/session.ts          | 45 ++++++++++++++++++++++++++++++-----------
 3 files changed, 35 insertions(+), 41 deletions(-)

diff --git a/src/features/Console.ts b/src/features/Console.ts
index 673c79108b..e1ad971813 100644
--- a/src/features/Console.ts
+++ b/src/features/Console.ts
@@ -10,12 +10,9 @@ import { ICheckboxQuickPickItem, showCheckboxQuickPick } from "../controls/check
 import { Logger } from "../logging";
 import Settings = require("../settings");
 import { LanguageClientConsumer } from "../languageClientConsumer";
-import { SessionManager } from "../session";
 
 export const EvaluateRequestType = new RequestType<IEvaluateRequestArguments, void, void>("evaluate");
 export const OutputNotificationType = new NotificationType<IOutputNotificationBody>("output");
-export const ExecutionStatusChangedNotificationType =
-    new NotificationType<ExecutionStatus>("powerShell/executionStatusChanged");
 
 export const ShowChoicePromptRequestType =
     new RequestType<IShowChoicePromptRequestArgs,
@@ -62,13 +59,6 @@ interface IShowInputPromptResponseBody {
     promptCancelled: boolean;
 }
 
-enum ExecutionStatus {
-    Pending,
-    Running,
-    Failed,
-    Aborted,
-    Completed,
-}
 
 function showChoicePrompt(
     promptDetails: IShowChoicePromptRequestArgs,
@@ -182,9 +172,8 @@ function onInputEntered(responseText: string): IShowInputPromptResponseBody {
 export class ConsoleFeature extends LanguageClientConsumer {
     private commands: vscode.Disposable[];
     private handlers: vscode.Disposable[];
-    private resolveStatusBarPromise: (value?: {} | PromiseLike<{}>) => void;
 
-    constructor(private log: Logger, private sessionManager: SessionManager) {
+    constructor(private log: Logger) {
         super();
         this.commands = [
             vscode.commands.registerCommand("PowerShell.RunSelection", async () => {
@@ -242,22 +231,6 @@ export class ConsoleFeature extends LanguageClientConsumer {
             this.languageClient.onRequest(
                 ShowInputPromptRequestType,
                 (promptDetails) => showInputPrompt(promptDetails)),
-
-            // Set up status bar alerts for when PowerShell is executing a script.
-            this.languageClient.onNotification(
-                ExecutionStatusChangedNotificationType,
-                (executionStatusDetails) => {
-                    switch (executionStatusDetails) {
-                        case ExecutionStatus.Running:
-                            this.sessionManager.setSessionBusyStatus();
-                            break;
-                        case ExecutionStatus.Completed:
-                        case ExecutionStatus.Aborted:
-                        case ExecutionStatus.Failed:
-                            this.sessionManager.setSessionRunningStatus();
-                            break;
-                    }
-                })
         ]
     }
 }
diff --git a/src/main.ts b/src/main.ts
index d5dfd9ba74..eaccd6e405 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -149,7 +149,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<IPower
 
     // Features and command registrations that require language client
     languageClientConsumers = [
-        new ConsoleFeature(logger, sessionManager),
+        new ConsoleFeature(logger),
         new ExpandAliasFeature(logger),
         new GetCommandsFeature(logger),
         new ShowHelpFeature(logger),
diff --git a/src/session.ts b/src/session.ts
index 718eabab31..1c6179dc1e 100644
--- a/src/session.ts
+++ b/src/session.ts
@@ -73,6 +73,9 @@ export type IReadSessionFileCallback = (details: IEditorServicesSessionDetails)
 export const SendKeyPressNotificationType =
     new NotificationType<void>("powerShell/sendKeyPress");
 
+export const ExecutionBusyStatusNotificationType =
+    new NotificationType<boolean>("powerShell/executionBusyStatus");
+
 export const PowerShellVersionRequestType =
     new RequestType0<IPowerShellVersionDetails, void>(
         "powerShell/getVersion");
@@ -530,7 +533,7 @@ Type 'help' to get help.
         this.languageServerProcess.onExited(
             async () => {
                 if (this.sessionStatus === SessionStatus.Running) {
-                    this.setSessionStatus("Session Exited", SessionStatus.Failed);
+                    this.setSessionStatus("Session Exited!", SessionStatus.Failed);
                     await this.promptForRestart();
                 }
             });
@@ -658,6 +661,14 @@ Type 'help' to get help.
             this.languageClient.onNotification(
                 SendKeyPressNotificationType,
                 () => { this.languageServerProcess.sendKeyPress(); }),
+
+            this.languageClient.onNotification(
+                ExecutionBusyStatusNotificationType,
+                (isBusy: boolean) => {
+                    if (isBusy) { this.setSessionBusyStatus(); }
+                    else { this.setSessionRunningStatus(); }
+                }
+            ),
         ]
 
         try {
@@ -668,7 +679,7 @@ Type 'help' to get help.
         }
 
         this.versionDetails = await this.languageClient.sendRequest(PowerShellVersionRequestType);
-        this.setSessionRunningStatus(); // This requires the version details to be set.
+        this.setSessionRunningStatus();
         this.sendTelemetryEvent("powershellVersionCheck", { powershellVersion: this.versionDetails.version });
 
         // We haven't "started" until we're done getting the version information.
@@ -716,11 +727,26 @@ Type 'help' to get help.
         this.languageStatusItem = vscode.languages.createLanguageStatusItem("powershell", this.documentSelector);
         this.languageStatusItem.command = { title: statusTitle, command: this.ShowSessionMenuCommandName };
         this.languageStatusItem.text = "$(terminal-powershell)";
+        this.languageStatusItem.detail = "PowerShell";
     }
 
     private setSessionStatus(statusText: string, status: SessionStatus): void {
         this.sessionStatus = status;
-        this.languageStatusItem.detail = "PowerShell " + statusText;
+        this.languageStatusItem.detail = "PowerShell";
+
+        if (this.versionDetails !== undefined) {
+            const version = this.versionDetails.architecture === "x86"
+                ? `${this.versionDetails.displayVersion} (${this.versionDetails.architecture})`
+                : this.versionDetails.displayVersion;
+
+            this.languageStatusItem.text = "$(terminal-powershell) " + version;
+            this.languageStatusItem.detail += " " + version;
+        }
+
+        if (statusText) {
+            this.languageStatusItem.detail += ": " + statusText;
+        }
+
         switch (status) {
             case SessionStatus.Running:
             case SessionStatus.NeverStarted:
@@ -745,21 +771,16 @@ Type 'help' to get help.
 
     }
 
-    public setSessionRunningStatus(): void {
-        const version = this.versionDetails.architecture === "x86"
-            ? `${this.versionDetails.displayVersion} (${this.versionDetails.architecture})`
-            : this.versionDetails.displayVersion;
-
-        this.languageStatusItem.text = "$(terminal-powershell) " + version;
-        this.setSessionStatus(version, SessionStatus.Running);
+    private setSessionRunningStatus(): void {
+        this.setSessionStatus("", SessionStatus.Running);
     }
 
-    public setSessionBusyStatus(): void {
+    private setSessionBusyStatus(): void {
         this.setSessionStatus("Executing...", SessionStatus.Busy);
     }
 
     private async setSessionFailure(message: string, ...additionalMessages: string[]) {
-        this.setSessionStatus("Initialization Error", SessionStatus.Failed);
+        this.setSessionStatus("Initialization Error!", SessionStatus.Failed);
         await this.log.writeAndShowError(message, ...additionalMessages);
     }