Implement textDocument/codeAction (PowerShell#1004)

* Add initial handler

* Add working codeAction implementation

* Crash

* Make tests work

* Fix issues

* Make tests work in WinPS

Author: rjmholt

src/PowerShellEditorServices.Engine/LanguageServer/OmnisharpLanguageServer.cs

@@ -94,7 +94,8 @@ public async Task StartAsync()
                     .WithHandler<DocumentSymbolHandler>()
                     .WithHandler<DocumentHighlightHandler>()
                     .WithHandler<PSHostProcessAndRunspaceHandlers>()
-                    .WithHandler<CodeLensHandlers>();
+                    .WithHandler<CodeLensHandlers>()
+                    .WithHandler<CodeActionHandler>();
 
                     logger.LogInformation("Handlers added");
             });

src/PowerShellEditorServices.Engine/PowerShellEditorServices.Engine.csproj

@@ -12,7 +12,7 @@
 
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.FileSystemGlobbing" Version="2.2.0" />
-    <PackageReference Include="OmniSharp.Extensions.LanguageServer" Version="0.13.0-*" />
+    <PackageReference Include="OmniSharp.Extensions.LanguageServer" Version="0.13.2-*" />
     <PackageReference Include="PowerShellStandard.Library" Version="5.1.1" />
     <PackageReference Include="Serilog.Extensions.Logging" Version="2.0.4" />
     <PackageReference Include="Serilog.Sinks.Console" Version="3.1.1" />

src/PowerShellEditorServices.Engine/Services/Analysis/AnalysisService.cs

@@ -15,14 +15,15 @@
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server;
 using System.Threading;
+using System.Collections.Concurrent;
 
 namespace Microsoft.PowerShell.EditorServices
 {
     /// <summary>
     /// Provides a high-level service for performing semantic analysis
     /// of PowerShell scripts.
     /// </summary>
-    public class AnalysisService : IDisposable
+    internal class AnalysisService : IDisposable
     {
         #region Static fields
 
@@ -54,9 +55,6 @@ public class AnalysisService : IDisposable
 
         private static readonly string[] s_emptyGetRuleResult = new string[0];
 
-        private Dictionary<string, Dictionary<string, MarkerCorrection>> codeActionsPerFile =
-            new Dictionary<string, Dictionary<string, MarkerCorrection>>();
-
         private static CancellationTokenSource s_existingRequestCancellation;
 
         /// <summary>
@@ -96,8 +94,11 @@ public class AnalysisService : IDisposable
         private PSModuleInfo _pssaModuleInfo;
 
         private readonly ILanguageServer _languageServer;
+
         private readonly ConfigurationService _configurationService;
 
+        private readonly ConcurrentDictionary<string, (SemaphoreSlim, Dictionary<string, MarkerCorrection>)> _mostRecentCorrectionsByFile;
+
         #endregion // Private Fields
 
         #region Properties
@@ -149,6 +150,7 @@ private AnalysisService(
             _configurationService = configurationService;
             _logger = logger;
             _pssaModuleInfo = pssaModuleInfo;
+            _mostRecentCorrectionsByFile = new ConcurrentDictionary<string, (SemaphoreSlim, Dictionary<string, MarkerCorrection>)>();
         }
 
         #endregion // constructors
@@ -813,26 +815,57 @@ private void PublishScriptDiagnostics(
             ScriptFile scriptFile,
             List<ScriptFileMarker> markers)
         {
-            List<Diagnostic> diagnostics = new List<Diagnostic>();
+            var diagnostics = new List<Diagnostic>();
 
-            // Hold on to any corrections that may need to be applied later
-            Dictionary<string, MarkerCorrection> fileCorrections =
-                new Dictionary<string, MarkerCorrection>();
+            // Create the entry for this file if it does not already exist
+            SemaphoreSlim fileLock;
+            Dictionary<string, MarkerCorrection> fileCorrections;
+            bool newEntryNeeded = false;
+            if (_mostRecentCorrectionsByFile.TryGetValue(scriptFile.DocumentUri, out (SemaphoreSlim, Dictionary<string, MarkerCorrection>) fileCorrectionsEntry))
+            {
+                fileLock = fileCorrectionsEntry.Item1;
+                fileCorrections = fileCorrectionsEntry.Item2;
+            }
+            else
+            {
+                fileLock = new SemaphoreSlim(initialCount: 1, maxCount: 1);
+                fileCorrections = new Dictionary<string, MarkerCorrection>();
+                newEntryNeeded = true;
+            }
 
-            foreach (var marker in markers)
+            fileLock.Wait();
+            try
             {
-                // Does the marker contain a correction?
-                Diagnostic markerDiagnostic = GetDiagnosticFromMarker(marker);
-                if (marker.Correction != null)
+                if (newEntryNeeded)
                 {
-                    string diagnosticId = GetUniqueIdFromDiagnostic(markerDiagnostic);
-                    fileCorrections[diagnosticId] = marker.Correction;
+                    // If we create a new entry, we should do it after acquiring the lock we just created
+                    // to ensure a competing thread can never acquire it first and read invalid information from it
+                    _mostRecentCorrectionsByFile[scriptFile.DocumentUri] = (fileLock, fileCorrections);
                 }
+                else
+                {
+                    // Otherwise we need to clear the stale corrections
+                    fileCorrections.Clear();
+                }
+
+                foreach (ScriptFileMarker marker in markers)
+                {
+                    // Does the marker contain a correction?
+                    Diagnostic markerDiagnostic = GetDiagnosticFromMarker(marker);
+                    if (marker.Correction != null)
+                    {
+                        string diagnosticId = GetUniqueIdFromDiagnostic(markerDiagnostic);
+                        fileCorrections[diagnosticId] = marker.Correction;
+                    }
 
-                diagnostics.Add(markerDiagnostic);
+                    diagnostics.Add(markerDiagnostic);
+                }
+            }
+            finally
+            {
+                fileLock.Release();
             }
 
-            codeActionsPerFile[scriptFile.DocumentUri] = fileCorrections;
 
             var uriBuilder = new UriBuilder()
             {
@@ -850,17 +883,42 @@ private void PublishScriptDiagnostics(
             });
         }
 
+        public async Task<IReadOnlyDictionary<string, MarkerCorrection>> GetMostRecentCodeActionsForFileAsync(string documentUri)
+        {
+            if (!_mostRecentCorrectionsByFile.TryGetValue(documentUri, out (SemaphoreSlim fileLock, Dictionary<string, MarkerCorrection> corrections) fileCorrectionsEntry))
+            {
+                return null;
+            }
+
+            await fileCorrectionsEntry.fileLock.WaitAsync();
+            // We must copy the dictionary for thread safety
+            var corrections = new Dictionary<string, MarkerCorrection>(fileCorrectionsEntry.corrections.Count);
+            try
+            {
+                foreach (KeyValuePair<string, MarkerCorrection> correction in fileCorrectionsEntry.corrections)
+                {
+                    corrections.Add(correction.Key, correction.Value);
+                }
+
+                return corrections;
+            }
+            finally
+            {
+                fileCorrectionsEntry.fileLock.Release();
+            }
+        }
+
         // Generate a unique id that is used as a key to look up the associated code action (code fix) when
         // we receive and process the textDocument/codeAction message.
-        private static string GetUniqueIdFromDiagnostic(Diagnostic diagnostic)
+        internal static string GetUniqueIdFromDiagnostic(Diagnostic diagnostic)
         {
             Position start = diagnostic.Range.Start;
             Position end = diagnostic.Range.End;
 
             var sb = new StringBuilder(256)
             .Append(diagnostic.Source ?? "?")
             .Append("_")
-            .Append(diagnostic.Code.ToString())
+            .Append(diagnostic.Code.IsString ? diagnostic.Code.String : diagnostic.Code.Long.ToString())
             .Append("_")
             .Append(diagnostic.Severity?.ToString() ?? "?")
             .Append("_")

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

@@ -0,0 +1,120 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.ComponentModel;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+using Newtonsoft.Json.Linq;
+using OmniSharp.Extensions.JsonRpc.Client;
+using OmniSharp.Extensions.LanguageServer.Protocol.Client.Capabilities;
+using OmniSharp.Extensions.LanguageServer.Protocol.Models;
+using OmniSharp.Extensions.LanguageServer.Protocol.Server;
+using PowerShellEditorServices.Engine.Services.Handlers;
+
+namespace Microsoft.PowerShell.EditorServices.TextDocument
+{
+    internal class CodeActionHandler : ICodeActionHandler
+    {
+        private static readonly CodeActionKind[] s_supportedCodeActions = new[]
+        {
+            CodeActionKind.QuickFix
+        };
+
+        private readonly CodeActionRegistrationOptions _registrationOptions;
+
+        private readonly ILogger _logger;
+
+        private readonly AnalysisService _analysisService;
+
+        private CodeActionCapability _capability;
+
+        public CodeActionHandler(ILoggerFactory factory, AnalysisService analysisService)
+        {
+            _logger = factory.CreateLogger<TextDocumentHandler>();
+            _analysisService = analysisService;
+            _registrationOptions = new CodeActionRegistrationOptions()
+            {
+                DocumentSelector = new DocumentSelector(new DocumentFilter() { Pattern = "**/*.ps*1" }),
+                CodeActionKinds = s_supportedCodeActions
+            };
+        }
+
+        public CodeActionRegistrationOptions GetRegistrationOptions()
+        {
+            return _registrationOptions;
+        }
+
+        public void SetCapability(CodeActionCapability capability)
+        {
+            _capability = capability;
+        }
+
+        public async Task<CommandOrCodeActionContainer> Handle(CodeActionParams request, CancellationToken cancellationToken)
+        {
+            IReadOnlyDictionary<string, MarkerCorrection> corrections = await _analysisService.GetMostRecentCodeActionsForFileAsync(request.TextDocument.Uri.ToString());
+
+            if (corrections == null)
+            {
+                // TODO: Find out if we can cache this empty value
+                return new CommandOrCodeActionContainer();
+            }
+
+            var codeActions = new List<CommandOrCodeAction>();
+
+            // If there are any code fixes, send these commands first so they appear at top of "Code Fix" menu in the client UI.
+            foreach (Diagnostic diagnostic in request.Context.Diagnostics)
+            {
+                if (diagnostic.Code.IsLong)
+                {
+                    _logger.LogWarning(
+                        $"textDocument/codeAction skipping diagnostic with non-string code {diagnostic.Code.Long}: {diagnostic.Source} {diagnostic.Message}");
+                }
+                else if (string.IsNullOrEmpty(diagnostic.Code.String))
+                {
+                    _logger.LogWarning(
+                        $"textDocument/codeAction skipping diagnostic with empty Code field: {diagnostic.Source} {diagnostic.Message}");
+
+                    continue;
+                }
+
+
+                string diagnosticId = AnalysisService.GetUniqueIdFromDiagnostic(diagnostic);
+                if (corrections.TryGetValue(diagnosticId, out MarkerCorrection correction))
+                {
+                    codeActions.Add(new Command()
+                    {
+                        Title = correction.Name,
+                        Name = "PowerShell.ApplyCodeActionEdits",
+                        Arguments = JArray.FromObject(correction.Edits)
+                    });
+                }
+            }
+
+            // Add "show documentation" commands last so they appear at the bottom of the client UI.
+            // These commands do not require code fixes. Sometimes we get a batch of diagnostics
+            // to create commands for. No need to create multiple show doc commands for the same rule.
+            var ruleNamesProcessed = new HashSet<string>();
+            foreach (Diagnostic diagnostic in request.Context.Diagnostics)
+            {
+                if (!diagnostic.Code.IsString || string.IsNullOrEmpty(diagnostic.Code.String)) { continue; }
+
+                if (string.Equals(diagnostic.Source, "PSScriptAnalyzer", StringComparison.OrdinalIgnoreCase) &&
+                    !ruleNamesProcessed.Contains(diagnostic.Code.String))
+                {
+                    ruleNamesProcessed.Add(diagnostic.Code.String);
+
+                    codeActions.Add(
+                        new Command
+                        {
+                            Title = $"Show documentation for \"{diagnostic.Code}\"",
+                            Name = "PowerShell.ShowCodeActionDocumentation",
+                            Arguments = JArray.FromObject(new[] { diagnostic.Code })
+                        });
+                }
+            }
+
+            return codeActions;
+        }
+    }
+}

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

@@ -10,7 +10,7 @@
 
 namespace PowerShellEditorServices.Engine.Services.Handlers
 {
-    public class DocumentFormattingHandler : IDocumentFormattingHandler
+    internal class DocumentFormattingHandler : IDocumentFormattingHandler
     {
         private readonly DocumentSelector _documentSelector = new DocumentSelector(
             new DocumentFilter()
@@ -88,7 +88,7 @@ public void SetCapability(DocumentFormattingCapability capability)
         }
     }
 
-    public class DocumentRangeFormattingHandler : IDocumentRangeFormattingHandler
+    internal class DocumentRangeFormattingHandler : IDocumentRangeFormattingHandler
     {
         private readonly DocumentSelector _documentSelector = new DocumentSelector(
             new DocumentFilter()

src/PowerShellEditorServices.Engine/Services/TextDocument/ScriptFileMarker.cs

@@ -147,11 +147,11 @@ internal static ScriptFileMarker FromDiagnosticRecord(PSObject psObject)
                         new ScriptRegion(
                             diagnosticRecord.ScriptPath,
                             suggestedCorrection.Text,
-                            suggestedCorrection.StartLineNumber,
-                            suggestedCorrection.StartColumnNumber,
+                            startLineNumber: suggestedCorrection.StartLineNumber,
+                            startColumnNumber: suggestedCorrection.StartColumnNumber,
+                            endLineNumber: suggestedCorrection.EndLineNumber,
+                            endColumnNumber: suggestedCorrection.EndColumnNumber,
                             startOffset: -1,
-                            suggestedCorrection.EndLineNumber,
-                            suggestedCorrection.EndColumnNumber,
                             endOffset: -1));
 
                     correctionMessage = suggestedCorrection.Description;

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

@@ -10,7 +10,7 @@
 
 namespace Microsoft.PowerShell.EditorServices
 {
-    public class ConfigurationHandler : IDidChangeConfigurationHandler
+    internal class ConfigurationHandler : IDidChangeConfigurationHandler
     {
         private readonly ILogger _logger;
         private readonly AnalysisService _analysisService;

src/PowerShellEditorServices.Protocol/Server/LanguageServer.cs

@@ -1744,7 +1744,6 @@ await DelayThenInvokeDiagnosticsAsync(
                 cancellationToken);
         }
 
-
         private static async Task DelayThenInvokeDiagnosticsAsync(
             int delayMilliseconds,
             ScriptFile[] filesToAnalyze,