PR feedback

Author: MackinnonBuck

src/ModelContextProtocol.AspNetCore/HttpMcpServerBuilderExtensions.cs

@@ -35,7 +35,7 @@ public static IMcpServerBuilder WithHttpTransport(this IMcpServerBuilder builder
         builder.Services.AddHostedService<IdleTrackingBackgroundService>();
 
         builder.Services.TryAddEnumerable(ServiceDescriptor.Transient<IPostConfigureOptions<McpServerOptions>, AuthorizationFilterSetup>());
-        builder.Services.TryAddEnumerable(ServiceDescriptor.Transient<IPostConfigureOptions<HttpServerTransportOptions>, HttpServerTransportOptionsSetup>());
+        builder.Services.TryAddEnumerable(ServiceDescriptor.Transient<IConfigureOptions<HttpServerTransportOptions>, HttpServerTransportOptionsSetup>());
 
         if (configureOptions is not null)
         {
@@ -107,7 +107,7 @@ public static IMcpServerBuilder WithDistributedCacheEventStreamStore(this IMcpSe
     {
         ArgumentNullException.ThrowIfNull(builder);
 
-        builder.Services.TryAddSingleton<ISseEventStreamStore>(sp =>
+        builder.Services.AddSingleton<ISseEventStreamStore>(sp =>
         {
             var resolvedCache = cache ?? sp.GetRequiredService<IDistributedCache>();
             var options = sp.GetService<IOptions<DistributedCacheEventStreamStoreOptions>>()?.Value;

src/ModelContextProtocol.AspNetCore/HttpServerTransportOptionsSetup.cs

@@ -8,9 +8,9 @@ namespace ModelContextProtocol.AspNetCore;
 /// Post-configures <see cref="HttpServerTransportOptions"/> by resolving services from DI
 /// when they haven't been explicitly set on the options.
 /// </summary>
-internal sealed class HttpServerTransportOptionsSetup(IServiceProvider serviceProvider) : IPostConfigureOptions<HttpServerTransportOptions>
+internal sealed class HttpServerTransportOptionsSetup(IServiceProvider serviceProvider) : IConfigureOptions<HttpServerTransportOptions>
 {
-    public void PostConfigure(string? name, HttpServerTransportOptions options)
+    public void Configure(HttpServerTransportOptions options)
     {
         options.EventStreamStore ??= serviceProvider.GetService<ISseEventStreamStore>();
         options.SessionMigrationHandler ??= serviceProvider.GetService<ISessionMigrationHandler>();

src/ModelContextProtocol/McpServerOptionsPostSetup.cs

@@ -1,15 +1,18 @@
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 using ModelContextProtocol.Server;
 
 namespace ModelContextProtocol;
 
 /// <summary>
-/// Configures the McpServerOptions using addition services from DI.
+/// Configures the McpServerOptions using additional services from DI.
 /// </summary>
+/// <param name="serviceProvider">The service provider for resolving optional services.</param>
 /// <param name="serverTools">The individually registered tools.</param>
 /// <param name="serverPrompts">The individually registered prompts.</param>
 /// <param name="serverResources">The individually registered resources.</param>
 internal sealed class McpServerOptionsSetup(
+    IServiceProvider serviceProvider,
     IEnumerable<McpServerTool> serverTools,
     IEnumerable<McpServerPrompt> serverPrompts,
     IEnumerable<McpServerResource> serverResources) : IConfigureOptions<McpServerOptions>
@@ -23,6 +26,8 @@ public void Configure(McpServerOptions options)
     {
         Throw.IfNull(options);
 
+        options.TaskStore ??= serviceProvider.GetService<IMcpTaskStore>();
+
         // Collect all of the provided tools into a tools collection. If the options already has
         // a collection, add to it, otherwise create a new one. We want to maintain the identity
         // of an existing collection in case someone has provided their own derived type, wants

src/ModelContextProtocol/McpServerOptionsSetup.cs

@@ -21,7 +21,6 @@ public static IMcpServerBuilder AddMcpServer(this IServiceCollection services, A
     {
         services.AddOptions();
         services.TryAddEnumerable(ServiceDescriptor.Transient<IConfigureOptions<McpServerOptions>, McpServerOptionsSetup>());
-        services.TryAddEnumerable(ServiceDescriptor.Transient<IPostConfigureOptions<McpServerOptions>, McpServerOptionsPostSetup>());
         if (configureOptions is not null)
         {
             services.Configure(configureOptions);

src/ModelContextProtocol/McpServerServiceCollectionExtensions.cs

@@ -88,6 +88,23 @@ public void EventStreamStore_RemainsNull_WhenNothingIsRegistered()
         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()
     {

tests/ModelContextProtocol.AspNetCore.Tests/HttpMcpServerBuilderExtensionsTests.cs

@@ -283,4 +283,57 @@ public void Configure_WithCompleteHandler_CreatesCompletionsCapability()
         Assert.NotNull(options.Capabilities?.Completions);
     }
     #endregion
+
+    #region TaskStore Tests
+    [Fact]
+    public void TaskStore_IsPopulatedFromDI_WhenNotExplicitlySet()
+    {
+        var services = new ServiceCollection();
+        services.AddMcpServer();
+        services.AddSingleton<IMcpTaskStore, InMemoryMcpTaskStore>();
+
+        var options = services.BuildServiceProvider().GetRequiredService<IOptions<McpServerOptions>>().Value;
+
+        Assert.IsType<InMemoryMcpTaskStore>(options.TaskStore);
+    }
+
+    [Fact]
+    public void TaskStore_ExplicitOption_TakesPrecedenceOverDI()
+    {
+        var explicitStore = new InMemoryMcpTaskStore();
+
+        var services = new ServiceCollection();
+        services.AddMcpServer(options => options.TaskStore = explicitStore);
+        services.AddSingleton<IMcpTaskStore, InMemoryMcpTaskStore>();
+
+        var options = services.BuildServiceProvider().GetRequiredService<IOptions<McpServerOptions>>().Value;
+
+        Assert.Same(explicitStore, options.TaskStore);
+    }
+
+    [Fact]
+    public void TaskStore_RemainsNull_WhenNothingIsRegistered()
+    {
+        var services = new ServiceCollection();
+        services.AddMcpServer();
+
+        var options = services.BuildServiceProvider().GetRequiredService<IOptions<McpServerOptions>>().Value;
+
+        Assert.Null(options.TaskStore);
+    }
+
+    [Fact]
+    public void TaskStore_CanBeOverriddenToNull_AfterDIRegistration()
+    {
+        var services = new ServiceCollection();
+        services.AddMcpServer();
+        services.AddSingleton<IMcpTaskStore, InMemoryMcpTaskStore>();
+
+        services.Configure<McpServerOptions>(options => options.TaskStore = null);
+
+        var options = services.BuildServiceProvider().GetRequiredService<IOptions<McpServerOptions>>().Value;
+
+        Assert.Null(options.TaskStore);
+    }
+    #endregion
 }