PR feedback

Author: MackinnonBuck

src/ModelContextProtocol.AspNetCore/DistributedCacheEventStreamStoreOptionsSetup.cs

@@ -0,0 +1,17 @@
+using Microsoft.Extensions.Caching.Distributed;
+using Microsoft.Extensions.Options;
+using ModelContextProtocol.Server;
+
+namespace ModelContextProtocol.AspNetCore;
+
+/// <summary>
+/// Configures <see cref="DistributedCacheEventStreamStoreOptions"/> by resolving
+/// the <see cref="IDistributedCache"/> from DI when not explicitly set.
+/// </summary>
+internal sealed class DistributedCacheEventStreamStoreOptionsSetup(IDistributedCache cache) : IConfigureOptions<DistributedCacheEventStreamStoreOptions>
+{
+    public void Configure(DistributedCacheEventStreamStoreOptions options)
+    {
+        options.Cache ??= cache;
+    }
+}

src/ModelContextProtocol.AspNetCore/HttpMcpServerBuilderExtensions.cs

@@ -1,7 +1,6 @@
 using Microsoft.AspNetCore.Authorization;
 using Microsoft.Extensions.Caching.Distributed;
 using Microsoft.Extensions.DependencyInjection.Extensions;
-using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using ModelContextProtocol.AspNetCore;
 using ModelContextProtocol.Server;
@@ -69,8 +68,7 @@ public static IMcpServerBuilder AddAuthorizationFilters(this IMcpServerBuilder b
     }
 
     /// <summary>
-    /// Registers a <see cref="DistributedCacheEventStreamStore"/> as the <see cref="ISseEventStreamStore"/> for SSE resumability,
-    /// using the <see cref="IDistributedCache"/> already registered in the service collection.
+    /// Registers a <see cref="DistributedCacheEventStreamStore"/> as the <see cref="ISseEventStreamStore"/> for SSE resumability.
     /// </summary>
     /// <param name="builder">The builder instance.</param>
     /// <param name="configureOptions">An optional action to configure <see cref="DistributedCacheEventStreamStoreOptions"/>.</param>
@@ -79,41 +77,19 @@ public static IMcpServerBuilder AddAuthorizationFilters(this IMcpServerBuilder b
     /// <remarks>
     /// <para>
     /// An <see cref="IDistributedCache"/> implementation must be registered in the service collection before calling this method.
-    /// For example, call <c>builder.Services.AddDistributedMemoryCache()</c> for an in-memory cache or
-    /// <c>builder.Services.AddStackExchangeRedisCache(...)</c> for Redis.
+    /// The registered cache is automatically assigned to <see cref="DistributedCacheEventStreamStoreOptions.Cache"/>.
     /// </para>
     /// <para>
     /// To use a specific <see cref="IDistributedCache"/> instance instead of the one registered in DI,
-    /// use the overload that accepts an <see cref="IDistributedCache"/> parameter.
+    /// set the <see cref="DistributedCacheEventStreamStoreOptions.Cache"/> property in the <paramref name="configureOptions"/> callback.
     /// </para>
     /// </remarks>
     public static IMcpServerBuilder WithDistributedCacheEventStreamStore(this IMcpServerBuilder builder, Action<DistributedCacheEventStreamStoreOptions>? configureOptions = null)
-        => builder.WithDistributedCacheEventStreamStore(cache: null, configureOptions);
-
-    /// <summary>
-    /// Registers a <see cref="DistributedCacheEventStreamStore"/> as the <see cref="ISseEventStreamStore"/> for SSE resumability,
-    /// using the specified <see cref="IDistributedCache"/>.
-    /// </summary>
-    /// <param name="builder">The builder instance.</param>
-    /// <param name="cache">The <see cref="IDistributedCache"/> to use for event storage, or <see langword="null"/> to resolve from DI.</param>
-    /// <param name="configureOptions">An optional action to configure <see cref="DistributedCacheEventStreamStoreOptions"/>.</param>
-    /// <returns>The builder provided in <paramref name="builder"/>.</returns>
-    /// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
-    /// <remarks>
-    /// Use this overload when you have multiple <see cref="IDistributedCache"/> implementations or want to
-    /// use a specific instance rather than the one registered in DI.
-    /// </remarks>
-    public static IMcpServerBuilder WithDistributedCacheEventStreamStore(this IMcpServerBuilder builder, IDistributedCache? cache, Action<DistributedCacheEventStreamStoreOptions>? configureOptions = null)
     {
         ArgumentNullException.ThrowIfNull(builder);
 
-        builder.Services.AddSingleton<ISseEventStreamStore>(sp =>
-        {
-            var resolvedCache = cache ?? sp.GetRequiredService<IDistributedCache>();
-            var options = sp.GetService<IOptions<DistributedCacheEventStreamStoreOptions>>()?.Value;
-            var logger = sp.GetService<ILogger<DistributedCacheEventStreamStore>>();
-            return new DistributedCacheEventStreamStore(resolvedCache, options, logger);
-        });
+        builder.Services.TryAddEnumerable(ServiceDescriptor.Singleton<IConfigureOptions<DistributedCacheEventStreamStoreOptions>, DistributedCacheEventStreamStoreOptionsSetup>());
+        builder.Services.AddSingleton<ISseEventStreamStore, DistributedCacheEventStreamStore>();
 
         if (configureOptions is not null)
         {

src/ModelContextProtocol/McpServerOptionsSetup.cs

@@ -1,4 +1,3 @@
-using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 using ModelContextProtocol.Server;
 
@@ -7,15 +6,15 @@ namespace ModelContextProtocol;
 /// <summary>
 /// 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>
+/// <param name="taskStore">The optional task store registered in DI.</param>
 internal sealed class McpServerOptionsSetup(
-    IServiceProvider serviceProvider,
     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
@@ -26,7 +25,7 @@ public void Configure(McpServerOptions options)
     {
         Throw.IfNull(options);
 
-        options.TaskStore ??= serviceProvider.GetService<IMcpTaskStore>();
+        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

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,18 @@ 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. " +
+            $"When using 'WithDistributedCacheEventStreamStore()', the cache will be automatically configured from DI. " +
+            $"When constructing manually, set the '{nameof(DistributedCacheEventStreamStoreOptions.Cache)}' property on the options object.");
+        _options = optionsValue;
         _logger = logger ?? NullLogger<DistributedCacheEventStreamStore>.Instance;
     }
 

src/ModelContextProtocol/Server/DistributedCacheEventStreamStoreOptions.cs

@@ -1,10 +1,22 @@
+using Microsoft.Extensions.Caching.Distributed;
+
 namespace ModelContextProtocol.Server;
 
 /// <summary>
 /// Configuration options for <see cref="DistributedCacheEventStreamStore"/>.
 /// </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>

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 @@ protected override ValueTask<ISseEventStreamStore> CreateEventStreamStoreAsync()
             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);
     }
 }