Unify service configuration for ISseEventStreamStore, IMcpTaskStore, and ISessionMigrationHandler (#1362)

tests/ModelContextProtocol.AspNetCore.Tests/DistributedCacheResumabilityIntegrationTests.cs

@@ -22,17 +22,18 @@ namespace ModelContextProtocol.AspNetCore.Tests;
 /// </remarks>
 public class DistributedCacheResumabilityIntegrationTests(ITestOutputHelper testOutputHelper) : ResumabilityIntegrationTestsBase(testOutputHelper)
 {
-    private MemoryDistributedCache? _cache;
-
     /// <inheritdoc />
     protected override ValueTask<ISseEventStreamStore> CreateEventStreamStoreAsync()
     {
         // Create a new in-memory distributed cache for each test
-        _cache = new MemoryDistributedCache(Options.Create(new MemoryDistributedCacheOptions()));
+        var cache = new MemoryDistributedCache(Options.Create(new MemoryDistributedCacheOptions()));
 
         // Configure the store with shorter expiration times suitable for testing
         var options = new DistributedCacheEventStreamStoreOptions
         {
+            // Use the in-memory distributed cache
+            Cache = cache,
+
             // Use shorter polling interval for faster test execution
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(50),
 
@@ -43,7 +44,7 @@ public class DistributedCacheResumabilityIntegrationTests(ITestOutputHelper test
             MetadataAbsoluteExpiration = TimeSpan.FromMinutes(10),
         };
 
-        var store = new DistributedCacheEventStreamStore(_cache, options, LoggerFactory.CreateLogger<DistributedCacheEventStreamStore>());
+        var store = new DistributedCacheEventStreamStore(Options.Create(options), LoggerFactory.CreateLogger<DistributedCacheEventStreamStore>());
         return new ValueTask<ISseEventStreamStore>(store);
     }
 }

tests/ModelContextProtocol.AspNetCore.Tests/HttpMcpServerBuilderExtensionsTests.cs

@@ -0,0 +1,197 @@
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Caching.Distributed;
+using Microsoft.Extensions.Caching.Memory;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Options;
+using ModelContextProtocol.AspNetCore.Tests.Utils;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+
+namespace ModelContextProtocol.AspNetCore.Tests;
+
+public class HttpMcpServerBuilderExtensionsTests(ITestOutputHelper testOutputHelper) : KestrelInMemoryTest(testOutputHelper)
+{
+    [Fact]
+    public void WithDistributedCacheEventStreamStore_RegistersStoreInDI()
+    {
+        Builder.Services.AddDistributedMemoryCache();
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport()
+            .WithDistributedCacheEventStreamStore();
+
+        using var app = Builder.Build();
+
+        var store = app.Services.GetService<ISseEventStreamStore>();
+        Assert.IsType<DistributedCacheEventStreamStore>(store);
+    }
+
+    [Fact]
+    public void WithDistributedCacheEventStreamStore_ConfigureCallbackIsInvoked()
+    {
+        DistributedCacheEventStreamStoreOptions? capturedOptions = null;
+
+        Builder.Services.AddDistributedMemoryCache();
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport()
+            .WithDistributedCacheEventStreamStore(options => capturedOptions = options);
+
+        using var app = Builder.Build();
+
+        // Force options resolution to trigger the configure callback.
+        _ = app.Services.GetRequiredService<IOptions<DistributedCacheEventStreamStoreOptions>>().Value;
+
+        Assert.NotNull(capturedOptions);
+    }
+
+    [Fact]
+    public void WithDistributedCacheEventStreamStore_WorksWithoutDICache_WhenCacheSetViaCallback()
+    {
+        var explicitCache = new MemoryDistributedCache(Options.Create(new MemoryDistributedCacheOptions()));
+
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport()
+            .WithDistributedCacheEventStreamStore(options => options.Cache = explicitCache);
+
+        using var app = Builder.Build();
+
+        var store = app.Services.GetService<ISseEventStreamStore>();
+        Assert.IsType<DistributedCacheEventStreamStore>(store);
+    }
+
+    [Fact]
+    public void WithDistributedCacheEventStreamStore_ThrowsOptionsValidationException_WhenNoCacheConfigured()
+    {
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport()
+            .WithDistributedCacheEventStreamStore();
+
+        using var app = Builder.Build();
+
+        var ex = Assert.Throws<OptionsValidationException>(
+            () => app.Services.GetRequiredService<ISseEventStreamStore>());
+        Assert.StartsWith($"The '{nameof(DistributedCacheEventStreamStoreOptions)}.{nameof(DistributedCacheEventStreamStoreOptions.Cache)}'", ex.Message);
+    }
+
+    [Fact]
+    public void EventStreamStore_IsPopulatedFromDI_ViaPostConfigure()
+    {
+        Builder.Services.AddDistributedMemoryCache();
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport()
+            .WithDistributedCacheEventStreamStore();
+
+        using var app = Builder.Build();
+
+        var options = app.Services.GetRequiredService<IOptions<HttpServerTransportOptions>>().Value;
+        Assert.IsType<DistributedCacheEventStreamStore>(options.EventStreamStore);
+    }
+
+    [Fact]
+    public void EventStreamStore_ExplicitOption_TakesPrecedenceOverDI()
+    {
+        var explicitStore = new TestSseEventStreamStore();
+
+        Builder.Services.AddDistributedMemoryCache();
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport(options => options.EventStreamStore = explicitStore)
+            .WithDistributedCacheEventStreamStore();
+
+        using var app = Builder.Build();
+
+        var options = app.Services.GetRequiredService<IOptions<HttpServerTransportOptions>>().Value;
+        Assert.Same(explicitStore, options.EventStreamStore);
+    }
+
+    [Fact]
+    public void EventStreamStore_RemainsNull_WhenNothingIsRegistered()
+    {
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport();
+
+        using var app = Builder.Build();
+
+        var options = app.Services.GetRequiredService<IOptions<HttpServerTransportOptions>>().Value;
+        Assert.Null(options.EventStreamStore);
+    }
+
+    [Fact]
+    public void EventStreamStore_CanBeOverriddenToNull_AfterDIRegistration()
+    {
+        Builder.Services.AddDistributedMemoryCache();
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport()
+            .WithDistributedCacheEventStreamStore();
+
+        Builder.Services.Configure<HttpServerTransportOptions>(options => options.EventStreamStore = null);
+
+        using var app = Builder.Build();
+
+        var options = app.Services.GetRequiredService<IOptions<HttpServerTransportOptions>>().Value;
+        Assert.Null(options.EventStreamStore);
+    }
+
+    [Fact]
+    public void SessionMigrationHandler_IsPopulatedFromDI_ViaPostConfigure()
+    {
+        var handler = new StubSessionMigrationHandler();
+
+        Builder.Services.AddSingleton<ISessionMigrationHandler>(handler);
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport();
+
+        using var app = Builder.Build();
+
+        var options = app.Services.GetRequiredService<IOptions<HttpServerTransportOptions>>().Value;
+        Assert.Same(handler, options.SessionMigrationHandler);
+    }
+
+    [Fact]
+    public void SessionMigrationHandler_ExplicitOption_TakesPrecedenceOverDI()
+    {
+        var diHandler = new StubSessionMigrationHandler();
+        var explicitHandler = new StubSessionMigrationHandler();
+
+        Builder.Services.AddSingleton<ISessionMigrationHandler>(diHandler);
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport(options => options.SessionMigrationHandler = explicitHandler);
+
+        using var app = Builder.Build();
+
+        var options = app.Services.GetRequiredService<IOptions<HttpServerTransportOptions>>().Value;
+        Assert.Same(explicitHandler, options.SessionMigrationHandler);
+    }
+
+    [Fact]
+    public void SessionMigrationHandler_RemainsNull_WhenNothingIsRegistered()
+    {
+        Builder.Services
+            .AddMcpServer()
+            .WithHttpTransport();
+
+        using var app = Builder.Build();
+
+        var options = app.Services.GetRequiredService<IOptions<HttpServerTransportOptions>>().Value;
+        Assert.Null(options.SessionMigrationHandler);
+    }
+
+    private sealed class StubSessionMigrationHandler : ISessionMigrationHandler
+    {
+        public ValueTask<InitializeRequestParams?> AllowSessionMigrationAsync(
+            HttpContext context, string sessionId, CancellationToken cancellationToken = default)
+            => new((InitializeRequestParams?)null);
+
+        public ValueTask OnSessionInitializedAsync(
+            HttpContext context, string sessionId, InitializeRequestParams initializeParams, CancellationToken cancellationToken = default)
+            => default;
+    }
+}