Move subscription management to a dedicated class and add subscription cleanup on session end

Author: mikekistler

samples/EverythingServer/Program.cs

@@ -15,14 +15,24 @@
 
 var builder = WebApplication.CreateBuilder(args);
 
-// Subscriptions tracks resource URIs to McpServer instances
-// Use thread-safe data structures since handlers can run in parallel
-// even in the context of a single session.
-ConcurrentDictionary<string, ConcurrentBag<IMcpServer>> subscriptions = new();
-
 builder.Services
     .AddMcpServer()
-    .WithHttpTransport()
+    .WithHttpTransport(options =>
+    {
+        // Add a RunSessionHandler to remove all subscriptions for the session when it ends
+        options.RunSessionHandler = async (httpContext, mcpServer, token) =>
+        {
+            try
+            {
+                await mcpServer.RunAsync(token);
+            }
+            finally
+            {
+                // This code runs when the session ends
+                SubscriptionManager.RemoveAllSubscriptions(mcpServer);
+            }
+        };
+    })
     .WithTools<AddTool>()
     .WithTools<AnnotatedMessageTool>()
     .WithTools<EchoTool>()
@@ -35,10 +45,13 @@
     .WithResources<SimpleResourceType>()
     .WithSubscribeToResourcesHandler(async (ctx, ct) =>
     {
+        if (ctx.Server.SessionId == null)
+        {
+            throw new McpException("Cannot add subscription for server with null SessionId");
+        }
         if (ctx.Params?.Uri is { } uri)
         {
-            var bag = subscriptions.GetOrAdd(uri, _ => new ConcurrentBag<IMcpServer>());
-            bag.Add(ctx.Server);
+            SubscriptionManager.AddSubscription(uri, ctx.Server);
 
             await ctx.Server.SampleAsync([
                 new ChatMessage(ChatRole.System, "You are a helpful test server"),
@@ -56,14 +69,13 @@ await ctx.Server.SampleAsync([
     })
     .WithUnsubscribeFromResourcesHandler(async (ctx, ct) =>
     {
+        if (ctx.Server.SessionId == null)
+        {
+            throw new McpException("Cannot remove subscription for server with null SessionId");
+        }
         if (ctx.Params?.Uri is { } uri)
         {
-            if (subscriptions.TryGetValue(uri, out var bag))
-            {
-                // Remove ctx.Server from the subscription bag (ConcurrentBag does not support removal, so recreate)
-                var newBag = new ConcurrentBag<IMcpServer>(bag.Where(s => s != ctx.Server));
-                subscriptions[uri] = newBag;
-            }
+            SubscriptionManager.RemoveSubscription(uri, ctx.Server);
         }
         return new EmptyResult();
     })
@@ -143,7 +155,6 @@ await ctx.Server.SampleAsync([
     .WithLogging(b => b.SetResourceBuilder(resource))
     .UseOtlpExporter();
 
-builder.Services.AddSingleton(subscriptions);
 builder.Services.AddHostedService<SubscriptionMessageSender>();
 builder.Services.AddHostedService<LoggingUpdateMessageSender>();
 

samples/EverythingServer/SubscriptionManager.cs

@@ -0,0 +1,91 @@
+using ModelContextProtocol;
+using ModelContextProtocol.Server;
+
+// This class manages subscriptions to resources by McpServer instances.
+// The subscription information must be accessed in a thread-safe manner since handlers
+// can run in parallel even in the context of a single session.
+static class SubscriptionManager
+{
+    // Subscriptions tracks resource URIs to bags of McpServer instances (thread-safe via locking)
+    private static Dictionary<string, List<IMcpServer>> subscriptions = new();
+
+    // SessionSubscriptions is a secondary index to subscriptions to allow efficient removal of all
+    // subscriptions for a given session when it ends. (thread-safe via locking)
+    private static Dictionary<string /* sessionId */, List<string> /* uris */> sessionSubscriptions = new();
+
+    private static readonly object _subscriptionsLock = new();
+
+    public static void AddSubscription(string uri, IMcpServer server)
+    {
+        if (server.SessionId == null)
+        {
+            throw new McpException("Cannot add subscription for server with null SessionId");
+        }
+        lock (_subscriptionsLock)
+        {
+            subscriptions[uri] ??= new List<IMcpServer>();
+            subscriptions[uri].Add(server);
+            sessionSubscriptions[server.SessionId] ??= new List<string>();
+            sessionSubscriptions[server.SessionId].Add(uri);
+        }
+    }
+
+    public static void RemoveSubscription(string uri, IMcpServer server)
+    {
+        if (server.SessionId == null)
+        {
+            throw new McpException("Cannot remove subscription for server with null SessionId");
+        }
+        lock (_subscriptionsLock)
+        {
+            if (subscriptions.ContainsKey(uri))
+            {
+                // Remove the server from the list of subscriptions for the URI
+                subscriptions[uri] = subscriptions[uri].Where(s => s.SessionId != server.SessionId).ToList();
+                if (subscriptions[uri]?.Count == 0)
+                {
+                    subscriptions.Remove(uri);
+                }
+            }
+            // Remove the URI from the list of subscriptions for the session
+            sessionSubscriptions[server.SessionId]?.Remove(uri);
+            if (sessionSubscriptions[server.SessionId]?.Count == 0)
+            {
+                sessionSubscriptions.Remove(server.SessionId);
+            }
+        }
+    }
+
+    public static IDictionary<string, List<IMcpServer>> GetSubscriptions()
+    {
+        lock (_subscriptionsLock)
+        {
+            // Return a copy of the subscriptions dictionary to avoid external modification
+            return subscriptions.ToDictionary(entry => entry.Key,
+                                             entry => entry.Value.ToList());
+        }
+    }
+
+    public static void RemoveAllSubscriptions(IMcpServer server)
+    {
+        if (server.SessionId is { } sessionId)
+        {
+            lock (_subscriptionsLock)
+            {
+                // Remove all subscriptions for the session
+                if (sessionSubscriptions.TryGetValue(sessionId, out var uris))
+                {
+                    foreach (var uri in uris)
+                    {
+                        subscriptions[uri] = subscriptions[uri].Where(s => s.SessionId != sessionId).ToList();
+                        if (subscriptions[uri]?.Count == 0)
+                        {
+                            subscriptions.Remove(uri);
+                        }
+                    }
+                    sessionSubscriptions.Remove(sessionId);
+                }
+            }
+        }
+    }
+}

samples/EverythingServer/SubscriptionMessageSender.cs

@@ -2,7 +2,7 @@
 using ModelContextProtocol.Server;
 
 using System.Collections.Concurrent;
-internal class SubscriptionMessageSender(ConcurrentDictionary<string, ConcurrentBag<IMcpServer>> subscriptions) : BackgroundService
+internal class SubscriptionMessageSender() : BackgroundService
 {
     protected override async Task ExecuteAsync(CancellationToken stoppingToken)
     {
@@ -13,7 +13,7 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
         {
             try
             {
-                foreach (var (uri, servers) in subscriptions)
+                foreach (var (uri, servers) in SubscriptionManager.GetSubscriptions())
                 {
                     foreach (var server in servers)
                     {