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

828 lines
25 KiB
TypeScript

/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { afterEach, describe, expect, it, vi, beforeEach } from 'vitest';
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
import {
populateMcpServerCommand,
createTransport,
isEnabled,
discoverTools,
discoverPrompts,
hasValidTypes,
connectToMcpServer,
} from './mcp-client.js';
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
import * as SdkClientStdioLib from '@modelcontextprotocol/sdk/client/stdio.js';
import * as ClientLib from '@modelcontextprotocol/sdk/client/index.js';
import * as GenAiLib from '@google/genai';
import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
import { AuthProviderType } from '../config/config.js';
import { PromptRegistry } from '../prompts/prompt-registry.js';
import { DiscoveredMCPTool } from './mcp-tool.js';
import { WorkspaceContext } from '../utils/workspaceContext.js';
import { pathToFileURL } from 'node:url';
vi.mock('@modelcontextprotocol/sdk/client/stdio.js');
vi.mock('@modelcontextprotocol/sdk/client/index.js');
vi.mock('@google/genai');
vi.mock('../mcp/oauth-provider.js');
vi.mock('../mcp/oauth-token-storage.js');
vi.mock('./mcp-tool.js');
describe('mcp-client', () => {
afterEach(() => {
vi.restoreAllMocks();
});
describe('discoverTools', () => {
it('should discover tools', async () => {
const mockedClient = {} as unknown as ClientLib.Client;
const mockedMcpToTool = vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
tool: () => ({
functionDeclarations: [
{
name: 'testFunction',
},
],
}),
} as unknown as GenAiLib.CallableTool);
const tools = await discoverTools('test-server', {}, mockedClient);
expect(tools.length).toBe(1);
expect(mockedMcpToTool).toHaveBeenCalledOnce();
});
it('should log an error if there is an error discovering a tool', async () => {
const mockedClient = {} as unknown as ClientLib.Client;
const consoleErrorSpy = vi
.spyOn(console, 'error')
.mockImplementation(() => {});
const testError = new Error('Invalid tool name');
vi.mocked(DiscoveredMCPTool).mockImplementation(
(
_mcpCallableTool: GenAiLib.CallableTool,
_serverName: string,
name: string,
) => {
if (name === 'invalid tool name') {
throw testError;
}
return { name: 'validTool' } as DiscoveredMCPTool;
},
);
vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
tool: () =>
Promise.resolve({
functionDeclarations: [
{
name: 'validTool',
},
{
name: 'invalid tool name', // this will fail validation
},
],
}),
} as unknown as GenAiLib.CallableTool);
const tools = await discoverTools('test-server', {}, mockedClient);
expect(tools.length).toBe(1);
expect(tools[0].name).toBe('validTool');
expect(consoleErrorSpy).toHaveBeenCalledOnce();
expect(consoleErrorSpy).toHaveBeenCalledWith(
`Error discovering tool: 'invalid tool name' from MCP server 'test-server': ${testError.message}`,
);
});
it('should skip tools if a parameter is missing a type', async () => {
const mockedClient = {} as unknown as ClientLib.Client;
const consoleWarnSpy = vi
.spyOn(console, 'warn')
.mockImplementation(() => {});
vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
tool: () =>
Promise.resolve({
functionDeclarations: [
{
name: 'validTool',
parametersJsonSchema: {
type: 'object',
properties: {
param1: { type: 'string' },
},
},
},
{
name: 'invalidTool',
parametersJsonSchema: {
type: 'object',
properties: {
param1: { description: 'a param with no type' },
},
},
},
],
}),
} as unknown as GenAiLib.CallableTool);
const tools = await discoverTools('test-server', {}, mockedClient);
expect(tools.length).toBe(1);
expect(vi.mocked(DiscoveredMCPTool).mock.calls[0][2]).toBe('validTool');
expect(consoleWarnSpy).toHaveBeenCalledOnce();
expect(consoleWarnSpy).toHaveBeenCalledWith(
`Skipping tool 'invalidTool' from MCP server 'test-server' because it has ` +
`missing types in its parameter schema. Please file an issue with the owner of the MCP server.`,
);
consoleWarnSpy.mockRestore();
});
it('should skip tools if a nested parameter is missing a type', async () => {
const mockedClient = {} as unknown as ClientLib.Client;
const consoleWarnSpy = vi
.spyOn(console, 'warn')
.mockImplementation(() => {});
vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
tool: () =>
Promise.resolve({
functionDeclarations: [
{
name: 'invalidTool',
parametersJsonSchema: {
type: 'object',
properties: {
param1: {
type: 'object',
properties: {
nestedParam: {
description: 'a nested param with no type',
},
},
},
},
},
},
],
}),
} as unknown as GenAiLib.CallableTool);
const tools = await discoverTools('test-server', {}, mockedClient);
expect(tools.length).toBe(0);
expect(consoleWarnSpy).toHaveBeenCalledOnce();
expect(consoleWarnSpy).toHaveBeenCalledWith(
`Skipping tool 'invalidTool' from MCP server 'test-server' because it has ` +
`missing types in its parameter schema. Please file an issue with the owner of the MCP server.`,
);
consoleWarnSpy.mockRestore();
});
it('should skip tool if an array item is missing a type', async () => {
const mockedClient = {} as unknown as ClientLib.Client;
const consoleWarnSpy = vi
.spyOn(console, 'warn')
.mockImplementation(() => {});
vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
tool: () =>
Promise.resolve({
functionDeclarations: [
{
name: 'invalidTool',
parametersJsonSchema: {
type: 'object',
properties: {
param1: {
type: 'array',
items: {
description: 'an array item with no type',
},
},
},
},
},
],
}),
} as unknown as GenAiLib.CallableTool);
const tools = await discoverTools('test-server', {}, mockedClient);
expect(tools.length).toBe(0);
expect(consoleWarnSpy).toHaveBeenCalledOnce();
expect(consoleWarnSpy).toHaveBeenCalledWith(
`Skipping tool 'invalidTool' from MCP server 'test-server' because it has ` +
`missing types in its parameter schema. Please file an issue with the owner of the MCP server.`,
);
consoleWarnSpy.mockRestore();
});
it('should discover tool with no properties in schema', async () => {
const mockedClient = {} as unknown as ClientLib.Client;
const consoleWarnSpy = vi
.spyOn(console, 'warn')
.mockImplementation(() => {});
vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
tool: () =>
Promise.resolve({
functionDeclarations: [
{
name: 'validTool',
parametersJsonSchema: {
type: 'object',
},
},
],
}),
} as unknown as GenAiLib.CallableTool);
const tools = await discoverTools('test-server', {}, mockedClient);
expect(tools.length).toBe(1);
expect(vi.mocked(DiscoveredMCPTool).mock.calls[0][2]).toBe('validTool');
expect(consoleWarnSpy).not.toHaveBeenCalled();
consoleWarnSpy.mockRestore();
});
it('should discover tool with empty properties object in schema', async () => {
const mockedClient = {} as unknown as ClientLib.Client;
const consoleWarnSpy = vi
.spyOn(console, 'warn')
.mockImplementation(() => {});
vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
tool: () =>
Promise.resolve({
functionDeclarations: [
{
name: 'validTool',
parametersJsonSchema: {
type: 'object',
properties: {},
},
},
],
}),
} as unknown as GenAiLib.CallableTool);
const tools = await discoverTools('test-server', {}, mockedClient);
expect(tools.length).toBe(1);
expect(vi.mocked(DiscoveredMCPTool).mock.calls[0][2]).toBe('validTool');
expect(consoleWarnSpy).not.toHaveBeenCalled();
consoleWarnSpy.mockRestore();
});
});
describe('connectToMcpServer', () => {
it('should register a roots/list handler', async () => {
const mockedClient = {
registerCapabilities: vi.fn(),
setRequestHandler: 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,
);
const mockWorkspaceContext = {
getDirectories: vi
.fn()
.mockReturnValue(['/test/dir', '/another/project']),
} as unknown as WorkspaceContext;
await connectToMcpServer(
'test-server',
{
command: 'test-command',
},
false,
mockWorkspaceContext,
);
expect(mockedClient.registerCapabilities).toHaveBeenCalledWith({
roots: {},
});
expect(mockedClient.setRequestHandler).toHaveBeenCalledOnce();
const handler = mockedClient.setRequestHandler.mock.calls[0][1];
const roots = await handler();
expect(roots).toEqual({
roots: [
{
uri: pathToFileURL('/test/dir').toString(),
name: 'dir',
},
{
uri: pathToFileURL('/another/project').toString(),
name: 'project',
},
],
});
});
});
describe('discoverPrompts', () => {
const mockedPromptRegistry = {
registerPrompt: vi.fn(),
} as unknown as PromptRegistry;
it('should discover and log prompts', async () => {
const mockRequest = vi.fn().mockResolvedValue({
prompts: [
{ name: 'prompt1', description: 'desc1' },
{ name: 'prompt2' },
],
});
const mockGetServerCapabilities = vi.fn().mockReturnValue({
prompts: {},
});
const mockedClient = {
getServerCapabilities: mockGetServerCapabilities,
request: mockRequest,
} as unknown as ClientLib.Client;
await discoverPrompts('test-server', mockedClient, mockedPromptRegistry);
expect(mockGetServerCapabilities).toHaveBeenCalledOnce();
expect(mockRequest).toHaveBeenCalledWith(
{ method: 'prompts/list', params: {} },
expect.anything(),
);
});
it('should do nothing if no prompts are discovered', async () => {
const mockRequest = vi.fn().mockResolvedValue({
prompts: [],
});
const mockGetServerCapabilities = vi.fn().mockReturnValue({
prompts: {},
});
const mockedClient = {
getServerCapabilities: mockGetServerCapabilities,
request: mockRequest,
} as unknown as ClientLib.Client;
const consoleLogSpy = vi
.spyOn(console, 'debug')
.mockImplementation(() => {});
await discoverPrompts('test-server', mockedClient, mockedPromptRegistry);
expect(mockGetServerCapabilities).toHaveBeenCalledOnce();
expect(mockRequest).toHaveBeenCalledOnce();
expect(consoleLogSpy).not.toHaveBeenCalled();
consoleLogSpy.mockRestore();
});
it('should do nothing if the server has no prompt support', async () => {
const mockRequest = vi.fn().mockResolvedValue({
prompts: [],
});
const mockGetServerCapabilities = vi.fn().mockReturnValue({});
const mockedClient = {
getServerCapabilities: mockGetServerCapabilities,
request: mockRequest,
} as unknown as ClientLib.Client;
const consoleLogSpy = vi
.spyOn(console, 'debug')
.mockImplementation(() => {});
await discoverPrompts('test-server', mockedClient, mockedPromptRegistry);
expect(mockGetServerCapabilities).toHaveBeenCalledOnce();
expect(mockRequest).not.toHaveBeenCalled();
expect(consoleLogSpy).not.toHaveBeenCalled();
consoleLogSpy.mockRestore();
});
it('should log an error if discovery fails', async () => {
const testError = new Error('test error');
testError.message = 'test error';
const mockRequest = vi.fn().mockRejectedValue(testError);
const mockGetServerCapabilities = vi.fn().mockReturnValue({
prompts: {},
});
const mockedClient = {
getServerCapabilities: mockGetServerCapabilities,
request: mockRequest,
} as unknown as ClientLib.Client;
const consoleErrorSpy = vi
.spyOn(console, 'error')
.mockImplementation(() => {});
await discoverPrompts('test-server', mockedClient, mockedPromptRegistry);
expect(mockRequest).toHaveBeenCalledOnce();
expect(consoleErrorSpy).toHaveBeenCalledWith(
`Error discovering prompts from test-server: ${testError.message}`,
);
consoleErrorSpy.mockRestore();
});
});
describe('appendMcpServerCommand', () => {
it('should do nothing if no MCP servers or command are configured', () => {
const out = populateMcpServerCommand({}, undefined);
expect(out).toEqual({});
});
it('should discover tools via mcpServerCommand', () => {
const commandString = 'command --arg1 value1';
const out = populateMcpServerCommand({}, commandString);
expect(out).toEqual({
mcp: {
command: 'command',
args: ['--arg1', 'value1'],
},
});
});
it('should handle error if mcpServerCommand parsing fails', () => {
expect(() => populateMcpServerCommand({}, 'derp && herp')).toThrowError();
});
});
describe('createTransport', () => {
const originalEnv = process.env;
beforeEach(() => {
vi.resetModules();
process.env = {};
});
afterEach(() => {
process.env = originalEnv;
});
describe('should connect via httpUrl', () => {
it('without headers', async () => {
const transport = await createTransport(
'test-server',
{
httpUrl: 'http://test-server',
},
false,
);
expect(transport).toEqual(
new StreamableHTTPClientTransport(new URL('http://test-server'), {}),
);
});
it('with headers', async () => {
const transport = await createTransport(
'test-server',
{
httpUrl: 'http://test-server',
headers: { Authorization: 'derp' },
},
false,
);
expect(transport).toEqual(
new StreamableHTTPClientTransport(new URL('http://test-server'), {
requestInit: {
headers: { Authorization: 'derp' },
},
}),
);
});
});
describe('should connect via url', () => {
it('without headers', async () => {
const transport = await createTransport(
'test-server',
{
url: 'http://test-server',
},
false,
);
expect(transport).toEqual(
new SSEClientTransport(new URL('http://test-server'), {}),
);
});
it('with headers', async () => {
const transport = await createTransport(
'test-server',
{
url: 'http://test-server',
headers: { Authorization: 'derp' },
},
false,
);
expect(transport).toEqual(
new SSEClientTransport(new URL('http://test-server'), {
requestInit: {
headers: { Authorization: 'derp' },
},
}),
);
});
});
it('should connect via command', async () => {
const mockedTransport = vi
.spyOn(SdkClientStdioLib, 'StdioClientTransport')
.mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);
await createTransport(
'test-server',
{
command: 'test-command',
args: ['--foo', 'bar'],
env: { FOO: 'bar' },
cwd: 'test/cwd',
},
false,
);
expect(mockedTransport).toHaveBeenCalledWith({
command: 'test-command',
args: ['--foo', 'bar'],
cwd: 'test/cwd',
env: { FOO: 'bar' },
stderr: 'pipe',
});
});
describe('useGoogleCredentialProvider', () => {
it('should use GoogleCredentialProvider when specified', async () => {
const transport = await createTransport(
'test-server',
{
httpUrl: 'http://test-server',
authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
oauth: {
scopes: ['scope1'],
},
},
false,
);
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const authProvider = (transport as any)._authProvider;
expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
});
it('should use GoogleCredentialProvider with SSE transport', async () => {
const transport = await createTransport(
'test-server',
{
url: 'http://test-server',
authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
oauth: {
scopes: ['scope1'],
},
},
false,
);
expect(transport).toBeInstanceOf(SSEClientTransport);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const authProvider = (transport as any)._authProvider;
expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
});
it('should throw an error if no URL is provided with GoogleCredentialProvider', async () => {
await expect(
createTransport(
'test-server',
{
authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
oauth: {
scopes: ['scope1'],
},
},
false,
),
).rejects.toThrow(
'No URL configured for Google Credentials MCP server',
);
});
});
});
describe('isEnabled', () => {
const funcDecl = { name: 'myTool' };
const serverName = 'myServer';
it('should return true if no include or exclude lists are provided', () => {
const mcpServerConfig = {};
expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(true);
});
it('should return false if the tool is in the exclude list', () => {
const mcpServerConfig = { excludeTools: ['myTool'] };
expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(false);
});
it('should return true if the tool is in the include list', () => {
const mcpServerConfig = { includeTools: ['myTool'] };
expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(true);
});
it('should return true if the tool is in the include list with parentheses', () => {
const mcpServerConfig = { includeTools: ['myTool()'] };
expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(true);
});
it('should return false if the include list exists but does not contain the tool', () => {
const mcpServerConfig = { includeTools: ['anotherTool'] };
expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(false);
});
it('should return false if the tool is in both the include and exclude lists', () => {
const mcpServerConfig = {
includeTools: ['myTool'],
excludeTools: ['myTool'],
};
expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(false);
});
it('should return false if the function declaration has no name', () => {
const namelessFuncDecl = {};
const mcpServerConfig = {};
expect(isEnabled(namelessFuncDecl, serverName, mcpServerConfig)).toBe(
false,
);
});
});
describe('hasValidTypes', () => {
it('should return true for a valid schema with anyOf', () => {
const schema = {
anyOf: [{ type: 'string' }, { type: 'number' }],
};
expect(hasValidTypes(schema)).toBe(true);
});
it('should return false for an invalid schema with anyOf', () => {
const schema = {
anyOf: [{ type: 'string' }, { description: 'no type' }],
};
expect(hasValidTypes(schema)).toBe(false);
});
it('should return true for a valid schema with allOf', () => {
const schema = {
allOf: [
{ type: 'string' },
{ type: 'object', properties: { foo: { type: 'string' } } },
],
};
expect(hasValidTypes(schema)).toBe(true);
});
it('should return false for an invalid schema with allOf', () => {
const schema = {
allOf: [{ type: 'string' }, { description: 'no type' }],
};
expect(hasValidTypes(schema)).toBe(false);
});
it('should return true for a valid schema with oneOf', () => {
const schema = {
oneOf: [{ type: 'string' }, { type: 'number' }],
};
expect(hasValidTypes(schema)).toBe(true);
});
it('should return false for an invalid schema with oneOf', () => {
const schema = {
oneOf: [{ type: 'string' }, { description: 'no type' }],
};
expect(hasValidTypes(schema)).toBe(false);
});
it('should return true for a valid schema with nested subschemas', () => {
const schema = {
anyOf: [
{ type: 'string' },
{
allOf: [
{ type: 'object', properties: { a: { type: 'string' } } },
{ type: 'object', properties: { b: { type: 'number' } } },
],
},
],
};
expect(hasValidTypes(schema)).toBe(true);
});
it('should return false for an invalid schema with nested subschemas', () => {
const schema = {
anyOf: [
{ type: 'string' },
{
allOf: [
{ type: 'object', properties: { a: { type: 'string' } } },
{ description: 'no type' },
],
},
],
};
expect(hasValidTypes(schema)).toBe(false);
});
it('should return true for a schema with a type and subschemas', () => {
const schema = {
type: 'string',
anyOf: [{ minLength: 1 }, { maxLength: 5 }],
};
expect(hasValidTypes(schema)).toBe(true);
});
it('should return false for a schema with no type and no subschemas', () => {
const schema = {
description: 'a schema with no type',
};
expect(hasValidTypes(schema)).toBe(false);
});
it('should return true for a valid schema', () => {
const schema = {
type: 'object',
properties: {
param1: { type: 'string' },
},
};
expect(hasValidTypes(schema)).toBe(true);
});
it('should return false if a parameter is missing a type', () => {
const schema = {
type: 'object',
properties: {
param1: { description: 'a param with no type' },
},
};
expect(hasValidTypes(schema)).toBe(false);
});
it('should return false if a nested parameter is missing a type', () => {
const schema = {
type: 'object',
properties: {
param1: {
type: 'object',
properties: {
nestedParam: {
description: 'a nested param with no type',
},
},
},
},
};
expect(hasValidTypes(schema)).toBe(false);
});
it('should return false if an array item is missing a type', () => {
const schema = {
type: 'object',
properties: {
param1: {
type: 'array',
items: {
description: 'an array item with no type',
},
},
},
};
expect(hasValidTypes(schema)).toBe(false);
});
it('should return true for a schema with no properties', () => {
const schema = {
type: 'object',
};
expect(hasValidTypes(schema)).toBe(true);
});
it('should return true for a schema with an empty properties object', () => {
const schema = {
type: 'object',
properties: {},
};
expect(hasValidTypes(schema)).toBe(true);
});
});
});