Rework how we get generated documents

In light of "should we get rid of document context" convo this morning, figured this made sense to do (and I needed _something_ on IDocumentSnapshot in order to actually make this whole idea work)

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentHighlight/RemoteDocumentHighlightService.cs

@@ -68,7 +68,7 @@ internal sealed partial class RemoteDocumentHighlightService(in ServiceArgs args
         var csharpDocument = codeDocument.GetCSharpDocument();
         if (_documentMappingService.TryMapToGeneratedDocumentPosition(csharpDocument, index, out var mappedPosition, out _))
         {
-            var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+            var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
             var highlights = await DocumentHighlights.GetHighlightsAsync(generatedDocument, mappedPosition, cancellationToken).ConfigureAwait(false);
 

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/FoldingRanges/RemoteFoldingRangeService.cs

@@ -43,7 +43,7 @@ internal sealed class RemoteFoldingRangeService(in ServiceArgs args) : RazorDocu
         ImmutableArray<RemoteFoldingRange> htmlRanges,
         CancellationToken cancellationToken)
     {
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var csharpRanges = await ExternalHandlers.FoldingRanges.GetFoldingRangesAsync(generatedDocument, cancellationToken).ConfigureAwait(false);
 

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/GoToDefinition/RemoteGoToDefinitionService.cs

@@ -96,7 +96,7 @@ internal sealed class RemoteGoToDefinitionService(in ServiceArgs args) : RazorDo
         }
 
         // Finally, call into C#.
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var locations = await ExternalHandlers.GoToDefinition
             .GetDefinitionsAsync(

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/InlayHints/RemoteInlayHintService.cs

@@ -52,7 +52,7 @@ internal sealed partial class RemoteInlayHintService(in ServiceArgs args) : Razo
             return null;
         }
 
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var textDocument = inlayHintParams.TextDocument.WithUri(generatedDocument.CreateUri());
         var range = projectedLinePositionSpan.ToRange();
@@ -106,7 +106,7 @@ internal sealed partial class RemoteInlayHintService(in ServiceArgs args) : Razo
 
     private async ValueTask<InlayHint> ResolveInlayHintAsync(RemoteDocumentContext context, InlayHint inlayHint, CancellationToken cancellationToken)
     {
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         return await InlayHints.ResolveInlayHintAsync(generatedDocument, inlayHint, cancellationToken).ConfigureAwait(false);
     }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/ProjectSystem/DocumentContextExtensions.cs

@@ -1,48 +0,0 @@
-// Copyright (c) .NET Foundation. All rights reserved.
-// Licensed under the MIT license. See License.txt in the project root for license information.
-
-using System.Diagnostics;
-using System.Linq;
-using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.AspNetCore.Razor;
-using Microsoft.CodeAnalysis.Razor.ProjectSystem;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
-using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
-
-namespace Microsoft.CodeAnalysis.Remote.Razor;
-
-internal static class DocumentContextExtensions
-{
-    public static async Task<Document> GetGeneratedDocumentAsync(
-        this VersionedDocumentContext documentContext,
-        IFilePathService filePathService,
-        CancellationToken cancellationToken)
-    {
-        Debug.Assert(documentContext.Snapshot is RemoteDocumentSnapshot, "This method only works on document contexts created in the OOP process");
-
-        var snapshot = (RemoteDocumentSnapshot)documentContext.Snapshot;
-
-        return await snapshot.GetOrAddGeneratedDocumentAsync(
-            (snapshot, documentContext, filePathService, cancellationToken),
-            static async arg =>
-            {
-                var (snapshot, documentContext, filePathService, cancellationToken) = arg;
-
-                var razorDocument = snapshot.TextDocument;
-                var projectKey = snapshot.Project.Key;
-                var solution = razorDocument.Project.Solution;
-
-                // TODO: A real implementation needs to get the SourceGeneratedDocument from the solution
-
-                var generatedFilePath = filePathService.GetRazorCSharpFilePath(projectKey, razorDocument.FilePath.AssumeNotNull());
-                var generatedDocumentId = solution.GetDocumentIdsWithFilePath(generatedFilePath).First(d => d.ProjectId == razorDocument.Project.Id);
-                var generatedDocument = solution.GetRequiredDocument(generatedDocumentId);
-
-                var csharpSourceText = await documentContext.GetCSharpSourceTextAsync(cancellationToken).ConfigureAwait(false);
-
-                // HACK: We're not in the same solution fork as the LSP server that provides content for this document
-                return generatedDocument.WithText(csharpSourceText);
-            });
-    }
-}

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/ProjectSystem/RemoteDocumentContext.cs

@@ -8,7 +8,9 @@ namespace Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 
 internal class RemoteDocumentContext : VersionedDocumentContext
 {
-    public TextDocument TextDocument => ((RemoteDocumentSnapshot)Snapshot).TextDocument;
+    public TextDocument TextDocument => Snapshot.TextDocument;
+
+    public new RemoteDocumentSnapshot Snapshot => (RemoteDocumentSnapshot)base.Snapshot;
 
     public RemoteDocumentContext(Uri uri, RemoteDocumentSnapshot snapshot)
         // HACK: Need to revisit version and projectContext here I guess

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/ProjectSystem/RemoteDocumentSnapshot.cs

@@ -1,13 +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 System;
+using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
@@ -88,14 +89,31 @@ internal class RemoteDocumentSnapshot(TextDocument textDocument, RemoteProjectSn
         return result is not null;
     }
 
-    public async Task<Document> GetOrAddGeneratedDocumentAsync<TArg>(TArg arg, Func<TArg, Task<Document>> createGeneratedDocument)
+    public async Task<Document> GetGeneratedDocumentAsync(IFilePathService filePathService)
     {
         if (_generatedDocument is Document generatedDocument)
         {
             return generatedDocument;
         }
 
-        generatedDocument = await createGeneratedDocument(arg);
+        generatedDocument = await HACK_GenerateDocumentAsync(filePathService).ConfigureAwait(false);
         return InterlockedOperations.Initialize(ref _generatedDocument, generatedDocument);
     }
+
+    private async Task<Document> HACK_GenerateDocumentAsync(IFilePathService filePathService)
+    {
+        // TODO: A real implementation needs to get the SourceGeneratedDocument from the solution
+
+        var solution = TextDocument.Project.Solution;
+        var generatedFilePath = filePathService.GetRazorCSharpFilePath(Project.Key, FilePath.AssumeNotNull());
+        var projectId = TextDocument.Project.Id;
+        var generatedDocumentId = solution.GetDocumentIdsWithFilePath(generatedFilePath).First(d => d.ProjectId == projectId);
+        var generatedDocument = solution.GetRequiredDocument(generatedDocumentId);
+
+        var codeDocument = await this.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var csharpSourceText = codeDocument.GetCSharpSourceText();
+
+        // HACK: We're not in the same solution fork as the LSP server that provides content for this document
+        return generatedDocument.WithText(csharpSourceText);
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Rename/RemoteRenameService.cs

@@ -53,7 +53,7 @@ internal sealed class RemoteRenameService(in ServiceArgs args) : RazorDocumentSe
             return NoFurtherHandling;
         }
 
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var razorEdit = await _renameService.TryGetRazorRenameEditsAsync(context, positionInfo, newName, cancellationToken).ConfigureAwait(false);
         if (razorEdit is not null)

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/SemanticTokens/RemoteCSharpSemanticTokensProvider.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Immutable;
 using System.Composition;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Telemetry;
@@ -11,6 +12,7 @@ using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.SemanticTokens;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor.SemanticTokens;
@@ -27,7 +29,9 @@ internal class RemoteCSharpSemanticTokensProvider(IFilePathService filePathServi
         using var _ = _telemetryReporter.TrackLspRequest(nameof(SemanticTokensRange.GetSemanticTokensAsync), Constants.ExternalAccessServerName, correlationId);
 
         // We have a razor document, lets find the generated C# document
-        var generatedDocument = await documentContext.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        Debug.Assert(documentContext is RemoteDocumentContext, "This method only works on document snapshots created in the OOP process");
+        var snapshot = (RemoteDocumentSnapshot)documentContext.Snapshot;
+        var generatedDocument = await snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         var data = await SemanticTokensRange.GetSemanticTokensAsync(
             generatedDocument,

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/SignatureHelp/RemoteSignatureHelpService.cs

@@ -11,11 +11,10 @@ using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.LanguageServer.Protocol;
 using ExternalHandlers = Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers;
+using LspSignatureHelp = Roslyn.LanguageServer.Protocol.SignatureHelp;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
-using SignatureHelp = Roslyn.LanguageServer.Protocol.SignatureHelp;
-
 internal sealed class RemoteSignatureHelpService(in ServiceArgs args) : RazorDocumentServiceBase(in args), IRemoteSignatureHelpService
 {
     internal sealed class Factory : FactoryBase<IRemoteSignatureHelpService>
@@ -26,20 +25,20 @@ internal sealed class RemoteSignatureHelpService(in ServiceArgs args) : RazorDoc
 
     private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
 
-    public ValueTask<SignatureHelp?> GetSignatureHelpAsync(JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo, JsonSerializableDocumentId documentId, Position position, CancellationToken cancellationToken)
+    public ValueTask<LspSignatureHelp?> GetSignatureHelpAsync(JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo, JsonSerializableDocumentId documentId, Position position, CancellationToken cancellationToken)
         => RunServiceAsync(
             solutionInfo,
             documentId,
             context => GetSignatureHelpsAsync(context, position, cancellationToken),
             cancellationToken);
 
-    private async ValueTask<SignatureHelp?> GetSignatureHelpsAsync(RemoteDocumentContext context, Position position, CancellationToken cancellationToken)
+    private async ValueTask<LspSignatureHelp?> GetSignatureHelpsAsync(RemoteDocumentContext context, Position position, CancellationToken cancellationToken)
     {
         var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
         var linePosition = new LinePosition(position.Line, position.Character);
         var absoluteIndex = codeDocument.Source.Text.GetRequiredAbsoluteIndex(linePosition);
 
-        var generatedDocument = await context.GetGeneratedDocumentAsync(_filePathService, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(_filePathService).ConfigureAwait(false);
 
         if (DocumentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), absoluteIndex, out var mappedPosition, out _))
         {