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

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 = null) : IConfigureOptions<DistributedCacheEventStreamStoreOptions>
+{
+    public void Configure(DistributedCacheEventStreamStoreOptions options)
+    {
+        options.Cache ??= cache;
+    }
+}

src/ModelContextProtocol.AspNetCore/DistributedCacheEventStreamStoreOptionsValidator.cs

@@ -0,0 +1,25 @@
+using Microsoft.Extensions.Caching.Distributed;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Options;
+using ModelContextProtocol.Server;
+
+namespace ModelContextProtocol.AspNetCore;
+
+/// <summary>
+/// Validates that <see cref="DistributedCacheEventStreamStoreOptions.Cache"/> is set.
+/// </summary>
+internal sealed class DistributedCacheEventStreamStoreOptionsValidator : IValidateOptions<DistributedCacheEventStreamStoreOptions>
+{
+    public ValidateOptionsResult Validate(string? name, DistributedCacheEventStreamStoreOptions options)
+    {
+        if (options.Cache is null)
+        {
+            return ValidateOptionsResult.Fail(
+                $"The '{nameof(DistributedCacheEventStreamStoreOptions)}.{nameof(DistributedCacheEventStreamStoreOptions.Cache)}' property must be set. " +
+                $"Register an {nameof(IDistributedCache)} in DI or set the {nameof(DistributedCacheEventStreamStoreOptions.Cache)} property " +
+                $"in the '{nameof(HttpMcpServerBuilderExtensions.WithDistributedCacheEventStreamStore)}' configure callback.");
+        }
+
+        return ValidateOptionsResult.Success;
+    }
+}

src/ModelContextProtocol.AspNetCore/HttpMcpServerBuilderExtensions.cs

@@ -1,4 +1,5 @@
 using Microsoft.AspNetCore.Authorization;
+using Microsoft.Extensions.Caching.Distributed;
 using Microsoft.Extensions.DependencyInjection.Extensions;
 using Microsoft.Extensions.Options;
 using ModelContextProtocol.AspNetCore;
@@ -33,6 +34,7 @@ public static class HttpMcpServerBuilderExtensions
         builder.Services.AddHostedService<IdleTrackingBackgroundService>();
 
         builder.Services.TryAddEnumerable(ServiceDescriptor.Transient<IPostConfigureOptions<McpServerOptions>, AuthorizationFilterSetup>());
+        builder.Services.TryAddEnumerable(ServiceDescriptor.Transient<IConfigureOptions<HttpServerTransportOptions>, HttpServerTransportOptionsSetup>());
 
         if (configureOptions is not null)
         {
@@ -64,4 +66,37 @@ public static class HttpMcpServerBuilderExtensions
 
         return builder;
     }
+
+    /// <summary>
+    /// 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>
+    /// <returns>The builder provided in <paramref name="builder"/>.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
+    /// <remarks>
+    /// <para>
+    /// An <see cref="IDistributedCache"/> implementation must be registered in the service collection before calling this method.
+    /// 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,
+    /// 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)
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+
+        builder.Services.TryAddEnumerable(ServiceDescriptor.Singleton<IConfigureOptions<DistributedCacheEventStreamStoreOptions>, DistributedCacheEventStreamStoreOptionsSetup>());
+        builder.Services.TryAddEnumerable(ServiceDescriptor.Singleton<IValidateOptions<DistributedCacheEventStreamStoreOptions>, DistributedCacheEventStreamStoreOptionsValidator>());
+        builder.Services.AddSingleton<ISseEventStreamStore, DistributedCacheEventStreamStore>();
+
+        if (configureOptions is not null)
+        {
+            builder.Services.Configure(configureOptions);
+        }
+
+        return builder;
+    }
 }

src/ModelContextProtocol.AspNetCore/HttpServerTransportOptions.cs

@@ -55,9 +55,29 @@ public class HttpServerTransportOptions
     /// <item><description>Replay missed events when a client reconnects with a Last-Event-ID header</description></item>
     /// <item><description>Send priming events to establish resumability before any actual messages</description></item>
     /// </list>
+    /// <para>
+    /// This can be set directly, or an <see cref="ISseEventStreamStore"/> can be registered in DI.
+    /// If this property is not set, the server will attempt to resolve an <see cref="ISseEventStreamStore"/> from DI.
+    /// </para>
     /// </remarks>
     public ISseEventStreamStore? EventStreamStore { get; set; }
 
+    /// <summary>
+    /// Gets or sets the session migration handler for cross-instance session migration.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// When configured, the server will support session migration between instances.
+    /// If a request arrives with a session ID that is not found locally, the handler
+    /// is consulted to determine if the session can be migrated from another instance.
+    /// </para>
+    /// <para>
+    /// This can be set directly, or an <see cref="ISessionMigrationHandler"/> can be registered in DI.
+    /// If this property is not set, the server will attempt to resolve an <see cref="ISessionMigrationHandler"/> from DI.
+    /// </para>
+    /// </remarks>
+    public ISessionMigrationHandler? SessionMigrationHandler { get; set; }
+
     /// <summary>
     /// Gets or sets a value that indicates whether the server uses a single execution context for the entire session.
     /// </summary>

src/ModelContextProtocol.AspNetCore/HttpServerTransportOptionsSetup.cs

@@ -0,0 +1,18 @@
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Options;
+using ModelContextProtocol.Server;
+
+namespace ModelContextProtocol.AspNetCore;
+
+/// <summary>
+/// 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) : IConfigureOptions<HttpServerTransportOptions>
+{
+    public void Configure(HttpServerTransportOptions options)
+    {
+        options.EventStreamStore ??= serviceProvider.GetService<ISseEventStreamStore>();
+        options.SessionMigrationHandler ??= serviceProvider.GetService<ISessionMigrationHandler>();
+    }
+}

src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs

@@ -21,8 +21,7 @@ internal sealed class StreamableHttpHandler(
     StatefulSessionManager sessionManager,
     IHostApplicationLifetime hostApplicationLifetime,
     IServiceProvider applicationServices,
-    ILoggerFactory loggerFactory,
+    ILoggerFactory loggerFactory)
-    ISessionMigrationHandler? sessionMigrationHandler = null)
 {
     private const string McpSessionIdHeaderName = "Mcp-Session-Id";
     private const string McpProtocolVersionHeaderName = "MCP-Protocol-Version";
@@ -255,7 +254,7 @@ internal sealed class StreamableHttpHandler(
 
     private async ValueTask<StreamableHttpSession?> TryMigrateSessionAsync(HttpContext context, string sessionId)
     {
-        if (sessionMigrationHandler is not { } handler)
+        if (HttpServerTransportOptions.SessionMigrationHandler is not { } handler)
         {
             return null;
         }
@@ -336,7 +335,7 @@ internal sealed class StreamableHttpHandler(
                 SessionId = sessionId,
                 FlowExecutionContextFromRequests = !HttpServerTransportOptions.PerSessionExecutionContext,
                 EventStreamStore = HttpServerTransportOptions.EventStreamStore,
-                OnSessionInitialized = sessionMigrationHandler is { } handler
+                OnSessionInitialized = HttpServerTransportOptions.SessionMigrationHandler is { } handler
                     ? (initParams, ct) => handler.OnSessionInitializedAsync(context, sessionId, initParams, ct)
                     : null,
             };

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">Configuration options for the store, including the <see cref="IDistributedCache"/> to use.</param>
-    /// <param name="options">Optional configuration options for the store.</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);
+        Throw.IfNull(options);
-        _cache = cache;
+
-        _options = options ?? new();
+        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>

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;
+    }
+}

tests/ModelContextProtocol.ConformanceServer/Program.cs

@@ -1,9 +1,6 @@
 using ConformanceServer.Prompts;
 using ConformanceServer.Resources;
 using ConformanceServer.Tools;
-using Microsoft.Extensions.Caching.Distributed;
-using Microsoft.Extensions.Caching.Memory;
-using Microsoft.Extensions.Options;
 using ModelContextProtocol.Protocol;
 using ModelContextProtocol.Server;
 using System.Collections.Concurrent;
@@ -28,14 +25,11 @@ public class Program
         // because .NET does not have a built-in concurrent HashSet
         ConcurrentDictionary<string, ConcurrentDictionary<string, byte>> subscriptions = new();
 
+        builder.Services.AddDistributedMemoryCache();
         builder.Services
             .AddMcpServer()
-            .WithHttpTransport(options =>
+            .WithHttpTransport()
-            {
+            .WithDistributedCacheEventStreamStore()
-                // Enable resumability for SSE polling conformance test
-                options.EventStreamStore = new DistributedCacheEventStreamStore(
-                    new MemoryDistributedCache(Options.Create(new MemoryDistributedCacheOptions())));
-            })
             .WithTools<ConformanceTools>()
             .WithTools([ConformanceTools.CreateJsonSchema202012Tool()])
             .WithRequestFilters(filters => filters.AddCallToolFilter(next => async (request, cancellationToken) =>

tests/ModelContextProtocol.Tests/Configuration/McpServerOptionsSetupTests.cs

@@ -283,4 +283,57 @@ public class McpServerOptionsSetupTests
         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
 }

tests/ModelContextProtocol.Tests/Server/DistributedCacheEventStreamStoreTests.cs

@@ -21,10 +21,25 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
         return new MemoryDistributedCache(options);
     }
 
-    [Fact]
+    private static DistributedCacheEventStreamStore CreateStore(IDistributedCache? cache = null, DistributedCacheEventStreamStoreOptions? storeOptions = null)
-    public void Constructor_ThrowsArgumentNullException_WhenCacheIsNull()
     {
-        Assert.Throws<ArgumentNullException>("cache", () => new DistributedCacheEventStreamStore(null!));
+        storeOptions ??= new();
+        storeOptions.Cache ??= cache ?? CreateMemoryCache();
+        return new DistributedCacheEventStreamStore(Options.Create(storeOptions));
+    }
+
+    [Fact]
+    public void Constructor_ThrowsArgumentNullException_WhenOptionsIsNull()
+    {
+        Assert.Throws<ArgumentNullException>("options", () => new DistributedCacheEventStreamStore(null!));
+    }
+
+    [Fact]
+    public void Constructor_ThrowsInvalidOperationException_WhenCacheIsNull()
+    {
+        var options = Options.Create(new DistributedCacheEventStreamStoreOptions());
+        var ex = Assert.Throws<InvalidOperationException>(() => new DistributedCacheEventStreamStore(options));
+        Assert.StartsWith($"The '{nameof(DistributedCacheEventStreamStoreOptions)}.{nameof(DistributedCacheEventStreamStoreOptions.Cache)}'", ex.Message);
     }
 
     [Fact]
@@ -32,7 +47,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
 
         // Act & Assert
         await Assert.ThrowsAsync<ArgumentNullException>("options",
@@ -44,7 +59,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -67,7 +82,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -90,7 +105,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -113,7 +128,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -135,7 +150,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -187,7 +202,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -224,7 +239,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -254,7 +269,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
         {
             EventSlidingExpiration = TimeSpan.FromMinutes(15)
         };
-        var store = new DistributedCacheEventStreamStore(mockCache, customOptions);
+        var store = CreateStore(mockCache, customOptions);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -282,7 +297,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
         {
             EventAbsoluteExpiration = TimeSpan.FromHours(3)
         };
-        var store = new DistributedCacheEventStreamStore(mockCache, customOptions);
+        var store = CreateStore(mockCache, customOptions);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -306,7 +321,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var mockCache = new TestDistributedCache();
-        var store = new DistributedCacheEventStreamStore(mockCache);
+        var store = CreateStore(mockCache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -328,7 +343,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var mockCache = new TestDistributedCache();
-        var store = new DistributedCacheEventStreamStore(mockCache);
+        var store = CreateStore(mockCache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -354,7 +369,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(10)
         };
-        var store = new DistributedCacheEventStreamStore(cache, customOptions);
+        var store = CreateStore(cache, customOptions);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -390,7 +405,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -425,7 +440,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -447,7 +462,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var mockCache = new TestDistributedCache();
-        var store = new DistributedCacheEventStreamStore(mockCache);
+        var store = CreateStore(mockCache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -469,7 +484,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
 
         // Act & Assert
         await Assert.ThrowsAsync<ArgumentNullException>("lastEventId",
@@ -481,7 +496,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
 
         // Act - Try various invalid event ID formats
         var result1 = await store.GetStreamReaderAsync("invalid-format", CancellationToken);
@@ -499,7 +514,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
 
         // Create a valid-looking event ID for a stream that doesn't exist
         var fakeEventId = DistributedCacheEventIdFormatter.Format("nonexistent-session", "nonexistent-stream", 1);
@@ -516,7 +531,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "my-session",
@@ -542,7 +557,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -579,7 +594,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -610,7 +625,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -648,7 +663,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -681,7 +696,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -717,7 +732,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -751,7 +766,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -780,7 +795,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -812,7 +827,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache, new DistributedCacheEventStreamStoreOptions
+        var store = CreateStore(cache, new DistributedCacheEventStreamStoreOptions
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(50)
         });
@@ -869,7 +884,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache, new DistributedCacheEventStreamStoreOptions
+        var store = CreateStore(cache, new DistributedCacheEventStreamStoreOptions
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(50)
         });
@@ -927,7 +942,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache, new DistributedCacheEventStreamStoreOptions
+        var store = CreateStore(cache, new DistributedCacheEventStreamStoreOptions
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(50)
         });
@@ -965,7 +980,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache, new DistributedCacheEventStreamStoreOptions
+        var store = CreateStore(cache, new DistributedCacheEventStreamStoreOptions
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(50)
         });
@@ -1029,7 +1044,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache, new DistributedCacheEventStreamStoreOptions
+        var store = CreateStore(cache, new DistributedCacheEventStreamStoreOptions
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(50)
         });
@@ -1074,7 +1089,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange - Start in default mode, write some events, switch to polling, reader should return remaining events
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache, new DistributedCacheEventStreamStoreOptions
+        var store = CreateStore(cache, new DistributedCacheEventStreamStoreOptions
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(50)
         });
@@ -1120,7 +1135,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
 
         // Create two streams with different session/stream IDs
         var writer1 = await store.CreateStreamAsync(new SseEventStreamOptions
@@ -1178,7 +1193,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
 
         // Create two streams with same session but different stream IDs
         var writer1 = await store.CreateStreamAsync(new SseEventStreamOptions
@@ -1231,7 +1246,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange
         var cache = CreateMemoryCache();
-        var store = new DistributedCacheEventStreamStore(cache);
+        var store = CreateStore(cache);
 
         var writer1 = await store.CreateStreamAsync(new SseEventStreamOptions
         {
@@ -1267,7 +1282,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
         {
             EventSlidingExpiration = TimeSpan.FromMinutes(30)
         };
-        var store = new DistributedCacheEventStreamStore(mockCache, customOptions);
+        var store = CreateStore(mockCache, customOptions);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -1295,7 +1310,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
         {
             EventAbsoluteExpiration = TimeSpan.FromHours(6)
         };
-        var store = new DistributedCacheEventStreamStore(mockCache, customOptions);
+        var store = CreateStore(mockCache, customOptions);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -1325,7 +1340,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
             MetadataSlidingExpiration = TimeSpan.FromMinutes(45),
             MetadataAbsoluteExpiration = TimeSpan.FromHours(12)
         };
-        var store = new DistributedCacheEventStreamStore(mockCache, customOptions);
+        var store = CreateStore(mockCache, customOptions);
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
         {
             SessionId = "session-1",
@@ -1366,7 +1381,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(10) // Fast polling to detect the bug quickly
         };
-        var store = new DistributedCacheEventStreamStore(trackingCache, customOptions);
+        var store = CreateStore(trackingCache, customOptions);
 
         // Create a stream and write an event
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
@@ -1405,7 +1420,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
     {
         // Arrange - Use a cache that allows us to simulate event expiration
         var trackingCache = new TestDistributedCache();
-        var store = new DistributedCacheEventStreamStore(trackingCache);
+        var store = CreateStore(trackingCache);
 
         // Create a stream and write multiple events
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions
@@ -1450,7 +1465,7 @@ public class DistributedCacheEventStreamStoreTests(ITestOutputHelper testOutputH
         {
             StreamReaderPollingInterval = TimeSpan.FromMilliseconds(10)
         };
-        var store = new DistributedCacheEventStreamStore(trackingCache, customOptions);
+        var store = CreateStore(trackingCache, customOptions);
 
         // Create a stream in POLLING mode - this allows the reader to exit after reading available events
         var writer = await store.CreateStreamAsync(new SseEventStreamOptions