Add tool list changed handling to Client

chipgpt · chipgpt · commit e50874b48229 · 2025-11-01T14:27:47.000-05:00
diff --git a/src/client/index.test.ts b/src/client/index.test.ts
@@ -16,7 +16,8 @@ import {
     CreateMessageRequestSchema,
     ElicitRequestSchema,
     ListRootsRequestSchema,
-    ErrorCode
+    ErrorCode,
+    Tool
 } from '../types.js';
 import { Transport } from '../shared/transport.js';
 import { Server } from '../server/index.js';
@@ -773,6 +774,173 @@ test('should handle request timeout', async () => {
     });
 });
 
+/***
+ * Test: Handle Tool List Changed Notifications with Auto Refresh
+ */
+test('should handle tool list changed notification with auto refresh', async () => {
+    // List changed notifications
+    const notifications: [Error | null, Tool[] | null][] = [];
+
+    const server = new Server(
+        {
+            name: 'test-server',
+            version: '1.0.0'
+        },
+        {
+            capabilities: {
+                tools: {
+                    listChanged: true
+                }
+            }
+        }
+    );
+
+    // Set up server handlers
+    server.setRequestHandler(InitializeRequestSchema, async request => ({
+        protocolVersion: request.params.protocolVersion,
+        capabilities: {
+            tools: {
+                listChanged: true
+            }
+        },
+        serverInfo: {
+            name: 'test-server',
+            version: '1.0.0'
+        }
+    }));
+
+    server.setRequestHandler(ListToolsRequestSchema, async () => ({
+        tools: []
+    }));
+
+    const client = new Client({
+        name: 'test-client',
+        version: '1.0.0',
+    }, {
+        toolListChangedOptions: {
+            autoRefresh: true,
+            onToolListChanged: (err, tools) => {
+                notifications.push([err, tools]);
+            }
+        }
+    });
+
+    const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+
+    await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
+
+    const result1 = await client.listTools();
+    expect(result1.tools).toHaveLength(0);
+
+    // Update the tools list
+    server.setRequestHandler(ListToolsRequestSchema, async () => ({
+        tools: [
+            {
+                name: 'test-tool',
+                description: 'A test tool',
+                inputSchema: {
+                    type: 'object',
+                    properties: {}
+                }
+                // No outputSchema
+            }
+        ]
+    }));
+    await server.sendToolListChanged();
+
+    // Wait for the debounced notifications to be processed
+    await new Promise(resolve => setTimeout(resolve, 1000));
+
+    // Should be 1 notification with 1 tool because autoRefresh is true
+    expect(notifications).toHaveLength(1);
+    expect(notifications[0][0]).toBeNull();
+    expect(notifications[0][1]).toHaveLength(1);
+    expect(notifications[0][1]?.[0].name).toBe('test-tool');
+});
+
+/***
+ * Test: Handle Tool List Changed Notifications with Manual Refresh
+ */
+test('should handle tool list changed notification with manual refresh', async () => {
+    // List changed notifications
+    const notifications: [Error | null, Tool[] | null][] = [];
+
+    const server = new Server(
+        {
+            name: 'test-server',
+            version: '1.0.0'
+        },
+        {
+            capabilities: {
+                tools: {
+                    listChanged: true
+                }
+            }
+        }
+    );
+
+    // Set up server handlers
+    server.setRequestHandler(InitializeRequestSchema, async request => ({
+        protocolVersion: request.params.protocolVersion,
+        capabilities: {
+            tools: {
+                listChanged: true
+            }
+        },
+        serverInfo: {
+            name: 'test-server',
+            version: '1.0.0'
+        }
+    }));
+
+    server.setRequestHandler(ListToolsRequestSchema, async () => ({
+        tools: []
+    }));
+
+    const client = new Client({
+        name: 'test-client',
+        version: '1.0.0',
+    }, {
+        toolListChangedOptions: {
+            autoRefresh: false,
+            onToolListChanged: (err, tools) => {
+                notifications.push([err, tools]);
+            }
+        }
+    });
+
+    const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+
+    await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
+
+    const result1 = await client.listTools();
+    expect(result1.tools).toHaveLength(0);
+
+    // Update the tools list
+    server.setRequestHandler(ListToolsRequestSchema, async () => ({
+        tools: [
+            {
+                name: 'test-tool',
+                description: 'A test tool',
+                inputSchema: {
+                    type: 'object',
+                    properties: {}
+                }
+                // No outputSchema
+            }
+        ]
+    }));
+    await server.sendToolListChanged();
+
+    // Wait for the debounced notifications to be processed
+    await new Promise(resolve => setTimeout(resolve, 1000));
+
+    // Should be 1 notification with no tool data because autoRefresh is false
+    expect(notifications).toHaveLength(1);
+    expect(notifications[0][0]).toBeNull();
+    expect(notifications[0][1]).toBeNull();
+});
+
 describe('outputSchema validation', () => {
     /***
      * Test: Validate structuredContent Against outputSchema
diff --git a/src/client/index.ts b/src/client/index.ts
@@ -36,7 +36,9 @@ import {
     SUPPORTED_PROTOCOL_VERSIONS,
     type SubscribeRequest,
     type Tool,
-    type UnsubscribeRequest
+    type UnsubscribeRequest,
+    ToolListChangedNotificationSchema,
+    ToolListChangedOptions
 } from '../types.js';
 import { AjvJsonSchemaValidator } from '../validation/ajv-provider.js';
 import type { JsonSchemaType, JsonSchemaValidator, jsonSchemaValidator } from '../validation/types.js';
@@ -77,6 +79,41 @@ export type ClientOptions = ProtocolOptions & {
      * ```
      */
     jsonSchemaValidator?: jsonSchemaValidator;
+
+    /**
+     * Configure automatic refresh behavior for tool list changed notifications
+     *
+     * @example 
+     * ```ts
+     * {
+     *   autoRefresh: true,
+     *   debounceMs: 300,
+     *   onToolListChanged: (err, tools) => {
+     *     if (err) {
+     *       console.error('Failed to refresh tool list:', err);
+     *       return;
+     *     }
+     *     // Use the updated tool list
+     *     console.log('Tool list changed:', tools);
+     *   }
+     * }
+     * ```
+     *
+     * @example
+     * ```ts
+     * {
+     *   autoRefresh: false,
+     *   onToolListChanged: (err, tools) => {
+     *     // err is always null when autoRefresh is false
+     * 
+     *     // Manually refresh the tool list
+     *     const result = await this.listTools();
+     *     console.log('Tool list changed:', result.tools);
+     *   }
+     * }
+     * ```
+     */
+    toolListChangedOptions?: ToolListChangedOptions;
 };
 
 /**
@@ -115,6 +152,8 @@ export class Client<
     private _instructions?: string;
     private _jsonSchemaValidator: jsonSchemaValidator;
     private _cachedToolOutputValidators: Map<string, JsonSchemaValidator<unknown>> = new Map();
+    private _toolListChangedOptions: ToolListChangedOptions | null = null;
+    private _toolListChangedDebounceTimer?: ReturnType<typeof setTimeout>;
 
     /**
      * Initializes this client with the given name and version information.
@@ -126,6 +165,9 @@ export class Client<
         super(options);
         this._capabilities = options?.capabilities ?? {};
         this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new AjvJsonSchemaValidator();
+
+        // Set up tool list changed options
+        this.setToolListChangedOptions(options?.toolListChangedOptions || null);
     }
 
     /**
@@ -434,6 +476,60 @@ export class Client<
         return result;
     }
 
+    /**
+     * Updates the tool list changed options
+     *
+     * Set to null to disable tool list changed notifications
+     */
+    public setToolListChangedOptions(options: ToolListChangedOptions | null): void {
+        // Set up tool list changed options and add notification handler
+        if (options) {
+            const toolListChangedOptions: ToolListChangedOptions = {
+                autoRefresh: !!options.autoRefresh,
+                debounceMs: options.debounceMs ?? 300,
+                onToolListChanged: options.onToolListChanged,
+            };
+            this._toolListChangedOptions = toolListChangedOptions;
+            this.setNotificationHandler(ToolListChangedNotificationSchema, () => {
+                // If autoRefresh is false, call the callback for the notification, but without tools data
+                if (!toolListChangedOptions.autoRefresh) {
+                    toolListChangedOptions.onToolListChanged?.(null, null);
+                    return;
+                }
+        
+                // Clear any pending debounce timer
+                if (this._toolListChangedDebounceTimer) {
+                    clearTimeout(this._toolListChangedDebounceTimer);
+                }
+        
+                // Set up debounced refresh
+                this._toolListChangedDebounceTimer = setTimeout(async () => {
+                    let tools: Tool[] | null = null;
+                    let error: Error | null = null;
+                    try {
+                        const result = await this.listTools();
+                        tools = result.tools;
+                    } catch (e) {
+                        error = e instanceof Error ? e : new Error(String(e));
+                    }
+                    toolListChangedOptions.onToolListChanged?.(error, tools);
+                }, toolListChangedOptions.debounceMs);
+            });
+        }
+        // Reset tool list changed options and remove notification handler
+        else {
+            this._toolListChangedOptions = null;
+            this.removeNotificationHandler(ToolListChangedNotificationSchema.shape.method.value);
+        }
+    }
+
+    /**
+     * Gets the current tool list changed options
+     */
+    public getToolListChangedOptions(): ToolListChangedOptions | null {
+        return this._toolListChangedOptions;
+    }
+
     async sendRootsListChanged() {
         return this.notification({ method: 'notifications/roots/list_changed' });
     }
diff --git a/src/types.ts b/src/types.ts
@@ -1040,6 +1040,36 @@ export const ToolListChangedNotificationSchema = NotificationSchema.extend({
     method: z.literal('notifications/tools/list_changed')
 });
 
+/**
+ * Client Options for tool list changed notifications.
+ */
+export const ToolListChangedOptionsSchema = z.object({
+    /**
+     * If true, the tool list will be refreshed automatically when a tool list changed notification is received.
+     *
+     * If `onToolListChanged` is also provided, it will be called after the tool list is auto refreshed.
+     *
+     * @default false
+     */
+    autoRefresh: z.boolean().optional(),
+    /**
+     * Debounce time in milliseconds for tool list changed notification processing.
+     *
+     * Multiple notifications received within this timeframe will only trigger one refresh.
+     *
+     * @default 300
+     */
+    debounceMs: z.number().int().optional(),
+    /**
+     * This callback is always called when the server sends a tool list changed notification.
+     *
+     * If `autoRefresh` is true, this callback will be called with updated tool list.
+     */
+    onToolListChanged: z.function(z.tuple([z.instanceof(Error).nullable(), z.array(ToolSchema).nullable()]), z.void()),
+});
+
+export type ToolListChangedOptions = z.infer<typeof ToolListChangedOptionsSchema>;
+
 /* Logging */
 /**
  * The severity of a log message.
