diff --git a/colaflow-api/src/ColaFlow.API/Program.cs b/colaflow-api/src/ColaFlow.API/Program.cs index 16cb7ad..f01e7cf 100644 --- a/colaflow-api/src/ColaFlow.API/Program.cs +++ b/colaflow-api/src/ColaFlow.API/Program.cs @@ -7,6 +7,7 @@ using ColaFlow.Modules.Identity.Application; using ColaFlow.Modules.Identity.Infrastructure; using ColaFlow.Modules.Identity.Infrastructure.Persistence; using ColaFlow.Modules.Mcp.Infrastructure.Extensions; +using ColaFlow.Modules.Mcp.Infrastructure.Hubs; using ColaFlow.Modules.ProjectManagement.Infrastructure.Persistence; using Microsoft.AspNetCore.Authentication.JwtBearer; using Microsoft.EntityFrameworkCore; diff --git a/colaflow-api/src/ColaFlow.API/Hubs/McpNotificationHub.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Hubs/McpNotificationHub.cs similarity index 59% rename from colaflow-api/src/ColaFlow.API/Hubs/McpNotificationHub.cs rename to colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Hubs/McpNotificationHub.cs index 593ee57..3901163 100644 --- a/colaflow-api/src/ColaFlow.API/Hubs/McpNotificationHub.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Hubs/McpNotificationHub.cs @@ -1,14 +1,15 @@ using Microsoft.AspNetCore.Authorization; using Microsoft.AspNetCore.SignalR; +using Microsoft.Extensions.Logging; -namespace ColaFlow.API.Hubs; +namespace ColaFlow.Modules.Mcp.Infrastructure.Hubs; /// /// SignalR Hub for MCP real-time notifications /// Supports notifying AI agents and users about PendingChange status updates /// [Authorize] -public class McpNotificationHub : BaseHub +public class McpNotificationHub : Hub { private readonly ILogger _logger; @@ -19,15 +20,26 @@ public class McpNotificationHub : BaseHub public override async Task OnConnectedAsync() { - var connectionId = Context.ConnectionId; - var userId = GetCurrentUserId(); - var tenantId = GetCurrentTenantId(); + try + { + var connectionId = Context.ConnectionId; + var userId = GetCurrentUserId(); + var tenantId = GetCurrentTenantId(); - _logger.LogInformation( - "MCP client connected - ConnectionId={ConnectionId}, UserId={UserId}, TenantId={TenantId}", - connectionId, userId, tenantId); + // Automatically join tenant group (tenant isolation) + await Groups.AddToGroupAsync(Context.ConnectionId, GetTenantGroupName(tenantId)); - await base.OnConnectedAsync(); + _logger.LogInformation( + "MCP client connected - ConnectionId={ConnectionId}, UserId={UserId}, TenantId={TenantId}", + connectionId, userId, tenantId); + + await base.OnConnectedAsync(); + } + catch (Exception ex) + { + _logger.LogError(ex, "MCP client connection error"); + Context.Abort(); + } } public override async Task OnDisconnectedAsync(Exception? exception) @@ -78,11 +90,37 @@ public class McpNotificationHub : BaseHub Context.ConnectionId, groupName, pendingChangeId); } - /// - /// Get the SignalR group name for a pending change - /// - /// The pending change ID - /// The group name + // Helper methods (copied from BaseHub to avoid cross-layer dependency) + protected Guid GetCurrentUserId() + { + var userIdClaim = Context.User?.FindFirst("sub") + ?? Context.User?.FindFirst("user_id"); + + if (userIdClaim == null || !Guid.TryParse(userIdClaim.Value, out var userId)) + { + throw new UnauthorizedAccessException("User ID not found in token"); + } + + return userId; + } + + protected Guid GetCurrentTenantId() + { + var tenantIdClaim = Context.User?.FindFirst("tenant_id"); + + if (tenantIdClaim == null || !Guid.TryParse(tenantIdClaim.Value, out var tenantId)) + { + throw new UnauthorizedAccessException("Tenant ID not found in token"); + } + + return tenantId; + } + + protected string GetTenantGroupName(Guid tenantId) + { + return $"tenant-{tenantId}"; + } + private static string GetPendingChangeGroupName(Guid pendingChangeId) { return $"pending-change-{pendingChangeId}"; diff --git a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Services/McpNotificationService.cs b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Services/McpNotificationService.cs index 0e9527c..0b7a71f 100644 --- a/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Services/McpNotificationService.cs +++ b/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Infrastructure/Services/McpNotificationService.cs @@ -1,6 +1,6 @@ -using ColaFlow.API.Hubs; using ColaFlow.Modules.Mcp.Application.DTOs.Notifications; using ColaFlow.Modules.Mcp.Application.Services; +using ColaFlow.Modules.Mcp.Infrastructure.Hubs; using Microsoft.AspNetCore.SignalR; using Microsoft.Extensions.Logging; diff --git a/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Services/McpNotificationServiceTests.cs b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Services/McpNotificationServiceTests.cs index 5473bb7..8e77157 100644 --- a/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Services/McpNotificationServiceTests.cs +++ b/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Services/McpNotificationServiceTests.cs @@ -1,5 +1,5 @@ -using ColaFlow.API.Hubs; using ColaFlow.Modules.Mcp.Application.DTOs.Notifications; +using ColaFlow.Modules.Mcp.Infrastructure.Hubs; using ColaFlow.Modules.Mcp.Infrastructure.Services; using Microsoft.AspNetCore.SignalR; using Microsoft.Extensions.Logging;