Fix provisional completion corrupting other requests

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/GeneratedDocumentPublisher.cs

@@ -98,10 +98,11 @@ public void PublishCSharp(ProjectKey projectKey, string filePath, SourceText sou
         {
             HostDocumentFilePath = filePath,
             ProjectKeyId = projectKey.Id,
-            Changes = textChanges.Select(static t => t.ToRazorTextChange()).ToArray(),
+            Changes = [.. textChanges.Select(static t => t.ToRazorTextChange())],
             HostDocumentVersion = hostDocumentVersion,
+            PreviousHostDocumentVersion = previouslyPublishedData.HostDocumentVersion,
             PreviousWasEmpty = previouslyPublishedData.SourceText.Length == 0,
-            Checksum = Convert.ToBase64String(sourceText.GetChecksum().ToArray()),
+            Checksum = Convert.ToBase64String([.. sourceText.GetChecksum()]),
             ChecksumAlgorithm = sourceText.ChecksumAlgorithm,
             SourceEncodingCodePage = sourceText.Encoding?.CodePage
         };
@@ -145,10 +146,10 @@ public void PublishHtml(ProjectKey projectKey, string filePath, SourceText sourc
         {
             HostDocumentFilePath = filePath,
             ProjectKeyId = projectKey.Id,
-            Changes = textChanges.Select(static t => t.ToRazorTextChange()).ToArray(),
+            Changes = [.. textChanges.Select(static t => t.ToRazorTextChange())],
             HostDocumentVersion = hostDocumentVersion,
             PreviousWasEmpty = previouslyPublishedData.SourceText.Length == 0,
-            Checksum = Convert.ToBase64String(sourceText.GetChecksum().ToArray()),
+            Checksum = Convert.ToBase64String([.. sourceText.GetChecksum()]),
             ChecksumAlgorithm = sourceText.ChecksumAlgorithm,
             SourceEncodingCodePage = sourceText.Encoding?.CodePage
         };

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Protocol/UpdateBufferRequest.cs

@@ -11,6 +11,9 @@ internal sealed class UpdateBufferRequest
     [JsonPropertyName("hostDocumentVersion")]
     public int? HostDocumentVersion { get; set; }
 
+    [JsonPropertyName("previousHostDocumentVersion")]
+    public int? PreviousHostDocumentVersion { get; set; }
+
     [JsonPropertyName("projectKeyId")]
     public string? ProjectKeyId { get; set; }
 

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Endpoints/RazorCustomMessageTarget_Completion.cs

@@ -6,6 +6,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.PooledObjects;
+using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Protocol.Completion;
 using Microsoft.CodeAnalysis.Razor.Workspaces.Telemetry;
@@ -109,7 +110,11 @@ internal partial class RazorCustomMessageTarget
             await _joinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
 
             var provisionalChange = new VisualStudioTextChange(provisionalTextEdit, virtualDocumentSnapshot.Snapshot);
-            UpdateVirtualDocument(provisionalChange, request.ProjectedKind, request.Identifier.Version, hostDocumentUri, virtualDocumentSnapshot.Uri);
+            // We update to a negative version number so that if a request comes in for v6, it won't see our modified document. We revert the version back
+            // later, don't worry.
+            UpdateVirtualDocument(provisionalChange, request.ProjectedKind, -1 * request.Identifier.Version, hostDocumentUri, virtualDocumentSnapshot.Uri);
+
+            _logger.LogDebug($"Updated for provisional completion to version -{request.Identifier.Version} of {virtualDocumentSnapshot!.Uri}.");
 
             // We want the delegation to continue on the captured context because we're currently on the `main` thread and we need to get back to the
             // main thread in order to update the virtual buffer with the reverted text edit.
@@ -163,6 +168,8 @@ internal partial class RazorCustomMessageTarget
         {
             if (provisionalTextEdit is not null)
             {
+                _logger.LogDebug($"Reverting the update for provisional completion back to {request.Identifier.Version} of {virtualDocumentSnapshot!.Uri}.");
+
                 var revertedProvisionalTextEdit = BuildRevertedEdit(provisionalTextEdit);
                 var revertedProvisionalChange = new VisualStudioTextChange(revertedProvisionalTextEdit, virtualDocumentSnapshot.Snapshot);
                 UpdateVirtualDocument(revertedProvisionalChange, request.ProjectedKind, request.Identifier.Version, hostDocumentUri, virtualDocumentSnapshot.Uri);

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Endpoints/RazorCustomMessageTarget_UpdateCSharpBuffer.cs

@@ -9,6 +9,7 @@
 using Microsoft.AspNetCore.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Protocol;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
 using StreamJsonRpc;
 
 namespace Microsoft.VisualStudio.Razor.LanguageClient.Endpoints;
@@ -24,20 +25,33 @@ public async Task UpdateCSharpBufferAsync(UpdateBufferRequest request, Cancellat
             throw new ArgumentNullException(nameof(request));
         }
 
-        await _joinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-
-        await UpdateCSharpBufferCoreAsync(request, cancellationToken);
+        // We're going to try updating the C# buffer, and it has to happen on the UI thread. That is a single shared resource
+        // that can only be updated on a specific thread, and so we can easily hit contention. In particular with provisional
+        // completion we can end up with a state where we are waiting for completion to finish before we can update the document
+        // but other features are waiting on us to update the document with some changes. Normally this is fine, and we handle it
+        // with our sync system and simply cancel the task and let the next one in. When updating buffers from the server though
+        // we can't do that, as the server assumes we can always apply the update. Our only option here is to keep trying until
+        // we get a successful update.
+        // It's worth noting we only try again specifically for provisional completion, so we don't get random deadocks.
+        var tryAgain = true;
+        while (tryAgain)
+        {
+            _logger.LogDebug($"Trying a call to UpdateCSharpBufferCoreAsync for v{request.HostDocumentVersion}");
+            tryAgain = await UpdateCSharpBufferCoreAsync(request, cancellationToken).ConfigureAwait(false);
+        }
     }
 
-    // Internal for testing
-    internal async Task UpdateCSharpBufferCoreAsync(UpdateBufferRequest request, CancellationToken cancellationToken)
+    private async Task<bool> UpdateCSharpBufferCoreAsync(UpdateBufferRequest request, CancellationToken cancellationToken)
     {
-        if (request is null || request.HostDocumentFilePath is null || request.HostDocumentVersion is null)
+        if (request.HostDocumentFilePath is null || request.HostDocumentVersion is null)
         {
-            return;
+            return false;
         }
 
-        var hostDocumentUri = new Uri(request.HostDocumentFilePath);
+        await _joinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+
+        var identifier = CreateTextDocumentIdentifier(request);
+        var hostDocumentUri = identifier.Uri;
 
         _logger.LogDebug($"UpdateCSharpBuffer for {request.HostDocumentVersion} of {hostDocumentUri} in {request.ProjectKeyId}");
 
@@ -68,12 +82,27 @@ request.ProjectKeyId is not null &&
                     // Sadly there isn't anything we can do here to, we're just in a state where the server and client are out of
                     // sync with their understanding of the document contents, and since changes come in as a list of changes,
                     // the user experience is broken. All we can do is hope the user closes and re-opens the document.
-                    Debug.Fail($"Server wants to update {hostDocumentUri} in {request.ProjectKeyId} but we don't know about the document being in any projects");
                     _logger.LogError($"Server wants to update {hostDocumentUri} in {request.ProjectKeyId} by we only know about that document in misc files. Server and client are now out of sync.");
-                    return;
+                    Debug.Fail($"Server wants to update {hostDocumentUri} in {request.ProjectKeyId} but we don't know about the document being in any projects");
+                    return false;
                 }
             }
 
+            // First we need to make sure we're synced to the previous version, or the changes won't apply properly. This should no-op in most cases, as this
+            // is (almost) the only thing that actually moves documents forward, we're really just validating we're in a good state.
+            // We're specifically checking here for provisional completion in flight, which is a case where we update the C# document
+            // in a way that doesn't represent the actual state, so we can't let server updates through while in this state (represented
+            // by a negative version number). Due to provisional completion being on the UI thread, we have to return from this method
+            // and try again so we get to the back of the UI thread queue.
+            if (request.PreviousHostDocumentVersion is { } previousVersion &&
+                await TrySynchronizeVirtualDocumentAsync<CSharpVirtualDocumentSnapshot>(previousVersion, identifier, cancellationToken, rejectOnNewerParallelRequest: false) is { } synchronizedResult &&
+                !synchronizedResult.Synchronized &&
+                synchronizedResult.VirtualSnapshot?.HostDocumentSyncVersion < 0)
+            {
+                _logger.LogError($"Request to update C# buffer from {previousVersion} to {request.HostDocumentVersion} failed because the server Roslyn and Razor are out of sync. Server version is {synchronizedResult.VirtualSnapshot?.HostDocumentSyncVersion}. Will try again as provisional completion is in flight which is an expected cause of de-sync, which we recover from.");
+                return true;
+            }
+
             foreach (var virtualDocument in virtualDocuments)
             {
                 if (virtualDocument.ProjectKey.Equals(new ProjectKey(request.ProjectKeyId)))
@@ -89,7 +118,7 @@ request.ProjectKeyId is not null &&
 
                     _logger.LogDebug($"UpdateCSharpBuffer finished updating doc for {request.HostDocumentVersion} of {virtualDocument.Uri}. New lines: {GetLineCountOfVirtualDocument(hostDocumentUri, virtualDocument)}");
 
-                    return;
+                    return false;
                 }
             }
 
@@ -108,7 +137,7 @@ request.ProjectKeyId is not null &&
 
             // Don't know about document, no-op. This can happen if the language server found a project.razor.bin from an old build
             // and is sending us updates.
-            return;
+            return false;
         }
 
         _logger.LogDebug($"UpdateCSharpBuffer fallback for {request.HostDocumentVersion} of {hostDocumentUri}");
@@ -118,6 +147,26 @@ request.ProjectKeyId is not null &&
             request.Changes.Select(change => change.ToVisualStudioTextChange()).ToArray(),
             request.HostDocumentVersion.Value,
             state: request.PreviousWasEmpty);
+
+        return false;
+    }
+
+    private static TextDocumentIdentifier CreateTextDocumentIdentifier(UpdateBufferRequest request)
+    {
+        var hostDocumentUri = new Uri(request.HostDocumentFilePath);
+        if (request.ProjectKeyId is { } id)
+        {
+            return new VSTextDocumentIdentifier
+            {
+                Uri = hostDocumentUri,
+                ProjectContext = new VSProjectContext
+                {
+                    Id = id,
+                }
+            };
+        }
+
+        return new TextDocumentIdentifier { Uri = hostDocumentUri };
     }
 
     private int GetLineCountOfVirtualDocument(Uri hostDocumentUri, CSharpVirtualDocumentSnapshot virtualDocument)

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/LanguageClient/RazorCustomMessageTargetTest.cs

@@ -42,7 +42,7 @@ public RazorCustomMessageTargetTest(ITestOutputHelper testOutput)
         _editorSettingsManager = new ClientSettingsManager(Array.Empty<IClientSettingsChangedTrigger>());
     }
 
-    [Fact]
+    [UIFact]
     public async Task UpdateCSharpBuffer_CannotLookupDocument_NoopsGracefully()
     {
         // Arrange
@@ -78,10 +78,10 @@ public async Task UpdateCSharpBuffer_CannotLookupDocument_NoopsGracefully()
         };
 
         // Act & Assert
-        await target.UpdateCSharpBufferCoreAsync(request, DisposalToken);
+        await target.UpdateCSharpBufferAsync(request, DisposalToken);
     }
 
-    [Fact]
+    [UIFact]
     public async Task UpdateCSharpBuffer_UpdatesDocument()
     {
         // Arrange
@@ -124,13 +124,13 @@ public async Task UpdateCSharpBuffer_UpdatesDocument()
         };
 
         // Act
-        await target.UpdateCSharpBufferCoreAsync(request, DisposalToken);
+        await target.UpdateCSharpBufferAsync(request, DisposalToken);
 
         // Assert
         documentManager.VerifyAll();
     }
 
-    [Fact]
+    [UIFact]
     public async Task UpdateCSharpBuffer_UpdatesCorrectDocument()
     {
         // Arrange
@@ -142,8 +142,8 @@ public async Task UpdateCSharpBuffer_UpdatesCorrectDocument()
         var document = Mock.Of<LSPDocumentSnapshot>(d => d.VirtualDocuments == documents, MockBehavior.Strict);
         var documentManager = new Mock<TrackingLSPDocumentManager>(MockBehavior.Strict);
         documentManager
-          .Setup(manager => manager.TryGetDocument(It.IsAny<Uri>(), out document))
-          .Returns(true);
+            .Setup(manager => manager.TryGetDocument(It.IsAny<Uri>(), out document))
+            .Returns(true);
         documentManager
             .Setup(manager => manager.UpdateVirtualDocument<CSharpVirtualDocument>(
                 It.IsAny<Uri>(),
@@ -181,7 +181,7 @@ public async Task UpdateCSharpBuffer_UpdatesCorrectDocument()
         };
 
         // Act
-        await target.UpdateCSharpBufferCoreAsync(request, DisposalToken);
+        await target.UpdateCSharpBufferAsync(request, DisposalToken);
 
         // Assert
         documentManager.VerifyAll();

src/Razor/test/Microsoft.VisualStudio.Razor.IntegrationTests/AbstractRazorEditorTest.cs

@@ -141,6 +141,9 @@ private static string CreateTemporaryPath()
 
     public override async Task DisposeAsync()
     {
+        // TODO: Would be good to have this as a last ditch check, but need to improve the detection and reporting here to be more robust
+        //await TestServices.Editor.ValidateNoDiscoColorsAsync(HangMitigatingCancellationToken);
+
         _testLogger!.LogInformation($"#### Razor integration test dispose.");
 
         TestServices.Output.ClearIntegrationTestLogger();

src/Razor/test/Microsoft.VisualStudio.Razor.IntegrationTests/CompletionIntegrationTests.cs

@@ -294,33 +294,6 @@ protected override void OnAfterRender(bool firstRender)
             commitChar: '\t');
     }
 
-    private async Task VerifyTypeAndCommitCompletionAsync(string input, string output, string search, string[] stringsToType, char? commitChar = null, string? expectedSelectedItemLabel = null)
-    {
-        await TestServices.SolutionExplorer.AddFileAsync(
-            RazorProjectConstants.BlazorProjectName,
-            "Test.razor",
-            input,
-            open: true,
-            ControlledHangMitigatingCancellationToken);
-
-        await TestServices.Editor.WaitForComponentClassificationAsync(ControlledHangMitigatingCancellationToken);
-
-        await TestServices.Editor.PlaceCaretAsync(search, charsOffset: 1, ControlledHangMitigatingCancellationToken);
-        foreach (var stringToType in stringsToType)
-        {
-            TestServices.Input.Send(stringToType);
-        }
-
-        if (expectedSelectedItemLabel is not null)
-        {
-            await CommitCompletionAndVerifyAsync(output, expectedSelectedItemLabel, commitChar);
-        }
-        else
-        {
-            await CommitCompletionAndVerifyAsync(output, commitChar);
-        }
-    }
-
     [IdeFact]
     public async Task SnippetCompletion_DoesntCommitOnSpace()
     {
@@ -477,6 +450,78 @@ public enum MyEnum
             """);
     }
 
+    [IdeFact, WorkItem("https://github.com/dotnet/razor/issues/11385")]
+    public async Task ProvisionalCompletion_DoesntBreakSemanticTokens()
+    {
+        await TestServices.SolutionExplorer.AddFileAsync(
+            RazorProjectConstants.BlazorProjectName,
+            "Test.razor",
+            """
+            @page "/counter"
+
+            <PageTitle>Counter</PageTitle>
+
+            <h1>Counter</h1>
+
+            <p role="status">Current count: @currentCount</p>
+
+            @DateTime
+
+            <button class="btn btn-primary" @onclick="IncrementCount">Click me</button>
+
+            @code {
+                private int currentCount = 0;
+
+                public void IncrementCount()
+                {
+                    currentCount++;
+                }
+            }
+            """,
+            open: true,
+            ControlledHangMitigatingCancellationToken);
+
+        await TestServices.Editor.WaitForComponentClassificationAsync(ControlledHangMitigatingCancellationToken);
+
+        await TestServices.Editor.PlaceCaretAsync("@DateTime", charsOffset: 1, ControlledHangMitigatingCancellationToken);
+
+        await Task.Delay(500, HangMitigatingCancellationToken);
+
+        TestServices.Input.Send(".");
+        TestServices.Input.Send("n");
+
+        await Task.Delay(500, HangMitigatingCancellationToken);
+
+        await TestServices.Editor.ValidateNoDiscoColorsAsync(HangMitigatingCancellationToken);
+    }
+
+    private async Task VerifyTypeAndCommitCompletionAsync(string input, string output, string search, string[] stringsToType, char? commitChar = null, string? expectedSelectedItemLabel = null)
+    {
+        await TestServices.SolutionExplorer.AddFileAsync(
+            RazorProjectConstants.BlazorProjectName,
+            "Test.razor",
+            input,
+            open: true,
+            ControlledHangMitigatingCancellationToken);
+
+        await TestServices.Editor.WaitForComponentClassificationAsync(ControlledHangMitigatingCancellationToken);
+
+        await TestServices.Editor.PlaceCaretAsync(search, charsOffset: 1, ControlledHangMitigatingCancellationToken);
+        foreach (var stringToType in stringsToType)
+        {
+            TestServices.Input.Send(stringToType);
+        }
+
+        if (expectedSelectedItemLabel is not null)
+        {
+            await CommitCompletionAndVerifyAsync(output, expectedSelectedItemLabel, commitChar);
+        }
+        else
+        {
+            await CommitCompletionAndVerifyAsync(output, commitChar);
+        }
+    }
+
     private async Task CommitCompletionAndVerifyAsync(string expected, char? commitChar = null)
     {
         var session = await TestServices.Editor.WaitForCompletionSessionAsync(HangMitigatingCancellationToken);