feat(core): enforce server prefixes for MCP tools in agent definitions (#17574)

This commit is contained in:
Abhi
2026-01-26 23:53:05 -05:00
committed by GitHub
parent 9fcdc0cdc1
commit eccc200f4f
5 changed files with 94 additions and 5 deletions

View File

@@ -17,6 +17,10 @@ import { debugLogger } from '../utils/debugLogger.js';
import { LocalAgentExecutor, type ActivityCallback } from './local-executor.js';
import { makeFakeConfig } from '../test-utils/config.js';
import { ToolRegistry } from '../tools/tool-registry.js';
import {
DiscoveredMCPTool,
MCP_QUALIFIED_NAME_SEPARATOR,
} from '../tools/mcp-tool.js';
import { LSTool } from '../tools/ls.js';
import { LS_TOOL_NAME, READ_FILE_TOOL_NAME } from '../tools/tool-names.js';
import {
@@ -31,6 +35,7 @@ import {
type Content,
type PartListUnion,
type Tool,
type CallableTool,
} from '@google/genai';
import type { Config } from '../config/config.js';
import { MockTool } from '../test-utils/mock-tool.js';
@@ -493,6 +498,56 @@ describe('LocalAgentExecutor', () => {
// Should exclude subagent
expect(agentRegistry.getTool(subAgentName)).toBeUndefined();
});
it('should enforce qualified names for MCP tools in agent definitions', async () => {
const serverName = 'mcp-server';
const toolName = 'mcp-tool';
const qualifiedName = `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}`;
const mockMcpTool = {
tool: vi.fn(),
callTool: vi.fn(),
} as unknown as CallableTool;
const mcpTool = new DiscoveredMCPTool(
mockMcpTool,
serverName,
toolName,
'description',
{},
mockConfig.getMessageBus(),
);
// Mock getTool to return our real DiscoveredMCPTool instance
const getToolSpy = vi
.spyOn(parentToolRegistry, 'getTool')
.mockImplementation((name) => {
if (name === toolName || name === qualifiedName) {
return mcpTool;
}
return undefined;
});
// 1. Qualified name works and registers the tool (using short name per status quo)
const definition = createTestDefinition([qualifiedName]);
const executor = await LocalAgentExecutor.create(
definition,
mockConfig,
onActivity,
);
const agentRegistry = executor['toolRegistry'];
// Registry shortening logic means it's registered as 'mcp-tool' internally
expect(agentRegistry.getTool(toolName)).toBeDefined();
// 2. Unqualified name for MCP tool THROWS
const badDefinition = createTestDefinition([toolName]);
await expect(
LocalAgentExecutor.create(badDefinition, mockConfig, onActivity),
).rejects.toThrow(/must be requested with its server prefix/);
getToolSpy.mockRestore();
});
});
describe('run (Execution Loop and Logic)', () => {

View File

@@ -16,6 +16,10 @@ import type {
Schema,
} from '@google/genai';
import { ToolRegistry } from '../tools/tool-registry.js';
import {
DiscoveredMCPTool,
MCP_QUALIFIED_NAME_SEPARATOR,
} from '../tools/mcp-tool.js';
import { CompressionStatus } from '../core/turn.js';
import { type ToolCallRequestInfo } from '../scheduler/types.js';
import { ChatCompressionService } from '../services/chatCompressionService.js';
@@ -129,6 +133,14 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// registry and register it with the agent's isolated registry.
const tool = parentToolRegistry.getTool(toolName);
if (tool) {
if (
tool instanceof DiscoveredMCPTool &&
!toolName.includes(MCP_QUALIFIED_NAME_SEPARATOR)
) {
throw new Error(
`MCP tool '${toolName}' must be requested with its server prefix (e.g., '${tool.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}') in agent '${definition.name}'.`,
);
}
agentToolRegistry.registerTool(tool);
}
};

View File

@@ -23,6 +23,12 @@ import { ToolErrorType } from './tool-error.js';
import type { Config } from '../config/config.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
/**
* The separator used to qualify MCP tool names with their server prefix.
* e.g. "server_name__tool_name"
*/
export const MCP_QUALIFIED_NAME_SEPARATOR = '__';
type ToolParams = Record<string, unknown>;
// Discriminated union for MCP Content Blocks to ensure type safety.
@@ -82,7 +88,7 @@ export class DiscoveredMCPToolInvocation extends BaseToolInvocation<
super(
params,
messageBus,
`${serverName}__${serverToolName}`,
`${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${serverToolName}`,
displayName,
serverName,
);
@@ -261,7 +267,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
}
getFullyQualifiedPrefix(): string {
return `${this.serverName}__`;
return `${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}`;
}
getFullyQualifiedName(): string {

View File

@@ -13,7 +13,7 @@ import { ApprovalMode } from '../policy/types.js';
import { ToolRegistry, DiscoveredTool } from './tool-registry.js';
import { DISCOVERED_TOOL_PREFIX } from './tool-names.js';
import { DiscoveredMCPTool } from './mcp-tool.js';
import { DiscoveredMCPTool, MCP_QUALIFIED_NAME_SEPARATOR } from './mcp-tool.js';
import type { FunctionDeclaration, CallableTool } from '@google/genai';
import { mcpToTool } from '@google/genai';
import { spawn } from 'node:child_process';
@@ -568,6 +568,22 @@ describe('ToolRegistry', () => {
expect(retrievedTool).toBeDefined();
expect(retrievedTool?.name).toBe(validToolName);
});
it('should resolve qualified names in getFunctionDeclarationsFiltered', () => {
const serverName = 'my-server';
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');
toolRegistry.registerTool(mcpTool);
const fullyQualifiedName = `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}`;
const declarations = toolRegistry.getFunctionDeclarationsFiltered([
fullyQualifiedName,
]);
expect(declarations).toHaveLength(1);
expect(declarations[0].name).toBe(toolName);
});
});
describe('DiscoveredToolInvocation', () => {

View File

@@ -488,8 +488,8 @@ export class ToolRegistry {
getFunctionDeclarationsFiltered(toolNames: string[]): FunctionDeclaration[] {
const declarations: FunctionDeclaration[] = [];
for (const name of toolNames) {
const tool = this.allKnownTools.get(name);
if (tool && this.isActiveTool(tool)) {
const tool = this.getTool(name);
if (tool) {
declarations.push(tool.schema);
}
}