Fix various folding range, and general VS Code issues

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

@@ -4,7 +4,6 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
@@ -15,7 +14,7 @@ internal class DefaultDocumentVersionCache : DocumentVersionCache
     internal const int MaxDocumentTrackingCount = 20;
 
     // Internal for testing
-    internal readonly Dictionary<DocumentKey, List<DocumentEntry>> DocumentLookup_NeedsLock;
+    internal readonly Dictionary<string, List<DocumentEntry>> DocumentLookup_NeedsLock;
     private readonly ReadWriterLocker _lock = new();
     private ProjectSnapshotManagerBase? _projectSnapshotManager;
 
@@ -24,7 +23,7 @@ private ProjectSnapshotManagerBase ProjectSnapshotManager
 
     public DefaultDocumentVersionCache()
     {
-        DocumentLookup_NeedsLock = new Dictionary<DocumentKey, List<DocumentEntry>>();
+        DocumentLookup_NeedsLock = new Dictionary<string, List<DocumentEntry>>(FilePathComparer.Instance);
     }
 
     public override void TrackDocumentVersion(IDocumentSnapshot documentSnapshot, int version)
@@ -43,7 +42,7 @@ private void TrackDocumentVersion(IDocumentSnapshot documentSnapshot, int versio
         // Need to ensure the write lock covers all uses of documentEntries, not just DocumentLookup
         using (upgradeableReadLock.EnterWriteLock())
         {
-            var key = new DocumentKey(documentSnapshot.Project.Key, documentSnapshot.FilePath.AssumeNotNull());
+            var key = documentSnapshot.FilePath.AssumeNotNull();
             if (!DocumentLookup_NeedsLock.TryGetValue(key, out var documentEntries))
             {
                 documentEntries = new List<DocumentEntry>();
@@ -73,7 +72,7 @@ public override bool TryGetDocumentVersion(IDocumentSnapshot documentSnapshot, [
 
         using var _ = _lock.EnterReadLock();
 
-        var key = new DocumentKey(documentSnapshot.Project.Key, documentSnapshot.FilePath.AssumeNotNull());
+        var key = documentSnapshot.FilePath.AssumeNotNull();
         if (!DocumentLookup_NeedsLock.TryGetValue(key, out var documentEntries))
         {
             version = null;
@@ -127,20 +126,13 @@ private void ProjectSnapshotManager_Changed(object? sender, ProjectChangeEventAr
         {
             case ProjectChangeKind.DocumentChanged:
                 var documentFilePath = args.DocumentFilePath!;
-
-                if (!ProjectSnapshotManager.IsDocumentOpen(documentFilePath))
+                if (DocumentLookup_NeedsLock.ContainsKey(documentFilePath) &&
+                    !ProjectSnapshotManager.IsDocumentOpen(documentFilePath))
                 {
                     using (upgradeableLock.EnterWriteLock())
                     {
-                        // Document closed, evict all entries.
-                        var keys = DocumentLookup_NeedsLock.Keys.ToArray();
-                        foreach (var key in keys)
-                        {
-                            if (key.DocumentFilePath == documentFilePath)
-                            {
-                                DocumentLookup_NeedsLock.Remove(key);
-                            }
-                        }
+                        // Document closed, evict entry.
+                        DocumentLookup_NeedsLock.Remove(documentFilePath);
                     }
                 }
 
@@ -170,8 +162,7 @@ private void CaptureProjectDocumentsAsLatest(IProjectSnapshot projectSnapshot, R
     {
         foreach (var documentPath in projectSnapshot.DocumentFilePaths)
         {
-            var key = new DocumentKey(projectSnapshot.Key, documentPath);
-            if (DocumentLookup_NeedsLock.ContainsKey(key) &&
+            if (DocumentLookup_NeedsLock.ContainsKey(documentPath) &&
                 projectSnapshot.GetDocument(documentPath) is { } document)
             {
                 MarkAsLatestVersion(document, upgradeableReadLock);
@@ -181,8 +172,7 @@ private void CaptureProjectDocumentsAsLatest(IProjectSnapshot projectSnapshot, R
 
     private void MarkAsLatestVersion(IDocumentSnapshot document, ReadWriterLocker.UpgradeableReadLock upgradeableReadLock)
     {
-        var key = new DocumentKey(document.Project.Key, document.FilePath.AssumeNotNull());
-        if (!DocumentLookup_NeedsLock.TryGetValue(key, out var documentEntries))
+        if (!DocumentLookup_NeedsLock.TryGetValue(document.FilePath.AssumeNotNull(), out var documentEntries))
         {
             return;
         }

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

@@ -189,6 +189,16 @@ public override void PublishHtml(ProjectKey projectKey, string filePath, SourceT
             HostDocumentVersion = hostDocumentVersion,
         };
 
+        // HACK: We know about a document being in multiple projects, but despite having ProjectKeyId in the request, currently the other end
+        // of this LSP message only knows about a single document file path. To prevent confusing them, we just send an update for the first project
+        // in the list.
+        if (_projectSnapshotManager is { } projectSnapshotManager &&
+            projectSnapshotManager.GetLoadedProject(projectKey) is { } projectSnapshot &&
+            projectSnapshotManager.GetAllProjectKeys(projectSnapshot.FilePath).First() != projectKey)
+        {
+            return;
+        }
+
         _ = _server.SendNotificationAsync(CustomMessageNames.RazorUpdateHtmlBufferEndpoint, request, CancellationToken.None);
     }
 

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Folding/FoldingRangeEndpoint.cs

@@ -88,7 +88,7 @@ public TextDocumentIdentifier GetTextDocumentIdentifier(FoldingRangeParams reque
         return foldingRanges;
     }
 
-    private async Task<IEnumerable<FoldingRange>?> HandleCoreAsync(RazorFoldingRangeRequestParam requestParams, DocumentContext documentContext, CancellationToken cancellationToken)
+    private async Task<List<FoldingRange>?> HandleCoreAsync(RazorFoldingRangeRequestParam requestParams, DocumentContext documentContext, CancellationToken cancellationToken)
     {
         var foldingResponse = await _languageServer.SendRequestAsync<RazorFoldingRangeRequestParam, RazorFoldingRangeResponse?>(
             CustomMessageNames.RazorFoldingRangeEndpoint,
@@ -132,7 +132,7 @@ public TextDocumentIdentifier GetTextDocumentIdentifier(FoldingRangeParams reque
         return finalRanges;
     }
 
-    private static IEnumerable<FoldingRange> FinalizeFoldingRanges(List<FoldingRange> mappedRanges, RazorCodeDocument codeDocument)
+    private List<FoldingRange> FinalizeFoldingRanges(List<FoldingRange> mappedRanges, RazorCodeDocument codeDocument)
     {
         // Don't allow ranges to be reported if they aren't spanning at least one line
         var validRanges = mappedRanges.Where(r => r.StartLine < r.EndLine);
@@ -144,14 +144,14 @@ private static IEnumerable<FoldingRange> FinalizeFoldingRanges(List<FoldingRange
             .Select(ranges => ranges.OrderByDescending(r => r.EndLine).First());
 
         // Fix the starting range so the "..." is shown at the end
-        return reducedRanges.Select(r => FixFoldingRangeStart(r, codeDocument));
+        return reducedRanges.Select(r => FixFoldingRangeStart(r, codeDocument)).ToList();
     }
 
     /// <summary>
     /// Fixes the start of a range so that the offset of the first line is the last character on that line. This makes
     /// it so collapsing will still show the text instead of just "..."
     /// </summary>
-    private static FoldingRange FixFoldingRangeStart(FoldingRange range, RazorCodeDocument codeDocument)
+    private FoldingRange FixFoldingRangeStart(FoldingRange range, RazorCodeDocument codeDocument)
     {
         Debug.Assert(range.StartLine < range.EndLine);
 
@@ -165,6 +165,15 @@ private static FoldingRange FixFoldingRangeStart(FoldingRange range, RazorCodeDo
 
         var sourceText = codeDocument.GetSourceText();
         var startLine = range.StartLine;
+
+        if (startLine > sourceText.Lines.Count)
+        {
+            // Sometimes VS Code seems to send us wildly out-of-range folding ranges for Html, so log a warning,
+            // but prevent a toast from appearing from an exception.
+            _logger.LogWarning("Got a folding range of ({StartLine}-{EndLine}) but Razor document {filePath} only has {count} lines.", range.StartLine, range.EndLine, codeDocument.Source.FilePath, sourceText.Lines.Count);
+            return range;
+        }
+
         var lineSpan = sourceText.Lines[startLine].Span;
 
         // Search from the end of the line to the beginning for the first non whitespace character. We want that

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/ProjectSystem/DefaultRazorProjectService.cs

@@ -117,8 +117,19 @@ void AddDocumentToProject(IProjectSnapshot projectSnapshot, string textDocumentP
             var hostDocument = new HostDocument(textDocumentPath, normalizedTargetFilePath);
             var textLoader = _remoteTextLoaderFactory.Create(textDocumentPath);
 
+            var projectSnapshotManager = _projectSnapshotManagerAccessor.Instance;
             _logger.LogInformation("Adding document '{filePath}' to project '{projectSnapshotFilePath}'.", filePath, projectSnapshot.FilePath);
-            _projectSnapshotManagerAccessor.Instance.DocumentAdded(projectSnapshot.Key, hostDocument, textLoader);
+            projectSnapshotManager.DocumentAdded(projectSnapshot.Key, hostDocument, textLoader);
+
+            // Adding a document to a project could also happen because a target was added to a project, or we're moving a document
+            // from Misc Project to a real one, and means the newly added document could actually already be open.
+            // If it is, we need to make sure we start generating it so we're ready to handle requests that could start coming in.
+            if (projectSnapshotManager.IsDocumentOpen(textDocumentPath) &&
+                projectSnapshotManager.GetLoadedProject(projectSnapshot.Key) is { } project &&
+                project.GetDocument(textDocumentPath) is { } document)
+            {
+                _ = document.GetGeneratedOutputAsync();
+            }
         }
     }
 

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

@@ -3,7 +3,9 @@
 
 #nullable disable
 
+using System;
 using System.Linq;
+using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Test.Common;
 using Microsoft.CodeAnalysis;
 using Xunit;
@@ -158,7 +160,7 @@ public void TrackDocumentVersion_AddsFirstEntry()
 
         // Assert
         var kvp = Assert.Single(documentVersionCache.DocumentLookup_NeedsLock);
-        Assert.Equal(document.FilePath, kvp.Key.DocumentFilePath);
+        Assert.Equal(document.FilePath, kvp.Key);
         var entry = Assert.Single(kvp.Value);
         Assert.True(entry.Document.TryGetTarget(out var actualDocument));
         Assert.Same(document, actualDocument);
@@ -233,4 +235,50 @@ public void TryGetDocumentVersion_KnownDocument_ReturnsTrue()
         Assert.True(result);
         Assert.Equal(1337, version);
     }
+
+    [Fact]
+    public void ProjectSnapshotManager_KnownDocumentAdded_TracksNewDocument()
+    {
+        // Arrange
+        var documentVersionCache = new DefaultDocumentVersionCache();
+        var projectSnapshotManager = TestProjectSnapshotManager.Create(ErrorReporter);
+        projectSnapshotManager.AllowNotifyListeners = true;
+        documentVersionCache.Initialize(projectSnapshotManager);
+
+        var project1 = TestProjectSnapshot.Create("C:/path/to/project1.csproj", intermediateOutputPath: "C:/path/to/obj1", documentFilePaths: Array.Empty<string>(), RazorConfiguration.Default, projectWorkspaceState: null);
+        projectSnapshotManager.ProjectAdded(project1.HostProject);
+        var document1 = projectSnapshotManager.CreateAndAddDocument(project1, @"C:\path\to\file.razor");
+
+        // Act
+        documentVersionCache.TrackDocumentVersion(document1, 1337);
+
+        // Assert
+        var kvp = Assert.Single(documentVersionCache.DocumentLookup_NeedsLock);
+        Assert.Equal(document1.FilePath, kvp.Key);
+        var entry = Assert.Single(kvp.Value);
+        Assert.True(entry.Document.TryGetTarget(out var actualDocument));
+        Assert.Same(document1, actualDocument);
+        Assert.Equal(1337, entry.Version);
+
+        // Act II
+        var project2 = TestProjectSnapshot.Create("C:/path/to/project2.csproj", intermediateOutputPath: "C:/path/to/obj2", documentFilePaths: Array.Empty<string>(), RazorConfiguration.Default, projectWorkspaceState: null);
+        projectSnapshotManager.ProjectAdded(project2.HostProject);
+        projectSnapshotManager.CreateAndAddDocument(project2, @"C:\path\to\file.razor");
+
+        var document2 = projectSnapshotManager.GetLoadedProject(project2.Key).GetDocument(document1.FilePath);
+
+        // Assert II
+        kvp = Assert.Single(documentVersionCache.DocumentLookup_NeedsLock);
+        Assert.Equal(document1.FilePath, kvp.Key);
+        Assert.Equal(2, kvp.Value.Count);
+
+        // Should still be tracking document 1 with no changes
+        Assert.True(kvp.Value[0].Document.TryGetTarget(out actualDocument));
+        Assert.Same(document1, actualDocument);
+        Assert.Equal(1337, kvp.Value[0].Version);
+
+        Assert.True(kvp.Value[1].Document.TryGetTarget(out actualDocument));
+        Assert.Same(document2, actualDocument);
+        Assert.Equal(1337, kvp.Value[1].Version);
+    }
 }

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

@@ -507,6 +507,7 @@ public void OpenDocument_OpensAndAddsDocumentToOwnerProject()
             },
             TestProjectSnapshot.Create("C:/__MISC_PROJECT__"));
         var projectSnapshotManager = new Mock<ProjectSnapshotManagerBase>(MockBehavior.Strict);
+        projectSnapshotManager.Setup(manager => manager.IsDocumentOpen(It.IsAny<string>())).Returns(false);
         projectSnapshotManager.Setup(manager => manager.DocumentAdded(It.IsAny<ProjectKey>(), It.IsAny<HostDocument>(), It.IsAny<TextLoader>()))
             .Callback<ProjectKey, HostDocument, TextLoader>((projectKey, hostDocument, loader) =>
             {
@@ -566,6 +567,7 @@ public void AddDocument_AddsDocumentToOwnerProject()
             },
             TestProjectSnapshot.Create("C:/__MISC_PROJECT__"));
         var projectSnapshotManager = new Mock<ProjectSnapshotManagerBase>(MockBehavior.Strict);
+        projectSnapshotManager.Setup(manager => manager.IsDocumentOpen(It.IsAny<string>())).Returns(false);
         projectSnapshotManager.Setup(manager => manager.DocumentAdded(It.IsAny<ProjectKey>(), It.IsAny<HostDocument>(), It.IsAny<TextLoader>()))
             .Callback<ProjectKey, HostDocument, TextLoader>((projectKey, hostDocument, loader) =>
             {
@@ -590,6 +592,7 @@ public void AddDocument_AddsDocumentToMiscellaneousProject()
         var miscellaneousProject = TestProjectSnapshot.Create("C:/__MISC_PROJECT__");
         var projectResolver = new TestSnapshotResolver(new Dictionary<string, IProjectSnapshot>(), miscellaneousProject);
         var projectSnapshotManager = new Mock<ProjectSnapshotManagerBase>(MockBehavior.Strict);
+        projectSnapshotManager.Setup(manager => manager.IsDocumentOpen(It.IsAny<string>())).Returns(false);
         projectSnapshotManager.Setup(manager => manager.DocumentAdded(It.IsAny<ProjectKey>(), It.IsAny<HostDocument>(), It.IsAny<TextLoader>()))
             .Callback<ProjectKey, HostDocument, TextLoader>((projectKey, hostDocument, loader) =>
             {