Cohosting: Update Html virtual document from OOP

eng/targets/Services.props

@@ -16,6 +16,7 @@
   -->
   <ItemGroup>
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.SemanticTokens" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteSemanticTokensServiceFactory" />
+    <ServiceHubService Include="Microsoft.VisualStudio.Razor.HtmlDocument" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteHtmlDocumentServiceFactory" />
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.TagHelperProvider" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteTagHelperProviderServiceFactory"/>
     <ServiceHubService Include="Microsoft.VisualStudio.Razor.ClientInitialization" ClassName="Microsoft.CodeAnalysis.Remote.Razor.RemoteClientInitializationServiceFactory" />
   </ItemGroup>

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

@@ -8,21 +8,16 @@
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer;
 
-internal class GeneratedDocumentSynchronizer : DocumentProcessedListener
+internal class GeneratedDocumentSynchronizer(
+    IGeneratedDocumentPublisher publisher,
+    IDocumentVersionCache documentVersionCache,
+    ProjectSnapshotManagerDispatcher dispatcher,
+    LanguageServerFeatureOptions languageServerFeatureOptions) : DocumentProcessedListener
 {
-    private readonly IGeneratedDocumentPublisher _publisher;
-    private readonly IDocumentVersionCache _documentVersionCache;
-    private readonly ProjectSnapshotManagerDispatcher _dispatcher;
-
-    public GeneratedDocumentSynchronizer(
-        IGeneratedDocumentPublisher publisher,
-        IDocumentVersionCache documentVersionCache,
-        ProjectSnapshotManagerDispatcher dispatcher)
-    {
-        _publisher = publisher;
-        _documentVersionCache = documentVersionCache;
-        _dispatcher = dispatcher;
-    }
+    private readonly IGeneratedDocumentPublisher _publisher = publisher;
+    private readonly IDocumentVersionCache _documentVersionCache = documentVersionCache;
+    private readonly ProjectSnapshotManagerDispatcher _dispatcher = dispatcher;
+    private readonly LanguageServerFeatureOptions _languageServerFeatureOptions = languageServerFeatureOptions;
 
     public override void Initialize(IProjectSnapshotManager projectManager)
     {
@@ -40,9 +35,13 @@ public override void DocumentProcessed(RazorCodeDocument codeDocument, IDocument
 
         var filePath = document.FilePath.AssumeNotNull();
 
-        var htmlText = codeDocument.GetHtmlSourceText();
+        // If cohosting is on, then it is responsible for updating the Html buffer
+        if (!_languageServerFeatureOptions.UseRazorCohostServer)
+        {
+            var htmlText = codeDocument.GetHtmlSourceText();
 
-        _publisher.PublishHtml(document.Project.Key, filePath, htmlText, hostDocumentVersion.Value);
+            _publisher.PublishHtml(document.Project.Key, filePath, htmlText, hostDocumentVersion.Value);
+        }
 
         var csharpText = codeDocument.GetCSharpSourceText();
 

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteHtmlDocumentService.cs

@@ -0,0 +1,13 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+
+namespace Microsoft.CodeAnalysis.Razor.Remote;
+
+internal interface IRemoteHtmlDocumentService
+{
+    ValueTask<string?> GetHtmlDocumentTextAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, CancellationToken cancellationToken);
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/RazorServices.cs

@@ -22,5 +22,6 @@ internal static class RazorServices
             (typeof(IRemoteTagHelperProviderService), null),
             (typeof(IRemoteClientInitializationService), null),
             (typeof(IRemoteSemanticTokensService), null),
+            (typeof(IRemoteHtmlDocumentService), null),
         ]);
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/HtmlDocuments/RemoteHtmlDocumentService.cs

@@ -0,0 +1,42 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Api;
+using Microsoft.CodeAnalysis.Razor.Remote;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
+using Microsoft.ServiceHub.Framework;
+
+namespace Microsoft.CodeAnalysis.Remote.Razor;
+
+internal sealed class RemoteHtmlDocumentService(
+    IServiceBroker serviceBroker,
+    DocumentSnapshotFactory documentSnapshotFactory)
+    : RazorServiceBase(serviceBroker), IRemoteHtmlDocumentService
+{
+    private readonly DocumentSnapshotFactory _documentSnapshotFactory = documentSnapshotFactory;
+
+    public ValueTask<string?> GetHtmlDocumentTextAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, CancellationToken cancellationToken)
+       => RazorBrokeredServiceImplementation.RunServiceAsync(
+            solutionInfo,
+            ServiceBrokerClient,
+            solution => GetHtmlDocumentTextAsync(solution, razorDocumentId, cancellationToken),
+            cancellationToken);
+
+    private async ValueTask<string?> GetHtmlDocumentTextAsync(Solution solution, DocumentId razorDocumentId, CancellationToken _)
+    {
+        var razorDocument = solution.GetAdditionalDocument(razorDocumentId);
+        if (razorDocument is null)
+        {
+            return null;
+        }
+
+        var documentSnapshot = _documentSnapshotFactory.GetOrCreate(razorDocument);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync();
+
+        return codeDocument.GetHtmlSourceText().ToString();
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/HtmlDocuments/RemoteHtmlDocumentServiceFactory.cs

@@ -0,0 +1,24 @@
+// 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.Razor.Remote;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
+using Microsoft.ServiceHub.Framework;
+using Microsoft.VisualStudio.Composition;
+
+namespace Microsoft.CodeAnalysis.Remote.Razor;
+
+internal sealed class RemoteHtmlDocumentServiceFactory : RazorServiceFactoryBase<IRemoteHtmlDocumentService>
+{
+    // WARNING: We must always have a parameterless constructor in order to be properly handled by ServiceHub.
+    public RemoteHtmlDocumentServiceFactory()
+        : base(RazorServices.Descriptors)
+    {
+    }
+
+    protected override IRemoteHtmlDocumentService CreateService(IServiceBroker serviceBroker, ExportProvider exportProvider)
+    {
+        var documentSnapshotFactory = exportProvider.GetExportedValue<DocumentSnapshotFactory>();
+        return new RemoteHtmlDocumentService(serviceBroker, documentSnapshotFactory);
+    }
+}

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/Cohost/CohostTextDocumentSyncHandler.cs

@@ -0,0 +1,85 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System;
+using System.Composition;
+using System.Diagnostics;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
+using Microsoft.CodeAnalysis.Razor.Logging;
+using Microsoft.CodeAnalysis.Razor.Remote;
+using Microsoft.Extensions.Logging;
+using Microsoft.VisualStudio.LanguageServer.ContainedLanguage;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
+using Microsoft.VisualStudio.Threading;
+
+namespace Microsoft.VisualStudio.LanguageServerClient.Razor.Cohost;
+
+[Export(typeof(IRazorCohostTextDocumentSyncHandler)), Shared]
+[method: ImportingConstructor]
+internal class CohostTextDocumentSyncHandler(
+    IRemoteClientProvider remoteClientProvider,
+    LSPDocumentManager documentManager,
+    JoinableTaskContext joinableTaskContext,
+    IRazorLoggerFactory loggerFactory) : IRazorCohostTextDocumentSyncHandler
+{
+    private readonly IRemoteClientProvider _remoteClientProvider = remoteClientProvider;
+    private readonly JoinableTaskContext _joinableTaskContext = joinableTaskContext;
+    private readonly TrackingLSPDocumentManager _documentManager = documentManager as TrackingLSPDocumentManager ?? throw new InvalidOperationException("Expected TrackingLSPDocumentManager");
+    private readonly ILogger _logger = loggerFactory.CreateLogger<CohostTextDocumentSyncHandler>();
+
+    public async Task HandleAsync(int version, RazorCohostRequestContext context, CancellationToken cancellationToken)
+    {
+        // For didClose, we don't need to do anything. We can't close the virtual document, because that requires the buffer
+        // so we just no-op and let our VS components handle closure.
+        if (context.Method == Methods.TextDocumentDidCloseName)
+        {
+            return;
+        }
+
+        var textDocument = context.TextDocument.AssumeNotNull();
+        var textDocumentPath = context.TextDocument.FilePath.AssumeNotNull();
+
+        _logger.LogDebug("[Cohost] {method} of '{document}' with version {version}.", context.Method, textDocumentPath, version);
+
+        var client = await _remoteClientProvider.TryGetClientAsync(cancellationToken);
+        if (client is null)
+        {
+            _logger.LogError("[Cohost] Couldn't get remote client for {method} of '{document}'. Html document contents will be stale", context.Method, textDocumentPath);
+            return;
+        }
+
+        var htmlText = await client.TryInvokeAsync<IRemoteHtmlDocumentService, string?>(textDocument.Project.Solution,
+            (service, solutionInfo, ct) => service.GetHtmlDocumentTextAsync(solutionInfo, textDocument.Id, ct),
+            cancellationToken).ConfigureAwait(false);
+
+        if (!htmlText.HasValue || htmlText.Value is null)
+        {
+            _logger.LogError("[Cohost] Couldn't get Html text for {method} of '{document}'. Html document contents will be stale", context.Method, textDocumentPath);
+            return;
+        }
+
+        // Eventually, for VS Code, the following piece of logic needs to make an LSP call rather than directly update the buffer
+
+        // Roslyn might have got changes from the LSP server, but we just get the actual source text, so we just construct one giant change
+        // from that. Guaranteed no sync issues, though we are passing long strings around unfortunately.
+        var uri = textDocument.CreateUri();
+        if (!_documentManager.TryGetDocument(uri, out var documentSnapshot) ||
+            !documentSnapshot.TryGetVirtualDocument<HtmlVirtualDocumentSnapshot>(out var htmlDocument))
+        {
+            Debug.Fail("Got an LSP text document update before getting informed of the VS buffer");
+            _logger.LogError("[Cohost] Couldn't get an Html document for {method} of '{document}'. Html document contents will be stale (or non-existent?)", context.Method, textDocumentPath);
+            return;
+        }
+
+        await _joinableTaskContext.Factory.SwitchToMainThreadAsync(cancellationToken);
+
+        VisualStudioTextChange[] changes = [new(0, htmlDocument.Snapshot.Length, htmlText.Value)];
+        _documentManager.UpdateVirtualDocument<HtmlVirtualDocument>(uri, changes, version, state: null);
+
+        _logger.LogDebug("[Cohost] Exiting {method} of '{document}' with version {version}.", context.Method, textDocumentPath, version);
+    }
+}

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/GeneratedDocumentSynchronizerTest.cs

@@ -6,6 +6,7 @@
 using Microsoft.AspNetCore.Razor.Test.Common;
 using Microsoft.AspNetCore.Razor.Test.Common.LanguageServer;
 using Microsoft.AspNetCore.Razor.Test.Common.ProjectSystem;
+using Microsoft.AspNetCore.Razor.Test.Common.Workspaces;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
 using Xunit;
@@ -27,7 +28,7 @@ public GeneratedDocumentSynchronizerTest(ITestOutputHelper testOutput)
         var projectManager = StrictMock.Of<IProjectSnapshotManager>();
         _cache = new DocumentVersionCache(projectManager);
         _publisher = new TestGeneratedDocumentPublisher();
-        _synchronizer = new GeneratedDocumentSynchronizer(_publisher, _cache, Dispatcher);
+        _synchronizer = new GeneratedDocumentSynchronizer(_publisher, _cache, Dispatcher, TestLanguageServerFeatureOptions.Instance);
         _document = TestDocumentSnapshot.Create("C:/path/to/file.razor");
         _codeDocument = CreateCodeDocument("<p>Hello World</p>");
     }