From db40e4f50f25b465b1938ac01a0c31cf2827d976 Mon Sep 17 00:00:00 2001 From: thomasluizon Date: Fri, 7 Aug 2026 12:00:55 -0300 Subject: [PATCH 1/7] chore: initialize ORB-8 review From 7cfd4d1bd3bbc4fa1fd357cfeec1b34f230f84cb Mon Sep 17 00:00:00 2001 From: thomasluizon Date: Fri, 7 Aug 2026 12:05:13 -0300 Subject: [PATCH 2/7] fix: rate limit MCP calls by API key --- .../Extensions/WebApplicationExtensions.cs | 130 ++++++++++++++- .../Services/DistributedRateLimitService.cs | 4 +- .../WebApplicationExtensionsMcpTests.cs | 148 ++++++++++++++++++ .../DistributedRateLimitServiceTests.cs | 15 ++ 4 files changed, 295 insertions(+), 2 deletions(-) diff --git a/src/Orbit.Api/Extensions/WebApplicationExtensions.cs b/src/Orbit.Api/Extensions/WebApplicationExtensions.cs index f15433d2..605c1744 100644 --- a/src/Orbit.Api/Extensions/WebApplicationExtensions.cs +++ b/src/Orbit.Api/Extensions/WebApplicationExtensions.cs @@ -1,4 +1,5 @@ using System.Net; +using System.Security.Claims; using System.Text.Json; using Microsoft.AspNetCore.Authentication; using Microsoft.AspNetCore.Diagnostics.HealthChecks; @@ -18,6 +19,12 @@ namespace Orbit.Api.Extensions; public static partial class WebApplicationExtensions { + internal const string McpRateLimitPolicy = "mcp"; + internal const string McpAiRateLimitPolicy = "mcp-ai"; + + private static readonly HashSet AiBearingMcpTools = + new(StringComparer.OrdinalIgnoreCase) { "get_daily_summary", "get_retrospective" }; + public static async Task ConfigureOrbitPipeline(this WebApplication app) { if (!BuildTimeDocumentGeneration.IsActive) @@ -124,7 +131,17 @@ private static async Task HandleMcpRequestAsync(HttpContext context, Func if (!await TryAuthenticateMcpRequestAsync(context)) return; - if (!TryGetMcpToolCall(root, out var toolName, out var requestId, out var operationId, out var operationFingerprint)) + var isToolCall = TryGetMcpToolCall( + root, + out var toolName, + out var requestId, + out var operationId, + out var operationFingerprint); + + if (!await TryApplyMcpRateLimitsAsync(context, toolName, requestId)) + return; + + if (!isToolCall) { await next(); return; @@ -137,6 +154,55 @@ await HandleMcpToolCallAsync( new McpToolCallRequest(toolName!, requestId, operationId, operationFingerprint)); } + internal static async Task TryApplyMcpRateLimitsAsync( + HttpContext context, + string? toolName, + JsonElement? requestId) + { + if (context.User.Identity?.IsAuthenticated != true) + return true; + + var apiKeyId = context.User.FindFirstValue("api_key_id"); + if (string.IsNullOrWhiteSpace(apiKeyId)) + throw new InvalidOperationException("Authenticated MCP principal is missing the api_key_id claim."); + + var partitionKey = $"api-key:{apiKeyId}"; + var service = context.RequestServices.GetRequiredService(); + + var decision = await service.TryAcquireAsync( + McpRateLimitPolicy, + partitionKey, + context.RequestAborted); + if (!decision.Allowed) + { + await WriteMcpRateLimitErrorAsync( + context, + requestId, + McpRateLimitPolicy, + partitionKey, + decision); + return false; + } + + if (toolName is null || !AiBearingMcpTools.Contains(toolName)) + return true; + + var aiDecision = await service.TryAcquireAsync( + McpAiRateLimitPolicy, + partitionKey, + context.RequestAborted); + if (aiDecision.Allowed) + return true; + + await WriteMcpRateLimitErrorAsync( + context, + requestId, + McpAiRateLimitPolicy, + partitionKey, + aiDecision); + return false; + } + internal static JsonDocument? TryParseMcpBody(string body) { try @@ -434,6 +500,55 @@ await context.Response.WriteAsJsonAsync(new }); } + private static async Task WriteMcpRateLimitErrorAsync( + HttpContext context, + JsonElement? requestId, + string policyName, + string partitionKey, + DistributedRateLimitDecision decision) + { + var now = context.RequestServices.GetRequiredService().GetUtcNow().UtcDateTime; + var retryAfterSeconds = Math.Max( + 1, + (int)Math.Ceiling((decision.WindowEndsAtUtc - now).TotalSeconds)); + + context.Response.StatusCode = StatusCodes.Status429TooManyRequests; + context.Response.ContentType = "application/json"; + context.Response.Headers.RetryAfter = retryAfterSeconds.ToString(); + context.Response.Headers[HttpContextExtensions.RequestIdHeaderName] = context.GetRequestId(); + + var logger = context.RequestServices + .GetRequiredService() + .CreateLogger(typeof(WebApplicationExtensions)); + LogMcpRateLimitRejected( + logger, + policyName, + partitionKey, + decision.CurrentCount, + decision.PermitLimit, + retryAfterSeconds, + context.GetRequestId()); + + await context.Response.WriteAsJsonAsync(new + { + jsonrpc = "2.0", + id = requestId, + error = new + { + code = -32002, + message = "rate_limit_exceeded", + data = new + { + reason = "rate_limit_exceeded", + policy = policyName, + limit = decision.PermitLimit, + count = decision.CurrentCount, + retryAfterUtc = decision.WindowEndsAtUtc + } + } + }); + } + private static async Task TryAuditLegacyMcpAsync( IAgentAuditService auditService, HttpContext context, @@ -474,6 +589,19 @@ await auditService.RecordAsync(new AgentAuditEntry( [LoggerMessage(EventId = 1, Level = LogLevel.Error, Message = "Failed to write MCP audit entry for {SourceName}. TraceId={TraceId}")] private static partial void LogLegacyMcpAuditWriteFailed(ILogger logger, Exception ex, string sourceName, string traceId); + [LoggerMessage( + EventId = 2, + Level = LogLevel.Warning, + Message = "MCP rate limit rejected. Policy={PolicyName} PartitionKey={PartitionKey} Count={CurrentCount}/{PermitLimit} RetryAfterSeconds={RetryAfterSeconds} RequestId={RequestId}")] + private static partial void LogMcpRateLimitRejected( + ILogger logger, + string policyName, + string partitionKey, + int currentCount, + int permitLimit, + int retryAfterSeconds, + string requestId); + private sealed record McpToolCallRequest( string ToolName, JsonElement? RequestId, diff --git a/src/Orbit.Infrastructure/Services/DistributedRateLimitService.cs b/src/Orbit.Infrastructure/Services/DistributedRateLimitService.cs index 05928ded..f67f3719 100644 --- a/src/Orbit.Infrastructure/Services/DistributedRateLimitService.cs +++ b/src/Orbit.Infrastructure/Services/DistributedRateLimitService.cs @@ -39,7 +39,9 @@ public class DistributedRateLimitService(OrbitDbContext dbContext, TimeProvider ["challenges"] = new(TimeSpan.FromHours(24), PermitLimit: 50, SegmentCount: 1), ["public-profile"] = new(TimeSpan.FromMinutes(1), PermitLimit: 30, SegmentCount: 4), ["admin-broadcast"] = new(TimeSpan.FromHours(1), PermitLimit: 5, SegmentCount: 1), - ["marketing-unsubscribe"] = new(TimeSpan.FromMinutes(1), PermitLimit: 20, SegmentCount: 4) + ["marketing-unsubscribe"] = new(TimeSpan.FromMinutes(1), PermitLimit: 20, SegmentCount: 4), + ["mcp"] = new(TimeSpan.FromMinutes(1), PermitLimit: 60, SegmentCount: 4), + ["mcp-ai"] = new(TimeSpan.FromMinutes(1), PermitLimit: 15, SegmentCount: 4) }; public async Task TryAcquireAsync( diff --git a/tests/Orbit.Infrastructure.Tests/Extensions/WebApplicationExtensionsMcpTests.cs b/tests/Orbit.Infrastructure.Tests/Extensions/WebApplicationExtensionsMcpTests.cs index a0d3c025..a31a9280 100644 --- a/tests/Orbit.Infrastructure.Tests/Extensions/WebApplicationExtensionsMcpTests.cs +++ b/tests/Orbit.Infrastructure.Tests/Extensions/WebApplicationExtensionsMcpTests.cs @@ -1,6 +1,14 @@ +using System.Security.Claims; using System.Text.Json; using FluentAssertions; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Logging.Abstractions; +using NSubstitute; using Orbit.Api.Extensions; +using Orbit.Domain.Interfaces; +using Orbit.Domain.Models; namespace Orbit.Infrastructure.Tests.Extensions; @@ -142,4 +150,144 @@ public void TryGetMcpToolCall_MalformedJson_ReturnsFalse() WebApplicationExtensions.TryGetMcpToolCall( document?.RootElement, out _, out _, out _, out _).Should().BeFalse(); } + + [Fact] + public async Task TryApplyMcpRateLimitsAsync_ExhaustedKey_DoesNotAffectSecondKey() + { + var service = Substitute.For(); + var counts = new Dictionary(StringComparer.Ordinal); + service.TryAcquireAsync( + WebApplicationExtensions.McpRateLimitPolicy, + Arg.Any(), + Arg.Any()) + .Returns(call => + { + var partitionKey = call.ArgAt(1); + counts.TryGetValue(partitionKey, out var count); + count++; + counts[partitionKey] = count; + return new DistributedRateLimitDecision( + count <= 1, + 1, + Math.Min(count, 1), + DateTime.UtcNow.AddMinutes(1)); + }); + + var firstKey = Guid.NewGuid(); + var secondKey = Guid.NewGuid(); + + (await WebApplicationExtensions.TryApplyMcpRateLimitsAsync( + CreateRateLimitContext(service, firstKey), + "get_habits", + requestId: null)).Should().BeTrue(); + (await WebApplicationExtensions.TryApplyMcpRateLimitsAsync( + CreateRateLimitContext(service, firstKey), + "get_habits", + requestId: null)).Should().BeFalse(); + (await WebApplicationExtensions.TryApplyMcpRateLimitsAsync( + CreateRateLimitContext(service, secondKey), + "get_habits", + requestId: null)).Should().BeTrue(); + } + + [Fact] + public async Task TryApplyMcpRateLimitsAsync_DailySummary_HitsAiLimitBeforeGeneralLimit() + { + var service = Substitute.For(); + var aiCount = 0; + service.TryAcquireAsync( + WebApplicationExtensions.McpRateLimitPolicy, + Arg.Any(), + Arg.Any()) + .Returns(new DistributedRateLimitDecision(true, 60, 1, DateTime.UtcNow.AddMinutes(1))); + service.TryAcquireAsync( + WebApplicationExtensions.McpAiRateLimitPolicy, + Arg.Any(), + Arg.Any()) + .Returns(_ => + { + aiCount++; + return new DistributedRateLimitDecision( + aiCount <= 1, + 1, + Math.Min(aiCount, 1), + DateTime.UtcNow.AddMinutes(1)); + }); + + var apiKeyId = Guid.NewGuid(); + var allowedContext = CreateRateLimitContext(service, apiKeyId); + var rejectedContext = CreateRateLimitContext(service, apiKeyId); + using var requestIdDocument = JsonDocument.Parse("42"); + + (await WebApplicationExtensions.TryApplyMcpRateLimitsAsync( + allowedContext, + "get_daily_summary", + requestIdDocument.RootElement.Clone())).Should().BeTrue(); + (await WebApplicationExtensions.TryApplyMcpRateLimitsAsync( + rejectedContext, + "get_daily_summary", + requestIdDocument.RootElement.Clone())).Should().BeFalse(); + + rejectedContext.Response.StatusCode.Should().Be(StatusCodes.Status429TooManyRequests); + rejectedContext.Response.Headers.RetryAfter.Should().NotBeEmpty(); + rejectedContext.Response.Body.Position = 0; + using var responseDocument = await JsonDocument.ParseAsync(rejectedContext.Response.Body); + var error = responseDocument.RootElement.GetProperty("error"); + error.GetProperty("message").GetString().Should().Be("rate_limit_exceeded"); + error.GetProperty("data").GetProperty("policy").GetString() + .Should().Be(WebApplicationExtensions.McpAiRateLimitPolicy); + + await service.Received(2).TryAcquireAsync( + WebApplicationExtensions.McpRateLimitPolicy, + $"api-key:{apiKeyId}", + Arg.Any()); + await service.Received(2).TryAcquireAsync( + WebApplicationExtensions.McpAiRateLimitPolicy, + $"api-key:{apiKeyId}", + Arg.Any()); + } + + [Fact] + public async Task TryApplyMcpRateLimitsAsync_UnauthenticatedRequest_DoesNotCallLimiter() + { + var service = Substitute.For(); + var context = CreateRateLimitContext(service, apiKeyId: null); + + var allowed = await WebApplicationExtensions.TryApplyMcpRateLimitsAsync( + context, + "get_daily_summary", + requestId: null); + + allowed.Should().BeTrue(); + await service.DidNotReceive().TryAcquireAsync( + Arg.Any(), + Arg.Any(), + Arg.Any()); + } + + private static DefaultHttpContext CreateRateLimitContext( + IDistributedRateLimitService service, + Guid? apiKeyId) + { + var services = new ServiceCollection() + .AddSingleton(service) + .AddSingleton(TimeProvider.System) + .AddSingleton(NullLoggerFactory.Instance) + .BuildServiceProvider(); + + var context = new DefaultHttpContext + { + RequestServices = services, + Response = { Body = new MemoryStream() } + }; + + if (apiKeyId.HasValue) + { + context.User = new ClaimsPrincipal(new ClaimsIdentity( + [new Claim("api_key_id", apiKeyId.Value.ToString())], + "ApiKey")); + } + + return context; + } } diff --git a/tests/Orbit.Infrastructure.Tests/Services/DistributedRateLimitServiceTests.cs b/tests/Orbit.Infrastructure.Tests/Services/DistributedRateLimitServiceTests.cs index d706a521..d4c90505 100644 --- a/tests/Orbit.Infrastructure.Tests/Services/DistributedRateLimitServiceTests.cs +++ b/tests/Orbit.Infrastructure.Tests/Services/DistributedRateLimitServiceTests.cs @@ -132,6 +132,21 @@ public async Task TryAcquireAsync_TagsPolicy_PartitionsPerUser() otherUser.CurrentCount.Should().Be(1); } + [Theory] + [InlineData("mcp", 60)] + [InlineData("mcp-ai", 15)] + public async Task TryAcquireAsync_McpPolicies_BlockAfterPermitLimit(string policyName, int permitLimit) + { + DistributedRateLimitDecision finalDecision = new(true, 0, 0, DateTime.UtcNow); + + for (var attempt = 0; attempt <= permitLimit; attempt++) + finalDecision = await _service.TryAcquireAsync(policyName, "api-key:one"); + + finalDecision.Allowed.Should().BeFalse(); + finalDecision.PermitLimit.Should().Be(permitLimit); + finalDecision.CurrentCount.Should().Be(permitLimit); + } + [Fact] public async Task TryAcquireAsync_RelationalProvider_RetriesSerializationConflictThenSucceeds() { From 046bada49eac24c017bf167e006c73f8f3c3cac8 Mon Sep 17 00:00:00 2001 From: thomasluizon Date: Fri, 7 Aug 2026 12:13:19 -0300 Subject: [PATCH 3/7] fix: honor MCP request cancellation --- src/Orbit.Api/Extensions/WebApplicationExtensions.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Orbit.Api/Extensions/WebApplicationExtensions.cs b/src/Orbit.Api/Extensions/WebApplicationExtensions.cs index 605c1744..f2e83fb0 100644 --- a/src/Orbit.Api/Extensions/WebApplicationExtensions.cs +++ b/src/Orbit.Api/Extensions/WebApplicationExtensions.cs @@ -546,7 +546,7 @@ await context.Response.WriteAsJsonAsync(new retryAfterUtc = decision.WindowEndsAtUtc } } - }); + }, context.RequestAborted); } private static async Task TryAuditLegacyMcpAsync( From 23931850ca93320e235789946090a4eaae389f51 Mon Sep 17 00:00:00 2001 From: thomasluizon Date: Fri, 7 Aug 2026 12:22:31 -0300 Subject: [PATCH 4/7] chore: refresh architecture map --- architecture.html | 2 +- architecture.json | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/architecture.html b/architecture.html index 2033cc4b..26ac4e28 100644 --- a/architecture.html +++ b/architecture.html @@ -47,7 +47,7 @@

Handlers with no endpoint

RequestHandler file

Entities

EntityDomain file
- +