From 8c51fa392b8331a513bdaa19b27cac294bff4930 Mon Sep 17 00:00:00 2001 From: Yaojia Wang Date: Sun, 23 Nov 2025 23:40:10 +0100 Subject: [PATCH] Refactoring --- .../Extensions/ModuleExtensions.cs | 4 ++- .../Handlers/ResourcesReadMethodHandler.cs | 18 +++++++--- .../Extensions/McpServiceExtensions.cs | 3 ++ .../McpApiKeyAuthenticationMiddleware.cs | 4 +-- .../McpExceptionHandlerMiddleware.cs | 7 ++++ .../Middleware/McpLoggingMiddleware.cs | 3 +- .../Middleware/McpMiddleware.cs | 4 +-- .../Common/Interfaces/ITenantContext.cs | 6 ++++ .../Persistence/PMDbContext.cs | 18 ++++------ .../Persistence/PMDbContextFactory.cs | 16 +++++---- .../Services/TenantContext.cs | 34 ++++++++++++++++--- 11 files changed, 84 insertions(+), 33 deletions(-) diff --git a/colaflow-api/src/ColaFlow.API/Extensions/ModuleExtensions.cs b/colaflow-api/src/ColaFlow.API/Extensions/ModuleExtensions.cs index 7371474..56706e2 100644 --- a/colaflow-api/src/ColaFlow.API/Extensions/ModuleExtensions.cs +++ b/colaflow-api/src/ColaFlow.API/Extensions/ModuleExtensions.cs @@ -66,11 +66,13 @@ public static class ModuleExtensions services.AddScoped(); - // Register MediatR handlers from Application assembly (v13.x syntax) + // Register MediatR handlers from Application assemblies (v13.x syntax) + // Consolidate all module handler registrations here to avoid duplicate registrations services.AddMediatR(cfg => { cfg.LicenseKey = configuration["MediatR:LicenseKey"]; cfg.RegisterServicesFromAssembly(typeof(CreateProjectCommand).Assembly); + cfg.RegisterServicesFromAssembly(typeof(ColaFlow.Modules.Mcp.Application.EventHandlers.PendingChangeApprovedEventHandler).Assembly); }); // Register FluentValidation validators diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Application/Handlers/ResourcesReadMethodHandler.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Application/Handlers/ResourcesReadMethodHandler.cs index 0405e74..235e4cd 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Application/Handlers/ResourcesReadMethodHandler.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Application/Handlers/ResourcesReadMethodHandler.cs @@ -9,10 +9,12 @@ namespace ColaFlow.Modules.Mcp.Application.Handlers; /// /// Handler for the 'resources/read' MCP method +/// Uses scoped IMcpResource instances from DI to avoid DbContext disposal issues /// public class ResourcesReadMethodHandler( ILogger logger, - IMcpResourceRegistry resourceRegistry) + IMcpResourceRegistry resourceRegistry, + IEnumerable scopedResources) : IMcpMethodHandler { public string MethodName => "resources/read"; @@ -32,13 +34,20 @@ public class ResourcesReadMethodHandler( logger.LogInformation("Reading resource: {Uri}", request.Uri); - // Find resource by URI - var resource = resourceRegistry.GetResourceByUri(request.Uri); - if (resource == null) + // Find resource descriptor from registry (for URI template matching) + var registryResource = resourceRegistry.GetResourceByUri(request.Uri); + if (registryResource == null) { throw new McpNotFoundException($"Resource not found: {request.Uri}"); } + // Get the scoped resource instance from DI (fresh DbContext) + var resource = scopedResources.FirstOrDefault(r => r.Uri == registryResource.Uri); + if (resource == null) + { + throw new McpNotFoundException($"Resource implementation not found: {registryResource.Uri}"); + } + // Parse URI and extract parameters var resourceRequest = ParseResourceRequest(request.Uri, resource.Uri); @@ -114,6 +123,7 @@ public class ResourcesReadMethodHandler( private class ResourceReadParams { + [System.Text.Json.Serialization.JsonPropertyName("uri")] public string Uri { get; set; } = string.Empty; } } diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Extensions/McpServiceExtensions.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Extensions/McpServiceExtensions.cs index ee86d06..1cb111d 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Extensions/McpServiceExtensions.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Extensions/McpServiceExtensions.cs @@ -78,6 +78,9 @@ public static class McpServiceExtensions 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; } diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpApiKeyAuthenticationMiddleware.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpApiKeyAuthenticationMiddleware.cs index 9284a7b..70514a5 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpApiKeyAuthenticationMiddleware.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpApiKeyAuthenticationMiddleware.cs @@ -15,8 +15,8 @@ public class McpApiKeyAuthenticationMiddleware( { public async Task InvokeAsync(HttpContext context, IMcpApiKeyService apiKeyService) { - // Only apply to /mcp endpoints - if (!context.Request.Path.StartsWithSegments("/mcp")) + // Only apply to /mcp-sdk endpoint (not /api/mcp/* which uses JWT auth) + if (!context.Request.Path.StartsWithSegments("/mcp-sdk")) { await next(context); return; diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpExceptionHandlerMiddleware.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpExceptionHandlerMiddleware.cs index 3f2b333..d9fbff0 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpExceptionHandlerMiddleware.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpExceptionHandlerMiddleware.cs @@ -16,6 +16,13 @@ public class McpExceptionHandlerMiddleware( { public async Task InvokeAsync(HttpContext context) { + // Only handle exceptions for MCP endpoints + if (!context.Request.Path.StartsWithSegments("/mcp-sdk")) + { + await next(context); + return; + } + try { await next(context); diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpLoggingMiddleware.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpLoggingMiddleware.cs index ec403cf..834e42d 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpLoggingMiddleware.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpLoggingMiddleware.cs @@ -64,11 +64,12 @@ public class McpLoggingMiddleware( /// /// Checks if this is an MCP request + /// Only applies to /mcp-sdk endpoint (not /api/mcp/* which uses JWT auth) /// private static bool IsMcpRequest(HttpContext context) { return context.Request.Method == "POST" - && context.Request.Path.StartsWithSegments("/mcp"); + && context.Request.Path.StartsWithSegments("/mcp-sdk"); } /// diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpMiddleware.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpMiddleware.cs index b4d912b..0508b71 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpMiddleware.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpMiddleware.cs @@ -13,8 +13,8 @@ public class McpMiddleware(RequestDelegate next, ILogger logger) { public async Task InvokeAsync(HttpContext context, IMcpProtocolHandler protocolHandler) { - // Only handle POST requests to /mcp endpoint - if (context.Request.Method != "POST" || !context.Request.Path.StartsWithSegments("/mcp")) + // Only handle POST requests to /mcp-sdk endpoint (not /api/mcp/* which uses JWT auth) + if (context.Request.Method != "POST" || !context.Request.Path.StartsWithSegments("/mcp-sdk")) { await next(context); return; diff --git a/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Application/Common/Interfaces/ITenantContext.cs b/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Application/Common/Interfaces/ITenantContext.cs index e773891..f0b675b 100644 --- a/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Application/Common/Interfaces/ITenantContext.cs +++ b/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Application/Common/Interfaces/ITenantContext.cs @@ -12,6 +12,12 @@ public interface ITenantContext /// Thrown when tenant context is not available Guid GetCurrentTenantId(); + /// + /// Tries to get the current tenant ID without throwing exceptions + /// + /// The current tenant ID or null if not available + Guid? TryGetCurrentTenantId(); + /// /// Gets the current user ID from claims (optional - may be null for system operations) /// diff --git a/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Persistence/PMDbContext.cs b/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Persistence/PMDbContext.cs index 08a5930..e7bcca9 100644 --- a/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Persistence/PMDbContext.cs +++ b/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Persistence/PMDbContext.cs @@ -1,5 +1,4 @@ using System.Reflection; -using Microsoft.AspNetCore.Http; using Microsoft.EntityFrameworkCore; using ColaFlow.Modules.ProjectManagement.Application.Common.Interfaces; using ColaFlow.Modules.ProjectManagement.Domain.Aggregates.ProjectAggregate; @@ -10,8 +9,9 @@ namespace ColaFlow.Modules.ProjectManagement.Infrastructure.Persistence; /// /// Project Management Module DbContext +/// Uses ITenantContext for centralized tenant resolution (supports both JWT and MCP API Key authentication) /// -public class PMDbContext(DbContextOptions options, IHttpContextAccessor httpContextAccessor) +public class PMDbContext(DbContextOptions options, ITenantContext tenantContext) : DbContext(options), IApplicationDbContext { public DbSet Projects => Set(); @@ -53,16 +53,12 @@ public class PMDbContext(DbContextOptions options, IHttpContextAcce private TenantId GetCurrentTenantId() { - var tenantIdClaim = httpContextAccessor?.HttpContext?.User - .FindFirst("tenant_id")?.Value; + // Use centralized tenant resolution from ITenantContext + // TryGetCurrentTenantId handles both JWT claims and MCP API Key authentication + var tenantId = tenantContext.TryGetCurrentTenantId(); - if (Guid.TryParse(tenantIdClaim, out var tenantId) && tenantId != Guid.Empty) - { - return TenantId.From(tenantId); - } - - // Return a dummy value for queries outside HTTP context (e.g., migrations) + // Return dummy value for queries outside tenant context (e.g., migrations) // These will return no results due to the filter - return TenantId.From(Guid.Empty); + return TenantId.From(tenantId ?? Guid.Empty); } } diff --git a/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Persistence/PMDbContextFactory.cs b/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Persistence/PMDbContextFactory.cs index 4112486..3822dc3 100644 --- a/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Persistence/PMDbContextFactory.cs +++ b/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Persistence/PMDbContextFactory.cs @@ -1,6 +1,6 @@ -using Microsoft.AspNetCore.Http; using Microsoft.EntityFrameworkCore; using Microsoft.EntityFrameworkCore.Design; +using ColaFlow.Modules.ProjectManagement.Application.Common.Interfaces; namespace ColaFlow.Modules.ProjectManagement.Infrastructure.Persistence; @@ -19,16 +19,18 @@ public class PMDbContextFactory : IDesignTimeDbContextFactory var optionsBuilder = new DbContextOptionsBuilder(); optionsBuilder.UseNpgsql(connectionString, b => b.MigrationsAssembly("ColaFlow.Modules.ProjectManagement.Infrastructure")); - // Create DbContext with a mock HttpContextAccessor (for migrations only) - return new PMDbContext(optionsBuilder.Options, new MockHttpContextAccessor()); + // Create DbContext with a mock TenantContext (for migrations only) + return new PMDbContext(optionsBuilder.Options, new MockTenantContext()); } /// - /// Mock HttpContextAccessor for design-time operations - /// Returns null HttpContext which PMDbContext handles gracefully + /// Mock TenantContext for design-time operations + /// Returns null/empty tenant ID which PMDbContext handles gracefully /// - private class MockHttpContextAccessor : IHttpContextAccessor + private class MockTenantContext : ITenantContext { - public HttpContext? HttpContext { get; set; } = null; + public Guid GetCurrentTenantId() => Guid.Empty; + public Guid? TryGetCurrentTenantId() => null; + public Guid? GetCurrentUserId() => null; } } diff --git a/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Services/TenantContext.cs b/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Services/TenantContext.cs index 5958433..5e590c1 100644 --- a/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Services/TenantContext.cs +++ b/colaflow-api/src/Modules/ProjectManagement/ColaFlow.Modules.ProjectManagement.Infrastructure/Services/TenantContext.cs @@ -5,23 +5,40 @@ using ColaFlow.Modules.ProjectManagement.Application.Common.Interfaces; namespace ColaFlow.Modules.ProjectManagement.Infrastructure.Services; /// -/// Implementation of ITenantContext that retrieves tenant ID from JWT claims +/// Implementation of ITenantContext that retrieves tenant ID from JWT claims or MCP API Key context +/// Supports both JWT authentication (claims) and MCP API Key authentication (HttpContext.Items) /// public sealed class TenantContext(IHttpContextAccessor httpContextAccessor) : ITenantContext { public Guid GetCurrentTenantId() + { + var tenantId = TryGetCurrentTenantId(); + if (tenantId == null) + throw new UnauthorizedAccessException("Tenant ID not found in claims"); + + return tenantId.Value; + } + + public Guid? TryGetCurrentTenantId() { var httpContext = httpContextAccessor.HttpContext; if (httpContext == null) - throw new InvalidOperationException("HTTP context is not available"); + return null; + // First, try to get from MCP API Key authentication (HttpContext.Items) + if (httpContext.Items.TryGetValue("McpTenantId", out var mcpTenantIdObj) && mcpTenantIdObj is Guid mcpTenantId) + { + return mcpTenantId; + } + + // Fallback to JWT claims var user = httpContext.User; var tenantClaim = user.FindFirst("tenant_id") ?? user.FindFirst("tenantId"); - if (tenantClaim == null || !Guid.TryParse(tenantClaim.Value, out var tenantId)) - throw new UnauthorizedAccessException("Tenant ID not found in claims"); + if (tenantClaim != null && Guid.TryParse(tenantClaim.Value, out var tenantId)) + return tenantId; - return tenantId; + return null; } public Guid? GetCurrentUserId() @@ -30,6 +47,13 @@ public sealed class TenantContext(IHttpContextAccessor httpContextAccessor) : IT if (httpContext == null) return null; + // First, try to get from MCP API Key authentication (HttpContext.Items) + if (httpContext.Items.TryGetValue("McpUserId", out var mcpUserIdObj) && mcpUserIdObj is Guid mcpUserId) + { + return mcpUserId; + } + + // Fallback to JWT claims var user = httpContext.User; var userIdClaim = user.FindFirst(ClaimTypes.NameIdentifier) ?? user.FindFirst("sub")