Add cancellation for textDocument/completion, textDocument/codeAction, textDocument/folding

src/PowerShellEditorServices/PowerShellEditorServices.csproj

@@ -39,7 +39,7 @@
     <PackageReference Include="Microsoft.CSharp" Version="4.7.0" />
     <PackageReference Include="Microsoft.Extensions.FileSystemGlobbing" Version="3.1.2" />
     <PackageReference Include="Microsoft.Extensions.Logging" Version="3.1.2" />
-    <PackageReference Include="OmniSharp.Extensions.LanguageServer" Version="0.15.0" />
+    <PackageReference Include="OmniSharp.Extensions.LanguageServer" Version="0.16.0-*" />
     <PackageReference Include="PowerShellStandard.Library" Version="5.1.1" />
     <PackageReference Include="Serilog" Version="2.9.0" />
     <PackageReference Include="Serilog.Extensions.Logging" Version="3.0.1" />
@@ -49,7 +49,7 @@
     <PackageReference Include="System.Security.Principal" Version="4.3.0" />
     <PackageReference Include="System.Security.Principal.Windows" Version="4.7.0" />
     <PackageReference Include="UnixConsoleEcho" Version="0.1.0" />
-    <PackageReference Include="OmniSharp.Extensions.DebugAdapter.Server" Version="0.15.0" />
+    <PackageReference Include="OmniSharp.Extensions.DebugAdapter.Server" Version="0.16.0-*" />
   </ItemGroup>
 
   <ItemGroup>

src/PowerShellEditorServices/Server/PsesDebugServer.cs

@@ -66,7 +66,7 @@ public async Task StartAsync()
             _jsonRpcServer = await JsonRpcServer.From(options =>
             {
                 options.Serializer = new DapProtocolSerializer();
-                options.Reciever = new DapReciever();
+                options.Receiver = new DapReceiver();
                 options.LoggerFactory = _loggerFactory;
                 ILogger logger = options.LoggerFactory.CreateLogger("DebugOptionsStartup");
 

src/PowerShellEditorServices/Services/TextDocument/Handlers/CodeActionHandler.cs

@@ -60,14 +60,18 @@ public void SetCapability(CodeActionCapability capability)
 
         public async Task<CommandOrCodeActionContainer> Handle(CodeActionParams request, CancellationToken cancellationToken)
         {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return Array.Empty<CommandOrCodeAction>();
+            }
+
             // On Windows, VSCode still gives us file URIs like "file:///c%3a/...", so we need to escape them
             IReadOnlyDictionary<string, MarkerCorrection> corrections = await _analysisService.GetMostRecentCodeActionsForFileAsync(
                 _workspaceService.GetFile(request.TextDocument.Uri)).ConfigureAwait(false);
 
             if (corrections == null)
             {
-                // TODO: Find out if we can cache this empty value
-                return new CommandOrCodeActionContainer();
+                return Array.Empty<CommandOrCodeAction>();
             }
 
             var codeActions = new List<CommandOrCodeAction>();
@@ -100,7 +104,7 @@ public async Task<CommandOrCodeActionContainer> Handle(CodeActionParams request,
                                         {
                                             Uri = request.TextDocument.Uri
                                         },
-                                        Edits = new Container<TextEdit>(correction.Edits.Select(ScriptRegion.ToTextEdit))
+                                        Edits = new TextEditContainer(correction.Edits.Select(ScriptRegion.ToTextEdit))
                                     }))
                         }
                     });

src/PowerShellEditorServices/Services/TextDocument/Handlers/CompletionHandler.cs

@@ -24,7 +24,7 @@ namespace Microsoft.PowerShell.EditorServices.Handlers
     internal class CompletionHandler : ICompletionHandler, ICompletionResolveHandler
     {
         const int DefaultWaitTimeoutMilliseconds = 5000;
-        private readonly CompletionItem[] s_emptyCompletionResult = Array.Empty<CompletionItem>();
+        private readonly SemaphoreSlim _completionLock = AsyncUtils.CreateSimpleLockingSemaphore();
 
         private readonly ILogger _logger;
         private readonly PowerShellContextService _powerShellContextService;
@@ -67,24 +67,39 @@ public async Task<CompletionList> Handle(CompletionParams request, CancellationT
 
             ScriptFile scriptFile = _workspaceService.GetFile(request.TextDocument.Uri);
 
-            CompletionResults completionResults =
-                await GetCompletionsInFileAsync(
-                    scriptFile,
-                    cursorLine,
-                    cursorColumn).ConfigureAwait(false);
+            await _completionLock.WaitAsync().ConfigureAwait(false);
 
-            CompletionItem[] completionItems = s_emptyCompletionResult;
-
-            if (completionResults != null)
+            try
             {
-                completionItems = new CompletionItem[completionResults.Completions.Length];
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    _logger.LogDebug("Completion request canceled for file: {0}", request.TextDocument.Uri);
+                    return Array.Empty<CompletionItem>();
+                }
+
+                CompletionResults completionResults =
+                    await GetCompletionsInFileAsync(
+                        scriptFile,
+                        cursorLine,
+                        cursorColumn).ConfigureAwait(false);
+
+                if (completionResults == null)
+                {
+                    return Array.Empty<CompletionItem>();
+                }
+
+                CompletionItem[] completionItems = new CompletionItem[completionResults.Completions.Length];
                 for (int i = 0; i < completionItems.Length; i++)
                 {
                     completionItems[i] = CreateCompletionItem(completionResults.Completions[i], completionResults.ReplacedRange, i + 1);
                 }
-            }
 
-            return new CompletionList(completionItems);
+                return completionItems;
+            }
+            finally
+            {
+                _completionLock.Release();
+            }
         }
 
         public bool CanResolve(CompletionItem value)

src/PowerShellEditorServices/Services/TextDocument/Handlers/DefinitionHandler.cs

@@ -35,9 +35,9 @@ public DefinitionHandler(
             _workspaceService = workspaceService;
         }
 
-        public TextDocumentRegistrationOptions GetRegistrationOptions()
+        public DefinitionRegistrationOptions GetRegistrationOptions()
         {
-            return new TextDocumentRegistrationOptions
+            return new DefinitionRegistrationOptions
             {
                 DocumentSelector = LspUtils.PowerShellDocumentSelector
             };

src/PowerShellEditorServices/Services/TextDocument/Handlers/DocumentHighlightHandler.cs

@@ -27,28 +27,25 @@ internal class DocumentHighlightHandler : IDocumentHighlightHandler
 
         private readonly SymbolsService _symbolsService;
 
-        private readonly TextDocumentRegistrationOptions _registrationOptions;
-
         private DocumentHighlightCapability _capability;
 
         public DocumentHighlightHandler(
             ILoggerFactory loggerFactory,
             WorkspaceService workspaceService,
             SymbolsService symbolService)
         {
-            _logger = loggerFactory.CreateLogger<OmniSharp.Extensions.LanguageServer.Protocol.Server.DocumentHighlightHandler>();
+            _logger = loggerFactory.CreateLogger<DocumentHighlightHandler>();
             _workspaceService = workspaceService;
             _symbolsService = symbolService;
-            _registrationOptions = new TextDocumentRegistrationOptions()
-            {
-                DocumentSelector = LspUtils.PowerShellDocumentSelector
-            };
             _logger.LogInformation("highlight handler loaded");
         }
 
-        public TextDocumentRegistrationOptions GetRegistrationOptions()
+        public DocumentHighlightRegistrationOptions GetRegistrationOptions()
         {
-            return _registrationOptions;
+            return new DocumentHighlightRegistrationOptions
+            {
+                DocumentSelector = LspUtils.PowerShellDocumentSelector
+            };
         }
 
         public Task<DocumentHighlightContainer> Handle(

src/PowerShellEditorServices/Services/TextDocument/Handlers/DocumentSymbolHandler.cs

@@ -33,7 +33,7 @@ internal class DocumentSymbolHandler : IDocumentSymbolHandler
 
         public DocumentSymbolHandler(ILoggerFactory factory, ConfigurationService configurationService, WorkspaceService workspaceService)
         {
-            _logger = factory.CreateLogger<FoldingRangeHandler>();
+            _logger = factory.CreateLogger<DocumentSymbolHandler>();
             _workspaceService = workspaceService;
             _providers = new IDocumentSymbolProvider[]
             {
@@ -43,11 +43,11 @@ public DocumentSymbolHandler(ILoggerFactory factory, ConfigurationService config
             };
         }
 
-        public TextDocumentRegistrationOptions GetRegistrationOptions()
+        public DocumentSymbolRegistrationOptions GetRegistrationOptions()
         {
-            return new TextDocumentRegistrationOptions
+            return new DocumentSymbolRegistrationOptions
             {
-                DocumentSelector = LspUtils.PowerShellDocumentSelector,
+                DocumentSelector = LspUtils.PowerShellDocumentSelector
             };
         }
 

src/PowerShellEditorServices/Services/TextDocument/Handlers/FoldingRangeHandler.cs

@@ -30,16 +30,22 @@ public FoldingRangeHandler(ILoggerFactory factory, ConfigurationService configur
             _configurationService = configurationService;
             _workspaceService = workspaceService;
         }
-        public TextDocumentRegistrationOptions GetRegistrationOptions()
+
+        public FoldingRangeRegistrationOptions GetRegistrationOptions()
         {
-            return new TextDocumentRegistrationOptions()
+            return new FoldingRangeRegistrationOptions
             {
-                DocumentSelector = LspUtils.PowerShellDocumentSelector,
+                DocumentSelector = LspUtils.PowerShellDocumentSelector
             };
         }
 
         public Task<Container<FoldingRange>> Handle(FoldingRangeRequestParam request, CancellationToken cancellationToken)
         {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return Task.FromResult(new Container<FoldingRange>());
+            }
+
             // TODO Should be using dynamic registrations
             if (!_configurationService.CurrentSettings.CodeFolding.Enable) { return null; }
 

src/PowerShellEditorServices/Services/TextDocument/Handlers/FormattingHandlers.cs

@@ -8,6 +8,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Utility;
+using OmniSharp.Extensions.LanguageServer.Protocol;
 using OmniSharp.Extensions.LanguageServer.Protocol.Client.Capabilities;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
@@ -37,9 +38,17 @@ public DocumentFormattingHandlers(
             _workspaceService = workspaceService;
         }
 
-        public TextDocumentRegistrationOptions GetRegistrationOptions()
+        public DocumentFormattingRegistrationOptions GetRegistrationOptions()
         {
-            return new TextDocumentRegistrationOptions
+            return new DocumentFormattingRegistrationOptions
+            {
+                DocumentSelector = LspUtils.PowerShellDocumentSelector
+            };
+        }
+
+        DocumentRangeFormattingRegistrationOptions IRegistration<DocumentRangeFormattingRegistrationOptions>.GetRegistrationOptions()
+        {
+            return new DocumentRangeFormattingRegistrationOptions
             {
                 DocumentSelector = LspUtils.PowerShellDocumentSelector
             };

src/PowerShellEditorServices/Services/TextDocument/Handlers/HoverHandler.cs

@@ -38,11 +38,11 @@ public HoverHandler(
             _powerShellContextService = powerShellContextService;
         }
 
-        public TextDocumentRegistrationOptions GetRegistrationOptions()
+        public HoverRegistrationOptions GetRegistrationOptions()
         {
-            return new TextDocumentRegistrationOptions
+            return new HoverRegistrationOptions
             {
-                DocumentSelector = LspUtils.PowerShellDocumentSelector,
+                DocumentSelector = LspUtils.PowerShellDocumentSelector
             };
         }
 

src/PowerShellEditorServices/Services/TextDocument/Handlers/ReferencesHandler.cs

@@ -22,7 +22,7 @@ class ReferencesHandler : IReferencesHandler
         private readonly ILogger _logger;
         private readonly SymbolsService _symbolsService;
         private readonly WorkspaceService _workspaceService;
-        private ReferencesCapability _capability;
+        private ReferenceCapability _capability;
 
         public ReferencesHandler(ILoggerFactory factory, SymbolsService symbolsService, WorkspaceService workspaceService)
         {
@@ -31,9 +31,9 @@ public ReferencesHandler(ILoggerFactory factory, SymbolsService symbolsService,
             _workspaceService = workspaceService;
         }
 
-        public TextDocumentRegistrationOptions GetRegistrationOptions()
+        public ReferenceRegistrationOptions GetRegistrationOptions()
         {
-            return new TextDocumentRegistrationOptions
+            return new ReferenceRegistrationOptions
             {
                 DocumentSelector = LspUtils.PowerShellDocumentSelector
             };
@@ -72,7 +72,7 @@ public async Task<LocationContainer> Handle(ReferenceParams request, Cancellatio
             return new LocationContainer(locations);
         }
 
-        public void SetCapability(ReferencesCapability capability)
+        public void SetCapability(ReferenceCapability capability)
         {
             _capability = capability;
         }

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

@@ -32,13 +32,12 @@ public WorkspaceSymbolsHandler(ILoggerFactory loggerFactory, SymbolsService symb
             _workspaceService = workspace;
         }
 
-        public object GetRegistrationOptions()
+        public WorkspaceSymbolRegistrationOptions GetRegistrationOptions()
         {
-            return null;
-            // throw new NotImplementedException();
+            return new WorkspaceSymbolRegistrationOptions();
         }
 
-        public Task<SymbolInformationContainer> Handle(WorkspaceSymbolParams request, CancellationToken cancellationToken)
+        public Task<Container<SymbolInformation>> Handle(WorkspaceSymbolParams request, CancellationToken cancellationToken)
         {
             var symbols = new List<SymbolInformation>();
 
@@ -75,7 +74,7 @@ public Task<SymbolInformationContainer> Handle(WorkspaceSymbolParams request, Ca
             }
             _logger.LogWarning("Logging in a handler works now.");
 
-            return Task.FromResult(new SymbolInformationContainer(symbols));
+            return Task.FromResult(new Container<SymbolInformation>(symbols));
         }
 
         public void SetCapability(WorkspaceSymbolCapability capability)

test/PowerShellEditorServices.Test.E2E/LanguageServerProtocolMessageTests.cs

@@ -124,7 +124,7 @@ function CanSendWorkspaceSymbolRequest {
 }
 ");
 
-            SymbolInformationContainer symbols = await LanguageClient.SendRequest<SymbolInformationContainer>(
+            Container<SymbolInformation> symbols = await LanguageClient.SendRequest<Container<SymbolInformation>>(
                 "workspace/symbol",
                 new WorkspaceSymbolParams
                 {

test/PowerShellEditorServices.Test.E2E/PowerShellEditorServices.Test.E2E.csproj

@@ -10,7 +10,7 @@
     <PackageReference Include="Microsoft.Extensions.Logging" Version="3.1.2" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="16.5.0" />
     <PackageReference Include="Newtonsoft.Json" Version="12.0.2" />
-    <PackageReference Include="OmniSharp.Extensions.LanguageClient" Version="0.15.0" />
+    <PackageReference Include="OmniSharp.Extensions.LanguageClient" Version="0.16.0-*" />
     <PackageReference Include="xunit" Version="2.4.0" />
     <PackageReference Include="xunit.runner.visualstudio" Version="2.4.0" />
   </ItemGroup>