diff --git a/samples/ProtectedMcpServer/Program.cs b/samples/ProtectedMcpServer/Program.cs index 97f8456a2..df81bb254 100644 --- a/samples/ProtectedMcpServer/Program.cs +++ b/samples/ProtectedMcpServer/Program.cs @@ -1,14 +1,18 @@ using Microsoft.AspNetCore.Authentication.JwtBearer; using Microsoft.IdentityModel.Tokens; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; using ModelContextProtocol.AspNetCore.Authentication; using ProtectedMcpServer.Tools; using System.Net.Http.Headers; using System.Security.Claims; +using System.Threading.RateLimiting; var builder = WebApplication.CreateBuilder(args); var serverUrl = "http://localhost:7071/"; var inMemoryOAuthServerUrl = "https://localhost:7029"; +const int maxConcurrentToolCallsPerUser = 10; builder.Services.AddAuthentication(options => { @@ -65,7 +69,36 @@ builder.Services.AddAuthorization(); builder.Services.AddHttpContextAccessor(); +builder.Services.AddSingleton>>(_ => + PartitionedRateLimiter.Create, string>(context => + RateLimitPartition.GetConcurrencyLimiter( + context.User?.FindFirstValue(ClaimTypes.NameIdentifier) ?? + context.User?.Identity?.Name ?? + context.JsonRpcMessage.Context?.RelatedTransport?.SessionId ?? + "anonymous", + _ => new ConcurrencyLimiterOptions + { + PermitLimit = maxConcurrentToolCallsPerUser, + QueueProcessingOrder = QueueProcessingOrder.OldestFirst, + QueueLimit = 0 + }))); builder.Services.AddMcpServer() + .WithRequestFilters(filters => filters.AddCallToolFilter(next => async (request, cancellationToken) => + { + var services = request.Services ?? throw new InvalidOperationException("Request context does not have an associated service provider."); + var toolCallConcurrencyLimiter = services.GetRequiredService>>(); + using var lease = await toolCallConcurrencyLimiter.AcquireAsync(request, 1, cancellationToken).ConfigureAwait(false); + if (!lease.IsAcquired) + { + return new CallToolResult + { + IsError = true, + Content = [new TextContentBlock { Text = $"Maximum concurrent tool calls ({maxConcurrentToolCallsPerUser}) exceeded. Try again after in-progress calls complete." }] + }; + } + + return await next(request, cancellationToken).ConfigureAwait(false); + })) .WithTools() .WithHttpTransport(); diff --git a/samples/ProtectedMcpServer/README.md b/samples/ProtectedMcpServer/README.md index ecbfee633..0db0a7437 100644 --- a/samples/ProtectedMcpServer/README.md +++ b/samples/ProtectedMcpServer/README.md @@ -7,6 +7,7 @@ This sample demonstrates how to create an MCP server that requires OAuth 2.0 aut The Protected MCP Server sample shows how to: - Create an MCP server with OAuth 2.0 protection - Configure JWT bearer token authentication +- Apply a per-user concurrent tool-call limit (10 in-flight calls) with a call-tool request filter - Implement protected MCP tools and resources - Integrate with ASP.NET Core authentication and authorization - Provide OAuth resource metadata for client discovery @@ -122,4 +123,4 @@ The weather tools use the National Weather Service API at `api.weather.gov` to f - `Program.cs`: Server setup with authentication and MCP configuration - `Tools/WeatherTools.cs`: Weather tool implementations - `Tools/HttpClientExt.cs`: HTTP client extensions -- `Properties/launchSettings.json`: Development launch configuration \ No newline at end of file +- `Properties/launchSettings.json`: Development launch configuration