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,
-    ISessionMigrationHandler? sessionMigrationHandler = null)
+    ILoggerFactory loggerFactory)
 {
     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,
             };