Add MCP Root change notifications (#6502)

packages/core/src/tools/mcp-client.test.ts

@@ -280,6 +280,46 @@ describe('mcp-client', () => {
   });
 
   describe('connectToMcpServer', () => {
+    it('should send a notification when directories change', async () => {
+      const mockedClient = {
+        registerCapabilities: vi.fn(),
+        setRequestHandler: vi.fn(),
+        notification: vi.fn(),
+        callTool: vi.fn(),
+        connect: vi.fn(),
+      };
+      vi.mocked(ClientLib.Client).mockReturnValue(
+        mockedClient as unknown as ClientLib.Client,
+      );
+      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
+        {} as SdkClientStdioLib.StdioClientTransport,
+      );
+      let onDirectoriesChangedCallback: () => void = () => {};
+      const mockWorkspaceContext = {
+        getDirectories: vi
+          .fn()
+          .mockReturnValue(['/test/dir', '/another/project']),
+        onDirectoriesChanged: vi.fn().mockImplementation((callback) => {
+          onDirectoriesChangedCallback = callback;
+        }),
+      } as unknown as WorkspaceContext;
+
+      await connectToMcpServer(
+        'test-server',
+        {
+          command: 'test-command',
+        },
+        false,
+        mockWorkspaceContext,
+      );
+
+      onDirectoriesChangedCallback();
+
+      expect(mockedClient.notification).toHaveBeenCalledWith({
+        method: 'notifications/roots/list_changed',
+      });
+    });
+
     it('should register a roots/list handler', async () => {
       const mockedClient = {
         registerCapabilities: vi.fn(),
@@ -297,6 +337,7 @@ describe('mcp-client', () => {
         getDirectories: vi
           .fn()
           .mockReturnValue(['/test/dir', '/another/project']),
+        onDirectoriesChanged: vi.fn(),
       } as unknown as WorkspaceContext;
 
       await connectToMcpServer(
@@ -309,7 +350,9 @@ describe('mcp-client', () => {
       );
 
       expect(mockedClient.registerCapabilities).toHaveBeenCalledWith({
-        roots: {},
+        roots: {
+          listChanged: true,
+        },
       });
       expect(mockedClient.setRequestHandler).toHaveBeenCalledOnce();
       const handler = mockedClient.setRequestHandler.mock.calls[0][1];

packages/core/src/tools/mcp-client.ts

@@ -36,7 +36,7 @@ import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
 import { getErrorMessage } from '../utils/errors.js';
 import { basename } from 'node:path';
 import { pathToFileURL } from 'node:url';
-import { WorkspaceContext } from '../utils/workspaceContext.js';
+import { Unsubscribe, WorkspaceContext } from '../utils/workspaceContext.js';
 
 export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
 
@@ -677,7 +677,9 @@ export async function connectToMcpServer(
   });
 
   mcpClient.registerCapabilities({
-    roots: {},
+    roots: {
+      listChanged: true,
+    },
   });
 
   mcpClient.setRequestHandler(ListRootsRequestSchema, async () => {
@@ -693,6 +695,32 @@ export async function connectToMcpServer(
     };
   });
 
+  let unlistenDirectories: Unsubscribe | undefined =
+    workspaceContext.onDirectoriesChanged(async () => {
+      try {
+        await mcpClient.notification({
+          method: 'notifications/roots/list_changed',
+        });
+      } catch (_) {
+        // If this fails, its almost certainly because the connection was closed
+        // and we should just stop listening for future directory changes.
+        unlistenDirectories?.();
+        unlistenDirectories = undefined;
+      }
+    });
+
+  // Attempt to pro-actively unsubscribe if the mcp client closes. This API is
+  // very brittle though so we don't have any guarantees, hence the try/catch
+  // above as well.
+  //
+  // Be a good steward and don't just bash over onclose.
+  const oldOnClose = mcpClient.onclose;
+  mcpClient.onclose = () => {
+    oldOnClose?.();
+    unlistenDirectories?.();
+    unlistenDirectories = undefined;
+  };
+
   // patch Client.callTool to use request timeout as genai McpCallTool.callTool does not do it
   // TODO: remove this hack once GenAI SDK does callTool with request options
   if ('callTool' in mcpClient) {

packages/core/src/utils/workspaceContext.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
 import * as fs from 'fs';
 import * as os from 'os';
 import * as path from 'path';
@@ -294,6 +294,88 @@ describe('WorkspaceContext with real filesystem', () => {
     });
   });
 
+  describe('onDirectoriesChanged', () => {
+    it('should call listener when adding a directory', () => {
+      const workspaceContext = new WorkspaceContext(cwd);
+      const listener = vi.fn();
+      workspaceContext.onDirectoriesChanged(listener);
+
+      workspaceContext.addDirectory(otherDir);
+
+      expect(listener).toHaveBeenCalledOnce();
+    });
+
+    it('should not call listener when adding a duplicate directory', () => {
+      const workspaceContext = new WorkspaceContext(cwd);
+      workspaceContext.addDirectory(otherDir);
+      const listener = vi.fn();
+      workspaceContext.onDirectoriesChanged(listener);
+
+      workspaceContext.addDirectory(otherDir);
+
+      expect(listener).not.toHaveBeenCalled();
+    });
+
+    it('should call listener when setting different directories', () => {
+      const workspaceContext = new WorkspaceContext(cwd);
+      const listener = vi.fn();
+      workspaceContext.onDirectoriesChanged(listener);
+
+      workspaceContext.setDirectories([otherDir]);
+
+      expect(listener).toHaveBeenCalledOnce();
+    });
+
+    it('should not call listener when setting same directories', () => {
+      const workspaceContext = new WorkspaceContext(cwd);
+      const listener = vi.fn();
+      workspaceContext.onDirectoriesChanged(listener);
+
+      workspaceContext.setDirectories([cwd]);
+
+      expect(listener).not.toHaveBeenCalled();
+    });
+
+    it('should support multiple listeners', () => {
+      const workspaceContext = new WorkspaceContext(cwd);
+      const listener1 = vi.fn();
+      const listener2 = vi.fn();
+      workspaceContext.onDirectoriesChanged(listener1);
+      workspaceContext.onDirectoriesChanged(listener2);
+
+      workspaceContext.addDirectory(otherDir);
+
+      expect(listener1).toHaveBeenCalledOnce();
+      expect(listener2).toHaveBeenCalledOnce();
+    });
+
+    it('should allow unsubscribing a listener', () => {
+      const workspaceContext = new WorkspaceContext(cwd);
+      const listener = vi.fn();
+      const unsubscribe = workspaceContext.onDirectoriesChanged(listener);
+
+      unsubscribe();
+      workspaceContext.addDirectory(otherDir);
+
+      expect(listener).not.toHaveBeenCalled();
+    });
+
+    it('should not fail if a listener throws an error', () => {
+      const workspaceContext = new WorkspaceContext(cwd);
+      const errorListener = () => {
+        throw new Error('test error');
+      };
+      const listener = vi.fn();
+      workspaceContext.onDirectoriesChanged(errorListener);
+      workspaceContext.onDirectoriesChanged(listener);
+
+      expect(() => {
+        workspaceContext.addDirectory(otherDir);
+      }).not.toThrow();
+      expect(listener).toHaveBeenCalledOnce();
+    });
+  });
+
   describe('getDirectories', () => {
     it('should return a copy of directories array', () => {
       const workspaceContext = new WorkspaceContext(cwd);

packages/core/src/utils/workspaceContext.ts

@@ -8,6 +8,8 @@ import { isNodeError } from '../utils/errors.js';
 import * as fs from 'fs';
 import * as path from 'path';
 
+export type Unsubscribe = () => void;
+
 /**
  * WorkspaceContext manages multiple workspace directories and validates paths
  * against them. This allows the CLI to operate on files from multiple directories
@@ -16,6 +18,7 @@ import * as path from 'path';
 export class WorkspaceContext {
   private directories = new Set<string>();
   private initialDirectories: Set<string>;
+  private onDirectoriesChangedListeners = new Set<() => void>();
 
   /**
    * Creates a new WorkspaceContext with the given initial directory and optional additional directories.
@@ -31,13 +34,42 @@ export class WorkspaceContext {
     this.initialDirectories = new Set(this.directories);
   }
 
+  /**
+   * Registers a listener that is called when the workspace directories change.
+   * @param listener The listener to call.
+   * @returns A function to unsubscribe the listener.
+   */
+  onDirectoriesChanged(listener: () => void): Unsubscribe {
+    this.onDirectoriesChangedListeners.add(listener);
+    return () => {
+      this.onDirectoriesChangedListeners.delete(listener);
+    };
+  }
+
+  private notifyDirectoriesChanged() {
+    // Iterate over a copy of the set in case a listener unsubscribes itself or others.
+    for (const listener of [...this.onDirectoriesChangedListeners]) {
+      try {
+        listener();
+      } catch (e) {
+        // Don't let one listener break others.
+        console.error('Error in WorkspaceContext listener:', e);
+      }
+    }
+  }
+
   /**
    * Adds a directory to the workspace.
    * @param directory The directory path to add (can be relative or absolute)
    * @param basePath Optional base path for resolving relative paths (defaults to cwd)
    */
   addDirectory(directory: string, basePath: string = process.cwd()): void {
-    this.directories.add(this.resolveAndValidateDir(directory, basePath));
+    const resolved = this.resolveAndValidateDir(directory, basePath);
+    if (this.directories.has(resolved)) {
+      return;
+    }
+    this.directories.add(resolved);
+    this.notifyDirectoriesChanged();
   }
 
   private resolveAndValidateDir(
@@ -72,9 +104,17 @@ export class WorkspaceContext {
   }
 
   setDirectories(directories: readonly string[]): void {
-    this.directories.clear();
+    const newDirectories = new Set<string>();
     for (const dir of directories) {
-      this.addDirectory(dir);
+      newDirectories.add(this.resolveAndValidateDir(dir));
+    }
+
+    if (
+      newDirectories.size !== this.directories.size ||
+      ![...newDirectories].every((d) => this.directories.has(d))
+    ) {
+      this.directories = newDirectories;
+      this.notifyDirectoriesChanged();
     }
   }