fix(everything): clean up subscriptions on session disconnect (#4712)

Export removeSubscriber() to drop a session from all URI subscriber sets
when a transport session closes, and call it from cleanup() alongside
the existing interval and logging teardown.

Fixes #4710

Co-authored-by: arimu1 <19286898+arimu1@users.noreply.github.com>

Author: arimu1

src/everything/__tests__/resources.test.ts

@@ -18,10 +18,14 @@ import {
   registerSessionResource,
 } from '../resources/session.js';
 import { registerFileResources } from '../resources/files.js';
+import {
+  SubscribeRequestSchema,
+} from '@modelcontextprotocol/sdk/types.js';
 import {
   setSubscriptionHandlers,
   beginSimulatedResourceUpdates,
   stopSimulatedResourceUpdates,
+  removeSubscriber,
 } from '../resources/subscriptions.js';
 
 describe('Resource Templates', () => {
@@ -298,10 +302,87 @@ describe('Subscriptions', () => {
     });
   });
 
+  describe('removeSubscriber', () => {
+    const testUri = 'demo://resource/dynamic/text/1';
+    const sessionId = 'disconnect-test-session';
+
+    let subscribeHandler: (
+      request: { params: { uri: string } },
+      extra: { sessionId: string }
+    ) => Promise<unknown>;
+
+    beforeEach(() => {
+      const handlers = new Map<unknown, typeof subscribeHandler>();
+      const mockServer = {
+        server: {
+          setRequestHandler: vi.fn((schema, handler) => {
+            handlers.set(schema, handler);
+          }),
+          notification: vi.fn(),
+        },
+        sendLoggingMessage: vi.fn(),
+      } as unknown as McpServer;
+
+      setSubscriptionHandlers(mockServer);
+      subscribeHandler = handlers.get(SubscribeRequestSchema)!;
+    });
+
+    afterEach(() => {
+      stopSimulatedResourceUpdates(sessionId);
+      removeSubscriber(sessionId);
+    });
+
+    it('should drop a disconnected session from all subscriptions', async () => {
+      const notification = vi.fn();
+      const mockServer = {
+        server: {
+          notification,
+        },
+      } as unknown as McpServer;
+
+      await subscribeHandler({ params: { uri: testUri } }, { sessionId });
+
+      beginSimulatedResourceUpdates(mockServer, sessionId);
+      expect(notification).toHaveBeenCalled();
+
+      notification.mockClear();
+      removeSubscriber(sessionId);
+      stopSimulatedResourceUpdates(sessionId);
+
+      beginSimulatedResourceUpdates(mockServer, sessionId);
+      expect(notification).not.toHaveBeenCalled();
+    });
+
+    it('should not affect other sessions subscribed to the same URI', async () => {
+      const otherSessionId = 'other-session';
+      const notification = vi.fn();
+      const mockServer = {
+        server: {
+          notification,
+        },
+      } as unknown as McpServer;
+
+      await subscribeHandler({ params: { uri: testUri } }, { sessionId });
+      await subscribeHandler(
+        { params: { uri: testUri } },
+        { sessionId: otherSessionId }
+      );
+
+      removeSubscriber(sessionId);
+
+      beginSimulatedResourceUpdates(mockServer, otherSessionId);
+      expect(notification).toHaveBeenCalled();
+
+      stopSimulatedResourceUpdates(otherSessionId);
+      removeSubscriber(otherSessionId);
+    });
+  });
+
   describe('simulated resource updates lifecycle', () => {
     afterEach(() => {
       // Clean up any intervals
       stopSimulatedResourceUpdates('lifecycle-test-session');
+      removeSubscriber('lifecycle-test-session');
     });
 
     it('should start and stop updates without errors', () => {

src/everything/resources/subscriptions.ts

@@ -166,3 +166,22 @@ export const stopSimulatedResourceUpdates = (sessionId?: string) => {
     subsUpdateIntervals.delete(sessionId);
   }
 };
+
+/**
+ * Removes a session from every URI's subscriber set, dropping any URI entry
+ * that ends up with no remaining subscribers.
+ *
+ * A session that disconnects without explicitly unsubscribing otherwise stays
+ * in `subscriptions` for the life of the process. Call this from the
+ * transport's `cleanup(sessionId)` when a session ends.
+ *
+ * @param {string} [sessionId]
+ */
+export const removeSubscriber = (sessionId?: string) => {
+  for (const [uri, subscribers] of subscriptions) {
+    subscribers.delete(sessionId);
+    if (subscribers.size === 0) {
+      subscriptions.delete(uri);
+    }
+  }
+};

src/everything/server/index.ts

@@ -6,6 +6,7 @@ import {
 import {
   setSubscriptionHandlers,
   stopSimulatedResourceUpdates,
+  removeSubscriber,
 } 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);
+      removeSubscriber(sessionId);
       // Clean up task store timers
       taskStore.cleanup();
       if (initializeTimeout) clearTimeout(initializeTimeout);