Add MCP Root change notifications (#6502)
This commit is contained in:
parent
465ac9f547
commit
3960ccf781
|
@ -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];
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in New Issue