ColaFlow/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Extensions/McpServiceExtensions.cs

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;

/// <summary>
/// Extension methods for registering MCP services
/// Enhanced with auto-discovery and dynamic registration
/// </summary>
public static class McpServiceExtensions
{
    /// <summary>
    /// Registers MCP module services with auto-discovery
    /// </summary>
    public static IServiceCollection AddMcpModule(this IServiceCollection services, IConfiguration configuration)
    {
        // Register DbContext
        services.AddDbContext<McpDbContext>(options =>
        {
            var connectionString = configuration.GetConnectionString("DefaultConnection");
            options.UseNpgsql(connectionString);
        });

        // Register repositories
        services.AddScoped<IMcpApiKeyRepository, McpApiKeyRepository>();
        services.AddScoped<IPendingChangeRepository, PendingChangeRepository>();
        services.AddScoped<ITaskLockRepository, TaskLockRepository>();

        // Register domain services
        services.AddScoped<ColaFlow.Modules.Mcp.Domain.Services.DiffPreviewService>();
        services.AddScoped<ColaFlow.Modules.Mcp.Domain.Services.TaskLockService>();
        services.AddScoped<IMcpApiKeyService, McpApiKeyService>();
        services.AddScoped<IPendingChangeService, PendingChangeService>();

        // Register notification service (SignalR real-time notifications)
        services.AddScoped<IMcpNotificationService, McpNotificationService>();

        // Register background services
        services.AddHostedService<PendingChangeExpirationBackgroundService>();

        // Register resource registry (Singleton - shared across all requests)
        services.AddSingleton<IMcpResourceRegistry, McpResourceRegistry>();

        // Register Resource Discovery Service (Singleton - used at startup)
        services.AddSingleton<IResourceDiscoveryService, ResourceDiscoveryService>();

        // Register MCP Resources (Scoped - manual registration for now, auto-discovery at startup)
        services.AddScoped<IMcpResource, ProjectsListResource>();
        services.AddScoped<IMcpResource, ProjectsGetResource>();
        services.AddScoped<IMcpResource, IssuesSearchResource>();
        services.AddScoped<IMcpResource, IssuesGetResource>();
        services.AddScoped<IMcpResource, SprintsCurrentResource>();
        services.AddScoped<IMcpResource, UsersListResource>();

        // Register protocol handler
        services.AddScoped<IMcpProtocolHandler, McpProtocolHandler>();

        // Register method handlers
        services.AddScoped<IMcpMethodHandler, InitializeMethodHandler>();
        services.AddScoped<IMcpMethodHandler, ResourcesListMethodHandler>();
        services.AddScoped<IMcpMethodHandler, ResourcesReadMethodHandler>();
        services.AddScoped<IMcpMethodHandler, ResourceHealthCheckHandler>(); // NEW: Health check handler
        services.AddScoped<IMcpMethodHandler, ToolsListMethodHandler>();
        services.AddScoped<IMcpMethodHandler, ToolsCallMethodHandler>();

        // Note: MediatR handlers for MCP are registered in ModuleExtensions.cs AddProjectManagementModule()
        // They are included via RegisterServicesFromAssembly for the MCP Application assembly

        return services;
    }

    /// <summary>
    /// Adds MCP API Key authentication scheme to the authentication builder.
    /// This enables the /mcp-sdk endpoint to use API Key authentication.
    /// </summary>
    public static AuthenticationBuilder AddMcpApiKeyAuthentication(
        this AuthenticationBuilder builder,
        Action<McpApiKeyAuthenticationOptions>? configureOptions = null)
    {
        return builder.AddScheme<McpApiKeyAuthenticationOptions, McpApiKeyAuthenticationHandler>(
            McpApiKeyAuthenticationOptions.DefaultScheme,
            configureOptions ?? (_ => { }));
    }

    /// <summary>
    /// 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)
    /// </summary>
    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<McpCorrelationIdMiddleware>();

        // 2. Exception handler (SECOND - catches all errors from downstream middleware)
        app.UseMiddleware<McpExceptionHandlerMiddleware>();

        // 3. Logging middleware (THIRD - logs request/response with correlation ID)
        app.UseMiddleware<McpLoggingMiddleware>();

        // 4. Authentication middleware (FOURTH - validates API key)
        app.UseMiddleware<McpApiKeyAuthenticationMiddleware>();

        // 5. MCP protocol handler (LAST - processes the actual MCP request)
        app.UseMiddleware<McpMiddleware>();

        return app;
    }

    /// <summary>
    /// Initialize resource registry by registering all resources
    /// Enhanced with auto-discovery using Assembly scanning
    /// This is called once at startup
    /// </summary>
    private static void InitializeResourceRegistry(IApplicationBuilder app)
    {
        using var scope = app.ApplicationServices.CreateScope();
        var registry = scope.ServiceProvider.GetRequiredService<IMcpResourceRegistry>();
        var discoveryService = scope.ServiceProvider.GetRequiredService<IResourceDiscoveryService>();
        var loggerFactory = scope.ServiceProvider.GetRequiredService<ILoggerFactory>();
        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));
    }
}