Add backpressure when rapidly creating new stateful Streamable HTTP sessions

Author: halter73

src/ModelContextProtocol.AspNetCore/HttpMcpSession.cs

@@ -1,20 +1,27 @@
 using ModelContextProtocol.AspNetCore.Stateless;
 using ModelContextProtocol.Protocol;
 using ModelContextProtocol.Server;
+using System.Diagnostics;
 using System.Security.Claims;
+using System.Threading;
 
 namespace ModelContextProtocol.AspNetCore;
 
 internal sealed class HttpMcpSession<TTransport>(
     string sessionId,
     TTransport transport,
     UserIdClaim? userId,
-    TimeProvider timeProvider) : IAsyncDisposable
+    TimeProvider timeProvider,
+    SemaphoreSlim? idleSessionSemaphore = null) : IAsyncDisposable
     where TTransport : ITransport
 {
     private int _referenceCount;
     private int _getRequestStarted;
-    private CancellationTokenSource _disposeCts = new();
+    private bool _isDisposed;
+
+    private readonly SemaphoreSlim? _idleSessionSemaphore = idleSessionSemaphore;
+    private readonly CancellationTokenSource _disposeCts = new();
+    private readonly object _referenceCountLock = new();
 
     public string Id { get; } = sessionId;
     public TTransport Transport { get; } = transport;
@@ -30,16 +37,39 @@ internal sealed class HttpMcpSession<TTransport>(
     public IMcpServer? Server { get; set; }
     public Task? ServerRunTask { get; set; }
 
-    public IDisposable AcquireReference()
+    public IAsyncDisposable AcquireReference()
     {
-        Interlocked.Increment(ref _referenceCount);
+        Debug.Assert(_idleSessionSemaphore is not null, "Only StreamableHttpHandler should call AcquireReference.");
+
+        lock (_referenceCountLock)
+        {
+            if (!_isDisposed && ++_referenceCount == 1)
+            {
+                // Non-idle sessions should not prevent session creation.
+                _idleSessionSemaphore.Release();
+            }
+        }
+
         return new UnreferenceDisposable(this);
     }
 
     public bool TryStartGetRequest() => Interlocked.Exchange(ref _getRequestStarted, 1) == 0;
 
     public async ValueTask DisposeAsync()
     {
+        bool shouldReleaseIdleSessionSemaphore;
+
+        lock (_referenceCountLock)
+        {
+            if (_isDisposed)
+            {
+                return;
+            }
+
+            _isDisposed = true;
+            shouldReleaseIdleSessionSemaphore = _referenceCount == 0;
+        }
+
         try
         {
             await _disposeCts.CancelAsync();
@@ -65,20 +95,39 @@ public async ValueTask DisposeAsync()
             {
                 await Transport.DisposeAsync();
                 _disposeCts.Dispose();
+
+                // If the session was disposed while it was inactive, we need to release the semaphore
+                // to allow new sessions to be created.
+                if (_idleSessionSemaphore is not null && shouldReleaseIdleSessionSemaphore)
+                {
+                    _idleSessionSemaphore.Release();
+                }
             }
         }
     }
 
-    public bool HasSameUserId(ClaimsPrincipal user)
-        => UserIdClaim == StreamableHttpHandler.GetUserIdClaim(user);
+    public bool HasSameUserId(ClaimsPrincipal user) => UserIdClaim == StreamableHttpHandler.GetUserIdClaim(user);
 
-    private sealed class UnreferenceDisposable(HttpMcpSession<TTransport> session) : IDisposable
+    private sealed class UnreferenceDisposable(HttpMcpSession<TTransport> session) : IAsyncDisposable
     {
-        public void Dispose()
+        public async ValueTask DisposeAsync()
         {
-            if (Interlocked.Decrement(ref session._referenceCount) == 0)
+            Debug.Assert(session._idleSessionSemaphore is not null, "Only StreamableHttpHandler should call AcquireReference.");
+
+            bool shouldMarkSessionIdle;
+
+            lock (session._referenceCountLock)
+            {
+                shouldMarkSessionIdle = !session._isDisposed && --session._referenceCount == 0;
+            }
+
+            if (shouldMarkSessionIdle)
             {
                 session.LastActivityTicks = session.TimeProvider.GetTimestamp();
+
+                // Acquire semaphore when session becomes inactive (reference count goes to 0) to slow
+                // down session creation until idle sessions are disposed by the background service.
+                await session._idleSessionSemaphore.WaitAsync();
             }
         }
     }

src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs

@@ -31,6 +31,10 @@ internal sealed class StreamableHttpHandler(
 
     public ConcurrentDictionary<string, HttpMcpSession<StreamableHttpServerTransport>> Sessions { get; } = new(StringComparer.Ordinal);
 
+    // Semaphore to control session creation backpressure when there are too many idle sessions
+    // Initial and max count is 10% more than MaxIdleSessionCount (or 100 more if that's higher)
+    private readonly SemaphoreSlim _idleSessionSemaphore = CreateIdleSessionSemaphore(httpServerTransportOptions.Value);
+
     public HttpServerTransportOptions HttpServerTransportOptions => httpServerTransportOptions.Value;
 
     private IDataProtector Protector { get; } = dataProtection.CreateProtector("Microsoft.AspNetCore.StreamableHttpHandler.StatelessSessionId");
@@ -58,7 +62,7 @@ await WriteJsonRpcErrorAsync(context,
 
         try
         {
-            using var _ = session.AcquireReference();
+            await using var _ = session.AcquireReference();
 
             InitializeSseResponse(context);
             var wroteResponse = await session.Transport.HandlePostRequest(new HttpDuplexPipe(context), context.RequestAborted);
@@ -106,7 +110,7 @@ await WriteJsonRpcErrorAsync(context,
             return;
         }
 
-        using var _ = session.AcquireReference();
+        await using var _ = session.AcquireReference();
         InitializeSseResponse(context);
 
         // We should flush headers to indicate a 200 success quickly, because the initialization response
@@ -184,6 +188,11 @@ private async ValueTask<HttpMcpSession<StreamableHttpServerTransport>> StartNewS
 
         if (!HttpServerTransportOptions.Stateless)
         {
+            // Acquire semaphore before creating stateful sessions to create backpressure.
+            // This semaphore represents "slots" for idle sessions, and we may need to wait on the
+            // IdleTrackingBackgroundService to dispose idle sessions before continuing.
+            await _idleSessionSemaphore.WaitAsync(context.RequestAborted);
+
             sessionId = MakeNewSessionId();
             transport = new()
             {
@@ -248,7 +257,8 @@ private async ValueTask<HttpMcpSession<StreamableHttpServerTransport>> CreateSes
         context.Features.Set(server);
 
         var userIdClaim = statelessId?.UserIdClaim ?? GetUserIdClaim(context.User);
-        var session = new HttpMcpSession<StreamableHttpServerTransport>(sessionId, transport, userIdClaim, HttpServerTransportOptions.TimeProvider)
+        var semaphore = HttpServerTransportOptions.Stateless ? null : _idleSessionSemaphore;
+        var session = new HttpMcpSession<StreamableHttpServerTransport>(sessionId, transport, userIdClaim, HttpServerTransportOptions.TimeProvider, semaphore)
         {
             Server = server,
         };
@@ -337,6 +347,13 @@ private static bool MatchesApplicationJsonMediaType(MediaTypeHeaderValue acceptH
     private static bool MatchesTextEventStreamMediaType(MediaTypeHeaderValue acceptHeaderValue)
         => acceptHeaderValue.MatchesMediaType("text/event-stream");
 
+    private static SemaphoreSlim CreateIdleSessionSemaphore(HttpServerTransportOptions options)
+    {
+        var maxIdleSessionCount = options.MaxIdleSessionCount;
+        var semaphoreCount = Math.Max(maxIdleSessionCount + 100, (int)(maxIdleSessionCount * 1.1));
+        return new SemaphoreSlim(semaphoreCount, semaphoreCount);
+    }
+
     private sealed class HttpDuplexPipe(HttpContext context) : IDuplexPipe
     {
         public PipeReader Input => context.Request.BodyReader;