Ensure diagnostics refresh when source generators run in balanced mode (

#77271)

Fixes an issue where diagnostics would not refresh when source
generators run in balanced mode (e.g. after a save triggers sg re-run,
diagnostics would not update based on the new sg contents).

checksum changes taken from #77273

src/LanguageServer/Protocol.TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

@@ -810,6 +810,12 @@ internal async Task WaitForDiagnosticsAsync()
                 await listenerProvider.GetWaiter(FeatureAttribute.DiagnosticService).ExpeditedWaitAsync();
             }
 
+            internal async Task WaitForSourceGeneratorsAsync()
+            {
+                var operations = TestWorkspace.ExportProvider.GetExportedValue<AsynchronousOperationListenerProvider>();
+                await operations.WaitAllAsync(TestWorkspace, [FeatureAttribute.Workspace, FeatureAttribute.SourceGenerators]);
+            }
+
             internal RequestExecutionQueue<RequestContext>.TestAccessor? GetQueueAccessor() => _languageServer.Value.GetTestAccessor().GetQueueAccessor();
 
             internal LspWorkspaceManager.TestAccessor GetManagerAccessor() => GetRequiredLspService<LspWorkspaceManager>().GetTestAccessor();

src/LanguageServer/Protocol/Handler/Diagnostics/DiagnosticsPullCache.cs

@@ -15,7 +15,7 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler.Diagnostics;
 internal abstract partial class AbstractPullDiagnosticHandler<TDiagnosticsParams, TReport, TReturn>
     where TDiagnosticsParams : IPartialResultParams<TReport>
 {
-    internal record struct DiagnosticsRequestState(Project Project, int GlobalStateVersion, RequestContext Context, IDiagnosticSource DiagnosticSource);
+    internal readonly record struct DiagnosticsRequestState(Project Project, int GlobalStateVersion, RequestContext Context, IDiagnosticSource DiagnosticSource);
 
     /// <summary>
     /// Cache where we store the data produced by prior requests so that they can be returned if nothing of significance 

src/LanguageServer/ProtocolUnitTests/Diagnostics/PullDiagnosticTests.cs

@@ -11,12 +11,14 @@
 using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.EditAndContinue;
 using Microsoft.CodeAnalysis.EditAndContinue.UnitTests;
+using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.LanguageServer.Handler.Diagnostics;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.SolutionCrawler;
 using Microsoft.CodeAnalysis.TaskList;
+using Microsoft.CodeAnalysis.Testing;
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.LanguageServer.Protocol;
 using Roslyn.Test.Utilities;
@@ -632,6 +634,94 @@ public async Task TestDocumentDiagnosticsFromLiveShareServer(bool useVSDiagnosti
         Assert.NotNull(results.Single().Diagnostics!.Single().CodeDescription!.Href);
     }
 
+    [Theory, CombinatorialData]
+    internal async Task TestDocumentDiagnosticsUpdatesSourceGeneratorDiagnostics(bool useVSDiagnostics, bool mutatingLspWorkspace, SourceGeneratorExecutionPreference executionPreference)
+    {
+        var markup = """
+            public {|insert:|}class C
+            {
+                public int F => _field;
+            }
+            """;
+
+        await using var testLspServer = await CreateTestWorkspaceWithDiagnosticsAsync(
+            markup, mutatingLspWorkspace, BackgroundAnalysisScope.OpenFiles, useVSDiagnostics);
+
+        // Set execution mode preference.
+        var configService = testLspServer.TestWorkspace.ExportProvider.GetExportedValue<TestWorkspaceConfigurationService>();
+        configService.Options = new WorkspaceConfigurationOptions(SourceGeneratorExecution: executionPreference);
+
+        var document = testLspServer.GetCurrentSolution().Projects.Single().Documents.Single();
+
+        // Add a generator to the solution that reports a diagnostic if the text matches original markup
+        var generator = new CallbackGenerator(onInit: (_) => { }, onExecute: context =>
+        {
+            var tree = context.Compilation.SyntaxTrees.Single();
+            var text = tree.GetText(CancellationToken.None).ToString();
+            if (text.Contains("public partial class C"))
+            {
+                context.AddSource("blah", """
+                    public partial class C
+                    {
+                        private readonly int _field = 1;
+                    }
+                    """);
+            }
+        });
+
+        testLspServer.TestWorkspace.OnAnalyzerReferenceAdded(
+            document.Project.Id,
+            new TestGeneratorReference(generator));
+        await testLspServer.WaitForSourceGeneratorsAsync();
+
+        await OpenDocumentAsync(testLspServer, document);
+
+        // First diagnostic request should report a diagnostic since the generator does not produce any source (text does not match).
+        var results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, document.GetURI(), useVSDiagnostics);
+        var diagnostic = AssertEx.Single(results.Single().Diagnostics);
+        Assert.Equal("CS0103", diagnostic.Code);
+
+        // Update the text to include the text trigger the source generator relies on.
+        var textLocation = testLspServer.GetLocations()["insert"].Single();
+
+        var textEdit = new LSP.TextEdit { Range = textLocation.Range, NewText = "partial " };
+        if (mutatingLspWorkspace)
+        {
+            // In the mutating workspace, we just need to update the LSP text (which will flow into the workspace).
+            await testLspServer.ReplaceTextAsync(textLocation.Uri, (textEdit.Range, textEdit.NewText));
+            await WaitForWorkspaceOperationsAsync(testLspServer.TestWorkspace);
+        }
+        else
+        {
+            // In the non-mutating workspace, we need to ensure that both the workspace and LSP view of the world are updated.
+            var workspaceText = await document.GetTextAsync(CancellationToken.None);
+            var textChange = ProtocolConversions.TextEditToTextChange(textEdit, workspaceText);
+            await testLspServer.TestWorkspace.ChangeDocumentAsync(document.Id, workspaceText.WithChanges(textChange));
+            await testLspServer.ReplaceTextAsync(textLocation.Uri, (textEdit.Range, textEdit.NewText));
+        }
+
+        await testLspServer.WaitForSourceGeneratorsAsync();
+        results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, document.GetURI(), useVSDiagnostics);
+
+        if (executionPreference == SourceGeneratorExecutionPreference.Automatic)
+        {
+            // In automatic mode, the diagnostic should be immediately removed as the source generator ran and produced the required field.
+            Assert.Empty(results.Single().Diagnostics!);
+        }
+        else
+        {
+            // In balanced mode, the diagnostic should remain until there is a manual source generator run that updates the sg text.
+            diagnostic = AssertEx.Single(results.Single().Diagnostics);
+            Assert.Equal("CS0103", diagnostic.Code);
+
+            testLspServer.TestWorkspace.EnqueueUpdateSourceGeneratorVersion(document.Project.Id, forceRegeneration: false);
+            await testLspServer.WaitForSourceGeneratorsAsync();
+
+            results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, document.GetURI(), useVSDiagnostics);
+            Assert.Empty(results.Single().Diagnostics!);
+        }
+    }
+
     [Theory, CombinatorialData]
     public async Task TestDocumentDiagnosticsIncludesSourceGeneratorDiagnostics(bool useVSDiagnostics, bool mutatingLspWorkspace)
     {
@@ -641,7 +731,10 @@ public async Task TestDocumentDiagnosticsIncludesSourceGeneratorDiagnostics(bool
 
         var document = testLspServer.GetCurrentSolution().Projects.Single().Documents.Single();
 
-        var generator = new DiagnosticProducingGenerator(context => Location.Create(context.Compilation.SyntaxTrees.Single(), new TextSpan(0, 10)));
+        var generator = new DiagnosticProducingGenerator(context =>
+        {
+            return Location.Create(context.Compilation.SyntaxTrees.Single(), new TextSpan(0, 10));
+        });
 
         testLspServer.TestWorkspace.OnAnalyzerReferenceAdded(
             document.Project.Id,

src/LanguageServer/ProtocolUnitTests/Workspaces/SourceGeneratedDocumentTests.cs

@@ -171,7 +171,7 @@ internal async Task TestReturnsGeneratedSourceWhenDocumentChanges(bool mutatingL
         // In automatic mode this should trigger generators to re-run.
         // In balanced mode generators should not re-run.
         await testLspServer.TestWorkspace.ChangeDocumentAsync(testLspServer.TestWorkspace.Documents.Single(d => !d.IsSourceGenerated).Id, SourceText.From("new text"));
-        await WaitForSourceGeneratorsAsync(testLspServer.TestWorkspace);
+        await testLspServer.WaitForSourceGeneratorsAsync();
 
         // Ask for the source generated text again.
         var secondRequest = await testLspServer.ExecuteRequestAsync<SourceGeneratorGetTextParams, SourceGeneratedDocumentText>(SourceGeneratedDocumentGetTextHandler.MethodName,
@@ -216,7 +216,7 @@ internal async Task TestReturnsGeneratedSourceWhenManuallyRefreshed(bool mutatin
 
         // Updating the execution version should trigger source generators to run in both automatic and balanced mode.
         testLspServer.TestWorkspace.EnqueueUpdateSourceGeneratorVersion(projectId: null, forceRegeneration: true);
-        await WaitForSourceGeneratorsAsync(testLspServer.TestWorkspace);
+        await testLspServer.WaitForSourceGeneratorsAsync();
 
         var secondRequest = await testLspServer.ExecuteRequestAsync<SourceGeneratorGetTextParams, SourceGeneratedDocumentText>(SourceGeneratedDocumentGetTextHandler.MethodName,
             new SourceGeneratorGetTextParams(new LSP.TextDocumentIdentifier { Uri = sourceGeneratorDocumentUri }, ResultId: text.ResultId), CancellationToken.None);
@@ -281,12 +281,6 @@ public async Task TestReturnsNullForRemovedOpenedGeneratedFile(bool mutatingLspW
         Assert.Null(secondRequest.Text);
     }
 
-    private static async Task WaitForSourceGeneratorsAsync(LspTestWorkspace workspace)
-    {
-        var operations = workspace.ExportProvider.GetExportedValue<AsynchronousOperationListenerProvider>();
-        await operations.WaitAllAsync(workspace, [FeatureAttribute.Workspace, FeatureAttribute.SourceGenerators]);
-    }
-
     private async Task<TestLspServer> CreateTestLspServerWithGeneratorAsync(bool mutatingLspWorkspace, string generatedDocumentText)
     {
         var testLspServer = await CreateTestLspServerAsync(string.Empty, mutatingLspWorkspace);

src/Workspaces/Core/Portable/Workspace/Solution/Project.cs

@@ -568,6 +568,8 @@ public Task<VersionStamp> GetSemanticVersionAsync(CancellationToken cancellation
     ///        <see cref="Project.GetDependentSemanticVersionAsync(CancellationToken)"/> changes as the project is removed, then added resulting in a version change.</item>
     /// </list>   
     /// </para>
+    /// This checksum is also affected by the <see cref="SourceGeneratorExecutionVersion"/> for this project.
+    /// As such, it is not usable across different sessions of a particular host.
     /// </remarks>
     internal Task<Checksum> GetDependentChecksumAsync(CancellationToken cancellationToken)
         => Solution.CompilationState.GetDependentChecksumAsync(this.Id, cancellationToken);

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.RegularCompilationTracker.cs

@@ -1096,40 +1096,50 @@ public Task<Checksum> GetDependentChecksumAsync(
                     Interlocked.CompareExchange(
                         ref _lazyDependentChecksum,
                         AsyncLazy.Create(static (arg, c) =>
-                            arg.self.ComputeDependentChecksumAsync(arg.SolutionState, c),
-                            arg: (self: this, compilationState.SolutionState)),
+                            arg.self.ComputeDependentChecksumAsync(arg.compilationState, c),
+                            arg: (self: this, compilationState)),
                         null);
                 }
 
                 return _lazyDependentChecksum.GetValueAsync(cancellationToken);
             }
 
-            private async Task<Checksum> ComputeDependentChecksumAsync(SolutionState solution, CancellationToken cancellationToken)
+            private async Task<Checksum> ComputeDependentChecksumAsync(
+                SolutionCompilationState solution, CancellationToken cancellationToken)
             {
                 using var _ = ArrayBuilder<Checksum>.GetInstance(out var tempChecksumArray);
 
+                // Mix in the SG information for this project.  That way if it changes, we will have a different
+                // checksum (since semantics could have changed because of this).
+                if (solution.SourceGeneratorExecutionVersionMap.Map.TryGetValue(this.ProjectState.Id, out var executionVersion))
+                    tempChecksumArray.Add(executionVersion.Checksum);
+
                 // Get the checksum for the project itself.
                 var projectChecksum = await this.ProjectState.GetChecksumAsync(cancellationToken).ConfigureAwait(false);
                 tempChecksumArray.Add(projectChecksum);
 
-                // Calculate a checksum this project and for each dependent project that could affect semantics for
-                // this project. Ensure that the checksum calculation orders the projects consistently so that we get
-                // the same checksum across sessions of VS.  Note: we use the project filepath+name as a unique way
-                // to reference a project.  This matches the logic in our persistence-service implemention as to how
-                // information is associated with a project.
-                var transitiveDependencies = solution.GetProjectDependencyGraph().GetProjectsThatThisProjectTransitivelyDependsOn(this.ProjectState.Id);
+                // Calculate a checksum this project and for each dependent project that could affect semantics for this
+                // project. We order the projects so that we are resilient to the underlying in-memory graph structure
+                // changing this arbitrarily.  We do not want that to cause us to change our semantic version.. Note: we
+                // use the project filepath+name as a unique way to reference a project.  This matches the logic in our
+                // persistence-service implementation as to how information is associated with a project.
+                var transitiveDependencies = solution.SolutionState.GetProjectDependencyGraph().GetProjectsThatThisProjectTransitivelyDependsOn(this.ProjectState.Id);
                 var orderedProjectIds = transitiveDependencies.OrderBy(id =>
                 {
-                    var depProject = solution.GetRequiredProjectState(id);
+                    var depProject = solution.SolutionState.GetRequiredProjectState(id);
                     return (depProject.FilePath, depProject.Name);
                 });
 
                 foreach (var projectId in orderedProjectIds)
                 {
-                    var referencedProject = solution.GetRequiredProjectState(projectId);
+                    // Mix in the SG information for the dependent project.  That way if it changes, we will have a
+                    // different checksum (since semantics could have changed because of this).
+                    if (solution.SourceGeneratorExecutionVersionMap.Map.TryGetValue(projectId, out executionVersion))
+                        tempChecksumArray.Add(executionVersion.Checksum);
 
                     // Note that these checksums should only actually be calculated once, if the project is unchanged
                     // the same checksum will be returned.
+                    var referencedProject = solution.SolutionState.GetRequiredProjectState(projectId);
                     var referencedProjectChecksum = await referencedProject.GetChecksumAsync(cancellationToken).ConfigureAwait(false);
                     tempChecksumArray.Add(referencedProjectChecksum);
                 }

src/Workspaces/Core/Portable/Workspace/Solution/SourceGeneratorExecutionVersion.cs

@@ -46,6 +46,8 @@ public static SourceGeneratorExecutionVersion ReadFrom(ObjectReader reader)
 
     public override string ToString()
         => $"{MajorVersion}.{MinorVersion}";
+
+    public Checksum Checksum => new(MajorVersion, MinorVersion);
 }
 
 /// <summary>