diff --git a/src/everything/__tests__/resources.test.ts b/src/everything/__tests__/resources.test.ts index a22b904175..074736324e 100644 --- a/src/everything/__tests__/resources.test.ts +++ b/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 = {}; 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); }); }); diff --git a/src/everything/__tests__/server.test.ts b/src/everything/__tests__/server.test.ts index e7985dd982..dcb54ccd77 100644 --- a/src/everything/__tests__/server.test.ts +++ b/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'); + } }); }); }); diff --git a/src/everything/resources/subscriptions.ts b/src/everything/resources/subscriptions.ts index 854a8633a2..6ebeb6fe5d 100644 --- a/src/everything/resources/subscriptions.ts +++ b/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; 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 +> => subscriptions; + /** * Sends simulated resource update notifications to the subscribed client. * diff --git a/src/everything/server/index.ts b/src/everything/server/index.ts index f1459cc812..ad9d2910a0 100644 --- a/src/everything/server/index.ts +++ b/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);