fix(everything): drop disconnected session from subscriptions map

src/everything/__tests__/resources.test.ts

@@ -22,6 +22,8 @@ import {
   setSubscriptionHandlers,
   beginSimulatedResourceUpdates,
   stopSimulatedResourceUpdates,
+  cleanupSubscriptions,
+  getSubscriptions,
 } from '../resources/subscriptions.js';
 
 describe('Resource Templates', () => {
@@ -283,10 +285,13 @@ describe('File Resources', () => {
 
 describe('Subscriptions', () => {
   describe('setSubscriptionHandlers', () => {
-    it('should set request handlers on server', () => {
+    it('should set request handlers on server and handle subscribe/unsubscribe/cleanup', async () => {
+      const handlers: Record<string, Function> = {};
       const mockServer = {
         server: {
-          setRequestHandler: vi.fn(),
+          setRequestHandler: vi.fn((schema: any, handler: Function) => {
+            handlers[schema.shape?.method?.value || 'handler'] = handler;
+          }),
         },
         sendLoggingMessage: vi.fn(),
       } as unknown as McpServer;
@@ -295,6 +300,82 @@ describe('Subscriptions', () => {
 
       // Should set both subscribe and unsubscribe handlers
       expect(mockServer.server.setRequestHandler).toHaveBeenCalledTimes(2);
+
+      const subscribeHandler = (mockServer.server.setRequestHandler as any).mock.calls[0][1];
+      const unsubscribeHandler = (mockServer.server.setRequestHandler as any).mock.calls[1][1];
+
+      // Subscribe session1 to uri1 and uri2
+      await subscribeHandler(
+        { method: 'resources/subscribe', params: { uri: 'test://uri1' } },
+        { sessionId: 'session1' }
+      );
+      await subscribeHandler(
+        { method: 'resources/subscribe', params: { uri: 'test://uri2' } },
+        { sessionId: 'session1' }
+      );
+      // Subscribe session2 to uri1
+      await subscribeHandler(
+        { method: 'resources/subscribe', params: { uri: 'test://uri1' } },
+        { sessionId: 'session2' }
+      );
+
+      const subs = getSubscriptions();
+      expect(subs.has('test://uri1')).toBe(true);
+      expect(subs.get('test://uri1')?.has('session1')).toBe(true);
+      expect(subs.get('test://uri1')?.has('session2')).toBe(true);
+      expect(subs.get('test://uri2')?.has('session1')).toBe(true);
+
+      // Unsubscribe session2 from uri1
+      await unsubscribeHandler(
+        { method: 'resources/unsubscribe', params: { uri: 'test://uri1' } },
+        { sessionId: 'session2' }
+      );
+      expect(subs.get('test://uri1')?.has('session2')).toBe(false);
+      expect(subs.get('test://uri1')?.has('session1')).toBe(true);
+
+      // Cleanup session1 on disconnect - should remove session1 from uri1 and uri2, deleting empty uri sets
+      cleanupSubscriptions('session1');
+      expect(subs.has('test://uri1')).toBe(false);
+      expect(subs.has('test://uri2')).toBe(false);
+    });
+  });
+
+  describe('cleanupSubscriptions', () => {
+    it('should cleanly remove subscriptions for disconnected session', async () => {
+      const mockServer = {
+        server: {
+          setRequestHandler: vi.fn(),
+        },
+        sendLoggingMessage: vi.fn(),
+      } as unknown as McpServer;
+
+      setSubscriptionHandlers(mockServer);
+      const subscribeHandler = (mockServer.server.setRequestHandler as any).mock.calls[0][1];
+
+      await subscribeHandler(
+        { method: 'resources/subscribe', params: { uri: 'test://shared' } },
+        { sessionId: 'userA' }
+      );
+      await subscribeHandler(
+        { method: 'resources/subscribe', params: { uri: 'test://shared' } },
+        { sessionId: 'userB' }
+      );
+      await subscribeHandler(
+        { method: 'resources/subscribe', params: { uri: 'test://privateA' } },
+        { sessionId: 'userA' }
+      );
+
+      const subs = getSubscriptions();
+      expect(subs.get('test://shared')?.size).toBe(2);
+      expect(subs.get('test://privateA')?.size).toBe(1);
+
+      cleanupSubscriptions('userA');
+      expect(subs.get('test://shared')?.has('userA')).toBe(false);
+      expect(subs.get('test://shared')?.has('userB')).toBe(true);
+      expect(subs.has('test://privateA')).toBe(false);
+
+      cleanupSubscriptions('userB');
+      expect(subs.has('test://shared')).toBe(false);
     });
   });
 

src/everything/__tests__/server.test.ts

@@ -29,13 +29,30 @@ describe('Server Factory', () => {
       expect(server.server.oninitialized).toBeDefined();
     });
 
-    it('should allow multiple servers to be created', () => {
-      const result1 = createServer();
-      const result2 = createServer();
-
-      expect(result1.server).toBeDefined();
-      expect(result2.server).toBeDefined();
-      expect(result1.server).not.toBe(result2.server);
+    it('should clean up subscriptions when cleanup is called for a session', async () => {
+      const { getSubscriptions, setSubscriptionHandlers } = await import(
+        '../resources/subscriptions.js'
+      );
+      const { createServer } = await import('../server/index.js');
+      const { server, cleanup } = createServer();
+
+      // Subscribe session 'cleanup-test-session'
+      const subscribeHandler = (server.server.setRequestHandler as any).mock?.calls?.find(
+        (call: any[]) => call[0]?.shape?.method?.value === 'resources/subscribe' || call[1]
+      )?.[1];
+
+      if (subscribeHandler) {
+        await subscribeHandler(
+          { method: 'resources/subscribe', params: { uri: 'test://cleanup-uri' } },
+          { sessionId: 'cleanup-test-session' }
+        );
+        expect(getSubscriptions().get('test://cleanup-uri')?.has('cleanup-test-session')).toBe(true);
+
+        cleanup('cleanup-test-session');
+        expect(getSubscriptions().get('test://cleanup-uri')?.has('cleanup-test-session')).toBeFalsy();
+      } else {
+        cleanup('cleanup-test-session');
+      }
     });
   });
 });

src/everything/resources/subscriptions.ts

@@ -90,12 +90,40 @@ export const setSubscriptionHandlers = (server: McpServer) => {
       if (subscriptions.has(uri)) {
         const subscribers = subscriptions.get(uri) as Set<string>;
         if (subscribers.has(sessionId)) subscribers.delete(sessionId);
+        if (subscribers.size === 0) {
+          subscriptions.delete(uri);
+        }
       }
       return {};
     }
   );
 };
 
+/**
+ * Removes all subscriptions associated with a given session ID.
+ *
+ * Iterates through all tracked subscription URIs, removes the session ID,
+ * and deletes the URI key if no subscribers remain.
+ *
+ * @param {string | undefined} [sessionId] - The session ID to clean up.
+ */
+export const cleanupSubscriptions = (sessionId?: string) => {
+  for (const [uri, subscribers] of subscriptions.entries()) {
+    subscribers.delete(sessionId);
+    if (subscribers.size === 0) {
+      subscriptions.delete(uri);
+    }
+  }
+};
+
+/**
+ * Returns the current subscriptions map (for testing/inspection).
+ */
+export const getSubscriptions = (): ReadonlyMap<
+  string,
+  ReadonlySet<string | undefined>
+> => subscriptions;
+
 /**
  * Sends simulated resource update notifications to the subscribed client.
  *

src/everything/server/index.ts

@@ -6,6 +6,7 @@ import {
 import {
   setSubscriptionHandlers,
   stopSimulatedResourceUpdates,
+  cleanupSubscriptions,
 } from "../resources/subscriptions.js";
 import { registerConditionalTools, registerTools } from "../tools/index.js";
 import { registerResources, readInstructions } from "../resources/index.js";
@@ -110,6 +111,7 @@ export const createServer: () => ServerFactoryResponse = () => {
       // Stop any simulated logging or resource updates that may have been initiated.
       stopSimulatedLogging(sessionId);
       stopSimulatedResourceUpdates(sessionId);
+      cleanupSubscriptions(sessionId);
       // Clean up task store timers
       taskStore.cleanup();
       if (initializeTimeout) clearTimeout(initializeTimeout);