diff options
| -rw-r--r-- | packages/core/src/tools/mcp-client.test.ts | 45 | ||||
| -rw-r--r-- | packages/core/src/tools/mcp-client.ts | 32 | ||||
| -rw-r--r-- | packages/core/src/utils/workspaceContext.test.ts | 84 | ||||
| -rw-r--r-- | packages/core/src/utils/workspaceContext.ts | 46 |
4 files changed, 200 insertions, 7 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) { diff --git a/packages/core/src/utils/workspaceContext.test.ts b/packages/core/src/utils/workspaceContext.test.ts index 5446c287..30c9ae16 100644 --- a/packages/core/src/utils/workspaceContext.test.ts +++ b/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); diff --git a/packages/core/src/utils/workspaceContext.ts b/packages/core/src/utils/workspaceContext.ts index be99a83c..924cd601 100644 --- a/packages/core/src/utils/workspaceContext.ts +++ b/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. @@ -32,12 +35,41 @@ export class WorkspaceContext { } /** + * 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(); } } |