Add support for policy engine in extensions (#20049)

Co-authored-by: Jerop Kipruto <jerop@google.com>

packages/cli/src/config/extension-manager.ts

@@ -52,6 +52,10 @@ import {
   applyAdminAllowlist,
   getAdminBlockedMcpServersMessage,
   CoreToolCallStatus,
+  loadExtensionPolicies,
+  isSubpath,
+  type PolicyRule,
+  type SafetyCheckerRule,
   HookType,
 } from '@google/gemini-cli-core';
 import { maybeRequestConsentOrFail } from './extensions/consent.js';
@@ -764,9 +768,18 @@ Would you like to attempt to install via "git clone" instead?`,
       }
 
       const contextFiles = getContextFileNames(config)
-        .map((contextFileName) =>
-          path.join(effectiveExtensionPath, contextFileName),
-        )
+        .map((contextFileName) => {
+          const contextFilePath = path.join(
+            effectiveExtensionPath,
+            contextFileName,
+          );
+          if (!isSubpath(effectiveExtensionPath, contextFilePath)) {
+            throw new Error(
+              `Invalid context file path: "${contextFileName}". Context files must be within the extension directory.`,
+            );
+          }
+          return contextFilePath;
+        })
         .filter((contextFilePath) => fs.existsSync(contextFilePath));
 
       const hydrationContext: VariableContext = {
@@ -820,6 +833,24 @@ Would you like to attempt to install via "git clone" instead?`,
         recursivelyHydrateStrings(skill, hydrationContext),
       );
 
+      let rules: PolicyRule[] | undefined;
+      let checkers: SafetyCheckerRule[] | undefined;
+
+      const policyDir = path.join(effectiveExtensionPath, 'policies');
+      if (fs.existsSync(policyDir)) {
+        const result = await loadExtensionPolicies(config.name, policyDir);
+        rules = result.rules;
+        checkers = result.checkers;
+
+        if (result.errors.length > 0) {
+          for (const error of result.errors) {
+            debugLogger.warn(
+              `[ExtensionManager] Error loading policies from ${config.name}: ${error.message}${error.details ? `\nDetails: ${error.details}` : ''}`,
+            );
+          }
+        }
+      }
+
       const agentLoadResult = await loadAgentsFromDirectory(
         path.join(effectiveExtensionPath, 'agents'),
       );
@@ -853,6 +884,8 @@ Would you like to attempt to install via "git clone" instead?`,
         skills,
         agents: agentLoadResult.agents,
         themes: config.themes,
+        rules,
+        checkers,
       };
     } catch (e) {
       debugLogger.error(

packages/cli/src/config/extension.test.ts

@@ -239,6 +239,27 @@ describe('extension tests', () => {
       expect(extensions[0].name).toBe('test-extension');
     });
 
+    it('should throw an error if a context file path is outside the extension directory', async () => {
+      const consoleSpy = vi
+        .spyOn(console, 'error')
+        .mockImplementation(() => {});
+      createExtension({
+        extensionsDir: userExtensionsDir,
+        name: 'traversal-extension',
+        version: '1.0.0',
+        contextFileName: '../secret.txt',
+      });
+
+      const extensions = await extensionManager.loadExtensions();
+      expect(extensions).toHaveLength(0);
+      expect(consoleSpy).toHaveBeenCalledWith(
+        expect.stringContaining(
+          'traversal-extension: Invalid context file path: "../secret.txt"',
+        ),
+      );
+      consoleSpy.mockRestore();
+    });
+
     it('should load context file path when GEMINI.md is present', async () => {
       createExtension({
         extensionsDir: userExtensionsDir,
@@ -363,6 +384,111 @@ describe('extension tests', () => {
       ]);
     });
 
+    it('should load extension policies from the policies directory', async () => {
+      const extDir = createExtension({
+        extensionsDir: userExtensionsDir,
+        name: 'policy-extension',
+        version: '1.0.0',
+      });
+
+      const policiesDir = path.join(extDir, 'policies');
+      fs.mkdirSync(policiesDir);
+
+      const policiesContent = `
+[[rule]]
+toolName = "deny_tool"
+decision = "deny"
+priority = 500
+
+[[rule]]
+toolName = "ask_tool"
+decision = "ask_user"
+priority = 100
+`;
+      fs.writeFileSync(
+        path.join(policiesDir, 'policies.toml'),
+        policiesContent,
+      );
+
+      const extensions = await extensionManager.loadExtensions();
+      expect(extensions).toHaveLength(1);
+      const extension = extensions[0];
+
+      expect(extension.rules).toBeDefined();
+      expect(extension.rules).toHaveLength(2);
+      expect(
+        extension.rules!.find((r) => r.toolName === 'deny_tool')?.decision,
+      ).toBe('deny');
+      expect(
+        extension.rules!.find((r) => r.toolName === 'ask_tool')?.decision,
+      ).toBe('ask_user');
+      // Verify source is prefixed
+      expect(extension.rules![0].source).toContain(
+        'Extension (policy-extension):',
+      );
+    });
+
+    it('should ignore ALLOW rules and YOLO mode from extension policies for security', async () => {
+      const consoleSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+      const extDir = createExtension({
+        extensionsDir: userExtensionsDir,
+        name: 'security-test-extension',
+        version: '1.0.0',
+      });
+
+      const policiesDir = path.join(extDir, 'policies');
+      fs.mkdirSync(policiesDir);
+
+      const policiesContent = `
+[[rule]]
+toolName = "allow_tool"
+decision = "allow"
+priority = 100
+
+[[rule]]
+toolName = "yolo_tool"
+decision = "ask_user"
+priority = 100
+modes = ["yolo"]
+
+[[safety_checker]]
+toolName = "yolo_check"
+priority = 100
+modes = ["yolo"]
+[safety_checker.checker]
+type = "external"
+name = "yolo-checker"
+`;
+      fs.writeFileSync(
+        path.join(policiesDir, 'policies.toml'),
+        policiesContent,
+      );
+
+      const extensions = await extensionManager.loadExtensions();
+      expect(extensions).toHaveLength(1);
+      const extension = extensions[0];
+
+      // ALLOW rules and YOLO rules/checkers should be filtered out
+      expect(extension.rules).toBeDefined();
+      expect(extension.rules).toHaveLength(0);
+      expect(extension.checkers).toBeDefined();
+      expect(extension.checkers).toHaveLength(0);
+
+      // Should have logged warnings
+      expect(consoleSpy).toHaveBeenCalledWith(
+        expect.stringContaining('attempted to contribute an ALLOW rule'),
+      );
+      expect(consoleSpy).toHaveBeenCalledWith(
+        expect.stringContaining('attempted to contribute a rule for YOLO mode'),
+      );
+      expect(consoleSpy).toHaveBeenCalledWith(
+        expect.stringContaining(
+          'attempted to contribute a safety checker for YOLO mode',
+        ),
+      );
+      consoleSpy.mockRestore();
+    });
+
     it('should hydrate ${extensionPath} correctly for linked extensions', async () => {
       const sourceExtDir = getRealPath(
         createExtension({
@@ -540,7 +666,7 @@ describe('extension tests', () => {
 
       // Bad extension
       const badExtDir = path.join(userExtensionsDir, 'bad-ext');
-      fs.mkdirSync(badExtDir);
+      fs.mkdirSync(badExtDir, { recursive: true });
       const badConfigPath = path.join(badExtDir, EXTENSIONS_CONFIG_FILENAME);
       fs.writeFileSync(badConfigPath, '{ "name": "bad-ext"'); // Malformed
 
@@ -548,7 +674,7 @@ describe('extension tests', () => {
 
       expect(extensions).toHaveLength(1);
       expect(extensions[0].name).toBe('good-ext');
-      expect(consoleSpy).toHaveBeenCalledExactlyOnceWith(
+      expect(consoleSpy).toHaveBeenCalledWith(
         expect.stringContaining(
           `Warning: Skipping extension in ${badExtDir}: Failed to load extension config from ${badConfigPath}`,
         ),
@@ -571,7 +697,7 @@ describe('extension tests', () => {
 
       // Bad extension
       const badExtDir = path.join(userExtensionsDir, 'bad-ext-no-name');
-      fs.mkdirSync(badExtDir);
+      fs.mkdirSync(badExtDir, { recursive: true });
       const badConfigPath = path.join(badExtDir, EXTENSIONS_CONFIG_FILENAME);
       fs.writeFileSync(badConfigPath, JSON.stringify({ version: '1.0.0' }));
 
@@ -579,7 +705,7 @@ describe('extension tests', () => {
 
       expect(extensions).toHaveLength(1);
       expect(extensions[0].name).toBe('good-ext');
-      expect(consoleSpy).toHaveBeenCalledExactlyOnceWith(
+      expect(consoleSpy).toHaveBeenCalledWith(
         expect.stringContaining(
           `Warning: Skipping extension in ${badExtDir}: Failed to load extension config from ${badConfigPath}: Invalid configuration in ${badConfigPath}: missing "name"`,
         ),

packages/cli/src/config/policy-engine.integration.test.ts

@@ -177,13 +177,13 @@ describe('Policy Engine Integration Tests', () => {
       );
       const engine = new PolicyEngine(config);
 
-      // MCP server allowed (priority 3.1) provides general allow for server
-      // MCP server allowed (priority 3.1) provides general allow for server
+      // MCP server allowed (priority 4.1) provides general allow for server
+      // MCP server allowed (priority 4.1) provides general allow for server
       expect(
         (await engine.check({ name: 'my-server__safe-tool' }, undefined))
           .decision,
       ).toBe(PolicyDecision.ALLOW);
-      // But specific tool exclude (priority 3.4) wins over server allow
+      // But specific tool exclude (priority 4.4) wins over server allow
       expect(
         (await engine.check({ name: 'my-server__dangerous-tool' }, undefined))
           .decision,
@@ -476,25 +476,25 @@ describe('Policy Engine Integration Tests', () => {
 
       // Find rules and verify their priorities
       const blockedToolRule = rules.find((r) => r.toolName === 'blocked-tool');
-      expect(blockedToolRule?.priority).toBe(3.4); // Command line exclude
+      expect(blockedToolRule?.priority).toBe(4.4); // Command line exclude
 
       const blockedServerRule = rules.find(
         (r) => r.toolName === 'blocked-server__*',
       );
-      expect(blockedServerRule?.priority).toBe(3.9); // MCP server exclude
+      expect(blockedServerRule?.priority).toBe(4.9); // MCP server exclude
 
       const specificToolRule = rules.find(
         (r) => r.toolName === 'specific-tool',
       );
-      expect(specificToolRule?.priority).toBe(3.3); // Command line allow
+      expect(specificToolRule?.priority).toBe(4.3); // Command line allow
 
       const trustedServerRule = rules.find(
         (r) => r.toolName === 'trusted-server__*',
       );
-      expect(trustedServerRule?.priority).toBe(3.2); // MCP trusted server
+      expect(trustedServerRule?.priority).toBe(4.2); // MCP trusted server
 
       const mcpServerRule = rules.find((r) => r.toolName === 'mcp-server__*');
-      expect(mcpServerRule?.priority).toBe(3.1); // MCP allowed server
+      expect(mcpServerRule?.priority).toBe(4.1); // MCP allowed server
 
       const readOnlyToolRule = rules.find((r) => r.toolName === 'glob');
       // Priority 70 in default tier → 1.07 (Overriding Plan Mode Deny)
@@ -641,16 +641,16 @@ describe('Policy Engine Integration Tests', () => {
 
       // Verify each rule has the expected priority
       const tool3Rule = rules.find((r) => r.toolName === 'tool3');
-      expect(tool3Rule?.priority).toBe(3.4); // Excluded tools (user tier)
+      expect(tool3Rule?.priority).toBe(4.4); // Excluded tools (user tier)
 
       const server2Rule = rules.find((r) => r.toolName === 'server2__*');
-      expect(server2Rule?.priority).toBe(3.9); // Excluded servers (user tier)
+      expect(server2Rule?.priority).toBe(4.9); // Excluded servers (user tier)
 
       const tool1Rule = rules.find((r) => r.toolName === 'tool1');
-      expect(tool1Rule?.priority).toBe(3.3); // Allowed tools (user tier)
+      expect(tool1Rule?.priority).toBe(4.3); // Allowed tools (user tier)
 
       const server1Rule = rules.find((r) => r.toolName === 'server1__*');
-      expect(server1Rule?.priority).toBe(3.1); // Allowed servers (user tier)
+      expect(server1Rule?.priority).toBe(4.1); // Allowed servers (user tier)
 
       const globRule = rules.find((r) => r.toolName === 'glob');
       // Priority 70 in default tier → 1.07