Fix debugger hang when you stop execution while input prompt is active

This change fixes an issue where the debugger will hang when the user
tries to stop execution while there is an active input or choice prompt.
The solution is to use a more robust method for cancelling the prompt
task which doesn't depend on the Console.ReadKey() call to return.

Fixes PowerShell#428.

Author: daviwil

src/PowerShellEditorServices/Console/ChoicePromptHandler.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Management.Automation;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -31,7 +32,7 @@ public enum PromptStyle
     }
 
     /// <summary>
-    /// Provides a base implementation for IPromptHandler classes 
+    /// Provides a base implementation for IPromptHandler classes
     /// that present the user a set of options from which a selection
     /// should be made.
     /// </summary>
@@ -41,6 +42,8 @@ public abstract class ChoicePromptHandler : PromptHandler
 
         private CancellationTokenSource promptCancellationTokenSource =
             new CancellationTokenSource();
+        private TaskCompletionSource<Dictionary<string, object>> cancelTask =
+            new TaskCompletionSource<Dictionary<string, object>>();
 
         #endregion
 
@@ -98,7 +101,7 @@ public abstract class ChoicePromptHandler : PromptHandler
         /// A Task instance that can be monitored for completion to get
         /// the user's choice.
         /// </returns>
-        public Task<int> PromptForChoice(
+        public async Task<int> PromptForChoice(
             string promptCaption,
             string promptMessage,
             ChoiceDetails[] choices,
@@ -120,7 +123,7 @@ public Task<int> PromptForChoice(
             cancellationToken.Register(this.CancelPrompt, true);
 
             // Convert the int[] result to int
-            return
+            return await this.WaitForTask(
                 this.StartPromptLoop(this.promptCancellationTokenSource.Token)
                     .ContinueWith(
                         task =>
@@ -135,7 +138,7 @@ public Task<int> PromptForChoice(
                             }
 
                             return this.GetSingleResult(task.Result);
-                        });
+                        }));
         }
 
         /// <summary>
@@ -161,7 +164,7 @@ public Task<int> PromptForChoice(
         /// A Task instance that can be monitored for completion to get
         /// the user's choices.
         /// </returns>
-        public Task<int[]> PromptForChoice(
+        public async Task<int[]> PromptForChoice(
             string promptCaption,
             string promptMessage,
             ChoiceDetails[] choices,
@@ -179,7 +182,24 @@ public Task<int[]> PromptForChoice(
             // Cancel the TaskCompletionSource if the caller cancels the task
             cancellationToken.Register(this.CancelPrompt, true);
 
-            return this.StartPromptLoop(this.promptCancellationTokenSource.Token);
+            return await this.WaitForTask(
+                this.StartPromptLoop(
+                    this.promptCancellationTokenSource.Token));
+        }
+
+        private async Task<T> WaitForTask<T>(Task<T> taskToWait)
+        {
+            Task finishedTask =
+                await Task.WhenAny(
+                    this.cancelTask.Task,
+                    taskToWait);
+
+            if (this.cancelTask.Task.IsCanceled)
+            {
+                throw new PipelineStoppedException();
+            }
+
+            return taskToWait.Result;
         }
 
         private async Task<int[]> StartPromptLoop(
@@ -232,7 +252,7 @@ private async Task<int[]> StartPromptLoop(
         /// </summary>
         /// <param name="responseString">The string representing the user's response.</param>
         /// <returns>
-        /// True if the prompt is complete, false if the prompt is 
+        /// True if the prompt is complete, false if the prompt is
         /// still waiting for a valid response.
         /// </returns>
         protected virtual int[] HandleResponse(string responseString)
@@ -280,6 +300,7 @@ protected override void OnPromptCancelled()
         {
             // Cancel the prompt task
             this.promptCancellationTokenSource.Cancel();
+            this.cancelTask.TrySetCanceled();
         }
 
         #endregion

src/PowerShellEditorServices/Console/ConsoleService.cs

@@ -447,11 +447,21 @@ private void PowerShellContext_ExecutionStatusChanged(object sender, ExecutionSt
         {
             if (this.EnableConsoleRepl)
             {
-                // Any command which writes output to the host will affect
-                // the display of the prompt
-                if (eventArgs.ExecutionOptions.WriteOutputToHost ||
+                if (eventArgs.ExecutionStatus == ExecutionStatus.Aborted)
+                {
+                    // When aborted, cancel any lingering prompts
+                    if (this.activePromptHandler != null)
+                    {
+                        this.activePromptHandler.CancelPrompt();
+                        this.WriteOutput(string.Empty);
+                    }
+                }
+                else if (
+                    eventArgs.ExecutionOptions.WriteOutputToHost ||
                     eventArgs.ExecutionOptions.InterruptCommandPrompt)
                 {
+                    // Any command which writes output to the host will affect
+                    // the display of the prompt
                     if (eventArgs.ExecutionStatus != ExecutionStatus.Running)
                     {
                         // Execution has completed, start the input prompt
@@ -461,6 +471,7 @@ private void PowerShellContext_ExecutionStatusChanged(object sender, ExecutionSt
                     {
                         // A new command was started, cancel the input prompt
                         this.CancelReadLoop();
+                        this.WriteOutput(string.Empty);
                     }
                 }
                 else if (
@@ -469,6 +480,7 @@ private void PowerShellContext_ExecutionStatusChanged(object sender, ExecutionSt
                      eventArgs.HadErrors))
                 {
                     this.CancelReadLoop();
+                    this.WriteOutput(string.Empty);
                     this.StartReadLoop();
                 }
             }

src/PowerShellEditorServices/Console/InputPromptHandler.cs

@@ -14,7 +14,7 @@
 namespace Microsoft.PowerShell.EditorServices.Console
 {
     /// <summary>
-    /// Provides a base implementation for IPromptHandler classes 
+    /// Provides a base implementation for IPromptHandler classes
     /// that present the user a set of fields for which values
     /// should be entered.
     /// </summary>
@@ -26,6 +26,8 @@ public abstract class InputPromptHandler : PromptHandler
         private FieldDetails currentField;
         private CancellationTokenSource promptCancellationTokenSource =
             new CancellationTokenSource();
+        private TaskCompletionSource<Dictionary<string, object>> cancelTask =
+            new TaskCompletionSource<Dictionary<string, object>>();
 
         #endregion
 
@@ -57,7 +59,7 @@ public Task<string> PromptForInput(
                     new FieldDetails[] { new FieldDetails("", "", typeof(string), false, "") },
                     cancellationToken);
 
-            return 
+            return
                 innerTask.ContinueWith<string>(
                     task =>
                     {
@@ -69,7 +71,7 @@ public Task<string> PromptForInput(
                         {
                             throw new TaskCanceledException(task);
                         }
-                        
+
                         // Return the value of the sole field
                         return (string)task.Result[""];
                     });
@@ -95,7 +97,7 @@ public Task<string> PromptForInput(
         /// A Task instance that can be monitored for completion to get
         /// the user's input.
         /// </returns>
-        public Task<Dictionary<string, object>> PromptForInput(
+        public async Task<Dictionary<string, object>> PromptForInput(
             string promptCaption,
             string promptMessage,
             FieldDetails[] fields,
@@ -108,7 +110,20 @@ public Task<Dictionary<string, object>> PromptForInput(
 
             this.ShowPromptMessage(promptCaption, promptMessage);
 
-            return this.StartPromptLoop(this.promptCancellationTokenSource.Token);
+            Task<Dictionary<string, object>> promptTask =
+                this.StartPromptLoop(this.promptCancellationTokenSource.Token);
+
+            Task finishedTask =
+                await Task.WhenAny(
+                    cancelTask.Task,
+                    promptTask);
+
+            if (this.cancelTask.Task.IsCanceled)
+            {
+                throw new PipelineStoppedException();
+            }
+
+            return promptTask.Result;
         }
 
         /// <summary>
@@ -128,7 +143,7 @@ public Task<SecureString> PromptForSecureInput(
                     new FieldDetails[] { new FieldDetails("", "", typeof(SecureString), false, "") },
                     cancellationToken);
 
-            return 
+            return
                 innerTask.ContinueWith(
                     task =>
                     {
@@ -140,7 +155,7 @@ public Task<SecureString> PromptForSecureInput(
                         {
                             throw new TaskCanceledException(task);
                         }
-                        
+
                         // Return the value of the sole field
                         return (SecureString)task.Result?[""];
                     });
@@ -153,6 +168,7 @@ protected override void OnPromptCancelled()
         {
             // Cancel the prompt task
             this.promptCancellationTokenSource.Cancel();
+            this.cancelTask.TrySetCanceled();
         }
 
         #endregion

src/PowerShellEditorServices/Session/PowerShellContext.cs

@@ -870,6 +870,11 @@ public void AbortExecution()
                 this.powerShell.BeginStop(null, null);
 
                 this.SessionState = PowerShellContextState.Aborting;
+
+                this.OnExecutionStatusChanged(
+                    ExecutionStatus.Aborted,
+                    null,
+                    false);
             }
             else
             {