Implement remote debugger waiting

Author: rjmholt

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

@@ -79,7 +79,11 @@ private IReadOnlyList<TResult> ExecuteNormally(CancellationToken cancellationTok
                     cancellationToken.ThrowIfCancellationRequested();
                 }
             }
-            catch (Exception e) when (e is PipelineStoppedException || e is PSRemotingDataStructureException)
+            catch (PipelineStoppedException)
+            {
+                throw new OperationCanceledException();
+            }
+            catch (PSRemotingDataStructureException e)
             {
                 string message = $"Pipeline stopped while executing command:{Environment.NewLine}{Environment.NewLine}{e}";
                 Logger.LogError(message);
@@ -131,7 +135,58 @@ private IReadOnlyList<TResult> ExecuteInDebugger(CancellationToken cancellationT
                     };
             }
 
-            DebuggerCommandResults debuggerResult = _pwsh.Runspace.Debugger.ProcessCommand(_psCommand, outputCollection);
+            DebuggerCommandResults debuggerResult = null;
+            try
+            {
+                debuggerResult = _pwsh.Runspace.Debugger.ProcessCommand(_psCommand, outputCollection);
+
+                if (_executionOptions.PropagateCancellationToCaller)
+                {
+                    cancellationToken.ThrowIfCancellationRequested();
+                }
+            }
+            catch (PipelineStoppedException)
+            {
+                throw new OperationCanceledException();
+            }
+            catch (PSRemotingDataStructureException e)
+            {
+                string message = $"Pipeline stopped while executing command:{Environment.NewLine}{Environment.NewLine}{e}";
+                Logger.LogError(message);
+                throw new ExecutionCanceledException(message, e);
+            }
+            catch (RuntimeException e)
+            {
+                Logger.LogWarning($"Runtime exception occurred while executing command:{Environment.NewLine}{Environment.NewLine}{e}");
+
+                if (!_executionOptions.WriteOutputToHost)
+                {
+                    throw;
+                }
+
+                var errorOutputCollection = new PSDataCollection<PSObject>();
+                errorOutputCollection.DataAdded += (object sender, DataAddedEventArgs args) =>
+                    {
+                        for (int i = args.Index; i < outputCollection.Count; i++)
+                        {
+                            _psHost.UI.WriteLine(outputCollection[i].ToString());
+                        }
+                    };
+
+                var command = new PSCommand()
+                    .AddDebugOutputCommand()
+                    .AddParameter("InputObject", e.ErrorRecord.AsPSObject());
+
+                _pwsh.Runspace.Debugger.ProcessCommand(command, errorOutputCollection);
+            }
+            finally
+            {
+                if (_pwsh.HadErrors)
+                {
+                    _pwsh.Streams.Error.Clear();
+                }
+            }
+
             _psRunspaceContext.ProcessDebuggerResult(debuggerResult);
 
             // Optimisation to save wasted computation if we're going to throw the output away anyway

src/PowerShellEditorServices/Services/PowerShell/PowerShellExecutionService.cs

@@ -649,8 +649,10 @@ private void RunPowerShellLoop(PowerShellFrameType powerShellFrameType)
 
         private void OnDebuggerStopped(object sender, DebuggerStopEventArgs debuggerStopEventArgs)
         {
+            CurrentPowerShell.WaitForRemoteOutputIfNeeded();
             _debuggingContext.OnDebuggerStop(sender, debuggerStopEventArgs);
             PushDebugPowerShell();
+            CurrentPowerShell.ResumeRemoteOutputIfNeeded();
         }
 
         private void SetDebuggerResuming(DebuggerResumeAction resumeAction)

src/PowerShellEditorServices/Services/PowerShell/Utility/PowerShellExtensions.cs

@@ -3,13 +3,36 @@
 using System.ComponentModel;
 using System.Management.Automation;
 using System.Management.Automation.Runspaces;
+using System.Reflection;
+using System.Runtime.CompilerServices;
 using System.Text;
 using SMA = System.Management.Automation;
 
 namespace Microsoft.PowerShell.EditorServices.Services.PowerShell.Utility
 {
     internal static class PowerShellExtensions
     {
+        private static readonly Action<SMA.PowerShell> s_waitForServicingComplete;
+
+        private static readonly Action<SMA.PowerShell> s_suspendIncomingData;
+
+        private static readonly Action<SMA.PowerShell> s_resumeIncomingData;
+
+        static PowerShellExtensions()
+        {
+            s_waitForServicingComplete = (Action<SMA.PowerShell>)Delegate.CreateDelegate(
+                typeof(Action<SMA.PowerShell>),
+                typeof(SMA.PowerShell).GetMethod("WaitForServicingComplete", BindingFlags.Instance | BindingFlags.NonPublic));
+
+            s_suspendIncomingData = (Action<SMA.PowerShell>)Delegate.CreateDelegate(
+                typeof(Action<SMA.PowerShell>),
+                typeof(SMA.PowerShell).GetMethod("SuspendIncomingData", BindingFlags.Instance | BindingFlags.NonPublic));
+
+            s_resumeIncomingData = (Action<SMA.PowerShell>)Delegate.CreateDelegate(
+                typeof(Action<SMA.PowerShell>),
+                typeof(SMA.PowerShell).GetMethod("ResumeIncomingData", BindingFlags.Instance | BindingFlags.NonPublic));
+        }
+
         public static Collection<TResult> InvokeAndClear<TResult>(this SMA.PowerShell pwsh)
         {
             try
@@ -46,6 +69,27 @@ public static void InvokeCommand(this SMA.PowerShell pwsh, PSCommand psCommand)
             pwsh.InvokeAndClear();
         }
 
+        public static void WaitForRemoteOutputIfNeeded(this SMA.PowerShell pwsh)
+        {
+            if (!pwsh.Runspace.RunspaceIsRemote)
+            {
+                return;
+            }
+
+            s_waitForServicingComplete(pwsh);
+            s_suspendIncomingData(pwsh);
+        }
+
+        public static void ResumeRemoteOutputIfNeeded(this SMA.PowerShell pwsh)
+        {
+            if (!pwsh.Runspace.RunspaceIsRemote)
+            {
+                return;
+            }
+
+            s_resumeIncomingData(pwsh);
+        }
+
         public static string GetErrorString(this SMA.PowerShell pwsh)
         {
             var sb = new StringBuilder(capacity: 1024)