diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index 353b4f05..09614442 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -9,7 +9,6 @@ import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/ import { populateMcpServerCommand, createTransport, - generateValidName, isEnabled, discoverTools, } from './mcp-client.js'; @@ -173,92 +172,6 @@ describe('mcp-client', () => { }); }); }); - describe('generateValidName', () => { - it('should return a valid name for a simple function', () => { - const funcDecl = { name: 'myFunction' }; - const serverName = 'myServer'; - const result = generateValidName(funcDecl, serverName); - expect(result).toBe('myServer__myFunction'); - }); - - it('should prepend the server name', () => { - const funcDecl = { name: 'anotherFunction' }; - const serverName = 'production-server'; - const result = generateValidName(funcDecl, serverName); - expect(result).toBe('production-server__anotherFunction'); - }); - - it('should replace invalid characters with underscores', () => { - const funcDecl = { name: 'invalid-name with spaces' }; - const serverName = 'test_server'; - const result = generateValidName(funcDecl, serverName); - expect(result).toBe('test_server__invalid-name_with_spaces'); - }); - - it('should truncate long names', () => { - const funcDecl = { - name: 'a_very_long_function_name_that_will_definitely_exceed_the_limit', - }; - const serverName = 'a_long_server_name'; - const result = generateValidName(funcDecl, serverName); - expect(result.length).toBe(63); - expect(result).toBe( - 'a_long_server_name__a_very_l___will_definitely_exceed_the_limit', - ); - }); - - it('should handle names with only invalid characters', () => { - const funcDecl = { name: '!@#$%^&*()' }; - const serverName = 'special-chars'; - const result = generateValidName(funcDecl, serverName); - expect(result).toBe('special-chars____________'); - }); - - it('should handle names that are already valid', () => { - const funcDecl = { name: 'already_valid' }; - const serverName = 'validator'; - const result = generateValidName(funcDecl, serverName); - expect(result).toBe('validator__already_valid'); - }); - - it('should handle names with leading/trailing invalid characters', () => { - const funcDecl = { name: '-_invalid-_' }; - const serverName = 'trim-test'; - const result = generateValidName(funcDecl, serverName); - expect(result).toBe('trim-test__-_invalid-_'); - }); - - it('should handle names that are exactly 63 characters long', () => { - const longName = 'a'.repeat(45); - const funcDecl = { name: longName }; - const serverName = 'server'; - const result = generateValidName(funcDecl, serverName); - expect(result).toBe(`server__${longName}`); - expect(result.length).toBe(53); - }); - - it('should handle names that are exactly 64 characters long', () => { - const longName = 'a'.repeat(55); - const funcDecl = { name: longName }; - const serverName = 'server'; - const result = generateValidName(funcDecl, serverName); - expect(result.length).toBe(63); - expect(result).toBe( - 'server__aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', - ); - }); - - it('should handle names that are longer than 64 characters', () => { - const longName = 'a'.repeat(100); - const funcDecl = { name: longName }; - const serverName = 'long-server'; - const result = generateValidName(funcDecl, serverName); - expect(result.length).toBe(63); - expect(result).toBe( - 'long-server__aaaaaaaaaaaaaaa___aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', - ); - }); - }); describe('isEnabled', () => { const funcDecl = { name: 'myTool' }; const serverName = 'myServer'; diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 5849884d..7e6b11c1 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -277,16 +277,13 @@ export async function discoverTools( continue; } - const toolNameForModel = generateValidName(funcDecl, mcpServerName); - discoveredTools.push( new DiscoveredMCPTool( mcpCallableTool, mcpServerName, - toolNameForModel, + funcDecl.name!, funcDecl.description ?? '', funcDecl.parametersJsonSchema ?? { type: 'object', properties: {} }, - funcDecl.name!, mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, mcpServerConfig.trust, ), @@ -427,26 +424,6 @@ export function createTransport( ); } -/** Visible for testing */ -export function generateValidName( - funcDecl: FunctionDeclaration, - mcpServerName: string, -) { - // Replace invalid characters (based on 400 error message from Gemini API) with underscores - let validToolname = funcDecl.name!.replace(/[^a-zA-Z0-9_.-]/g, '_'); - - // Prepend MCP server name to avoid conflicts with other tools - validToolname = mcpServerName + '__' + validToolname; - - // If longer than 63 characters, replace middle with '___' - // (Gemini API says max length 64, but actual limit seems to be 63) - if (validToolname.length > 63) { - validToolname = - validToolname.slice(0, 28) + '___' + validToolname.slice(-32); - } - return validToolname; -} - /** Visible for testing */ export function isEnabled( funcDecl: FunctionDeclaration, diff --git a/packages/core/src/tools/mcp-tool.test.ts b/packages/core/src/tools/mcp-tool.test.ts index 4411e674..2e700710 100644 --- a/packages/core/src/tools/mcp-tool.test.ts +++ b/packages/core/src/tools/mcp-tool.test.ts @@ -14,7 +14,7 @@ import { afterEach, Mocked, } from 'vitest'; -import { DiscoveredMCPTool } from './mcp-tool.js'; // Added getStringifiedResultForDisplay +import { DiscoveredMCPTool, generateValidName } from './mcp-tool.js'; // Added getStringifiedResultForDisplay import { ToolResult, ToolConfirmationOutcome } from './tools.js'; // Added ToolConfirmationOutcome import { CallableTool, Part } from '@google/genai'; @@ -29,9 +29,42 @@ const mockCallableToolInstance: Mocked = { // Add other methods if DiscoveredMCPTool starts using them }; +describe('generateValidName', () => { + it('should return a valid name for a simple function', () => { + expect(generateValidName('myFunction')).toBe('myFunction'); + }); + + it('should replace invalid characters with underscores', () => { + expect(generateValidName('invalid-name with spaces')).toBe( + 'invalid-name_with_spaces', + ); + }); + + it('should truncate long names', () => { + expect(generateValidName('x'.repeat(80))).toBe( + 'xxxxxxxxxxxxxxxxxxxxxxxxxxxx___xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx', + ); + }); + + it('should handle names with only invalid characters', () => { + expect(generateValidName('!@#$%^&*()')).toBe('__________'); + }); + + it('should handle names that are exactly 63 characters long', () => { + expect(generateValidName('a'.repeat(63)).length).toBe(63); + }); + + it('should handle names that are exactly 64 characters long', () => { + expect(generateValidName('a'.repeat(64)).length).toBe(63); + }); + + it('should handle names that are longer than 64 characters', () => { + expect(generateValidName('a'.repeat(80)).length).toBe(63); + }); +}); + describe('DiscoveredMCPTool', () => { 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: Record = { @@ -52,18 +85,17 @@ describe('DiscoveredMCPTool', () => { }); describe('constructor', () => { - it('should set properties correctly (non-generic server)', () => { + it('should set properties correctly', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, - serverName, // serverName is 'mock-mcp-server', not 'mcp' - toolNameForModel, + serverName, + serverToolName, baseDescription, inputSchema, - serverToolName, ); - expect(tool.name).toBe(toolNameForModel); - expect(tool.schema.name).toBe(toolNameForModel); + expect(tool.name).toBe(serverToolName); + expect(tool.schema.name).toBe(serverToolName); expect(tool.schema.description).toBe(baseDescription); expect(tool.schema.parameters).toBeUndefined(); expect(tool.schema.parametersJsonSchema).toEqual(inputSchema); @@ -71,30 +103,14 @@ describe('DiscoveredMCPTool', () => { expect(tool.timeout).toBeUndefined(); }); - it('should set properties correctly (generic "mcp" server)', () => { - const genericServerName = 'mcp'; - const tool = new DiscoveredMCPTool( - mockCallableToolInstance, - genericServerName, // serverName is 'mcp' - toolNameForModel, - baseDescription, - inputSchema, - serverToolName, - ); - expect(tool.schema.description).toBe(baseDescription); - expect(tool.schema.parameters).toBeUndefined(); - expect(tool.schema.parametersJsonSchema).toEqual(inputSchema); - }); - it('should accept and store a custom timeout', () => { const customTimeout = 5000; const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, customTimeout, ); expect(tool.timeout).toBe(customTimeout); @@ -106,10 +122,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, ); const params = { param: 'testValue' }; const mockToolSuccessResultObject = { @@ -146,10 +161,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, ); const params = { param: 'testValue' }; const mockMcpToolResponsePartsEmpty: Part[] = []; @@ -162,10 +176,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, ); const params = { param: 'failCase' }; const expectedError = new Error('MCP call failed'); @@ -182,10 +195,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, undefined, true, ); @@ -199,10 +211,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, ); expect( await tool.shouldConfirmExecute({}, new AbortController().signal), @@ -215,10 +226,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, ); expect( await tool.shouldConfirmExecute({}, new AbortController().signal), @@ -229,10 +239,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, ); const confirmation = await tool.shouldConfirmExecute( {}, @@ -260,10 +269,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, ); const confirmation = await tool.shouldConfirmExecute( {}, @@ -291,10 +299,9 @@ describe('DiscoveredMCPTool', () => { const tool = new DiscoveredMCPTool( mockCallableToolInstance, serverName, - toolNameForModel, + serverToolName, baseDescription, inputSchema, - serverToolName, ); const toolAllowlistKey = `${serverName}.${serverToolName}`; const confirmation = await tool.shouldConfirmExecute( diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index aadc484a..9916d7f9 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -28,15 +28,15 @@ export class DiscoveredMCPTool extends BaseTool { constructor( private readonly mcpTool: CallableTool, readonly serverName: string, - readonly name: string, - readonly description: string, - readonly parameterSchemaJson: unknown, readonly serverToolName: string, + description: string, + readonly parameterSchemaJson: unknown, readonly timeout?: number, readonly trust?: boolean, + nameOverride?: string, ) { super( - name, + nameOverride ?? generateValidName(serverToolName), `${serverToolName} (${serverName} MCP Server)`, description, Icon.Hammer, @@ -46,6 +46,19 @@ export class DiscoveredMCPTool extends BaseTool { ); } + asFullyQualifiedTool(): DiscoveredMCPTool { + return new DiscoveredMCPTool( + this.mcpTool, + this.serverName, + this.serverToolName, + this.description, + this.parameterSchemaJson, + this.timeout, + this.trust, + `${this.serverName}__${this.serverToolName}`, + ); + } + /** * Overrides the base schema to use parametersJsonSchema when building * FunctionDeclaration @@ -166,3 +179,17 @@ function getStringifiedResultForDisplay(result: Part[]) { return '```json\n' + JSON.stringify(processedResults, null, 2) + '\n```'; } + +/** Visible for testing */ +export function generateValidName(name: string) { + // Replace invalid characters (based on 400 error message from Gemini API) with underscores + let validToolname = name.replace(/[^a-zA-Z0-9_.-]/g, '_'); + + // If longer than 63 characters, replace middle with '___' + // (Gemini API says max length 64, but actual limit seems to be 63) + if (validToolname.length > 63) { + validToolname = + validToolname.slice(0, 28) + '___' + validToolname.slice(-32); + } + return validToolname; +} diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts index ec709a44..ab337252 100644 --- a/packages/core/src/tools/tool-registry.test.ts +++ b/packages/core/src/tools/tool-registry.test.ts @@ -210,35 +210,31 @@ describe('ToolRegistry', () => { const mcpTool1_c = new DiscoveredMCPTool( mockCallable, server1Name, - `${server1Name}__zebra-tool`, + 'zebra-tool', 'd1', {}, - 'zebra-tool', ); const mcpTool1_a = new DiscoveredMCPTool( mockCallable, server1Name, - `${server1Name}__apple-tool`, + 'apple-tool', 'd2', {}, - 'apple-tool', ); const mcpTool1_b = new DiscoveredMCPTool( mockCallable, server1Name, - `${server1Name}__banana-tool`, + 'banana-tool', 'd3', {}, - 'banana-tool', ); const mcpTool2 = new DiscoveredMCPTool( mockCallable, server2Name, - 'server2Name__tool-on-server2', + 'tool-on-server2', 'd4', {}, - 'tool-on-server2', ); const nonMcpTool = new MockTool('regular-tool'); @@ -253,11 +249,7 @@ describe('ToolRegistry', () => { // Assert that the array has the correct tools and is sorted by name expect(toolsFromServer1).toHaveLength(3); - expect(toolNames).toEqual([ - `${server1Name}__apple-tool`, - `${server1Name}__banana-tool`, - `${server1Name}__zebra-tool`, - ]); + expect(toolNames).toEqual(['apple-tool', 'banana-tool', 'zebra-tool']); // Assert that all returned tools are indeed from the correct server for (const tool of toolsFromServer1) { diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index e6a4121d..d6e84de3 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -18,7 +18,7 @@ type ToolParams = Record; export class DiscoveredTool extends BaseTool { constructor( private readonly config: Config, - readonly name: string, + name: string, readonly description: string, readonly parameterSchema: Record, ) { @@ -138,10 +138,14 @@ export class ToolRegistry { */ registerTool(tool: Tool): void { if (this.tools.has(tool.name)) { - // Decide on behavior: throw error, log warning, or allow overwrite - console.warn( - `Tool with name "${tool.name}" is already registered. Overwriting.`, - ); + if (tool instanceof DiscoveredMCPTool) { + tool = tool.asFullyQualifiedTool(); + } else { + // Decide on behavior: throw error, log warning, or allow overwrite + console.warn( + `Tool with name "${tool.name}" is already registered. Overwriting.`, + ); + } } this.tools.set(tool.name, tool); }