Complete WrapWithTag cohosting implementation with text edit fixing

Co-authored-by: davidwengier <[email protected]>

Author: Copilot

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

@@ -5,10 +5,13 @@
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.Protocol;
+using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
 internal interface IRemoteWrapWithTagService
 {
     ValueTask<RemoteResponse<bool>> IsValidWrapWithTagLocationAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, LspRange range, CancellationToken cancellationToken);
+    
+    ValueTask<RemoteResponse<TextEdit[]?>> FixHtmlTextEditsAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, TextEdit[] textEdits, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/WrapWithTag/RemoteWrapWithTagService.cs

@@ -6,12 +6,14 @@
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Microsoft.CodeAnalysis.Razor.Formatting;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
 using Response = Microsoft.CodeAnalysis.Razor.Remote.RemoteResponse<bool>;
+using TextEditResponse = Microsoft.CodeAnalysis.Razor.Remote.RemoteResponse<Microsoft.CodeAnalysis.Text.TextEdit[]?>;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
@@ -34,6 +36,17 @@ public ValueTask<Response> IsValidWrapWithTagLocationAsync(
             context => IsValidWrapWithTagLocationAsync(context, range, cancellationToken),
             cancellationToken);
 
+    public ValueTask<TextEditResponse> FixHtmlTextEditsAsync(
+        RazorPinnedSolutionInfoWrapper solutionInfo,
+        DocumentId razorDocumentId,
+        TextEdit[] textEdits,
+        CancellationToken cancellationToken)
+        => RunServiceAsync(
+            solutionInfo,
+            razorDocumentId,
+            context => FixHtmlTextEditsAsync(context, textEdits, cancellationToken),
+            cancellationToken);
+
     private async ValueTask<Response> IsValidWrapWithTagLocationAsync(
         RemoteDocumentContext context,
         LspRange range,
@@ -114,4 +127,21 @@ private async ValueTask<Response> IsValidWrapWithTagLocationAsync(
 
         return Response.Results(true);
     }
+
+    private async ValueTask<TextEditResponse> FixHtmlTextEditsAsync(
+        RemoteDocumentContext context,
+        TextEdit[] textEdits,
+        CancellationToken cancellationToken)
+    {
+        if (textEdits.Length == 0)
+        {
+            return TextEditResponse.Results(textEdits);
+        }
+
+        var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
+        var htmlSourceText = codeDocument.GetHtmlSourceText();
+
+        var fixedEdits = FormattingUtilities.FixHtmlTextEdits(htmlSourceText, textEdits);
+        return TextEditResponse.Results(fixedEdits);
+    }
 }

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

@@ -78,9 +78,19 @@ public ImmutableArray<Registration> GetRegistrations(VSInternalClientCapabilitie
             request,
             cancellationToken).ConfigureAwait(false);
 
-        // TODO: Consider if we need to fix HTML text edits in the cohost scenario
-        // The original language server implementation calls FormattingUtilities.FixHtmlTextEdits
-        // but this might not be necessary in the cohost context
+        if (htmlResponse?.TextEdits is not null)
+        {
+            // Fix the HTML text edits to handle any tilde characters in the generated document
+            var fixedEdits = await _remoteServiceInvoker.TryInvokeAsync<IRemoteWrapWithTagService, RemoteResponse<TextEdit[]?>>(
+                razorDocument.Project.Solution,
+                (service, solutionInfo, cancellationToken) => service.FixHtmlTextEditsAsync(solutionInfo, razorDocument.Id, htmlResponse.TextEdits, cancellationToken),
+                cancellationToken).ConfigureAwait(false);
+
+            if (fixedEdits.Result is not null)
+            {
+                htmlResponse.TextEdits = fixedEdits.Result;
+            }
+        }
 
         return htmlResponse;
     }