Add MCP Root change notifications (#6502)

2 files changed, 74 insertions, 3 deletions

diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts
index e54c3e0f..3467ad95 100644
--- a/packages/core/src/tools/mcp-client.test.ts
+++ b/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];
diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts
index 87d38815..e9001466 100644
--- a/packages/core/src/tools/mcp-client.ts
+++ b/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) {