diff --git a/src/Smith/McpExtensions.cs b/src/Smith/McpExtensions.cs index 5f35853..15ab6b1 100644 --- a/src/Smith/McpExtensions.cs +++ b/src/Smith/McpExtensions.cs @@ -1,5 +1,6 @@ using System.Text.Json; using Microsoft.Extensions.DependencyInjection; +using ModelContextProtocol.Protocol; using ModelContextProtocol.Server; namespace Smith; @@ -58,5 +59,33 @@ public IMcpServerBuilder WithTool(string name, string title, string description, return builder; } + + /// + /// Run the specified initializer function before listing tools (after client initialized). + /// + /// An initializer that has access to the server instance. + public IMcpServerBuilder WithInitializer(Func initializer) + { + // NOTE: we wrap existing list handler so we avoid breaking existing tools. + builder.Services.Configure(handlers => handlers.ListToolsHandler = + new InitializerListTools(initializer, handlers.ListToolsHandler).Execute); + + return builder; + } + } + + class InitializerListTools( + Func initializer, + Func, CancellationToken, ValueTask>? handler) + { + public async ValueTask Execute(RequestContext request, CancellationToken token) + { + await initializer(request.Server, token); + + if (handler != null) + return await handler(request, token); + + return new ListToolsResult(); + } } }