Improve document handling in VS Code and cohosting

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/HtmlDocumentServices/HtmlDocumentSynchronizer.SynchronizationRequest.cs

@@ -14,30 +14,30 @@ internal sealed partial class HtmlDocumentSynchronizer
     private class SynchronizationRequest(RazorDocumentVersion requestedVersion) : IDisposable
     {
         private readonly RazorDocumentVersion _requestedVersion = requestedVersion;
-        private readonly TaskCompletionSource<bool> _tcs = new();
+        private readonly TaskCompletionSource<SynchronizationResult> _tcs = new();
         private CancellationTokenSource? _cts;
 
-        public Task<bool> Task => _tcs.Task;
+        public Task<SynchronizationResult> Task => _tcs.Task;
 
         public RazorDocumentVersion RequestedVersion => _requestedVersion;
 
-        internal static SynchronizationRequest CreateAndStart(TextDocument document, RazorDocumentVersion requestedVersion, Func<TextDocument, CancellationToken, Task<bool>> syncFunction)
+        internal static SynchronizationRequest CreateAndStart(TextDocument document, RazorDocumentVersion requestedVersion, Func<TextDocument, RazorDocumentVersion, CancellationToken, Task<SynchronizationResult>> syncFunction)
         {
             var request = new SynchronizationRequest(requestedVersion);
             request.Start(document, syncFunction);
             return request;
         }
 
-        private void Start(TextDocument document, Func<TextDocument, CancellationToken, Task<bool>> syncFunction)
+        private void Start(TextDocument document, Func<TextDocument, RazorDocumentVersion, CancellationToken, Task<SynchronizationResult>> syncFunction)
         {
             _cts = new(TimeSpan.FromMinutes(1));
             _cts.Token.Register(Dispose);
-            _ = syncFunction.Invoke(document, _cts.Token).ContinueWith((t, state) =>
+            _ = syncFunction.Invoke(document, _requestedVersion, _cts.Token).ContinueWith((t, state) =>
             {
-                var tcs = (TaskCompletionSource<bool>)state.AssumeNotNull();
+                var tcs = (TaskCompletionSource<SynchronizationResult>)state.AssumeNotNull();
                 if (t.IsCanceled)
                 {
-                    tcs.SetResult(false);
+                    tcs.SetResult(default);
                 }
                 else if (t.Exception is { } ex)
                 {
@@ -58,7 +58,7 @@ public void Dispose()
             _cts?.Cancel();
             _cts?.Dispose();
             _cts = null;
-            _tcs.TrySetResult(false);
+            _tcs.TrySetResult(default);
         }
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/HtmlDocumentServices/HtmlDocumentSynchronizer.cs

@@ -42,7 +42,7 @@ public void DocumentRemoved(Uri razorFileUri)
         }
     }
 
-    public async Task<bool> TrySynchronizeAsync(TextDocument document, CancellationToken cancellationToken)
+    public async Task<SynchronizationResult> TrySynchronizeAsync(TextDocument document, CancellationToken cancellationToken)
     {
         var requestedVersion = await RazorDocumentVersion.CreateAsync(document, cancellationToken).ConfigureAwait(false);
 
@@ -56,7 +56,7 @@ public async Task<bool> TrySynchronizeAsync(TextDocument document, CancellationT
         return await GetSynchronizationRequestTaskAsync(document, requestedVersion).ConfigureAwait(false);
     }
 
-    private Task<bool> GetSynchronizationRequestTaskAsync(TextDocument document, RazorDocumentVersion requestedVersion)
+    private Task<SynchronizationResult> GetSynchronizationRequestTaskAsync(TextDocument document, RazorDocumentVersion requestedVersion)
     {
         lock (_gate)
         {
@@ -71,7 +71,7 @@ private Task<bool> GetSynchronizationRequestTaskAsync(TextDocument document, Raz
 
 #pragma warning disable VSTHRD103 // Use await instead of .Result
 #pragma warning disable VSTHRD003 // Avoid awaiting foreign Tasks
-                    if (request.Task.IsCompleted && request.Task.Result == false)
+                    if (request.Task.IsCompleted && request.Task.Result.Equals(default))
                     {
                         _logger.LogDebug($"Already finished that version for {document.FilePath}, but was unsuccessful, so will recompute");
                         request.Dispose();
@@ -92,7 +92,7 @@ private Task<bool> GetSynchronizationRequestTaskAsync(TextDocument document, Raz
                     // for a different checksum, but the same workspace version, we assume the new request is the newer document.
 
                     _logger.LogDebug($"We've already seen {request.RequestedVersion} for {document.FilePath} so that's a no from me");
-                    return SpecializedTasks.False;
+                    return SpecializedTasks.Default<SynchronizationResult>();
                 }
                 else if (!request.Task.IsCompleted)
                 {
@@ -110,7 +110,7 @@ private Task<bool> GetSynchronizationRequestTaskAsync(TextDocument document, Raz
         }
     }
 
-    private async Task<bool> PublishHtmlDocumentAsync(TextDocument document, CancellationToken cancellationToken)
+    private async Task<SynchronizationResult> PublishHtmlDocumentAsync(TextDocument document, RazorDocumentVersion requestedVersion, CancellationToken cancellationToken)
     {
         string? htmlText;
         try
@@ -122,30 +122,31 @@ private async Task<bool> PublishHtmlDocumentAsync(TextDocument document, Cancell
         catch (Exception ex)
         {
             _logger.LogError(ex, $"Error getting Html text for {document.FilePath}. Html document contents will be stale");
-            return false;
+            return default;
         }
 
         if (cancellationToken.IsCancellationRequested)
         {
             // Checking cancellation before logging, as a new request coming in doesn't count as "Couldn't get Html"
-            return false;
+            return default;
         }
 
         if (htmlText is null)
         {
             _logger.LogError($"Couldn't get Html text for {document.FilePath}. Html document contents will be stale");
-            return false;
+            return default;
         }
 
         try
         {
-            await _htmlDocumentPublisher.PublishAsync(document, htmlText, cancellationToken).ConfigureAwait(false);
-            return true;
+            var result = new SynchronizationResult(true, requestedVersion.Checksum);
+            await _htmlDocumentPublisher.PublishAsync(document, result, htmlText, cancellationToken).ConfigureAwait(false);
+            return result;
         }
         catch (Exception ex)
         {
             _logger.LogError(ex, $"Error publishing Html text for {document.FilePath}. Html document contents will be stale");
-            return false;
+            return default;
         }
     }
 
@@ -163,7 +164,7 @@ internal TestAccessor(HtmlDocumentSynchronizer instance)
             _instance = instance;
         }
 
-        public Task<bool> GetSynchronizationRequestTaskAsync(TextDocument document, RazorDocumentVersion requestedVersion)
+        public Task<SynchronizationResult> GetSynchronizationRequestTaskAsync(TextDocument document, RazorDocumentVersion requestedVersion)
             => _instance.GetSynchronizationRequestTaskAsync(document, requestedVersion);
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/HtmlDocumentServices/IHtmlDocumentPublisher.cs

@@ -9,5 +9,5 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
 internal interface IHtmlDocumentPublisher
 {
-    Task PublishAsync(TextDocument document, string htmlText, CancellationToken cancellationToken);
+    Task PublishAsync(TextDocument document, SynchronizationResult synchronizationResult, string htmlText, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/HtmlDocumentServices/IHtmlDocumentSynchronizer.cs

@@ -11,5 +11,5 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 internal interface IHtmlDocumentSynchronizer
 {
     void DocumentRemoved(Uri uri);
-    Task<bool> TrySynchronizeAsync(TextDocument document, CancellationToken cancellationToken);
+    Task<SynchronizationResult> TrySynchronizeAsync(TextDocument document, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/HtmlDocumentServices/SynchronizationResult.cs

@@ -0,0 +1,14 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+
+namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
+
+/// <summary>
+/// A result of a synchronization operation for a Html document
+/// </summary>
+/// <remarks>
+/// If <see cref="Synchronized" /> is <see langword="false" />, <see cref="Checksum" /> will be <see langword="default" />.
+/// </remarks>
+internal readonly record struct SynchronizationResult(bool Synchronized, ChecksumWrapper Checksum);

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/Microsoft.CodeAnalysis.Razor.CohostingShared.projitems

@@ -24,6 +24,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)DocumentSymbol\CohostDocumentSymbolEndpoint.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Formatting\CohostDocumentFormattingEndpoint.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)FindAllReferences\CohostFindAllReferencesEndpoint.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)HtmlDocumentServices\SynchronizationResult.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Navigation\CohostGoToDefinitionEndpoint.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Navigation\CohostGoToImplementationEndpoint.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)LinkedEditingRange\CohostLinkedEditingRangeEndpoint.cs" />

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Cohost/HtmlDocumentPublisher.cs

@@ -6,6 +6,7 @@
 using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.Logging;
@@ -25,11 +26,9 @@ internal sealed class HtmlDocumentPublisher(
     private readonly TrackingLSPDocumentManager _documentManager = documentManager as TrackingLSPDocumentManager ?? throw new InvalidOperationException("Expected TrackingLSPDocumentManager");
     private readonly ILogger _logger = loggerFactory.GetOrCreateLogger<HtmlDocumentPublisher>();
 
-    public async Task PublishAsync(TextDocument document, string htmlText, CancellationToken cancellationToken)
+    public async Task PublishAsync(TextDocument document, SynchronizationResult synchronizationResult, string htmlText, CancellationToken cancellationToken)
     {
-        // TODO: Eventually, for VS Code, the following piece of logic needs to make an LSP call rather than directly update the
-        // buffer, but the assembly this code currently lives in doesn't ship in VS Code, so we need to solve a few other things
-        // before we get there.
+        Assumed.True(synchronizationResult.Synchronized);
 
         var uri = document.CreateUri();
         if (!_documentManager.TryGetDocument(uri, out var documentSnapshot) ||
@@ -55,7 +54,7 @@ public async Task PublishAsync(TextDocument document, string htmlText, Cancellat
         }
 
         VisualStudioTextChange[] changes = [new(0, htmlDocument.Snapshot.Length, htmlText)];
-        _documentManager.UpdateVirtualDocument<HtmlVirtualDocument>(uri, changes, documentSnapshot.Version, state: null);
+        _documentManager.UpdateVirtualDocument<HtmlVirtualDocument>(uri, changes, documentSnapshot.Version, state: synchronizationResult.Checksum.ToString());
 
         _logger.LogDebug($"Finished Html document generation for {document.FilePath} (into {uri})");
     }

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Cohost/HtmlRequestInvoker.cs

@@ -6,6 +6,7 @@
 using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.Logging;
@@ -32,7 +33,7 @@ internal sealed class HtmlRequestInvoker(
     public async Task<TResponse?> MakeHtmlLspRequestAsync<TRequest, TResponse>(TextDocument razorDocument, string method, TRequest request, TimeSpan threshold, Guid correlationId, CancellationToken cancellationToken) where TRequest : notnull
     {
         var syncResult = await _htmlDocumentSynchronizer.TrySynchronizeAsync(razorDocument, cancellationToken).ConfigureAwait(false);
-        if (!syncResult)
+        if (!syncResult.Synchronized)
         {
             _logger.LogDebug($"Couldn't synchronize for {razorDocument.FilePath}");
             return default;
@@ -50,6 +51,12 @@ internal sealed class HtmlRequestInvoker(
             return default;
         }
 
+        if ((string)htmlDocument.State.AssumeNotNull() != syncResult.Checksum.ToString())
+        {
+            _logger.LogError($"Checksum for {snapshot.Uri}, {htmlDocument.State} doesn't match {syncResult.Checksum}.");
+            return default;
+        }
+
         // If the request is for a text document, we need to update the Uri to point to the Html document,
         // and most importantly set it back again before leaving the method in case a caller uses it.
         Uri? originalUri = null;
@@ -62,8 +69,7 @@ internal sealed class HtmlRequestInvoker(
 
         try
         {
-
-            _logger.LogDebug($"Making LSP request for {method} from {htmlDocument.Uri}{(request is ITextDocumentPositionParams positionParams ? $" at {positionParams.Position}" : "")}.");
+            _logger.LogDebug($"Making LSP request for {method} from {htmlDocument.Uri}{(request is ITextDocumentPositionParams positionParams ? $" at {positionParams.Position}" : "")}, checksum {syncResult.Checksum}.");
 
             // Passing Guid.Empty to this method will mean no tracking
             using var _ = _telemetryReporter.TrackLspRequest(Methods.TextDocumentCodeActionName, RazorLSPConstants.HtmlLanguageServerName, threshold, correlationId);

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/HtmlVirtualDocument.cs

@@ -10,5 +10,5 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient;
 internal class HtmlVirtualDocument(Uri uri, ITextBuffer textBuffer, ITelemetryReporter telemetryReporter)
     : GeneratedVirtualDocument<HtmlVirtualDocumentSnapshot>(uri, textBuffer, telemetryReporter)
 {
-    protected override HtmlVirtualDocumentSnapshot GetUpdatedSnapshot(object? state) => new(Uri, TextBuffer.CurrentSnapshot, HostDocumentVersion);
+    protected override HtmlVirtualDocumentSnapshot GetUpdatedSnapshot(object? state) => new(Uri, TextBuffer.CurrentSnapshot, HostDocumentVersion, state);
 }

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/HtmlVirtualDocumentSnapshot.cs

@@ -12,7 +12,8 @@ internal class HtmlVirtualDocumentSnapshot : VirtualDocumentSnapshot
     public HtmlVirtualDocumentSnapshot(
         Uri uri,
         ITextSnapshot snapshot,
-        long? hostDocumentSyncVersion)
+        long? hostDocumentSyncVersion,
+        object? state)
     {
         if (uri is null)
         {
@@ -27,11 +28,14 @@ public HtmlVirtualDocumentSnapshot(
         Uri = uri;
         Snapshot = snapshot;
         HostDocumentSyncVersion = hostDocumentSyncVersion;
+        State = state;
     }
 
     public override Uri Uri { get; }
 
     public override ITextSnapshot Snapshot { get; }
 
     public override long? HostDocumentSyncVersion { get; }
+
+    public object? State { get; }
 }

src/Razor/src/Microsoft.VisualStudioCode.RazorExtension/Endpoints/RazorDocumentClosedEndpoint.cs

@@ -0,0 +1,35 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Composition;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor;
+using Microsoft.AspNetCore.Razor.Threading;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Features;
+using Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
+
+namespace Microsoft.VisualStudioCode.RazorExtension.Endpoints;
+
+[Shared]
+[ExportRazorStatelessLspService(typeof(RazorDocumentClosedEndpoint))]
+[RazorEndpoint("razor/documentClosed")]
+[method: ImportingConstructor]
+internal class RazorDocumentClosedEndpoint(IHtmlDocumentSynchronizer htmlDocumentSynchronizer) : AbstractRazorCohostDocumentRequestHandler<TextDocumentIdentifier, VoidResult>
+{
+    private readonly IHtmlDocumentSynchronizer _htmlDocumentSynchronizer = htmlDocumentSynchronizer;
+
+    protected override bool MutatesSolutionState => false;
+
+    protected override bool RequiresLSPSolution => true;
+
+    protected override RazorTextDocumentIdentifier? GetRazorTextDocumentIdentifier(TextDocumentIdentifier request)
+        => request.ToRazorTextDocumentIdentifier();
+
+    protected override Task<VoidResult> HandleRequestAsync(TextDocumentIdentifier textDocument, RazorCohostRequestContext requestContext, CancellationToken cancellationToken)
+    {
+        _htmlDocumentSynchronizer.DocumentRemoved(requestContext.Uri.AssumeNotNull());
+        return SpecializedTasks.Default<VoidResult>();
+    }
+}

src/Razor/src/Microsoft.VisualStudioCode.RazorExtension/Services/HtmlDocumentPublisher.cs

@@ -20,9 +20,11 @@ internal sealed class HtmlDocumentPublisher(
 {
     private readonly RazorClientServerManagerProvider _razorClientServerManagerProvider = razorClientServerManagerProvider;
 
-    public async Task PublishAsync(TextDocument document, string htmlText, CancellationToken cancellationToken)
+    public async Task PublishAsync(TextDocument document, SynchronizationResult synchronizationResult, string htmlText, CancellationToken cancellationToken)
     {
-        var request = new HtmlUpdateParameters(new TextDocumentIdentifier { Uri = document.CreateUri() }, htmlText);
+        Assumed.True(synchronizationResult.Synchronized);
+
+        var request = new HtmlUpdateParameters(new TextDocumentIdentifier { Uri = document.CreateUri() }, synchronizationResult.Checksum.ToString(), htmlText);
 
         var clientConnection = _razorClientServerManagerProvider.ClientLanguageServerManager.AssumeNotNull();
         await clientConnection.SendRequestAsync("razor/updateHtml", request, cancellationToken).ConfigureAwait(false);
@@ -31,6 +33,8 @@ public async Task PublishAsync(TextDocument document, string htmlText, Cancellat
     private record HtmlUpdateParameters(
         [property: JsonPropertyName("textDocument")]
         TextDocumentIdentifier TextDocument,
+        [property: JsonPropertyName("checksum")]
+        string Checksum,
         [property: JsonPropertyName("text")]
         string Text);
 }

src/Razor/src/Microsoft.VisualStudioCode.RazorExtension/Services/HtmlRequestInvoker.cs

@@ -3,11 +3,13 @@
 
 using System;
 using System.ComponentModel.Composition;
+using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.LanguageServer;
 using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
@@ -27,14 +29,30 @@ internal sealed class HtmlRequestInvoker(
     public async Task<TResponse?> MakeHtmlLspRequestAsync<TRequest, TResponse>(TextDocument razorDocument, string method, TRequest request, TimeSpan threshold, Guid correlationId, CancellationToken cancellationToken) where TRequest : notnull
     {
         var syncResult = await _htmlDocumentSynchronizer.TrySynchronizeAsync(razorDocument, cancellationToken).ConfigureAwait(false);
-        if (!syncResult)
+        if (!syncResult.Synchronized)
         {
             return default;
         }
 
-        _logger.LogDebug($"Making Html request for {method} on {razorDocument.FilePath}");
+        _logger.LogDebug($"Making Html request for {method} on {razorDocument.FilePath}, checksum {syncResult.Checksum}");
+
+        var forwardedRequest = new HtmlForwardedRequest<TRequest>(
+            new TextDocumentIdentifier
+            {
+                Uri = razorDocument.CreateUri()
+            },
+            syncResult.Checksum.ToString(),
+            request);
 
         var clientConnection = _razorClientServerManagerProvider.ClientLanguageServerManager.AssumeNotNull();
-        return await clientConnection.SendRequestAsync<TRequest, TResponse>(method, request, cancellationToken).ConfigureAwait(false);
+        return await clientConnection.SendRequestAsync<HtmlForwardedRequest<TRequest>, TResponse>(method, forwardedRequest, cancellationToken).ConfigureAwait(false);
     }
+
+    private record HtmlForwardedRequest<TRequest>(
+        [property: JsonPropertyName("textDocument")]
+        TextDocumentIdentifier TextDocument,
+        [property: JsonPropertyName("checksum")]
+        string Checksum,
+        [property: JsonPropertyName("request")]
+        TRequest Request);
 }