Ensure LspWorkspaceManager returns solutions without misc document when file moved (#80535)

* Add test demonstrating issue with document context in FBP

* Fix GetTextDocumentAsync returning misc document preferentially over host document

* fixes

* Alternative approach: Remove from misc files and re-lookup document

* Reword comment

* Review feedback

Author: dibarbet

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/FileBasedPrograms/FileBasedProgramsProjectSystem.cs

@@ -104,10 +104,10 @@ TextDocument AddPrimordialDocument(DocumentUri uri, SourceText documentText, str
         }
     }
 
-    public async ValueTask TryRemoveMiscellaneousDocumentAsync(DocumentUri uri)
+    public async ValueTask<bool> TryRemoveMiscellaneousDocumentAsync(DocumentUri uri)
     {
         var documentPath = GetDocumentFilePath(uri);
-        await UnloadProjectAsync(documentPath);
+        return await TryUnloadProjectAsync(documentPath);
     }
 
     protected override async Task<RemoteProjectLoadResult?> TryLoadProjectInMSBuildHostAsync(

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/HostWorkspace/LanguageServerProjectLoader.cs

@@ -461,15 +461,15 @@ protected async Task BeginLoadingProjectAsync(string projectPath, string? projec
 
     protected Task WaitForProjectsToFinishLoadingAsync() => _projectsToReload.WaitUntilCurrentBatchCompletesAsync();
 
-    protected async ValueTask UnloadProjectAsync(string projectPath)
+    protected async ValueTask<bool> TryUnloadProjectAsync(string projectPath)
     {
         using (await _gate.DisposableWaitAsync(CancellationToken.None))
         {
             if (!_loadedProjects.Remove(projectPath, out var loadState))
             {
                 // It is common to be called with a path to a project which is already not loaded.
                 // In this case, we should do nothing.
-                return;
+                return false;
             }
 
             if (loadState is ProjectLoadState.Primordial(var projectFactory, var projectId))
@@ -491,5 +491,6 @@ protected async ValueTask UnloadProjectAsync(string projectPath)
         }
 
         await OnProjectUnloadedAsync(projectPath);
+        return true;
     }
 }

src/LanguageServer/Protocol.TestUtilities/AbstractLspMiscellaneousFilesWorkspaceTests.cs

@@ -348,6 +348,7 @@ public async Task TestLspTransfersFromMiscellaneousFilesToHostWorkspaceAsync(boo
         AssertEx.NotNull(miscDocument);
         var miscText = await miscDocument.GetTextAsync(CancellationToken.None);
         Assert.Equal("More LSP textLSP text", miscText.ToString());
+        Assert.True(await testLspServer.GetManagerAccessor().IsMiscellaneousFilesDocumentAsync(miscDocument));
 
         // Update the registered workspace with the new document.
         var newDocumentId = (await AddDocumentAsync(testLspServer, newDocumentFilePath, "New Doc")).Id;
@@ -361,12 +362,16 @@ public async Task TestLspTransfersFromMiscellaneousFilesToHostWorkspaceAsync(boo
         // Verify we still are using the tracked LSP text for the document.
         var documentText = await document.GetTextAsync(CancellationToken.None);
         Assert.Equal("More LSP textLSP text", documentText.ToString());
+
+        // There should not be any other misc document in the solution anymore.
+        var matchingDocuments = await document.Project.Solution.GetTextDocumentsAsync(newDocumentUri, CancellationToken.None);
+        Assert.Single(matchingDocuments);
     }
 
     private protected abstract ValueTask<Document> AddDocumentAsync(TestLspServer testLspServer, string filePath, string content);
     private protected abstract Workspace GetHostWorkspace(TestLspServer testLspServer);
 
-    private static async Task<(Workspace? workspace, Document? document)> GetLspWorkspaceAndDocumentAsync(DocumentUri uri, TestLspServer testLspServer)
+    private protected static async Task<(Workspace? workspace, Document? document)> GetLspWorkspaceAndDocumentAsync(DocumentUri uri, TestLspServer testLspServer)
     {
         var (workspace, _, document) = await testLspServer.GetManager().GetLspDocumentInfoAsync(CreateTextDocumentIdentifier(uri), CancellationToken.None).ConfigureAwait(false);
         return (workspace, document as Document);

src/LanguageServer/Protocol/Workspaces/ILspMiscellaneousFilesWorkspaceProvider.cs

@@ -36,5 +36,6 @@ internal interface ILspMiscellaneousFilesWorkspaceProvider : ILspService
     /// Note that the implementation of this method should not depend on anything expensive such as RPC calls.
     /// async is used here to allow taking locks asynchronously and "relatively fast" stuff like that.
     /// </summary>
-    ValueTask TryRemoveMiscellaneousDocumentAsync(DocumentUri uri);
+    /// <returns><see langword="true"/> when a document was found and removed</returns>
+    ValueTask<bool> TryRemoveMiscellaneousDocumentAsync(DocumentUri uri);
 }

src/LanguageServer/Protocol/Workspaces/LspMiscellaneousFilesWorkspaceProvider.cs

@@ -83,7 +83,7 @@ public ValueTask<bool> IsMiscellaneousFilesDocumentAsync(TextDocument document,
     /// Calls to this method and <see cref="AddMiscellaneousDocument(DocumentUri, SourceText, string, ILspLogger)"/> are made
     /// from LSP text sync request handling which do not run concurrently.
     /// </summary>
-    public ValueTask TryRemoveMiscellaneousDocumentAsync(DocumentUri uri)
+    public ValueTask<bool> TryRemoveMiscellaneousDocumentAsync(DocumentUri uri)
     {
         // We'll only ever have a single document matching this URI in the misc solution.
         var matchingDocument = CurrentSolution.GetDocumentIds(uri).SingleOrDefault();
@@ -102,9 +102,11 @@ public ValueTask TryRemoveMiscellaneousDocumentAsync(DocumentUri uri)
             // so it should never have other documents in it.
             var project = CurrentSolution.GetRequiredProject(matchingDocument.ProjectId);
             OnProjectRemoved(project.Id);
+
+            return ValueTask.FromResult(true);
         }
 
-        return ValueTask.CompletedTask;
+        return ValueTask.FromResult(false);
     }
 
     public ValueTask UpdateTextIfPresentAsync(DocumentId documentId, SourceText sourceText, CancellationToken cancellationToken)

src/LanguageServer/Protocol/Workspaces/LspWorkspaceManager.cs

@@ -251,28 +251,29 @@ public void UpdateTrackedDocument(DocumentUri uri, SourceText newSourceText, int
         foreach (var (workspace, lspSolution, isForked) in lspSolutions)
         {
             var documents = await lspSolution.GetTextDocumentsAsync(textDocumentIdentifier.DocumentUri, cancellationToken).ConfigureAwait(false);
+
             if (documents.Length > 0)
             {
-                // We have at least one document, so find the one in the right project context
+                // We have at least one document, so find the one in the right project context.
                 var document = documents.FindDocumentInProjectContext(textDocumentIdentifier, (sln, id) => sln.GetRequiredTextDocument(id));
 
                 if (_lspMiscellaneousFilesWorkspaceProvider is not null)
                 {
-                    // If we started with multiple documents and didn't have specific context information, it's possible we picked a miscellaneous files document when
-                    // we could have picked a real one.
-                    if (documents.Length > 1 && await _lspMiscellaneousFilesWorkspaceProvider.IsMiscellaneousFilesDocumentAsync(document, cancellationToken).ConfigureAwait(false))
-                    {
-                        // Pick a different one; our choice here is arbitrary, since if we had a specified context in the first place we would have picked the right one.
-                        document = documents.First(d => d != document);
-                    }
-
-                    // If we found the document in a non-misc workspace (either immediately or by the correction above), also attempt to remove it from the misc workspace
-                    // if it happens to be in there as well.
-                    if (_lspMiscellaneousFilesWorkspaceProvider is not null && !await _lspMiscellaneousFilesWorkspaceProvider.IsMiscellaneousFilesDocumentAsync(document, cancellationToken).ConfigureAwait(false))
+                    // It is possible that a document that was previously a misc file is now part of a real workspace (e.g. project system told us about a file we already had open).
+                    // If we found a non-misc document, we should clean up any references to it in the misc provider.
+                    var foundNonMiscDocument = await documents
+                        .AnyAsync(async doc => !await _lspMiscellaneousFilesWorkspaceProvider.IsMiscellaneousFilesDocumentAsync(doc, cancellationToken).ConfigureAwait(false))
+                        .ConfigureAwait(false);
+                    if (foundNonMiscDocument)
                     {
                         try
                         {
-                            await _lspMiscellaneousFilesWorkspaceProvider.TryRemoveMiscellaneousDocumentAsync(uri).ConfigureAwait(false);
+                            var didRemove = await _lspMiscellaneousFilesWorkspaceProvider.TryRemoveMiscellaneousDocumentAsync(uri).ConfigureAwait(false);
+                            if (didRemove)
+                            {
+                                // If we actually removed something, lookup the document again to ensure we return updated solutions without the misc document.
+                                return await GetLspDocumentInfoAsync(textDocumentIdentifier, cancellationToken).ConfigureAwait(false);
+                            }
                         }
                         catch (Exception ex) when (FatalError.ReportAndCatch(ex))
                         {