diff --git a/src/everything/__tests__/resources.test.ts b/src/everything/__tests__/resources.test.ts index a22b904175..c70f991340 100644 --- a/src/everything/__tests__/resources.test.ts +++ b/src/everything/__tests__/resources.test.ts @@ -22,6 +22,7 @@ import { setSubscriptionHandlers, beginSimulatedResourceUpdates, stopSimulatedResourceUpdates, + removeSubscriber, } from '../resources/subscriptions.js'; describe('Resource Templates', () => { @@ -324,4 +325,65 @@ describe('Subscriptions', () => { expect(true).toBe(true); }); }); + + describe('removeSubscriber', () => { + const flushAsync = () => new Promise((resolve) => setTimeout(resolve, 10)); + + it('should stop sending resource update notifications to a removed session', async () => { + const notification = vi.fn(); + const setRequestHandler = vi.fn(); + const mockServer = { + server: { setRequestHandler, notification }, + sendLoggingMessage: vi.fn(), + } as unknown as McpServer; + + setSubscriptionHandlers(mockServer); + const subscribeHandler = setRequestHandler.mock.calls[0][1]; + + const uri = 'demo://resource/dynamic/text/removeSubscriber-1'; + await subscribeHandler( + { params: { uri } }, + { sessionId: 'session-removed' } + ); + + removeSubscriber('session-removed'); + + beginSimulatedResourceUpdates(mockServer, 'session-removed'); + await flushAsync(); + + expect(notification).not.toHaveBeenCalled(); + + stopSimulatedResourceUpdates('session-removed'); + }); + + it('should leave other sessions subscribed to the same URI unaffected', async () => { + const notification = vi.fn(); + const setRequestHandler = vi.fn(); + const mockServer = { + server: { setRequestHandler, notification }, + sendLoggingMessage: vi.fn(), + } as unknown as McpServer; + + setSubscriptionHandlers(mockServer); + const subscribeHandler = setRequestHandler.mock.calls[0][1]; + + const uri = 'demo://resource/dynamic/text/removeSubscriber-2'; + await subscribeHandler({ params: { uri } }, { sessionId: 'session-a' }); + await subscribeHandler({ params: { uri } }, { sessionId: 'session-b' }); + + removeSubscriber('session-a'); + + beginSimulatedResourceUpdates(mockServer, 'session-b'); + await flushAsync(); + + expect(notification).toHaveBeenCalledWith( + expect.objectContaining({ + method: 'notifications/resources/updated', + params: { uri }, + }) + ); + + stopSimulatedResourceUpdates('session-b'); + }); + }); }); diff --git a/src/everything/resources/subscriptions.ts b/src/everything/resources/subscriptions.ts index 854a8633a2..9d2044147c 100644 --- a/src/everything/resources/subscriptions.ts +++ b/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); + } + } +}; diff --git a/src/everything/server/index.ts b/src/everything/server/index.ts index f1459cc812..182a5a084a 100644 --- a/src/everything/server/index.ts +++ b/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,8 @@ export const createServer: () => ServerFactoryResponse = () => { // Stop any simulated logging or resource updates that may have been initiated. stopSimulatedLogging(sessionId); stopSimulatedResourceUpdates(sessionId); + // Drop this session from any resource subscriptions it left open. + removeSubscriber(sessionId); // Clean up task store timers taskStore.cleanup(); if (initializeTimeout) clearTimeout(initializeTimeout);