Add support for additional file diagnostics in workspace pull

src/EditorFeatures/Core/LanguageServer/Handlers/CodeActions/CodeActionHelpers.cs

@@ -154,7 +154,7 @@ static VSInternalCodeAction[] GenerateNestedVSCodeActions(
 
             static LSP.Diagnostic[]? GetApplicableDiagnostics(LSP.CodeActionContext context, IUnifiedSuggestedAction action)
             {
-                if (action is UnifiedCodeFixSuggestedAction codeFixAction)
+                if (action is UnifiedCodeFixSuggestedAction codeFixAction && context.Diagnostics != null)
                 {
                     // Associate the diagnostics from the request that match the diagnostic fixed by the code action by ID.
                     // The request diagnostics are already restricted to the code fix location by the request.

src/EditorFeatures/TestUtilities/IsExternalInit.cs

@@ -0,0 +1,11 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+namespace System.Runtime.CompilerServices
+{
+    // Used to compile against C# 9 in a netstandard2.0
+    internal class IsExternalInit
+    {
+    }
+}

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.InitializationOptions.cs

@@ -0,0 +1,22 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using Microsoft.CodeAnalysis.LanguageServer;
+using Microsoft.CodeAnalysis.Options;
+using LSP = Microsoft.VisualStudio.LanguageServer.Protocol;
+
+namespace Roslyn.Test.Utilities
+{
+    public abstract partial class AbstractLanguageServerProtocolTests
+    {
+        internal record struct InitializationOptions()
+        {
+            internal string[] SourceGeneratedMarkups { get; init; } = Array.Empty<string>();
+            internal LSP.ClientCapabilities ClientCapabilities { get; init; } = new LSP.ClientCapabilities();
+            internal WellKnownLspServerKinds ServerKind { get; init; } = WellKnownLspServerKinds.AlwaysActiveVSLspServer;
+            internal Action<IGlobalOptionService>? OptionUpdater { get; init; } = null;
+        }
+    }
+}

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

@@ -99,6 +99,9 @@ protected class OrderLocations : Comparer<LSP.Location>
 
         protected virtual TestComposition Composition => s_composition;
 
+        private protected virtual TestAnalyzerReferenceByLanguage TestAnalyzerReferences
+            => new(DiagnosticExtensions.GetCompilerDiagnosticAnalyzersMap());
+
         protected static LSP.ClientCapabilities CapabilitiesWithVSExtensions => new LSP.VSInternalClientCapabilities { SupportsVisualStudioExtensions = true };
 
         /// <summary>
@@ -282,50 +285,42 @@ protected static LSP.TextEdit GenerateTextEdit(string newText, int startLine, in
         private protected static CodeActionResolveData CreateCodeActionResolveData(string uniqueIdentifier, LSP.Location location, IEnumerable<string>? customTags = null)
             => new CodeActionResolveData(uniqueIdentifier, customTags.ToImmutableArrayOrEmpty(), location.Range, CreateTextDocumentIdentifier(location.Uri));
 
-        /// <summary>
-        /// Creates an LSP server backed by a workspace instance with a solution containing the markup.
-        /// </summary>
-        protected Task<TestLspServer> CreateTestLspServerAsync(string markup, LSP.ClientCapabilities? clientCapabilities = null)
-            => CreateTestLspServerAsync(new string[] { markup }, Array.Empty<string>(), LanguageNames.CSharp, clientCapabilities);
-
-        private protected Task<TestLspServer> CreateVisualBasicTestLspServerAsync(string markup, LSP.ClientCapabilities? clientCapabilities = null, WellKnownLspServerKinds serverKind = WellKnownLspServerKinds.AlwaysActiveVSLspServer)
-            => CreateTestLspServerAsync(new string[] { markup }, Array.Empty<string>(), LanguageNames.VisualBasic, clientCapabilities, serverKind);
-
-        protected Task<TestLspServer> CreateMultiProjectLspServerAsync(string xmlMarkup, LSP.ClientCapabilities? clientCapabilities = null)
-            => CreateTestLspServerAsync(TestWorkspace.Create(xmlMarkup, composition: Composition), clientCapabilities, WellKnownLspServerKinds.AlwaysActiveVSLspServer);
+        private protected Task<TestLspServer> CreateTestLspServerAsync(string markup, LSP.ClientCapabilities clientCapabilities)
+            => CreateTestLspServerAsync(new string[] { markup }, LanguageNames.CSharp, new InitializationOptions { ClientCapabilities = clientCapabilities });
 
-        /// <summary>
-        /// Creates an LSP server backed by a workspace instance with a solution containing the specified documents.
-        /// </summary>
-        protected Task<TestLspServer> CreateTestLspServerAsync(string[] markups, LSP.ClientCapabilities? clientCapabilities = null)
-            => CreateTestLspServerAsync(markups, Array.Empty<string>(), LanguageNames.CSharp, clientCapabilities);
+        private protected Task<TestLspServer> CreateTestLspServerAsync(string markup, InitializationOptions? initializationOptions = null)
+            => CreateTestLspServerAsync(new string[] { markup }, LanguageNames.CSharp, initializationOptions);
 
-        private protected Task<TestLspServer> CreateTestLspServerAsync(string markup, LSP.ClientCapabilities clientCapabilities, WellKnownLspServerKinds serverKind)
-            => CreateTestLspServerAsync(new string[] { markup }, Array.Empty<string>(), LanguageNames.CSharp, clientCapabilities, serverKind);
+        private protected Task<TestLspServer> CreateTestLspServerAsync(string[] markups, InitializationOptions? initializationOptions = null)
+            => CreateTestLspServerAsync(markups, LanguageNames.CSharp, initializationOptions);
 
-        /// <summary>
-        /// Creates an LSP server backed by a workspace instance with a solution containing the specified documents.
-        /// </summary>
-        protected Task<TestLspServer> CreateTestLspServerAsync(string[] markups, string[] sourceGeneratedMarkups, LSP.ClientCapabilities? clientCapabilities = null)
-            => CreateTestLspServerAsync(markups, sourceGeneratedMarkups, LanguageNames.CSharp, clientCapabilities);
+        private protected Task<TestLspServer> CreateVisualBasicTestLspServerAsync(string markup, InitializationOptions? initializationOptions = null)
+            => CreateTestLspServerAsync(new string[] { markup }, LanguageNames.VisualBasic, initializationOptions);
 
-        private Task<TestLspServer> CreateTestLspServerAsync(string[] markups, string[] sourceGeneratedMarkups, string languageName, LSP.ClientCapabilities? clientCapabilities, WellKnownLspServerKinds serverKind = WellKnownLspServerKinds.AlwaysActiveVSLspServer)
+        private Task<TestLspServer> CreateTestLspServerAsync(string[] markups, string languageName, InitializationOptions? initializationOptions)
         {
+            var lspOptions = initializationOptions ?? new InitializationOptions();
             var exportProvider = Composition.ExportProviderFactory.CreateExportProvider();
             var workspaceConfigurationService = exportProvider.GetExportedValue<TestWorkspaceConfigurationService>();
             workspaceConfigurationService.Options = new WorkspaceConfigurationOptions(EnableOpeningSourceGeneratedFiles: true);
 
+            if (lspOptions.OptionUpdater != null)
+            {
+                var globalOptions = exportProvider.GetExportedValue<IGlobalOptionService>();
+                lspOptions.OptionUpdater(globalOptions);
+            }
+
             var workspace = languageName switch
             {
-                LanguageNames.CSharp => TestWorkspace.CreateCSharp(markups, sourceGeneratedMarkups, exportProvider: exportProvider),
-                LanguageNames.VisualBasic => TestWorkspace.CreateVisualBasic(markups, sourceGeneratedMarkups, exportProvider: exportProvider),
+                LanguageNames.CSharp => TestWorkspace.CreateCSharp(markups, lspOptions.SourceGeneratedMarkups, exportProvider: exportProvider),
+                LanguageNames.VisualBasic => TestWorkspace.CreateVisualBasic(markups, lspOptions.SourceGeneratedMarkups, exportProvider: exportProvider),
                 _ => throw new ArgumentException($"language name {languageName} is not valid for a test workspace"),
             };
 
-            return CreateTestLspServerAsync(workspace, clientCapabilities, serverKind);
+            return CreateTestLspServerAsync(workspace, lspOptions);
         }
 
-        private static async Task<TestLspServer> CreateTestLspServerAsync(TestWorkspace workspace, LSP.ClientCapabilities? clientCapabilities, WellKnownLspServerKinds serverKind)
+        private async Task<TestLspServer> CreateTestLspServerAsync(TestWorkspace workspace, InitializationOptions initializationOptions)
         {
             var solution = workspace.CurrentSolution;
 
@@ -346,29 +341,38 @@ private static async Task<TestLspServer> CreateTestLspServerAsync(TestWorkspace
                 solution = solution.WithProjectFilePath(project.Id, GetDocumentFilePathFromName(project.FilePath));
             }
 
+            solution = solution.WithAnalyzerReferences(new[] { TestAnalyzerReferences });
             workspace.ChangeSolution(solution);
 
             // Important: We must wait for workspace creation operations to finish.
             // Otherwise we could have a race where workspace change events triggered by creation are changing the state
             // created by the initial test steps. This can interfere with the expected test state.
             await WaitForWorkspaceOperationsAsync(workspace);
 
-            return await TestLspServer.CreateAsync(workspace, clientCapabilities ?? new LSP.ClientCapabilities(), serverKind);
+            return await TestLspServer.CreateAsync(workspace, initializationOptions);
         }
 
         private protected async Task<TestLspServer> CreateXmlTestLspServerAsync(
             string xmlContent,
             string? workspaceKind = null,
-            LSP.ClientCapabilities? clientCapabilities = null,
-            WellKnownLspServerKinds serverKind = WellKnownLspServerKinds.AlwaysActiveVSLspServer)
+            InitializationOptions? initializationOptions = null)
         {
-            var workspace = TestWorkspace.Create(XElement.Parse(xmlContent), openDocuments: false, composition: Composition, workspaceKind: workspaceKind);
+            var lspOptions = initializationOptions ?? new InitializationOptions();
+            var exportProvider = Composition.ExportProviderFactory.CreateExportProvider();
+            if (lspOptions.OptionUpdater != null)
+            {
+                var globalOptions = exportProvider.GetExportedValue<IGlobalOptionService>();
+                lspOptions.OptionUpdater(globalOptions);
+            }
+
+            var workspace = TestWorkspace.Create(XElement.Parse(xmlContent), openDocuments: false, exportProvider: exportProvider, workspaceKind: workspaceKind);
+            workspace.TryApplyChanges(workspace.CurrentSolution.WithAnalyzerReferences(new[] { TestAnalyzerReferences }));
 
             // Important: We must wait for workspace creation operations to finish.
             // Otherwise we could have a race where workspace change events triggered by creation are changing the state
             // created by the initial test steps. This can interfere with the expected test state.
             await WaitForWorkspaceOperationsAsync(workspace);
-            return await TestLspServer.CreateAsync(workspace, clientCapabilities ?? new LSP.ClientCapabilities(), serverKind);
+            return await TestLspServer.CreateAsync(workspace, lspOptions);
         }
 
         /// <summary>
@@ -472,7 +476,7 @@ private static LSP.DidCloseTextDocumentParams CreateDidCloseTextDocumentParams(U
                }
            };
 
-        public sealed class TestLspServer : IDisposable
+        internal sealed class TestLspServer : IDisposable
         {
             public readonly TestWorkspace TestWorkspace;
             private readonly Dictionary<string, IList<LSP.Location>> _locations;
@@ -534,14 +538,14 @@ private static JsonMessageFormatter CreateJsonMessageFormatter()
                 return messageFormatter;
             }
 
-            internal static async Task<TestLspServer> CreateAsync(TestWorkspace testWorkspace, LSP.ClientCapabilities clientCapabilities, WellKnownLspServerKinds serverKind)
+            internal static async Task<TestLspServer> CreateAsync(TestWorkspace testWorkspace, InitializationOptions initializationOptions)
             {
                 var locations = await GetAnnotatedLocationsAsync(testWorkspace, testWorkspace.CurrentSolution);
-                var server = new TestLspServer(testWorkspace, locations, clientCapabilities, serverKind);
+                var server = new TestLspServer(testWorkspace, locations, initializationOptions.ClientCapabilities, initializationOptions.ServerKind);
 
                 await server.ExecuteRequestAsync<LSP.InitializeParams, LSP.InitializeResult>(LSP.Methods.InitializeName, new LSP.InitializeParams
                 {
-                    Capabilities = clientCapabilities,
+                    Capabilities = initializationOptions.ClientCapabilities,
                 }, CancellationToken.None);
 
                 return server;
@@ -639,22 +643,6 @@ public Task CloseDocumentAsync(Uri documentUri)
 
             public Solution GetCurrentSolution() => TestWorkspace.CurrentSolution;
 
-            internal void InitializeDiagnostics(BackgroundAnalysisScope scope, DiagnosticMode diagnosticMode, TestAnalyzerReferenceByLanguage references)
-            {
-                TestWorkspace.GlobalOptions.SetGlobalOption(new OptionKey(SolutionCrawlerOptionsStorage.BackgroundAnalysisScopeOption, LanguageNames.CSharp), scope);
-                TestWorkspace.GlobalOptions.SetGlobalOption(new OptionKey(SolutionCrawlerOptionsStorage.BackgroundAnalysisScopeOption, LanguageNames.VisualBasic), scope);
-                TestWorkspace.GlobalOptions.SetGlobalOption(new OptionKey(SolutionCrawlerOptionsStorage.BackgroundAnalysisScopeOption, InternalLanguageNames.TypeScript), scope);
-                TestWorkspace.GlobalOptions.SetGlobalOption(new OptionKey(InternalDiagnosticsOptions.NormalDiagnosticMode), diagnosticMode);
-
-                TestWorkspace.TryApplyChanges(TestWorkspace.CurrentSolution.WithAnalyzerReferences(new[] { references }));
-
-                var registrationService = TestWorkspace.Services.GetRequiredService<ISolutionCrawlerRegistrationService>();
-                registrationService.Register(TestWorkspace);
-
-                var diagnosticService = (DiagnosticService)TestWorkspace.ExportProvider.GetExportedValue<IDiagnosticService>();
-                diagnosticService.Register(new TestHostDiagnosticUpdateSource(TestWorkspace));
-            }
-
             internal async Task WaitForDiagnosticsAsync()
             {
                 var listenerProvider = TestWorkspace.GetService<IAsynchronousOperationListenerProvider>();

src/Features/LanguageServer/Protocol/Extensions/Extensions.cs

@@ -6,6 +6,7 @@
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Reflection.Metadata;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.FindUsages;
@@ -36,8 +37,8 @@ public static ImmutableArray<Document> GetDocuments(this Solution solution, Uri
         {
             var documentIds = GetDocumentIds(solution, documentUri);
 
-            var documents = documentIds.SelectAsArray(id => solution.GetRequiredDocument(id));
-
+            // We don't call GetRequiredDocument here as the id could be referring to an additional document.
+            var documents = documentIds.Select(solution.GetDocument).WhereNotNull().ToImmutableArray();
             return documents;
         }
 
@@ -103,6 +104,14 @@ public static Document FindDocumentInProjectContext(this ImmutableArray<Document
         public static Project? GetProject(this Solution solution, TextDocumentIdentifier projectIdentifier)
             => solution.Projects.Where(project => project.FilePath == projectIdentifier.Uri.LocalPath).SingleOrDefault();
 
+        public static TextDocument? GetAdditionalDocument(this Solution solution, TextDocumentIdentifier documentIdentifier)
+        {
+            var documentIds = GetDocumentIds(solution, documentIdentifier.Uri);
+
+            // We don't call GetRequiredAdditionalDocument as the id could be referring to a regular document.
+            return documentIds.Select(solution.GetAdditionalDocument).WhereNotNull().SingleOrDefault();
+        }
+
         public static async Task<int> GetPositionFromLinePositionAsync(this TextDocument document, LinePosition linePosition, CancellationToken cancellationToken)
         {
             var text = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);

src/Features/LanguageServer/Protocol/Handler/Diagnostics/AbstractPullDiagnosticHandler.cs

@@ -89,10 +89,19 @@ protected AbstractPullDiagnosticHandler(
         protected abstract ValueTask<ImmutableArray<IDiagnosticSource>> GetOrderedDiagnosticSourcesAsync(RequestContext context, CancellationToken cancellationToken);
 
         /// <summary>
-        /// Creates the <see cref="VSInternalDiagnosticReport"/> instance we'll report back to clients to let them know our
-        /// progress.  Subclasses can fill in data specific to their needs as appropriate.
+        /// Creates the appropriate LSP type to report a new set of diagnostics and resultId.
         /// </summary>
-        protected abstract TReport CreateReport(TextDocumentIdentifier identifier, LSP.Diagnostic[]? diagnostics, string? resultId);
+        protected abstract TReport CreateReport(TextDocumentIdentifier identifier, LSP.Diagnostic[] diagnostics, string resultId);
+
+        /// <summary>
+        /// Creates the appropriate LSP type to report unchanged diagnostics.
+        /// </summary>
+        protected abstract TReport CreateUnchangedReport(TextDocumentIdentifier identifier, string resultId);
+
+        /// <summary>
+        /// Creates the appropriate LSP type to report a removed file.
+        /// </summary>
+        protected abstract TReport CreateRemovedReport(TextDocumentIdentifier identifier);
 
         protected abstract TReturn? CreateReturn(BufferedProgress<TReport> progress);
 
@@ -157,7 +166,7 @@ protected AbstractPullDiagnosticHandler(
                     // same-result-id) response to the client as that means they should just preserve the current
                     // diagnostics they have for this file.
                     var previousParams = documentToPreviousDiagnosticParams[diagnosticSource.GetId()];
-                    progress.Report(CreateReport(previousParams.TextDocument, diagnostics: null, previousParams.PreviousResultId));
+                    progress.Report(CreateUnchangedReport(previousParams.TextDocument, previousParams.PreviousResultId));
                 }
             }
 
@@ -209,6 +218,12 @@ private static Dictionary<ProjectOrDocumentId, PreviousPullResult> GetIdToPrevio
                     return new ProjectOrDocumentId(project.Id);
                 }
 
+                var additionalDocument = solution.GetAdditionalDocument(textDocumentIdentifier);
+                if (additionalDocument != null)
+                {
+                    return new ProjectOrDocumentId(additionalDocument.Id);
+                }
+
                 return null;
             }
         }
@@ -254,7 +269,7 @@ private void HandleRemovedDocuments(RequestContext context, ImmutableArray<Previ
                 // the workspace). Report a (null-diagnostics, null-result-id) response to the client as that
                 // means they should just consider the file deleted and should remove all diagnostics
                 // information they've cached for it.
-                progress.Report(CreateReport(removedResult.TextDocument, diagnostics: null, resultId: null));
+                progress.Report(CreateRemovedReport(removedResult.TextDocument));
             }
         }