Merge pull request #121 from PowerShell/daviwil/input-prompts

Add text input prompt support

Author: daviwil

src/PowerShellEditorServices.Channel.WebSocket/WebsocketServerChannel.cs

@@ -110,7 +110,7 @@ protected EditorServiceWebSocketConnection()
             Channel = new WebSocketServerChannel(this);
         }
 
-        protected ProtocolServer Server { get; set; }
+        protected ProtocolEndpoint Server { get; set; }
 
         protected WebSocketServerChannel Channel { get; private set; }
 

src/PowerShellEditorServices.Host/Program.cs

@@ -3,6 +3,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using Microsoft.PowerShell.EditorServices.Protocol.MessageProtocol;
 using Microsoft.PowerShell.EditorServices.Protocol.Server;
 using Microsoft.PowerShell.EditorServices.Utility;
 using System;
@@ -41,7 +42,7 @@ static void Main(string[] args)
             // Catch unhandled exceptions for logging purposes
             AppDomain.CurrentDomain.UnhandledException += CurrentDomain_UnhandledException;
 
-            ProtocolServer server = null;
+            ProtocolEndpoint server = null;
             if (runDebugAdapter)
             {
                 logPath = logPath ?? "DebugAdapter.log";
@@ -105,7 +106,7 @@ static void Main(string[] args)
                     fileVersionInfo.FileVersion));
 
             // Start the server
-            server.Start();
+            server.Start().Wait();
             Logger.Write(LogLevel.Normal, "PowerShell Editor Services Host started!");
 
             // Wait for the server to finish

src/PowerShellEditorServices.Protocol/Client/DebugAdapterClientBase.cs

@@ -10,7 +10,7 @@
 
 namespace Microsoft.PowerShell.EditorServices.Protocol.Client
 {
-    public class DebugAdapterClient : ProtocolClient
+    public class DebugAdapterClient : ProtocolEndpoint
     {
         public DebugAdapterClient(ChannelBase clientChannel)
             : base(clientChannel, MessageProtocolType.DebugAdapter)

src/PowerShellEditorServices.Protocol/Client/LanguageClientBase.cs

@@ -13,7 +13,7 @@ namespace Microsoft.PowerShell.EditorServices.Protocol.Client
     /// <summary>
     /// Provides a base implementation for language server clients.
     /// </summary>
-    public abstract class LanguageClientBase : ProtocolClient
+    public abstract class LanguageClientBase : ProtocolEndpoint
     {
         /// <summary>
         /// Initializes an instance of the language client using the

src/PowerShellEditorServices.Protocol/Server/IEventWriter.cs renamed to src/PowerShellEditorServices.Protocol/MessageProtocol/IMessageSender.cs

@@ -3,16 +3,20 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
-using Microsoft.PowerShell.EditorServices.Protocol.MessageProtocol;
 using System.Threading.Tasks;
 
-namespace Microsoft.PowerShell.EditorServices.Protocol.Server
+namespace Microsoft.PowerShell.EditorServices.Protocol.MessageProtocol
 {
-    internal interface IEventWriter
+    internal interface IMessageSender
     {
         Task SendEvent<TParams>(
             EventType<TParams> eventType,
             TParams eventParams);
+
+        Task<TResult> SendRequest<TParams, TResult>(
+            RequestType<TParams, TResult> requestType,
+            TParams requestParams,
+            bool waitForResponse);
     }
 }
 

src/PowerShellEditorServices.Protocol/MessageProtocol/MessageDispatcher.cs

@@ -234,12 +234,14 @@ protected async Task DispatchMessage(
             Message messageToDispatch, 
             MessageWriter messageWriter)
         {
+            Task handlerToAwait = null;
+
             if (messageToDispatch.MessageType == MessageType.Request)
             {
                 Func<Message, MessageWriter, Task> requestHandler = null;
                 if (this.requestHandlers.TryGetValue(messageToDispatch.Method, out requestHandler))
                 {
-                    await requestHandler(messageToDispatch, messageWriter);
+                    handlerToAwait = requestHandler(messageToDispatch, messageWriter);
                 }
                 else
                 {
@@ -258,7 +260,7 @@ protected async Task DispatchMessage(
                 Func<Message, MessageWriter, Task> eventHandler = null;
                 if (this.eventHandlers.TryGetValue(messageToDispatch.Method, out eventHandler))
                 {
-                    await eventHandler(messageToDispatch, messageWriter);
+                    handlerToAwait = eventHandler(messageToDispatch, messageWriter);
                 }
                 else
                 {
@@ -269,6 +271,28 @@ protected async Task DispatchMessage(
             {
                 // TODO: Return message not supported
             }
+
+            if (handlerToAwait != null)
+            {
+                try
+                {
+                    await handlerToAwait;
+                }
+                catch (TaskCanceledException)
+                {
+                    // Some tasks may be cancelled due to legitimate
+                    // timeouts so don't let those exceptions go higher.
+                }
+                catch (AggregateException e)
+                {
+                    if (!(e.InnerExceptions[0] is TaskCanceledException))
+                    {
+                        // Cancelled tasks aren't a problem, so rethrow
+                        // anything that isn't a TaskCanceledException
+                        throw e;
+                    }
+                }
+            }
         }
 
         private void OnListenTaskCompleted(Task listenTask)

src/PowerShellEditorServices.Protocol/Client/ProtocolClient.cs renamed to src/PowerShellEditorServices.Protocol/MessageProtocol/ProtocolEndpoint.cs

@@ -11,31 +11,41 @@
 using System.Threading;
 using System.Threading.Tasks;
 
-namespace Microsoft.PowerShell.EditorServices.Protocol.Client
+namespace Microsoft.PowerShell.EditorServices.Protocol.MessageProtocol
 {
-    public class ProtocolClient
+    /// <summary>
+    /// Provides behavior for a client or server endpoint that
+    /// communicates using the specified protocol.
+    /// </summary>
+    public class ProtocolEndpoint : IMessageSender
     {
         private bool isStarted;
         private int currentMessageId;
-        private ChannelBase clientChannel;
+        private ChannelBase protocolChannel;
         private MessageProtocolType messageProtocolType;
+        private TaskCompletionSource<bool> endpointExitedTask;
         private SynchronizationContext originalSynchronizationContext;
 
         private Dictionary<string, TaskCompletionSource<Message>> pendingRequests =
             new Dictionary<string, TaskCompletionSource<Message>>();
 
         /// <summary>
-        /// Initializes an instance of the protocol client using the
+        /// Initializes an instance of the protocol server using the
         /// specified channel for communication.
         /// </summary>
-        /// <param name="clientChannel">The channel to use for communication with the server.</param>
-        /// <param name="messageProtocolType">The type of message protocol used by the server.</param>
-        public ProtocolClient(
-            ChannelBase clientChannel,
+        /// <param name="protocolChannel">
+        /// The channel to use for communication with the connected endpoint.
+        /// </param>
+        /// <param name="messageProtocolType">
+        /// The type of message protocol used by the endpoint.
+        /// </param>
+        public ProtocolEndpoint(
+            ChannelBase protocolChannel,
             MessageProtocolType messageProtocolType)
         {
-            this.clientChannel = clientChannel;
+            this.protocolChannel = protocolChannel;
             this.messageProtocolType = messageProtocolType;
+            this.originalSynchronizationContext = SynchronizationContext.Current;
         }
 
         /// <summary>
@@ -46,35 +56,50 @@ public async Task Start()
         {
             if (!this.isStarted)
             {
-                // Start the provided client channel
-                this.clientChannel.Start(this.messageProtocolType);
+                // Start the provided protocol channel
+                this.protocolChannel.Start(this.messageProtocolType);
 
                 // Set the handler for any message responses that come back
-                this.clientChannel.MessageDispatcher.SetResponseHandler(this.HandleResponse);
+                this.protocolChannel.MessageDispatcher.SetResponseHandler(this.HandleResponse);
 
                 // Listen for unhandled exceptions from the dispatcher
-                this.clientChannel.MessageDispatcher.UnhandledException += MessageDispatcher_UnhandledException;
+                this.protocolChannel.MessageDispatcher.UnhandledException += MessageDispatcher_UnhandledException;
 
-                // Notify implementation about client start
+                // Notify implementation about endpoint start
                 await this.OnStart();
 
-                // Client is now started
+                // Endpoint is now started
                 this.isStarted = true;
             }
         }
 
+        public void WaitForExit()
+        {
+            this.endpointExitedTask = new TaskCompletionSource<bool>();
+            this.endpointExitedTask.Task.Wait();
+        }
+
         public async Task Stop()
         {
             if (this.isStarted)
             {
+                // Make sure no future calls try to stop the endpoint during shutdown
+                this.isStarted = false;
+
                 // Stop the implementation first
                 await this.OnStop();
+                this.protocolChannel.Stop();
 
-                this.clientChannel.Stop();
-                this.isStarted = false;
+                // Notify anyone waiting for exit
+                if (this.endpointExitedTask != null)
+                {
+                    this.endpointExitedTask.SetResult(true);
+                }
             }
         }
 
+        #region Message Sending
+
         /// <summary>
         /// Sends a request to the server
         /// </summary>
@@ -107,7 +132,7 @@ public async Task<TResult> SendRequest<TParams, TResult>(
                     responseTask);
             }
 
-            await this.clientChannel.MessageWriter.WriteRequest<TParams, TResult>(
+            await this.protocolChannel.MessageWriter.WriteRequest<TParams, TResult>(
                 requestType, 
                 requestParams, 
                 this.currentMessageId);
@@ -128,19 +153,64 @@ await this.clientChannel.MessageWriter.WriteRequest<TParams, TResult>(
             }
         }
 
-        public async Task SendEvent<TParams>(EventType<TParams> eventType, TParams eventParams)
+        /// <summary>
+        /// Sends an event to the channel's endpoint.
+        /// </summary>
+        /// <typeparam name="TParams">The event parameter type.</typeparam>
+        /// <param name="eventType">The type of event being sent.</param>
+        /// <param name="eventParams">The event parameters being sent.</param>
+        /// <returns>A Task that tracks completion of the send operation.</returns>
+        public Task SendEvent<TParams>(
+            EventType<TParams> eventType,
+            TParams eventParams)
+        {
+            // Some events could be raised from a different thread.
+            // To ensure that messages are written serially, dispatch
+            // dispatch the SendEvent call to the message loop thread.
+
+            if (!this.protocolChannel.MessageDispatcher.InMessageLoopThread)
+            {
+                TaskCompletionSource<bool> writeTask = new TaskCompletionSource<bool>();
+
+                this.protocolChannel.MessageDispatcher.SynchronizationContext.Post(
+                    async (obj) =>
+                    {
+                        await this.protocolChannel.MessageWriter.WriteEvent(
+                            eventType,
+                            eventParams);
+
+                        writeTask.SetResult(true);
+                    }, null);
+
+                return writeTask.Task;
+            }
+            else
+            {
+                return this.protocolChannel.MessageWriter.WriteEvent(
+                    eventType,
+                    eventParams);
+            }
+        }
+
+        #endregion
+
+        #region Message Handling
+
+        public void SetRequestHandler<TParams, TResult>(
+            RequestType<TParams, TResult> requestType,
+            Func<TParams, RequestContext<TResult>, Task> requestHandler)
         {
-            await this.clientChannel.MessageWriter.WriteMessage(
-                Message.Event(
-                    eventType.MethodName,
-                    JToken.FromObject(eventParams)));
+            this.protocolChannel.MessageDispatcher.SetRequestHandler(
+                requestType,
+                requestHandler);
         }
 
+
         public void SetEventHandler<TParams>(
             EventType<TParams> eventType,
             Func<TParams, EventContext, Task> eventHandler)
         {
-            this.clientChannel.MessageDispatcher.SetEventHandler(
+            this.protocolChannel.MessageDispatcher.SetEventHandler(
                 eventType,
                 eventHandler,
                 false);
@@ -151,20 +221,12 @@ public void SetEventHandler<TParams>(
             Func<TParams, EventContext, Task> eventHandler,
             bool overrideExisting)
         {
-            this.clientChannel.MessageDispatcher.SetEventHandler(
+            this.protocolChannel.MessageDispatcher.SetEventHandler(
                 eventType,
                 eventHandler,
                 overrideExisting);
         }
 
-        private void MessageDispatcher_UnhandledException(object sender, Exception e)
-        {
-            if (this.originalSynchronizationContext != null)
-            {
-                this.originalSynchronizationContext.Post(o => { throw e; }, null);
-            }
-        }
-
         private void HandleResponse(Message responseMessage)
         {
             TaskCompletionSource<Message> pendingRequestTask = null;
@@ -176,6 +238,10 @@ private void HandleResponse(Message responseMessage)
             }
         }
 
+        #endregion
+
+        #region Subclass Lifetime Methods
+
         protected virtual Task OnStart()
         {
             return Task.FromResult(true);
@@ -185,6 +251,25 @@ protected virtual Task OnStop()
         {
             return Task.FromResult(true);
         }
+
+        #endregion
+
+        #region Event Handlers
+
+        private void MessageDispatcher_UnhandledException(object sender, Exception e)
+        {
+            if (this.endpointExitedTask != null)
+            {
+                this.endpointExitedTask.SetException(e);
+            }
+
+            else if (this.originalSynchronizationContext != null)
+            {
+                this.originalSynchronizationContext.Post(o => { throw e; }, null);
+            }
+        }
+
+        #endregion
     }
 }