feat(core): implement context-aware persistent policy approvals (#23257)

packages/core/src/policy/config.ts

@@ -533,7 +533,6 @@ export async function createPolicyEngineConfig(
     disableAlwaysAllow: settings.disableAlwaysAllow,
   };
 }
-
 interface TomlRule {
   toolName?: string;
   mcpName?: string;
@@ -542,10 +541,64 @@ interface TomlRule {
   commandPrefix?: string | string[];
   argsPattern?: string;
   allowRedirection?: boolean;
+  modes?: ApprovalMode[];
   // Index signature to satisfy Record type if needed for toml.stringify
   [key: string]: unknown;
 }
 
+/**
+ * Finds a rule in the rule array that matches the given criteria.
+ */
+function findMatchingRule(
+  rules: TomlRule[],
+  criteria: {
+    toolName: string;
+    mcpName?: string;
+    commandPrefix?: string | string[];
+    argsPattern?: string;
+  },
+): TomlRule | undefined {
+  return rules.find(
+    (r) =>
+      r.toolName === criteria.toolName &&
+      r.mcpName === criteria.mcpName &&
+      JSON.stringify(r.commandPrefix) ===
+        JSON.stringify(criteria.commandPrefix) &&
+      r.argsPattern === criteria.argsPattern,
+  );
+}
+
+/**
+ * Creates a new TOML rule object from the given tool name and message.
+ */
+function createTomlRule(toolName: string, message: UpdatePolicy): TomlRule {
+  const rule: TomlRule = {
+    decision: 'allow',
+    priority: getAlwaysAllowPriorityFraction(),
+    toolName,
+  };
+
+  if (message.mcpName) {
+    rule.mcpName = message.mcpName;
+  }
+
+  if (message.commandPrefix) {
+    rule.commandPrefix = message.commandPrefix;
+  } else if (message.argsPattern) {
+    rule.argsPattern = message.argsPattern;
+  }
+
+  if (message.allowRedirection !== undefined) {
+    rule.allowRedirection = message.allowRedirection;
+  }
+
+  if (message.modes) {
+    rule.modes = message.modes;
+  }
+
+  return rule;
+}
+
 export function createPolicyUpdater(
   policyEngine: PolicyEngine,
   messageBus: MessageBus,
@@ -585,6 +638,7 @@ export function createPolicyUpdater(
               priority,
               argsPattern: new RegExp(pattern),
               mcpName: message.mcpName,
+              modes: message.modes,
               source: 'Dynamic (Confirmed)',
               allowRedirection: message.allowRedirection,
             });
@@ -622,6 +676,7 @@ export function createPolicyUpdater(
           priority,
           argsPattern,
           mcpName: message.mcpName,
+          modes: message.modes,
           source: 'Dynamic (Confirmed)',
           allowRedirection: message.allowRedirection,
         });
@@ -662,39 +717,36 @@ export function createPolicyUpdater(
               existingData.rule = [];
             }
 
-            // Create new rule object
-            const newRule: TomlRule = {
-              decision: 'allow',
-              priority: getAlwaysAllowPriorityFraction(),
-            };
-
+            // Normalize tool name for MCP
+            let normalizedToolName = toolName;
             if (message.mcpName) {
-              newRule.mcpName = message.mcpName;
-
               const expectedPrefix = `${MCP_TOOL_PREFIX}${message.mcpName}_`;
               if (toolName.startsWith(expectedPrefix)) {
-                newRule.toolName = toolName.slice(expectedPrefix.length);
-              } else {
-                newRule.toolName = toolName;
+                normalizedToolName = toolName.slice(expectedPrefix.length);
+              }
+            }
+
+            // Look for an existing rule to update
+            const existingRule = findMatchingRule(existingData.rule, {
+              toolName: normalizedToolName,
+              mcpName: message.mcpName,
+              commandPrefix: message.commandPrefix,
+              argsPattern: message.argsPattern,
+            });
+
+            if (existingRule) {
+              if (message.allowRedirection !== undefined) {
+                existingRule.allowRedirection = message.allowRedirection;
+              }
+              if (message.modes) {
+                existingRule.modes = message.modes;
               }
             } else {
-              newRule.toolName = toolName;
+              existingData.rule.push(
+                createTomlRule(normalizedToolName, message),
+              );
             }
 
-            if (message.commandPrefix) {
-              newRule.commandPrefix = message.commandPrefix;
-            } else if (message.argsPattern) {
-              // message.argsPattern was already validated above
-              newRule.argsPattern = message.argsPattern;
-            }
-
-            if (message.allowRedirection !== undefined) {
-              newRule.allowRedirection = message.allowRedirection;
-            }
-
-            // Add to rules
-            existingData.rule.push(newRule);
-
             // Serialize back to TOML
             // @iarna/toml stringify might not produce beautiful output but it handles escaping correctly
             // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion

packages/core/src/policy/persistence.test.ts

@@ -242,4 +242,57 @@ decision = "deny"
     const content = memfs.readFileSync(policyFile, 'utf-8') as string;
     expect(content).toContain('toolName = "test_tool"');
   });
+
+  it('should include modes if provided', async () => {
+    createPolicyUpdater(policyEngine, messageBus, mockStorage);
+
+    const policyFile = '/mock/user/.gemini/policies/auto-saved.toml';
+    vi.spyOn(mockStorage, 'getAutoSavedPolicyPath').mockReturnValue(policyFile);
+
+    await messageBus.publish({
+      type: MessageBusType.UPDATE_POLICY,
+      toolName: 'test_tool',
+      persist: true,
+      modes: [ApprovalMode.DEFAULT, ApprovalMode.YOLO],
+    });
+
+    await vi.advanceTimersByTimeAsync(100);
+
+    const content = memfs.readFileSync(policyFile, 'utf-8') as string;
+    expect(content).toContain('modes = [ "default", "yolo" ]');
+  });
+
+  it('should update existing rule modes instead of appending redundant rule', async () => {
+    createPolicyUpdater(policyEngine, messageBus, mockStorage);
+
+    const policyFile = '/mock/user/.gemini/policies/auto-saved.toml';
+    vi.spyOn(mockStorage, 'getAutoSavedPolicyPath').mockReturnValue(policyFile);
+
+    const existingContent = `
+[[rule]]
+decision = "allow"
+priority = 950
+toolName = "test_tool"
+modes = [ "autoEdit", "yolo" ]
+`;
+    const dir = path.dirname(policyFile);
+    memfs.mkdirSync(dir, { recursive: true });
+    memfs.writeFileSync(policyFile, existingContent);
+
+    // Now grant in DEFAULT mode, which should include [default, autoEdit, yolo]
+    await messageBus.publish({
+      type: MessageBusType.UPDATE_POLICY,
+      toolName: 'test_tool',
+      persist: true,
+      modes: [ApprovalMode.DEFAULT, ApprovalMode.AUTO_EDIT, ApprovalMode.YOLO],
+    });
+
+    await vi.advanceTimersByTimeAsync(100);
+
+    const content = memfs.readFileSync(policyFile, 'utf-8') as string;
+    // Should NOT have two [[rule]] entries for test_tool
+    const ruleCount = (content.match(/\[\[rule\]\]/g) || []).length;
+    expect(ruleCount).toBe(1);
+    expect(content).toContain('modes = [ "default", "autoEdit", "yolo" ]');
+  });
 });

packages/core/src/policy/types.ts

@@ -52,6 +52,18 @@ export enum ApprovalMode {
   PLAN = 'plan',
 }
 
+/**
+ * The order of permissiveness for approval modes.
+ * Tools allowed in a less permissive mode should also be allowed
+ * in more permissive modes.
+ */
+export const MODES_BY_PERMISSIVENESS = [
+  ApprovalMode.PLAN,
+  ApprovalMode.DEFAULT,
+  ApprovalMode.AUTO_EDIT,
+  ApprovalMode.YOLO,
+];
+
 /**
  * Configuration for the built-in allowed-path checker.
  */