feat(protocol): Debounce notifications to improve network efficiancy (#746)
This commit is contained in:
parent
f3584f2a41
commit
af61a086ae
3 changed files with 269 additions and 0 deletions
37
README.md
37
README.md
|
|
@ -913,6 +913,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:
|
||||
|
|
|
|||
|
|
@ -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<Request, Notification, Result> {
|
||||
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<Request, Notification, Result> {
|
||||
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<Request, Notification, Result> {
|
||||
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<Request, Notification, Result> {
|
||||
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<Request, Notification, Result> {
|
||||
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<Request, Notification, Result> {
|
||||
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<Request, Notification, Result> {
|
||||
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", () => {
|
||||
|
|
|
|||
|
|
@ -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<number, ProgressCallback> = new Map();
|
||||
private _timeoutInfo: Map<number, TimeoutInfo> = new Map();
|
||||
private _pendingDebouncedNotifications = new Set<string>();
|
||||
|
||||
/**
|
||||
* 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",
|
||||
|
|
|
|||
Loading…
Reference in a new issue