diff --git a/src/everything/__tests__/resources.test.ts b/src/everything/__tests__/resources.test.ts index a22b904175..33b992db78 100644 --- a/src/everything/__tests__/resources.test.ts +++ b/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; + + beforeEach(() => { + const handlers = new Map(); + 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', () => { 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..da130ce391 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,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);