Skip to content
83 changes: 83 additions & 0 deletions packages/core/src/policy/policy-engine.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,43 @@ vi.mock('../utils/shell-utils.js', async (importOriginal) => {
};
});

// Mock tool-names to provide a consistent alias for testing

vi.mock('../tools/tool-names.js', async (importOriginal) => {
const actual =
await importOriginal<typeof import('../tools/tool-names.js')>();

const mockedAliases: Record<string, string> = {
...actual.TOOL_LEGACY_ALIASES,

legacy_test_tool: 'current_test_tool',

another_legacy_test_tool: 'current_test_tool',
};

return {
...actual,

TOOL_LEGACY_ALIASES: mockedAliases,

getToolAliases: vi.fn().mockImplementation((name: 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);
}),
};
});

describe('PolicyEngine', () => {
let engine: PolicyEngine;
let mockCheckerRunner: CheckerRunner;
Expand Down Expand Up @@ -187,6 +224,52 @@ describe('PolicyEngine', () => {
);
});

it('should match current tool call against legacy tool name rules', async () => {
const legacyName = 'legacy_test_tool';
const currentName = 'current_test_tool';

const rules: PolicyRule[] = [
{ toolName: legacyName, decision: PolicyDecision.DENY },
];

engine = new PolicyEngine({ rules });

// Call using the CURRENT name, should be denied because of legacy rule
const { decision } = await engine.check({ name: currentName }, undefined);
expect(decision).toBe(PolicyDecision.DENY);
});

it('should match legacy tool call against current tool name rules (for skills support)', async () => {
const legacyName = 'legacy_test_tool';
const currentName = 'current_test_tool';

const rules: PolicyRule[] = [
{ toolName: currentName, decision: PolicyDecision.ALLOW },
];

engine = new PolicyEngine({ rules });

// Call using the LEGACY name (from a skill), should be allowed because of current rule
const { decision } = await engine.check({ name: legacyName }, undefined);
expect(decision).toBe(PolicyDecision.ALLOW);
});

it('should match tool call using one legacy name against policy for another legacy name (same canonical tool)', async () => {
const legacyName1 = 'legacy_test_tool';
const legacyName2 = 'another_legacy_test_tool';

const rules: PolicyRule[] = [
{ toolName: legacyName2, decision: PolicyDecision.DENY },
];

engine = new PolicyEngine({ rules });

// Call using legacyName1, should be denied because legacyName2 has a deny rule
// and they both point to the same canonical tool.
const { decision } = await engine.check({ name: legacyName1 }, undefined);
expect(decision).toBe(PolicyDecision.DENY);
});

it('should apply wildcard rules (no toolName)', async () => {
const rules: PolicyRule[] = [
{ decision: PolicyDecision.DENY }, // Applies to all tools
Expand Down
25 changes: 19 additions & 6 deletions packages/core/src/policy/policy-engine.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import {
splitCommands,
hasRedirection,
} from '../utils/shell-utils.js';
import { getToolAliases } from '../tools/tool-names.js';

function ruleMatches(
rule: PolicyRule | SafetyCheckerRule,
Expand Down Expand Up @@ -308,12 +309,24 @@ export class PolicyEngine {

// For tools with a server name, we want to try matching both the
// original name and the fully qualified name (server__tool).
const toolCallsToTry: FunctionCall[] = [toolCall];
if (serverName && toolCall.name && !toolCall.name.includes('__')) {
toolCallsToTry.push({
...toolCall,
name: `${serverName}__${toolCall.name}`,
});
// We also want to check legacy aliases for the tool name.
const toolNamesToTry = new Set<string>();
if (toolCall.name) {
const aliases = getToolAliases(toolCall.name);
for (const alias of aliases) {
toolNamesToTry.add(alias);
}
}

const toolCallsToTry: FunctionCall[] = [];
for (const name of toolNamesToTry) {
toolCallsToTry.push({ ...toolCall, name });
if (serverName && !name.includes('__')) {
toolCallsToTry.push({
...toolCall,
name: `${serverName}__${name}`,
});
}
}

for (const rule of this.rules) {
Expand Down
60 changes: 59 additions & 1 deletion packages/core/src/tools/tool-names.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,44 @@
* SPDX-License-Identifier: Apache-2.0
*/

import { describe, it, expect } from 'vitest';
import { describe, it, expect, vi } from 'vitest';
import {
isValidToolName,
getToolAliases,
ALL_BUILTIN_TOOL_NAMES,
DISCOVERED_TOOL_PREFIX,
LS_TOOL_NAME,
} from './tool-names.js';

// 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',
another_legacy_test_tool: 'current_test_tool',
};
return {
...actual,
TOOL_LEGACY_ALIASES: mockedAliases,
isValidToolName: vi.fn().mockImplementation((name: string, options) => {
if (mockedAliases[name]) return true;
return actual.isValidToolName(name, options);
}),
getToolAliases: vi.fn().mockImplementation((name: 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);
}),
};
});

describe('tool-names', () => {
describe('isValidToolName', () => {
it('should validate built-in tool names', () => {
Expand All @@ -30,6 +60,13 @@ describe('tool-names', () => {
expect(isValidToolName('my-server__my-tool')).toBe(true);
});

it('should validate legacy tool aliases', async () => {
const { TOOL_LEGACY_ALIASES } = await import('./tool-names.js');
for (const legacyName of Object.keys(TOOL_LEGACY_ALIASES)) {
expect(isValidToolName(legacyName)).toBe(true);
}
});

it('should reject invalid tool names', () => {
expect(isValidToolName('')).toBe(false);
expect(isValidToolName('invalid-name')).toBe(false);
Expand All @@ -54,4 +91,25 @@ describe('tool-names', () => {
);
});
});

describe('getToolAliases', () => {
it('should return all associated names for a current tool', () => {
const aliases = getToolAliases('current_test_tool');
expect(aliases).toContain('current_test_tool');
expect(aliases).toContain('legacy_test_tool');
expect(aliases).toContain('another_legacy_test_tool');
});

it('should return all associated names for a legacy tool', () => {
const aliases = getToolAliases('legacy_test_tool');
expect(aliases).toContain('current_test_tool');
expect(aliases).toContain('legacy_test_tool');
expect(aliases).toContain('another_legacy_test_tool');
});

it('should return only the name itself if no aliases exist', () => {
const aliases = getToolAliases('unknown_tool');
expect(aliases).toEqual(['unknown_tool']);
});
});
});
38 changes: 37 additions & 1 deletion packages/core/src/tools/tool-names.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,38 @@ export const EDIT_TOOL_NAMES = new Set([EDIT_TOOL_NAME, WRITE_FILE_TOOL_NAME]);
export const ASK_USER_TOOL_NAME = 'ask_user';
export const ASK_USER_DISPLAY_NAME = 'Ask User';

/** Prefix used for tools discovered via the toolDiscoveryCommand. */
/**
* Mapping of legacy tool names to their current names.
* This ensures backward compatibility for user-defined policies, skills, and hooks.
*/
export const TOOL_LEGACY_ALIASES: Record<string, string> = {
// Add future renames here, e.g.:
// 'search_file_content': GREP_TOOL_NAME,
};

/**
* Returns all associated names for a tool (including legacy aliases and current name).
* This ensures that if multiple legacy names point to the same tool, we consider all of them
* for policy application.
*/
export function getToolAliases(name: string): string[] {
const aliases = new Set<string>([name]);

// Determine the canonical (current) name
const canonicalName = TOOL_LEGACY_ALIASES[name] ?? name;
aliases.add(canonicalName);

// Find all other legacy aliases that point to the same canonical name
for (const [legacyName, currentName] of Object.entries(TOOL_LEGACY_ALIASES)) {
if (currentName === canonicalName) {
aliases.add(legacyName);
}
}

return Array.from(aliases);
}

/** Prefix used for tools discovered via the tool DiscoveryCommand. */
export const DISCOVERED_TOOL_PREFIX = 'discovered_tool_';

/**
Expand Down Expand Up @@ -76,6 +107,11 @@ export function isValidToolName(
return true;
}

// Legacy aliases
if (TOOL_LEGACY_ALIASES[name]) {
return true;
}

// Discovered tools
if (name.startsWith(DISCOVERED_TOOL_PREFIX)) {
return true;
Expand Down
29 changes: 29 additions & 0 deletions packages/core/src/tools/tool-registry.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,18 @@ 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')>();
return {
...actual,
TOOL_LEGACY_ALIASES: {
...actual.TOOL_LEGACY_ALIASES,
legacy_test_tool: 'current_test_tool',
},
};
});

// Helper to create a mock CallableTool for specific test needs
const createMockCallableTool = (
toolDeclarations: FunctionDeclaration[],
Expand Down Expand Up @@ -584,6 +596,23 @@ describe('ToolRegistry', () => {
expect(declarations).toHaveLength(1);
expect(declarations[0].name).toBe(toolName);
});

it('should retrieve a tool using its legacy alias', async () => {
const legacyName = 'legacy_test_tool';
const currentName = 'current_test_tool';

const mockTool = new MockTool({
name: currentName,
description: 'Test Tool',
messageBus: mockMessageBus,
});

toolRegistry.registerTool(mockTool);

const retrievedTool = toolRegistry.getTool(legacyName);
expect(retrievedTool).toBeDefined();
expect(retrievedTool?.name).toBe(currentName);
});
});

describe('DiscoveredToolInvocation', () => {
Expand Down
14 changes: 13 additions & 1 deletion packages/core/src/tools/tool-registry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ 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 } from './tool-names.js';
import { DISCOVERED_TOOL_PREFIX, TOOL_LEGACY_ALIASES } from './tool-names.js';

type ToolParams = Record<string, unknown>;

Expand Down Expand Up @@ -531,6 +531,18 @@ export class ToolRegistry {
*/
getTool(name: string): AnyDeclarativeTool | undefined {
let tool = this.allKnownTools.get(name);

// If not found, check legacy aliases
if (!tool && TOOL_LEGACY_ALIASES[name]) {
const currentName = TOOL_LEGACY_ALIASES[name];
tool = this.allKnownTools.get(currentName);
if (tool) {
debugLogger.debug(
`Resolved legacy tool name "${name}" to current name "${currentName}"`,
);
}
}

if (!tool && name.includes('__')) {
for (const t of this.allKnownTools.values()) {
if (t instanceof DiscoveredMCPTool) {
Expand Down
Loading