Generalize the execution busy status to all PowerShell tasks (#1928)

We generalized the "execution busy status" from just editor commands to any running PowerShell task.

Author: andyleejordan

src/PowerShellEditorServices/Services/Extension/ExtensionService.cs

@@ -14,37 +14,6 @@
 
 namespace Microsoft.PowerShell.EditorServices.Services.Extension
 {
-    /// Enumerates the possible execution results that can occur after
-    /// executing a command or script.
-    /// </summary>
-    internal enum ExecutionStatus
-    {
-        /// <summary>
-        /// Indicates that execution has not yet started.
-        /// </summary>
-        Pending,
-
-        /// <summary>
-        /// Indicates that the command is executing.
-        /// </summary>
-        Running,
-
-        /// <summary>
-        /// Indicates that execution has failed.
-        /// </summary>
-        Failed,
-
-        /// <summary>
-        /// Indicates that execution was aborted by the user.
-        /// </summary>
-        Aborted,
-
-        /// <summary>
-        /// Indicates that execution completed successfully.
-        /// </summary>
-        Completed
-    }
-
     /// <summary>
     /// Provides a high-level service which enables PowerShell scripts
     /// and modules to extend the behavior of the host editor.
@@ -154,7 +123,6 @@ internal Task InitializeAsync()
         /// <exception cref="KeyNotFoundException">The command being invoked was not registered.</exception>
         public Task InvokeCommandAsync(string commandName, EditorContext editorContext, CancellationToken cancellationToken)
         {
-            _languageServer?.SendNotification("powerShell/executionStatusChanged", ExecutionStatus.Pending);
             if (editorCommands.TryGetValue(commandName, out EditorCommand editorCommand))
             {
                 PSCommand executeCommand = new PSCommand()
@@ -163,7 +131,6 @@ public Task InvokeCommandAsync(string commandName, EditorContext editorContext,
                     .AddParameter("ArgumentList", new object[] { editorContext });
 
                 // This API is used for editor command execution so it requires the foreground.
-                _languageServer?.SendNotification("powerShell/executionStatusChanged", ExecutionStatus.Running);
                 return ExecutionService.ExecutePSCommandAsync(
                     executeCommand,
                     cancellationToken,
@@ -173,27 +140,9 @@ public Task InvokeCommandAsync(string commandName, EditorContext editorContext,
                         WriteOutputToHost = !editorCommand.SuppressOutput,
                         AddToHistory = !editorCommand.SuppressOutput,
                         ThrowOnError = false,
-                    }).ContinueWith((Task executeTask) =>
-                    {
-                        ExecutionStatus status = ExecutionStatus.Failed;
-                        if (executeTask.IsCompleted)
-                        {
-                            status = ExecutionStatus.Completed;
-                        }
-                        else if (executeTask.IsCanceled)
-                        {
-                            status = ExecutionStatus.Aborted;
-                        }
-                        else if (executeTask.IsFaulted)
-                        {
-                            status = ExecutionStatus.Failed;
-                        }
-
-                        _languageServer?.SendNotification("powerShell/executionStatusChanged", status);
-                    }, TaskScheduler.Default);
+                    });
             }
 
-            _languageServer?.SendNotification("powerShell/executionStatusChanged", ExecutionStatus.Failed);
             throw new KeyNotFoundException($"Editor command not found: '{commandName}'");
         }
 

src/PowerShellEditorServices/Services/PowerShell/Execution/SynchronousPowerShellTask.cs

@@ -16,7 +16,14 @@
 
 namespace Microsoft.PowerShell.EditorServices.Services.PowerShell.Execution
 {
-    internal class SynchronousPowerShellTask<TResult> : SynchronousTask<IReadOnlyList<TResult>>
+    internal interface ISynchronousPowerShellTask
+    {
+        PowerShellExecutionOptions PowerShellExecutionOptions { get; }
+
+        void MaybeAddToHistory();
+    }
+
+    internal class SynchronousPowerShellTask<TResult> : SynchronousTask<IReadOnlyList<TResult>>, ISynchronousPowerShellTask
     {
         private static readonly PowerShellExecutionOptions s_defaultPowerShellExecutionOptions = new();
 
@@ -353,7 +360,7 @@ private void CancelNormalExecution()
             }
         }
 
-        internal void MaybeAddToHistory()
+        public void MaybeAddToHistory()
         {
             // Do not add PSES internal commands to history. Also exclude input that came from the
             // REPL (e.g. PSReadLine) as it handles history itself in that scenario.

src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs

@@ -295,6 +295,8 @@ public void SetExit()
 
         internal void ForceSetExit() => _shouldExit = true;
 
+        private void SetBusy(bool busy) => _languageServer?.SendNotification("powerShell/executionBusyStatus", busy);
+
         private bool CancelForegroundAndPrepend(ISynchronousTask task, bool isIdle = false)
         {
             // NOTE: This causes foreground tasks to act like they have `ExecutionPriority.Next`.
@@ -313,9 +315,9 @@ private bool CancelForegroundAndPrepend(ISynchronousTask task, bool isIdle = fal
 
             _skipNextPrompt = true;
 
-            if (task is SynchronousPowerShellTask<PSObject> psTask)
+            if (task is ISynchronousPowerShellTask t)
             {
-                psTask.MaybeAddToHistory();
+                t.MaybeAddToHistory();
             }
 
             using (_taskQueue.BlockConsumers())
@@ -334,6 +336,32 @@ private bool CancelForegroundAndPrepend(ISynchronousTask task, bool isIdle = fal
             return true;
         }
 
+        // This handles executing the task while also notifying the client that the pipeline is
+        // currently busy with a PowerShell task. The extension indicates this with a spinner.
+        private void ExecuteTaskSynchronously(ISynchronousTask task, CancellationToken cancellationToken)
+        {
+            // TODO: Simplify this logic.
+            bool busy = false;
+            if (task is ISynchronousPowerShellTask t
+                && (t.PowerShellExecutionOptions.AddToHistory
+                    || t.PowerShellExecutionOptions.FromRepl))
+            {
+                busy = true;
+                SetBusy(true);
+            }
+            try
+            {
+                task.ExecuteSynchronously(cancellationToken);
+            }
+            finally
+            {
+                if (busy)
+                {
+                    SetBusy(false);
+                }
+            }
+        }
+
         public Task<T> InvokeTaskOnPipelineThreadAsync<T>(SynchronousTask<T> task)
         {
             if (CancelForegroundAndPrepend(task))
@@ -769,8 +797,13 @@ private void RunExecutionLoop(bool isForDebug = false)
                 {
                     try
                     {
-                        task.ExecuteSynchronously(cancellationScope.CancellationToken);
+                        ExecuteTaskSynchronously(task, cancellationScope.CancellationToken);
                     }
+                    // Our flaky extension command test seems to be such because sometimes another
+                    // task gets queued, and since it runs in the foreground it cancels that task.
+                    // Interactively, this happens in the first loop (with DoOneRepl) which catches
+                    // the cancellation exception, but when under test that is a no-op, so  it
+                    // happens in this second loop. Hence we need to catch it here too.
                     catch (OperationCanceledException e)
                     {
                         _logger.LogDebug(e, "Task {Task} was canceled!", task);
@@ -935,19 +968,27 @@ private string InvokeReadLine(CancellationToken cancellationToken)
             }
         }
 
+        // TODO: Should we actually be directly invoking input versus queueing it as a task like everything else?
         private void InvokeInput(string input, CancellationToken cancellationToken)
         {
-            PSCommand command = new PSCommand().AddScript(input, useLocalScope: false);
-            InvokePSCommand(
-                command,
-                new PowerShellExecutionOptions
-                {
-                    AddToHistory = true,
-                    ThrowOnError = false,
-                    WriteOutputToHost = true,
-                    FromRepl = true,
-                },
-                cancellationToken);
+            SetBusy(true);
+            try
+            {
+                InvokePSCommand(
+                    new PSCommand().AddScript(input, useLocalScope: false),
+                    new PowerShellExecutionOptions
+                    {
+                        AddToHistory = true,
+                        ThrowOnError = false,
+                        WriteOutputToHost = true,
+                        FromRepl = true,
+                    },
+                    cancellationToken);
+            }
+            finally
+            {
+                SetBusy(false);
+            }
         }
 
         private void AddRunspaceEventHandlers(Runspace runspace)
@@ -1076,16 +1117,18 @@ private void OnPowerShellIdle(CancellationToken idleCancellationToken)
                 while (!cancellationScope.CancellationToken.IsCancellationRequested
                     && _taskQueue.TryTake(out ISynchronousTask task))
                 {
+                    // Tasks which require the foreground cannot run under this idle handler, so the
+                    // current foreground tasks gets canceled, the new task gets prepended, and this
+                    // handler returns.
                     if (CancelForegroundAndPrepend(task, isIdle: true))
                     {
                         return;
                     }
 
-                    // If we're executing a task, we don't need to run an extra pipeline later for events
-                    // TODO: This may not be a PowerShell task, so ideally we can differentiate that here.
-                    //       For now it's mostly true and an easy assumption to make.
-                    runPipelineForEventProcessing = false;
-                    task.ExecuteSynchronously(cancellationScope.CancellationToken);
+                    // If we're executing a PowerShell task, we don't need to run an extra pipeline
+                    // later for events.
+                    runPipelineForEventProcessing = task is not ISynchronousPowerShellTask;
+                    ExecuteTaskSynchronously(task, cancellationScope.CancellationToken);
                 }
             }