diff --git a/README.md b/README.md index ac10e8cb..452fe7dc 100644 --- a/README.md +++ b/README.md @@ -902,6 +902,43 @@ const transport = new StdioServerTransport(); await server.connect(transport); ``` +### Improving Network Efficiency with Notification Debouncing + +When performing bulk updates that trigger notifications (e.g., enabling or disabling multiple tools in a loop), the SDK can send a large number of messages in a short period. To improve performance and reduce network traffic, you can enable notification debouncing. + +This feature coalesces multiple, rapid calls for the same notification type into a single message. For example, if you disable five tools in a row, only one `notifications/tools/list_changed` message will be sent instead of five. + +> [!IMPORTANT] +> This feature is designed for "simple" notifications that do not carry unique data in their parameters. To prevent silent data loss, debouncing is **automatically bypassed** for any notification that contains a `params` object or a `relatedRequestId`. Such notifications will always be sent immediately. + +This is an opt-in feature configured during server initialization. + +```typescript +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; + +const server = new McpServer( + { + name: "efficient-server", + version: "1.0.0" + }, + { + // Enable notification debouncing for specific methods + debouncedNotificationMethods: [ + 'notifications/tools/list_changed', + 'notifications/resources/list_changed', + 'notifications/prompts/list_changed' + ] + } +); + +// Now, any rapid changes to tools, resources, or prompts will result +// in a single, consolidated notification for each type. +server.registerTool("tool1", ...).disable(); +server.registerTool("tool2", ...).disable(); +server.registerTool("tool3", ...).disable(); +// Only one 'notifications/tools/list_changed' is sent. +``` + ### Low-Level Server For more control, you can use the low-level Server class directly: diff --git a/src/shared/protocol.test.ts b/src/shared/protocol.test.ts index b16db73f..f4e74c8b 100644 --- a/src/shared/protocol.test.ts +++ b/src/shared/protocol.test.ts @@ -466,6 +466,189 @@ describe("protocol tests", () => { await expect(requestPromise).resolves.toEqual({ result: "success" }); }); }); + + describe("Debounced Notifications", () => { + // We need to flush the microtask queue to test the debouncing logic. + // This helper function does that. + const flushMicrotasks = () => new Promise(resolve => setImmediate(resolve)); + + it("should NOT debounce a notification that has parameters", async () => { + // ARRANGE + protocol = new (class extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + })({ debouncedNotificationMethods: ['test/debounced_with_params'] }); + await protocol.connect(transport); + + // ACT + // These notifications are configured for debouncing but contain params, so they should be sent immediately. + await protocol.notification({ method: 'test/debounced_with_params', params: { data: 1 } }); + await protocol.notification({ method: 'test/debounced_with_params', params: { data: 2 } }); + + // ASSERT + // Both should have been sent immediately to avoid data loss. + expect(sendSpy).toHaveBeenCalledTimes(2); + expect(sendSpy).toHaveBeenCalledWith(expect.objectContaining({ params: { data: 1 } }), undefined); + expect(sendSpy).toHaveBeenCalledWith(expect.objectContaining({ params: { data: 2 } }), undefined); + }); + + it("should NOT debounce a notification that has a relatedRequestId", async () => { + // ARRANGE + protocol = new (class extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + })({ debouncedNotificationMethods: ['test/debounced_with_options'] }); + await protocol.connect(transport); + + // ACT + await protocol.notification({ method: 'test/debounced_with_options' }, { relatedRequestId: 'req-1' }); + await protocol.notification({ method: 'test/debounced_with_options' }, { relatedRequestId: 'req-2' }); + + // ASSERT + expect(sendSpy).toHaveBeenCalledTimes(2); + expect(sendSpy).toHaveBeenCalledWith(expect.any(Object), { relatedRequestId: 'req-1' }); + expect(sendSpy).toHaveBeenCalledWith(expect.any(Object), { relatedRequestId: 'req-2' }); + }); + + it("should clear pending debounced notifications on connection close", async () => { + // ARRANGE + protocol = new (class extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + })({ debouncedNotificationMethods: ['test/debounced'] }); + await protocol.connect(transport); + + // ACT + // Schedule a notification but don't flush the microtask queue. + protocol.notification({ method: 'test/debounced' }); + + // Close the connection. This should clear the pending set. + await protocol.close(); + + // Now, flush the microtask queue. + await flushMicrotasks(); + + // ASSERT + // The send should never have happened because the transport was cleared. + expect(sendSpy).not.toHaveBeenCalled(); + }); + + it("should debounce multiple synchronous calls when params property is omitted", async () => { + // ARRANGE + protocol = new (class extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + })({ debouncedNotificationMethods: ['test/debounced'] }); + await protocol.connect(transport); + + // ACT + // This is the more idiomatic way to write a notification with no params. + protocol.notification({ method: 'test/debounced' }); + protocol.notification({ method: 'test/debounced' }); + protocol.notification({ method: 'test/debounced' }); + + expect(sendSpy).not.toHaveBeenCalled(); + await flushMicrotasks(); + + // ASSERT + expect(sendSpy).toHaveBeenCalledTimes(1); + // The final sent object might not even have the `params` key, which is fine. + // We can check that it was called and that the params are "falsy". + const sentNotification = sendSpy.mock.calls[0][0]; + expect(sentNotification.method).toBe('test/debounced'); + expect(sentNotification.params).toBeUndefined(); + }); + + it("should debounce calls when params is explicitly undefined", async () => { + // ARRANGE + protocol = new (class extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + })({ debouncedNotificationMethods: ['test/debounced'] }); + await protocol.connect(transport); + + // ACT + protocol.notification({ method: 'test/debounced', params: undefined }); + protocol.notification({ method: 'test/debounced', params: undefined }); + await flushMicrotasks(); + + // ASSERT + expect(sendSpy).toHaveBeenCalledTimes(1); + expect(sendSpy).toHaveBeenCalledWith( + expect.objectContaining({ + method: 'test/debounced', + params: undefined + }), + undefined + ); + }); + + it("should send non-debounced notifications immediately and multiple times", async () => { + // ARRANGE + protocol = new (class extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + })({ debouncedNotificationMethods: ['test/debounced'] }); // Configure for a different method + await protocol.connect(transport); + + // ACT + // Call a non-debounced notification method multiple times. + await protocol.notification({ method: 'test/immediate' }); + await protocol.notification({ method: 'test/immediate' }); + + // ASSERT + // Since this method is not in the debounce list, it should be sent every time. + expect(sendSpy).toHaveBeenCalledTimes(2); + }); + + it("should not debounce any notifications if the option is not provided", async () => { + // ARRANGE + // Use the default protocol from beforeEach, which has no debounce options. + await protocol.connect(transport); + + // ACT + await protocol.notification({ method: 'any/method' }); + await protocol.notification({ method: 'any/method' }); + + // ASSERT + // Without the config, behavior should be immediate sending. + expect(sendSpy).toHaveBeenCalledTimes(2); + }); + + it("should handle sequential batches of debounced notifications correctly", async () => { + // ARRANGE + protocol = new (class extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + })({ debouncedNotificationMethods: ['test/debounced'] }); + await protocol.connect(transport); + + // ACT (Batch 1) + protocol.notification({ method: 'test/debounced' }); + protocol.notification({ method: 'test/debounced' }); + await flushMicrotasks(); + + // ASSERT (Batch 1) + expect(sendSpy).toHaveBeenCalledTimes(1); + + // ACT (Batch 2) + // After the first batch has been sent, a new batch should be possible. + protocol.notification({ method: 'test/debounced' }); + protocol.notification({ method: 'test/debounced' }); + await flushMicrotasks(); + + // ASSERT (Batch 2) + // The total number of sends should now be 2. + expect(sendSpy).toHaveBeenCalledTimes(2); + }); + }); }); describe("mergeCapabilities", () => { diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 50bdcc3c..6142140d 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -45,6 +45,13 @@ export type ProtocolOptions = { * Currently this defaults to false, for backwards compatibility with SDK versions that did not advertise capabilities correctly. In future, this will default to true. */ enforceStrictCapabilities?: boolean; + /** + * An array of notification method names that should be automatically debounced. + * Any notifications with a method in this list will be coalesced if they + * occur in the same tick of the event loop. + * e.g., ['notifications/tools/list_changed'] + */ + debouncedNotificationMethods?: string[]; }; /** @@ -191,6 +198,7 @@ export abstract class Protocol< > = new Map(); private _progressHandlers: Map = new Map(); private _timeoutInfo: Map = new Map(); + private _pendingDebouncedNotifications = new Set(); /** * Callback for when the connection is closed for any reason. @@ -321,6 +329,7 @@ export abstract class Protocol< const responseHandlers = this._responseHandlers; this._responseHandlers = new Map(); this._progressHandlers.clear(); + this._pendingDebouncedNotifications.clear(); this._transport = undefined; this.onclose?.(); @@ -632,6 +641,46 @@ export abstract class Protocol< this.assertNotificationCapability(notification.method); + const debouncedMethods = this._options?.debouncedNotificationMethods ?? []; + // A notification can only be debounced if it's in the list AND it's "simple" + // (i.e., has no parameters and no related request ID that could be lost). + const canDebounce = debouncedMethods.includes(notification.method) + && !notification.params + && !(options?.relatedRequestId); + + if (canDebounce) { + // If a notification of this type is already scheduled, do nothing. + if (this._pendingDebouncedNotifications.has(notification.method)) { + return; + } + + // Mark this notification type as pending. + this._pendingDebouncedNotifications.add(notification.method); + + // Schedule the actual send to happen in the next microtask. + // This allows all synchronous calls in the current event loop tick to be coalesced. + Promise.resolve().then(() => { + // Un-mark the notification so the next one can be scheduled. + this._pendingDebouncedNotifications.delete(notification.method); + + // SAFETY CHECK: If the connection was closed while this was pending, abort. + if (!this._transport) { + return; + } + + const jsonrpcNotification: JSONRPCNotification = { + ...notification, + jsonrpc: "2.0", + }; + // Send the notification, but don't await it here to avoid blocking. + // Handle potential errors with a .catch(). + this._transport?.send(jsonrpcNotification, options).catch(error => this._onerror(error)); + }); + + // Return immediately. + return; + } + const jsonrpcNotification: JSONRPCNotification = { ...notification, jsonrpc: "2.0",