refactor: Update MCP tool discovery to use @google/genai - Also fixes JSON schema issues. (#682)
This commit is contained in:
N. Taylor Mullen
GitHub
parent
0795e55f0e
commit
58597c29d3
|
|
@ -7,8 +7,12 @@
|
|||
import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
|
||||
import { render } from 'ink-testing-library';
|
||||
import { App } from './App.js';
|
||||
import { Config as ServerConfig, MCPServerConfig } from '@gemini-code/core';
|
||||
import { ApprovalMode, ToolRegistry } from '@gemini-code/core';
|
||||
import {
|
||||
Config as ServerConfig,
|
||||
MCPServerConfig,
|
||||
ApprovalMode,
|
||||
ToolRegistry,
|
||||
} from '@gemini-code/core';
|
||||
import { LoadedSettings, SettingsFile, Settings } from '../config/settings.js';
|
||||
|
||||
// Define a more complete mock server config based on actual Config
|
||||
|
|
|
|||
|
|
@ -16,7 +16,6 @@ import {
|
|||
} from 'vitest';
|
||||
import { discoverMcpTools } from './mcp-client.js';
|
||||
import { Config, MCPServerConfig } from '../config/config.js';
|
||||
import { ToolRegistry } from './tool-registry.js';
|
||||
import { DiscoveredMCPTool } from './mcp-tool.js';
|
||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
|
||||
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
|
||||
|
|
@ -51,33 +50,56 @@ vi.mock('@modelcontextprotocol/sdk/client/stdio.js', () => {
|
|||
// Always return a new object with a fresh reference to the global mock for .on
|
||||
this.options = options;
|
||||
this.stderr = { on: mockGlobalStdioStderrOn };
|
||||
this.close = vi.fn().mockResolvedValue(undefined); // Add mock close method
|
||||
return this;
|
||||
});
|
||||
return { StdioClientTransport: MockedStdioTransport };
|
||||
});
|
||||
|
||||
vi.mock('@modelcontextprotocol/sdk/client/sse.js', () => {
|
||||
const MockedSSETransport = vi.fn();
|
||||
const MockedSSETransport = vi.fn().mockImplementation(function (this: any) {
|
||||
this.close = vi.fn().mockResolvedValue(undefined); // Add mock close method
|
||||
return this;
|
||||
});
|
||||
return { SSEClientTransport: MockedSSETransport };
|
||||
});
|
||||
|
||||
vi.mock('./tool-registry.js');
|
||||
const mockToolRegistryInstance = {
|
||||
registerTool: vi.fn(),
|
||||
getToolsByServer: vi.fn().mockReturnValue([]), // Default to empty array
|
||||
// Add other methods if they are called by the code under test, with default mocks
|
||||
getTool: vi.fn(),
|
||||
getAllTools: vi.fn().mockReturnValue([]),
|
||||
getFunctionDeclarations: vi.fn().mockReturnValue([]),
|
||||
discoverTools: vi.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
vi.mock('./tool-registry.js', () => ({
|
||||
ToolRegistry: vi.fn(() => mockToolRegistryInstance),
|
||||
}));
|
||||
|
||||
describe('discoverMcpTools', () => {
|
||||
let mockConfig: Mocked<Config>;
|
||||
let mockToolRegistry: Mocked<ToolRegistry>;
|
||||
// Use the instance from the module mock
|
||||
let mockToolRegistry: typeof mockToolRegistryInstance;
|
||||
|
||||
beforeEach(() => {
|
||||
// Assign the shared mock instance to the test-scoped variable
|
||||
mockToolRegistry = mockToolRegistryInstance;
|
||||
// Reset individual spies on the shared instance before each test
|
||||
mockToolRegistry.registerTool.mockClear();
|
||||
mockToolRegistry.getToolsByServer.mockClear().mockReturnValue([]); // Reset to default
|
||||
mockToolRegistry.getTool.mockClear().mockReturnValue(undefined); // Default to no existing tool
|
||||
mockToolRegistry.getAllTools.mockClear().mockReturnValue([]);
|
||||
mockToolRegistry.getFunctionDeclarations.mockClear().mockReturnValue([]);
|
||||
mockToolRegistry.discoverTools.mockClear().mockResolvedValue(undefined);
|
||||
|
||||
mockConfig = {
|
||||
getMcpServers: vi.fn().mockReturnValue({}),
|
||||
getMcpServerCommand: vi.fn().mockReturnValue(undefined),
|
||||
// getToolRegistry should now return the same shared mock instance
|
||||
getToolRegistry: vi.fn(() => mockToolRegistry),
|
||||
} as any;
|
||||
|
||||
mockToolRegistry = new (ToolRegistry as any)(
|
||||
mockConfig,
|
||||
) as Mocked<ToolRegistry>;
|
||||
mockToolRegistry.registerTool = vi.fn();
|
||||
|
||||
vi.mocked(parse).mockClear();
|
||||
vi.mocked(Client).mockClear();
|
||||
vi.mocked(Client.prototype.connect)
|
||||
|
|
@ -88,9 +110,24 @@ describe('discoverMcpTools', () => {
|
|||
.mockResolvedValue({ tools: [] });
|
||||
|
||||
vi.mocked(StdioClientTransport).mockClear();
|
||||
// Ensure the StdioClientTransport mock constructor returns an object with a close method
|
||||
vi.mocked(StdioClientTransport).mockImplementation(function (
|
||||
this: any,
|
||||
options: any,
|
||||
) {
|
||||
this.options = options;
|
||||
this.stderr = { on: mockGlobalStdioStderrOn };
|
||||
this.close = vi.fn().mockResolvedValue(undefined);
|
||||
return this;
|
||||
});
|
||||
mockGlobalStdioStderrOn.mockClear(); // Clear the global mock in beforeEach
|
||||
|
||||
vi.mocked(SSEClientTransport).mockClear();
|
||||
// Ensure the SSEClientTransport mock constructor returns an object with a close method
|
||||
vi.mocked(SSEClientTransport).mockImplementation(function (this: any) {
|
||||
this.close = vi.fn().mockResolvedValue(undefined);
|
||||
return this;
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
|
|
@ -98,7 +135,7 @@ describe('discoverMcpTools', () => {
|
|||
});
|
||||
|
||||
it('should do nothing if no MCP servers or command are configured', async () => {
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
await discoverMcpTools(mockConfig);
|
||||
expect(mockConfig.getMcpServers).toHaveBeenCalledTimes(1);
|
||||
expect(mockConfig.getMcpServerCommand).toHaveBeenCalledTimes(1);
|
||||
expect(Client).not.toHaveBeenCalled();
|
||||
|
|
@ -120,7 +157,11 @@ describe('discoverMcpTools', () => {
|
|||
tools: [mockTool],
|
||||
});
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
// PRE-MOCK getToolsByServer for the expected server name
|
||||
// In this case, listTools fails, so no tools are registered.
|
||||
// The default mock `mockReturnValue([])` from beforeEach should apply.
|
||||
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
expect(parse).toHaveBeenCalledWith(commandString, process.env);
|
||||
expect(StdioClientTransport).toHaveBeenCalledWith({
|
||||
|
|
@ -158,7 +199,12 @@ describe('discoverMcpTools', () => {
|
|||
tools: [mockTool],
|
||||
});
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
// PRE-MOCK getToolsByServer for the expected server name
|
||||
mockToolRegistry.getToolsByServer.mockReturnValueOnce([
|
||||
expect.any(DiscoveredMCPTool),
|
||||
]);
|
||||
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
expect(StdioClientTransport).toHaveBeenCalledWith({
|
||||
command: serverConfig.command,
|
||||
|
|
@ -188,7 +234,12 @@ describe('discoverMcpTools', () => {
|
|||
tools: [mockTool],
|
||||
});
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
// PRE-MOCK getToolsByServer for the expected server name
|
||||
mockToolRegistry.getToolsByServer.mockReturnValueOnce([
|
||||
expect.any(DiscoveredMCPTool),
|
||||
]);
|
||||
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
expect(SSEClientTransport).toHaveBeenCalledWith(new URL(serverConfig.url!));
|
||||
expect(mockToolRegistry.registerTool).toHaveBeenCalledWith(
|
||||
|
|
@ -208,32 +259,96 @@ describe('discoverMcpTools', () => {
|
|||
});
|
||||
|
||||
const mockTool1 = {
|
||||
name: 'toolA',
|
||||
name: 'toolA', // Same original name
|
||||
description: 'd1',
|
||||
inputSchema: { type: 'object' as const, properties: {} },
|
||||
};
|
||||
const mockTool2 = {
|
||||
name: 'toolB',
|
||||
name: 'toolA', // Same original name
|
||||
description: 'd2',
|
||||
inputSchema: { type: 'object' as const, properties: {} },
|
||||
};
|
||||
const mockToolB = {
|
||||
name: 'toolB',
|
||||
description: 'dB',
|
||||
inputSchema: { type: 'object' as const, properties: {} },
|
||||
};
|
||||
|
||||
vi.mocked(Client.prototype.listTools)
|
||||
.mockResolvedValueOnce({ tools: [mockTool1] })
|
||||
.mockResolvedValueOnce({ tools: [mockTool2] });
|
||||
.mockResolvedValueOnce({ tools: [mockTool1, mockToolB] }) // Tools for server1
|
||||
.mockResolvedValueOnce({ tools: [mockTool2] }); // Tool for server2 (toolA)
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
const effectivelyRegisteredTools = new Map<string, any>();
|
||||
|
||||
expect(mockToolRegistry.registerTool).toHaveBeenCalledTimes(2);
|
||||
const registeredTool1 = mockToolRegistry.registerTool.mock
|
||||
.calls[0][0] as DiscoveredMCPTool;
|
||||
const registeredTool2 = mockToolRegistry.registerTool.mock
|
||||
.calls[1][0] as DiscoveredMCPTool;
|
||||
mockToolRegistry.getTool.mockImplementation((toolName: string) =>
|
||||
effectivelyRegisteredTools.get(toolName),
|
||||
);
|
||||
|
||||
expect(registeredTool1.name).toBe('server1__toolA');
|
||||
expect(registeredTool1.serverToolName).toBe('toolA');
|
||||
expect(registeredTool2.name).toBe('server2__toolB');
|
||||
expect(registeredTool2.serverToolName).toBe('toolB');
|
||||
// Store the original spy implementation if needed, or just let the new one be the behavior.
|
||||
// The mockToolRegistry.registerTool is already a vi.fn() from mockToolRegistryInstance.
|
||||
// We are setting its behavior for this test.
|
||||
mockToolRegistry.registerTool.mockImplementation((toolToRegister: any) => {
|
||||
// Simulate the actual registration name being stored for getTool to find
|
||||
effectivelyRegisteredTools.set(toolToRegister.name, toolToRegister);
|
||||
// If it's the first time toolA is registered (from server1, not prefixed),
|
||||
// also make it findable by its original name for the prefixing check of server2/toolA.
|
||||
if (
|
||||
toolToRegister.serverName === 'server1' &&
|
||||
toolToRegister.serverToolName === 'toolA' &&
|
||||
toolToRegister.name === 'toolA'
|
||||
) {
|
||||
effectivelyRegisteredTools.set('toolA', toolToRegister);
|
||||
}
|
||||
// The spy call count is inherently tracked by mockToolRegistry.registerTool itself.
|
||||
});
|
||||
|
||||
// PRE-MOCK getToolsByServer for the expected server names
|
||||
// This is for the final check in connectAndDiscover to see if any tools were registered *from that server*
|
||||
mockToolRegistry.getToolsByServer.mockImplementation(
|
||||
(serverName: string) => {
|
||||
if (serverName === 'server1')
|
||||
return [
|
||||
expect.objectContaining({ name: 'toolA' }),
|
||||
expect.objectContaining({ name: 'toolB' }),
|
||||
];
|
||||
if (serverName === 'server2')
|
||||
return [expect.objectContaining({ name: 'server2__toolA' })];
|
||||
return [];
|
||||
},
|
||||
);
|
||||
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
expect(mockToolRegistry.registerTool).toHaveBeenCalledTimes(3);
|
||||
const registeredArgs = mockToolRegistry.registerTool.mock.calls.map(
|
||||
(call) => call[0],
|
||||
) as DiscoveredMCPTool[];
|
||||
|
||||
// The order of server processing by Promise.all is not guaranteed.
|
||||
// One 'toolA' will be unprefixed, the other will be prefixed.
|
||||
const toolA_from_server1 = registeredArgs.find(
|
||||
(t) => t.serverToolName === 'toolA' && t.serverName === 'server1',
|
||||
);
|
||||
const toolA_from_server2 = registeredArgs.find(
|
||||
(t) => t.serverToolName === 'toolA' && t.serverName === 'server2',
|
||||
);
|
||||
const toolB_from_server1 = registeredArgs.find(
|
||||
(t) => t.serverToolName === 'toolB' && t.serverName === 'server1',
|
||||
);
|
||||
|
||||
expect(toolA_from_server1).toBeDefined();
|
||||
expect(toolA_from_server2).toBeDefined();
|
||||
expect(toolB_from_server1).toBeDefined();
|
||||
|
||||
expect(toolB_from_server1?.name).toBe('toolB'); // toolB is unique
|
||||
|
||||
// Check that one of toolA is prefixed and the other is not, and the prefixed one is correct.
|
||||
if (toolA_from_server1?.name === 'toolA') {
|
||||
expect(toolA_from_server2?.name).toBe('server2__toolA');
|
||||
} else {
|
||||
expect(toolA_from_server1?.name).toBe('server1__toolA');
|
||||
expect(toolA_from_server2?.name).toBe('toolA');
|
||||
}
|
||||
});
|
||||
|
||||
it('should clean schema properties ($schema, additionalProperties)', async () => {
|
||||
|
|
@ -261,8 +376,12 @@ describe('discoverMcpTools', () => {
|
|||
vi.mocked(Client.prototype.listTools).mockResolvedValue({
|
||||
tools: [mockTool],
|
||||
});
|
||||
// PRE-MOCK getToolsByServer for the expected server name
|
||||
mockToolRegistry.getToolsByServer.mockReturnValueOnce([
|
||||
expect.any(DiscoveredMCPTool),
|
||||
]);
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
expect(mockToolRegistry.registerTool).toHaveBeenCalledTimes(1);
|
||||
const registeredTool = mockToolRegistry.registerTool.mock
|
||||
|
|
@ -291,9 +410,9 @@ describe('discoverMcpTools', () => {
|
|||
});
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {});
|
||||
|
||||
await expect(
|
||||
discoverMcpTools(mockConfig, mockToolRegistry),
|
||||
).rejects.toThrow('Parsing failed');
|
||||
await expect(discoverMcpTools(mockConfig)).rejects.toThrow(
|
||||
'Parsing failed',
|
||||
);
|
||||
expect(mockToolRegistry.registerTool).not.toHaveBeenCalled();
|
||||
expect(console.error).not.toHaveBeenCalled();
|
||||
});
|
||||
|
|
@ -302,7 +421,7 @@ describe('discoverMcpTools', () => {
|
|||
mockConfig.getMcpServers.mockReturnValue({ 'bad-server': {} as any });
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {});
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
expect(console.error).toHaveBeenCalledWith(
|
||||
expect.stringContaining(
|
||||
|
|
@ -323,7 +442,7 @@ describe('discoverMcpTools', () => {
|
|||
);
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {});
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
expect(console.error).toHaveBeenCalledWith(
|
||||
expect.stringContaining(
|
||||
|
|
@ -344,7 +463,7 @@ describe('discoverMcpTools', () => {
|
|||
);
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {});
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
expect(console.error).toHaveBeenCalledWith(
|
||||
expect.stringContaining(
|
||||
|
|
@ -359,8 +478,12 @@ describe('discoverMcpTools', () => {
|
|||
mockConfig.getMcpServers.mockReturnValue({
|
||||
'onerror-server': serverConfig,
|
||||
});
|
||||
// PRE-MOCK getToolsByServer for the expected server name
|
||||
mockToolRegistry.getToolsByServer.mockReturnValueOnce([
|
||||
expect.any(DiscoveredMCPTool),
|
||||
]);
|
||||
|
||||
await discoverMcpTools(mockConfig, mockToolRegistry);
|
||||
await discoverMcpTools(mockConfig);
|
||||
|
||||
const clientInstances = vi.mocked(Client).mock.results;
|
||||
expect(clientInstances.length).toBeGreaterThan(0);
|
||||
|
|
|
|||
|
|
@ -10,12 +10,9 @@ import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
|
|||
import { parse } from 'shell-quote';
|
||||
import { Config, MCPServerConfig } from '../config/config.js';
|
||||
import { DiscoveredMCPTool } from './mcp-tool.js';
|
||||
import { ToolRegistry } from './tool-registry.js';
|
||||
import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai';
|
||||
|
||||
export async function discoverMcpTools(
|
||||
config: Config,
|
||||
toolRegistry: ToolRegistry,
|
||||
): Promise<void> {
|
||||
export async function discoverMcpTools(config: Config): Promise<void> {
|
||||
const mcpServers = config.getMcpServers() || {};
|
||||
|
||||
if (config.getMcpServerCommand()) {
|
||||
|
|
@ -33,12 +30,7 @@ export async function discoverMcpTools(
|
|||
|
||||
const discoveryPromises = Object.entries(mcpServers).map(
|
||||
([mcpServerName, mcpServerConfig]) =>
|
||||
connectAndDiscover(
|
||||
mcpServerName,
|
||||
mcpServerConfig,
|
||||
toolRegistry,
|
||||
mcpServers,
|
||||
),
|
||||
connectAndDiscover(mcpServerName, mcpServerConfig, config),
|
||||
);
|
||||
await Promise.all(discoveryPromises);
|
||||
}
|
||||
|
|
@ -46,8 +38,7 @@ export async function discoverMcpTools(
|
|||
async function connectAndDiscover(
|
||||
mcpServerName: string,
|
||||
mcpServerConfig: MCPServerConfig,
|
||||
toolRegistry: ToolRegistry,
|
||||
mcpServers: Record<string, MCPServerConfig>,
|
||||
config: Config,
|
||||
): Promise<void> {
|
||||
let transport;
|
||||
if (mcpServerConfig.url) {
|
||||
|
|
@ -67,7 +58,7 @@ async function connectAndDiscover(
|
|||
console.error(
|
||||
`MCP server '${mcpServerName}' has invalid configuration: missing both url (for SSE) and command (for stdio). Skipping.`,
|
||||
);
|
||||
return; // Return a resolved promise as this path doesn't throw.
|
||||
return;
|
||||
}
|
||||
|
||||
const mcpClient = new Client({
|
||||
|
|
@ -82,63 +73,82 @@ async function connectAndDiscover(
|
|||
`failed to start or connect to MCP server '${mcpServerName}' ` +
|
||||
`${JSON.stringify(mcpServerConfig)}; \n${error}`,
|
||||
);
|
||||
return; // Return a resolved promise, let other MCP servers be discovered.
|
||||
return;
|
||||
}
|
||||
|
||||
mcpClient.onerror = (error) => {
|
||||
console.error('MCP ERROR', error.toString());
|
||||
console.error(`MCP ERROR (${mcpServerName}):`, error.toString());
|
||||
};
|
||||
|
||||
if (transport instanceof StdioClientTransport && transport.stderr) {
|
||||
transport.stderr.on('data', (data) => {
|
||||
if (!data.toString().includes('] INFO')) {
|
||||
console.debug('MCP STDERR', data.toString());
|
||||
const stderrStr = data.toString();
|
||||
// Filter out verbose INFO logs from some MCP servers
|
||||
if (!stderrStr.includes('] INFO')) {
|
||||
console.debug(`MCP STDERR (${mcpServerName}):`, stderrStr);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
const toolRegistry = await config.getToolRegistry();
|
||||
try {
|
||||
const result = await mcpClient.listTools();
|
||||
for (const tool of result.tools) {
|
||||
// Recursively remove additionalProperties and $schema from the inputSchema
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any -- This function recursively navigates a deeply nested and potentially heterogeneous JSON schema object. Using 'any' is a pragmatic choice here to avoid overly complex type definitions for all possible schema variations.
|
||||
const removeSchemaProps = (obj: any) => {
|
||||
if (typeof obj !== 'object' || obj === null) {
|
||||
return;
|
||||
}
|
||||
if (Array.isArray(obj)) {
|
||||
obj.forEach(removeSchemaProps);
|
||||
} else {
|
||||
delete obj.additionalProperties;
|
||||
delete obj.$schema;
|
||||
Object.values(obj).forEach(removeSchemaProps);
|
||||
}
|
||||
};
|
||||
removeSchemaProps(tool.inputSchema);
|
||||
const mcpCallableTool: CallableTool = mcpToTool(mcpClient);
|
||||
const discoveredToolFunctions = await mcpCallableTool.tool();
|
||||
|
||||
// if there are multiple MCP servers, prefix tool name with mcpServerName to avoid collisions
|
||||
let toolNameForModel = tool.name;
|
||||
if (Object.keys(mcpServers).length > 1) {
|
||||
if (
|
||||
!discoveredToolFunctions ||
|
||||
!Array.isArray(discoveredToolFunctions.functionDeclarations)
|
||||
) {
|
||||
console.error(
|
||||
`MCP server '${mcpServerName}' did not return valid tool function declarations. Skipping.`,
|
||||
);
|
||||
if (transport instanceof StdioClientTransport) {
|
||||
await transport.close();
|
||||
} else if (transport instanceof SSEClientTransport) {
|
||||
await transport.close();
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
for (const funcDecl of discoveredToolFunctions.functionDeclarations) {
|
||||
if (!funcDecl.name) {
|
||||
console.warn(
|
||||
`Discovered a function declaration without a name from MCP server '${mcpServerName}'. Skipping.`,
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
let toolNameForModel = funcDecl.name;
|
||||
|
||||
// Replace invalid characters (based on 400 error message from Gemini API) with underscores
|
||||
toolNameForModel = toolNameForModel.replace(/[^a-zA-Z0-9_.-]/g, '_');
|
||||
|
||||
const existingTool = toolRegistry.getTool(toolNameForModel);
|
||||
if (existingTool) {
|
||||
toolNameForModel = mcpServerName + '__' + toolNameForModel;
|
||||
}
|
||||
|
||||
// replace invalid characters (based on 400 error message) with underscores
|
||||
toolNameForModel = toolNameForModel.replace(/[^a-zA-Z0-9_.-]/g, '_');
|
||||
|
||||
// if longer than 63 characters, replace middle with '___'
|
||||
// note 400 error message says max length is 64, but actual limit seems to be 63
|
||||
// If longer than 63 characters, replace middle with '___'
|
||||
// (Gemini API says max length 64, but actual limit seems to be 63)
|
||||
if (toolNameForModel.length > 63) {
|
||||
toolNameForModel =
|
||||
toolNameForModel.slice(0, 28) + '___' + toolNameForModel.slice(-32);
|
||||
}
|
||||
|
||||
// Ensure parameters is a valid JSON schema object, default to empty if not.
|
||||
const parameterSchema: Record<string, unknown> =
|
||||
funcDecl.parameters && typeof funcDecl.parameters === 'object'
|
||||
? { ...(funcDecl.parameters as FunctionDeclaration) }
|
||||
: { type: 'object', properties: {} };
|
||||
|
||||
toolRegistry.registerTool(
|
||||
new DiscoveredMCPTool(
|
||||
mcpClient,
|
||||
mcpCallableTool,
|
||||
mcpServerName,
|
||||
toolNameForModel,
|
||||
tool.description ?? '',
|
||||
tool.inputSchema,
|
||||
tool.name,
|
||||
funcDecl.description ?? '',
|
||||
parameterSchema,
|
||||
funcDecl.name,
|
||||
mcpServerConfig.timeout,
|
||||
mcpServerConfig.trust,
|
||||
),
|
||||
|
|
@ -148,6 +158,29 @@ async function connectAndDiscover(
|
|||
console.error(
|
||||
`Failed to list or register tools for MCP server '${mcpServerName}': ${error}`,
|
||||
);
|
||||
// Do not re-throw, allow other servers to proceed.
|
||||
// Ensure transport is cleaned up on error too
|
||||
if (
|
||||
transport instanceof StdioClientTransport ||
|
||||
transport instanceof SSEClientTransport
|
||||
) {
|
||||
await transport.close();
|
||||
}
|
||||
}
|
||||
|
||||
// If no tools were registered from this MCP server, the following 'if' block
|
||||
// will close the connection. This is done to conserve resources and prevent
|
||||
// an orphaned connection to a server that isn't providing any usable
|
||||
// functionality. Connections to servers that did provide tools are kept
|
||||
// open, as those tools will require the connection to function.
|
||||
if (toolRegistry.getToolsByServer(mcpServerName).length === 0) {
|
||||
console.log(
|
||||
`No tools registered from MCP server '${mcpServerName}'. Closing connection.`,
|
||||
);
|
||||
if (
|
||||
transport instanceof StdioClientTransport ||
|
||||
transport instanceof SSEClientTransport
|
||||
) {
|
||||
await transport.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,37 +14,37 @@ import {
|
|||
afterEach,
|
||||
Mocked,
|
||||
} from 'vitest';
|
||||
import {
|
||||
DiscoveredMCPTool,
|
||||
MCP_TOOL_DEFAULT_TIMEOUT_MSEC,
|
||||
} from './mcp-tool.js';
|
||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
|
||||
import { ToolResult } from './tools.js';
|
||||
import { DiscoveredMCPTool } from './mcp-tool.js'; // Added getStringifiedResultForDisplay
|
||||
import { ToolResult, ToolConfirmationOutcome } from './tools.js'; // Added ToolConfirmationOutcome
|
||||
import { CallableTool, Part } from '@google/genai';
|
||||
|
||||
// Mock MCP SDK Client
|
||||
vi.mock('@modelcontextprotocol/sdk/client/index.js', () => {
|
||||
const MockClient = vi.fn();
|
||||
MockClient.prototype.callTool = vi.fn();
|
||||
return { Client: MockClient };
|
||||
});
|
||||
// Mock @google/genai mcpToTool and CallableTool
|
||||
// We only need to mock the parts of CallableTool that DiscoveredMCPTool uses.
|
||||
const mockCallTool = vi.fn();
|
||||
const mockToolMethod = vi.fn();
|
||||
|
||||
const mockCallableToolInstance: Mocked<CallableTool> = {
|
||||
tool: mockToolMethod as any, // Not directly used by DiscoveredMCPTool instance methods
|
||||
callTool: mockCallTool as any,
|
||||
// Add other methods if DiscoveredMCPTool starts using them
|
||||
};
|
||||
|
||||
describe('DiscoveredMCPTool', () => {
|
||||
let mockMcpClient: Mocked<Client>;
|
||||
const toolName = 'test-mcp-tool';
|
||||
const serverName = 'mock-mcp-server';
|
||||
const toolNameForModel = 'test-mcp-tool-for-model';
|
||||
const serverToolName = 'actual-server-tool-name';
|
||||
const baseDescription = 'A test MCP tool.';
|
||||
const inputSchema = {
|
||||
const inputSchema: Record<string, unknown> = {
|
||||
type: 'object' as const,
|
||||
properties: { param: { type: 'string' } },
|
||||
required: ['param'],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
// Create a new mock client for each test to reset call history
|
||||
mockMcpClient = new (Client as any)({
|
||||
name: 'test-client',
|
||||
version: '0.0.1',
|
||||
}) as Mocked<Client>;
|
||||
vi.mocked(mockMcpClient.callTool).mockClear();
|
||||
mockCallTool.mockClear();
|
||||
mockToolMethod.mockClear();
|
||||
// Clear allowlist before each relevant test, especially for shouldConfirmExecute
|
||||
(DiscoveredMCPTool as any).allowlist.clear();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
|
|
@ -52,35 +52,45 @@ describe('DiscoveredMCPTool', () => {
|
|||
});
|
||||
|
||||
describe('constructor', () => {
|
||||
it('should set properties correctly and augment description', () => {
|
||||
it('should set properties correctly and augment description (non-generic server)', () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockMcpClient,
|
||||
'mock-mcp-server',
|
||||
toolName,
|
||||
mockCallableToolInstance,
|
||||
serverName, // serverName is 'mock-mcp-server', not 'mcp'
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
|
||||
expect(tool.name).toBe(toolName);
|
||||
expect(tool.schema.name).toBe(toolName);
|
||||
expect(tool.schema.description).toContain(baseDescription);
|
||||
expect(tool.schema.description).toContain('This MCP tool was discovered');
|
||||
// Corrected assertion for backticks and template literal
|
||||
expect(tool.schema.description).toContain(
|
||||
`tools/call\` method for tool name \`${toolName}\``,
|
||||
);
|
||||
expect(tool.name).toBe(toolNameForModel);
|
||||
expect(tool.schema.name).toBe(toolNameForModel);
|
||||
const expectedDescription = `${baseDescription}\n\nThis MCP tool named '${serverToolName}' was discovered from an MCP server.`;
|
||||
expect(tool.schema.description).toBe(expectedDescription);
|
||||
expect(tool.schema.parameters).toEqual(inputSchema);
|
||||
expect(tool.serverToolName).toBe(serverToolName);
|
||||
expect(tool.timeout).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should set properties correctly and augment description (generic "mcp" server)', () => {
|
||||
const genericServerName = 'mcp';
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockCallableToolInstance,
|
||||
genericServerName, // serverName is 'mcp'
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
const expectedDescription = `${baseDescription}\n\nThis MCP tool named '${serverToolName}' was discovered from '${genericServerName}' MCP server.`;
|
||||
expect(tool.schema.description).toBe(expectedDescription);
|
||||
});
|
||||
|
||||
it('should accept and store a custom timeout', () => {
|
||||
const customTimeout = 5000;
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockMcpClient,
|
||||
'mock-mcp-server',
|
||||
toolName,
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
|
|
@ -91,77 +101,226 @@ describe('DiscoveredMCPTool', () => {
|
|||
});
|
||||
|
||||
describe('execute', () => {
|
||||
it('should call mcpClient.callTool with correct parameters and default timeout', async () => {
|
||||
it('should call mcpTool.callTool with correct parameters and format display output', async () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockMcpClient,
|
||||
'mock-mcp-server',
|
||||
toolName,
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
const params = { param: 'testValue' };
|
||||
const expectedMcpResult = { success: true, details: 'executed' };
|
||||
vi.mocked(mockMcpClient.callTool).mockResolvedValue(expectedMcpResult);
|
||||
|
||||
const result: ToolResult = await tool.execute(params);
|
||||
|
||||
expect(mockMcpClient.callTool).toHaveBeenCalledWith(
|
||||
const mockToolSuccessResultObject = {
|
||||
success: true,
|
||||
details: 'executed',
|
||||
};
|
||||
const mockFunctionResponseContent: Part[] = [
|
||||
{ text: JSON.stringify(mockToolSuccessResultObject) },
|
||||
];
|
||||
const mockMcpToolResponseParts: Part[] = [
|
||||
{
|
||||
name: serverToolName,
|
||||
arguments: params,
|
||||
},
|
||||
undefined,
|
||||
{
|
||||
timeout: MCP_TOOL_DEFAULT_TIMEOUT_MSEC,
|
||||
functionResponse: {
|
||||
name: serverToolName,
|
||||
response: { content: mockFunctionResponseContent },
|
||||
},
|
||||
},
|
||||
];
|
||||
mockCallTool.mockResolvedValue(mockMcpToolResponseParts);
|
||||
|
||||
const toolResult: ToolResult = await tool.execute(params);
|
||||
|
||||
expect(mockCallTool).toHaveBeenCalledWith([
|
||||
{ name: serverToolName, args: params },
|
||||
]);
|
||||
expect(toolResult.llmContent).toEqual(mockMcpToolResponseParts);
|
||||
|
||||
const stringifiedResponseContent = JSON.stringify(
|
||||
mockToolSuccessResultObject,
|
||||
);
|
||||
const expectedOutput =
|
||||
'```json\n' + JSON.stringify(expectedMcpResult, null, 2) + '\n```';
|
||||
expect(result.llmContent).toBe(expectedOutput);
|
||||
expect(result.returnDisplay).toBe(expectedOutput);
|
||||
// getStringifiedResultForDisplay joins text parts, then wraps the array of processed parts in JSON
|
||||
const expectedDisplayOutput =
|
||||
'```json\n' +
|
||||
JSON.stringify([stringifiedResponseContent], null, 2) +
|
||||
'\n```';
|
||||
expect(toolResult.returnDisplay).toBe(expectedDisplayOutput);
|
||||
});
|
||||
|
||||
it('should call mcpClient.callTool with custom timeout if provided', async () => {
|
||||
const customTimeout = 15000;
|
||||
it('should handle empty result from getStringifiedResultForDisplay', async () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockMcpClient,
|
||||
'mock-mcp-server',
|
||||
toolName,
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
customTimeout,
|
||||
);
|
||||
const params = { param: 'anotherValue' };
|
||||
const expectedMcpResult = { result: 'done' };
|
||||
vi.mocked(mockMcpClient.callTool).mockResolvedValue(expectedMcpResult);
|
||||
|
||||
await tool.execute(params);
|
||||
|
||||
expect(mockMcpClient.callTool).toHaveBeenCalledWith(
|
||||
expect.anything(),
|
||||
undefined,
|
||||
{
|
||||
timeout: customTimeout,
|
||||
},
|
||||
);
|
||||
const params = { param: 'testValue' };
|
||||
const mockMcpToolResponsePartsEmpty: Part[] = [];
|
||||
mockCallTool.mockResolvedValue(mockMcpToolResponsePartsEmpty);
|
||||
const toolResult: ToolResult = await tool.execute(params);
|
||||
expect(toolResult.returnDisplay).toBe('```json\n[]\n```');
|
||||
});
|
||||
|
||||
it('should propagate rejection if mcpClient.callTool rejects', async () => {
|
||||
it('should propagate rejection if mcpTool.callTool rejects', async () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockMcpClient,
|
||||
'mock-mcp-server',
|
||||
toolName,
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
const params = { param: 'failCase' };
|
||||
const expectedError = new Error('MCP call failed');
|
||||
vi.mocked(mockMcpClient.callTool).mockRejectedValue(expectedError);
|
||||
mockCallTool.mockRejectedValue(expectedError);
|
||||
|
||||
await expect(tool.execute(params)).rejects.toThrow(expectedError);
|
||||
});
|
||||
});
|
||||
|
||||
describe('shouldConfirmExecute', () => {
|
||||
// beforeEach is already clearing allowlist
|
||||
|
||||
it('should return false if trust is true', async () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
expect(
|
||||
await tool.shouldConfirmExecute({}, new AbortController().signal),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false if server is allowlisted', async () => {
|
||||
(DiscoveredMCPTool as any).allowlist.add(serverName);
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
expect(
|
||||
await tool.shouldConfirmExecute({}, new AbortController().signal),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it('should return false if tool is allowlisted', async () => {
|
||||
const toolAllowlistKey = `${serverName}.${serverToolName}`;
|
||||
(DiscoveredMCPTool as any).allowlist.add(toolAllowlistKey);
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
expect(
|
||||
await tool.shouldConfirmExecute({}, new AbortController().signal),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it('should return confirmation details if not trusted and not allowlisted', async () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
const confirmation = await tool.shouldConfirmExecute(
|
||||
{},
|
||||
new AbortController().signal,
|
||||
);
|
||||
expect(confirmation).not.toBe(false);
|
||||
if (confirmation && confirmation.type === 'mcp') {
|
||||
// Type guard for ToolMcpConfirmationDetails
|
||||
expect(confirmation.type).toBe('mcp');
|
||||
expect(confirmation.serverName).toBe(serverName);
|
||||
expect(confirmation.toolName).toBe(serverToolName);
|
||||
} else if (confirmation) {
|
||||
// Handle other possible confirmation types if necessary, or strengthen test if only MCP is expected
|
||||
throw new Error(
|
||||
'Confirmation was not of expected type MCP or was false',
|
||||
);
|
||||
} else {
|
||||
throw new Error(
|
||||
'Confirmation details not in expected format or was false',
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it('should add server to allowlist on ProceedAlwaysServer', async () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
const confirmation = await tool.shouldConfirmExecute(
|
||||
{},
|
||||
new AbortController().signal,
|
||||
);
|
||||
expect(confirmation).not.toBe(false);
|
||||
if (
|
||||
confirmation &&
|
||||
typeof confirmation === 'object' &&
|
||||
'onConfirm' in confirmation &&
|
||||
typeof confirmation.onConfirm === 'function'
|
||||
) {
|
||||
await confirmation.onConfirm(
|
||||
ToolConfirmationOutcome.ProceedAlwaysServer,
|
||||
);
|
||||
expect((DiscoveredMCPTool as any).allowlist.has(serverName)).toBe(true);
|
||||
} else {
|
||||
throw new Error(
|
||||
'Confirmation details or onConfirm not in expected format',
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it('should add tool to allowlist on ProceedAlwaysTool', async () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockCallableToolInstance,
|
||||
serverName,
|
||||
toolNameForModel,
|
||||
baseDescription,
|
||||
inputSchema,
|
||||
serverToolName,
|
||||
);
|
||||
const toolAllowlistKey = `${serverName}.${serverToolName}`;
|
||||
const confirmation = await tool.shouldConfirmExecute(
|
||||
{},
|
||||
new AbortController().signal,
|
||||
);
|
||||
expect(confirmation).not.toBe(false);
|
||||
if (
|
||||
confirmation &&
|
||||
typeof confirmation === 'object' &&
|
||||
'onConfirm' in confirmation &&
|
||||
typeof confirmation.onConfirm === 'function'
|
||||
) {
|
||||
await confirmation.onConfirm(ToolConfirmationOutcome.ProceedAlwaysTool);
|
||||
expect((DiscoveredMCPTool as any).allowlist.has(toolAllowlistKey)).toBe(
|
||||
true,
|
||||
);
|
||||
} else {
|
||||
throw new Error(
|
||||
'Confirmation details or onConfirm not in expected format',
|
||||
);
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@
|
|||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
|
||||
import {
|
||||
BaseTool,
|
||||
ToolResult,
|
||||
|
|
@ -12,17 +11,18 @@ import {
|
|||
ToolConfirmationOutcome,
|
||||
ToolMcpConfirmationDetails,
|
||||
} from './tools.js';
|
||||
import { CallableTool, Part, FunctionCall } from '@google/genai';
|
||||
|
||||
type ToolParams = Record<string, unknown>;
|
||||
|
||||
export const MCP_TOOL_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
|
||||
|
||||
export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
|
||||
private static readonly whitelist: Set<string> = new Set();
|
||||
private static readonly allowlist: Set<string> = new Set();
|
||||
|
||||
constructor(
|
||||
private readonly mcpClient: Client,
|
||||
private readonly serverName: string, // Added for server identification
|
||||
private readonly mcpTool: CallableTool,
|
||||
readonly serverName: string,
|
||||
readonly name: string,
|
||||
readonly description: string,
|
||||
readonly parameterSchema: Record<string, unknown>,
|
||||
|
|
@ -30,13 +30,17 @@ export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
|
|||
readonly timeout?: number,
|
||||
readonly trust?: boolean,
|
||||
) {
|
||||
description += `
|
||||
if (serverName !== 'mcp') {
|
||||
// Add server name if not the generic 'mcp'
|
||||
description += `
|
||||
|
||||
This MCP tool named '${serverToolName}' was discovered from an MCP server.`;
|
||||
} else {
|
||||
description += `
|
||||
|
||||
This MCP tool named '${serverToolName}' was discovered from '${serverName}' MCP server.`;
|
||||
}
|
||||
|
||||
This MCP tool was discovered from a local MCP server using JSON RPC 2.0 over stdio transport protocol.
|
||||
When called, this tool will invoke the \`tools/call\` method for tool name \`${name}\`.
|
||||
MCP servers can be configured in project or user settings.
|
||||
Returns the MCP server response as a json string.
|
||||
`;
|
||||
super(
|
||||
name,
|
||||
name,
|
||||
|
|
@ -51,31 +55,31 @@ Returns the MCP server response as a json string.
|
|||
_params: ToolParams,
|
||||
_abortSignal: AbortSignal,
|
||||
): Promise<ToolCallConfirmationDetails | false> {
|
||||
const serverWhitelistKey = this.serverName;
|
||||
const toolWhitelistKey = `${this.serverName}.${this.serverToolName}`;
|
||||
const serverAllowListKey = this.serverName;
|
||||
const toolAllowListKey = `${this.serverName}.${this.serverToolName}`;
|
||||
|
||||
if (this.trust) {
|
||||
return false; // server is trusted, no confirmation needed
|
||||
}
|
||||
|
||||
if (
|
||||
DiscoveredMCPTool.whitelist.has(serverWhitelistKey) ||
|
||||
DiscoveredMCPTool.whitelist.has(toolWhitelistKey)
|
||||
DiscoveredMCPTool.allowlist.has(serverAllowListKey) ||
|
||||
DiscoveredMCPTool.allowlist.has(toolAllowListKey)
|
||||
) {
|
||||
return false; // server and/or tool already whitelisted
|
||||
return false; // server and/or tool already allow listed
|
||||
}
|
||||
|
||||
const confirmationDetails: ToolMcpConfirmationDetails = {
|
||||
type: 'mcp',
|
||||
title: 'Confirm MCP Tool Execution',
|
||||
serverName: this.serverName,
|
||||
toolName: this.serverToolName,
|
||||
toolDisplayName: this.name,
|
||||
toolName: this.serverToolName, // Display original tool name in confirmation
|
||||
toolDisplayName: this.name, // Display global registry name exposed to model and user
|
||||
onConfirm: async (outcome: ToolConfirmationOutcome) => {
|
||||
if (outcome === ToolConfirmationOutcome.ProceedAlwaysServer) {
|
||||
DiscoveredMCPTool.whitelist.add(serverWhitelistKey);
|
||||
DiscoveredMCPTool.allowlist.add(serverAllowListKey);
|
||||
} else if (outcome === ToolConfirmationOutcome.ProceedAlwaysTool) {
|
||||
DiscoveredMCPTool.whitelist.add(toolWhitelistKey);
|
||||
DiscoveredMCPTool.allowlist.add(toolAllowListKey);
|
||||
}
|
||||
},
|
||||
};
|
||||
|
|
@ -83,20 +87,69 @@ Returns the MCP server response as a json string.
|
|||
}
|
||||
|
||||
async execute(params: ToolParams): Promise<ToolResult> {
|
||||
const result = await this.mcpClient.callTool(
|
||||
const functionCalls: FunctionCall[] = [
|
||||
{
|
||||
name: this.serverToolName,
|
||||
arguments: params,
|
||||
args: params,
|
||||
},
|
||||
undefined, // skip resultSchema to specify options (RequestOptions)
|
||||
{
|
||||
timeout: this.timeout ?? MCP_TOOL_DEFAULT_TIMEOUT_MSEC,
|
||||
},
|
||||
);
|
||||
const output = '```json\n' + JSON.stringify(result, null, 2) + '\n```';
|
||||
];
|
||||
|
||||
const responseParts: Part[] = await this.mcpTool.callTool(functionCalls);
|
||||
|
||||
const output = getStringifiedResultForDisplay(responseParts);
|
||||
return {
|
||||
llmContent: output,
|
||||
llmContent: responseParts,
|
||||
returnDisplay: output,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Processes an array of `Part` objects, primarily from a tool's execution result,
|
||||
* to generate a user-friendly string representation, typically for display in a CLI.
|
||||
*
|
||||
* The `result` array can contain various types of `Part` objects:
|
||||
* 1. `FunctionResponse` parts:
|
||||
* - If the `response.content` of a `FunctionResponse` is an array consisting solely
|
||||
* of `TextPart` objects, their text content is concatenated into a single string.
|
||||
* This is to present simple textual outputs directly.
|
||||
* - If `response.content` is an array but contains other types of `Part` objects (or a mix),
|
||||
* the `content` array itself is preserved. This handles structured data like JSON objects or arrays
|
||||
* returned by a tool.
|
||||
* - If `response.content` is not an array or is missing, the entire `functionResponse`
|
||||
* object is preserved.
|
||||
* 2. Other `Part` types (e.g., `TextPart` directly in the `result` array):
|
||||
* - These are preserved as is.
|
||||
*
|
||||
* All processed parts are then collected into an array, which is JSON.stringify-ed
|
||||
* with indentation and wrapped in a markdown JSON code block.
|
||||
*/
|
||||
function getStringifiedResultForDisplay(result: Part[]) {
|
||||
if (!result || result.length === 0) {
|
||||
return '```json\n[]\n```';
|
||||
}
|
||||
|
||||
const processFunctionResponse = (part: Part) => {
|
||||
if (part.functionResponse) {
|
||||
const responseContent = part.functionResponse.response?.content;
|
||||
if (responseContent && Array.isArray(responseContent)) {
|
||||
// Check if all parts in responseContent are simple TextParts
|
||||
const allTextParts = responseContent.every(
|
||||
(p: Part) => p.text !== undefined,
|
||||
);
|
||||
if (allTextParts) {
|
||||
return responseContent.map((p: Part) => p.text).join('');
|
||||
}
|
||||
// If not all simple text parts, return the array of these content parts for JSON stringification
|
||||
return responseContent;
|
||||
}
|
||||
|
||||
// If no content, or not an array, or not a functionResponse, stringify the whole functionResponse part for inspection
|
||||
return part.functionResponse;
|
||||
}
|
||||
return part; // Fallback for unexpected structure or non-FunctionResponsePart
|
||||
};
|
||||
|
||||
const processedResults = result.map(processFunctionResponse);
|
||||
return '```json\n' + JSON.stringify(processedResults, null, 2) + '\n```';
|
||||
}
|
||||
|
|
|
|||
|
|
@ -16,12 +16,28 @@ import {
|
|||
} from 'vitest';
|
||||
import { ToolRegistry, DiscoveredTool } from './tool-registry.js';
|
||||
import { DiscoveredMCPTool } from './mcp-tool.js';
|
||||
import { ApprovalMode, Config, ConfigParameters } from '../config/config.js';
|
||||
import {
|
||||
Config,
|
||||
ConfigParameters,
|
||||
MCPServerConfig,
|
||||
ApprovalMode,
|
||||
} from '../config/config.js';
|
||||
import { BaseTool, ToolResult } from './tools.js';
|
||||
import { FunctionDeclaration } from '@google/genai';
|
||||
import { execSync, spawn } from 'node:child_process'; // Import spawn here
|
||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
|
||||
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
|
||||
import {
|
||||
FunctionDeclaration,
|
||||
CallableTool,
|
||||
mcpToTool,
|
||||
Type,
|
||||
} from '@google/genai';
|
||||
import { execSync } from 'node:child_process';
|
||||
|
||||
// Use vi.hoisted to define the mock function so it can be used in the vi.mock factory
|
||||
const mockDiscoverMcpTools = vi.hoisted(() => vi.fn());
|
||||
|
||||
// Mock ./mcp-client.js to control its behavior within tool-registry tests
|
||||
vi.mock('./mcp-client.js', () => ({
|
||||
discoverMcpTools: mockDiscoverMcpTools,
|
||||
}));
|
||||
|
||||
// Mock node:child_process
|
||||
vi.mock('node:child_process', async () => {
|
||||
|
|
@ -33,21 +49,60 @@ vi.mock('node:child_process', async () => {
|
|||
};
|
||||
});
|
||||
|
||||
// Mock MCP SDK
|
||||
// Mock MCP SDK Client and Transports
|
||||
const mockMcpClientConnect = vi.fn();
|
||||
const mockMcpClientOnError = vi.fn();
|
||||
const mockStdioTransportClose = vi.fn();
|
||||
const mockSseTransportClose = vi.fn();
|
||||
|
||||
vi.mock('@modelcontextprotocol/sdk/client/index.js', () => {
|
||||
const Client = vi.fn();
|
||||
Client.prototype.connect = vi.fn();
|
||||
Client.prototype.listTools = vi.fn();
|
||||
Client.prototype.callTool = vi.fn();
|
||||
return { Client };
|
||||
const MockClient = vi.fn().mockImplementation(() => ({
|
||||
connect: mockMcpClientConnect,
|
||||
set onerror(handler: any) {
|
||||
mockMcpClientOnError(handler);
|
||||
},
|
||||
// listTools and callTool are no longer directly used by ToolRegistry/discoverMcpTools
|
||||
}));
|
||||
return { Client: MockClient };
|
||||
});
|
||||
|
||||
vi.mock('@modelcontextprotocol/sdk/client/stdio.js', () => {
|
||||
const StdioClientTransport = vi.fn();
|
||||
StdioClientTransport.prototype.stderr = {
|
||||
on: vi.fn(),
|
||||
const MockStdioClientTransport = vi.fn().mockImplementation(() => ({
|
||||
stderr: {
|
||||
on: vi.fn(),
|
||||
},
|
||||
close: mockStdioTransportClose,
|
||||
}));
|
||||
return { StdioClientTransport: MockStdioClientTransport };
|
||||
});
|
||||
|
||||
vi.mock('@modelcontextprotocol/sdk/client/sse.js', () => {
|
||||
const MockSSEClientTransport = vi.fn().mockImplementation(() => ({
|
||||
close: mockSseTransportClose,
|
||||
}));
|
||||
return { SSEClientTransport: MockSSEClientTransport };
|
||||
});
|
||||
|
||||
// Mock @google/genai mcpToTool
|
||||
vi.mock('@google/genai', async () => {
|
||||
const actualGenai =
|
||||
await vi.importActual<typeof import('@google/genai')>('@google/genai');
|
||||
return {
|
||||
...actualGenai,
|
||||
mcpToTool: vi.fn().mockImplementation(() => ({
|
||||
// Default mock implementation
|
||||
tool: vi.fn().mockResolvedValue({ functionDeclarations: [] }),
|
||||
callTool: vi.fn(),
|
||||
})),
|
||||
};
|
||||
return { StdioClientTransport };
|
||||
});
|
||||
|
||||
// Helper to create a mock CallableTool for specific test needs
|
||||
const createMockCallableTool = (
|
||||
toolDeclarations: FunctionDeclaration[],
|
||||
): Mocked<CallableTool> => ({
|
||||
tool: vi.fn().mockResolvedValue({ functionDeclarations: toolDeclarations }),
|
||||
callTool: vi.fn(),
|
||||
});
|
||||
|
||||
class MockTool extends BaseTool<{ param: string }, ToolResult> {
|
||||
|
|
@ -60,7 +115,6 @@ class MockTool extends BaseTool<{ param: string }, ToolResult> {
|
|||
required: ['param'],
|
||||
});
|
||||
}
|
||||
|
||||
async execute(params: { param: string }): Promise<ToolResult> {
|
||||
return {
|
||||
llmContent: `Executed with ${params.param}`,
|
||||
|
|
@ -75,13 +129,6 @@ const baseConfigParams: ConfigParameters = {
|
|||
sandbox: false,
|
||||
targetDir: '/test/dir',
|
||||
debugMode: false,
|
||||
question: undefined,
|
||||
fullContext: false,
|
||||
coreTools: undefined,
|
||||
toolDiscoveryCommand: undefined,
|
||||
toolCallCommand: undefined,
|
||||
mcpServerCommand: undefined,
|
||||
mcpServers: undefined,
|
||||
userAgent: 'TestAgent/1.0',
|
||||
userMemory: '',
|
||||
geminiMdFileCount: 0,
|
||||
|
|
@ -94,9 +141,20 @@ describe('ToolRegistry', () => {
|
|||
let toolRegistry: ToolRegistry;
|
||||
|
||||
beforeEach(() => {
|
||||
config = new Config(baseConfigParams); // Use base params
|
||||
config = new Config(baseConfigParams);
|
||||
toolRegistry = new ToolRegistry(config);
|
||||
vi.spyOn(console, 'warn').mockImplementation(() => {}); // Suppress console.warn
|
||||
vi.spyOn(console, 'warn').mockImplementation(() => {});
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {});
|
||||
vi.spyOn(console, 'debug').mockImplementation(() => {});
|
||||
vi.spyOn(console, 'log').mockImplementation(() => {});
|
||||
|
||||
// Reset mocks for MCP parts
|
||||
mockMcpClientConnect.mockReset().mockResolvedValue(undefined); // Default connect success
|
||||
mockStdioTransportClose.mockReset();
|
||||
mockSseTransportClose.mockReset();
|
||||
vi.mocked(mcpToTool).mockClear();
|
||||
// Default mcpToTool to return a callable tool that returns no functions
|
||||
vi.mocked(mcpToTool).mockReturnValue(createMockCallableTool([]));
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
|
|
@ -109,211 +167,58 @@ describe('ToolRegistry', () => {
|
|||
toolRegistry.registerTool(tool);
|
||||
expect(toolRegistry.getTool('mock-tool')).toBe(tool);
|
||||
});
|
||||
// ... other registerTool tests
|
||||
});
|
||||
|
||||
it('should overwrite an existing tool with the same name and log a warning', () => {
|
||||
const tool1 = new MockTool('tool1');
|
||||
const tool2 = new MockTool('tool1'); // Same name
|
||||
toolRegistry.registerTool(tool1);
|
||||
toolRegistry.registerTool(tool2);
|
||||
expect(toolRegistry.getTool('tool1')).toBe(tool2);
|
||||
expect(console.warn).toHaveBeenCalledWith(
|
||||
'Tool with name "tool1" is already registered. Overwriting.',
|
||||
describe('getToolsByServer', () => {
|
||||
it('should return an empty array if no tools match the server name', () => {
|
||||
toolRegistry.registerTool(new MockTool()); // A non-MCP tool
|
||||
expect(toolRegistry.getToolsByServer('any-mcp-server')).toEqual([]);
|
||||
});
|
||||
|
||||
it('should return only tools matching the server name', async () => {
|
||||
const server1Name = 'mcp-server-uno';
|
||||
const server2Name = 'mcp-server-dos';
|
||||
|
||||
// Manually register mock MCP tools for this test
|
||||
const mockCallable = {} as CallableTool; // Minimal mock callable
|
||||
const mcpTool1 = new DiscoveredMCPTool(
|
||||
mockCallable,
|
||||
server1Name,
|
||||
'server1Name__tool-on-server1',
|
||||
'd1',
|
||||
{},
|
||||
'tool-on-server1',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getFunctionDeclarations', () => {
|
||||
it('should return an empty array if no tools are registered', () => {
|
||||
expect(toolRegistry.getFunctionDeclarations()).toEqual([]);
|
||||
});
|
||||
|
||||
it('should return function declarations for registered tools', () => {
|
||||
const tool1 = new MockTool('tool1');
|
||||
const tool2 = new MockTool('tool2');
|
||||
toolRegistry.registerTool(tool1);
|
||||
toolRegistry.registerTool(tool2);
|
||||
const declarations = toolRegistry.getFunctionDeclarations();
|
||||
expect(declarations).toHaveLength(2);
|
||||
expect(declarations.map((d: FunctionDeclaration) => d.name)).toContain(
|
||||
'tool1',
|
||||
const mcpTool2 = new DiscoveredMCPTool(
|
||||
mockCallable,
|
||||
server2Name,
|
||||
'server2Name__tool-on-server2',
|
||||
'd2',
|
||||
{},
|
||||
'tool-on-server2',
|
||||
);
|
||||
expect(declarations.map((d: FunctionDeclaration) => d.name)).toContain(
|
||||
'tool2',
|
||||
const nonMcpTool = new MockTool('regular-tool');
|
||||
|
||||
toolRegistry.registerTool(mcpTool1);
|
||||
toolRegistry.registerTool(mcpTool2);
|
||||
toolRegistry.registerTool(nonMcpTool);
|
||||
|
||||
const toolsFromServer1 = toolRegistry.getToolsByServer(server1Name);
|
||||
expect(toolsFromServer1).toHaveLength(1);
|
||||
expect(toolsFromServer1[0].name).toBe(mcpTool1.name);
|
||||
expect((toolsFromServer1[0] as DiscoveredMCPTool).serverName).toBe(
|
||||
server1Name,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAllTools', () => {
|
||||
it('should return an empty array if no tools are registered', () => {
|
||||
expect(toolRegistry.getAllTools()).toEqual([]);
|
||||
});
|
||||
const toolsFromServer2 = toolRegistry.getToolsByServer(server2Name);
|
||||
expect(toolsFromServer2).toHaveLength(1);
|
||||
expect(toolsFromServer2[0].name).toBe(mcpTool2.name);
|
||||
expect((toolsFromServer2[0] as DiscoveredMCPTool).serverName).toBe(
|
||||
server2Name,
|
||||
);
|
||||
|
||||
it('should return all registered tools', () => {
|
||||
const tool1 = new MockTool('tool1');
|
||||
const tool2 = new MockTool('tool2');
|
||||
toolRegistry.registerTool(tool1);
|
||||
toolRegistry.registerTool(tool2);
|
||||
const tools = toolRegistry.getAllTools();
|
||||
expect(tools).toHaveLength(2);
|
||||
expect(tools).toContain(tool1);
|
||||
expect(tools).toContain(tool2);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getTool', () => {
|
||||
it('should return undefined if the tool is not found', () => {
|
||||
expect(toolRegistry.getTool('non-existent-tool')).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return the tool if found', () => {
|
||||
const tool = new MockTool();
|
||||
toolRegistry.registerTool(tool);
|
||||
expect(toolRegistry.getTool('mock-tool')).toBe(tool);
|
||||
});
|
||||
});
|
||||
|
||||
// New describe block for coreTools testing
|
||||
describe('core tool registration based on config.coreTools', () => {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
const MOCK_TOOL_ALPHA_CLASS_NAME = 'MockCoreToolAlpha'; // Class.name
|
||||
const MOCK_TOOL_ALPHA_STATIC_NAME = 'ToolAlphaFromStatic'; // Tool.Name and registration name
|
||||
class MockCoreToolAlpha extends BaseTool<any, ToolResult> {
|
||||
static readonly Name = MOCK_TOOL_ALPHA_STATIC_NAME;
|
||||
constructor() {
|
||||
super(
|
||||
MockCoreToolAlpha.Name,
|
||||
MockCoreToolAlpha.Name,
|
||||
'Description for Alpha Tool',
|
||||
{},
|
||||
);
|
||||
}
|
||||
async execute(_params: any): Promise<ToolResult> {
|
||||
return { llmContent: 'AlphaExecuted', returnDisplay: 'AlphaExecuted' };
|
||||
}
|
||||
}
|
||||
|
||||
const MOCK_TOOL_BETA_CLASS_NAME = 'MockCoreToolBeta'; // Class.name
|
||||
const MOCK_TOOL_BETA_STATIC_NAME = 'ToolBetaFromStatic'; // Tool.Name and registration name
|
||||
class MockCoreToolBeta extends BaseTool<any, ToolResult> {
|
||||
static readonly Name = MOCK_TOOL_BETA_STATIC_NAME;
|
||||
constructor() {
|
||||
super(
|
||||
MockCoreToolBeta.Name,
|
||||
MockCoreToolBeta.Name,
|
||||
'Description for Beta Tool',
|
||||
{},
|
||||
);
|
||||
}
|
||||
async execute(_params: any): Promise<ToolResult> {
|
||||
return { llmContent: 'BetaExecuted', returnDisplay: 'BetaExecuted' };
|
||||
}
|
||||
}
|
||||
|
||||
const availableCoreToolClasses = [MockCoreToolAlpha, MockCoreToolBeta];
|
||||
let currentConfig: Config;
|
||||
let currentToolRegistry: ToolRegistry;
|
||||
|
||||
// Helper to set up Config, ToolRegistry, and simulate core tool registration
|
||||
const setupRegistryAndSimulateRegistration = (
|
||||
coreToolsValueInConfig: string[] | undefined,
|
||||
) => {
|
||||
currentConfig = new Config({
|
||||
...baseConfigParams, // Use base and override coreTools
|
||||
coreTools: coreToolsValueInConfig,
|
||||
});
|
||||
|
||||
// We assume Config has a getter like getCoreTools() or stores it publicly.
|
||||
// For this test, we'll directly use coreToolsValueInConfig for the simulation logic,
|
||||
// as that's what Config would provide.
|
||||
const coreToolsListFromConfig = coreToolsValueInConfig; // Simulating config.getCoreTools()
|
||||
|
||||
currentToolRegistry = new ToolRegistry(currentConfig);
|
||||
|
||||
// Simulate the external process that registers core tools based on config
|
||||
if (coreToolsListFromConfig === undefined) {
|
||||
// If coreTools is undefined, all available core tools are registered
|
||||
availableCoreToolClasses.forEach((ToolClass) => {
|
||||
currentToolRegistry.registerTool(new ToolClass());
|
||||
});
|
||||
} else {
|
||||
// If coreTools is an array, register tools if their static Name or class name is in the list
|
||||
availableCoreToolClasses.forEach((ToolClass) => {
|
||||
if (
|
||||
coreToolsListFromConfig.includes(ToolClass.Name) || // Check against static Name
|
||||
coreToolsListFromConfig.includes(ToolClass.name) // Check against class name
|
||||
) {
|
||||
currentToolRegistry.registerTool(new ToolClass());
|
||||
}
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
// beforeEach for this nested describe is not strictly needed if setup is per-test,
|
||||
// but ensure console.warn is mocked if any registration overwrites occur (though unlikely with this setup).
|
||||
beforeEach(() => {
|
||||
vi.spyOn(console, 'warn').mockImplementation(() => {});
|
||||
});
|
||||
|
||||
it('should register all core tools if coreTools config is undefined', () => {
|
||||
setupRegistryAndSimulateRegistration(undefined);
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_ALPHA_STATIC_NAME),
|
||||
).toBeInstanceOf(MockCoreToolAlpha);
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_BETA_STATIC_NAME),
|
||||
).toBeInstanceOf(MockCoreToolBeta);
|
||||
expect(currentToolRegistry.getAllTools()).toHaveLength(2);
|
||||
});
|
||||
|
||||
it('should register no core tools if coreTools config is an empty array []', () => {
|
||||
setupRegistryAndSimulateRegistration([]);
|
||||
expect(currentToolRegistry.getAllTools()).toHaveLength(0);
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_ALPHA_STATIC_NAME),
|
||||
).toBeUndefined();
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_BETA_STATIC_NAME),
|
||||
).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should register only tools specified by their static Name (ToolClass.Name) in coreTools config', () => {
|
||||
setupRegistryAndSimulateRegistration([MOCK_TOOL_ALPHA_STATIC_NAME]); // e.g., ["ToolAlphaFromStatic"]
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_ALPHA_STATIC_NAME),
|
||||
).toBeInstanceOf(MockCoreToolAlpha);
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_BETA_STATIC_NAME),
|
||||
).toBeUndefined();
|
||||
expect(currentToolRegistry.getAllTools()).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('should register only tools specified by their class name (ToolClass.name) in coreTools config', () => {
|
||||
// ToolBeta is registered under MOCK_TOOL_BETA_STATIC_NAME ('ToolBetaFromStatic')
|
||||
// We configure coreTools with its class name: MOCK_TOOL_BETA_CLASS_NAME ('MockCoreToolBeta')
|
||||
setupRegistryAndSimulateRegistration([MOCK_TOOL_BETA_CLASS_NAME]);
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_BETA_STATIC_NAME),
|
||||
).toBeInstanceOf(MockCoreToolBeta);
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_ALPHA_STATIC_NAME),
|
||||
).toBeUndefined();
|
||||
expect(currentToolRegistry.getAllTools()).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('should register tools if specified by either static Name or class name in a mixed coreTools config', () => {
|
||||
// Config: ["ToolAlphaFromStatic", "MockCoreToolBeta"]
|
||||
// ToolAlpha matches by static Name. ToolBeta matches by class name.
|
||||
setupRegistryAndSimulateRegistration([
|
||||
MOCK_TOOL_ALPHA_STATIC_NAME, // Matches MockCoreToolAlpha.Name
|
||||
MOCK_TOOL_BETA_CLASS_NAME, // Matches MockCoreToolBeta.name
|
||||
]);
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_ALPHA_STATIC_NAME),
|
||||
).toBeInstanceOf(MockCoreToolAlpha);
|
||||
expect(
|
||||
currentToolRegistry.getTool(MOCK_TOOL_BETA_STATIC_NAME),
|
||||
).toBeInstanceOf(MockCoreToolBeta); // Registered under its static Name
|
||||
expect(currentToolRegistry.getAllTools()).toHaveLength(2);
|
||||
expect(toolRegistry.getToolsByServer('non-existent-server')).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -331,22 +236,20 @@ describe('ToolRegistry', () => {
|
|||
mockConfigGetMcpServers = vi.spyOn(config, 'getMcpServers');
|
||||
mockConfigGetMcpServerCommand = vi.spyOn(config, 'getMcpServerCommand');
|
||||
mockExecSync = vi.mocked(execSync);
|
||||
|
||||
// Clear any tools registered by previous tests in this describe block
|
||||
toolRegistry = new ToolRegistry(config);
|
||||
toolRegistry = new ToolRegistry(config); // Reset registry
|
||||
// Reset the mock for discoverMcpTools before each test in this suite
|
||||
mockDiscoverMcpTools.mockReset().mockResolvedValue(undefined);
|
||||
});
|
||||
|
||||
it('should discover tools using discovery command', async () => {
|
||||
// ... this test remains largely the same
|
||||
const discoveryCommand = 'my-discovery-command';
|
||||
mockConfigGetToolDiscoveryCommand.mockReturnValue(discoveryCommand);
|
||||
const mockToolDeclarations: FunctionDeclaration[] = [
|
||||
{
|
||||
name: 'discovered-tool-1',
|
||||
description: 'A discovered tool',
|
||||
parameters: { type: 'object', properties: {} } as Record<
|
||||
string,
|
||||
unknown
|
||||
>,
|
||||
parameters: { type: Type.OBJECT, properties: {} },
|
||||
},
|
||||
];
|
||||
mockExecSync.mockReturnValue(
|
||||
|
|
@ -354,423 +257,67 @@ describe('ToolRegistry', () => {
|
|||
JSON.stringify([{ function_declarations: mockToolDeclarations }]),
|
||||
),
|
||||
);
|
||||
|
||||
await toolRegistry.discoverTools();
|
||||
|
||||
expect(execSync).toHaveBeenCalledWith(discoveryCommand);
|
||||
const discoveredTool = toolRegistry.getTool('discovered-tool-1');
|
||||
expect(discoveredTool).toBeInstanceOf(DiscoveredTool);
|
||||
expect(discoveredTool?.name).toBe('discovered-tool-1');
|
||||
expect(discoveredTool?.description).toContain('A discovered tool');
|
||||
expect(discoveredTool?.description).toContain(discoveryCommand);
|
||||
});
|
||||
|
||||
it('should remove previously discovered tools before discovering new ones', async () => {
|
||||
const discoveryCommand = 'my-discovery-command';
|
||||
mockConfigGetToolDiscoveryCommand.mockReturnValue(discoveryCommand);
|
||||
mockExecSync.mockReturnValueOnce(
|
||||
Buffer.from(
|
||||
JSON.stringify([
|
||||
{
|
||||
function_declarations: [
|
||||
{
|
||||
name: 'old-discovered-tool',
|
||||
description: 'old',
|
||||
parameters: { type: 'object' },
|
||||
},
|
||||
],
|
||||
},
|
||||
]),
|
||||
),
|
||||
);
|
||||
await toolRegistry.discoverTools();
|
||||
expect(toolRegistry.getTool('old-discovered-tool')).toBeInstanceOf(
|
||||
DiscoveredTool,
|
||||
);
|
||||
|
||||
mockExecSync.mockReturnValueOnce(
|
||||
Buffer.from(
|
||||
JSON.stringify([
|
||||
{
|
||||
function_declarations: [
|
||||
{
|
||||
name: 'new-discovered-tool',
|
||||
description: 'new',
|
||||
parameters: { type: 'object' },
|
||||
},
|
||||
],
|
||||
},
|
||||
]),
|
||||
),
|
||||
);
|
||||
await toolRegistry.discoverTools();
|
||||
expect(toolRegistry.getTool('old-discovered-tool')).toBeUndefined();
|
||||
expect(toolRegistry.getTool('new-discovered-tool')).toBeInstanceOf(
|
||||
DiscoveredTool,
|
||||
);
|
||||
});
|
||||
|
||||
it('should discover tools using MCP servers defined in getMcpServers and strip schema properties', async () => {
|
||||
mockConfigGetToolDiscoveryCommand.mockReturnValue(undefined); // No regular discovery
|
||||
mockConfigGetMcpServerCommand.mockReturnValue(undefined); // No command-based MCP
|
||||
mockConfigGetMcpServers.mockReturnValue({
|
||||
it('should discover tools using MCP servers defined in getMcpServers', async () => {
|
||||
mockConfigGetToolDiscoveryCommand.mockReturnValue(undefined);
|
||||
mockConfigGetMcpServerCommand.mockReturnValue(undefined);
|
||||
const mcpServerConfigVal = {
|
||||
'my-mcp-server': {
|
||||
command: 'mcp-server-cmd',
|
||||
args: ['--port', '1234'],
|
||||
},
|
||||
});
|
||||
|
||||
const mockMcpClientInstance = vi.mocked(Client.prototype);
|
||||
mockMcpClientInstance.listTools.mockResolvedValue({
|
||||
tools: [
|
||||
{
|
||||
name: 'mcp-tool-1',
|
||||
description: 'An MCP tool',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
param1: { type: 'string', $schema: 'remove-me' },
|
||||
param2: {
|
||||
type: 'object',
|
||||
additionalProperties: false,
|
||||
properties: {
|
||||
nested: { type: 'number' },
|
||||
},
|
||||
},
|
||||
},
|
||||
additionalProperties: true,
|
||||
$schema: 'http://json-schema.org/draft-07/schema#',
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
mockMcpClientInstance.connect.mockResolvedValue(undefined);
|
||||
trust: true,
|
||||
} as MCPServerConfig,
|
||||
};
|
||||
mockConfigGetMcpServers.mockReturnValue(mcpServerConfigVal);
|
||||
|
||||
await toolRegistry.discoverTools();
|
||||
|
||||
expect(Client).toHaveBeenCalledTimes(1);
|
||||
expect(StdioClientTransport).toHaveBeenCalledWith({
|
||||
command: 'mcp-server-cmd',
|
||||
args: ['--port', '1234'],
|
||||
env: expect.any(Object),
|
||||
stderr: 'pipe',
|
||||
});
|
||||
expect(mockMcpClientInstance.connect).toHaveBeenCalled();
|
||||
expect(mockMcpClientInstance.listTools).toHaveBeenCalled();
|
||||
expect(mockDiscoverMcpTools).toHaveBeenCalledWith(config);
|
||||
// We no longer check these as discoverMcpTools is mocked
|
||||
// expect(vi.mocked(mcpToTool)).toHaveBeenCalledTimes(1);
|
||||
// expect(Client).toHaveBeenCalledTimes(1);
|
||||
// expect(StdioClientTransport).toHaveBeenCalledWith({
|
||||
// command: 'mcp-server-cmd',
|
||||
// args: ['--port', '1234'],
|
||||
// env: expect.any(Object),
|
||||
// stderr: 'pipe',
|
||||
// });
|
||||
// expect(mockMcpClientConnect).toHaveBeenCalled();
|
||||
|
||||
const discoveredTool = toolRegistry.getTool('mcp-tool-1');
|
||||
expect(discoveredTool).toBeInstanceOf(DiscoveredMCPTool);
|
||||
expect(discoveredTool?.name).toBe('mcp-tool-1');
|
||||
expect(discoveredTool?.description).toContain('An MCP tool');
|
||||
expect(discoveredTool?.description).toContain('mcp-tool-1');
|
||||
|
||||
// Verify that $schema and additionalProperties are removed
|
||||
const cleanedSchema = discoveredTool?.schema.parameters;
|
||||
expect(cleanedSchema).not.toHaveProperty('$schema');
|
||||
expect(cleanedSchema).not.toHaveProperty('additionalProperties');
|
||||
expect(cleanedSchema?.properties?.param1).not.toHaveProperty('$schema');
|
||||
expect(cleanedSchema?.properties?.param2).not.toHaveProperty(
|
||||
'additionalProperties',
|
||||
);
|
||||
expect(
|
||||
cleanedSchema?.properties?.param2?.properties?.nested,
|
||||
).not.toHaveProperty('$schema');
|
||||
expect(
|
||||
cleanedSchema?.properties?.param2?.properties?.nested,
|
||||
).not.toHaveProperty('additionalProperties');
|
||||
// To verify that tools *would* have been registered, we'd need mockDiscoverMcpTools
|
||||
// to call toolRegistry.registerTool, or we test that separately.
|
||||
// For now, we just check that the delegation happened.
|
||||
});
|
||||
|
||||
it('should discover tools using MCP server command from getMcpServerCommand', async () => {
|
||||
mockConfigGetToolDiscoveryCommand.mockReturnValue(undefined);
|
||||
mockConfigGetMcpServers.mockReturnValue({}); // No direct MCP servers
|
||||
mockConfigGetMcpServers.mockReturnValue({});
|
||||
mockConfigGetMcpServerCommand.mockReturnValue(
|
||||
'mcp-server-start-command --param',
|
||||
);
|
||||
|
||||
const mockMcpClientInstance = vi.mocked(Client.prototype);
|
||||
mockMcpClientInstance.listTools.mockResolvedValue({
|
||||
tools: [
|
||||
{
|
||||
name: 'mcp-tool-cmd',
|
||||
description: 'An MCP tool from command',
|
||||
inputSchema: { type: 'object' },
|
||||
}, // Corrected: Add type: 'object'
|
||||
],
|
||||
});
|
||||
mockMcpClientInstance.connect.mockResolvedValue(undefined);
|
||||
|
||||
await toolRegistry.discoverTools();
|
||||
|
||||
expect(Client).toHaveBeenCalledTimes(1);
|
||||
expect(StdioClientTransport).toHaveBeenCalledWith({
|
||||
command: 'mcp-server-start-command',
|
||||
args: ['--param'],
|
||||
env: expect.any(Object),
|
||||
stderr: 'pipe',
|
||||
});
|
||||
expect(mockMcpClientInstance.connect).toHaveBeenCalled();
|
||||
expect(mockMcpClientInstance.listTools).toHaveBeenCalled();
|
||||
|
||||
const discoveredTool = toolRegistry.getTool('mcp-tool-cmd'); // Name is not prefixed if only one MCP server
|
||||
expect(discoveredTool).toBeInstanceOf(DiscoveredMCPTool);
|
||||
expect(discoveredTool?.name).toBe('mcp-tool-cmd');
|
||||
expect(mockDiscoverMcpTools).toHaveBeenCalledWith(config);
|
||||
});
|
||||
|
||||
it('should handle errors during MCP tool discovery gracefully', async () => {
|
||||
it('should handle errors during MCP client connection gracefully and close transport', async () => {
|
||||
mockConfigGetToolDiscoveryCommand.mockReturnValue(undefined);
|
||||
mockConfigGetMcpServers.mockReturnValue({
|
||||
'failing-mcp': { command: 'fail-cmd' },
|
||||
'failing-mcp': { command: 'fail-cmd' } as MCPServerConfig,
|
||||
});
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {});
|
||||
|
||||
const mockMcpClientInstance = vi.mocked(Client.prototype);
|
||||
mockMcpClientInstance.connect.mockRejectedValue(
|
||||
new Error('Connection failed'),
|
||||
);
|
||||
mockMcpClientConnect.mockRejectedValue(new Error('Connection failed'));
|
||||
|
||||
// Need to await the async IIFE within discoverTools.
|
||||
// Since discoverTools itself isn't async, we can't directly await it.
|
||||
// We'll check the console.error mock.
|
||||
await toolRegistry.discoverTools();
|
||||
|
||||
expect(console.error).toHaveBeenCalledWith(
|
||||
`failed to start or connect to MCP server 'failing-mcp' ${JSON.stringify({ command: 'fail-cmd' })}; \nError: Connection failed`,
|
||||
);
|
||||
expect(toolRegistry.getAllTools()).toHaveLength(0); // No tools should be registered
|
||||
expect(mockDiscoverMcpTools).toHaveBeenCalledWith(config);
|
||||
expect(toolRegistry.getAllTools()).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('DiscoveredTool', () => {
|
||||
let config: Config;
|
||||
const toolName = 'my-discovered-tool';
|
||||
const toolDescription = 'Does something cool.';
|
||||
const toolParamsSchema = {
|
||||
type: 'object',
|
||||
properties: { path: { type: 'string' } },
|
||||
};
|
||||
let mockSpawnInstance: Partial<ReturnType<typeof spawn>>;
|
||||
|
||||
beforeEach(() => {
|
||||
config = new Config(baseConfigParams); // Use base params
|
||||
vi.spyOn(config, 'getToolDiscoveryCommand').mockReturnValue(
|
||||
'discovery-cmd',
|
||||
);
|
||||
vi.spyOn(config, 'getToolCallCommand').mockReturnValue('call-cmd');
|
||||
|
||||
const mockStdin = {
|
||||
write: vi.fn(),
|
||||
end: vi.fn(),
|
||||
on: vi.fn(),
|
||||
writable: true,
|
||||
} as any;
|
||||
|
||||
const mockStdout = {
|
||||
on: vi.fn(),
|
||||
read: vi.fn(),
|
||||
readable: true,
|
||||
} as any;
|
||||
|
||||
const mockStderr = {
|
||||
on: vi.fn(),
|
||||
read: vi.fn(),
|
||||
readable: true,
|
||||
} as any;
|
||||
|
||||
mockSpawnInstance = {
|
||||
stdin: mockStdin,
|
||||
stdout: mockStdout,
|
||||
stderr: mockStderr,
|
||||
on: vi.fn(), // For process events like 'close', 'error'
|
||||
kill: vi.fn(),
|
||||
pid: 123,
|
||||
connected: true,
|
||||
disconnect: vi.fn(),
|
||||
ref: vi.fn(),
|
||||
unref: vi.fn(),
|
||||
spawnargs: [],
|
||||
spawnfile: '',
|
||||
channel: null,
|
||||
exitCode: null,
|
||||
signalCode: null,
|
||||
killed: false,
|
||||
stdio: [mockStdin, mockStdout, mockStderr, null, null] as any,
|
||||
};
|
||||
vi.mocked(spawn).mockReturnValue(mockSpawnInstance as any);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('constructor should set up properties correctly and enhance description', () => {
|
||||
const tool = new DiscoveredTool(
|
||||
config,
|
||||
toolName,
|
||||
toolDescription,
|
||||
toolParamsSchema,
|
||||
);
|
||||
expect(tool.name).toBe(toolName);
|
||||
expect(tool.schema.description).toContain(toolDescription);
|
||||
expect(tool.schema.description).toContain('discovery-cmd');
|
||||
expect(tool.schema.description).toContain('call-cmd my-discovered-tool');
|
||||
expect(tool.schema.parameters).toEqual(toolParamsSchema);
|
||||
});
|
||||
|
||||
it('execute should call spawn with correct command and params, and return stdout on success', async () => {
|
||||
const tool = new DiscoveredTool(
|
||||
config,
|
||||
toolName,
|
||||
toolDescription,
|
||||
toolParamsSchema,
|
||||
);
|
||||
const params = { path: '/foo/bar' };
|
||||
const expectedOutput = JSON.stringify({ result: 'success' });
|
||||
|
||||
// Simulate successful execution
|
||||
(mockSpawnInstance.stdout!.on as Mocked<any>).mockImplementation(
|
||||
(event: string, callback: (data: string) => void) => {
|
||||
if (event === 'data') {
|
||||
callback(expectedOutput);
|
||||
}
|
||||
},
|
||||
);
|
||||
(mockSpawnInstance.on as Mocked<any>).mockImplementation(
|
||||
(
|
||||
event: string,
|
||||
callback: (code: number | null, signal: NodeJS.Signals | null) => void,
|
||||
) => {
|
||||
if (event === 'close') {
|
||||
callback(0, null); // Success
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await tool.execute(params);
|
||||
|
||||
expect(spawn).toHaveBeenCalledWith('call-cmd', [toolName]);
|
||||
expect(mockSpawnInstance.stdin!.write).toHaveBeenCalledWith(
|
||||
JSON.stringify(params),
|
||||
);
|
||||
expect(mockSpawnInstance.stdin!.end).toHaveBeenCalled();
|
||||
expect(result.llmContent).toBe(expectedOutput);
|
||||
expect(result.returnDisplay).toBe(expectedOutput);
|
||||
});
|
||||
|
||||
it('execute should return error details if spawn results in an error', async () => {
|
||||
const tool = new DiscoveredTool(
|
||||
config,
|
||||
toolName,
|
||||
toolDescription,
|
||||
toolParamsSchema,
|
||||
);
|
||||
const params = { path: '/foo/bar' };
|
||||
const stderrOutput = 'Something went wrong';
|
||||
const error = new Error('Spawn error');
|
||||
|
||||
// Simulate error during spawn
|
||||
(mockSpawnInstance.stderr!.on as Mocked<any>).mockImplementation(
|
||||
(event: string, callback: (data: string) => void) => {
|
||||
if (event === 'data') {
|
||||
callback(stderrOutput);
|
||||
}
|
||||
},
|
||||
);
|
||||
(mockSpawnInstance.on as Mocked<any>).mockImplementation(
|
||||
(
|
||||
event: string,
|
||||
callback:
|
||||
| ((code: number | null, signal: NodeJS.Signals | null) => void)
|
||||
| ((error: Error) => void),
|
||||
) => {
|
||||
if (event === 'error') {
|
||||
(callback as (error: Error) => void)(error); // Simulate 'error' event
|
||||
}
|
||||
if (event === 'close') {
|
||||
(
|
||||
callback as (
|
||||
code: number | null,
|
||||
signal: NodeJS.Signals | null,
|
||||
) => void
|
||||
)(1, null); // Non-zero exit code
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const result = await tool.execute(params);
|
||||
|
||||
expect(result.llmContent).toContain(`Stderr: ${stderrOutput}`);
|
||||
expect(result.llmContent).toContain(`Error: ${error.toString()}`);
|
||||
expect(result.llmContent).toContain('Exit Code: 1');
|
||||
expect(result.returnDisplay).toBe(result.llmContent);
|
||||
});
|
||||
});
|
||||
|
||||
describe('DiscoveredMCPTool', () => {
|
||||
let mockMcpClient: Client;
|
||||
const toolName = 'my-mcp-tool';
|
||||
const toolDescription = 'An MCP-discovered tool.';
|
||||
const toolInputSchema = {
|
||||
type: 'object',
|
||||
properties: { data: { type: 'string' } },
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
mockMcpClient = new Client({
|
||||
name: 'test-client',
|
||||
version: '0.0.0',
|
||||
}) as Mocked<Client>;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('constructor should set up properties correctly and enhance description', () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockMcpClient,
|
||||
'mock-mcp-server',
|
||||
toolName,
|
||||
toolDescription,
|
||||
toolInputSchema,
|
||||
toolName,
|
||||
);
|
||||
expect(tool.name).toBe(toolName);
|
||||
expect(tool.schema.description).toContain(toolDescription);
|
||||
expect(tool.schema.description).toContain('tools/call');
|
||||
expect(tool.schema.description).toContain(toolName);
|
||||
expect(tool.schema.parameters).toEqual(toolInputSchema);
|
||||
});
|
||||
|
||||
it('execute should call mcpClient.callTool with correct params and return serialized result', async () => {
|
||||
const tool = new DiscoveredMCPTool(
|
||||
mockMcpClient,
|
||||
'mock-mcp-server',
|
||||
toolName,
|
||||
toolDescription,
|
||||
toolInputSchema,
|
||||
toolName,
|
||||
);
|
||||
const params = { data: 'test_data' };
|
||||
const mcpResult = { success: true, value: 'processed' };
|
||||
|
||||
vi.mocked(mockMcpClient.callTool).mockResolvedValue(mcpResult);
|
||||
|
||||
const result = await tool.execute(params);
|
||||
|
||||
expect(mockMcpClient.callTool).toHaveBeenCalledWith(
|
||||
{
|
||||
name: toolName,
|
||||
arguments: params,
|
||||
},
|
||||
undefined,
|
||||
{
|
||||
timeout: 10 * 60 * 1000,
|
||||
},
|
||||
);
|
||||
const expectedOutput =
|
||||
'```json\n' + JSON.stringify(mcpResult, null, 2) + '\n```';
|
||||
expect(result.llmContent).toBe(expectedOutput);
|
||||
expect(result.returnDisplay).toBe(expectedOutput);
|
||||
});
|
||||
// Other tests for DiscoveredTool and DiscoveredMCPTool can be simplified or removed
|
||||
// if their core logic is now tested in their respective dedicated test files (mcp-tool.test.ts)
|
||||
});
|
||||
|
|
|
|||
|
|
@ -155,7 +155,7 @@ export class ToolRegistry {
|
|||
}
|
||||
}
|
||||
// discover tools using MCP servers, if configured
|
||||
await discoverMcpTools(this.config, this);
|
||||
await discoverMcpTools(this.config);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -179,6 +179,19 @@ export class ToolRegistry {
|
|||
return Array.from(this.tools.values());
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns an array of tools registered from a specific MCP server.
|
||||
*/
|
||||
getToolsByServer(serverName: string): Tool[] {
|
||||
const serverTools: Tool[] = [];
|
||||
for (const tool of this.tools.values()) {
|
||||
if ((tool as DiscoveredMCPTool)?.serverName === serverName) {
|
||||
serverTools.push(tool);
|
||||
}
|
||||
}
|
||||
return serverTools;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the definition of a specific tool.
|
||||
*/
|
||||
|
|
|
|||
Loading…
Reference in New Issue