fix(core): sanitize and length-check MCP tool qualified names (#20987)

This commit is contained in:
Abhi
2026-03-03 16:38:52 -05:00
committed by GitHub
parent fdf5b18cfb
commit 28e79831ac
4 changed files with 115 additions and 14 deletions

View File

@@ -54,7 +54,7 @@ describe('generateValidName', () => {
it('should truncate long names', () => {
expect(generateValidName('x'.repeat(80))).toBe(
'xxxxxxxxxxxxxxxxxxxxxxxxxxxx___xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx...xxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
);
});
@@ -933,3 +933,82 @@ describe('DiscoveredMCPTool', () => {
});
});
});
describe('MCP Tool Naming Regression Fixes', () => {
describe('generateValidName', () => {
it('should replace spaces with underscores', () => {
expect(generateValidName('My Tool')).toBe('My_Tool');
});
it('should allow colons', () => {
expect(generateValidName('namespace:tool')).toBe('namespace:tool');
});
it('should ensure name starts with a letter or underscore', () => {
expect(generateValidName('123-tool')).toBe('_123-tool');
expect(generateValidName('-tool')).toBe('_-tool');
expect(generateValidName('.tool')).toBe('_.tool');
});
it('should handle very long names by truncating in the middle', () => {
const longName = 'a'.repeat(40) + '__' + 'b'.repeat(40);
const result = generateValidName(longName);
expect(result.length).toBeLessThanOrEqual(63);
expect(result).toMatch(/^a{30}\.\.\.b{30}$/);
});
it('should handle very long names starting with a digit', () => {
const longName = '1' + 'a'.repeat(80);
const result = generateValidName(longName);
expect(result.length).toBeLessThanOrEqual(63);
expect(result.startsWith('_1')).toBe(true);
});
});
describe('DiscoveredMCPTool qualified names', () => {
it('should generate a valid qualified name even with spaces in server name', () => {
const tool = new DiscoveredMCPTool(
{} as any,
'My Server',
'my-tool',
'desc',
{},
{} as any,
);
const qn = tool.getFullyQualifiedName();
expect(qn).toBe('My_Server__my-tool');
});
it('should handle long server and tool names in qualified name', () => {
const serverName = 'a'.repeat(40);
const toolName = 'b'.repeat(40);
const tool = new DiscoveredMCPTool(
{} as any,
serverName,
toolName,
'desc',
{},
{} as any,
);
const qn = tool.getFullyQualifiedName();
expect(qn.length).toBeLessThanOrEqual(63);
expect(qn).toContain('...');
});
it('should handle server names starting with digits', () => {
const tool = new DiscoveredMCPTool(
{} as any,
'123-server',
'tool',
'desc',
{},
{} as any,
);
const qn = tool.getFullyQualifiedName();
expect(qn).toBe('_123-server__tool');
});
});
});

View File

@@ -96,14 +96,17 @@ export class DiscoveredMCPToolInvocation extends BaseToolInvocation<
) {
// Use composite format for policy checks: serverName__toolName
// This enables server wildcards (e.g., "google-workspace__*")
// while still allowing specific tool rules
// while still allowing specific tool rules.
// We use the same sanitized names as the registry to ensure policy matches.
super(
params,
messageBus,
`${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${serverToolName}`,
generateValidName(
`${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${serverToolName}`,
),
displayName,
serverName,
generateValidName(serverName),
toolAnnotationsData,
);
}
@@ -273,7 +276,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
private readonly _toolAnnotations?: Record<string, unknown>,
) {
super(
nameOverride ?? generateValidName(serverToolName),
generateValidName(nameOverride ?? serverToolName),
`${serverToolName} (${serverName} MCP Server)`,
description,
Kind.Other,
@@ -305,7 +308,9 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
}
getFullyQualifiedName(): string {
return `${this.getFullyQualifiedPrefix()}${generateValidName(this.serverToolName)}`;
return generateValidName(
`${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${this.serverToolName}`,
);
}
asFullyQualifiedTool(): DiscoveredMCPTool {
@@ -482,16 +487,29 @@ function getStringifiedResultForDisplay(rawResponse: Part[]): string {
return displayParts.join('\n');
}
/**
* Maximum length for a function name in the Gemini API.
* @see https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-reference/function-calling#functiondeclaration
*/
const MAX_FUNCTION_NAME_LENGTH = 64;
/** 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, '_');
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);
// Ensure it starts with a letter or underscore
if (/^[^a-zA-Z_]/.test(validToolname)) {
validToolname = `_${validToolname}`;
}
// If longer than the API limit, replace middle with '...'
// Note: We use 63 instead of 64 to be safe, as some environments have off-by-one behaviors.
const safeLimit = MAX_FUNCTION_NAME_LENGTH - 1;
if (validToolname.length > safeLimit) {
validToolname =
validToolname.slice(0, 30) + '...' + validToolname.slice(-30);
}
return validToolname;
}

View File

@@ -58,6 +58,8 @@ describe('tool-names', () => {
it('should validate MCP tool names (server__tool)', () => {
expect(isValidToolName('server__tool')).toBe(true);
expect(isValidToolName('my-server__my-tool')).toBe(true);
expect(isValidToolName('my.server__my:tool')).toBe(true);
expect(isValidToolName('my-server...truncated__tool')).toBe(true);
});
it('should validate legacy tool aliases', async () => {

View File

@@ -260,8 +260,10 @@ export function isValidToolName(
return !!options.allowWildcards;
}
// Basic slug validation for server and tool names
const slugRegex = /^[a-z0-9-_]+$/i;
// Basic slug validation for server and tool names.
// We allow dots (.) and colons (:) as they are valid in function names and
// used for truncation markers.
const slugRegex = /^[a-z0-9_.:-]+$/i;
return slugRegex.test(server) && slugRegex.test(tool);
}