Make ExecuteCommandAsync cancellable

This lets us cancel the debugger more reliably.

Author: andyleejordan

src/PowerShellEditorServices/Services/DebugAdapter/Handlers/LaunchAndAttachHandler.cs

@@ -303,11 +303,13 @@ await _powerShellContextService.ExecuteScriptStringAsync(
             string debugRunspaceCmd;
             if (request.RunspaceName != null)
             {
-                IEnumerable<int?> ids = await _powerShellContextService.ExecuteCommandAsync<int?>(new PSCommand()
+                IEnumerable<int?> ids = await _powerShellContextService.ExecuteCommandAsync<int?>(
+                    new PSCommand()
                     .AddCommand("Microsoft.PowerShell.Utility\\Get-Runspace")
                     .AddParameter("Name", request.RunspaceName)
                     .AddCommand("Microsoft.PowerShell.Utility\\Select-Object")
-                    .AddParameter("ExpandProperty", "Id")).ConfigureAwait(false);
+                    .AddParameter("ExpandProperty", "Id"), cancellationToken: cancellationToken).ConfigureAwait(false);
+
                 foreach (var id in ids)
                 {
                     _debugStateService.RunspaceId = id;

src/PowerShellEditorServices/Services/PowerShellContext/Console/ConsoleReadLine.cs

@@ -208,9 +208,8 @@ internal async Task<string> InvokeLegacyReadLineAsync(bool isCommandLine, Cancel
                                 command.AddParameter("CursorColumn", currentCursorIndex);
                                 command.AddParameter("Options", null);
 
-                                var results = await this.powerShellContext
-                                    .ExecuteCommandAsync<CommandCompletion>(command, sendOutputToHost: false, sendErrorToHost: false)
-                                    .ConfigureAwait(false);
+                                var results = await this.powerShellContext.ExecuteCommandAsync<CommandCompletion>(
+                                    command, sendOutputToHost: false, sendErrorToHost: false, cancellationToken).ConfigureAwait(false);
 
                                 currentCompletion = results.FirstOrDefault();
                             }
@@ -327,8 +326,8 @@ internal async Task<string> InvokeLegacyReadLineAsync(bool isCommandLine, Cancel
                             PSCommand command = new PSCommand();
                             command.AddCommand("Get-History");
 
-                            currentHistory = await this.powerShellContext.ExecuteCommandAsync<PSObject>(command, sendOutputToHost: false, sendErrorToHost: false)
-                                .ConfigureAwait(false)
+                            currentHistory = await this.powerShellContext.ExecuteCommandAsync<PSObject>(
+                                command, sendOutputToHost: false, sendErrorToHost: false, cancellationToken).ConfigureAwait(false)
                                 as Collection<PSObject>;
 
                             if (currentHistory != null)

src/PowerShellEditorServices/Services/PowerShellContext/Handlers/ExpandAliasHandler.cs

@@ -69,7 +69,8 @@ function __Expand-Alias {
                 .AddStatement()
                 .AddCommand("__Expand-Alias")
                 .AddArgument(request.Text);
-            var result = await _powerShellContextService.ExecuteCommandAsync<string>(psCommand).ConfigureAwait(false);
+            var result = await _powerShellContextService.ExecuteCommandAsync<string>(
+                psCommand, cancellationToken: cancellationToken).ConfigureAwait(false);
 
             return new ExpandAliasResult
             {

src/PowerShellEditorServices/Services/PowerShellContext/Handlers/GetCommandHandler.cs

@@ -53,7 +53,8 @@ public async Task<List<PSCommandMessage>> Handle(GetCommandParams request, Cance
                 .AddCommand("Microsoft.PowerShell.Utility\\Sort-Object")
                     .AddParameter("Property", "Name");
 
-            IEnumerable<CommandInfo> result = await _powerShellContextService.ExecuteCommandAsync<CommandInfo>(psCommand).ConfigureAwait(false);
+            IEnumerable<CommandInfo> result = await _powerShellContextService.ExecuteCommandAsync<CommandInfo>(
+                psCommand, cancellationToken: cancellationToken).ConfigureAwait(false);
 
             var commandList = new List<PSCommandMessage>();
             if (result != null)

src/PowerShellEditorServices/Services/PowerShellContext/Handlers/PSHostProcessAndRunspaceHandlers.cs

@@ -87,7 +87,8 @@ public async Task<RunspaceResponse[]> Handle(GetRunspaceParams request, Cancella
                 var psCommand = new PSCommand().AddCommand("Microsoft.PowerShell.Utility\\Get-Runspace");
                 var sb = new StringBuilder();
                 // returns (not deserialized) Runspaces. For simpler code, we use PSObject and rely on dynamic later.
-                runspaces = await _powerShellContextService.ExecuteCommandAsync<PSObject>(psCommand, sb).ConfigureAwait(false);
+                runspaces = await _powerShellContextService.ExecuteCommandAsync<PSObject>(
+                    psCommand, sb, cancellationToken: cancellationToken).ConfigureAwait(false);
             }
 
             var runspaceResponses = new List<RunspaceResponse>();

src/PowerShellEditorServices/Services/PowerShellContext/Handlers/ShowHelpHandler.cs

@@ -72,7 +72,8 @@ public async Task<Unit> Handle(ShowHelpParams request, CancellationToken cancell
 
             // TODO: Rather than print the help in the console, we should send the string back
             //       to VSCode to display in a help pop-up (or similar)
-            await _powerShellContextService.ExecuteCommandAsync<PSObject>(checkHelpPSCommand, sendOutputToHost: true).ConfigureAwait(false);
+            await _powerShellContextService.ExecuteCommandAsync<PSObject>(
+                checkHelpPSCommand, sendOutputToHost: true, cancellationToken: cancellationToken).ConfigureAwait(false);
             return Unit.Value;
         }
     }

src/PowerShellEditorServices/Services/PowerShellContext/PowerShellContextService.cs

@@ -519,9 +519,11 @@ public Task<RunspaceHandle> GetRunspaceHandleAsync(CancellationToken cancellatio
         public Task<IEnumerable<TResult>> ExecuteCommandAsync<TResult>(
             PSCommand psCommand,
             bool sendOutputToHost = false,
-            bool sendErrorToHost = true)
+            bool sendErrorToHost = true,
+            CancellationToken cancellationToken = default)
         {
-            return ExecuteCommandAsync<TResult>(psCommand, errorMessages: null, sendOutputToHost, sendErrorToHost);
+            return this.ExecuteCommandAsync<TResult>(
+                psCommand, errorMessages: null, sendOutputToHost, sendErrorToHost, cancellationToken: cancellationToken);
         }
 
         /// <summary>
@@ -549,7 +551,8 @@ public Task<IEnumerable<TResult>> ExecuteCommandAsync<TResult>(
             StringBuilder errorMessages,
             bool sendOutputToHost = false,
             bool sendErrorToHost = true,
-            bool addToHistory = false)
+            bool addToHistory = false,
+            CancellationToken cancellationToken = default)
         {
             return
                 this.ExecuteCommandAsync<TResult>(
@@ -560,7 +563,8 @@ public Task<IEnumerable<TResult>> ExecuteCommandAsync<TResult>(
                         WriteOutputToHost = sendOutputToHost,
                         WriteErrorsToHost = sendErrorToHost,
                         AddToHistory = addToHistory
-                    });
+                    },
+                    cancellationToken);
         }
 
 
@@ -581,7 +585,8 @@ public Task<IEnumerable<TResult>> ExecuteCommandAsync<TResult>(
         public async Task<IEnumerable<TResult>> ExecuteCommandAsync<TResult>(
             PSCommand psCommand,
             StringBuilder errorMessages,
-            ExecutionOptions executionOptions)
+            ExecutionOptions executionOptions,
+            CancellationToken cancellationToken = default)
         {
             Validate.IsNotNull(nameof(psCommand), psCommand);
             Validate.IsNotNull(nameof(executionOptions), executionOptions);
@@ -759,8 +764,10 @@ public async Task<IEnumerable<TResult>> ExecuteCommandAsync<TResult>(
                         return shell.Invoke<TResult>(null, invocationSettings);
                     }
 
-                    // May need a cancellation token here
-                    return await Task.Run<IEnumerable<TResult>>(() => shell.Invoke<TResult>(input: null, invocationSettings), CancellationToken.None).ConfigureAwait(false);
+                    // This is the primary reason that ExecuteCommandAsync takes a CancellationToken
+                    return await Task.Run<IEnumerable<TResult>>(
+                        () => shell.Invoke<TResult>(input: null, invocationSettings), cancellationToken)
+                        .ConfigureAwait(false);
                 }
                 finally
                 {
@@ -872,7 +879,7 @@ public async Task<IEnumerable<TResult>> ExecuteCommandAsync<TResult>(
                         // will exist already so we need to create one and then use it
                         if (runspaceHandle == null)
                         {
-                            runspaceHandle = await this.GetRunspaceHandleAsync().ConfigureAwait(false);
+                            runspaceHandle = await this.GetRunspaceHandleAsync(cancellationToken).ConfigureAwait(false);
                         }
 
                         sessionDetails = this.GetSessionDetailsInRunspace(runspaceHandle.Runspace);
@@ -1938,7 +1945,7 @@ private void PowerShellContext_ExecutionStatusChangedAsync(object sender, Execut
 
         private IEnumerable<TResult> ExecuteCommandInDebugger<TResult>(PSCommand psCommand, bool sendOutputToHost)
         {
-            this.logger.LogTrace($"Attempting to execute command(s)a in the debugger: {GetStringForPSCommand(psCommand)}");
+            this.logger.LogTrace($"Attempting to execute command(s) in the debugger: {GetStringForPSCommand(psCommand)}");
 
             IEnumerable<TResult> output =
                 this.versionSpecificOperations.ExecuteCommandInDebugger<TResult>(

src/PowerShellEditorServices/Services/PowerShellContext/Session/Host/EditorServicesPSHostUserInterface.cs

@@ -731,7 +731,8 @@ private async Task WritePromptStringToHostAsync(CancellationToken cancellationTo
 
             cancellationToken.ThrowIfCancellationRequested();
             string promptString =
-                (await this.powerShellContext.ExecuteCommandAsync<PSObject>(promptCommand, false, false).ConfigureAwait(false))
+                (await this.powerShellContext.ExecuteCommandAsync<PSObject>(
+                    promptCommand, false, false, cancellationToken).ConfigureAwait(false))
                     .Select(pso => pso.BaseObject)
                     .OfType<string>()
                     .FirstOrDefault() ?? "PS> ";

src/PowerShellEditorServices/Services/PowerShellContext/Session/PSReadLinePromptContext.cs

@@ -143,7 +143,8 @@ public async Task<string> InvokeReadLineAsync(bool isCommandLine, CancellationTo
             IEnumerable<string> readLineResults = await _powerShellContext.ExecuteCommandAsync<string>(
                 readLineCommand,
                 errorMessages: null,
-                s_psrlExecutionOptions).ConfigureAwait(false);
+                s_psrlExecutionOptions,
+                cancellationToken).ConfigureAwait(false);
 
             string line = readLineResults.FirstOrDefault();