Fix cross project span and edit mapping

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/CodeActions/CohostCodeActionsEndpoint.cs

@@ -103,7 +103,14 @@ public ImmutableArray<Registration> GetRegistrations(VSInternalClientCapabilitie
 
     private async Task<RazorVSInternalCodeAction[]> GetCSharpCodeActionsAsync(TextDocument razorDocument, VSCodeActionParams request, Guid correlationId, CancellationToken cancellationToken)
     {
-        var generatedDocument = await razorDocument.Project.TryGetCSharpDocumentFromGeneratedDocumentUriAsync(request.TextDocument.DocumentUri.GetRequiredParsedUri(), cancellationToken).ConfigureAwait(false);
+        var solution = razorDocument.Project.Solution;
+        if (!solution.TryGetSourceGeneratedDocumentIdentity(request.TextDocument.DocumentUri.GetRequiredParsedUri(), out var identity) ||
+            !solution.TryGetProject(identity.DocumentId.ProjectId, out var project))
+        {
+            return [];
+        }
+
+        var generatedDocument = await project.TryGetCSharpDocumentForGeneratedDocumentAsync(identity, cancellationToken).ConfigureAwait(false);
         if (generatedDocument is null)
         {
             return [];

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/CodeActions/CohostCodeActionsResolveEndpoint.cs

@@ -104,7 +104,14 @@ private async Task<CodeAction> ResolveCSharpCodeActionAsync(TextDocument razorDo
 
             var uri = resolveParams.DelegatedDocumentUri.AssumeNotNull();
 
-            var generatedDocument = await razorDocument.Project.TryGetCSharpDocumentFromGeneratedDocumentUriAsync(uri, cancellationToken).ConfigureAwait(false);
+            var solution = razorDocument.Project.Solution;
+            if (!solution.TryGetSourceGeneratedDocumentIdentity(uri, out var identity) ||
+                !solution.TryGetProject(identity.DocumentId.ProjectId, out var project))
+            {
+                return codeAction;
+            }
+
+            var generatedDocument = await project.TryGetCSharpDocumentForGeneratedDocumentAsync(identity, cancellationToken).ConfigureAwait(false);
             if (generatedDocument is null)
             {
                 return codeAction;

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/ProjectExtensions.cs

@@ -3,13 +3,12 @@
 
 using System;
 using System.Collections.Immutable;
-using System.Diagnostics.CodeAnalysis;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Language;
-using Microsoft.AspNetCore.Razor.Threading;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.NET.Sdk.Razor.SourceGenerators;
@@ -44,12 +43,10 @@ public static async ValueTask<TagHelperCollection> GetTagHelpersAsync(
         return discoveryService.GetTagHelpers(compilation, Options, cancellationToken);
     }
 
-    public static Task<SourceGeneratedDocument?> TryGetCSharpDocumentFromGeneratedDocumentUriAsync(this Project project, Uri generatedDocumentUri, CancellationToken cancellationToken)
+    public static Task<SourceGeneratedDocument?> TryGetCSharpDocumentForGeneratedDocumentAsync(this Project project, RazorGeneratedDocumentIdentity identity, CancellationToken cancellationToken)
     {
-        if (!TryGetHintNameFromGeneratedDocumentUri(project, generatedDocumentUri, out var hintName))
-        {
-            return SpecializedTasks.Null<SourceGeneratedDocument>();
-        }
+        Debug.Assert(identity.DocumentId.ProjectId == project.Id, "Generated document URI does not belong to this project.");
+        var hintName = identity.HintName;
 
         return TryGetSourceGeneratedDocumentFromHintNameAsync(project, hintName, cancellationToken);
     }
@@ -138,28 +135,4 @@ public static async ValueTask<TagHelperCollection> GetTagHelpersAsync(
             return RazorSourceGenerator.GetIdentifierFromPath(relativeDocumentPath);
         }
     }
-
-    /// <summary>
-    /// Finds source generated documents by iterating through all of them. In OOP there are better options!
-    /// </summary>
-    public static bool TryGetHintNameFromGeneratedDocumentUri(this Project project, Uri generatedDocumentUri, [NotNullWhen(true)] out string? hintName)
-    {
-        if (!RazorUri.IsGeneratedDocumentUri(generatedDocumentUri))
-        {
-            hintName = null;
-            return false;
-        }
-
-        var identity = RazorUri.GetIdentityOfGeneratedDocument(project.Solution, generatedDocumentUri);
-
-        if (!identity.IsRazorSourceGeneratedDocument())
-        {
-            // This is not a Razor source generated document, so we don't know the hint name.
-            hintName = null;
-            return false;
-        }
-
-        hintName = identity.HintName;
-        return true;
-    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/SolutionExtensions.cs

@@ -93,4 +93,22 @@ public static Project GetRequiredProject(this Solution solution, ProjectKey proj
         return solution.GetProject(projectKey)
             ?? ThrowHelper.ThrowInvalidOperationException<Project>($"The project {projectKey} did not exist in {solution}.");
     }
+
+    public static bool TryGetSourceGeneratedDocumentIdentity(this Solution solution, Uri generatedDocumentUri, out RazorGeneratedDocumentIdentity identity)
+    {
+        identity = default;
+        if (!RazorUri.IsGeneratedDocumentUri(generatedDocumentUri))
+        {
+            return false;
+        }
+
+        identity = RazorUri.GetIdentityOfGeneratedDocument(solution, generatedDocumentUri);
+
+        if (!identity.IsRazorSourceGeneratedDocument())
+        {
+            return false;
+        }
+
+        return true;
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentMapping/RemoteDocumentMappingService.cs

@@ -5,6 +5,7 @@
 using System.Composition;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
@@ -42,8 +43,19 @@ internal sealed class RemoteDocumentMappingService(
             return (generatedDocumentUri, generatedDocumentRange);
         }
 
-        var project = originSnapshot.TextDocument.Project;
-        var razorCodeDocument = await _snapshotManager.GetSnapshot(project).TryGetCodeDocumentFromGeneratedDocumentUriAsync(generatedDocumentUri, cancellationToken).ConfigureAwait(false);
+        var solution = originSnapshot.TextDocument.Project.Solution;
+        if (!solution.TryGetSourceGeneratedDocumentIdentity(generatedDocumentUri, out var identity))
+        {
+            return (generatedDocumentUri, generatedDocumentRange);
+        }
+
+        var project = solution.GetProject(identity.DocumentId.ProjectId);
+        if (project is null)
+        {
+            return (generatedDocumentUri, generatedDocumentRange);
+        }
+
+        var razorCodeDocument = await _snapshotManager.GetSnapshot(project).TryGetCodeDocumentForGeneratedDocumentAsync(identity, cancellationToken).ConfigureAwait(false);
         if (razorCodeDocument is null)
         {
             return (generatedDocumentUri, generatedDocumentRange);
@@ -57,7 +69,7 @@ internal sealed class RemoteDocumentMappingService(
 
         // If the position is unmappable, but was in a generated Razor, we have one last check to see if Roslyn wants to navigate
         // to the class declaration, in which case we'll map to (0,0) in the Razor document itself.
-        if (await TryGetCSharpClassDeclarationSpanAsync(generatedDocumentUri, project, cancellationToken).ConfigureAwait(false) is { } classDeclSpan &&
+        if (await TryGetCSharpClassDeclarationSpanAsync(identity, project, cancellationToken).ConfigureAwait(false) is { } classDeclSpan &&
             generatedDocumentRange.Start == classDeclSpan.Start &&
                 (generatedDocumentRange.End == generatedDocumentRange.Start ||
                 generatedDocumentRange.End == classDeclSpan.End))
@@ -68,9 +80,9 @@ internal sealed class RemoteDocumentMappingService(
         return (generatedDocumentUri, generatedDocumentRange);
     }
 
-    private static async Task<LinePositionSpan?> TryGetCSharpClassDeclarationSpanAsync(Uri generatedDocumentUri, Project project, CancellationToken cancellationToken)
+    private static async Task<LinePositionSpan?> TryGetCSharpClassDeclarationSpanAsync(RazorGeneratedDocumentIdentity identity, Project project, CancellationToken cancellationToken)
     {
-        var generatedDocument = await project.TryGetCSharpDocumentFromGeneratedDocumentUriAsync(generatedDocumentUri, cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await project.TryGetCSharpDocumentForGeneratedDocumentAsync(identity, cancellationToken).ConfigureAwait(false);
         if (generatedDocument is null)
         {
             return null;

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentMapping/RemoteEditMappingService.cs

@@ -51,8 +51,14 @@ protected override bool TryGetDocumentContext(IDocumentSnapshot contextDocumentS
             throw new InvalidOperationException("RemoteEditMappingService can only be used with RemoteDocumentSnapshot instances.");
         }
 
-        var project = originSnapshot.TextDocument.Project;
-        var razorCodeDocument = await _snapshotManager.GetSnapshot(project).TryGetCodeDocumentFromGeneratedDocumentUriAsync(generatedDocumentUri, cancellationToken).ConfigureAwait(false);
+        var solution = originSnapshot.TextDocument.Project.Solution;
+        if (!solution.TryGetSourceGeneratedDocumentIdentity(generatedDocumentUri, out var identity) ||
+            !solution.TryGetProject(identity.DocumentId.ProjectId, out var project))
+        {
+            return null;
+        }
+
+        var razorCodeDocument = await _snapshotManager.GetSnapshot(project).TryGetCodeDocumentForGeneratedDocumentAsync(identity, cancellationToken).ConfigureAwait(false);
         if (razorCodeDocument is null)
         {
             return null;

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentMapping/RemoteSpanMappingService.cs

@@ -202,6 +202,6 @@ private async ValueTask<ImmutableArray<RazorMappedEditResult>> MapTextChangesAsy
 
         var projectSnapshot = _snapshotManager.GetSnapshot(generatedDocument.Project);
 
-        return await projectSnapshot.TryGetRazorDocumentFromGeneratedHintNameAsync(identity.HintName, cancellationToken).ConfigureAwait(false);
+        return await projectSnapshot.TryGetRazorDocumentForGeneratedDocumentAsync(identity, cancellationToken).ConfigureAwait(false);
     }
 }

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

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Threading;
@@ -12,6 +13,7 @@
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Utilities;
@@ -163,17 +165,15 @@ internal async Task<SourceGeneratedDocument> GetRequiredGeneratedDocumentAsync(I
         return await generatorResult.GetRequiredSourceGeneratedDocumentForRazorFilePathAsync(documentSnapshot.FilePath, cancellationToken).ConfigureAwait(false);
     }
 
-    public async Task<RazorCodeDocument?> TryGetCodeDocumentFromGeneratedDocumentUriAsync(Uri generatedDocumentUri, CancellationToken cancellationToken)
+    public async Task<RazorCodeDocument?> TryGetCodeDocumentForGeneratedDocumentAsync(RazorGeneratedDocumentIdentity identity, CancellationToken cancellationToken)
     {
-        if (!_project.TryGetHintNameFromGeneratedDocumentUri(generatedDocumentUri, out var hintName))
-        {
-            return null;
-        }
+        Debug.Assert(identity.DocumentId.ProjectId == _project.Id, "Generated document does not belong to this project.");
+        var hintName = identity.HintName;
 
         return await TryGetCodeDocumentFromGeneratedHintNameAsync(hintName, cancellationToken).ConfigureAwait(false);
     }
 
-    public async Task<RazorCodeDocument?> TryGetCodeDocumentFromGeneratedHintNameAsync(string generatedDocumentHintName, CancellationToken cancellationToken)
+    private async Task<RazorCodeDocument?> TryGetCodeDocumentFromGeneratedHintNameAsync(string generatedDocumentHintName, CancellationToken cancellationToken)
     {
         var generatorResult = await GeneratorRunResult.CreateAsync(throwIfNotFound: false, _project, SolutionSnapshot.SnapshotManager, cancellationToken).ConfigureAwait(false);
         if (generatorResult.IsDefault)
@@ -186,15 +186,18 @@ internal async Task<SourceGeneratedDocument> GetRequiredGeneratedDocumentAsync(I
             : null;
     }
 
-    public async Task<TextDocument?> TryGetRazorDocumentFromGeneratedHintNameAsync(string generatedDocumentHintName, CancellationToken cancellationToken)
+    public async Task<TextDocument?> TryGetRazorDocumentForGeneratedDocumentAsync(RazorGeneratedDocumentIdentity identity, CancellationToken cancellationToken)
     {
+        Debug.Assert(identity.DocumentId.ProjectId == _project.Id, "Generated document does not belong to this project.");
+        var hintName = identity.HintName;
+
         var generatorResult = await GeneratorRunResult.CreateAsync(throwIfNotFound: false, _project, SolutionSnapshot.SnapshotManager, cancellationToken).ConfigureAwait(false);
         if (generatorResult.IsDefault)
         {
             return null;
         }
 
-        return generatorResult.GetRazorFilePathFromHintName(generatedDocumentHintName) is { } razorFilePath &&
+        return generatorResult.GetRazorFilePathFromHintName(hintName) is { } razorFilePath &&
             generatorResult.TryGetRazorDocument(razorFilePath, out var razorDocument)
                 ? razorDocument
                 : null;

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

@@ -73,7 +73,14 @@ public ImmutableArray<Registration> GetRegistrations(VSInternalClientCapabilitie
             return null;
         }
 
-        var generatedDocument = await razorDocument.Project.TryGetCSharpDocumentFromGeneratedDocumentUriAsync(generatedDocumentUri, cancellationToken).ConfigureAwait(false);
+        var solution = razorDocument.Project.Solution;
+        if (!solution.TryGetSourceGeneratedDocumentIdentity(generatedDocumentUri, out var identity) ||
+            !solution.TryGetProject(identity.DocumentId.ProjectId, out var project))
+        {
+            return null;
+        }
+
+        var generatedDocument = await project.TryGetCSharpDocumentForGeneratedDocumentAsync(identity, cancellationToken).ConfigureAwait(false);
         if (generatedDocument is null)
         {
             return null;

src/Razor/test/Microsoft.AspNetCore.Razor.Test.Common.Cohosting/CohostTestBase.cs

@@ -179,10 +179,21 @@ private protected TextDocument CreateProjectAndRazorDocument(
 
     private protected static TextDocument CreateProjectAndRazorDocument(CodeAnalysis.Workspace workspace, ProjectId projectId, bool miscellaneousFile, DocumentId documentId, string documentFilePath, string contents, (string fileName, string contents)[]? additionalFiles, bool inGlobalNamespace, bool addDefaultImports, Action<RazorProjectBuilder>? projectConfigure)
     {
-        return AddProjectAndRazorDocument(workspace.CurrentSolution, TestProjectData.SomeProject.FilePath, projectId, miscellaneousFile, documentId, documentFilePath, contents, additionalFiles, inGlobalNamespace, addDefaultImports, projectConfigure);
+        return AddProjectAndRazorDocument(workspace.CurrentSolution, TestProjectData.SomeProject.FilePath, projectId, documentId, documentFilePath, contents, miscellaneousFile, additionalFiles, inGlobalNamespace, addDefaultImports, projectConfigure);
     }
 
-    private protected static TextDocument AddProjectAndRazorDocument(Solution solution, [DisallowNull] string? projectFilePath, ProjectId projectId, bool miscellaneousFile, DocumentId documentId, string documentFilePath, string contents, (string fileName, string contents)[]? additionalFiles, bool inGlobalNamespace, bool addDefaultImports, Action<RazorProjectBuilder>? projectConfigure)
+    private protected static TextDocument AddProjectAndRazorDocument(
+        Solution solution,
+        [DisallowNull] string? projectFilePath,
+        ProjectId projectId,
+        DocumentId documentId,
+        string documentFilePath,
+        string contents,
+        bool miscellaneousFile = false,
+        (string fileName, string contents)[]? additionalFiles = null,
+        bool inGlobalNamespace = false,
+        bool addDefaultImports = true,
+        Action<RazorProjectBuilder>? projectConfigure = null)
     {
         var builder = new RazorProjectBuilder(projectId);
 

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

@@ -72,7 +72,7 @@ public async Task HoverRequest_MultipleProjects_ReturnsResults()
         var projectId = ProjectId.CreateNewId(debugName: TestProjectData.SomeProject.DisplayName);
         var documentFilePath = TestProjectData.AnotherProjectComponentFile1.FilePath;
         var documentId = DocumentId.CreateNewId(projectId, debugName: documentFilePath);
-        var otherDocument = AddProjectAndRazorDocument(document.Project.Solution, TestProjectData.AnotherProject.FilePath, projectId, miscellaneousFile: false, documentId, documentFilePath, otherInput.Text, additionalFiles: null, inGlobalNamespace: false, addDefaultImports: true, projectConfigure: null);
+        var otherDocument = AddProjectAndRazorDocument(document.Project.Solution, TestProjectData.AnotherProject.FilePath, projectId, documentId, documentFilePath, otherInput.Text);
 
         // Make sure we have the document from our new fork
         document = otherDocument.Project.Solution.GetAdditionalDocument(document.Id).AssumeNotNull();