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

src/ModelContextProtocol/McpServerOptionsSetup.cs

@@ -9,10 +9,12 @@ namespace ModelContextProtocol;
 /// <param name="serverTools">The individually registered tools.</param>
 /// <param name="serverPrompts">The individually registered prompts.</param>
 /// <param name="serverResources">The individually registered resources.</param>
+/// <param name="taskStore">The optional task store registered in DI.</param>
 internal sealed class McpServerOptionsSetup(
     IEnumerable<McpServerTool> serverTools,
     IEnumerable<McpServerPrompt> serverPrompts,
-    IEnumerable<McpServerResource> serverResources) : IConfigureOptions<McpServerOptions>
+    IEnumerable<McpServerResource> serverResources,
+    IMcpTaskStore? taskStore = null) : IConfigureOptions<McpServerOptions>
 {
     /// <summary>
     /// Configures the given McpServerOptions instance by setting server information
@@ -23,6 +25,8 @@ internal sealed class McpServerOptionsSetup(
     {
         Throw.IfNull(options);
 
+        options.TaskStore ??= taskStore;
+
         // 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/McpServerServiceCollectionExtensions.cs

@@ -26,17 +26,6 @@ public static class McpServerServiceCollectionExtensions
             services.Configure(configureOptions);
         }
 
-        // Register IMcpTaskStore from options if not already registered.
-        // This allows users to either:
-        // 1. Register IMcpTaskStore directly in DI (takes precedence)
-        // 2. Set options.TaskStore in the configuration callback (used as fallback)
-        // If neither is done, resolving IMcpTaskStore will throw.
-        services.TryAddSingleton<IMcpTaskStore>(sp =>
-        {
-            var options = sp.GetRequiredService<IOptions<McpServerOptions>>().Value;
-            return options.TaskStore ?? throw new InvalidOperationException("No IMcpTaskStore has been configured. Either register an IMcpTaskStore in the service collection or set McpServerOptions.TaskStore when configuring the MCP server.");
-        });
-
         return new DefaultMcpServerBuilder(services);
     }
 }

src/ModelContextProtocol/Server/DistributedCacheEventStreamStore.cs

@@ -1,6 +1,7 @@
 using Microsoft.Extensions.Caching.Distributed;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
+using Microsoft.Extensions.Options;
 using ModelContextProtocol.Protocol;
 using System.Net.ServerSentEvents;
 using System.Runtime.CompilerServices;
@@ -31,14 +32,16 @@ public sealed partial class DistributedCacheEventStreamStore : ISseEventStreamSt
     /// <summary>
     /// Initializes a new instance of the <see cref="DistributedCacheEventStreamStore"/> class.
     /// </summary>
-    /// <param name="cache">The distributed cache to use for storage.</param>
-    /// <param name="options">Optional configuration options for the store.</param>
+    /// <param name="options">Configuration options for the store, including the <see cref="IDistributedCache"/> to use.</param>
     /// <param name="logger">Optional logger for diagnostic output.</param>
-    public DistributedCacheEventStreamStore(IDistributedCache cache, DistributedCacheEventStreamStoreOptions? options = null, ILogger<DistributedCacheEventStreamStore>? logger = null)
+    public DistributedCacheEventStreamStore(IOptions<DistributedCacheEventStreamStoreOptions> options, ILogger<DistributedCacheEventStreamStore>? logger = null)
     {
-        Throw.IfNull(cache);
-        _cache = cache;
-        _options = options ?? new();
+        Throw.IfNull(options);
+
+        var optionsValue = options.Value;
+        _cache = optionsValue.Cache ?? throw new InvalidOperationException(
+            $"The '{nameof(DistributedCacheEventStreamStoreOptions)}.{nameof(DistributedCacheEventStreamStoreOptions.Cache)}' property must be set.");
+        _options = optionsValue;
         _logger = logger ?? NullLogger<DistributedCacheEventStreamStore>.Instance;
     }
 

src/ModelContextProtocol/Server/DistributedCacheEventStreamStoreOptions.cs

@@ -1,3 +1,5 @@
+using Microsoft.Extensions.Caching.Distributed;
+
 namespace ModelContextProtocol.Server;
 
 /// <summary>
@@ -5,6 +7,16 @@ namespace ModelContextProtocol.Server;
 /// </summary>
 public sealed class DistributedCacheEventStreamStoreOptions
 {
+    /// <summary>
+    /// Gets or sets the <see cref="IDistributedCache"/> to use for event storage.
+    /// </summary>
+    /// <remarks>
+    /// When using dependency injection with <c>WithDistributedCacheEventStreamStore()</c>, this is
+    /// automatically populated from the <see cref="IDistributedCache"/> registered in DI.
+    /// Set this property explicitly to use a specific cache instance.
+    /// </remarks>
+    public IDistributedCache? Cache { get; set; }
+
     /// <summary>
     /// Gets or sets the sliding expiration for individual events in the cache.
     /// </summary>