Add LSP handler for copilot related documents

src/LanguageServer/Protocol/Extensions/ProtocolConversions.cs

@@ -309,13 +309,14 @@ public static LSP.TextDocumentPositionParams PositionToTextDocumentPositionParam
         }
 
         public static LSP.TextDocumentIdentifier DocumentToTextDocumentIdentifier(TextDocument document)
-            => new LSP.TextDocumentIdentifier { Uri = document.GetURI() };
+            => new() { Uri = document.GetURI() };
 
         public static LSP.VersionedTextDocumentIdentifier DocumentToVersionedTextDocumentIdentifier(Document document)
-            => new LSP.VersionedTextDocumentIdentifier { Uri = document.GetURI() };
+            => new() { Uri = document.GetURI() };
 
         public static LinePosition PositionToLinePosition(LSP.Position position)
-            => new LinePosition(position.Line, position.Character);
+            => new(position.Line, position.Character);
+
         public static LinePositionSpan RangeToLinePositionSpan(LSP.Range range)
             => new(PositionToLinePosition(range.Start), PositionToLinePosition(range.End));
 

src/LanguageServer/Protocol/Handler/RelatedDocuments/RelatedDocumentsHandler.cs

@@ -0,0 +1,147 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Composition;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.RelatedDocuments;
+using Microsoft.CodeAnalysis.Serialization;
+using Microsoft.CodeAnalysis.Shared.Extensions;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.CommonLanguageServerProtocol.Framework;
+using Roslyn.LanguageServer.Protocol;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.LanguageServer.Handler.RelatedDocuments;
+
+[ExportCSharpVisualBasicLspServiceFactory(typeof(RelatedDocumentsHandler)), Shared]
+[method: ImportingConstructor]
+[method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+internal sealed class RelatedDocumentsHandlerFactory() : ILspServiceFactory
+{
+    public ILspService CreateILspService(LspServices lspServices, WellKnownLspServerKinds serverKind)
+        => new RelatedDocumentsHandler();
+}
+
+[Method(VSInternalMethods.CopilotRelatedDocumentsName)]
+internal sealed class RelatedDocumentsHandler
+    : ILspServiceRequestHandler<VSInternalRelatedDocumentParams, VSInternalRelatedDocumentReport[]?>,
+      ITextDocumentIdentifierHandler<VSInternalRelatedDocumentParams, TextDocumentIdentifier>
+{
+    /// <summary>
+    /// Cache where we store the data produced by prior requests so that they can be returned if nothing of significance
+    /// changed. The version key is produced by combining the checksums for project options <see
+    /// cref="ProjectState.GetParseOptionsChecksum"/> and <see cref="DocumentStateChecksums.Text"/>
+    /// </summary>
+    private readonly VersionedPullCache<(Checksum parseOptionsChecksum, Checksum textChecksum)?> _versionedCache = new(nameof(RelatedDocumentsHandler));
+
+    public bool MutatesSolutionState => false;
+    public bool RequiresLSPSolution => true;
+
+    private static async Task<(Checksum parseOptionsChecksum, Checksum textChecksum)> ComputeChecksumsAsync(Document document, CancellationToken cancellationToken)
+    {
+        var project = document.Project;
+        var parseOptionsChecksum = project.State.GetParseOptionsChecksum();
+
+        var documentChecksumState = await document.State.GetStateChecksumsAsync(cancellationToken).ConfigureAwait(false);
+        var textChecksum = documentChecksumState.Text;
+
+        return (parseOptionsChecksum, textChecksum);
+    }
+
+    public TextDocumentIdentifier GetTextDocumentIdentifier(VSInternalRelatedDocumentParams requestParams)
+        => requestParams.TextDocument;
+
+    /// <summary>
+    /// Retrieve the previous results we reported.  Used so we can avoid resending data for unchanged files.
+    /// </summary>
+    private static ImmutableArray<PreviousPullResult>? GetPreviousResults(VSInternalRelatedDocumentParams requestParams)
+        => requestParams.PreviousResultId != null && requestParams.TextDocument != null
+            ? [new PreviousPullResult(requestParams.PreviousResultId, requestParams.TextDocument)]
+            // The client didn't provide us with a previous result to look for, so we can't lookup anything.
+            : null;
+
+    public async Task<VSInternalRelatedDocumentReport[]?> HandleRequestAsync(
+        VSInternalRelatedDocumentParams requestParams, RequestContext context, CancellationToken cancellationToken)
+    {
+        context.TraceInformation($"{this.GetType()} started getting related documents");
+
+        // The progress object we will stream reports to.
+        using var progress = BufferedProgress.Create(requestParams.PartialResultToken);
+
+        context.TraceInformation($"PreviousResultId={requestParams.PreviousResultId}");
+
+        var solution = context.Solution;
+        var document = context.Document;
+        Contract.ThrowIfNull(solution);
+        Contract.ThrowIfNull(document);
+
+        context.TraceInformation($"Processing: {document.FilePath}");
+
+        var relatedDocumentsService = document.GetLanguageService<IRelatedDocumentsService>();
+        if (relatedDocumentsService == null)
+        {
+            context.TraceInformation($"Ignoring document '{document.FilePath}' because it does not support related documents");
+        }
+        else
+        {
+            var documentToPreviousParams = new Dictionary<Document, PreviousPullResult>();
+            if (requestParams.PreviousResultId != null)
+                documentToPreviousParams.Add(document, new PreviousPullResult(requestParams.PreviousResultId, requestParams.TextDocument));
+
+            var newResultId = await _versionedCache.GetNewResultIdAsync(
+                documentToPreviousParams,
+                document,
+                computeVersionAsync: async () => await ComputeChecksumsAsync(document, cancellationToken).ConfigureAwait(false),
+                cancellationToken).ConfigureAwait(false);
+            if (newResultId != null)
+            {
+                context.TraceInformation($"Version was changed for document: {document.FilePath}");
+
+                var linePosition = requestParams.Position is null
+                    ? new LinePosition(0, 0)
+                    : ProtocolConversions.PositionToLinePosition(requestParams.Position);
+
+                var text = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
+                var position = text.Lines.GetPosition(linePosition);
+
+                await relatedDocumentsService.GetRelatedDocumentIdsAsync(
+                    document,
+                    position,
+                    (relatedDocumentIds, cancellationToken) =>
+                    {
+                        // As the related docs services reports document ids to us, stream those immediately through our
+                        // progress reporter.
+                        progress.Report(new VSInternalRelatedDocumentReport
+                        {
+                            ResultId = newResultId,
+                            FilePaths = relatedDocumentIds.Select(id => solution.GetRequiredDocument(id).FilePath).WhereNotNull().ToArray(),
+                        });
+
+                        return ValueTaskFactory.CompletedTask;
+                    },
+                    cancellationToken).ConfigureAwait(false);
+            }
+            else
+            {
+                context.TraceInformation($"Version was unchanged for document: {document.FilePath}");
+
+                // Nothing changed between the last request and this one.  Report a (null-file-paths, same-result-id)
+                // response to the client as that means they should just preserve the current related file paths they
+                // have for this file.
+                progress.Report(new VSInternalRelatedDocumentReport { ResultId = requestParams.PreviousResultId });
+            }
+        }
+
+        // If we had a progress object, then we will have been reporting to that.  Otherwise, take what we've been
+        // collecting and return that.
+        context.TraceInformation($"{this.GetType()} finished getting related documents");
+        return progress.GetFlattenedValues();
+    }
+}

src/LanguageServer/Protocol/Microsoft.CodeAnalysis.LanguageServer.Protocol.csproj

@@ -117,6 +117,10 @@
     </EmbeddedResource>
   </ItemGroup>
 
+  <ItemGroup>
+    <Folder Include="Handler\RelatedDocuments\" />
+  </ItemGroup>
+
   <Import Project="..\Microsoft.CommonLanguageServerProtocol.Framework\Microsoft.CommonLanguageServerProtocol.Framework.Shared.projitems" Label="Shared" />
 
 </Project>

src/LanguageServer/Protocol/Protocol/Internal/VSInternalDocumentSpellCheckableParams.cs

@@ -8,9 +8,9 @@ namespace Roslyn.LanguageServer.Protocol
     using System.Text.Json.Serialization;
 
     /// <summary>
-    /// Parameter for tD/_vs_spellCheckableRanges.
+    /// Parameter for textDocument/_vs_spellCheckableRanges.
     /// </summary>
-    internal class VSInternalDocumentSpellCheckableParams : VSInternalStreamingParams, IPartialResultParams<VSInternalSpellCheckableRangeReport[]>
+    internal sealed class VSInternalDocumentSpellCheckableParams : VSInternalStreamingParams, IPartialResultParams<VSInternalSpellCheckableRangeReport[]>
     {
         /// <inheritdoc/>
         [JsonPropertyName(Methods.PartialResultTokenName)]

src/LanguageServer/Protocol/Protocol/Internal/VSInternalMethods.cs

@@ -9,6 +9,11 @@ namespace Roslyn.LanguageServer.Protocol
     /// </summary>
     internal static class VSInternalMethods
     {
+        /// <summary>
+        /// Method name for 'copilot/_related_documents'.
+        /// </summary>
+        public const string CopilotRelatedDocumentsName = "copilot/_related_documents";
+
         /// <summary>
         /// Method name for 'textDocument/foldingRange/_vs_refresh'.
         /// </summary>

src/LanguageServer/Protocol/Protocol/Internal/VSInternalRelatedDocumentParams.cs

@@ -0,0 +1,44 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+namespace Roslyn.LanguageServer.Protocol
+{
+    using System;
+    using System.Text.Json.Serialization;
+
+    /// <summary>
+    /// Parameter for copilot/_related_documents.
+    /// </summary>
+    internal sealed class VSInternalRelatedDocumentParams : VSInternalStreamingParams, IPartialResultParams<VSInternalRelatedDocumentReport[]>
+    {
+        /// <summary>
+        /// Gets or sets the value which indicates the position within the document.
+        /// </summary>
+        [JsonPropertyName("position")]
+        [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+        public Position? Position { get; set; }
+
+        /// <inheritdoc/>
+        [JsonPropertyName(Methods.PartialResultTokenName)]
+        [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+        public IProgress<VSInternalRelatedDocumentReport[]>? PartialResultToken { get; set; }
+    }
+
+    internal sealed class VSInternalRelatedDocumentReport
+    {
+        /// <summary>
+        /// Gets or sets the server-generated version number for the related documents result. This is treated as a
+        /// black box by the client: it is stored on the client for each textDocument and sent back to the server when
+        /// requesting related documents. The server can use this result ID to avoid resending results
+        /// that had previously been sent.
+        /// </summary>
+        [JsonPropertyName("_vs_resultId")]
+        [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+        public string? ResultId { get; set; }
+
+        [JsonPropertyName("_vs_file_paths")]
+        [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+        public string[]? FilePaths { get; set; }
+    }
+}

src/LanguageServer/ProtocolUnitTests/RelatedDocuments/RelatedDocumentsTests.cs

@@ -0,0 +1,137 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.LanguageServer.Handler;
+using Roslyn.LanguageServer.Protocol;
+using Roslyn.Test.Utilities;
+using Roslyn.Utilities;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests.RelatedDocuments;
+
+public sealed class RelatedDocumentsTests(ITestOutputHelper testOutputHelper)
+    : AbstractLanguageServerProtocolTests(testOutputHelper)
+{
+    private static async Task<VSInternalRelatedDocumentReport[]> RunGetRelatedDocumentsAsync(
+        TestLspServer testLspServer,
+        Uri uri,
+        string? previousResultId = null,
+        bool useProgress = false)
+    {
+        BufferedProgress<VSInternalRelatedDocumentReport[]>? progress = useProgress ? BufferedProgress.Create<VSInternalRelatedDocumentReport[]>(null) : null;
+        var spans = await testLspServer.ExecuteRequestAsync<VSInternalRelatedDocumentParams, VSInternalRelatedDocumentReport[]>(
+            VSInternalMethods.CopilotRelatedDocumentsName,
+            new VSInternalRelatedDocumentParams
+            {
+                TextDocument = new TextDocumentIdentifier { Uri = uri },
+                PreviousResultId = previousResultId,
+                PartialResultToken = progress,
+            },
+            CancellationToken.None).ConfigureAwait(false);
+
+        if (useProgress)
+        {
+            Assert.Null(spans);
+            spans = progress!.Value.GetFlattenedValues();
+        }
+
+        AssertEx.NotNull(spans);
+        return spans;
+    }
+
+    [Theory, CombinatorialData]
+    public async Task ReferenceNoDocuments(bool mutatingLspWorkspace, bool useProgress)
+    {
+        var markup1 = """
+            class X
+            {
+            }
+            """;
+
+        var markup2 = """
+            class Y
+            {
+            }
+            """;
+
+        await using var testLspServer = await CreateTestLspServerAsync([markup1, markup2], mutatingLspWorkspace);
+
+        var project = testLspServer.TestWorkspace.CurrentSolution.Projects.Single();
+        var results = await RunGetRelatedDocumentsAsync(
+            testLspServer,
+            project.Documents.First().GetURI(),
+            useProgress: useProgress);
+
+        Assert.Equal(0, results.Length);
+    }
+
+    [Theory, CombinatorialData]
+    public async Task ReferenceSingleOtherDocument(bool mutatingLspWorkspace, bool useProgress)
+    {
+        var markup1 = """
+            class X
+            {
+                Y y;
+            }
+            """;
+
+        var markup2 = """
+            class Y
+            {
+            }
+            """;
+
+        await using var testLspServer = await CreateTestLspServerAsync([markup1, markup2], mutatingLspWorkspace);
+
+        var project = testLspServer.TestWorkspace.CurrentSolution.Projects.Single();
+        var results = await RunGetRelatedDocumentsAsync(
+            testLspServer,
+            project.Documents.First().GetURI(),
+            useProgress: useProgress);
+
+        Assert.Equal(1, results.Length);
+        Assert.Equal(project.Documents.Last().FilePath, results[0].FilePaths.Single());
+    }
+
+    [Theory, CombinatorialData]
+    public async Task ReferenceMultipleOtherDocument(bool mutatingLspWorkspace, bool useProgress)
+    {
+        var markup1 = """
+            class X
+            {
+                Y y;
+                Z z;
+            }
+            """;
+
+        var markup2 = """
+            class Y
+            {
+            }
+            """;
+
+        var markup3 = """
+            class Z
+            {
+            }
+            """;
+
+        await using var testLspServer = await CreateTestLspServerAsync([markup1, markup2, markup3], mutatingLspWorkspace);
+
+        var project = testLspServer.TestWorkspace.CurrentSolution.Projects.Single();
+        var results = await RunGetRelatedDocumentsAsync(
+            testLspServer,
+            project.Documents.First().GetURI(),
+            useProgress: useProgress);
+
+        Assert.Equal(1, results.Length);
+        Assert.Equal(2, results[0].FilePaths!.Length);
+        AssertEx.SetEqual([.. project.Documents.Skip(1).Select(d => d.FilePath)], results[0].FilePaths);
+    }
+}

src/LanguageServer/ProtocolUnitTests/SpellCheck/SpellCheckTests.cs

@@ -18,11 +18,9 @@
 
 namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests.SpellCheck
 {
-    public class SpellCheckTests : AbstractLanguageServerProtocolTests
+    public sealed class SpellCheckTests(ITestOutputHelper testOutputHelper)
+        : AbstractLanguageServerProtocolTests(testOutputHelper)
     {
-        public SpellCheckTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper)
-        {
-        }
 
         #region Document