mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-02-01 22:48:03 +00:00
feat(core): enforce server prefixes for MCP tools in agent definitions (#17574)
This commit is contained in:
@@ -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)', () => {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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', () => {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user