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