using ColaFlow.Modules.Mcp.Application.Handlers; using ColaFlow.Modules.Mcp.Application.Resources; using ColaFlow.Modules.Mcp.Application.Services; using ColaFlow.Modules.Mcp.Contracts.Resources; using ColaFlow.Modules.Mcp.Domain.Repositories; using ColaFlow.Modules.Mcp.Infrastructure.Authentication; using ColaFlow.Modules.Mcp.Infrastructure.BackgroundServices; using ColaFlow.Modules.Mcp.Infrastructure.Middleware; using ColaFlow.Modules.Mcp.Infrastructure.Persistence; using ColaFlow.Modules.Mcp.Infrastructure.Persistence.Repositories; using ColaFlow.Modules.Mcp.Infrastructure.Services; using Microsoft.AspNetCore.Authentication; using Microsoft.AspNetCore.Builder; using Microsoft.EntityFrameworkCore; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Logging; namespace ColaFlow.Modules.Mcp.Infrastructure.Extensions; /// /// Extension methods for registering MCP services /// Enhanced with auto-discovery and dynamic registration /// public static class McpServiceExtensions { /// /// Registers MCP module services with auto-discovery /// public static IServiceCollection AddMcpModule(this IServiceCollection services, IConfiguration configuration) { // Register DbContext services.AddDbContext(options => { var connectionString = configuration.GetConnectionString("DefaultConnection"); options.UseNpgsql(connectionString); }); // Register repositories services.AddScoped(); services.AddScoped(); services.AddScoped(); // Register domain services services.AddScoped(); services.AddScoped(); services.AddScoped(); services.AddScoped(); // Register notification service (SignalR real-time notifications) services.AddScoped(); // Register background services services.AddHostedService(); // Register resource registry (Singleton - shared across all requests) services.AddSingleton(); // Register Resource Discovery Service (Singleton - used at startup) services.AddSingleton(); // Register MCP Resources (Scoped - manual registration for now, auto-discovery at startup) services.AddScoped(); services.AddScoped(); services.AddScoped(); services.AddScoped(); services.AddScoped(); services.AddScoped(); // Register protocol handler services.AddScoped(); // Register method handlers services.AddScoped(); services.AddScoped(); services.AddScoped(); services.AddScoped(); // NEW: Health check handler services.AddScoped(); services.AddScoped(); // Note: MediatR handlers for MCP are registered in ModuleExtensions.cs AddProjectManagementModule() // They are included via RegisterServicesFromAssembly for the MCP Application assembly return services; } /// /// Adds MCP API Key authentication scheme to the authentication builder. /// This enables the /mcp-sdk endpoint to use API Key authentication. /// public static AuthenticationBuilder AddMcpApiKeyAuthentication( this AuthenticationBuilder builder, Action? configureOptions = null) { return builder.AddScheme( McpApiKeyAuthenticationOptions.DefaultScheme, configureOptions ?? (_ => { })); } /// /// Adds MCP middleware to the application pipeline /// IMPORTANT: Middleware order matters - must be in this sequence: /// 1. Correlation ID (for request tracking) /// 2. Exception Handler (catches all errors) /// 3. Logging (logs requests/responses) /// 4. Authentication (validates API key) /// 5. MCP Protocol Handler (processes MCP requests) /// public static IApplicationBuilder UseMcpMiddleware(this IApplicationBuilder app) { // Initialize resource registry (register all resources with auto-discovery) InitializeResourceRegistry(app); // 1. Correlation ID middleware (FIRST - needed for all subsequent logging) app.UseMiddleware(); // 2. Exception handler (SECOND - catches all errors from downstream middleware) app.UseMiddleware(); // 3. Logging middleware (THIRD - logs request/response with correlation ID) app.UseMiddleware(); // 4. Authentication middleware (FOURTH - validates API key) app.UseMiddleware(); // 5. MCP protocol handler (LAST - processes the actual MCP request) app.UseMiddleware(); return app; } /// /// Initialize resource registry by registering all resources /// Enhanced with auto-discovery using Assembly scanning /// This is called once at startup /// private static void InitializeResourceRegistry(IApplicationBuilder app) { using var scope = app.ApplicationServices.CreateScope(); var registry = scope.ServiceProvider.GetRequiredService(); var discoveryService = scope.ServiceProvider.GetRequiredService(); var loggerFactory = scope.ServiceProvider.GetRequiredService(); var logger = loggerFactory.CreateLogger("McpResourceRegistry"); logger.LogInformation("Initializing MCP Resource Registry with auto-discovery..."); // Auto-discover and instantiate all resources var resources = discoveryService.DiscoverAndInstantiateResources(scope.ServiceProvider); // Register all discovered resources foreach (var resource in resources) { registry.RegisterResource(resource); } var categories = registry.GetCategories(); logger.LogInformation("MCP Resource Registry initialized: {ResourceCount} resources in {CategoryCount} categories: {Categories}", resources.Count, categories.Count, string.Join(", ", categories)); } }