feat(policy): centralize plan mode tool visibility in policy engine (#20178)

Co-authored-by: Mahima Shanware <mshanware@google.com>

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.

packages/core/src/tools/tool-registry.test.ts

@@ -659,6 +659,76 @@ describe('ToolRegistry', () => {
     });
   });
 
+  describe('plan mode', () => {
+    it('should only return policy-allowed tools in plan mode', () => {
+      // Register several tools
+      const globTool = new MockTool({ name: 'glob', displayName: 'Glob' });
+      const readFileTool = new MockTool({
+        name: 'read_file',
+        displayName: 'ReadFile',
+      });
+      const shellTool = new MockTool({ name: 'shell', displayName: 'Shell' });
+      const writeTool = new MockTool({
+        name: 'write_file',
+        displayName: 'WriteFile',
+      });
+
+      toolRegistry.registerTool(globTool);
+      toolRegistry.registerTool(readFileTool);
+      toolRegistry.registerTool(shellTool);
+      toolRegistry.registerTool(writeTool);
+
+      // Mock config in PLAN mode: exclude shell and write_file
+      mockConfigGetExcludedTools.mockReturnValue(
+        new Set(['shell', 'write_file']),
+      );
+
+      const allTools = toolRegistry.getAllTools();
+      const toolNames = allTools.map((t) => t.name);
+
+      expect(toolNames).toContain('glob');
+      expect(toolNames).toContain('read_file');
+      expect(toolNames).not.toContain('shell');
+      expect(toolNames).not.toContain('write_file');
+    });
+
+    it('should include read-only MCP tools when allowed by policy in plan mode', () => {
+      const readOnlyMcp = createMCPTool(
+        'test-server',
+        'read-only-tool',
+        'A read-only MCP tool',
+      );
+      // Set readOnlyHint to true via toolAnnotations
+      Object.defineProperty(readOnlyMcp, 'isReadOnly', { value: true });
+
+      toolRegistry.registerTool(readOnlyMcp);
+
+      // Policy allows this tool (not in excluded set)
+      mockConfigGetExcludedTools.mockReturnValue(new Set());
+
+      const allTools = toolRegistry.getAllTools();
+      const toolNames = allTools.map((t) => t.name);
+      expect(toolNames).toContain('read-only-tool');
+    });
+
+    it('should exclude non-read-only MCP tools when denied by policy in plan mode', () => {
+      const writeMcp = createMCPTool(
+        'test-server',
+        'write-mcp-tool',
+        'A write MCP tool',
+      );
+
+      toolRegistry.registerTool(writeMcp);
+
+      // Policy excludes this tool
+      mockConfigGetExcludedTools.mockReturnValue(new Set(['write-mcp-tool']));
+
+      const allTools = toolRegistry.getAllTools();
+      const toolNames = allTools.map((t) => t.name);
+      expect(toolNames).not.toContain('write-mcp-tool');
+    });
+  });
+
   describe('DiscoveredToolInvocation', () => {
     it('should return the stringified params from getDescription', () => {
       const tool = new DiscoveredTool(

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';
@@ -445,7 +444,13 @@ export class ToolRegistry {
     const toolMetadata = new Map<string, Record<string, unknown>>();
     for (const [name, tool] of this.allKnownTools) {
       if (tool.toolAnnotations) {
-        toolMetadata.set(name, tool.toolAnnotations);
+        const metadata: Record<string, unknown> = { ...tool.toolAnnotations };
+        // Include server name so the policy engine can resolve composite
+        // wildcard patterns (e.g. "*__*") against unqualified tool names.
+        if (tool instanceof DiscoveredMCPTool) {
+          metadata['_serverName'] = tool.serverName;
+        }
+        toolMetadata.set(name, metadata);
       }
     }
     return toolMetadata;
@@ -456,9 +461,10 @@ export class ToolRegistry {
    */
   private getActiveTools(): AnyDeclarativeTool[] {
     const toolMetadata = this.buildToolMetadata();
+    const allKnownNames = new Set(this.allKnownTools.keys());
     const excludedTools =
       this.expandExcludeToolsWithAliases(
-        this.config.getExcludeTools(toolMetadata),
+        this.config.getExcludeTools(toolMetadata, allKnownNames),
       ) ?? new Set([]);
     const activeTools: AnyDeclarativeTool[] = [];
     for (const tool of this.allKnownTools.values()) {
@@ -500,33 +506,12 @@ export class ToolRegistry {
   ): boolean {
     excludeTools ??=
       this.expandExcludeToolsWithAliases(
-        this.config.getExcludeTools(this.buildToolMetadata()),
+        this.config.getExcludeTools(
+          this.buildToolMetadata(),
+          new Set(this.allKnownTools.keys()),
+        ),
       ) ?? new Set([]);
 
-    // Filter tools in Plan Mode to only allow approved read-only tools.
-    const isPlanMode =
-      typeof this.config.getApprovalMode === 'function' &&
-      this.config.getApprovalMode() === ApprovalMode.PLAN;
-    if (isPlanMode) {
-      const allowedToolNames = new Set<string>(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);
-      }
-
-      if (!allowedToolNames.has(tool.name)) {
-        return false;
-      }
-    }
-
     const normalizedClassName = tool.constructor.name.replace(/^_+/, '');
     const possibleNames = [tool.name, normalizedClassName];
     if (tool instanceof DiscoveredMCPTool) {