From 3e5a250b6fb47bc145c31b8cb86f556bf673b13f Mon Sep 17 00:00:00 2001 From: Mahima Shanware Date: Tue, 24 Feb 2026 00:02:43 +0000 Subject: [PATCH] feat(tools): delegate Plan Mode tool filtering to PolicyEngine - Removes hardcoded `PLAN_MODE_TOOLS` from `ToolRegistry` and delegates permission checks to the `PolicyEngine` using the new `isToolPotentiallyAllowed` method. - Preserves Plan Mode bypass logic for read-only MCP tools while ensuring they respect explicit/wildcard policy blocks. - Makes `getActiveTools` public with TSDoc documentation for use in prompt generation. - Adds regression tests to `tool-registry.test.ts` to verify policy-driven tool availability. --- packages/core/src/tools/tool-names.ts | 16 -------- packages/core/src/tools/tool-registry.test.ts | 40 +++++++++++++++++++ packages/core/src/tools/tool-registry.ts | 40 ++++++++++--------- 3 files changed, 61 insertions(+), 35 deletions(-) diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index 5cc1dc6e3a..9905fb44b3 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -112,22 +112,6 @@ export const ALL_BUILTIN_TOOL_NAMES = [ EXIT_PLAN_MODE_TOOL_NAME, ] as const; -/** - * Read-only tools available in Plan Mode. - * This list is used to dynamically generate the Plan Mode prompt, - * filtered by what tools are actually enabled in the current configuration. - */ -export const PLAN_MODE_TOOLS = [ - GLOB_TOOL_NAME, - GREP_TOOL_NAME, - READ_FILE_TOOL_NAME, - LS_TOOL_NAME, - WEB_SEARCH_TOOL_NAME, - ASK_USER_TOOL_NAME, - ACTIVATE_SKILL_TOOL_NAME, - EXIT_PLAN_MODE_TOOL_NAME, -] as const; - /** * Validates if a tool name is syntactically valid. * Checks against built-in tools, discovered tools, and MCP naming conventions. diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts index 963830200d..cec1016813 100644 --- a/packages/core/src/tools/tool-registry.test.ts +++ b/packages/core/src/tools/tool-registry.test.ts @@ -379,6 +379,46 @@ describe('ToolRegistry', () => { }); }); + describe('getActiveTools', () => { + it('should include tools in Plan Mode when allowed by policy', () => { + // Setup Plan Mode + vi.spyOn(config, 'getApprovalMode').mockReturnValue(ApprovalMode.PLAN); + + const mockTool = new MockTool({ name: 'mock-tool' }); + toolRegistry.registerTool(mockTool); + + // Mock PolicyEngine to explicitly allow this tool + const policyEngine = config.getPolicyEngine(); + vi.spyOn(policyEngine, 'isToolPotentiallyAllowed').mockImplementation( + (toolName) => toolName === 'mock-tool', + ); + + const activeTools = toolRegistry.getActiveTools(); + const activeToolNames = activeTools.map((t) => t.name); + + expect(activeToolNames).toContain('mock-tool'); + }); + + it('should exclude tools in Plan Mode when denied by policy', () => { + // Setup Plan Mode + vi.spyOn(config, 'getApprovalMode').mockReturnValue(ApprovalMode.PLAN); + + const mockTool = new MockTool({ name: 'mock-tool' }); + toolRegistry.registerTool(mockTool); + + // Mock PolicyEngine to explicitly deny this tool + const policyEngine = config.getPolicyEngine(); + vi.spyOn(policyEngine, 'isToolPotentiallyAllowed').mockImplementation( + (toolName) => toolName !== 'mock-tool', + ); + + const activeTools = toolRegistry.getActiveTools(); + const activeToolNames = activeTools.map((t) => t.name); + + expect(activeToolNames).not.toContain('mock-tool'); + }); + }); + describe('getAllToolNames', () => { it('should return all registered tool names', () => { // Register tools with displayNames in non-alphabetical order diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index 95bac200be..905f1b7d27 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -26,7 +26,6 @@ import { DISCOVERED_TOOL_PREFIX, TOOL_LEGACY_ALIASES, getToolAliases, - PLAN_MODE_TOOLS, WRITE_FILE_TOOL_NAME, EDIT_TOOL_NAME, } from './tool-names.js'; @@ -442,9 +441,12 @@ export class ToolRegistry { } /** + * Retrieves the list of currently active tools. + * This is filtered by the current policies and approval modes (e.g. Plan Mode). + * It is public so that PromptProvider can dynamically construct accurate tool lists. * @returns All the tools that are not excluded. */ - private getActiveTools(): AnyDeclarativeTool[] { + getActiveTools(): AnyDeclarativeTool[] { const excludedTools = this.expandExcludeToolsWithAliases(this.config.getExcludeTools()) ?? new Set([]); @@ -490,28 +492,28 @@ export class ToolRegistry { this.expandExcludeToolsWithAliases(this.config.getExcludeTools()) ?? new Set([]); - // Filter tools in Plan Mode to only allow approved read-only tools. + let serverName: string | undefined; + const isMcpTool = tool instanceof DiscoveredMCPTool; + if (isMcpTool) { + serverName = tool.getFullyQualifiedPrefix().slice(0, -2); + } + const isPlanMode = typeof this.config.getApprovalMode === 'function' && this.config.getApprovalMode() === ApprovalMode.PLAN; - if (isPlanMode) { - const allowedToolNames = new Set(PLAN_MODE_TOOLS); - // We allow write_file and replace for writing plans specifically. - allowedToolNames.add(WRITE_FILE_TOOL_NAME); - allowedToolNames.add(EDIT_TOOL_NAME); - // Discovered MCP tools are allowed if they are read-only. - if ( - tool instanceof DiscoveredMCPTool && - tool.isReadOnly && - !allowedToolNames.has(tool.name) - ) { - allowedToolNames.add(tool.name); - } + const isReadOnlyMcp = isMcpTool && tool.isReadOnly; - if (!allowedToolNames.has(tool.name)) { - return false; - } + // Read-only MCP tools bypass the Plan Mode global deny policy, + // but still respect specific tool or wildcard rules. + const ignoreGlobalRules = isPlanMode && isReadOnlyMcp; + + if ( + !this.config + .getPolicyEngine() + .isToolPotentiallyAllowed(tool.name, serverName, ignoreGlobalRules) + ) { + return false; } const normalizedClassName = tool.constructor.name.replace(/^_+/, '');