MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-26 13:04:49 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: Add support for MCP Resources (#13178)

Browse Source 
Co-authored-by: Jack Wotherspoon <jackwoth@google.com>
This commit is contained in:
Alex Gavrilescu
2025-12-09 03:43:12 +01:00
committed by GitHub
parent 720b31cb8b
commit 560550f5df
20 changed files with 1146 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/config/config.ts
+7
View File
@@ -17,6 +17,7 @@ import {
createContentGeneratorConfig,
} from '../core/contentGenerator.js';
import { PromptRegistry } from '../prompts/prompt-registry.js';
import { ResourceRegistry } from '../resources/resource-registry.js';
import { ToolRegistry } from '../tools/tool-registry.js';
import { LSTool } from '../tools/ls.js';
import { ReadFileTool } from '../tools/read-file.js';
@@ -331,6 +332,7 @@ export class Config {
private allowedMcpServers: string[];
private blockedMcpServers: string[];
private promptRegistry!: PromptRegistry;
private resourceRegistry!: ResourceRegistry;
private agentRegistry!: AgentRegistry;
private sessionId: string;
private fileSystemService: FileSystemService;
@@ -656,6 +658,7 @@ export class Config {
await this.getGitService();
}
this.promptRegistry = new PromptRegistry();
this.resourceRegistry = new ResourceRegistry();
this.agentRegistry = new AgentRegistry(this);
await this.agentRegistry.initialize();
@@ -921,6 +924,10 @@ export class Config {
return this.promptRegistry;
}
getResourceRegistry(): ResourceRegistry {
return this.resourceRegistry;
}
getDebugMode(): boolean {
return this.debugMode;
}
packages/core/src/index.ts
+1
View File
@@ -103,6 +103,7 @@ export * from './tools/tools.js';
export * from './tools/tool-error.js';
export * from './tools/tool-registry.js';
export * from './tools/tool-names.js';
export * from './resources/resource-registry.js';
// Export prompt logic
export * from './prompts/mcp-prompts.js';
packages/core/src/resources/resource-registry.test.ts
+76
View File
@@ -0,0 +1,76 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, expect, it, beforeEach } from 'vitest';
import type { Resource } from '@modelcontextprotocol/sdk/types.js';
import { ResourceRegistry } from './resource-registry.js';
describe('ResourceRegistry', () => {
let registry: ResourceRegistry;
beforeEach(() => {
registry = new ResourceRegistry();
});
const createResource = (overrides: Partial<Resource> = {}): Resource => ({
uri: 'file:///tmp/foo.txt',
name: 'foo',
description: 'example resource',
mimeType: 'text/plain',
...overrides,
});
it('stores resources per server', () => {
registry.setResourcesForServer('a', [createResource()]);
registry.setResourcesForServer('b', [createResource({ uri: 'foo' })]);
expect(
registry.getAllResources().filter((res) => res.serverName === 'a'),
).toHaveLength(1);
expect(
registry.getAllResources().filter((res) => res.serverName === 'b'),
).toHaveLength(1);
});
it('clears resources for server before adding new ones', () => {
registry.setResourcesForServer('a', [
createResource(),
createResource({ uri: 'bar' }),
]);
registry.setResourcesForServer('a', [createResource({ uri: 'baz' })]);
const resources = registry
.getAllResources()
.filter((res) => res.serverName === 'a');
expect(resources).toHaveLength(1);
expect(resources[0].uri).toBe('baz');
});
it('finds resources by serverName:uri identifier', () => {
registry.setResourcesForServer('a', [createResource()]);
registry.setResourcesForServer('b', [
createResource({ uri: 'file:///tmp/bar.txt' }),
]);
expect(
registry.findResourceByUri('b:file:///tmp/bar.txt')?.serverName,
).toBe('b');
expect(
registry.findResourceByUri('a:file:///tmp/foo.txt')?.serverName,
).toBe('a');
expect(registry.findResourceByUri('a:file:///tmp/bar.txt')).toBeUndefined();
expect(registry.findResourceByUri('nonexistent')).toBeUndefined();
});
it('clears resources for a server', () => {
registry.setResourcesForServer('a', [createResource()]);
registry.removeResourcesByServer('a');
expect(
registry.getAllResources().filter((res) => res.serverName === 'a'),
).toHaveLength(0);
});
});
packages/core/src/resources/resource-registry.ts
+72
View File
@@ -0,0 +1,72 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Resource } from '@modelcontextprotocol/sdk/types.js';
const resourceKey = (serverName: string, uri: string): string =>
`${serverName}::${uri}`;
export interface MCPResource extends Resource {
serverName: string;
discoveredAt: number;
}
export type DiscoveredMCPResource = MCPResource;
/**
* Tracks resources discovered from MCP servers so other
* components can query or include them in conversations.
*/
export class ResourceRegistry {
private resources: Map<string, MCPResource> = new Map();
/**
* Replace the resources for a specific server.
*/
setResourcesForServer(serverName: string, resources: Resource[]): void {
this.removeResourcesByServer(serverName);
const discoveredAt = Date.now();
for (const resource of resources) {
if (!resource.uri) {
continue;
}
this.resources.set(resourceKey(serverName, resource.uri), {
serverName,
discoveredAt,
...resource,
});
}
}
getAllResources(): MCPResource[] {
return Array.from(this.resources.values());
}
/**
* Find a resource by its identifier.
* Format: serverName:uri (e.g., "myserver:file:///data.txt")
*/
findResourceByUri(identifier: string): MCPResource | undefined {
const colonIndex = identifier.indexOf(':');
if (colonIndex <= 0) {
return undefined;
}
const serverName = identifier.substring(0, colonIndex);
const uri = identifier.substring(colonIndex + 1);
return this.resources.get(resourceKey(serverName, uri));
}
removeResourcesByServer(serverName: string): void {
for (const key of Array.from(this.resources.keys())) {
if (key.startsWith(`${serverName}::`)) {
this.resources.delete(key);
}
}
}
clear(): void {
this.resources.clear();
}
}
packages/core/src/tools/mcp-client-manager.test.ts
+13 -10
View File
@@ -29,6 +29,7 @@ vi.mock('./mcp-client.js', async () => {
describe('McpClientManager', () => {
let mockedMcpClient: MockedObject<McpClient>;
let mockConfig: MockedObject<Config>;
let toolRegistry: ToolRegistry;
beforeEach(() => {
mockedMcpClient = vi.mockObject({
@@ -43,6 +44,7 @@ describe('McpClientManager', () => {
isTrustedFolder: vi.fn().mockReturnValue(true),
getMcpServers: vi.fn().mockReturnValue({}),
getPromptRegistry: () => {},
getResourceRegistry: () => {},
getDebugMode: () => false,
getWorkspaceContext: () => {},
getAllowedMcpServers: vi.fn().mockReturnValue([]),
@@ -52,6 +54,7 @@ describe('McpClientManager', () => {
isInitialized: vi.fn(),
}),
} as unknown as Config);
toolRegistry = {} as ToolRegistry;
});
afterEach(() => {
@@ -62,7 +65,7 @@ describe('McpClientManager', () => {
mockConfig.getMcpServers.mockReturnValue({
'test-server': {},
});
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
@@ -73,7 +76,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockConfig.isTrustedFolder.mockReturnValue(false);
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
@@ -84,7 +87,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockConfig.getBlockedMcpServers.mockReturnValue(['test-server']);
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
@@ -96,14 +99,14 @@ describe('McpClientManager', () => {
'another-server': {},
});
mockConfig.getAllowedMcpServers.mockReturnValue(['another-server']);
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
});
it('should start servers from extensions', async () => {
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startExtension({
name: 'test-extension',
mcpServers: {
@@ -120,7 +123,7 @@ describe('McpClientManager', () => {
});
it('should not start servers from disabled extensions', async () => {
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startExtension({
name: 'test-extension',
mcpServers: {
@@ -141,7 +144,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockConfig.getBlockedMcpServers.mockReturnValue(['test-server']);
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(manager.getBlockedMcpServers()).toEqual([
{ name: 'test-server', extensionName: '' },
@@ -154,7 +157,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockedMcpClient.getServerConfig.mockReturnValue({});
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(1);
@@ -173,7 +176,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockedMcpClient.getServerConfig.mockReturnValue({});
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(1);
@@ -187,7 +190,7 @@ describe('McpClientManager', () => {
});
it('should throw an error if the server does not exist', async () => {
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await expect(manager.restartServer('non-existent')).rejects.toThrow(
'No MCP server registered with the name "non-existent"',
);
packages/core/src/tools/mcp-client-manager.ts
+5
View File
@@ -52,6 +52,10 @@ export class McpClientManager {
return this.blockedMcpServers;
}
getClient(serverName: string): McpClient | undefined {
return this.clients.get(serverName);
}
/**
* For all the MCP servers associated with this extension:
*
@@ -174,6 +178,7 @@ export class McpClientManager {
config,
this.toolRegistry,
this.cliConfig.getPromptRegistry(),
this.cliConfig.getResourceRegistry(),
this.cliConfig.getWorkspaceContext(),
this.cliConfig,
this.cliConfig.getDebugMode(),
packages/core/src/tools/mcp-client.test.ts
+243 -8
View File
@@ -30,6 +30,7 @@ import {
populateMcpServerCommand,
} from './mcp-client.js';
import type { ToolRegistry } from './tool-registry.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
import * as fs from 'node:fs';
import * as os from 'node:os';
import * as path from 'node:path';
@@ -77,6 +78,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
listTools: vi.fn().mockResolvedValue({
tools: [
@@ -105,13 +107,22 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -135,6 +146,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
listTools: vi.fn().mockResolvedValue({
@@ -175,13 +187,22 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -201,6 +222,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
listTools: vi.fn().mockResolvedValue({ tools: [] }),
listPrompts: vi.fn().mockRejectedValue(new Error('Test error')),
@@ -216,20 +238,29 @@ describe('mcp-client', () => {
registerTool: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
);
await client.connect();
await expect(client.discover({} as Config)).rejects.toThrow(
'No prompts or tools found on the server.',
'No prompts, tools, or resources found on the server.',
);
expect(coreEvents.emitFeedback).toHaveBeenCalledWith(
'error',
@@ -246,6 +277,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
request: vi.fn().mockResolvedValue({}),
@@ -261,20 +293,29 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
);
await client.connect();
await expect(client.discover({} as Config)).rejects.toThrow(
'No prompts or tools found on the server.',
'No prompts, tools, or resources found on the server.',
);
});
@@ -286,6 +327,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
listTools: vi.fn().mockResolvedValue({
tools: [
@@ -310,13 +352,22 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -334,6 +385,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
listTools: vi.fn().mockResolvedValue({
tools: [
@@ -373,13 +425,22 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -405,6 +466,165 @@ describe('mcp-client', () => {
});
});
it('should discover resources when a server only exposes resources', async () => {
const mockedClient = {
connect: vi.fn(),
discover: vi.fn(),
disconnect: vi.fn(),
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ resources: {} }),
request: vi.fn().mockImplementation(({ method }) => {
if (method === 'resources/list') {
return Promise.resolve({
resources: [
{
uri: 'file:///tmp/resource.txt',
name: 'resource',
description: 'Test Resource',
mimeType: 'text/plain',
},
],
});
}
return Promise.resolve({ prompts: [] });
}),
} as unknown as ClientLib.Client;
vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
{} as SdkClientStdioLib.StdioClientTransport,
);
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
);
await client.connect();
await client.discover({} as Config);
expect(resourceRegistry.setResourcesForServer).toHaveBeenCalledWith(
'test-server',
[
expect.objectContaining({
uri: 'file:///tmp/resource.txt',
name: 'resource',
}),
],
);
});
it('refreshes registry when resource list change notification is received', async () => {
let listCallCount = 0;
let resourceListHandler:
| ((notification: unknown) => Promise<void> | void)
| undefined;
const mockedClient = {
connect: vi.fn(),
discover: vi.fn(),
disconnect: vi.fn(),
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn((_, handler) => {
resourceListHandler = handler;
}),
getServerCapabilities: vi
.fn()
.mockReturnValue({ resources: { listChanged: true } }),
request: vi.fn().mockImplementation(({ method }) => {
if (method === 'resources/list') {
listCallCount += 1;
if (listCallCount === 1) {
return Promise.resolve({
resources: [
{
uri: 'file:///tmp/one.txt',
},
],
});
}
return Promise.resolve({
resources: [
{
uri: 'file:///tmp/two.txt',
},
],
});
}
return Promise.resolve({ prompts: [] });
}),
} as unknown as ClientLib.Client;
vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
{} as SdkClientStdioLib.StdioClientTransport,
);
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
);
await client.connect();
await client.discover({} as Config);
expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce();
expect(resourceListHandler).toBeDefined();
await resourceListHandler?.({
method: 'notifications/resources/list_changed',
});
expect(resourceRegistry.setResourcesForServer).toHaveBeenLastCalledWith(
'test-server',
[expect.objectContaining({ uri: 'file:///tmp/two.txt' })],
);
expect(coreEvents.emitFeedback).toHaveBeenCalledWith(
'info',
'Resources updated for server: test-server',
);
});
it('should remove tools and prompts on disconnect', async () => {
const mockedClient = {
connect: vi.fn(),
@@ -412,6 +632,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi
.fn()
.mockReturnValue({ tools: {}, prompts: {} }),
@@ -447,6 +668,10 @@ describe('mcp-client', () => {
unregisterPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
@@ -454,6 +679,7 @@ describe('mcp-client', () => {
},
mockedToolRegistry,
mockedPromptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -469,6 +695,7 @@ describe('mcp-client', () => {
expect(mockedClient.close).toHaveBeenCalledOnce();
expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledOnce();
expect(mockedPromptRegistry.removePromptsByServer).toHaveBeenCalledOnce();
expect(resourceRegistry.removeResourcesByServer).toHaveBeenCalledOnce();
});
});
@@ -501,6 +728,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
{} as ToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -536,6 +764,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
{} as ToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -590,6 +819,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -659,6 +889,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -728,6 +959,7 @@ describe('mcp-client', () => {
{ command: 'cmd-a' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -739,6 +971,7 @@ describe('mcp-client', () => {
{ command: 'cmd-b' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -820,6 +1053,7 @@ describe('mcp-client', () => {
{ command: 'test-command', timeout: 100 },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -883,6 +1117,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
packages/core/src/tools/mcp-client.ts
+192 -28
View File
@@ -20,9 +20,14 @@ import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
import type {
GetPromptResult,
Prompt,
ReadResourceResult,
Resource,
} from '@modelcontextprotocol/sdk/types.js';
import {
ListResourcesResultSchema,
ListRootsRequestSchema,
ReadResourceResultSchema,
ResourceListChangedNotificationSchema,
ToolListChangedNotificationSchema,
type Tool as McpTool,
} from '@modelcontextprotocol/sdk/types.js';
@@ -54,6 +59,7 @@ import type { ToolRegistry } from './tool-registry.js';
import { debugLogger } from '../utils/debugLogger.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { coreEvents } from '../utils/events.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
@@ -98,14 +104,17 @@ export class McpClient {
private client: Client | undefined;
private transport: Transport | undefined;
private status: MCPServerStatus = MCPServerStatus.DISCONNECTED;
private isRefreshing: boolean = false;
private pendingRefresh: boolean = false;
private isRefreshingTools: boolean = false;
private pendingToolRefresh: boolean = false;
private isRefreshingResources: boolean = false;
private pendingResourceRefresh: boolean = false;
constructor(
private readonly serverName: string,
private readonly serverConfig: MCPServerConfig,
private readonly toolRegistry: ToolRegistry,
private readonly promptRegistry: PromptRegistry,
private readonly resourceRegistry: ResourceRegistry,
private readonly workspaceContext: WorkspaceContext,
private readonly cliConfig: Config,
private readonly debugMode: boolean,
@@ -130,24 +139,8 @@ export class McpClient {
this.workspaceContext,
);
// setup dynamic tool listener
const capabilities = this.client.getServerCapabilities();
this.registerNotificationHandlers();
if (capabilities?.tools?.listChanged) {
debugLogger.log(
`Server '${this.serverName}' supports tool updates. Listening for changes...`,
);
this.client.setNotificationHandler(
ToolListChangedNotificationSchema,
async () => {
debugLogger.log(
`🔔 Received tool update notification from '${this.serverName}'`,
);
await this.refreshTools();
},
);
}
const originalOnError = this.client.onerror;
this.client.onerror = (error) => {
if (this.status !== MCPServerStatus.CONNECTED) {
@@ -176,9 +169,11 @@ export class McpClient {
const prompts = await this.discoverPrompts();
const tools = await this.discoverTools(cliConfig);
const resources = await this.discoverResources();
this.updateResourceRegistry(resources);
if (prompts.length === 0 && tools.length === 0) {
throw new Error('No prompts or tools found on the server.');
if (prompts.length === 0 && tools.length === 0 && resources.length === 0) {
throw new Error('No prompts, tools, or resources found on the server.');
}
for (const tool of tools) {
@@ -196,6 +191,7 @@ export class McpClient {
}
this.toolRegistry.removeMcpToolsByServer(this.serverName);
this.promptRegistry.removePromptsByServer(this.serverName);
this.resourceRegistry.removeResourcesByServer(this.serverName);
this.updateStatus(MCPServerStatus.DISCONNECTING);
const client = this.client;
this.client = undefined;
@@ -250,6 +246,128 @@ export class McpClient {
return discoverPrompts(this.serverName, this.client!, this.promptRegistry);
}
private async discoverResources(): Promise<Resource[]> {
this.assertConnected();
return discoverResources(this.serverName, this.client!);
}
private updateResourceRegistry(resources: Resource[]): void {
this.resourceRegistry.setResourcesForServer(this.serverName, resources);
}
async readResource(uri: string): Promise<ReadResourceResult> {
this.assertConnected();
return this.client!.request(
{
method: 'resources/read',
params: { uri },
},
ReadResourceResultSchema,
);
}
/**
* Registers notification handlers for dynamic updates from the MCP server.
* This includes handlers for tool list changes and resource list changes.
*/
private registerNotificationHandlers(): void {
if (!this.client) {
return;
}
const capabilities = this.client.getServerCapabilities();
if (capabilities?.tools?.listChanged) {
debugLogger.log(
`Server '${this.serverName}' supports tool updates. Listening for changes...`,
);
this.client.setNotificationHandler(
ToolListChangedNotificationSchema,
async () => {
debugLogger.log(
`🔔 Received tool update notification from '${this.serverName}'`,
);
await this.refreshTools();
},
);
}
if (capabilities?.resources?.listChanged) {
debugLogger.log(
`Server '${this.serverName}' supports resource updates. Listening for changes...`,
);
this.client.setNotificationHandler(
ResourceListChangedNotificationSchema,
async () => {
debugLogger.log(
`🔔 Received resource update notification from '${this.serverName}'`,
);
await this.refreshResources();
},
);
}
}
/**
* Refreshes the resources for this server by re-querying the MCP `resources/list` endpoint.
*
* This method implements a **Coalescing Pattern** to handle rapid bursts of notifications
* (e.g., during server startup or bulk updates) without overwhelming the server or
* creating race conditions in the ResourceRegistry.
*/
private async refreshResources(): Promise<void> {
if (this.isRefreshingResources) {
debugLogger.log(
`Resource refresh for '${this.serverName}' is already in progress. Pending update.`,
);
this.pendingResourceRefresh = true;
return;
}
this.isRefreshingResources = true;
try {
do {
this.pendingResourceRefresh = false;
if (this.status !== MCPServerStatus.CONNECTED || !this.client) break;
const timeoutMs = this.serverConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC;
const abortController = new AbortController();
const timeoutId = setTimeout(() => abortController.abort(), timeoutMs);
let newResources;
try {
newResources = await this.discoverResources();
} catch (err) {
debugLogger.error(
`Resource discovery failed during refresh: ${getErrorMessage(err)}`,
);
clearTimeout(timeoutId);
break;
}
this.updateResourceRegistry(newResources);
clearTimeout(timeoutId);
coreEvents.emitFeedback(
'info',
`Resources updated for server: ${this.serverName}`,
);
} while (this.pendingResourceRefresh);
} catch (error) {
debugLogger.error(
`Critical error in resource refresh loop for ${this.serverName}: ${getErrorMessage(error)}`,
);
} finally {
this.isRefreshingResources = false;
this.pendingResourceRefresh = false;
}
}
getServerConfig(): MCPServerConfig {
return this.serverConfig;
}
@@ -266,19 +384,19 @@ export class McpClient {
* creating race conditions in the global ToolRegistry.
*/
private async refreshTools(): Promise<void> {
if (this.isRefreshing) {
if (this.isRefreshingTools) {
debugLogger.log(
`Tool refresh for '${this.serverName}' is already in progress. Pending update.`,
);
this.pendingRefresh = true;
this.pendingToolRefresh = true;
return;
}
this.isRefreshing = true;
this.isRefreshingTools = true;
try {
do {
this.pendingRefresh = false;
this.pendingToolRefresh = false;
if (this.status !== MCPServerStatus.CONNECTED || !this.client) break;
@@ -316,14 +434,14 @@ export class McpClient {
'info',
`Tools updated for server: ${this.serverName}`,
);
} while (this.pendingRefresh);
} while (this.pendingToolRefresh);
} catch (error) {
debugLogger.error(
`Critical error in refresh loop for ${this.serverName}: ${getErrorMessage(error)}`,
);
} finally {
this.isRefreshing = false;
this.pendingRefresh = false;
this.isRefreshingTools = false;
this.pendingToolRefresh = false;
}
}
}
@@ -944,6 +1062,52 @@ export async function discoverPrompts(
}
}
export async function discoverResources(
mcpServerName: string,
mcpClient: Client,
): Promise<Resource[]> {
if (mcpClient.getServerCapabilities()?.resources == null) {
return [];
}
const resources = await listResources(mcpServerName, mcpClient);
return resources;
}
async function listResources(
mcpServerName: string,
mcpClient: Client,
): Promise<Resource[]> {
const resources: Resource[] = [];
let cursor: string | undefined;
try {
do {
const response = await mcpClient.request(
{
method: 'resources/list',
params: cursor ? { cursor } : {},
},
ListResourcesResultSchema,
);
resources.push(...(response.resources ?? []));
cursor = response.nextCursor ?? undefined;
} while (cursor);
} catch (error) {
if (error instanceof Error && error.message?.includes('Method not found')) {
return [];
}
coreEvents.emitFeedback(
'error',
`Error discovering resources from ${mcpServerName}: ${getErrorMessage(
error,
)}`,
error,
);
throw error;
}
return resources;
}
/**
* Invokes a prompt on a connected MCP client.
*