diff --git a/colaflow-api/src/ColaFlow.API/Program.cs b/colaflow-api/src/ColaFlow.API/Program.cs index 5a14f35..bd031e5 100644 --- a/colaflow-api/src/ColaFlow.API/Program.cs +++ b/colaflow-api/src/ColaFlow.API/Program.cs @@ -12,10 +12,28 @@ using Microsoft.AspNetCore.Authentication.JwtBearer; using Microsoft.EntityFrameworkCore; using Microsoft.IdentityModel.Tokens; using Scalar.AspNetCore; +using Serilog; using System.Text; var builder = WebApplication.CreateBuilder(args); +// Configure Serilog +builder.Host.UseSerilog((context, services, configuration) => +{ + configuration + .ReadFrom.Configuration(context.Configuration) + .Enrich.FromLogContext() + .Enrich.WithProperty("Application", "ColaFlow") + .Enrich.WithProperty("Environment", context.HostingEnvironment.EnvironmentName) + .WriteTo.Console( + outputTemplate: "[{Timestamp:HH:mm:ss} {Level:u3}] {CorrelationId} {Message:lj}{NewLine}{Exception}") + .WriteTo.File( + path: "logs/colaflow-.log", + rollingInterval: RollingInterval.Day, + retainedFileCountLimit: 30, + outputTemplate: "[{Timestamp:yyyy-MM-dd HH:mm:ss.fff zzz}] [{Level:u3}] {CorrelationId} {Message:lj}{NewLine}{Exception}"); +}); + // Register ProjectManagement Module builder.Services.AddProjectManagementModule(builder.Configuration, builder.Environment); diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Contracts/JsonRpc/JsonRpcError.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Contracts/JsonRpc/JsonRpcError.cs index 38f4ed3..ca0c85c 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Contracts/JsonRpc/JsonRpcError.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Contracts/JsonRpc/JsonRpcError.cs @@ -26,6 +26,13 @@ public class JsonRpcError [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)] public object? Data { get; set; } + /// + /// Parameterless constructor for JSON deserialization + /// + public JsonRpcError() + { + } + /// /// Creates a new JSON-RPC error /// diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpException.cs new file mode 100644 index 0000000..16c5d86 --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpException.cs @@ -0,0 +1,68 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Base exception class for all MCP-related exceptions +/// Maps to JSON-RPC 2.0 error responses +/// +public abstract class McpException : Exception +{ + /// + /// JSON-RPC error code + /// + public JsonRpcErrorCode ErrorCode { get; } + + /// + /// Additional error data (optional, can be any JSON-serializable object) + /// + public object? ErrorData { get; } + + /// + /// Initializes a new instance of the class + /// + /// JSON-RPC error code + /// Error message + /// Additional error data (optional) + /// Inner exception (optional) + protected McpException( + JsonRpcErrorCode errorCode, + string message, + object? errorData = null, + Exception? innerException = null) + : base(message, innerException) + { + ErrorCode = errorCode; + ErrorData = errorData; + } + + /// + /// Converts this exception to a JSON-RPC error object + /// + /// JSON-RPC error object + public JsonRpcError ToJsonRpcError() + { + return new JsonRpcError(ErrorCode, Message, ErrorData); + } + + /// + /// Gets the HTTP status code that should be returned for this error + /// + /// HTTP status code + public virtual int GetHttpStatusCode() + { + return ErrorCode switch + { + JsonRpcErrorCode.Unauthorized => 401, + JsonRpcErrorCode.Forbidden => 403, + JsonRpcErrorCode.NotFound => 404, + JsonRpcErrorCode.ValidationFailed => 422, + JsonRpcErrorCode.ParseError => 400, + JsonRpcErrorCode.InvalidRequest => 400, + JsonRpcErrorCode.MethodNotFound => 404, + JsonRpcErrorCode.InvalidParams => 400, + JsonRpcErrorCode.InternalError => 500, + _ => 500 + }; + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpForbiddenException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpForbiddenException.cs new file mode 100644 index 0000000..7018d64 --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpForbiddenException.cs @@ -0,0 +1,21 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Exception thrown when authorization fails (authenticated but not allowed) +/// Maps to JSON-RPC error code -32002 (Forbidden) +/// HTTP 403 status code +/// +public class McpForbiddenException : McpException +{ + /// + /// Initializes a new instance of the class + /// + /// Error message + /// Additional error data (optional) + public McpForbiddenException(string message = "Forbidden", object? errorData = null) + : base(JsonRpcErrorCode.Forbidden, message, errorData) + { + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpInvalidParamsException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpInvalidParamsException.cs new file mode 100644 index 0000000..744e27f --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpInvalidParamsException.cs @@ -0,0 +1,20 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Exception thrown when invalid method parameters are provided +/// Maps to JSON-RPC error code -32602 (InvalidParams) +/// +public class McpInvalidParamsException : McpException +{ + /// + /// Initializes a new instance of the class + /// + /// Error message + /// Additional error data (optional) + public McpInvalidParamsException(string message = "Invalid params", object? errorData = null) + : base(JsonRpcErrorCode.InvalidParams, message, errorData) + { + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpInvalidRequestException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpInvalidRequestException.cs new file mode 100644 index 0000000..5f14168 --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpInvalidRequestException.cs @@ -0,0 +1,20 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Exception thrown when the JSON sent is not a valid Request object +/// Maps to JSON-RPC error code -32600 (InvalidRequest) +/// +public class McpInvalidRequestException : McpException +{ + /// + /// Initializes a new instance of the class + /// + /// Error message + /// Additional error data (optional) + public McpInvalidRequestException(string message = "Invalid Request", object? errorData = null) + : base(JsonRpcErrorCode.InvalidRequest, message, errorData) + { + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpMethodNotFoundException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpMethodNotFoundException.cs new file mode 100644 index 0000000..e649010 --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpMethodNotFoundException.cs @@ -0,0 +1,19 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Exception thrown when the requested method does not exist or is not available +/// Maps to JSON-RPC error code -32601 (MethodNotFound) +/// +public class McpMethodNotFoundException : McpException +{ + /// + /// Initializes a new instance of the class + /// + /// The method name that was not found + public McpMethodNotFoundException(string method) + : base(JsonRpcErrorCode.MethodNotFound, $"Method not found: {method}") + { + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpNotFoundException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpNotFoundException.cs new file mode 100644 index 0000000..9ddc417 --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpNotFoundException.cs @@ -0,0 +1,31 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Exception thrown when a requested resource is not found +/// Maps to JSON-RPC error code -32003 (NotFound) +/// HTTP 404 status code +/// +public class McpNotFoundException : McpException +{ + /// + /// Initializes a new instance of the class + /// + /// Type of resource (e.g., "Task", "Epic") + /// ID of the resource + public McpNotFoundException(string resourceType, string resourceId) + : base(JsonRpcErrorCode.NotFound, $"{resourceType} not found: {resourceId}") + { + } + + /// + /// Initializes a new instance of the class with custom message + /// + /// Error message + /// Additional error data (optional) + public McpNotFoundException(string message, object? errorData = null) + : base(JsonRpcErrorCode.NotFound, message, errorData) + { + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpParseException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpParseException.cs new file mode 100644 index 0000000..ccb380a --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpParseException.cs @@ -0,0 +1,31 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Exception thrown when invalid JSON is received by the server +/// Maps to JSON-RPC error code -32700 (ParseError) +/// +public class McpParseException : McpException +{ + /// + /// Initializes a new instance of the class + /// + /// Error message + /// Inner exception (optional) + public McpParseException(string message = "Parse error", Exception? innerException = null) + : base(JsonRpcErrorCode.ParseError, message, null, innerException) + { + } + + /// + /// Initializes a new instance of the class with additional error data + /// + /// Error message + /// Additional error data + /// Inner exception (optional) + public McpParseException(string message, object? errorData, Exception? innerException = null) + : base(JsonRpcErrorCode.ParseError, message, errorData, innerException) + { + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpUnauthorizedException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpUnauthorizedException.cs new file mode 100644 index 0000000..14e757b --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpUnauthorizedException.cs @@ -0,0 +1,21 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Exception thrown when authentication fails +/// Maps to JSON-RPC error code -32001 (Unauthorized) +/// HTTP 401 status code +/// +public class McpUnauthorizedException : McpException +{ + /// + /// Initializes a new instance of the class + /// + /// Error message + /// Additional error data (optional) + public McpUnauthorizedException(string message = "Unauthorized", object? errorData = null) + : base(JsonRpcErrorCode.Unauthorized, message, errorData) + { + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpValidationException.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpValidationException.cs new file mode 100644 index 0000000..125386d --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Domain/Exceptions/McpValidationException.cs @@ -0,0 +1,21 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; + +namespace ColaFlow.Modules.Mcp.Domain.Exceptions; + +/// +/// Exception thrown when request validation fails +/// Maps to JSON-RPC error code -32004 (ValidationFailed) +/// HTTP 422 status code +/// +public class McpValidationException : McpException +{ + /// + /// Initializes a new instance of the class + /// + /// Error message + /// Additional error data (optional, e.g., validation errors dictionary) + public McpValidationException(string message = "Validation failed", object? errorData = null) + : base(JsonRpcErrorCode.ValidationFailed, message, errorData) + { + } +} diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/ColaFlow.Modules.Mcp.Infrastructure.csproj b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/ColaFlow.Modules.Mcp.Infrastructure.csproj index c91d253..63212f3 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/ColaFlow.Modules.Mcp.Infrastructure.csproj +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/ColaFlow.Modules.Mcp.Infrastructure.csproj @@ -23,6 +23,7 @@ all runtime; build; native; contentfiles; analyzers; buildtransitive + 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 f9cf3e8..3df2d8b 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 @@ -48,14 +48,28 @@ public static class McpServiceExtensions /// /// Adds MCP middleware to the application pipeline - /// IMPORTANT: Add authentication middleware BEFORE MCP middleware + /// 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) { - // Authentication middleware MUST come first + // 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(); - // Then the MCP protocol handler + // 5. MCP protocol handler (LAST - processes the actual MCP request) app.UseMiddleware(); return app; diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpCorrelationIdMiddleware.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpCorrelationIdMiddleware.cs new file mode 100644 index 0000000..32cb4a5 --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpCorrelationIdMiddleware.cs @@ -0,0 +1,63 @@ +using Microsoft.AspNetCore.Http; +using Serilog.Context; + +namespace ColaFlow.Modules.Mcp.Infrastructure.Middleware; + +/// +/// Middleware that generates or extracts correlation ID for request tracking +/// Adds the correlation ID to: +/// - HttpContext.Items (for access in downstream middleware/handlers) +/// - Response headers (for client-side tracking) +/// - Serilog LogContext (for structured logging) +/// +public class McpCorrelationIdMiddleware +{ + private readonly RequestDelegate _next; + private const string CorrelationIdHeaderName = "X-Correlation-Id"; + + public McpCorrelationIdMiddleware(RequestDelegate next) + { + _next = next; + } + + public async Task InvokeAsync(HttpContext context) + { + // Try to get correlation ID from request header, otherwise generate new one + var correlationId = GetOrCreateCorrelationId(context); + + // Store in HttpContext.Items for access in downstream middleware/handlers + context.Items["CorrelationId"] = correlationId; + + // Add to response header so client can track the request + context.Response.OnStarting(() => + { + if (!context.Response.Headers.ContainsKey(CorrelationIdHeaderName)) + { + context.Response.Headers.TryAdd(CorrelationIdHeaderName, correlationId); + } + return Task.CompletedTask; + }); + + // Add to Serilog LogContext so it appears in all log entries for this request + using (LogContext.PushProperty("CorrelationId", correlationId)) + { + await _next(context); + } + } + + /// + /// Gets correlation ID from request header or generates a new one + /// + private static string GetOrCreateCorrelationId(HttpContext context) + { + // Check if client provided a correlation ID + if (context.Request.Headers.TryGetValue(CorrelationIdHeaderName, out var correlationId) + && !string.IsNullOrWhiteSpace(correlationId)) + { + return correlationId.ToString(); + } + + // Generate new correlation ID + return Guid.NewGuid().ToString("N"); // Use "N" format (no hyphens) for shorter IDs + } +} 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 new file mode 100644 index 0000000..6e09fc5 --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpExceptionHandlerMiddleware.cs @@ -0,0 +1,111 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; +using ColaFlow.Modules.Mcp.Domain.Exceptions; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using System.Text.Json; + +namespace ColaFlow.Modules.Mcp.Infrastructure.Middleware; + +/// +/// Global exception handler middleware for MCP requests +/// Catches all unhandled exceptions and converts them to JSON-RPC error responses +/// +public class McpExceptionHandlerMiddleware +{ + private readonly RequestDelegate _next; + private readonly ILogger _logger; + + public McpExceptionHandlerMiddleware( + RequestDelegate next, + ILogger logger) + { + _next = next; + _logger = logger; + } + + public async Task InvokeAsync(HttpContext context) + { + try + { + await _next(context); + } + catch (McpException mcpEx) + { + await HandleMcpExceptionAsync(context, mcpEx); + } + catch (Exception ex) + { + await HandleUnexpectedExceptionAsync(context, ex); + } + } + + /// + /// Handles known MCP exceptions by converting them to JSON-RPC error responses + /// + private async Task HandleMcpExceptionAsync(HttpContext context, McpException mcpEx) + { + var correlationId = context.Items["CorrelationId"] as string; + var tenantId = context.Items["TenantId"]?.ToString(); + var apiKeyId = context.Items["ApiKeyId"]?.ToString(); + + // Log the error with structured data + _logger.LogError(mcpEx, + "MCP Error: {ErrorCode} - {Message} | CorrelationId: {CorrelationId} | TenantId: {TenantId} | ApiKeyId: {ApiKeyId}", + mcpEx.ErrorCode, mcpEx.Message, correlationId, tenantId, apiKeyId); + + // Create JSON-RPC error response + var response = new JsonRpcResponse + { + JsonRpc = "2.0", + Error = mcpEx.ToJsonRpcError(), + Id = context.Items["McpRequestId"] // Preserve request ID from parsed request + }; + + // Set HTTP status code based on error type + context.Response.StatusCode = mcpEx.GetHttpStatusCode(); + context.Response.ContentType = "application/json"; + + // Write response + await context.Response.WriteAsync( + JsonSerializer.Serialize(response, new JsonSerializerOptions + { + PropertyNamingPolicy = JsonNamingPolicy.CamelCase + })); + } + + /// + /// Handles unexpected exceptions (not MCP-specific) by converting them to InternalError + /// IMPORTANT: Never expose internal exception details to clients + /// + private async Task HandleUnexpectedExceptionAsync(HttpContext context, Exception ex) + { + var correlationId = context.Items["CorrelationId"] as string; + var tenantId = context.Items["TenantId"]?.ToString(); + var apiKeyId = context.Items["ApiKeyId"]?.ToString(); + + // Log the full exception with stack trace + _logger.LogError(ex, + "Unexpected error in MCP Server | CorrelationId: {CorrelationId} | TenantId: {TenantId} | ApiKeyId: {ApiKeyId}", + correlationId, tenantId, apiKeyId); + + // Create generic InternalError response (DO NOT expose exception details) + var response = new JsonRpcResponse + { + JsonRpc = "2.0", + Error = new JsonRpcError( + JsonRpcErrorCode.InternalError, + "Internal server error", + null), // No error data to avoid leaking internals + Id = context.Items["McpRequestId"] + }; + + context.Response.StatusCode = 500; + context.Response.ContentType = "application/json"; + + await context.Response.WriteAsync( + JsonSerializer.Serialize(response, new JsonSerializerOptions + { + PropertyNamingPolicy = JsonNamingPolicy.CamelCase + })); + } +} 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 new file mode 100644 index 0000000..a5ea8bf --- /dev/null +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Middleware/McpLoggingMiddleware.cs @@ -0,0 +1,161 @@ +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using System.Diagnostics; +using System.Text; +using System.Text.RegularExpressions; + +namespace ColaFlow.Modules.Mcp.Infrastructure.Middleware; + +/// +/// Middleware that logs all MCP requests and responses +/// Includes performance timing and sensitive data sanitization +/// +public class McpLoggingMiddleware +{ + private readonly RequestDelegate _next; + private readonly ILogger _logger; + + // Patterns for sanitizing sensitive data + private static readonly Regex ApiKeyHashPattern = new(@"""keyHash"":\s*""[^""]+""", RegexOptions.Compiled); + private static readonly Regex ApiKeyPattern = new(@"""apiKey"":\s*""[^""]+""", RegexOptions.Compiled); + private static readonly Regex PasswordPattern = new(@"""password"":\s*""[^""]+""", RegexOptions.Compiled); + + public McpLoggingMiddleware( + RequestDelegate next, + ILogger logger) + { + _next = next; + _logger = logger; + } + + public async Task InvokeAsync(HttpContext context) + { + // Only log MCP requests (POST to /mcp endpoint) + if (!IsMcpRequest(context)) + { + await _next(context); + return; + } + + var correlationId = context.Items["CorrelationId"] as string; + var stopwatch = Stopwatch.StartNew(); + + // Enable request buffering so we can read the body multiple times + context.Request.EnableBuffering(); + + // Log request + await LogRequestAsync(context, correlationId); + + // Capture response + var originalBodyStream = context.Response.Body; + using var responseBodyStream = new MemoryStream(); + context.Response.Body = responseBodyStream; + + try + { + // Execute the rest of the pipeline + await _next(context); + + stopwatch.Stop(); + + // Log response + await LogResponseAsync(context, correlationId, stopwatch.ElapsedMilliseconds); + + // Copy response to original stream + responseBodyStream.Seek(0, SeekOrigin.Begin); + await responseBodyStream.CopyToAsync(originalBodyStream); + } + finally + { + context.Response.Body = originalBodyStream; + } + } + + /// + /// Checks if this is an MCP request + /// + private static bool IsMcpRequest(HttpContext context) + { + return context.Request.Method == "POST" + && context.Request.Path.StartsWithSegments("/mcp"); + } + + /// + /// Logs the incoming MCP request + /// + private async Task LogRequestAsync(HttpContext context, string? correlationId) + { + context.Request.Body.Position = 0; + var bodyText = await new StreamReader(context.Request.Body, Encoding.UTF8).ReadToEndAsync(); + context.Request.Body.Position = 0; // Reset for next middleware + + var tenantId = context.Items["TenantId"]; + var apiKeyId = context.Items["ApiKeyId"]; + var userId = context.Items["UserId"]; + + // Sanitize sensitive data before logging + var sanitizedBody = SanitizeSensitiveData(bodyText); + + _logger.LogDebug( + "MCP Request | Method: {Method} | Path: {Path} | CorrelationId: {CorrelationId} | " + + "TenantId: {TenantId} | ApiKeyId: {ApiKeyId} | UserId: {UserId}\nBody: {Body}", + context.Request.Method, + context.Request.Path, + correlationId, + tenantId, + apiKeyId, + userId, + sanitizedBody); + } + + /// + /// Logs the outgoing MCP response + /// + private async Task LogResponseAsync( + HttpContext context, + string? correlationId, + long elapsedMs) + { + context.Response.Body.Position = 0; + var bodyText = await new StreamReader(context.Response.Body, Encoding.UTF8).ReadToEndAsync(); + context.Response.Body.Position = 0; + + var statusCode = context.Response.StatusCode; + var logLevel = statusCode >= 400 ? LogLevel.Error : LogLevel.Debug; + + _logger.Log(logLevel, + "MCP Response | StatusCode: {StatusCode} | CorrelationId: {CorrelationId} | " + + "Duration: {Duration}ms\nBody: {Body}", + statusCode, + correlationId, + elapsedMs, + bodyText); + + // Also log performance metrics + if (elapsedMs > 1000) // Log slow requests (> 1 second) + { + _logger.LogWarning( + "Slow MCP Request | CorrelationId: {CorrelationId} | Duration: {Duration}ms", + correlationId, + elapsedMs); + } + } + + /// + /// Sanitizes sensitive data from request/response bodies + /// IMPORTANT: Prevents logging API keys, passwords, etc. + /// + private static string SanitizeSensitiveData(string body) + { + if (string.IsNullOrWhiteSpace(body)) + return body; + + // Replace sensitive fields with [REDACTED] + var sanitized = body; + sanitized = ApiKeyHashPattern.Replace(sanitized, "\"keyHash\":\"[REDACTED]\""); + sanitized = ApiKeyPattern.Replace(sanitized, "\"apiKey\":\"[REDACTED]\""); + sanitized = PasswordPattern.Replace(sanitized, "\"password\":\"[REDACTED]\""); + + return sanitized; + } +} diff --git a/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/ColaFlow.Modules.Mcp.Tests.csproj b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/ColaFlow.Modules.Mcp.Tests.csproj index c5e1d3e..ad7c88c 100644 --- a/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/ColaFlow.Modules.Mcp.Tests.csproj +++ b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/ColaFlow.Modules.Mcp.Tests.csproj @@ -10,6 +10,7 @@ + @@ -21,6 +22,7 @@ + diff --git a/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Domain/Exceptions/McpExceptionTests.cs b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Domain/Exceptions/McpExceptionTests.cs new file mode 100644 index 0000000..4031270 --- /dev/null +++ b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Domain/Exceptions/McpExceptionTests.cs @@ -0,0 +1,158 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; +using ColaFlow.Modules.Mcp.Domain.Exceptions; +using FluentAssertions; + +namespace ColaFlow.Modules.Mcp.Tests.Domain.Exceptions; + +/// +/// Unit tests for MCP exception classes +/// +public class McpExceptionTests +{ + [Fact] + public void McpParseException_ShouldHaveCorrectErrorCode() + { + // Arrange & Act + var exception = new McpParseException("Invalid JSON"); + + // Assert + exception.ErrorCode.Should().Be(JsonRpcErrorCode.ParseError); + exception.Message.Should().Be("Invalid JSON"); + exception.GetHttpStatusCode().Should().Be(400); + } + + [Fact] + public void McpInvalidRequestException_ShouldHaveCorrectErrorCode() + { + // Arrange & Act + var exception = new McpInvalidRequestException("Missing required field"); + + // Assert + exception.ErrorCode.Should().Be(JsonRpcErrorCode.InvalidRequest); + exception.Message.Should().Be("Missing required field"); + exception.GetHttpStatusCode().Should().Be(400); + } + + [Fact] + public void McpMethodNotFoundException_ShouldHaveCorrectErrorCode() + { + // Arrange & Act + var exception = new McpMethodNotFoundException("unknown_method"); + + // Assert + exception.ErrorCode.Should().Be(JsonRpcErrorCode.MethodNotFound); + exception.Message.Should().Be("Method not found: unknown_method"); + exception.GetHttpStatusCode().Should().Be(404); + } + + [Fact] + public void McpInvalidParamsException_ShouldHaveCorrectErrorCode() + { + // Arrange & Act + var exception = new McpInvalidParamsException("Invalid parameter type"); + + // Assert + exception.ErrorCode.Should().Be(JsonRpcErrorCode.InvalidParams); + exception.Message.Should().Be("Invalid parameter type"); + exception.GetHttpStatusCode().Should().Be(400); + } + + [Fact] + public void McpUnauthorizedException_ShouldHaveCorrectErrorCode() + { + // Arrange & Act + var exception = new McpUnauthorizedException("Invalid API key"); + + // Assert + exception.ErrorCode.Should().Be(JsonRpcErrorCode.Unauthorized); + exception.Message.Should().Be("Invalid API key"); + exception.GetHttpStatusCode().Should().Be(401); + } + + [Fact] + public void McpForbiddenException_ShouldHaveCorrectErrorCode() + { + // Arrange & Act + var exception = new McpForbiddenException("Insufficient permissions"); + + // Assert + exception.ErrorCode.Should().Be(JsonRpcErrorCode.Forbidden); + exception.Message.Should().Be("Insufficient permissions"); + exception.GetHttpStatusCode().Should().Be(403); + } + + [Fact] + public void McpNotFoundException_ShouldHaveCorrectErrorCode() + { + // Arrange & Act + var exception = new McpNotFoundException("Task", "task-123"); + + // Assert + exception.ErrorCode.Should().Be(JsonRpcErrorCode.NotFound); + exception.Message.Should().Be("Task not found: task-123"); + exception.GetHttpStatusCode().Should().Be(404); + } + + [Fact] + public void McpValidationException_ShouldHaveCorrectErrorCode() + { + // Arrange & Act + var errorData = new { field = "name", error = "required" }; + var exception = new McpValidationException("Validation failed", errorData); + + // Assert + exception.ErrorCode.Should().Be(JsonRpcErrorCode.ValidationFailed); + exception.Message.Should().Be("Validation failed"); + exception.ErrorData.Should().Be(errorData); + exception.GetHttpStatusCode().Should().Be(422); + } + + [Fact] + public void ToJsonRpcError_ShouldConvertToJsonRpcError() + { + // Arrange + var errorData = new { detail = "test" }; + var exception = new McpValidationException("Test error", errorData); + + // Act + var jsonRpcError = exception.ToJsonRpcError(); + + // Assert + jsonRpcError.Code.Should().Be((int)JsonRpcErrorCode.ValidationFailed); + jsonRpcError.Message.Should().Be("Test error"); + jsonRpcError.Data.Should().Be(errorData); + } + + [Fact] + public void McpParseException_WithInnerException_ShouldPreserveInnerException() + { + // Arrange + var innerException = new InvalidOperationException("Root cause"); + + // Act + var exception = new McpParseException("Parse failed", innerException); + + // Assert + exception.InnerException.Should().Be(innerException); + exception.InnerException.Message.Should().Be("Root cause"); + } + + [Fact] + public void McpException_WithErrorData_ShouldPreserveErrorData() + { + // Arrange + var errorData = new + { + field = "email", + error = "invalid format", + value = "not-an-email" + }; + + // Act + var exception = new McpValidationException("Validation failed", errorData); + + // Assert + exception.ErrorData.Should().NotBeNull(); + exception.ErrorData.Should().Be(errorData); + } +} diff --git a/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpCorrelationIdMiddlewareTests.cs b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpCorrelationIdMiddlewareTests.cs new file mode 100644 index 0000000..5837980 --- /dev/null +++ b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpCorrelationIdMiddlewareTests.cs @@ -0,0 +1,129 @@ +using ColaFlow.Modules.Mcp.Infrastructure.Middleware; +using FluentAssertions; +using Microsoft.AspNetCore.Http; +using NSubstitute; + +namespace ColaFlow.Modules.Mcp.Tests.Infrastructure.Middleware; + +/// +/// Unit tests for McpCorrelationIdMiddleware +/// +public class McpCorrelationIdMiddlewareTests +{ + [Fact] + public async Task InvokeAsync_ShouldGenerateCorrelationId_WhenNotProvidedInRequest() + { + // Arrange + var context = new DefaultHttpContext(); + var nextCalled = false; + RequestDelegate next = (HttpContext ctx) => + { + nextCalled = true; + return Task.CompletedTask; + }; + + var middleware = new McpCorrelationIdMiddleware(next); + + // Act + await middleware.InvokeAsync(context); + + // Assert + nextCalled.Should().BeTrue(); + context.Items.Should().ContainKey("CorrelationId"); + context.Items["CorrelationId"].Should().NotBeNull(); + context.Items["CorrelationId"].Should().BeOfType(); + + var correlationId = context.Items["CorrelationId"] as string; + correlationId.Should().NotBeNullOrWhiteSpace(); + + // Should be a valid GUID without hyphens (format "N") + Guid.TryParse(correlationId, out _).Should().BeTrue(); + } + + [Fact] + public async Task InvokeAsync_ShouldUseProvidedCorrelationId_WhenPresentInRequestHeader() + { + // Arrange + var expectedCorrelationId = Guid.NewGuid().ToString("N"); + var context = new DefaultHttpContext(); + context.Request.Headers["X-Correlation-Id"] = expectedCorrelationId; + + RequestDelegate next = (HttpContext ctx) => Task.CompletedTask; + var middleware = new McpCorrelationIdMiddleware(next); + + // Act + await middleware.InvokeAsync(context); + + // Assert + context.Items["CorrelationId"].Should().Be(expectedCorrelationId); + } + + [Fact] + public async Task InvokeAsync_ShouldAddCorrelationIdToResponseHeader() + { + // Arrange + var context = new DefaultHttpContext(); + context.Response.Body = new MemoryStream(); // Required for response headers + + RequestDelegate next = (HttpContext ctx) => + { + // Trigger response to start (which adds headers) + ctx.Response.StatusCode = 200; + return Task.CompletedTask; + }; + + var middleware = new McpCorrelationIdMiddleware(next); + + // Act + await middleware.InvokeAsync(context); + + // Assert - Note: Headers are added via OnStarting callback, which executes when response starts + // In this test, we can verify the correlation ID was stored in context items + context.Items.Should().ContainKey("CorrelationId"); + } + + [Fact] + public async Task InvokeAsync_ShouldCallNextMiddleware() + { + // Arrange + var context = new DefaultHttpContext(); + var nextCalled = false; + + RequestDelegate next = (HttpContext ctx) => + { + nextCalled = true; + return Task.CompletedTask; + }; + + var middleware = new McpCorrelationIdMiddleware(next); + + // Act + await middleware.InvokeAsync(context); + + // Assert + nextCalled.Should().BeTrue(); + } + + [Fact] + public async Task InvokeAsync_ShouldGenerateUniqueCorrelationIds() + { + // Arrange + var context1 = new DefaultHttpContext(); + var context2 = new DefaultHttpContext(); + + RequestDelegate next = (HttpContext ctx) => Task.CompletedTask; + var middleware = new McpCorrelationIdMiddleware(next); + + // Act + await middleware.InvokeAsync(context1); + await middleware.InvokeAsync(context2); + + // Assert + var correlationId1 = context1.Items["CorrelationId"] as string; + var correlationId2 = context2.Items["CorrelationId"] as string; + + correlationId1.Should().NotBeNullOrWhiteSpace(); + correlationId2.Should().NotBeNullOrWhiteSpace(); + correlationId1.Should().NotBe(correlationId2); + } +} diff --git a/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpExceptionHandlerMiddlewareTests.cs b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpExceptionHandlerMiddlewareTests.cs new file mode 100644 index 0000000..f66026d --- /dev/null +++ b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpExceptionHandlerMiddlewareTests.cs @@ -0,0 +1,196 @@ +using ColaFlow.Modules.Mcp.Contracts.JsonRpc; +using ColaFlow.Modules.Mcp.Domain.Exceptions; +using ColaFlow.Modules.Mcp.Infrastructure.Middleware; +using FluentAssertions; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using NSubstitute; +using System.Text.Json; + +namespace ColaFlow.Modules.Mcp.Tests.Infrastructure.Middleware; + +/// +/// Unit tests for McpExceptionHandlerMiddleware +/// +public class McpExceptionHandlerMiddlewareTests +{ + private readonly ILogger _logger; + + public McpExceptionHandlerMiddlewareTests() + { + _logger = Substitute.For>(); + } + + [Fact] + public async Task InvokeAsync_ShouldCallNextMiddleware_WhenNoExceptionThrown() + { + // Arrange + var context = new DefaultHttpContext(); + var nextCalled = false; + + RequestDelegate next = (HttpContext ctx) => + { + nextCalled = true; + return Task.CompletedTask; + }; + + var middleware = new McpExceptionHandlerMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + nextCalled.Should().BeTrue(); + } + + [Fact] + public async Task InvokeAsync_ShouldHandleMcpException_AndReturnJsonRpcErrorResponse() + { + // Arrange + var context = new DefaultHttpContext(); + context.Response.Body = new MemoryStream(); + context.Items["CorrelationId"] = "test-correlation-id"; + context.Items["McpRequestId"] = "test-request-id"; + + var expectedException = new McpNotFoundException("Task", "task-123"); + + RequestDelegate next = (HttpContext ctx) => + { + throw expectedException; + }; + + var middleware = new McpExceptionHandlerMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + context.Response.StatusCode.Should().Be(404); + context.Response.ContentType.Should().Be("application/json"); + + context.Response.Body.Seek(0, SeekOrigin.Begin); + var responseBody = await new StreamReader(context.Response.Body).ReadToEndAsync(); + var response = JsonSerializer.Deserialize(responseBody, new JsonSerializerOptions + { + PropertyNameCaseInsensitive = true + }); + + response.Should().NotBeNull(); + response!.JsonRpc.Should().Be("2.0"); + response.Error.Should().NotBeNull(); + response.Error!.Code.Should().Be((int)JsonRpcErrorCode.NotFound); + response.Error.Message.Should().Be("Task not found: task-123"); + response.Id.Should().NotBeNull(); + response.Id!.ToString().Should().Be("test-request-id"); + } + + [Fact] + public async Task InvokeAsync_ShouldHandleUnexpectedException_AndReturnInternalError() + { + // Arrange + var context = new DefaultHttpContext(); + context.Response.Body = new MemoryStream(); + context.Items["CorrelationId"] = "test-correlation-id"; + + RequestDelegate next = (HttpContext ctx) => + { + throw new InvalidOperationException("Unexpected error"); + }; + + var middleware = new McpExceptionHandlerMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + context.Response.StatusCode.Should().Be(500); + context.Response.ContentType.Should().Be("application/json"); + + context.Response.Body.Seek(0, SeekOrigin.Begin); + var responseBody = await new StreamReader(context.Response.Body).ReadToEndAsync(); + var response = JsonSerializer.Deserialize(responseBody, new JsonSerializerOptions + { + PropertyNameCaseInsensitive = true + }); + + response.Should().NotBeNull(); + response!.Error.Should().NotBeNull(); + response.Error!.Code.Should().Be((int)JsonRpcErrorCode.InternalError); + response.Error.Message.Should().Be("Internal server error"); + // Should NOT expose exception details + response.Error.Data.Should().BeNull(); + } + + [Theory] + [InlineData(typeof(McpUnauthorizedException), 401)] + [InlineData(typeof(McpForbiddenException), 403)] + [InlineData(typeof(McpNotFoundException), 404)] + [InlineData(typeof(McpValidationException), 422)] + [InlineData(typeof(McpParseException), 400)] + [InlineData(typeof(McpInvalidRequestException), 400)] + [InlineData(typeof(McpMethodNotFoundException), 404)] + [InlineData(typeof(McpInvalidParamsException), 400)] + public async Task InvokeAsync_ShouldMapExceptionToCorrectHttpStatusCode(Type exceptionType, int expectedStatusCode) + { + // Arrange + var context = new DefaultHttpContext(); + context.Response.Body = new MemoryStream(); + + McpException exception = exceptionType.Name switch + { + nameof(McpUnauthorizedException) => new McpUnauthorizedException(), + nameof(McpForbiddenException) => new McpForbiddenException(), + nameof(McpNotFoundException) => new McpNotFoundException("Resource", "123"), + nameof(McpValidationException) => new McpValidationException(), + nameof(McpParseException) => new McpParseException(), + nameof(McpInvalidRequestException) => new McpInvalidRequestException(), + nameof(McpMethodNotFoundException) => new McpMethodNotFoundException("test"), + nameof(McpInvalidParamsException) => new McpInvalidParamsException(), + _ => throw new ArgumentException("Unknown exception type") + }; + + RequestDelegate next = (HttpContext ctx) => + { + throw exception; + }; + + var middleware = new McpExceptionHandlerMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + context.Response.StatusCode.Should().Be(expectedStatusCode); + } + + [Fact] + public async Task InvokeAsync_ShouldLogErrorWithStructuredData() + { + // Arrange + var context = new DefaultHttpContext(); + context.Response.Body = new MemoryStream(); + context.Items["CorrelationId"] = "test-correlation-id"; + context.Items["TenantId"] = "tenant-123"; + context.Items["ApiKeyId"] = "key-456"; + + var exception = new McpValidationException("Test validation error"); + + RequestDelegate next = (HttpContext ctx) => + { + throw exception; + }; + + var middleware = new McpExceptionHandlerMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + _logger.Received(1).Log( + LogLevel.Error, + Arg.Any(), + Arg.Is(o => o.ToString()!.Contains("MCP Error")), + exception, + Arg.Any>()); + } +} diff --git a/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpLoggingMiddlewareTests.cs b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpLoggingMiddlewareTests.cs new file mode 100644 index 0000000..e318804 --- /dev/null +++ b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Infrastructure/Middleware/McpLoggingMiddlewareTests.cs @@ -0,0 +1,241 @@ +using ColaFlow.Modules.Mcp.Infrastructure.Middleware; +using FluentAssertions; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Logging; +using NSubstitute; +using System.Text; + +namespace ColaFlow.Modules.Mcp.Tests.Infrastructure.Middleware; + +/// +/// Unit tests for McpLoggingMiddleware +/// +public class McpLoggingMiddlewareTests +{ + private readonly ILogger _logger; + + public McpLoggingMiddlewareTests() + { + _logger = Substitute.For>(); + } + + [Fact] + public async Task InvokeAsync_ShouldSkipLogging_ForNonMcpRequests() + { + // Arrange + var context = new DefaultHttpContext(); + context.Request.Method = "GET"; + context.Request.Path = "/api/tasks"; + + var nextCalled = false; + RequestDelegate next = (HttpContext ctx) => + { + nextCalled = true; + return Task.CompletedTask; + }; + + var middleware = new McpLoggingMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + nextCalled.Should().BeTrue(); + // Logger should not be called for non-MCP requests + _logger.DidNotReceive().Log( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any>()); + } + + [Fact] + public async Task InvokeAsync_ShouldLogRequestAndResponse_ForMcpRequests() + { + // Arrange + var context = new DefaultHttpContext(); + context.Request.Method = "POST"; + context.Request.Path = "/mcp"; + context.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes("{\"jsonrpc\":\"2.0\",\"method\":\"initialize\"}")); + context.Response.Body = new MemoryStream(); + context.Items["CorrelationId"] = "test-correlation-id"; + + RequestDelegate next = (HttpContext ctx) => + { + ctx.Response.StatusCode = 200; + var responseBytes = Encoding.UTF8.GetBytes("{\"jsonrpc\":\"2.0\",\"result\":{}}"); + ctx.Response.Body.Write(responseBytes, 0, responseBytes.Length); + return Task.CompletedTask; + }; + + var middleware = new McpLoggingMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + // Should log request (Debug level) + _logger.Received().Log( + LogLevel.Debug, + Arg.Any(), + Arg.Is(o => o.ToString()!.Contains("MCP Request")), + Arg.Any(), + Arg.Any>()); + + // Should log response (Debug level for 2xx status) + _logger.Received().Log( + LogLevel.Debug, + Arg.Any(), + Arg.Is(o => o.ToString()!.Contains("MCP Response")), + Arg.Any(), + Arg.Any>()); + } + + [Fact] + public async Task InvokeAsync_ShouldLogErrorLevel_ForErrorResponses() + { + // Arrange + var context = new DefaultHttpContext(); + context.Request.Method = "POST"; + context.Request.Path = "/mcp"; + context.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes("{}")); + context.Response.Body = new MemoryStream(); + context.Items["CorrelationId"] = "test-correlation-id"; + + RequestDelegate next = (HttpContext ctx) => + { + ctx.Response.StatusCode = 500; // Error status + return Task.CompletedTask; + }; + + var middleware = new McpLoggingMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + // Should log response at Error level for 5xx status + _logger.Received().Log( + LogLevel.Error, + Arg.Any(), + Arg.Is(o => o.ToString()!.Contains("MCP Response")), + Arg.Any(), + Arg.Any>()); + } + + [Fact] + public async Task InvokeAsync_ShouldLogWarning_ForSlowRequests() + { + // Arrange + var context = new DefaultHttpContext(); + context.Request.Method = "POST"; + context.Request.Path = "/mcp"; + context.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes("{}")); + context.Response.Body = new MemoryStream(); + context.Items["CorrelationId"] = "test-correlation-id"; + + RequestDelegate next = async (HttpContext ctx) => + { + // Simulate slow request (> 1 second) + await Task.Delay(1100); + ctx.Response.StatusCode = 200; + }; + + var middleware = new McpLoggingMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + // Should log warning for slow requests + _logger.Received().Log( + LogLevel.Warning, + Arg.Any(), + Arg.Is(o => o.ToString()!.Contains("Slow MCP Request")), + Arg.Any(), + Arg.Any>()); + } + + [Fact] + public async Task InvokeAsync_ShouldSanitizeSensitiveData() + { + // Arrange + var requestBody = "{\"keyHash\":\"secret-key-hash\",\"password\":\"my-password\"}"; + var context = new DefaultHttpContext(); + context.Request.Method = "POST"; + context.Request.Path = "/mcp"; + context.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes(requestBody)); + context.Response.Body = new MemoryStream(); + context.Items["CorrelationId"] = "test-correlation-id"; + + var loggedRequest = string.Empty; + _logger.When(x => x.Log( + LogLevel.Debug, + Arg.Any(), + Arg.Is(o => o.ToString()!.Contains("MCP Request")), + Arg.Any(), + Arg.Any>())) + .Do(callInfo => + { + var state = callInfo.ArgAt(2); + loggedRequest = state.ToString() ?? ""; + }); + + RequestDelegate next = (HttpContext ctx) => + { + ctx.Response.StatusCode = 200; + return Task.CompletedTask; + }; + + var middleware = new McpLoggingMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + loggedRequest.Should().Contain("[REDACTED]"); + loggedRequest.Should().NotContain("secret-key-hash"); + loggedRequest.Should().NotContain("my-password"); + } + + [Fact] + public async Task InvokeAsync_ShouldIncludePerformanceMetrics() + { + // Arrange + var context = new DefaultHttpContext(); + context.Request.Method = "POST"; + context.Request.Path = "/mcp"; + context.Request.Body = new MemoryStream(Encoding.UTF8.GetBytes("{}")); + context.Response.Body = new MemoryStream(); + context.Items["CorrelationId"] = "test-correlation-id"; + + var loggedResponse = string.Empty; + _logger.When(x => x.Log( + LogLevel.Debug, + Arg.Any(), + Arg.Is(o => o.ToString()!.Contains("MCP Response")), + Arg.Any(), + Arg.Any>())) + .Do(callInfo => + { + var state = callInfo.ArgAt(2); + loggedResponse = state.ToString() ?? ""; + }); + + RequestDelegate next = (HttpContext ctx) => + { + ctx.Response.StatusCode = 200; + return Task.CompletedTask; + }; + + var middleware = new McpLoggingMiddleware(next, _logger); + + // Act + await middleware.InvokeAsync(context); + + // Assert + loggedResponse.Should().Contain("Duration:"); + loggedResponse.Should().MatchRegex(@"Duration:\s*\d+ms"); + } +}