PsesLanguageServer.cs

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using System.IO;
using System.Threading.Tasks;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Microsoft.PowerShell.EditorServices.Handlers;
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 Newtonsoft.Json.Linq;
using OmniSharp.Extensions.LanguageServer.Protocol.Models;
using OmniSharp.Extensions.LanguageServer.Protocol.Server;
using OmniSharp.Extensions.LanguageServer.Server;
using Serilog;

namespace Microsoft.PowerShell.EditorServices.Server
{
    /// <summary>
    /// Server runner class for handling LSP messages for Editor Services.
    /// </summary>
    internal class PsesLanguageServer
    {
        internal ILoggerFactory LoggerFactory { get; }
        internal ILanguageServer LanguageServer { get; private set; }
        private readonly LogLevel _minimumLogLevel;
        private readonly Stream _inputStream;
        private readonly Stream _outputStream;
        private readonly HostStartupInfo _hostDetails;
        private readonly TaskCompletionSource<bool> _serverStart;
        private PsesInternalHost _psesHost;

        /// <summary>
        /// Create a new language server instance.
        /// </summary>
        /// <remarks>
        /// This class is only ever instantiated via <see
        /// cref="EditorServicesServerFactory.CreateLanguageServer"/>. It is essentially a
        /// singleton. The factory hides the logger.
        /// </remarks>
        /// <param name="factory">Factory to create loggers with.</param>
        /// <param name="inputStream">Protocol transport input stream.</param>
        /// <param name="outputStream">Protocol transport output stream.</param>
        /// <param name="hostStartupInfo">Host configuration to instantiate the server and services
        /// with.</param>
        public PsesLanguageServer(
            ILoggerFactory factory,
            Stream inputStream,
            Stream outputStream,
            HostStartupInfo hostStartupInfo)
        {
            LoggerFactory = factory;
            _minimumLogLevel = (LogLevel)hostStartupInfo.LogLevel;
            _inputStream = inputStream;
            _outputStream = outputStream;
            _hostDetails = hostStartupInfo;
            _serverStart = new TaskCompletionSource<bool>();
        }

        /// <summary>
        /// Start the server listening for input.
        /// </summary>
        /// <remarks>
        /// For the services (including the <see cref="PowerShellContextService">
        /// context wrapper around PowerShell itself) see <see
        /// cref="PsesServiceCollectionExtensions.AddPsesLanguageServices"/>.
        /// </remarks>
        /// <returns>A task that completes when the server is ready and listening.</returns>
        public async Task StartAsync()
        {
            LanguageServer = await OmniSharp.Extensions.LanguageServer.Server.LanguageServer.From(options =>
            {
                options
                    .WithInput(_inputStream)
                    .WithOutput(_outputStream)
                    .WithServices(serviceCollection =>
                    {
                        // NOTE: This adds a lot of services!
                        serviceCollection.AddPsesLanguageServices(_hostDetails);
                    })
                    .ConfigureLogging(builder => builder
                        .AddSerilog(Log.Logger) // TODO: Set dispose to true?
                        .AddLanguageProtocolLogging()
                        .SetMinimumLevel(_minimumLogLevel))
                    // TODO: Consider replacing all WithHandler with AddSingleton
                    .WithHandler<PsesWorkspaceSymbolsHandler>()
                    .WithHandler<PsesTextDocumentHandler>()
                    .WithHandler<GetVersionHandler>()
                    .WithHandler<PsesConfigurationHandler>()
                    .WithHandler<PsesFoldingRangeHandler>()
                    .WithHandler<PsesDocumentFormattingHandler>()
                    .WithHandler<PsesDocumentRangeFormattingHandler>()
                    .WithHandler<PsesReferencesHandler>()
                    .WithHandler<PsesDocumentSymbolHandler>()
                    .WithHandler<PsesDocumentHighlightHandler>()
                    .WithHandler<PSHostProcessAndRunspaceHandlers>()
                    .WithHandler<PsesCodeLensHandlers>()
                    .WithHandler<PsesCodeActionHandler>()
                    .WithHandler<InvokeExtensionCommandHandler>()
                    .WithHandler<PsesCompletionHandler>()
                    .WithHandler<PsesHoverHandler>()
                    .WithHandler<PsesSignatureHelpHandler>()
                    .WithHandler<PsesDefinitionHandler>()
                    .WithHandler<TemplateHandlers>()
                    .WithHandler<GetCommentHelpHandler>()
                    .WithHandler<EvaluateHandler>()
                    .WithHandler<GetCommandHandler>()
                    .WithHandler<ShowHelpHandler>()
                    .WithHandler<ExpandAliasHandler>()
                    .WithHandler<PsesSemanticTokensHandler>()
                    // NOTE: The OnInitialize delegate gets run when we first receive the
                    // _Initialize_ request:
                    // https://microsoft.github.io/language-server-protocol/specifications/specification-current/#initialize
                    .OnInitialize(
                        (languageServer, initializeParams, cancellationToken) =>
                        {
                            // Set the workspace path from the parameters.
                            WorkspaceService workspaceService = languageServer.Services.GetService<WorkspaceService>();
                            if (initializeParams.WorkspaceFolders is not null)
                            {
                                // TODO: Support multi-workspace.
                                foreach (WorkspaceFolder workspaceFolder in initializeParams.WorkspaceFolders)
                                {
                                    workspaceService.WorkspacePath = workspaceFolder.Uri.GetFileSystemPath();
                                    break;
                                }
                            }

                            // Parse initialization options.
                            JObject initializationOptions = initializeParams.InitializationOptions as JObject;
                            HostStartOptions hostStartOptions = new()
                            {
                                LoadProfiles = initializationOptions?.GetValue("EnableProfileLoading")?.Value<bool>() ?? false,
                                // TODO: Consider deprecating the setting which sets this and
                                // instead use WorkspacePath exclusively.
                                InitialWorkingDirectory = initializationOptions?.GetValue("InitialWorkingDirectory")?.Value<string>() ?? workspaceService.WorkspacePath
                            };

                            _psesHost = languageServer.Services.GetService<PsesInternalHost>();
                            return _psesHost.TryStartAsync(hostStartOptions, cancellationToken);
                        });
            }).ConfigureAwait(false);

            _serverStart.SetResult(true);
        }

        /// <summary>
        /// Get a task that completes when the server is shut down.
        /// </summary>
        /// <returns>A task that completes when the server is shut down.</returns>
        public async Task WaitForShutdown()
        {
            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);
        }
    }
}