Update liveshare callers to specify thread context.

src/Features/LanguageServer/Protocol/Handler/IRequestHandler.cs

@@ -15,6 +15,15 @@ internal interface IRequestHandler
 
     internal interface IRequestHandler<RequestType, ResponseType> : IRequestHandler
     {
+        /// <summary>
+        /// Handles an LSP request.
+        /// </summary>
+        /// <param name="solution">the solution to apply the request to.</param>
+        /// <param name="request">the lsp request.</param>
+        /// <param name="clientCapabilities">the client capabilities for the request.</param>
+        /// <param name="cancellationToken">a cancellation token.</param>
+        /// <param name="keepThreadContext">a value to set if the threading context in the handler should be kept from the caller.</param>
+        /// <returns>the lps response.</returns>
         Task<ResponseType> HandleRequestAsync(Solution solution, RequestType request, ClientCapabilities clientCapabilities, CancellationToken cancellationToken, bool keepThreadContext = false);
     }
 }

src/VisualStudio/LiveShare/Impl/PreviewCodeActionsHandler.cs

@@ -34,6 +34,7 @@ public async Task<LSP.TextEdit[]> HandleAsync(RunCodeActionParams request, Reque
             var codeActions = await GetCodeActionsAsync(solution,
                                                         request.CodeActionParams.TextDocument.Uri,
                                                         request.CodeActionParams.Range,
+                                                        keepThreadContext: false,
                                                         cancellationToken).ConfigureAwait(false);
 
             var actionToRun = codeActions?.FirstOrDefault(a => a.Title == request.Title);

src/VisualStudio/LiveShare/Impl/Shims/AbstractLiveShareHandlerShim.cs

@@ -28,6 +28,15 @@ public virtual Task<ResponseType> HandleAsync(RequestType param, RequestContext<
             return ((IRequestHandler<RequestType, ResponseType>)LazyRequestHandler.Value).HandleRequestAsync(requestContext.Context, param, requestContext.ClientCapabilities, cancellationToken);
         }
 
+        /// <summary>
+        /// Certain implementations require that the processing be done on the UI thread.
+        /// So allow the handler to specifiy that the thread context should be preserved.
+        /// </summary>
+        protected Task<ResponseType> HandleAsyncPreserveThreadContext(RequestType param, RequestContext<Solution> requestContext, CancellationToken cancellationToken)
+        {
+            return ((IRequestHandler<RequestType, ResponseType>)LazyRequestHandler.Value).HandleRequestAsync(requestContext.Context, param, requestContext.ClientCapabilities, cancellationToken, keepThreadContext: true);
+        }
+
         protected Lazy<IRequestHandler, IRequestHandlerMetadata> GetRequestHandler(IEnumerable<Lazy<IRequestHandler, IRequestHandlerMetadata>> requestHandlers, string methodName)
         {
             return requestHandlers.First(handler => handler.Metadata.MethodName == methodName);

src/VisualStudio/LiveShare/Impl/Shims/FindImplementationsHandlerShim.cs

@@ -29,7 +29,7 @@ public override async Task<object> HandleAsync(TextDocumentPositionParams param,
         {
             // TypeScript requires this call to be on the UI thread.
             await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-            return await base.HandleAsync(param, requestContext, cancellationToken).ConfigureAwait(false);
+            return await base.HandleAsyncPreserveThreadContext(param, requestContext, cancellationToken).ConfigureAwait(false);
         }
     }
 }

src/VisualStudio/LiveShare/Impl/Shims/FormatDocumentHandlerShim.cs

@@ -29,7 +29,7 @@ public override async Task<TextEdit[]> HandleAsync(DocumentFormattingParams para
         {
             // To get the formatting options, TypeScript expects to be called on the UI thread.
             await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-            return await base.HandleAsync(param, requestContext, cancellationToken).ConfigureAwait(false);
+            return await base.HandleAsyncPreserveThreadContext(param, requestContext, cancellationToken).ConfigureAwait(false);
         }
     }
 }

src/VisualStudio/LiveShare/Impl/Shims/FormatDocumentOnTypeHandlerShim.cs

@@ -29,7 +29,7 @@ public override async Task<TextEdit[]> HandleAsync(DocumentOnTypeFormattingParam
         {
             // To get the formatting options, TypeScript expects to be called on the UI thread.
             await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-            return await base.HandleAsync(param, requestContext, cancellationToken).ConfigureAwait(false);
+            return await base.HandleAsyncPreserveThreadContext(param, requestContext, cancellationToken).ConfigureAwait(false);
         }
     }
 }

src/VisualStudio/LiveShare/Impl/Shims/FormatDocumentRangeHandlerShim.cs

@@ -29,7 +29,7 @@ public override async Task<TextEdit[]> HandleAsync(DocumentRangeFormattingParams
         {
             // To get the formatting options, TypeScript expects to be called on the UI thread.
             await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-            return await base.HandleAsync(param, requestContext, cancellationToken).ConfigureAwait(false);
+            return await base.HandleAsyncPreserveThreadContext(param, requestContext, cancellationToken).ConfigureAwait(false);
         }
     }
 }

src/VisualStudio/LiveShare/Impl/Shims/RunCodeActionsHandlerShim.cs

@@ -51,7 +51,7 @@ public async override Task<object> HandleAsync(LSP.ExecuteCommandParams request,
             {
                 // Code actions must be applied from the UI thread in the VS context.
                 await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-                return await base.HandleAsync(request, requestContext, cancellationToken).ConfigureAwait(true);
+                return await base.HandleAsyncPreserveThreadContext(request, requestContext, cancellationToken).ConfigureAwait(false);
             }
             catch (ArgumentException ex)
             {