Route startup exceptions through and ensure startup has completed before proceeding

Author: rjmholt

src/PowerShellEditorServices/Server/PsesDebugServer.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
@@ -28,7 +29,9 @@ internal class PsesDebugServer : IDisposable
 
         private DebugAdapterServer _debugAdapterServer;
 
-        private PowerShellDebugContext _debugContext;
+        private PsesInternalHost _psesHost;
+
+        private bool _startedPses;
 
         protected readonly ILoggerFactory _loggerFactory;
 
@@ -61,8 +64,8 @@ public async Task StartAsync()
             {
                 // We need to let the PowerShell Context Service know that we are in a debug session
                 // so that it doesn't send the powerShell/startDebugger message.
-                _debugContext = ServiceProvider.GetService<PsesInternalHost>().DebugContext;
-                _debugContext.IsDebugServerActive = true;
+                _psesHost = ServiceProvider.GetService<PsesInternalHost>();
+                _psesHost.DebugContext.IsDebugServerActive = true;
 
                 options
                     .WithInput(_inputStream)
@@ -88,8 +91,11 @@ public async Task StartAsync()
                     // The OnInitialize delegate gets run when we first receive the _Initialize_ request:
                     // https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Initialize
                     .OnInitialize(async (server, request, cancellationToken) => {
+                        // We need to make sure the host has been started
+                        _startedPses = !(await _psesHost.TryStartAsync(new HostStartOptions(), CancellationToken.None).ConfigureAwait(false));
+
                         // Ensure the debugger mode is set correctly - this is required for remote debugging to work
-                        _debugContext.EnableDebugMode();
+                        _psesHost.DebugContext.EnableDebugMode();
 
                         var breakpointService = server.GetService<BreakpointService>();
                         // Clear any existing breakpoints before proceeding
@@ -115,7 +121,7 @@ public void Dispose()
             // Note that the lifetime of the DebugContext is longer than the debug server;
             // It represents the debugger on the PowerShell process we're in,
             // while a new debug server is spun up for every debugging session
-            _debugContext.IsDebugServerActive = false;
+            _psesHost.DebugContext.IsDebugServerActive = false;
             _debugAdapterServer.Dispose();
             _inputStream.Dispose();
             _outputStream.Dispose();
@@ -126,6 +132,13 @@ public void Dispose()
         public async Task WaitForShutdown()
         {
             await _serverStopped.Task.ConfigureAwait(false);
+
+            // If we started the host, we need to ensure any errors are marshalled back to us like this
+            if (_startedPses)
+            {
+                _psesHost.TriggerShutdown();
+                await _psesHost.Shutdown.ConfigureAwait(false);
+            }
         }
 
         #region Events

src/PowerShellEditorServices/Server/PsesLanguageServer.cs

@@ -10,6 +10,7 @@
 using Microsoft.PowerShell.EditorServices.Hosting;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Services.Extension;
+using Microsoft.PowerShell.EditorServices.Services.PowerShell.Host;
 using Microsoft.PowerShell.EditorServices.Services.Template;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
 using OmniSharp.Extensions.LanguageServer.Server;
@@ -32,6 +33,8 @@ internal class PsesLanguageServer
         private readonly HostStartupInfo _hostDetails;
         private readonly TaskCompletionSource<bool> _serverStart;
 
+        private PsesInternalHost _psesHost;
+
         /// <summary>
         /// Create a new language server instance.
         /// </summary>
@@ -75,8 +78,12 @@ public async Task StartAsync()
                 options
                     .WithInput(_inputStream)
                     .WithOutput(_outputStream)
-                    .WithServices(serviceCollection => serviceCollection
-                        .AddPsesLanguageServices(_hostDetails)) // NOTE: This adds a lot of services!
+                    .WithServices(serviceCollection =>
+                    {
+
+                        // NOTE: This adds a lot of services!
+                        serviceCollection.AddPsesLanguageServices(_hostDetails);
+                    })
                     .ConfigureLogging(builder => builder
                         .AddSerilog(Log.Logger) // TODO: Set dispose to true?
                         .AddLanguageProtocolLogging()
@@ -116,6 +123,8 @@ public async Task StartAsync()
 
                             IServiceProvider serviceProvider = languageServer.Services;
 
+                            _psesHost = serviceProvider.GetService<PsesInternalHost>();
+
                             var workspaceService = serviceProvider.GetService<WorkspaceService>();
 
                             // Grab the workspace path from the parameters
@@ -150,6 +159,10 @@ public async Task WaitForShutdown()
             Log.Logger.Debug("Shutting down OmniSharp Language Server");
             await _serverStart.Task.ConfigureAwait(false);
             await LanguageServer.WaitForExit.ConfigureAwait(false);
+
+            // Doing this means we're able to route through any exceptions experienced on the pipeline thread
+            _psesHost.TriggerShutdown();
+            await _psesHost.Shutdown.ConfigureAwait(false);
         }
     }
 }

src/PowerShellEditorServices/Services/PowerShell/Debugging/PowerShellDebugContext.cs

@@ -107,7 +107,12 @@ public void StepOver()
         public void SetDebugResuming(DebuggerResumeAction debuggerResumeAction)
         {
             _psesHost.SetExit();
-            LastStopEventArgs.ResumeAction = debuggerResumeAction;
+
+            if (LastStopEventArgs is not null)
+            {
+                LastStopEventArgs.ResumeAction = debuggerResumeAction;
+            }
+
             // We need to tell whatever is happening right now in the debug prompt to wrap up so we can continue
             _psesHost.CancelCurrentTask();
         }

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

@@ -57,10 +57,16 @@ internal class PsesInternalHost : PSHost, IHostSupportsInteractiveSession, IRuns
 
         private readonly IdempotentLatch _isRunningLatch = new();
 
+        private readonly TaskCompletionSource<bool> _started = new();
+
+        private readonly TaskCompletionSource<bool> _stopped = new();
+
         private EngineIntrinsics _mainRunspaceEngineIntrinsics;
 
         private bool _shouldExit = false;
 
+        private int _shuttingDown = 0;
+
         private string _localComputerName;
 
         private ConsoleKeyInfo? _lastKey;
@@ -128,12 +134,16 @@ public PsesInternalHost(
 
         public string InitialWorkingDirectory { get; private set; }
 
+        public Task Shutdown => _stopped.Task;
+
         IRunspaceInfo IRunspaceContext.CurrentRunspace => CurrentRunspace;
 
         private PowerShellContextFrame CurrentFrame => _psFrameStack.Peek();
 
         public event Action<object, RunspaceChangedEventArgs> RunspaceChanged;
 
+        private bool ShouldExitExecutionLoop => _shouldExit || _shuttingDown != 0;
+
         public override void EnterNestedPrompt()
         {
             PushPowerShellAndRunLoop(CreateNestedPowerShell(CurrentRunspace), PowerShellFrameType.Nested);
@@ -172,13 +182,22 @@ public override void SetShouldExit(int exitCode)
             SetExit();
         }
 
-        public async Task StartAsync(HostStartOptions startOptions, CancellationToken cancellationToken)
+        /// <summary>
+        /// Try to start the PowerShell loop in the host.
+        /// If the host is already started, this is idempotent.
+        /// Returns when the host is in a valid initialized state.
+        /// </summary>
+        /// <param name="startOptions">Options to configure host startup.</param>
+        /// <param name="cancellationToken">A token to cancel startup.</param>
+        /// <returns>A task that resolves when the host has finished startup, with the value true if the caller started the host, and false otherwise.</returns>
+        public async Task<bool> TryStartAsync(HostStartOptions startOptions, CancellationToken cancellationToken)
         {
             _logger.LogInformation("Host starting");
             if (!_isRunningLatch.TryEnter())
             {
-                _logger.LogDebug("Host start requested after already started");
-                return;
+                _logger.LogDebug("Host start requested after already started.");
+                await _started.Task.ConfigureAwait(false);
+                return false;
             }
 
             _pipelineThread.Start();
@@ -198,6 +217,15 @@ await ExecuteDelegateAsync(
             {
                 await SetInitialWorkingDirectoryAsync(startOptions.InitialWorkingDirectory, CancellationToken.None).ConfigureAwait(false);
             }
+
+            await _started.Task.ConfigureAwait(false);
+            return true;
+        }
+
+        public void TriggerShutdown()
+        {
+            Interlocked.Exchange(ref _shuttingDown, 1);
+            _cancellationContext.CancelCurrentTaskStack();
         }
 
         public void SetExit()
@@ -349,11 +377,19 @@ public Task SetInitialWorkingDirectoryAsync(string path, CancellationToken cance
 
         private void Run()
         {
-            (PowerShell pwsh, RunspaceInfo localRunspaceInfo, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShellSession();
-            _mainRunspaceEngineIntrinsics = engineIntrinsics;
-            _localComputerName = localRunspaceInfo.SessionDetails.ComputerName;
-            _runspaceStack.Push(new RunspaceFrame(pwsh.Runspace, localRunspaceInfo));
-            PushPowerShellAndRunLoop(pwsh, PowerShellFrameType.Normal, localRunspaceInfo);
+            try
+            {
+                (PowerShell pwsh, RunspaceInfo localRunspaceInfo, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShellSession();
+                _mainRunspaceEngineIntrinsics = engineIntrinsics;
+                _localComputerName = localRunspaceInfo.SessionDetails.ComputerName;
+                _runspaceStack.Push(new RunspaceFrame(pwsh.Runspace, localRunspaceInfo));
+                PushPowerShellAndRunLoop(pwsh, PowerShellFrameType.Normal, localRunspaceInfo);
+            }
+            catch (Exception e)
+            {
+                _started.TrySetException(e);
+                _stopped.TrySetException(e);
+            }
         }
 
         private (PowerShell, RunspaceInfo, EngineIntrinsics) CreateInitialPowerShellSession()
@@ -465,25 +501,40 @@ private void PopPowerShell(RunspaceChangeAction runspaceChangeAction = RunspaceC
 
         private void RunTopLevelExecutionLoop()
         {
-            // Make sure we execute any startup tasks first
-            while (_taskQueue.TryTake(out ISynchronousTask task))
+            try
             {
-                task.ExecuteSynchronously(CancellationToken.None);
-            }
+                // Make sure we execute any startup tasks first
+                while (_taskQueue.TryTake(out ISynchronousTask task))
+                {
+                    task.ExecuteSynchronously(CancellationToken.None);
+                }
 
-            if (_hostInfo.ConsoleReplEnabled)
-            {
-                RunExecutionLoop();
+                // Signal that we are ready for outside services to use
+                _started.TrySetResult(true);
+
+                if (_hostInfo.ConsoleReplEnabled)
+                {
+                    RunExecutionLoop();
+                }
+                else
+                {
+                    RunNoPromptExecutionLoop();
+                }
             }
-            else
+            catch (Exception e)
             {
-                RunNoPromptExecutionLoop();
+                _logger.LogError(e, "PSES pipeline thread loop experienced an unexpected top-level exception");
+                _stopped.TrySetException(e);
+                return;
             }
+
+            _logger.LogInformation("PSES pipeline thread loop shutting down");
+            _stopped.SetResult(true);
         }
 
         private void RunNoPromptExecutionLoop()
         {
-            while (!_shouldExit)
+            while (!ShouldExitExecutionLoop)
             {
                 using (CancellationScope cancellationScope = _cancellationContext.EnterScope(isIdleScope: false))
                 {
@@ -521,13 +572,13 @@ private void RunDebugExecutionLoop()
 
         private void RunExecutionLoop()
         {
-            while (!_shouldExit)
+            while (!ShouldExitExecutionLoop)
             {
                 using (CancellationScope cancellationScope = _cancellationContext.EnterScope(isIdleScope: false))
                 {
                     DoOneRepl(cancellationScope.CancellationToken);
 
-                    while (!_shouldExit
+                    while (!ShouldExitExecutionLoop
                         && !cancellationScope.CancellationToken.IsCancellationRequested
                         && _taskQueue.TryTake(out ISynchronousTask task))
                     {
@@ -806,7 +857,7 @@ private void OnBreakpointUpdated(object sender, BreakpointUpdatedEventArgs break
 
         private void OnRunspaceStateChanged(object sender, RunspaceStateEventArgs runspaceStateEventArgs)
         {
-            if (!_shouldExit && !_resettingRunspace && !runspaceStateEventArgs.RunspaceStateInfo.IsUsable())
+            if (!ShouldExitExecutionLoop && !_resettingRunspace && !runspaceStateEventArgs.RunspaceStateInfo.IsUsable())
             {
                 _resettingRunspace = true;
                 PopOrReinitializeRunspaceAsync().HandleErrorsAsync(_logger);

src/PowerShellEditorServices/Services/Workspace/Handlers/ConfigurationHandler.cs

@@ -92,7 +92,7 @@ public override async Task<Unit> Handle(DidChangeConfigurationParams request, Ca
                     _logger.LogTrace("Loading profiles...");
                 }
 
-                await _psesHost.StartAsync(new HostStartOptions(), CancellationToken.None).ConfigureAwait(false);
+                await _psesHost.TryStartAsync(new HostStartOptions { LoadProfiles = loadProfiles }, CancellationToken.None).ConfigureAwait(false);
 
                 if (loadProfiles)
                 {

test/PowerShellEditorServices.Test.E2E/DebugAdapterProtocolMessageTests.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Reflection;
@@ -44,6 +45,13 @@ public async Task InitializeAsync()
             await _psesProcess.Start().ConfigureAwait(false);
 
             var initialized = new TaskCompletionSource<bool>();
+
+            _psesProcess.ProcessExited += (sender, args) =>
+            {
+                initialized.TrySetException(new ProcessExitedException("Initialization failed due to process failure", args.ExitCode, args.ErrorMessage));
+                Started.TrySetException(new ProcessExitedException("Startup failed due to process failure", args.ExitCode, args.ErrorMessage));
+            };
+
             PsesDebugAdapterClient = DebugAdapterClient.Create(options =>
             {
                 options
@@ -61,6 +69,11 @@ public async Task InitializeAsync()
                         initialized.SetResult(true);
                         return Task.CompletedTask;
                     });
+
+                options.OnUnhandledException = (exception) => {
+                    initialized.SetException(exception);
+                    Started.SetException(exception);
+                };
             });
 
             // PSES follows the following flow:

test/PowerShellEditorServices.Test.E2E/LSPTestsFixures.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Reflection;
@@ -44,6 +45,7 @@ public async Task InitializeAsync()
             var factory = new LoggerFactory();
             _psesProcess = new PsesStdioProcess(factory, IsDebugAdapterTests);
             await _psesProcess.Start().ConfigureAwait(false);
+            //Debugger.Launch();
 
             Diagnostics = new List<Diagnostic>();
             TelemetryEvents = new List<PsesTelemetryEvent>();

test/PowerShellEditorServices.Test.E2E/LanguageServerProtocolMessageTests.cs

@@ -41,6 +41,7 @@ public class LanguageServerProtocolMessageTests : IClassFixture<LSPTestsFixture>
         private readonly List<Diagnostic> Diagnostics;
         private readonly List<PsesTelemetryEvent> TelemetryEvents;
         private readonly string PwshExe;
+        private readonly LSPTestsFixture _fixture;
 
         public LanguageServerProtocolMessageTests(ITestOutputHelper output, LSPTestsFixture data)
         {
@@ -50,6 +51,7 @@ public LanguageServerProtocolMessageTests(ITestOutputHelper output, LSPTestsFixt
             Diagnostics.Clear();
             TelemetryEvents = data.TelemetryEvents;
             TelemetryEvents.Clear();
+            _fixture = data;
 
             PwshExe = PsesStdioProcess.PwshExe;
         }