Allow LSP and cohosting to provide specialized methods to get a syntax tree (#10765)

This PR is a classic self-nerd-snipe, and resolves this comment:
https://github.com/dotnet/razor/pull/10750#discussion_r1719087058

Also inadvertently found a slight "bug" with the existing go to def code
in cohosting. Bug is in quotes because the actual user behaviour
probably was identical in 99% of cases.

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Definition/RazorComponentDefinitionService.cs

@@ -1,11 +1,6 @@
 // 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.Linq;
-using System.Text;
-using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.GoToDefinition;
 using Microsoft.CodeAnalysis.Razor.Logging;

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/GoToDefinition/AbstractRazorComponentDefinitionService.cs

@@ -69,10 +69,8 @@ internal abstract class AbstractRazorComponentDefinitionService(
         {
             _logger.LogInformation($"Attempting to get definition from an attribute directly.");
 
-            var originCodeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
-
             var range = await RazorComponentDefinitionHelpers
-                .TryGetPropertyRangeAsync(originCodeDocument, attributeDescriptor.GetPropertyName(), _documentMappingService, _logger, cancellationToken)
+                .TryGetPropertyRangeAsync(documentSnapshot, attributeDescriptor.GetPropertyName(), _documentMappingService, _logger, cancellationToken)
                 .ConfigureAwait(false);
 
             if (range is not null)

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/GoToDefinition/RazorComponentDefinitionHelpers.cs

@@ -8,10 +8,10 @@ using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
-using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
+using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using LspRange = Microsoft.VisualStudio.LanguageServer.Protocol.Range;
@@ -130,13 +130,13 @@ internal static class RazorComponentDefinitionHelpers
     }
 
     public static async Task<LspRange?> TryGetPropertyRangeAsync(
-        RazorCodeDocument codeDocument,
+        IDocumentSnapshot documentSnapshot,
         string propertyName,
         IDocumentMappingService documentMappingService,
         ILogger logger,
         CancellationToken cancellationToken)
     {
-        // Parse the C# file and find the property that matches the name.
+        // Process the C# tree and find the property that matches the name.
         // We don't worry about parameter attributes here for two main reasons:
         //   1. We don't have symbolic information, so the best we could do would be checking for any
         //      attribute named Parameter, regardless of which namespace. It also means we would have
@@ -147,9 +147,10 @@ internal static class RazorComponentDefinitionHelpers
         //      tag helper attribute. If they don't have the [Parameter] attribute then the Razor compiler
         //      will error, but allowing them to Go To Def on that property regardless, actually helps
         //      them fix the error.
-        var csharpText = codeDocument.GetCSharpSourceText();
-        var syntaxTree = CSharpSyntaxTree.ParseText(csharpText, cancellationToken: cancellationToken);
-        var root = await syntaxTree.GetRootAsync(cancellationToken).ConfigureAwait(false);
+
+        var csharpSyntaxTree = await documentSnapshot.GetCSharpSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
+        var root = await csharpSyntaxTree.GetRootAsync(cancellationToken).ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
 
         // Since we know how the compiler generates the C# source we can be a little specific here, and avoid
         // long tree walks. If the compiler ever changes how they generate their code, the tests for this will break
@@ -169,6 +170,7 @@ internal static class RazorComponentDefinitionHelpers
                 return null;
             }
 
+            var csharpText = codeDocument.GetCSharpSourceText();
             var range = csharpText.GetRange(property.Identifier.Span);
             if (documentMappingService.TryMapToHostDocumentRange(codeDocument.GetCSharpDocument(), range, out var originalRange))
             {

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentSnapshot.cs

@@ -3,8 +3,10 @@
 
 using System;
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -62,4 +64,11 @@ internal class DocumentSnapshot : IDocumentSnapshot
     {
         return new DocumentSnapshot(ProjectInternal, State.WithText(text, VersionStamp.Create()));
     }
+
+    public async Task<SyntaxTree> GetCSharpSyntaxTreeAsync(CancellationToken cancellationToken)
+    {
+        var codeDocument = await GetGeneratedOutputAsync().ConfigureAwait(false);
+        var csharpText = codeDocument.GetCSharpSourceText();
+        return CSharpSyntaxTree.ParseText(csharpText, cancellationToken: cancellationToken);
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/IDocumentSnapshot.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Text;
@@ -20,6 +21,12 @@ internal interface IDocumentSnapshot
     Task<VersionStamp> GetTextVersionAsync();
     Task<RazorCodeDocument> GetGeneratedOutputAsync();
 
+    /// <summary>
+    /// Gets the Roslyn syntax tree for the generated C# for this Razor document
+    /// </summary>
+    /// <remarks>Using this from the LSP server side of things is not ideal. Use sparingly :)</remarks>
+    Task<SyntaxTree> GetCSharpSyntaxTreeAsync(CancellationToken cancellationToken);
+
     bool TryGetText([NotNullWhen(true)] out SourceText? result);
     bool TryGetTextVersion(out VersionStamp result);
     bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? result);

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/ImportDocumentSnapshot.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Text;
@@ -75,4 +76,7 @@ internal class ImportDocumentSnapshot : IDocumentSnapshot
 
     public IDocumentSnapshot WithText(SourceText text)
         => throw new NotSupportedException();
+
+    public Task<SyntaxTree> GetCSharpSyntaxTreeAsync(CancellationToken cancellationToken)
+        => throw new NotSupportedException();
 }

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

@@ -11,7 +11,6 @@ using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Protocol.DocumentHighlight;
 using Microsoft.CodeAnalysis.Razor.Remote;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
 using static Microsoft.VisualStudio.LanguageServer.Protocol.VsLspExtensions;
@@ -29,7 +28,6 @@ internal sealed partial class RemoteDocumentHighlightService(in ServiceArgs args
     }
 
     private readonly IDocumentMappingService _documentMappingService = args.ExportProvider.GetExportedValue<IDocumentMappingService>();
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
 
     public ValueTask<Response> GetHighlightsAsync(
         RazorPinnedSolutionInfoWrapper solutionInfo,
@@ -68,7 +66,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().ConfigureAwait(false);
 
             var highlights = await DocumentHighlights.GetHighlightsAsync(generatedDocument, mappedPosition, cancellationToken).ConfigureAwait(false);
 

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

@@ -9,7 +9,6 @@ using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.FoldingRanges;
 using Microsoft.CodeAnalysis.Razor.Protocol.Folding;
 using Microsoft.CodeAnalysis.Razor.Remote;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using ExternalHandlers = Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers;
@@ -25,7 +24,6 @@ internal sealed class RemoteFoldingRangeService(in ServiceArgs args) : RazorDocu
     }
 
     private readonly IFoldingRangeService _foldingRangeService = args.ExportProvider.GetExportedValue<IFoldingRangeService>();
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
 
     public ValueTask<ImmutableArray<RemoteFoldingRange>> GetFoldingRangesAsync(
         RazorPinnedSolutionInfoWrapper solutionInfo,
@@ -43,7 +41,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().ConfigureAwait(false);
 
         var csharpRanges = await ExternalHandlers.FoldingRanges.GetFoldingRangesAsync(generatedDocument, cancellationToken).ConfigureAwait(false);
 

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

@@ -10,7 +10,6 @@ using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.GoToDefinition;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Remote;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
@@ -33,7 +32,6 @@ internal sealed class RemoteGoToDefinitionService(in ServiceArgs args) : RazorDo
     }
 
     private readonly IRazorComponentDefinitionService _componentDefinitionService = args.ExportProvider.GetExportedValue<IRazorComponentDefinitionService>();
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
 
     protected override IDocumentPositionInfoStrategy DocumentPositionInfoStrategy => PreferAttributeNameDocumentPositionInfoStrategy.Instance;
 
@@ -62,14 +60,6 @@ internal sealed class RemoteGoToDefinitionService(in ServiceArgs args) : RazorDo
 
         var positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex);
 
-        // First, see if this is a Razor component.
-        var componentLocation = await _componentDefinitionService.GetDefinitionAsync(context.Snapshot, positionInfo, ignoreAttributes: false, cancellationToken).ConfigureAwait(false);
-        if (componentLocation is not null)
-        {
-            // Convert from VS LSP Location to Roslyn. This can be removed when Razor moves fully onto Roslyn's LSP types.
-            return Results([RoslynLspFactory.CreateLocation(componentLocation.Uri, componentLocation.Range.ToLinePositionSpan())]);
-        }
-
         if (positionInfo.LanguageKind == RazorLanguageKind.Html)
         {
             // Sometimes Html can actually be mapped to C#, like for example component attributes, which map to
@@ -83,9 +73,17 @@ internal sealed class RemoteGoToDefinitionService(in ServiceArgs args) : RazorDo
             }
         }
 
-        // If it isn't a Razor component, and it isn't C#, let the server know to delegate to HTML.
-        if (positionInfo.LanguageKind != RazorLanguageKind.CSharp)
+        if (positionInfo.LanguageKind is RazorLanguageKind.Html or RazorLanguageKind.Razor)
         {
+            // First, see if this is a Razor component. We ignore attributes here, because they're better served by the C# handler.
+            var componentLocation = await _componentDefinitionService.GetDefinitionAsync(context.Snapshot, positionInfo, ignoreAttributes: true, cancellationToken).ConfigureAwait(false);
+            if (componentLocation is not null)
+            {
+                // Convert from VS LSP Location to Roslyn. This can be removed when Razor moves fully onto Roslyn's LSP types.
+                return Results([RoslynLspFactory.CreateLocation(componentLocation.Uri, componentLocation.Range.ToLinePositionSpan())]);
+            }
+
+            // If it isn't a Razor component, and it isn't C#, let the server know to delegate to HTML.
             return CallHtml;
         }
 
@@ -96,7 +94,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().ConfigureAwait(false);
 
         var locations = await ExternalHandlers.GoToDefinition
             .GetDefinitionsAsync(
@@ -121,7 +119,7 @@ internal sealed class RemoteGoToDefinitionService(in ServiceArgs args) : RazorDo
             var (uri, range) = location;
 
             var (mappedDocumentUri, mappedRange) = await DocumentMappingService
-                .MapToHostDocumentUriAndRangeAsync((RemoteDocumentSnapshot)context.Snapshot, uri, range.ToLinePositionSpan(), cancellationToken)
+                .MapToHostDocumentUriAndRangeAsync(context.Snapshot, uri, range.ToLinePositionSpan(), cancellationToken)
                 .ConfigureAwait(false);
 
             var mappedLocation = RoslynLspFactory.CreateLocation(mappedDocumentUri, mappedRange);

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

@@ -26,8 +26,6 @@ internal sealed partial class RemoteInlayHintService(in ServiceArgs args) : Razo
             => new RemoteInlayHintService(in args);
     }
 
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
-
     public ValueTask<InlayHint[]?> GetInlayHintsAsync(JsonSerializableRazorPinnedSolutionInfoWrapper solutionInfo, JsonSerializableDocumentId razorDocumentId, InlayHintParams inlayHintParams, bool displayAllOverride, CancellationToken cancellationToken)
        => RunServiceAsync(
             solutionInfo,
@@ -52,7 +50,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().ConfigureAwait(false);
 
         var textDocument = inlayHintParams.TextDocument.WithUri(generatedDocument.CreateUri());
         var range = projectedLinePositionSpan.ToRange();
@@ -106,7 +104,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().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/DocumentSnapshotFactory.cs

@@ -4,26 +4,28 @@
 using System;
 using System.Composition;
 using System.Runtime.CompilerServices;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 
 [Export(typeof(DocumentSnapshotFactory)), Shared]
 [method: ImportingConstructor]
-internal class DocumentSnapshotFactory(Lazy<ProjectSnapshotFactory> projectSnapshotFactory)
+internal class DocumentSnapshotFactory(Lazy<ProjectSnapshotFactory> projectSnapshotFactory, IFilePathService filePathService)
 {
-    private static readonly ConditionalWeakTable<TextDocument, RemoteDocumentSnapshot> _documentSnapshots = new();
+    private static readonly ConditionalWeakTable<TextDocument, RemoteDocumentSnapshot> s_documentSnapshots = new();
 
     private readonly Lazy<ProjectSnapshotFactory> _projectSnapshotFactory = projectSnapshotFactory;
+    private readonly IFilePathService _filePathService = filePathService;
 
     public RemoteDocumentSnapshot GetOrCreate(TextDocument textDocument)
     {
-        lock (_documentSnapshots)
+        lock (s_documentSnapshots)
         {
-            if (!_documentSnapshots.TryGetValue(textDocument, out var documentSnapshot))
+            if (!s_documentSnapshots.TryGetValue(textDocument, out var documentSnapshot))
             {
                 var projectSnapshot = _projectSnapshotFactory.Value.GetOrCreate(textDocument.Project);
-                documentSnapshot = new RemoteDocumentSnapshot(textDocument, projectSnapshot);
-                _documentSnapshots.Add(textDocument, documentSnapshot);
+                documentSnapshot = new RemoteDocumentSnapshot(textDocument, projectSnapshot, _filePathService);
+                s_documentSnapshots.Add(textDocument, documentSnapshot);
             }
 
             return documentSnapshot;

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

@@ -11,19 +11,19 @@ namespace Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 [method: ImportingConstructor]
 internal class ProjectSnapshotFactory(DocumentSnapshotFactory documentSnapshotFactory, ITelemetryReporter telemetryReporter)
 {
-    private static readonly ConditionalWeakTable<Project, RemoteProjectSnapshot> _projectSnapshots = new();
+    private static readonly ConditionalWeakTable<Project, RemoteProjectSnapshot> s_projectSnapshots = new();
 
     private readonly DocumentSnapshotFactory _documentSnapshotFactory = documentSnapshotFactory;
     private readonly ITelemetryReporter _telemetryReporter = telemetryReporter;
 
     public RemoteProjectSnapshot GetOrCreate(Project project)
     {
-        lock (_projectSnapshots)
+        lock (s_projectSnapshots)
         {
-            if (!_projectSnapshots.TryGetValue(project, out var projectSnapshot))
+            if (!s_projectSnapshots.TryGetValue(project, out var projectSnapshot))
             {
                 projectSnapshot = new RemoteProjectSnapshot(project, _documentSnapshotFactory, _telemetryReporter);
-                _projectSnapshots.Add(project, projectSnapshot);
+                s_projectSnapshots.Add(project, projectSnapshot);
             }
 
             return projectSnapshot;

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,21 +1,23 @@
 // 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;
 
-internal class RemoteDocumentSnapshot(TextDocument textDocument, RemoteProjectSnapshot projectSnapshot) : IDocumentSnapshot
+internal class RemoteDocumentSnapshot(TextDocument textDocument, RemoteProjectSnapshot projectSnapshot, IFilePathService filePathService) : IDocumentSnapshot
 {
     private readonly TextDocument _textDocument = textDocument;
     private readonly RemoteProjectSnapshot _projectSnapshot = projectSnapshot;
+    private readonly IFilePathService _filePathService = filePathService;
 
     // TODO: Delete this field when the source generator is hooked up
     private Document? _generatedDocument;
@@ -79,7 +81,7 @@ internal class RemoteDocumentSnapshot(TextDocument textDocument, RemoteProjectSn
         var id = _textDocument.Id;
         var newDocument = _textDocument.Project.Solution.WithAdditionalDocumentText(id, text).GetAdditionalDocument(id).AssumeNotNull();
 
-        return new RemoteDocumentSnapshot(newDocument, _projectSnapshot);
+        return new RemoteDocumentSnapshot(newDocument, _projectSnapshot, _filePathService);
     }
 
     public bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? result)
@@ -88,14 +90,38 @@ 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()
     {
         if (_generatedDocument is Document generatedDocument)
         {
             return generatedDocument;
         }
 
-        generatedDocument = await createGeneratedDocument(arg);
+        generatedDocument = await HACK_GenerateDocumentAsync().ConfigureAwait(false);
         return InterlockedOperations.Initialize(ref _generatedDocument, generatedDocument);
     }
+
+    private async Task<Document> HACK_GenerateDocumentAsync()
+    {
+        // 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);
+    }
+
+    public async Task<SyntaxTree> GetCSharpSyntaxTreeAsync(CancellationToken cancellationToken)
+    {
+        var document = await GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var tree = await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
+        return tree.AssumeNotNull();
+    }
 }

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

@@ -25,7 +25,6 @@ internal sealed class RemoteRenameService(in ServiceArgs args) : RazorDocumentSe
     }
 
     private readonly IRenameService _renameService = args.ExportProvider.GetExportedValue<IRenameService>();
-    private readonly IFilePathService _filePathService = args.ExportProvider.GetExportedValue<IFilePathService>();
     private readonly IEditMappingService _editMappingService = args.ExportProvider.GetExportedValue<IEditMappingService>();
 
     public ValueTask<RemoteResponse<WorkspaceEdit?>> GetRenameEditAsync(
@@ -53,7 +52,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().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().ConfigureAwait(false);
 
         var data = await SemanticTokensRange.GetSemanticTokensAsync(
             generatedDocument,

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

@@ -6,16 +6,14 @@ using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.Remote;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 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>
@@ -24,22 +22,20 @@ internal sealed class RemoteSignatureHelpService(in ServiceArgs args) : RazorDoc
             => new RemoteSignatureHelpService(in args);
     }
 
-    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().ConfigureAwait(false);
 
         if (DocumentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), absoluteIndex, out var mappedPosition, out _))
         {

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/Definition/RazorComponentDefinitionHelpersTest.cs

@@ -6,6 +6,7 @@ using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.LanguageServer.Completion;
 using Microsoft.AspNetCore.Razor.Test.Common;
 using Microsoft.AspNetCore.Razor.Test.Common.LanguageServer;
+using Microsoft.AspNetCore.Razor.Test.Common.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.GoToDefinition;
 using Microsoft.CodeAnalysis.Testing;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
@@ -389,10 +390,11 @@ public class RazorComponentDefinitionHelpersTest(ITestOutputHelper testOutput) :
 
         var codeDocument = CreateCodeDocument(content);
         var expectedRange = codeDocument.Source.Text.GetRange(selection);
+        var snapshot = TestDocumentSnapshot.Create("test.razor", content).With(codeDocument);
 
         var documentMappingService = new LspDocumentMappingService(FilePathService, new TestDocumentContextFactory(), LoggerFactory);
 
-        var range = await RazorComponentDefinitionHelpers.TryGetPropertyRangeAsync(codeDocument, propertyName, documentMappingService, Logger, DisposalToken);
+        var range = await RazorComponentDefinitionHelpers.TryGetPropertyRangeAsync(snapshot, propertyName, documentMappingService, Logger, DisposalToken);
         Assert.NotNull(range);
         Assert.Equal(expectedRange, range);
     }

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CohostGoToDefinitionEndpointTest.cs

@@ -124,6 +124,51 @@ public class CohostGoToDefinitionEndpointTest(ITestOutputHelper testOutputHelper
         await VerifyGoToDefinitionAsync(input, FileKinds.Component);
     }
 
+    [Fact]
+    public async Task Component()
+    {
+        TestCode input = """
+            <Surv$$eyPrompt Title="InputValue" />
+            """;
+
+        TestCode surveyPrompt = """
+            [||]@namespace SomeProject
+
+            <div></div>
+
+            @code
+            {
+                [Parameter]
+                public string Title { get; set; }
+            }
+            """;
+
+        TestCode surveyPromptGeneratedCode = """
+            using Microsoft.AspNetCore.Components;
+
+            namespace SomeProject
+            {
+                public partial class SurveyPrompt : ComponentBase
+                {
+                    [Parameter]
+                    public string Title { get; set; }
+                }
+            }
+            """;
+
+        var result = await GetGoToDefinitionResultAsync(input, FileKinds.Component,
+            (FileName("SurveyPrompt.razor"), surveyPrompt.Text),
+            (FileName("SurveyPrompt.razor.g.cs"), surveyPromptGeneratedCode.Text));
+
+        Assert.NotNull(result.Value.Second);
+        var locations = result.Value.Second;
+        var location = Assert.Single(locations);
+
+        var text = SourceText.From(surveyPrompt.Text);
+        var range = RoslynLspExtensions.GetRange(text, surveyPrompt.Span);
+        Assert.Equal(range, location.Range);
+    }
+
     [Theory]
     [InlineData("Ti$$tle")]
     [InlineData("$$@bind-Title")]

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/RazorComponentDefinitionServiceTest.cs

@@ -0,0 +1,99 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.IO;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.AspNetCore.Razor.Test.Common;
+using Microsoft.CodeAnalysis.Razor.DocumentMapping;
+using Microsoft.CodeAnalysis.Razor.GoToDefinition;
+using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
+
+public class RazorComponentDefinitionServiceTest(ITestOutputHelper testOutputHelper) : CohostEndpointTestBase(testOutputHelper)
+{
+    // PREAMBLE: Right now these tests are about ensuring we don't accidentally introduce a future bug
+    //           in the generated document handling code in the RazorComponentDefinitionService in OOP.
+    //           Right now in cohosting none of the code under test is actually used. This is because
+    //           the logic for manually looking up properties from attributes is only necessary when
+    //           "Single Server Mode" is off, which is currently only VS Code. When cohosting comes to
+    //           VS Code, that will no longer be true, and VS Code will use the same code paths as VS,
+    //           even then these tests will be exercising uncalled code.
+    //           The tests, and the "ignoreAttributes" parameter in the call to GetDefinitionAsync, should
+    //           be deleted entirely at that point. "ignoreAttributes" will essentially always be true,
+    //           as directly calling Roslyn provides better results.
+
+    [Fact]
+    public async Task Do()
+    {
+        TestCode input = """
+            <SurveyPrompt Ti$$tle="InputValue" />
+
+            @code
+            {
+                private string? InputValue { get; set; }
+
+                private void BindAfter()
+                {
+                }
+            }
+            """;
+
+        TestCode surveyPrompt = """
+            @namespace SomeProject
+
+            <div></div>
+
+            @code
+            {
+                [Parameter]
+                public string [|Title|] { get; set; }
+            }
+            """;
+
+        TestCode surveyPromptGeneratedCode = """
+            using Microsoft.AspNetCore.Components;
+
+            namespace SomeProject
+            {
+                public partial class SurveyPrompt : ComponentBase
+                {
+                    [Parameter]
+                    public string Title { get; set; }
+                }
+            }
+            """;
+
+        await VerifyDefinitionAsync(input, surveyPrompt, (FileName("SurveyPrompt.razor"), surveyPrompt.Text),
+            (FileName("SurveyPrompt.razor.g.cs"), surveyPromptGeneratedCode.Text));
+    }
+
+    private async Task VerifyDefinitionAsync(TestCode input, TestCode expectedDocument, params (string fileName, string contents)[]? additionalFiles)
+    {
+        var document = CreateProjectAndRazorDocument(input.Text, FileKinds.Component, additionalFiles);
+
+        var service = OOPExportProvider.GetExportedValue<IRazorComponentDefinitionService>();
+        var documentSnapshotFactory = OOPExportProvider.GetExportedValue<DocumentSnapshotFactory>();
+        var documentMappingService = OOPExportProvider.GetExportedValue<IDocumentMappingService>();
+
+        var documentSnapshot = documentSnapshotFactory.GetOrCreate(document);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync();
+        var positionInfo = documentMappingService.GetPositionInfo(codeDocument, input.Position);
+
+        var location = await service.GetDefinitionAsync(documentSnapshot, positionInfo, ignoreAttributes: false, DisposalToken);
+
+        Assert.NotNull(location);
+
+        var text = SourceText.From(expectedDocument.Text);
+        var range = text.GetRange(expectedDocument.Span);
+        Assert.Equal(range, location.Range);
+    }
+
+    private static string FileName(string projectRelativeFileName)
+        => Path.Combine(TestProjectData.SomeProjectPath, projectRelativeFileName);
+}