Cohosting test infrastructure v2

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteClientInitializationService.cs

@@ -1,12 +1,13 @@
 // 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.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
-internal interface IRemoteClientInitializationService
+internal interface IRemoteClientInitializationService : IDisposable
 {
     ValueTask InitializeAsync(RemoteClientInitializationOptions initializationOptions, CancellationToken cancellationToken);
 

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteFoldingRangeService.cs

@@ -1,6 +1,7 @@
 // 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.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
@@ -9,7 +10,7 @@
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
-internal interface IRemoteFoldingRangeService
+internal interface IRemoteFoldingRangeService : IDisposable
 {
     ValueTask<ImmutableArray<RemoteFoldingRange>> GetFoldingRangesAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId documentId, ImmutableArray<RemoteFoldingRange> htmlRanges, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteHtmlDocumentService.cs

@@ -1,13 +1,14 @@
 // 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.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
-internal interface IRemoteHtmlDocumentService
+internal interface IRemoteHtmlDocumentService : IDisposable
 {
     ValueTask<string?> GetHtmlDocumentTextAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteLinkedEditingRangeService.cs

@@ -1,14 +1,15 @@
 // 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.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
-internal interface IRemoteLinkedEditingRangeService
+internal interface IRemoteLinkedEditingRangeService : IDisposable
 {
     ValueTask<LinePositionSpan[]?> GetRangesAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, LinePosition linePosition, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteSemanticTokensService.cs

@@ -9,7 +9,7 @@
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
-internal interface IRemoteSemanticTokensService
+internal interface IRemoteSemanticTokensService : IDisposable
 {
     ValueTask<int[]?> GetSemanticTokensDataAsync(
         RazorPinnedSolutionInfoWrapper solutionInfo,

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteServiceProvider.cs

@@ -17,5 +17,5 @@ internal interface IRemoteServiceProvider
         CancellationToken cancellationToken,
         [CallerFilePath] string? callerFilePath = null,
         [CallerMemberName] string? callerMemberName = null)
-        where TService : class;
+        where TService : class, IDisposable;
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteTagHelperProviderService.cs

@@ -1,6 +1,7 @@
 // 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.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
@@ -10,7 +11,7 @@
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
-internal interface IRemoteTagHelperProviderService
+internal interface IRemoteTagHelperProviderService : IDisposable
 {
     ValueTask<TagHelperDeltaResult> GetTagHelpersDeltaAsync(
         RazorPinnedSolutionInfoWrapper solutionInfo,

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteUriPresentationService.cs

@@ -9,7 +9,7 @@
 
 namespace Microsoft.CodeAnalysis.Razor.Remote;
 
-internal interface IRemoteUriPresentationService
+internal interface IRemoteUriPresentationService : IDisposable
 {
     ValueTask<TextChange?> GetPresentationAsync(RazorPinnedSolutionInfoWrapper solutionInfo, DocumentId razorDocumentId, LinePositionSpan span, Uri[]? uris, CancellationToken cancellationToken);
 }

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

@@ -11,13 +11,12 @@
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
-using Microsoft.ServiceHub.Framework;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
 internal sealed class RemoteFoldingRangeService(
-    IServiceBroker serviceBroker,
+    IRazorServiceBroker serviceBroker,
     IFoldingRangeService foldingRangeService,
     DocumentSnapshotFactory documentSnapshotFactory,
     IFilePathService filePathService)

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

@@ -5,7 +5,6 @@
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
-using Microsoft.ServiceHub.Framework;
 using Microsoft.VisualStudio.Composition;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
@@ -18,7 +17,7 @@ public RemoteFoldingRangeServiceFactory()
     {
     }
 
-    protected override IRemoteFoldingRangeService CreateService(IServiceBroker serviceBroker, ExportProvider exportProvider)
+    protected override IRemoteFoldingRangeService CreateService(IRazorServiceBroker serviceBroker, ExportProvider exportProvider)
     {
         var infoService = exportProvider.GetExportedValue<IFoldingRangeService>();
         var documentSnapshotFactory = exportProvider.GetExportedValue<DocumentSnapshotFactory>();

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/HtmlDocuments/RemoteHtmlDocumentService.cs

@@ -7,12 +7,11 @@
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
-using Microsoft.ServiceHub.Framework;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
 internal sealed class RemoteHtmlDocumentService(
-    IServiceBroker serviceBroker,
+    IRazorServiceBroker serviceBroker,
     DocumentSnapshotFactory documentSnapshotFactory)
     : RazorDocumentServiceBase(serviceBroker, documentSnapshotFactory), IRemoteHtmlDocumentService
 {

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/HtmlDocuments/RemoteHtmlDocumentServiceFactory.cs

@@ -3,7 +3,6 @@
 
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
-using Microsoft.ServiceHub.Framework;
 using Microsoft.VisualStudio.Composition;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
@@ -16,7 +15,7 @@ public RemoteHtmlDocumentServiceFactory()
     {
     }
 
-    protected override IRemoteHtmlDocumentService CreateService(IServiceBroker serviceBroker, ExportProvider exportProvider)
+    protected override IRemoteHtmlDocumentService CreateService(IRazorServiceBroker serviceBroker, ExportProvider exportProvider)
     {
         var documentSnapshotFactory = exportProvider.GetExportedValue<DocumentSnapshotFactory>();
         return new RemoteHtmlDocumentService(serviceBroker, documentSnapshotFactory);

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/IBrokeredServiceInterceptor.cs → src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/IRazorServiceBroker.cs

@@ -12,7 +12,7 @@ namespace Microsoft.CodeAnalysis.Remote.Razor;
 /// <summary>
 /// An abstraction to avoid calling the static <see cref="RazorBrokeredServiceImplementation"/> helper defined in Roslyn
 /// </summary>
-internal interface IBrokeredServiceInterceptor
+internal interface IRazorServiceBroker : IDisposable
 {
     ValueTask RunServiceAsync(Func<CancellationToken, ValueTask> implementation, CancellationToken cancellationToken);
 

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Initialization/RemoteClientInitializationService.cs

@@ -5,12 +5,11 @@
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Remote.Razor.SemanticTokens;
-using Microsoft.ServiceHub.Framework;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
 internal sealed class RemoteClientInitializationService(
-    IServiceBroker serviceBroker)
+    IRazorServiceBroker serviceBroker)
     : RazorServiceBase(serviceBroker), IRemoteClientInitializationService
 {
     public ValueTask InitializeAsync(RemoteClientInitializationOptions options, CancellationToken cancellationToken)

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Initialization/RemoteClientInitializationServiceFactory.cs

@@ -2,7 +2,6 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using Microsoft.CodeAnalysis.Razor.Remote;
-using Microsoft.ServiceHub.Framework;
 using Microsoft.VisualStudio.Composition;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
@@ -15,7 +14,7 @@ public RemoteClientInitializationServiceFactory()
     {
     }
 
-    protected override IRemoteClientInitializationService CreateService(IServiceBroker serviceBroker, ExportProvider exportProvider)
+    protected override IRemoteClientInitializationService CreateService(IRazorServiceBroker serviceBroker, ExportProvider exportProvider)
     {
         return new RemoteClientInitializationService(serviceBroker);
     }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/LinkedEditingRange/RemoteLinkedEditingRangeService.cs

@@ -9,12 +9,11 @@
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
-using Microsoft.ServiceHub.Framework;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
 internal sealed class RemoteLinkedEditingRangeService(
-    IServiceBroker serviceBroker,
+    IRazorServiceBroker serviceBroker,
     DocumentSnapshotFactory documentSnapshotFactory,
     ILoggerFactory loggerFactory)
     : RazorDocumentServiceBase(serviceBroker, documentSnapshotFactory), IRemoteLinkedEditingRangeService

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/LinkedEditingRange/RemoteLinkedEditingRangeServiceFactory.cs

@@ -4,7 +4,6 @@
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
-using Microsoft.ServiceHub.Framework;
 using Microsoft.VisualStudio.Composition;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
@@ -17,7 +16,7 @@ public RemoteLinkedEditingRangeServiceFactory()
     {
     }
 
-    protected override IRemoteLinkedEditingRangeService CreateService(IServiceBroker serviceBroker, ExportProvider exportProvider)
+    protected override IRemoteLinkedEditingRangeService CreateService(IRazorServiceBroker serviceBroker, ExportProvider exportProvider)
     {
         var documentSnapshotFactory = exportProvider.GetExportedValue<DocumentSnapshotFactory>();
         var loggerFactory = exportProvider.GetExportedValue<ILoggerFactory>();

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/RazorDocumentServiceBase.cs

@@ -6,12 +6,11 @@
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
-using Microsoft.ServiceHub.Framework;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
 internal abstract class RazorDocumentServiceBase(
-    IServiceBroker serviceBroker,
+    IRazorServiceBroker serviceBroker,
     DocumentSnapshotFactory documentSnapshotFactory)
     : RazorServiceBase(serviceBroker)
 {

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/RazorServiceBase.cs

@@ -5,30 +5,26 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
-using Microsoft.CodeAnalysis.ExternalAccess.Razor.Api;
-using Microsoft.ServiceHub.Framework;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
 internal abstract class RazorServiceBase : IDisposable
 {
-    private readonly ServiceBrokerClient _serviceBrokerClient;
-    private readonly IBrokeredServiceInterceptor? _brokeredServiceInterceptor;
+    private readonly IRazorServiceBroker _razorServiceBroker;
 
-    public RazorServiceBase(IServiceBroker serviceBroker)
+    public RazorServiceBase(IRazorServiceBroker razorServiceBroker)
     {
-        _brokeredServiceInterceptor = serviceBroker as IBrokeredServiceInterceptor;
-        _serviceBrokerClient = new ServiceBrokerClient(serviceBroker, joinableTaskFactory: null);
+        _razorServiceBroker = razorServiceBroker;
     }
 
     protected ValueTask RunServiceAsync(Func<CancellationToken, ValueTask> implementation, CancellationToken cancellationToken)
-        => _brokeredServiceInterceptor?.RunServiceAsync(implementation, cancellationToken) ?? RazorBrokeredServiceImplementation.RunServiceAsync(implementation, cancellationToken);
+        => _razorServiceBroker.RunServiceAsync(implementation, cancellationToken);
 
     protected ValueTask<T> RunServiceAsync<T>(RazorPinnedSolutionInfoWrapper solutionInfo, Func<Solution, ValueTask<T>> implementation, CancellationToken cancellationToken)
-        => _brokeredServiceInterceptor?.RunServiceAsync(solutionInfo, implementation, cancellationToken) ?? RazorBrokeredServiceImplementation.RunServiceAsync(solutionInfo, _serviceBrokerClient, implementation, cancellationToken);
+        => _razorServiceBroker.RunServiceAsync(solutionInfo, implementation, cancellationToken);
 
     public void Dispose()
     {
-        _serviceBrokerClient.Dispose();
+        _razorServiceBroker.Dispose();
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/RazorServiceBroker.cs

@@ -0,0 +1,32 @@
+// 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.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Api;
+using Microsoft.ServiceHub.Framework;
+
+namespace Microsoft.CodeAnalysis.Remote.Razor;
+
+internal class RazorServiceBroker : IRazorServiceBroker
+{
+    private readonly ServiceBrokerClient _serviceBrokerClient;
+
+    public RazorServiceBroker(IServiceBroker serviceBroker)
+    {
+        _serviceBrokerClient = new ServiceBrokerClient(serviceBroker, joinableTaskFactory: null);
+    }
+
+    public ValueTask RunServiceAsync(Func<CancellationToken, ValueTask> implementation, CancellationToken cancellationToken)
+        => RazorBrokeredServiceImplementation.RunServiceAsync(implementation, cancellationToken);
+
+    public ValueTask<T> RunServiceAsync<T>(RazorPinnedSolutionInfoWrapper solutionInfo, Func<Solution, ValueTask<T>> implementation, CancellationToken cancellationToken)
+        => RazorBrokeredServiceImplementation.RunServiceAsync(solutionInfo, _serviceBrokerClient, implementation, cancellationToken);
+
+    public void Dispose()
+    {
+        _serviceBrokerClient.Dispose();
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/RazorServiceFactoryBase.cs

@@ -35,11 +35,11 @@ public RazorServiceFactoryBase(RazorServiceDescriptorsWrapper razorServiceDescri
     }
 
     public async Task<object> CreateAsync(
-       Stream stream,
-       IServiceProvider hostProvidedServices,
-       ServiceActivationOptions serviceActivationOptions,
-       IServiceBroker serviceBroker,
-       AuthorizationServiceClient? authorizationServiceClient)
+        Stream stream,
+        IServiceProvider hostProvidedServices,
+        ServiceActivationOptions serviceActivationOptions,
+        IServiceBroker serviceBroker,
+        AuthorizationServiceClient? authorizationServiceClient)
     {
         // Dispose the AuthorizationServiceClient since we won't be using it
         authorizationServiceClient?.Dispose();
@@ -53,13 +53,26 @@ public async Task<object> CreateAsync(
 
         var exportProvider = await RemoteMefComposition.GetExportProviderAsync().ConfigureAwait(false);
 
-        var service = CreateService(serviceBroker, exportProvider);
+        var razorServiceBroker = new RazorServiceBroker(serviceBroker);
+
+        var service = CreateService(razorServiceBroker, exportProvider);
 
         serverConnection.AddLocalRpcTarget(service);
         serverConnection.StartListening();
 
         return service;
     }
 
-    protected abstract TService CreateService(IServiceBroker serviceBroker, ExportProvider exportProvider);
+    protected abstract TService CreateService(IRazorServiceBroker serviceBroker, ExportProvider exportProvider);
+
+    internal TestAccessor GetTestAccessor()
+    {
+        return new TestAccessor(this);
+    }
+
+    internal readonly struct TestAccessor(RazorServiceFactoryBase<TService> instance)
+    {
+        public TService CreateService(IRazorServiceBroker serviceBroker, ExportProvider exportProvider)
+            => instance.CreateService(serviceBroker, exportProvider);
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/RemoteMefComposition.cs

@@ -20,8 +20,9 @@ internal static class RemoteMefComposition
     public static Task<ExportProvider> GetExportProviderAsync()
         => s_exportProviderLazy.GetValueAsync();
 
+    // Internal for testing
     // Inspired by https://github.com/dotnet/roslyn/blob/25aa74d725e801b8232dbb3e5abcda0fa72da8c5/src/Workspaces/Remote/ServiceHub/Host/RemoteWorkspaceManager.cs#L77
-    private static async Task<ExportProvider> CreateExportProviderAsync()
+    internal static async Task<ExportProvider> CreateExportProviderAsync()
     {
         var resolver = new Resolver(SimpleAssemblyLoader.Instance);
         var discovery = new AttributedPartDiscovery(resolver, isNonPublicSupported: true); // MEFv2 only

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/SemanticTokens/RemoteSemanticTokensService.cs

@@ -9,12 +9,11 @@
 using Microsoft.CodeAnalysis.Razor.SemanticTokens;
 using Microsoft.CodeAnalysis.Remote.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Text;
-using Microsoft.ServiceHub.Framework;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
 internal sealed class RemoteSemanticTokensService(
-    IServiceBroker serviceBroker,
+    IRazorServiceBroker serviceBroker,
     IRazorSemanticTokensInfoService razorSemanticTokensInfoService,
     DocumentSnapshotFactory documentSnapshotFactory)
     : RazorDocumentServiceBase(serviceBroker, documentSnapshotFactory), IRemoteSemanticTokensService