Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 14 additions & 10 deletions packages/core/src/agents/local-executor.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -501,7 +501,7 @@ describe('LocalAgentExecutor', () => {
expect(agentRegistry.getTool(subAgentName)).toBeUndefined();
});

it('should enforce qualified names for MCP tools in agent definitions', async () => {
it('should automatically qualify MCP tools in agent definitions', async () => {
const serverName = 'mcp-server';
const toolName = 'mcp-tool';
const qualifiedName = `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}`;
Expand Down Expand Up @@ -530,7 +530,7 @@ describe('LocalAgentExecutor', () => {
return undefined;
});

// 1. Qualified name works and registers the tool (using short name per status quo)
// 1. Qualified name works and registers the tool (using qualified name)
const definition = createTestDefinition([qualifiedName]);
const executor = await LocalAgentExecutor.create(
definition,
Expand All @@ -539,14 +539,18 @@ describe('LocalAgentExecutor', () => {
);

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/);
// It should be registered as the qualified name
expect(agentRegistry.getTool(qualifiedName)).toBeDefined();

// 2. Unqualified name for MCP tool now also works (and gets upgraded to qualified)
const definition2 = createTestDefinition([toolName]);
const executor2 = await LocalAgentExecutor.create(
definition2,
mockConfig,
onActivity,
);
const agentRegistry2 = executor2['toolRegistry'];
expect(agentRegistry2.getTool(qualifiedName)).toBeDefined();

getToolSpy.mockRestore();
});
Expand Down
20 changes: 8 additions & 12 deletions packages/core/src/agents/local-executor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,7 @@ 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 { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { CompressionStatus } from '../core/turn.js';
import { type ToolCallRequestInfo } from '../scheduler/types.js';
import { ChatCompressionService } from '../services/chatCompressionService.js';
Expand Down Expand Up @@ -142,15 +139,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}'.`,
);
if (tool instanceof DiscoveredMCPTool) {
// Subagents MUST use fully qualified names for MCP tools to ensure
// unambiguous tool calls and to comply with policy requirements.
// We automatically "upgrade" any MCP tool to its qualified version.
agentToolRegistry.registerTool(tool.asFullyQualifiedTool());
} else {
agentToolRegistry.registerTool(tool);
}
agentToolRegistry.registerTool(tool);
}
};

Expand Down
58 changes: 51 additions & 7 deletions packages/core/src/tools/tool-registry.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -380,20 +380,36 @@ describe('ToolRegistry', () => {
});

describe('getAllToolNames', () => {
it('should return all registered tool names', () => {
it('should return all registered tool names with qualified names for MCP tools', () => {
// Register tools with displayNames in non-alphabetical order
const toolC = new MockTool({ name: 'c-tool', displayName: 'Tool C' });
const toolA = new MockTool({ name: 'a-tool', displayName: 'Tool A' });
const toolB = new MockTool({ name: 'b-tool', displayName: 'Tool B' });
const mcpTool = createMCPTool('my-server', 'my-tool', 'desc');

toolRegistry.registerTool(toolC);
toolRegistry.registerTool(toolA);
toolRegistry.registerTool(toolB);
toolRegistry.registerTool(mcpTool);

const toolNames = toolRegistry.getAllToolNames();

// Assert that the returned array contains all tool names
expect(toolNames).toEqual(['c-tool', 'a-tool', 'b-tool']);
// Assert that the returned array contains all tool names, with MCP qualified
expect(toolNames).toContain('c-tool');
expect(toolNames).toContain('a-tool');
expect(toolNames).toContain('my-server__my-tool');
expect(toolNames).toHaveLength(3);
});

it('should deduplicate tool names', () => {
const serverName = 'my-server';
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'desc');

// Register same MCP tool twice (one as alias, one as qualified)
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());

const toolNames = toolRegistry.getAllToolNames();
expect(toolNames).toEqual([`${serverName}__${toolName}`]);
});
});

Expand Down Expand Up @@ -465,8 +481,8 @@ describe('ToolRegistry', () => {
'builtin-1',
'builtin-2',
DISCOVERED_TOOL_PREFIX + 'discovered-1',
'mcp-apple',
'mcp-zebra',
'apple-server__mcp-apple',
'zebra-server__mcp-zebra',
]);
});
});
Expand Down Expand Up @@ -659,6 +675,34 @@ describe('ToolRegistry', () => {
});
});

describe('getFunctionDeclarations', () => {
it('should use fully qualified names for MCP tools in declarations', () => {
const serverName = 'my-server';
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');

toolRegistry.registerTool(mcpTool);

const declarations = toolRegistry.getFunctionDeclarations();
expect(declarations).toHaveLength(1);
expect(declarations[0].name).toBe(`${serverName}__${toolName}`);
});

it('should deduplicate MCP tools in declarations', () => {
const serverName = 'my-server';
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');

// Register both alias and qualified
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());

const declarations = toolRegistry.getFunctionDeclarations();
expect(declarations).toHaveLength(1);
expect(declarations[0].name).toBe(`${serverName}__${toolName}`);
});
});

describe('plan mode', () => {
it('should only return policy-allowed tools in plan mode', () => {
// Register several tools
Expand Down
55 changes: 49 additions & 6 deletions packages/core/src/tools/tool-registry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -539,11 +539,32 @@ export class ToolRegistry {
const plansDir = this.config.storage.getPlansDir();

const declarations: FunctionDeclaration[] = [];
const seenNames = new Set<string>();

this.getActiveTools().forEach((tool) => {
const toolName =
tool instanceof DiscoveredMCPTool
? tool.getFullyQualifiedName()
: tool.name;

if (seenNames.has(toolName)) {
return;
}
seenNames.add(toolName);

let schema = tool.getSchema(modelId);

// Ensure the schema name matches the qualified name for MCP tools
if (tool instanceof DiscoveredMCPTool) {
schema = {
...schema,
name: toolName,
};
}

if (
isPlanMode &&
(tool.name === WRITE_FILE_TOOL_NAME || tool.name === EDIT_TOOL_NAME)
(toolName === WRITE_FILE_TOOL_NAME || toolName === EDIT_TOOL_NAME)
) {
schema = {
...schema,
Expand Down Expand Up @@ -576,20 +597,42 @@ export class ToolRegistry {
}

/**
* Returns an array of all registered and discovered tool names which are not
* excluded via configuration.
* Returns an array of names for all active tools.
* For MCP tools, this returns their fully qualified names.
* The list is deduplicated.
*/
getAllToolNames(): string[] {
return this.getActiveTools().map((tool) => tool.name);
const names = new Set<string>();
for (const tool of this.getActiveTools()) {
if (tool instanceof DiscoveredMCPTool) {
names.add(tool.getFullyQualifiedName());
} else {
names.add(tool.name);
}
}
return Array.from(names);
}

/**
* Returns an array of all registered and discovered tool instances.
*/
getAllTools(): AnyDeclarativeTool[] {
return this.getActiveTools().sort((a, b) =>
const seen = new Set<string>();
const tools: AnyDeclarativeTool[] = [];

for (const tool of this.getActiveTools().sort((a, b) =>
a.displayName.localeCompare(b.displayName),
);
)) {
const name =
tool instanceof DiscoveredMCPTool
? tool.getFullyQualifiedName()
: tool.name;
if (!seen.has(name)) {
seen.add(name);
tools.push(tool);
}
}
return tools;
}

/**
Expand Down
Loading