diff --git a/src/HotChocolate/Adapters/src/Adapters.Mcp/Extensions/RequestExecutorBuilderExtensions.cs b/src/HotChocolate/Adapters/src/Adapters.Mcp/Extensions/RequestExecutorBuilderExtensions.cs index 2936171082a..d394864921c 100644 --- a/src/HotChocolate/Adapters/src/Adapters.Mcp/Extensions/RequestExecutorBuilderExtensions.cs +++ b/src/HotChocolate/Adapters/src/Adapters.Mcp/Extensions/RequestExecutorBuilderExtensions.cs @@ -1,6 +1,4 @@ -#if !NET9_0_OR_GREATER using System.Diagnostics.CodeAnalysis; -#endif using HotChocolate.Adapters.Mcp.Directives; using HotChocolate.Adapters.Mcp.Storage; using HotChocolate.Execution.Configuration; @@ -29,17 +27,24 @@ public static IRequestExecutorBuilder AddMcp( builder.AddDirectiveType(); - builder.ConfigureOnRequestExecutorCreatedAsync( - async (executor, cancellationToken) => - { - var schema = executor.Schema; - var storageObserver = schema.Services.GetRequiredService(); - await storageObserver.StartAsync(cancellationToken); - }); - return builder; } + public static IRequestExecutorBuilder AddMcpStorage< + [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors)] T>( + this IRequestExecutorBuilder builder, + Func factory, + Func? skipIf = null) + where T : class, IMcpStorage + { + ArgumentNullException.ThrowIfNull(builder); + ArgumentNullException.ThrowIfNull(factory); + + builder.ConfigureSchemaServices(s => s.AddSingleton(factory)); + + return builder.AddMcpStorageWarmupTask(skipIf); + } + public static IRequestExecutorBuilder AddMcpStorage( this IRequestExecutorBuilder builder, IMcpStorage storage) @@ -49,6 +54,20 @@ public static IRequestExecutorBuilder AddMcpStorage( builder.ConfigureSchemaServices(s => s.AddSingleton(storage)); + return builder.AddMcpStorageWarmupTask(); + } + + private static IRequestExecutorBuilder AddMcpStorageWarmupTask( + this IRequestExecutorBuilder builder, + Func? skipIf = null) + { + builder.AddWarmupTask(async (executor, cancellationToken) => + { + var schema = executor.Schema; + var storageObserver = schema.Services.GetRequiredService(); + await storageObserver.StartAsync(cancellationToken); + }, skipIf); + return builder; } } diff --git a/src/HotChocolate/Adapters/src/Adapters.Mcp/HotChocolate.Adapters.Mcp.csproj b/src/HotChocolate/Adapters/src/Adapters.Mcp/HotChocolate.Adapters.Mcp.csproj index 6d167c021a3..4a7025ea3e0 100644 --- a/src/HotChocolate/Adapters/src/Adapters.Mcp/HotChocolate.Adapters.Mcp.csproj +++ b/src/HotChocolate/Adapters/src/Adapters.Mcp/HotChocolate.Adapters.Mcp.csproj @@ -6,6 +6,7 @@ + diff --git a/src/HotChocolate/Adapters/src/Fusion.Adapters.Mcp/Extensions/FusionGatewayBuilderExtensions.cs b/src/HotChocolate/Adapters/src/Fusion.Adapters.Mcp/Extensions/FusionGatewayBuilderExtensions.cs index 062ed29edb0..cbb8efb64d3 100644 --- a/src/HotChocolate/Adapters/src/Fusion.Adapters.Mcp/Extensions/FusionGatewayBuilderExtensions.cs +++ b/src/HotChocolate/Adapters/src/Fusion.Adapters.Mcp/Extensions/FusionGatewayBuilderExtensions.cs @@ -1,6 +1,4 @@ -#if !NET9_0_OR_GREATER using System.Diagnostics.CodeAnalysis; -#endif using HotChocolate.Adapters.Mcp.Storage; using HotChocolate.Fusion.Configuration; using Microsoft.Extensions.DependencyInjection; @@ -26,16 +24,24 @@ public static IFusionGatewayBuilder AddMcp( builder.ConfigureSchemaServices( (_, services) => services.AddMcpSchemaServices(configureServerOptions, configureServer)); - builder.AddWarmupTask(async (executor, cancellationToken) => - { - var schema = executor.Schema; - var storageObserver = schema.Services.GetRequiredService(); - await storageObserver.StartAsync(cancellationToken); - }); - return builder; } + public static IFusionGatewayBuilder AddMcpStorage< + [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors)] T>( + this IFusionGatewayBuilder builder, + Func factory, + Func? skipIf = null) + where T : class, IMcpStorage + { + ArgumentNullException.ThrowIfNull(builder); + ArgumentNullException.ThrowIfNull(factory); + + builder.ConfigureSchemaServices((_, s) => s.AddSingleton(factory)); + + return builder.AddMcpStorageWarmupTask(skipIf); + } + public static IFusionGatewayBuilder AddMcpStorage( this IFusionGatewayBuilder builder, IMcpStorage storage) @@ -45,6 +51,20 @@ public static IFusionGatewayBuilder AddMcpStorage( builder.ConfigureSchemaServices((_, s) => s.AddSingleton(storage)); + return builder.AddMcpStorageWarmupTask(); + } + + private static IFusionGatewayBuilder AddMcpStorageWarmupTask( + this IFusionGatewayBuilder builder, + Func? skipIf = null) + { + builder.AddWarmupTask(async (executor, cancellationToken) => + { + var schema = executor.Schema; + var storageObserver = schema.Services.GetRequiredService(); + await storageObserver.StartAsync(cancellationToken); + }, skipIf); + return builder; } }