Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
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
39 changes: 36 additions & 3 deletions packages/core/src/tools/tool-registry.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -84,11 +84,24 @@ vi.mock('@google/genai', async () => {
// Mock tool-names to provide a consistent alias for testing
vi.mock('./tool-names.js', async (importOriginal) => {
const actual = await importOriginal<typeof import('./tool-names.js')>();
const mockedAliases: Record<string, string> = {
...actual.TOOL_LEGACY_ALIASES,
legacy_test_tool: 'current_test_tool',
};
return {
...actual,
TOOL_LEGACY_ALIASES: {
...actual.TOOL_LEGACY_ALIASES,
legacy_test_tool: 'current_test_tool',
TOOL_LEGACY_ALIASES: mockedAliases,
// Override getToolAliases to use the mocked aliases map
getToolAliases: (name: string): string[] => {
const aliases = new Set<string>([name]);
const canonicalName = mockedAliases[name] ?? name;
aliases.add(canonicalName);
for (const [legacyName, currentName] of Object.entries(mockedAliases)) {
if (currentName === canonicalName) {
aliases.add(legacyName);
}
}
return Array.from(aliases);
},
};
});
Expand Down Expand Up @@ -290,6 +303,26 @@ describe('ToolRegistry', () => {
tools: [excludedTool],
excludedTools: ['ExcludedMockTool'],
},
{
name: 'should exclude a tool when its legacy alias is in excludeTools',
tools: [
new MockTool({
name: 'current_test_tool',
displayName: 'Current Test Tool',
}),
],
excludedTools: ['legacy_test_tool'],
},
{
name: 'should exclude a tool when its current name is in excludeTools and tool is registered under current name',
tools: [
new MockTool({
name: 'current_test_tool',
displayName: 'Current Test Tool',
}),
],
excludedTools: ['current_test_tool'],
},
])('$name', ({ tools, excludedTools }) => {
toolRegistry.registerTool(allowedTool);
for (const tool of tools) {
Expand Down
34 changes: 31 additions & 3 deletions packages/core/src/tools/tool-registry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,11 @@ import { safeJsonStringify } from '../utils/safeJsonStringify.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { debugLogger } from '../utils/debugLogger.js';
import { coreEvents } from '../utils/events.js';
import { DISCOVERED_TOOL_PREFIX, TOOL_LEGACY_ALIASES } from './tool-names.js';
import {
DISCOVERED_TOOL_PREFIX,
TOOL_LEGACY_ALIASES,
getToolAliases,
} from './tool-names.js';

type ToolParams = Record<string, unknown>;

Expand Down Expand Up @@ -431,7 +435,9 @@ export class ToolRegistry {
* @returns All the tools that are not excluded.
*/
private getActiveTools(): AnyDeclarativeTool[] {
const excludedTools = this.config.getExcludeTools() ?? new Set([]);
const excludedTools =
this.expandExcludeToolsWithAliases(this.config.getExcludeTools()) ??
new Set([]);
const activeTools: AnyDeclarativeTool[] = [];
for (const tool of this.allKnownTools.values()) {
if (this.isActiveTool(tool, excludedTools)) {
Expand All @@ -441,6 +447,26 @@ export class ToolRegistry {
return activeTools;
}

/**
* Expands an excludeTools set to include all legacy aliases.
* For example, if 'search_file_content' is excluded and it's an alias for
* 'grep_search', both names will be in the returned set.
*/
private expandExcludeToolsWithAliases(
excludeTools: Set<string> | undefined,
): Set<string> | undefined {
if (!excludeTools || excludeTools.size === 0) {
return excludeTools;
}
const expanded = new Set<string>();
for (const name of excludeTools) {
for (const alias of getToolAliases(name)) {
expanded.add(alias);
}
}
return expanded;
}

/**
* @param tool
* @param excludeTools (optional, helps performance for repeated calls)
Expand All @@ -450,7 +476,9 @@ export class ToolRegistry {
tool: AnyDeclarativeTool,
excludeTools?: Set<string>,
): boolean {
excludeTools ??= this.config.getExcludeTools() ?? new Set([]);
excludeTools ??=
this.expandExcludeToolsWithAliases(this.config.getExcludeTools()) ??
new Set([]);
Comment thread
SandyTao520 marked this conversation as resolved.
Resolved
const normalizedClassName = tool.constructor.name.replace(/^_+/, '');
const possibleNames = [tool.name, normalizedClassName];
if (tool instanceof DiscoveredMCPTool) {
Expand Down
Loading