From a220874281e97f32220bd58e706e4d8aeaf766bf Mon Sep 17 00:00:00 2001 From: Spencer Date: Tue, 10 Mar 2026 13:01:41 -0400 Subject: [PATCH] feat(policy): support auto-add to policy by default and scoped persistence (#20361) --- docs/cli/settings.md | 1 + docs/reference/configuration.md | 5 + packages/cli/src/config/settingsSchema.ts | 12 + .../messages/ToolConfirmationMessage.test.tsx | 7 +- .../messages/ToolConfirmationMessage.tsx | 518 ++++++++++-------- .../ToolConfirmationMessage.test.tsx.snap | 16 + .../core/src/agents/browser/mcpToolWrapper.ts | 4 +- packages/core/src/config/config.ts | 63 ++- packages/core/src/config/storage.ts | 7 + packages/core/src/confirmation-bus/types.ts | 1 + packages/core/src/policy/config.ts | 80 ++- packages/core/src/policy/persistence.test.ts | 251 ++++----- .../core/src/policy/policy-updater.test.ts | 13 +- packages/core/src/policy/utils.test.ts | 33 +- packages/core/src/policy/utils.ts | 45 +- packages/core/src/scheduler/policy.test.ts | 100 +++- packages/core/src/scheduler/policy.ts | 43 +- packages/core/src/scheduler/scheduler.ts | 1 + packages/core/src/tools/edit.ts | 91 +-- packages/core/src/tools/glob.ts | 11 + packages/core/src/tools/grep.ts | 11 + packages/core/src/tools/ls.ts | 11 + packages/core/src/tools/mcp-tool.ts | 2 +- packages/core/src/tools/read-file.ts | 11 + packages/core/src/tools/read-many-files.ts | 13 + packages/core/src/tools/shell.ts | 2 +- packages/core/src/tools/tool-names.ts | 29 +- packages/core/src/tools/tools.ts | 11 +- packages/core/src/tools/web-fetch.ts | 18 + packages/core/src/tools/write-file.ts | 10 + schemas/settings.schema.json | 7 + 31 files changed, 929 insertions(+), 498 deletions(-) diff --git a/docs/cli/settings.md b/docs/cli/settings.md index 5565a5e1f6..0f4b44f159 100644 --- a/docs/cli/settings.md +++ b/docs/cli/settings.md @@ -125,6 +125,7 @@ they appear in the UI. | ------------------------------------- | ----------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------- | | Disable YOLO Mode | `security.disableYoloMode` | Disable YOLO mode, even if enabled by a flag. | `false` | | Allow Permanent Tool Approval | `security.enablePermanentToolApproval` | Enable the "Allow for all future sessions" option in tool confirmation dialogs. | `false` | +| Auto-add to Policy by Default | `security.autoAddToPolicyByDefault` | When enabled, the "Allow for all future sessions" option becomes the default choice for low-risk tools in trusted workspaces. | `false` | | Blocks extensions from Git | `security.blockGitExtensions` | Blocks installing and loading extensions from Git. | `false` | | Extension Source Regex Allowlist | `security.allowedExtensions` | List of Regex patterns for allowed extensions. If nonempty, only extensions that match the patterns in this list are allowed. Overrides the blockGitExtensions setting. | `[]` | | Folder Trust | `security.folderTrust.enabled` | Setting to track whether Folder trust is enabled. | `true` | diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index b1d1f7f021..39870262c9 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -872,6 +872,11 @@ their corresponding top-level category object in your `settings.json` file. confirmation dialogs. - **Default:** `false` +- **`security.autoAddToPolicyByDefault`** (boolean): + - **Description:** When enabled, the "Allow for all future sessions" option + becomes the default choice for low-risk tools in trusted workspaces. + - **Default:** `false` + - **`security.blockGitExtensions`** (boolean): - **Description:** Blocks installing and loading extensions from Git. - **Default:** `false` diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index bd1f9d82a4..0e96c88b24 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -1496,6 +1496,18 @@ const SETTINGS_SCHEMA = { 'Enable the "Allow for all future sessions" option in tool confirmation dialogs.', showInDialog: true, }, + autoAddToPolicyByDefault: { + type: 'boolean', + label: 'Auto-add to Policy by Default', + category: 'Security', + requiresRestart: false, + default: false, + description: oneLine` + When enabled, the "Allow for all future sessions" option becomes the + default choice for low-risk tools in trusted workspaces. + `, + showInDialog: true, + }, blockGitExtensions: { type: 'boolean', label: 'Blocks extensions from Git', diff --git a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.test.tsx b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.test.tsx index fec1228c63..ec623f69a4 100644 --- a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.test.tsx +++ b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.test.tsx @@ -411,7 +411,7 @@ describe('ToolConfirmationMessage', () => { unmount(); }); - it('should show "Allow for all future sessions" when setting is true', async () => { + it('should show "Allow for all future sessions" when trusted', async () => { const mockConfig = { isTrustedFolder: () => true, getIdeMode: () => false, @@ -434,7 +434,10 @@ describe('ToolConfirmationMessage', () => { ); await waitUntilReady(); - expect(lastFrame()).toContain('Allow for all future sessions'); + const output = lastFrame(); + expect(output).toContain('future sessions'); + // Verify it is the default selection (matching the indicator in the snapshot) + expect(output).toMatchSnapshot(); unmount(); }); }); diff --git a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx index 329d8e6262..113852cb8d 100644 --- a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx @@ -246,9 +246,9 @@ export const ToolConfirmationMessage: React.FC< }); if (allowPermanentApproval) { options.push({ - label: 'Allow for all future sessions', + label: 'Allow for this file in all future sessions', value: ToolConfirmationOutcome.ProceedAlwaysAndSave, - key: 'Allow for all future sessions', + key: 'Allow for this file in all future sessions', }); } } @@ -282,7 +282,7 @@ export const ToolConfirmationMessage: React.FC< }); if (allowPermanentApproval) { options.push({ - label: `Allow for all future sessions`, + label: `Allow this command for all future sessions`, value: ToolConfirmationOutcome.ProceedAlwaysAndSave, key: `Allow for all future sessions`, }); @@ -388,266 +388,301 @@ export const ToolConfirmationMessage: React.FC< return Math.max(availableTerminalHeight - surroundingElementsHeight, 1); }, [availableTerminalHeight, getOptions, handlesOwnUI]); - const { question, bodyContent, options, securityWarnings } = useMemo<{ - question: string; - bodyContent: React.ReactNode; - options: Array>; - securityWarnings: React.ReactNode; - }>(() => { - let bodyContent: React.ReactNode | null = null; - let securityWarnings: React.ReactNode | null = null; - let question = ''; - const options = getOptions(); + const { question, bodyContent, options, securityWarnings, initialIndex } = + useMemo<{ + question: string; + bodyContent: React.ReactNode; + options: Array>; + securityWarnings: React.ReactNode; + initialIndex: number; + }>(() => { + let bodyContent: React.ReactNode | null = null; + let securityWarnings: React.ReactNode | null = null; + let question = ''; + const options = getOptions(); - if (deceptiveUrlWarningText) { - securityWarnings = ; - } - - if (confirmationDetails.type === 'ask_user') { - bodyContent = ( - { - handleConfirm(ToolConfirmationOutcome.ProceedOnce, { answers }); - }} - onCancel={() => { - handleConfirm(ToolConfirmationOutcome.Cancel); - }} - width={terminalWidth} - availableHeight={availableBodyContentHeight()} - /> - ); - return { - question: '', - bodyContent, - options: [], - securityWarnings: null, - }; - } - - if (confirmationDetails.type === 'exit_plan_mode') { - bodyContent = ( - { - handleConfirm(ToolConfirmationOutcome.ProceedOnce, { - approved: true, - approvalMode, - }); - }} - onFeedback={(feedback) => { - handleConfirm(ToolConfirmationOutcome.ProceedOnce, { - approved: false, - feedback, - }); - }} - onCancel={() => { - handleConfirm(ToolConfirmationOutcome.Cancel); - }} - width={terminalWidth} - availableHeight={availableBodyContentHeight()} - /> - ); - return { question: '', bodyContent, options: [], securityWarnings: null }; - } - - if (confirmationDetails.type === 'edit') { - if (!confirmationDetails.isModifying) { - question = `Apply this change?`; + let initialIndex = 0; + if (isTrustedFolder && allowPermanentApproval) { + // It is safe to allow permanent approval for info, edit, and mcp tools + // in trusted folders because the generated policy rules are narrowed + // to specific files, patterns, or tools (rather than allowing all access). + const isSafeToPersist = + confirmationDetails.type === 'info' || + confirmationDetails.type === 'edit' || + confirmationDetails.type === 'mcp'; + if ( + isSafeToPersist && + settings.merged.security.autoAddToPolicyByDefault + ) { + const alwaysAndSaveIndex = options.findIndex( + (o) => o.value === ToolConfirmationOutcome.ProceedAlwaysAndSave, + ); + if (alwaysAndSaveIndex !== -1) { + initialIndex = alwaysAndSaveIndex; + } + } } - } else if (confirmationDetails.type === 'exec') { - const executionProps = confirmationDetails; - if (executionProps.commands && executionProps.commands.length > 1) { - question = `Allow execution of ${executionProps.commands.length} commands?`; - } else { - question = `Allow execution of: '${sanitizeForDisplay(executionProps.rootCommand)}'?`; + if (deceptiveUrlWarningText) { + securityWarnings = ; } - } else if (confirmationDetails.type === 'info') { - question = `Do you want to proceed?`; - } else if (confirmationDetails.type === 'mcp') { - // mcp tool confirmation - const mcpProps = confirmationDetails; - question = `Allow execution of MCP tool "${sanitizeForDisplay(mcpProps.toolName)}" from server "${sanitizeForDisplay(mcpProps.serverName)}"?`; - } - if (confirmationDetails.type === 'edit') { - if (!confirmationDetails.isModifying) { + if (confirmationDetails.type === 'ask_user') { bodyContent = ( - { + handleConfirm(ToolConfirmationOutcome.ProceedOnce, { answers }); + }} + onCancel={() => { + handleConfirm(ToolConfirmationOutcome.Cancel); + }} + width={terminalWidth} + availableHeight={availableBodyContentHeight()} /> ); - } - } else if (confirmationDetails.type === 'exec') { - const executionProps = confirmationDetails; - - const commandsToDisplay = - executionProps.commands && executionProps.commands.length > 1 - ? executionProps.commands - : [executionProps.command]; - const containsRedirection = commandsToDisplay.some((cmd) => - hasRedirection(cmd), - ); - - let bodyContentHeight = availableBodyContentHeight(); - let warnings: React.ReactNode = null; - - if (bodyContentHeight !== undefined) { - bodyContentHeight -= 2; // Account for padding; + return { + question: '', + bodyContent, + options: [], + securityWarnings: null, + initialIndex: 0, + }; } - if (containsRedirection) { - // Calculate lines needed for Note and Tip - const safeWidth = Math.max(terminalWidth, 1); - const tipText = `Toggle auto-edit (${formatCommand(Command.CYCLE_APPROVAL_MODE)}) to allow redirection in the future.`; + if (confirmationDetails.type === 'exit_plan_mode') { + bodyContent = ( + { + handleConfirm(ToolConfirmationOutcome.ProceedOnce, { + approved: true, + approvalMode, + }); + }} + onFeedback={(feedback) => { + handleConfirm(ToolConfirmationOutcome.ProceedOnce, { + approved: false, + feedback, + }); + }} + onCancel={() => { + handleConfirm(ToolConfirmationOutcome.Cancel); + }} + width={terminalWidth} + availableHeight={availableBodyContentHeight()} + /> + ); + return { + question: '', + bodyContent, + options: [], + securityWarnings: null, + initialIndex: 0, + }; + } - const noteLength = - REDIRECTION_WARNING_NOTE_LABEL.length + - REDIRECTION_WARNING_NOTE_TEXT.length; - const tipLength = REDIRECTION_WARNING_TIP_LABEL.length + tipText.length; + if (confirmationDetails.type === 'edit') { + if (!confirmationDetails.isModifying) { + question = `Apply this change?`; + } + } else if (confirmationDetails.type === 'exec') { + const executionProps = confirmationDetails; - const noteLines = Math.ceil(noteLength / safeWidth); - const tipLines = Math.ceil(tipLength / safeWidth); - const spacerLines = 1; - const warningHeight = noteLines + tipLines + spacerLines; + if (executionProps.commands && executionProps.commands.length > 1) { + question = `Allow execution of ${executionProps.commands.length} commands?`; + } else { + question = `Allow execution of: '${sanitizeForDisplay(executionProps.rootCommand)}'?`; + } + } else if (confirmationDetails.type === 'info') { + question = `Do you want to proceed?`; + } else if (confirmationDetails.type === 'mcp') { + // mcp tool confirmation + const mcpProps = confirmationDetails; + question = `Allow execution of MCP tool "${sanitizeForDisplay(mcpProps.toolName)}" from server "${sanitizeForDisplay(mcpProps.serverName)}"?`; + } + + if (confirmationDetails.type === 'edit') { + if (!confirmationDetails.isModifying) { + bodyContent = ( + + ); + } + } else if (confirmationDetails.type === 'exec') { + const executionProps = confirmationDetails; + + const commandsToDisplay = + executionProps.commands && executionProps.commands.length > 1 + ? executionProps.commands + : [executionProps.command]; + const containsRedirection = commandsToDisplay.some((cmd) => + hasRedirection(cmd), + ); + + let bodyContentHeight = availableBodyContentHeight(); + let warnings: React.ReactNode = null; if (bodyContentHeight !== undefined) { - bodyContentHeight = Math.max( - bodyContentHeight - warningHeight, - MINIMUM_MAX_HEIGHT, + bodyContentHeight -= 2; // Account for padding; + } + + if (containsRedirection) { + // Calculate lines needed for Note and Tip + const safeWidth = Math.max(terminalWidth, 1); + const noteLength = + REDIRECTION_WARNING_NOTE_LABEL.length + + REDIRECTION_WARNING_NOTE_TEXT.length; + const tipText = `Toggle auto-edit (${formatCommand(Command.CYCLE_APPROVAL_MODE)}) to allow redirection in the future.`; + const tipLength = + REDIRECTION_WARNING_TIP_LABEL.length + tipText.length; + + const noteLines = Math.ceil(noteLength / safeWidth); + const tipLines = Math.ceil(tipLength / safeWidth); + const spacerLines = 1; + const warningHeight = noteLines + tipLines + spacerLines; + + if (bodyContentHeight !== undefined) { + bodyContentHeight = Math.max( + bodyContentHeight - warningHeight, + MINIMUM_MAX_HEIGHT, + ); + } + + warnings = ( + <> + + + + {REDIRECTION_WARNING_NOTE_LABEL} + {REDIRECTION_WARNING_NOTE_TEXT} + + + + + {REDIRECTION_WARNING_TIP_LABEL} + {tipText} + + + ); } - warnings = ( - <> - - - - {REDIRECTION_WARNING_NOTE_LABEL} - {REDIRECTION_WARNING_NOTE_TEXT} + bodyContent = ( + + + + {commandsToDisplay.map((cmd, idx) => ( + + {colorizeCode({ + code: cmd, + language: 'bash', + maxWidth: Math.max(terminalWidth, 1), + settings, + hideLineNumbers: true, + })} + + ))} + + + {warnings} + + ); + } else if (confirmationDetails.type === 'info') { + const infoProps = confirmationDetails; + const displayUrls = + infoProps.urls && + !( + infoProps.urls.length === 1 && + infoProps.urls[0] === infoProps.prompt + ); + + bodyContent = ( + + + + + {displayUrls && infoProps.urls && infoProps.urls.length > 0 && ( + + URLs to fetch: + {infoProps.urls.map((urlString) => ( + + {' '} + - + + ))} + + )} + + ); + } else if (confirmationDetails.type === 'mcp') { + // mcp tool confirmation + const mcpProps = confirmationDetails; + + bodyContent = ( + + <> + + MCP Server: {sanitizeForDisplay(mcpProps.serverName)} - - - - {REDIRECTION_WARNING_TIP_LABEL} - {tipText} + + Tool: {sanitizeForDisplay(mcpProps.toolName)} - - + + {hasMcpToolDetails && ( + + MCP Tool Details: + {isMcpToolDetailsExpanded ? ( + <> + + (press {expandDetailsHintKey} to collapse MCP tool + details) + + {mcpToolDetailsText} + + ) : ( + + (press {expandDetailsHintKey} to expand MCP tool details) + + )} + + )} + ); } - bodyContent = ( - - - - {commandsToDisplay.map((cmd, idx) => ( - - {colorizeCode({ - code: cmd, - language: 'bash', - maxWidth: Math.max(terminalWidth, 1), - settings, - hideLineNumbers: true, - })} - - ))} - - - {warnings} - - ); - } else if (confirmationDetails.type === 'info') { - const infoProps = confirmationDetails; - const displayUrls = - infoProps.urls && - !( - infoProps.urls.length === 1 && infoProps.urls[0] === infoProps.prompt - ); - - bodyContent = ( - - - - - {displayUrls && infoProps.urls && infoProps.urls.length > 0 && ( - - URLs to fetch: - {infoProps.urls.map((urlString) => ( - - {' '} - - - - ))} - - )} - - ); - } else if (confirmationDetails.type === 'mcp') { - // mcp tool confirmation - const mcpProps = confirmationDetails; - - bodyContent = ( - - <> - - MCP Server: {sanitizeForDisplay(mcpProps.serverName)} - - - Tool: {sanitizeForDisplay(mcpProps.toolName)} - - - {hasMcpToolDetails && ( - - MCP Tool Details: - {isMcpToolDetailsExpanded ? ( - <> - - (press {expandDetailsHintKey} to collapse MCP tool details) - - {mcpToolDetailsText} - - ) : ( - - (press {expandDetailsHintKey} to expand MCP tool details) - - )} - - )} - - ); - } - - return { question, bodyContent, options, securityWarnings }; - }, [ - confirmationDetails, - getOptions, - availableBodyContentHeight, - terminalWidth, - handleConfirm, - deceptiveUrlWarningText, - isMcpToolDetailsExpanded, - hasMcpToolDetails, - mcpToolDetailsText, - expandDetailsHintKey, - getPreferredEditor, - settings, - ]); + return { question, bodyContent, options, securityWarnings, initialIndex }; + }, [ + confirmationDetails, + getOptions, + availableBodyContentHeight, + terminalWidth, + handleConfirm, + deceptiveUrlWarningText, + isMcpToolDetailsExpanded, + hasMcpToolDetails, + mcpToolDetailsText, + expandDetailsHintKey, + getPreferredEditor, + isTrustedFolder, + allowPermanentApproval, + settings, + ]); const bodyOverflowDirection: 'top' | 'bottom' = confirmationDetails.type === 'mcp' && isMcpToolDetailsExpanded @@ -710,6 +745,7 @@ export const ToolConfirmationMessage: React.FC< items={options} onSelect={handleSelect} isFocused={isFocused} + initialIndex={initialIndex} /> diff --git a/packages/cli/src/ui/components/messages/__snapshots__/ToolConfirmationMessage.test.tsx.snap b/packages/cli/src/ui/components/messages/__snapshots__/ToolConfirmationMessage.test.tsx.snap index 3f207df881..085d0bc445 100644 --- a/packages/cli/src/ui/components/messages/__snapshots__/ToolConfirmationMessage.test.tsx.snap +++ b/packages/cli/src/ui/components/messages/__snapshots__/ToolConfirmationMessage.test.tsx.snap @@ -1,5 +1,21 @@ // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html +exports[`ToolConfirmationMessage > enablePermanentToolApproval setting > should show "Allow for all future sessions" when trusted 1`] = ` +"╭──────────────────────────────────────────────────────────────────────────────╮ +│ │ +│ No changes detected. │ +│ │ +╰──────────────────────────────────────────────────────────────────────────────╯ +Apply this change? + +● 1. Allow once + 2. Allow for this session + 3. Allow for this file in all future sessions + 4. Modify with external editor + 5. No, suggest changes (esc) +" +`; + exports[`ToolConfirmationMessage > should display multiple commands for exec type when provided 1`] = ` "echo "hello" diff --git a/packages/core/src/agents/browser/mcpToolWrapper.ts b/packages/core/src/agents/browser/mcpToolWrapper.ts index 96b6aa9b68..f27d3462e6 100644 --- a/packages/core/src/agents/browser/mcpToolWrapper.ts +++ b/packages/core/src/agents/browser/mcpToolWrapper.ts @@ -70,7 +70,7 @@ class McpToolInvocation extends BaseToolInvocation< }; } - protected override getPolicyUpdateOptions( + override getPolicyUpdateOptions( _outcome: ToolConfirmationOutcome, ): PolicyUpdateOptions | undefined { return { @@ -177,7 +177,7 @@ class TypeTextInvocation extends BaseToolInvocation< }; } - protected override getPolicyUpdateOptions( + override getPolicyUpdateOptions( _outcome: ToolConfirmationOutcome, ): PolicyUpdateOptions | undefined { return { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 86cdf584b5..752ad25c4f 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -553,6 +553,7 @@ export interface ConfigParameters { truncateToolOutputThreshold?: number; eventEmitter?: EventEmitter; useWriteTodos?: boolean; + workspacePoliciesDir?: string; policyEngineConfig?: PolicyEngineConfig; directWebFetch?: boolean; policyUpdateConfirmationRequest?: PolicyUpdateConfirmationRequest; @@ -746,6 +747,7 @@ export class Config implements McpContext, AgentLoopContext { private readonly fileExclusions: FileExclusions; private readonly eventEmitter?: EventEmitter; private readonly useWriteTodos: boolean; + private readonly workspacePoliciesDir: string | undefined; private readonly _messageBus: MessageBus; private readonly policyEngine: PolicyEngine; private policyUpdateConfirmationRequest: @@ -956,6 +958,7 @@ export class Config implements McpContext, AgentLoopContext { this.useWriteTodos = isPreviewModel(this.model) ? false : (params.useWriteTodos ?? true); + this.workspacePoliciesDir = params.workspacePoliciesDir; this.enableHooksUI = params.enableHooksUI ?? true; this.enableHooks = params.enableHooks ?? true; this.disabledHooks = params.disabledHooks ?? []; @@ -1187,7 +1190,7 @@ export class Config implements McpContext, AgentLoopContext { if (this.getSkillManager().getSkills().length > 0) { this.getToolRegistry().unregisterTool(ActivateSkillTool.Name); this.getToolRegistry().registerTool( - new ActivateSkillTool(this, this._messageBus), + new ActivateSkillTool(this, this.messageBus), ); } } @@ -1999,6 +2002,10 @@ export class Config implements McpContext, AgentLoopContext { return this.geminiMdFilePaths; } + getWorkspacePoliciesDir(): string | undefined { + return this.workspacePoliciesDir; + } + setGeminiMdFilePaths(paths: string[]): void { this.geminiMdFilePaths = paths; } @@ -2621,7 +2628,7 @@ export class Config implements McpContext, AgentLoopContext { if (this.getSkillManager().getSkills().length > 0) { this.getToolRegistry().unregisterTool(ActivateSkillTool.Name); this.getToolRegistry().registerTool( - new ActivateSkillTool(this, this._messageBus), + new ActivateSkillTool(this, this.messageBus), ); } else { this.getToolRegistry().unregisterTool(ActivateSkillTool.Name); @@ -2805,7 +2812,7 @@ export class Config implements McpContext, AgentLoopContext { } async createToolRegistry(): Promise { - const registry = new ToolRegistry(this, this._messageBus); + const registry = new ToolRegistry(this, this.messageBus); // helper to create & register core tools that are enabled const maybeRegister = ( @@ -2835,10 +2842,10 @@ export class Config implements McpContext, AgentLoopContext { }; maybeRegister(LSTool, () => - registry.registerTool(new LSTool(this, this._messageBus)), + registry.registerTool(new LSTool(this, this.messageBus)), ); maybeRegister(ReadFileTool, () => - registry.registerTool(new ReadFileTool(this, this._messageBus)), + registry.registerTool(new ReadFileTool(this, this.messageBus)), ); if (this.getUseRipgrep()) { @@ -2851,85 +2858,81 @@ export class Config implements McpContext, AgentLoopContext { } if (useRipgrep) { maybeRegister(RipGrepTool, () => - registry.registerTool(new RipGrepTool(this, this._messageBus)), + registry.registerTool(new RipGrepTool(this, this.messageBus)), ); } else { logRipgrepFallback(this, new RipgrepFallbackEvent(errorString)); maybeRegister(GrepTool, () => - registry.registerTool(new GrepTool(this, this._messageBus)), + registry.registerTool(new GrepTool(this, this.messageBus)), ); } } else { maybeRegister(GrepTool, () => - registry.registerTool(new GrepTool(this, this._messageBus)), + registry.registerTool(new GrepTool(this, this.messageBus)), ); } maybeRegister(GlobTool, () => - registry.registerTool(new GlobTool(this, this._messageBus)), + registry.registerTool(new GlobTool(this, this.messageBus)), ); maybeRegister(ActivateSkillTool, () => - registry.registerTool(new ActivateSkillTool(this, this._messageBus)), + registry.registerTool(new ActivateSkillTool(this, this.messageBus)), ); maybeRegister(EditTool, () => - registry.registerTool(new EditTool(this, this._messageBus)), + registry.registerTool(new EditTool(this, this.messageBus)), ); maybeRegister(WriteFileTool, () => - registry.registerTool(new WriteFileTool(this, this._messageBus)), + registry.registerTool(new WriteFileTool(this, this.messageBus)), ); maybeRegister(WebFetchTool, () => - registry.registerTool(new WebFetchTool(this, this._messageBus)), + registry.registerTool(new WebFetchTool(this, this.messageBus)), ); maybeRegister(ShellTool, () => - registry.registerTool(new ShellTool(this, this._messageBus)), + registry.registerTool(new ShellTool(this, this.messageBus)), ); maybeRegister(MemoryTool, () => - registry.registerTool(new MemoryTool(this._messageBus)), + registry.registerTool(new MemoryTool(this.messageBus)), ); maybeRegister(WebSearchTool, () => - registry.registerTool(new WebSearchTool(this, this._messageBus)), + registry.registerTool(new WebSearchTool(this, this.messageBus)), ); maybeRegister(AskUserTool, () => - registry.registerTool(new AskUserTool(this._messageBus)), + registry.registerTool(new AskUserTool(this.messageBus)), ); if (this.getUseWriteTodos()) { maybeRegister(WriteTodosTool, () => - registry.registerTool(new WriteTodosTool(this._messageBus)), + registry.registerTool(new WriteTodosTool(this.messageBus)), ); } if (this.isPlanEnabled()) { maybeRegister(ExitPlanModeTool, () => - registry.registerTool(new ExitPlanModeTool(this, this._messageBus)), + registry.registerTool(new ExitPlanModeTool(this, this.messageBus)), ); maybeRegister(EnterPlanModeTool, () => - registry.registerTool(new EnterPlanModeTool(this, this._messageBus)), + registry.registerTool(new EnterPlanModeTool(this, this.messageBus)), ); } if (this.isTrackerEnabled()) { maybeRegister(TrackerCreateTaskTool, () => - registry.registerTool( - new TrackerCreateTaskTool(this, this._messageBus), - ), + registry.registerTool(new TrackerCreateTaskTool(this, this.messageBus)), ); maybeRegister(TrackerUpdateTaskTool, () => - registry.registerTool( - new TrackerUpdateTaskTool(this, this._messageBus), - ), + registry.registerTool(new TrackerUpdateTaskTool(this, this.messageBus)), ); maybeRegister(TrackerGetTaskTool, () => - registry.registerTool(new TrackerGetTaskTool(this, this._messageBus)), + registry.registerTool(new TrackerGetTaskTool(this, this.messageBus)), ); maybeRegister(TrackerListTasksTool, () => - registry.registerTool(new TrackerListTasksTool(this, this._messageBus)), + registry.registerTool(new TrackerListTasksTool(this, this.messageBus)), ); maybeRegister(TrackerAddDependencyTool, () => registry.registerTool( - new TrackerAddDependencyTool(this, this._messageBus), + new TrackerAddDependencyTool(this, this.messageBus), ), ); maybeRegister(TrackerVisualizeTool, () => - registry.registerTool(new TrackerVisualizeTool(this, this._messageBus)), + registry.registerTool(new TrackerVisualizeTool(this, this.messageBus)), ); } diff --git a/packages/core/src/config/storage.ts b/packages/core/src/config/storage.ts index 4c4ddaa2d9..b89c2bccbc 100644 --- a/packages/core/src/config/storage.ts +++ b/packages/core/src/config/storage.ts @@ -172,6 +172,13 @@ export class Storage { return path.join(this.getGeminiDir(), 'policies'); } + getWorkspaceAutoSavedPolicyPath(): string { + return path.join( + this.getWorkspacePoliciesDir(), + AUTO_SAVED_POLICY_FILENAME, + ); + } + getAutoSavedPolicyPath(): string { return path.join(Storage.getUserPoliciesDir(), AUTO_SAVED_POLICY_FILENAME); } diff --git a/packages/core/src/confirmation-bus/types.ts b/packages/core/src/confirmation-bus/types.ts index 277c821da3..99df9da616 100644 --- a/packages/core/src/confirmation-bus/types.ts +++ b/packages/core/src/confirmation-bus/types.ts @@ -122,6 +122,7 @@ export interface UpdatePolicy { type: MessageBusType.UPDATE_POLICY; toolName: string; persist?: boolean; + persistScope?: 'workspace' | 'user'; argsPattern?: string; commandPrefix?: string | string[]; mcpName?: string; diff --git a/packages/core/src/policy/config.ts b/packages/core/src/policy/config.ts index 8437cb9845..7085da7e3e 100644 --- a/packages/core/src/policy/config.ts +++ b/packages/core/src/policy/config.ts @@ -29,7 +29,7 @@ import { type MessageBus } from '../confirmation-bus/message-bus.js'; import { coreEvents } from '../utils/events.js'; import { debugLogger } from '../utils/debugLogger.js'; import { SHELL_TOOL_NAMES } from '../utils/shell-utils.js'; -import { SHELL_TOOL_NAME } from '../tools/tool-names.js'; +import { SHELL_TOOL_NAME, SENSITIVE_TOOLS } from '../tools/tool-names.js'; import { isNodeError } from '../utils/errors.js'; import { MCP_TOOL_PREFIX } from '../tools/mcp-tool.js'; @@ -46,13 +46,20 @@ export const WORKSPACE_POLICY_TIER = 3; export const USER_POLICY_TIER = 4; export const ADMIN_POLICY_TIER = 5; -// Specific priority offsets and derived priorities for dynamic/settings rules. -// These are added to the tier base (e.g., USER_POLICY_TIER). +/** + * The fractional priority of "Always allow" rules (e.g., 950/1000). + * Higher fraction within a tier wins. + */ +export const ALWAYS_ALLOW_PRIORITY_FRACTION = 950; -// Workspace tier (3) + high priority (950/1000) = ALWAYS_ALLOW_PRIORITY -// This ensures user "always allow" selections are high priority -// within the workspace tier but still lose to user/admin policies. -export const ALWAYS_ALLOW_PRIORITY = WORKSPACE_POLICY_TIER + 0.95; +/** + * The fractional priority offset for "Always allow" rules (e.g., 0.95). + * This ensures consistency between in-memory rules and persisted rules. + */ +export const ALWAYS_ALLOW_PRIORITY_OFFSET = + ALWAYS_ALLOW_PRIORITY_FRACTION / 1000; + +// Specific priority offsets and derived priorities for dynamic/settings rules. export const MCP_EXCLUDED_PRIORITY = USER_POLICY_TIER + 0.9; export const EXCLUDE_TOOLS_FLAG_PRIORITY = USER_POLICY_TIER + 0.4; @@ -60,6 +67,18 @@ export const ALLOWED_TOOLS_FLAG_PRIORITY = USER_POLICY_TIER + 0.3; export const TRUSTED_MCP_SERVER_PRIORITY = USER_POLICY_TIER + 0.2; export const ALLOWED_MCP_SERVER_PRIORITY = USER_POLICY_TIER + 0.1; +// These are added to the tier base (e.g., USER_POLICY_TIER). +// Workspace tier (3) + high priority (950/1000) = ALWAYS_ALLOW_PRIORITY +export const ALWAYS_ALLOW_PRIORITY = + WORKSPACE_POLICY_TIER + ALWAYS_ALLOW_PRIORITY_OFFSET; + +/** + * Returns the fractional priority of ALWAYS_ALLOW_PRIORITY scaled to 1000. + */ +export function getAlwaysAllowPriorityFraction(): number { + return Math.round((ALWAYS_ALLOW_PRIORITY % 1) * 1000); +} + /** * Gets the list of directories to search for policy files, in order of increasing priority * (Default -> Extension -> Workspace -> User -> Admin). @@ -492,6 +511,19 @@ export function createPolicyUpdater( if (message.commandPrefix) { // Convert commandPrefix(es) to argsPatterns for in-memory rules const patterns = buildArgsPatterns(undefined, message.commandPrefix); + const tier = + message.persistScope === 'user' + ? USER_POLICY_TIER + : WORKSPACE_POLICY_TIER; + const priority = tier + getAlwaysAllowPriorityFraction() / 1000; + + if (SENSITIVE_TOOLS.has(toolName) && !message.commandPrefix) { + debugLogger.warn( + `Attempted to update policy for sensitive tool '${toolName}' without a commandPrefix. Skipping.`, + ); + return; + } + for (const pattern of patterns) { if (pattern) { // Note: patterns from buildArgsPatterns are derived from escapeRegex, @@ -499,7 +531,7 @@ export function createPolicyUpdater( policyEngine.addRule({ toolName, decision: PolicyDecision.ALLOW, - priority: ALWAYS_ALLOW_PRIORITY, + priority, argsPattern: new RegExp(pattern), source: 'Dynamic (Confirmed)', }); @@ -518,10 +550,23 @@ export function createPolicyUpdater( ? new RegExp(message.argsPattern) : undefined; + const tier = + message.persistScope === 'user' + ? USER_POLICY_TIER + : WORKSPACE_POLICY_TIER; + const priority = tier + getAlwaysAllowPriorityFraction() / 1000; + + if (SENSITIVE_TOOLS.has(toolName) && !message.argsPattern) { + debugLogger.warn( + `Attempted to update policy for sensitive tool '${toolName}' without an argsPattern. Skipping.`, + ); + return; + } + policyEngine.addRule({ toolName, decision: PolicyDecision.ALLOW, - priority: ALWAYS_ALLOW_PRIORITY, + priority, argsPattern, source: 'Dynamic (Confirmed)', }); @@ -530,7 +575,10 @@ export function createPolicyUpdater( if (message.persist) { persistenceQueue = persistenceQueue.then(async () => { try { - const policyFile = storage.getAutoSavedPolicyPath(); + const policyFile = + message.persistScope === 'workspace' + ? storage.getWorkspaceAutoSavedPolicyPath() + : storage.getAutoSavedPolicyPath(); await fs.mkdir(path.dirname(policyFile), { recursive: true }); // Read existing file @@ -560,21 +608,19 @@ export function createPolicyUpdater( } // Create new rule object - const newRule: TomlRule = {}; + const newRule: TomlRule = { + decision: 'allow', + priority: getAlwaysAllowPriorityFraction(), + }; if (message.mcpName) { newRule.mcpName = message.mcpName; // Extract simple tool name - const simpleToolName = toolName.startsWith(`${message.mcpName}__`) + newRule.toolName = toolName.startsWith(`${message.mcpName}__`) ? toolName.slice(message.mcpName.length + 2) : toolName; - newRule.toolName = simpleToolName; - newRule.decision = 'allow'; - newRule.priority = 200; } else { newRule.toolName = toolName; - newRule.decision = 'allow'; - newRule.priority = 100; } if (message.commandPrefix) { diff --git a/packages/core/src/policy/persistence.test.ts b/packages/core/src/policy/persistence.test.ts index c5a71fdd93..da39160020 100644 --- a/packages/core/src/policy/persistence.test.ts +++ b/packages/core/src/policy/persistence.test.ts @@ -4,25 +4,22 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { - describe, - it, - expect, - vi, - beforeEach, - afterEach, - type Mock, -} from 'vitest'; -import * as fs from 'node:fs/promises'; +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import * as path from 'node:path'; -import { createPolicyUpdater, ALWAYS_ALLOW_PRIORITY } from './config.js'; +import { + createPolicyUpdater, + getAlwaysAllowPriorityFraction, +} from './config.js'; import { PolicyEngine } from './policy-engine.js'; import { MessageBus } from '../confirmation-bus/message-bus.js'; import { MessageBusType } from '../confirmation-bus/types.js'; import { Storage, AUTO_SAVED_POLICY_FILENAME } from '../config/storage.js'; import { ApprovalMode } from './types.js'; +import { vol, fs as memfs } from 'memfs'; + +// Use memfs for all fs operations in this test +vi.mock('node:fs/promises', () => import('memfs').then((m) => m.fs.promises)); -vi.mock('node:fs/promises'); vi.mock('../config/storage.js'); describe('createPolicyUpdater', () => { @@ -31,6 +28,8 @@ describe('createPolicyUpdater', () => { let mockStorage: Storage; beforeEach(() => { + vi.useFakeTimers(); + vol.reset(); policyEngine = new PolicyEngine({ rules: [], checkers: [], @@ -43,202 +42,184 @@ describe('createPolicyUpdater', () => { afterEach(() => { vi.restoreAllMocks(); + vi.useRealTimers(); }); it('should persist policy when persist flag is true', async () => { createPolicyUpdater(policyEngine, messageBus, mockStorage); - const userPoliciesDir = '/mock/user/.gemini/policies'; - const policyFile = path.join(userPoliciesDir, AUTO_SAVED_POLICY_FILENAME); + const policyFile = '/mock/user/.gemini/policies/auto-saved.toml'; vi.spyOn(mockStorage, 'getAutoSavedPolicyPath').mockReturnValue(policyFile); - (fs.mkdir as unknown as Mock).mockResolvedValue(undefined); - (fs.readFile as unknown as Mock).mockRejectedValue( - new Error('File not found'), - ); // Simulate new file - const mockFileHandle = { - writeFile: vi.fn().mockResolvedValue(undefined), - close: vi.fn().mockResolvedValue(undefined), - }; - (fs.open as unknown as Mock).mockResolvedValue(mockFileHandle); - (fs.rename as unknown as Mock).mockResolvedValue(undefined); - - const toolName = 'test_tool'; await messageBus.publish({ type: MessageBusType.UPDATE_POLICY, - toolName, + toolName: 'test_tool', persist: true, }); - // Wait for async operations (microtasks) - await new Promise((resolve) => setTimeout(resolve, 0)); + // Policy updater handles persistence asynchronously in a promise queue. + // We use advanceTimersByTimeAsync to yield to the microtask queue. + await vi.advanceTimersByTimeAsync(100); - expect(fs.mkdir).toHaveBeenCalledWith(userPoliciesDir, { - recursive: true, - }); + const fileExists = memfs.existsSync(policyFile); + expect(fileExists).toBe(true); - expect(fs.open).toHaveBeenCalledWith(expect.stringMatching(/\.tmp$/), 'wx'); - - // Check written content - const expectedContent = expect.stringContaining(`toolName = "test_tool"`); - expect(mockFileHandle.writeFile).toHaveBeenCalledWith( - expectedContent, - 'utf-8', - ); - expect(fs.rename).toHaveBeenCalledWith( - expect.stringMatching(/\.tmp$/), - policyFile, - ); + const content = memfs.readFileSync(policyFile, 'utf-8') as string; + expect(content).toContain('toolName = "test_tool"'); + expect(content).toContain('decision = "allow"'); + const expectedPriority = getAlwaysAllowPriorityFraction(); + expect(content).toContain(`priority = ${expectedPriority}`); }); it('should not persist policy when persist flag is false or undefined', 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', }); - await new Promise((resolve) => setTimeout(resolve, 0)); + await vi.advanceTimersByTimeAsync(100); - expect(fs.writeFile).not.toHaveBeenCalled(); - expect(fs.rename).not.toHaveBeenCalled(); + expect(memfs.existsSync(policyFile)).toBe(false); }); - it('should persist policy with commandPrefix when provided', async () => { + it('should append to existing policy file', async () => { createPolicyUpdater(policyEngine, messageBus, mockStorage); - const userPoliciesDir = '/mock/user/.gemini/policies'; - const policyFile = path.join(userPoliciesDir, AUTO_SAVED_POLICY_FILENAME); + const policyFile = '/mock/user/.gemini/policies/auto-saved.toml'; vi.spyOn(mockStorage, 'getAutoSavedPolicyPath').mockReturnValue(policyFile); - (fs.mkdir as unknown as Mock).mockResolvedValue(undefined); - (fs.readFile as unknown as Mock).mockRejectedValue( - new Error('File not found'), - ); - const mockFileHandle = { - writeFile: vi.fn().mockResolvedValue(undefined), - close: vi.fn().mockResolvedValue(undefined), - }; - (fs.open as unknown as Mock).mockResolvedValue(mockFileHandle); - (fs.rename as unknown as Mock).mockResolvedValue(undefined); - - const toolName = 'run_shell_command'; - const commandPrefix = 'git status'; + const existingContent = + '[[rule]]\ntoolName = "existing_tool"\ndecision = "allow"\n'; + const dir = path.dirname(policyFile); + memfs.mkdirSync(dir, { recursive: true }); + memfs.writeFileSync(policyFile, existingContent); await messageBus.publish({ type: MessageBusType.UPDATE_POLICY, - toolName, + toolName: 'new_tool', persist: true, - commandPrefix, }); - await new Promise((resolve) => setTimeout(resolve, 0)); + await vi.advanceTimersByTimeAsync(100); - // In-memory rule check (unchanged) - const rules = policyEngine.getRules(); - const addedRule = rules.find((r) => r.toolName === toolName); - expect(addedRule).toBeDefined(); - expect(addedRule?.priority).toBe(ALWAYS_ALLOW_PRIORITY); - expect(addedRule?.argsPattern).toEqual( - new RegExp(`"command":"git\\ status(?:[\\s"]|\\\\")`), - ); - - // Verify file written - expect(fs.open).toHaveBeenCalledWith(expect.stringMatching(/\.tmp$/), 'wx'); - expect(mockFileHandle.writeFile).toHaveBeenCalledWith( - expect.stringContaining(`commandPrefix = "git status"`), - 'utf-8', - ); + const content = memfs.readFileSync(policyFile, 'utf-8') as string; + expect(content).toContain('toolName = "existing_tool"'); + expect(content).toContain('toolName = "new_tool"'); }); - it('should persist policy with mcpName and toolName when provided', async () => { + it('should handle toml with multiple rules correctly', async () => { createPolicyUpdater(policyEngine, messageBus, mockStorage); - const userPoliciesDir = '/mock/user/.gemini/policies'; - const policyFile = path.join(userPoliciesDir, AUTO_SAVED_POLICY_FILENAME); + const policyFile = '/mock/user/.gemini/policies/auto-saved.toml'; vi.spyOn(mockStorage, 'getAutoSavedPolicyPath').mockReturnValue(policyFile); - (fs.mkdir as unknown as Mock).mockResolvedValue(undefined); - (fs.readFile as unknown as Mock).mockRejectedValue( - new Error('File not found'), - ); - const mockFileHandle = { - writeFile: vi.fn().mockResolvedValue(undefined), - close: vi.fn().mockResolvedValue(undefined), - }; - (fs.open as unknown as Mock).mockResolvedValue(mockFileHandle); - (fs.rename as unknown as Mock).mockResolvedValue(undefined); + const existingContent = ` +[[rule]] +toolName = "tool1" +decision = "allow" - const mcpName = 'my-jira-server'; - const simpleToolName = 'search'; - const toolName = `${mcpName}__${simpleToolName}`; +[[rule]] +toolName = "tool2" +decision = "deny" +`; + const dir = path.dirname(policyFile); + memfs.mkdirSync(dir, { recursive: true }); + memfs.writeFileSync(policyFile, existingContent); await messageBus.publish({ type: MessageBusType.UPDATE_POLICY, - toolName, + toolName: 'tool3', persist: true, - mcpName, }); - await new Promise((resolve) => setTimeout(resolve, 0)); + await vi.advanceTimersByTimeAsync(100); - // Verify file written - expect(fs.open).toHaveBeenCalledWith(expect.stringMatching(/\.tmp$/), 'wx'); - const writeCall = mockFileHandle.writeFile.mock.calls[0]; - const writtenContent = writeCall[0] as string; - expect(writtenContent).toContain(`mcpName = "${mcpName}"`); - expect(writtenContent).toContain(`toolName = "${simpleToolName}"`); - expect(writtenContent).toContain('priority = 200'); + const content = memfs.readFileSync(policyFile, 'utf-8') as string; + expect(content).toContain('toolName = "tool1"'); + expect(content).toContain('toolName = "tool2"'); + expect(content).toContain('toolName = "tool3"'); }); - it('should escape special characters in toolName and mcpName', async () => { + it('should include argsPattern if provided', async () => { createPolicyUpdater(policyEngine, messageBus, mockStorage); - const userPoliciesDir = '/mock/user/.gemini/policies'; - const policyFile = path.join(userPoliciesDir, AUTO_SAVED_POLICY_FILENAME); + const policyFile = '/mock/user/.gemini/policies/auto-saved.toml'; vi.spyOn(mockStorage, 'getAutoSavedPolicyPath').mockReturnValue(policyFile); - (fs.mkdir as unknown as Mock).mockResolvedValue(undefined); - (fs.readFile as unknown as Mock).mockRejectedValue( - new Error('File not found'), - ); - - const mockFileHandle = { - writeFile: vi.fn().mockResolvedValue(undefined), - close: vi.fn().mockResolvedValue(undefined), - }; - (fs.open as unknown as Mock).mockResolvedValue(mockFileHandle); - (fs.rename as unknown as Mock).mockResolvedValue(undefined); - - const mcpName = 'my"jira"server'; - const toolName = `my"jira"server__search"tool"`; await messageBus.publish({ type: MessageBusType.UPDATE_POLICY, - toolName, + toolName: 'test_tool', persist: true, - mcpName, + argsPattern: '^foo.*$', }); - await new Promise((resolve) => setTimeout(resolve, 0)); + await vi.advanceTimersByTimeAsync(100); - expect(fs.open).toHaveBeenCalledWith(expect.stringMatching(/\.tmp$/), 'wx'); - const writeCall = mockFileHandle.writeFile.mock.calls[0]; - const writtenContent = writeCall[0] as string; + const content = memfs.readFileSync(policyFile, 'utf-8') as string; + expect(content).toContain('argsPattern = "^foo.*$"'); + }); - // Verify escaping - should be valid TOML + it('should include mcpName 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: 'search"tool"', + persist: true, + mcpName: 'my"jira"server', + }); + + await vi.advanceTimersByTimeAsync(100); + + const writtenContent = memfs.readFileSync(policyFile, 'utf-8') as string; + + // Verify escaping - should be valid TOML and contain the values // Note: @iarna/toml optimizes for shortest representation, so it may use single quotes 'foo"bar' // instead of "foo\"bar\"" if there are no single quotes in the string. try { - expect(writtenContent).toContain(`mcpName = "my\\"jira\\"server"`); + expect(writtenContent).toContain('mcpName = "my\\"jira\\"server"'); } catch { - expect(writtenContent).toContain(`mcpName = 'my"jira"server'`); + expect(writtenContent).toContain('mcpName = \'my"jira"server\''); } try { - expect(writtenContent).toContain(`toolName = "search\\"tool\\""`); + expect(writtenContent).toContain('toolName = "search\\"tool\\""'); } catch { - expect(writtenContent).toContain(`toolName = 'search"tool"'`); + expect(writtenContent).toContain('toolName = \'search"tool"\''); } }); + + it('should persist to workspace when persistScope is workspace', async () => { + createPolicyUpdater(policyEngine, messageBus, mockStorage); + + const workspacePoliciesDir = '/mock/project/.gemini/policies'; + const policyFile = path.join( + workspacePoliciesDir, + AUTO_SAVED_POLICY_FILENAME, + ); + vi.spyOn(mockStorage, 'getWorkspaceAutoSavedPolicyPath').mockReturnValue( + policyFile, + ); + + await messageBus.publish({ + type: MessageBusType.UPDATE_POLICY, + toolName: 'test_tool', + persist: true, + persistScope: 'workspace', + }); + + await vi.advanceTimersByTimeAsync(100); + + expect(memfs.existsSync(policyFile)).toBe(true); + const content = memfs.readFileSync(policyFile, 'utf-8') as string; + expect(content).toContain('toolName = "test_tool"'); + }); }); diff --git a/packages/core/src/policy/policy-updater.test.ts b/packages/core/src/policy/policy-updater.test.ts index 3037667949..7aafcd5153 100644 --- a/packages/core/src/policy/policy-updater.test.ts +++ b/packages/core/src/policy/policy-updater.test.ts @@ -19,6 +19,7 @@ import { type PolicyUpdateOptions, } from '../tools/tools.js'; import * as shellUtils from '../utils/shell-utils.js'; +import { escapeRegex } from './utils.js'; vi.mock('node:fs/promises'); vi.mock('../config/storage.js'); @@ -75,7 +76,9 @@ describe('createPolicyUpdater', () => { expect.objectContaining({ toolName: 'run_shell_command', priority: ALWAYS_ALLOW_PRIORITY, - argsPattern: new RegExp('"command":"echo(?:[\\s"]|\\\\")'), + argsPattern: new RegExp( + escapeRegex('"command":"echo') + '(?:[\\s"]|\\\\")', + ), }), ); expect(policyEngine.addRule).toHaveBeenNthCalledWith( @@ -83,7 +86,9 @@ describe('createPolicyUpdater', () => { expect.objectContaining({ toolName: 'run_shell_command', priority: ALWAYS_ALLOW_PRIORITY, - argsPattern: new RegExp('"command":"ls(?:[\\s"]|\\\\")'), + argsPattern: new RegExp( + escapeRegex('"command":"ls') + '(?:[\\s"]|\\\\")', + ), }), ); }); @@ -103,7 +108,9 @@ describe('createPolicyUpdater', () => { expect.objectContaining({ toolName: 'run_shell_command', priority: ALWAYS_ALLOW_PRIORITY, - argsPattern: new RegExp('"command":"git(?:[\\s"]|\\\\")'), + argsPattern: new RegExp( + escapeRegex('"command":"git') + '(?:[\\s"]|\\\\")', + ), }), ); }); diff --git a/packages/core/src/policy/utils.test.ts b/packages/core/src/policy/utils.test.ts index 90f3c632c7..db6225827a 100644 --- a/packages/core/src/policy/utils.test.ts +++ b/packages/core/src/policy/utils.test.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect } from 'vitest'; +import { expect, describe, it } from 'vitest'; import { escapeRegex, buildArgsPatterns, isSafeRegExp } from './utils.js'; describe('policy/utils', () => { @@ -43,20 +43,20 @@ describe('policy/utils', () => { }); it('should return false for invalid regexes', () => { + expect(isSafeRegExp('[')).toBe(false); expect(isSafeRegExp('([a-z)')).toBe(false); expect(isSafeRegExp('*')).toBe(false); }); - it('should return false for extremely long regexes', () => { - expect(isSafeRegExp('a'.repeat(2049))).toBe(false); + it('should return false for long regexes', () => { + expect(isSafeRegExp('a'.repeat(3000))).toBe(false); }); - it('should return false for nested quantifiers (potential ReDoS)', () => { + it('should return false for nested quantifiers (ReDoS heuristic)', () => { expect(isSafeRegExp('(a+)+')).toBe(false); - expect(isSafeRegExp('(a+)*')).toBe(false); - expect(isSafeRegExp('(a*)+')).toBe(false); - expect(isSafeRegExp('(a*)*')).toBe(false); - expect(isSafeRegExp('(a|b+)+')).toBe(false); + expect(isSafeRegExp('(a|b)*')).toBe(true); + expect(isSafeRegExp('(.*)*')).toBe(false); + expect(isSafeRegExp('([a-z]+)+')).toBe(false); expect(isSafeRegExp('(.*)+')).toBe(false); }); }); @@ -69,14 +69,14 @@ describe('policy/utils', () => { it('should build pattern from a single commandPrefix', () => { const result = buildArgsPatterns(undefined, 'ls', undefined); - expect(result).toEqual(['"command":"ls(?:[\\s"]|\\\\")']); + expect(result).toEqual(['\\"command\\":\\"ls(?:[\\s"]|\\\\")']); }); it('should build patterns from an array of commandPrefixes', () => { - const result = buildArgsPatterns(undefined, ['ls', 'cd'], undefined); + const result = buildArgsPatterns(undefined, ['echo', 'ls'], undefined); expect(result).toEqual([ - '"command":"ls(?:[\\s"]|\\\\")', - '"command":"cd(?:[\\s"]|\\\\")', + '\\"command\\":\\"echo(?:[\\s"]|\\\\")', + '\\"command\\":\\"ls(?:[\\s"]|\\\\")', ]); }); @@ -87,7 +87,7 @@ describe('policy/utils', () => { it('should prioritize commandPrefix over commandRegex and argsPattern', () => { const result = buildArgsPatterns('raw', 'prefix', 'regex'); - expect(result).toEqual(['"command":"prefix(?:[\\s"]|\\\\")']); + expect(result).toEqual(['\\"command\\":\\"prefix(?:[\\s"]|\\\\")']); }); it('should prioritize commandRegex over argsPattern if no commandPrefix', () => { @@ -98,14 +98,15 @@ describe('policy/utils', () => { it('should escape characters in commandPrefix', () => { const result = buildArgsPatterns(undefined, 'git checkout -b', undefined); expect(result).toEqual([ - '"command":"git\\ checkout\\ \\-b(?:[\\s"]|\\\\")', + '\\"command\\":\\"git\\ checkout\\ \\-b(?:[\\s"]|\\\\")', ]); }); it('should correctly escape quotes in commandPrefix', () => { const result = buildArgsPatterns(undefined, 'git "fix"', undefined); expect(result).toEqual([ - '"command":"git\\ \\\\\\"fix\\\\\\"(?:[\\s"]|\\\\")', + // eslint-disable-next-line no-useless-escape + '\\\"command\\\":\\\"git\\ \\\\\\\"fix\\\\\\\"(?:[\\s\"]|\\\\\")', ]); }); @@ -142,7 +143,7 @@ describe('policy/utils', () => { const gitRegex = new RegExp(gitPatterns[0]!); // git\status -> {"command":"git\\status"} const gitAttack = '{"command":"git\\\\status"}'; - expect(gitRegex.test(gitAttack)).toBe(false); + expect(gitAttack).not.toMatch(gitRegex); }); }); }); diff --git a/packages/core/src/policy/utils.ts b/packages/core/src/policy/utils.ts index 3742ba3ed6..bec3e9e0cd 100644 --- a/packages/core/src/policy/utils.ts +++ b/packages/core/src/policy/utils.ts @@ -63,16 +63,22 @@ export function buildArgsPatterns( ? commandPrefix : [commandPrefix]; - // Expand command prefixes to multiple patterns. - // We append [\\s"] to ensure we match whole words only (e.g., "git" but not - // "github"). Since we match against JSON stringified args, the value is - // always followed by a space or a closing quote. return prefixes.map((prefix) => { - const jsonPrefix = JSON.stringify(prefix).slice(1, -1); + // JSON.stringify safely encodes the prefix in quotes. + // We remove ONLY the trailing quote to match it as an open prefix string. + const encodedPrefix = JSON.stringify(prefix); + const openQuotePrefix = encodedPrefix.substring( + 0, + encodedPrefix.length - 1, + ); + + // Escape the exact JSON literal segment we expect to see + const matchSegment = escapeRegex(`"command":${openQuotePrefix}`); + // We allow [\s], ["], or the specific sequence [\"] (for escaped quotes // in JSON). We do NOT allow generic [\\], which would match "git\status" // -> "gitstatus". - return `"command":"${escapeRegex(jsonPrefix)}(?:[\\s"]|\\\\")`; + return `${matchSegment}(?:[\\s"]|\\\\")`; }); } @@ -82,3 +88,30 @@ export function buildArgsPatterns( return [argsPattern]; } + +/** + * Builds a regex pattern to match a specific file path in tool arguments. + * This is used to narrow tool approvals for edit tools to specific files. + * + * @param filePath The relative path to the file. + * @returns A regex string that matches "file_path":"" in a JSON string. + */ +export function buildFilePathArgsPattern(filePath: string): string { + // JSON.stringify safely encodes the path (handling quotes, backslashes, etc) + // and wraps it in double quotes. We simply prepend the key name and escape + // the entire sequence for Regex matching without any slicing. + const encodedPath = JSON.stringify(filePath); + return escapeRegex(`"file_path":${encodedPath}`); +} + +/** + * Builds a regex pattern to match a specific "pattern" in tool arguments. + * This is used to narrow tool approvals for search tools like glob/grep to specific patterns. + * + * @param pattern The pattern to match. + * @returns A regex string that matches "pattern":"" in a JSON string. + */ +export function buildPatternArgsPattern(pattern: string): string { + const encodedPattern = JSON.stringify(pattern); + return escapeRegex(`"pattern":${encodedPattern}`); +} diff --git a/packages/core/src/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts index 9320893bd6..4bf2b32a46 100644 --- a/packages/core/src/scheduler/policy.test.ts +++ b/packages/core/src/scheduler/policy.test.ts @@ -16,8 +16,12 @@ import { import { checkPolicy, updatePolicy, getPolicyDenialError } from './policy.js'; import type { Config } from '../config/config.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; -import { MessageBusType } from '../confirmation-bus/types.js'; +import { + MessageBusType, + type SerializableConfirmationDetails, +} from '../confirmation-bus/types.js'; import { ApprovalMode, PolicyDecision } from '../policy/types.js'; +import { escapeRegex } from '../policy/utils.js'; import { ToolConfirmationOutcome, type AnyDeclarativeTool, @@ -219,6 +223,8 @@ describe('policy.ts', () => { it('should handle standard policy updates with persistence', async () => { const mockConfig = { + isTrustedFolder: vi.fn().mockReturnValue(false), + getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined), setApprovalMode: vi.fn(), } as unknown as Mocked; @@ -453,6 +459,8 @@ describe('policy.ts', () => { it('should handle MCP ProceedAlwaysAndSave (persist: true)', async () => { const mockConfig = { + isTrustedFolder: vi.fn().mockReturnValue(false), + getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined), setApprovalMode: vi.fn(), } as unknown as Mocked; @@ -487,6 +495,96 @@ describe('policy.ts', () => { }), ); }); + + it('should determine persistScope: workspace in trusted folders', async () => { + const mockConfig = { + isTrustedFolder: vi.fn().mockReturnValue(true), + getWorkspacePoliciesDir: vi + .fn() + .mockReturnValue('/mock/project/policies'), + setApprovalMode: vi.fn(), + } as unknown as Mocked; + const mockMessageBus = { + publish: vi.fn(), + } as unknown as Mocked; + const tool = { name: 'test-tool' } as AnyDeclarativeTool; + + await updatePolicy( + tool, + ToolConfirmationOutcome.ProceedAlwaysAndSave, + undefined, + { config: mockConfig, messageBus: mockMessageBus }, + ); + + expect(mockMessageBus.publish).toHaveBeenCalledWith( + expect.objectContaining({ + persistScope: 'workspace', + }), + ); + }); + + it('should determine persistScope: user in untrusted folders', async () => { + const mockConfig = { + isTrustedFolder: vi.fn().mockReturnValue(false), + getWorkspacePoliciesDir: vi + .fn() + .mockReturnValue('/mock/project/policies'), + setApprovalMode: vi.fn(), + } as unknown as Mocked; + const mockMessageBus = { + publish: vi.fn(), + } as unknown as Mocked; + const tool = { name: 'test-tool' } as AnyDeclarativeTool; + + await updatePolicy( + tool, + ToolConfirmationOutcome.ProceedAlwaysAndSave, + undefined, + { config: mockConfig, messageBus: mockMessageBus }, + ); + + expect(mockMessageBus.publish).toHaveBeenCalledWith( + expect.objectContaining({ + persistScope: 'user', + }), + ); + }); + + it('should narrow edit tools with argsPattern', async () => { + const mockConfig = { + isTrustedFolder: vi.fn().mockReturnValue(false), + getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined), + getTargetDir: vi.fn().mockReturnValue('/mock/dir'), + setApprovalMode: vi.fn(), + } as unknown as Mocked; + const mockMessageBus = { + publish: vi.fn(), + } as unknown as Mocked; + const tool = { name: 'write_file' } as AnyDeclarativeTool; + const details: SerializableConfirmationDetails = { + type: 'edit', + title: 'Edit', + filePath: 'src/foo.ts', + fileName: 'foo.ts', + fileDiff: '--- foo.ts\n+++ foo.ts\n@@ -1 +1 @@\n-old\n+new', + originalContent: 'old', + newContent: 'new', + }; + + await updatePolicy( + tool, + ToolConfirmationOutcome.ProceedAlwaysAndSave, + details, + { config: mockConfig, messageBus: mockMessageBus }, + ); + + expect(mockMessageBus.publish).toHaveBeenCalledWith( + expect.objectContaining({ + toolName: 'write_file', + argsPattern: escapeRegex('"file_path":"src/foo.ts"'), + }), + ); + }); }); describe('getPolicyDenialError', () => { diff --git a/packages/core/src/scheduler/policy.ts b/packages/core/src/scheduler/policy.ts index ad4aa745bb..1ac70a108b 100644 --- a/packages/core/src/scheduler/policy.ts +++ b/packages/core/src/scheduler/policy.ts @@ -20,8 +20,11 @@ import { import { ToolConfirmationOutcome, type AnyDeclarativeTool, + type AnyToolInvocation, type PolicyUpdateOptions, } from '../tools/tools.js'; +import { buildFilePathArgsPattern } from '../policy/utils.js'; +import { makeRelative } from '../utils/paths.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; import { EDIT_TOOL_NAMES } from '../tools/tool-names.js'; import type { ValidatingToolCall } from './types.js'; @@ -94,7 +97,11 @@ export async function updatePolicy( tool: AnyDeclarativeTool, outcome: ToolConfirmationOutcome, confirmationDetails: SerializableConfirmationDetails | undefined, - deps: { config: Config; messageBus: MessageBus }, + deps: { + config: Config; + messageBus: MessageBus; + toolInvocation?: AnyToolInvocation; + }, ): Promise { // Mode Transitions (AUTO_EDIT) if (isAutoEditTransition(tool, outcome)) { @@ -102,6 +109,20 @@ export async function updatePolicy( return; } + // Determine persist scope if we are persisting. + let persistScope: 'workspace' | 'user' | undefined; + if (outcome === ToolConfirmationOutcome.ProceedAlwaysAndSave) { + // If folder is trusted and workspace policies are enabled, we prefer workspace scope. + if ( + deps.config.isTrustedFolder() && + deps.config.getWorkspacePoliciesDir() !== undefined + ) { + persistScope = 'workspace'; + } else { + persistScope = 'user'; + } + } + // Specialized Tools (MCP) if (confirmationDetails?.type === 'mcp') { await handleMcpPolicyUpdate( @@ -109,6 +130,7 @@ export async function updatePolicy( outcome, confirmationDetails, deps.messageBus, + persistScope, ); return; } @@ -119,6 +141,9 @@ export async function updatePolicy( outcome, confirmationDetails, deps.messageBus, + persistScope, + deps.toolInvocation, + deps.config, ); } @@ -148,21 +173,31 @@ async function handleStandardPolicyUpdate( outcome: ToolConfirmationOutcome, confirmationDetails: SerializableConfirmationDetails | undefined, messageBus: MessageBus, + persistScope?: 'workspace' | 'user', + toolInvocation?: AnyToolInvocation, + config?: Config, ): Promise { if ( outcome === ToolConfirmationOutcome.ProceedAlways || outcome === ToolConfirmationOutcome.ProceedAlwaysAndSave ) { - const options: PolicyUpdateOptions = {}; + const options: PolicyUpdateOptions = + toolInvocation?.getPolicyUpdateOptions?.(outcome) || {}; - if (confirmationDetails?.type === 'exec') { + if (!options.commandPrefix && confirmationDetails?.type === 'exec') { options.commandPrefix = confirmationDetails.rootCommands; + } else if (!options.argsPattern && confirmationDetails?.type === 'edit') { + const filePath = config + ? makeRelative(confirmationDetails.filePath, config.getTargetDir()) + : confirmationDetails.filePath; + options.argsPattern = buildFilePathArgsPattern(filePath); } await messageBus.publish({ type: MessageBusType.UPDATE_POLICY, toolName: tool.name, persist: outcome === ToolConfirmationOutcome.ProceedAlwaysAndSave, + persistScope, ...options, }); } @@ -180,6 +215,7 @@ async function handleMcpPolicyUpdate( { type: 'mcp' } >, messageBus: MessageBus, + persistScope?: 'workspace' | 'user', ): Promise { const isMcpAlways = outcome === ToolConfirmationOutcome.ProceedAlways || @@ -204,5 +240,6 @@ async function handleMcpPolicyUpdate( toolName, mcpName: confirmationDetails.serverName, persist, + persistScope, }); } diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 613e23b2d6..187916623e 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -608,6 +608,7 @@ export class Scheduler { await updatePolicy(toolCall.tool, outcome, lastDetails, { config: this.config, messageBus: this.messageBus, + toolInvocation: toolCall.invocation, }); } diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index 214875c574..06f9657745 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -20,11 +20,14 @@ import { type ToolLocation, type ToolResult, type ToolResultDisplay, + type PolicyUpdateOptions, } from './tools.js'; +import { buildFilePathArgsPattern } from '../policy/utils.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { ToolErrorType } from './tool-error.js'; import { makeRelative, shortenPath } from '../utils/paths.js'; import { isNodeError } from '../utils/errors.js'; +import { correctPath } from '../utils/pathCorrector.js'; import type { Config } from '../config/config.js'; import { ApprovalMode } from '../policy/types.js'; import { CoreToolCallStatus } from '../scheduler/types.js'; @@ -44,7 +47,6 @@ import { logEditCorrectionEvent, } from '../telemetry/loggers.js'; -import { correctPath } from '../utils/pathCorrector.js'; import { EDIT_TOOL_NAME, READ_FILE_TOOL_NAME, @@ -442,6 +444,8 @@ class EditToolInvocation extends BaseToolInvocation implements ToolInvocation { + private readonly resolvedPath: string; + constructor( private readonly config: Config, params: EditToolParams, @@ -450,10 +454,31 @@ class EditToolInvocation displayName?: string, ) { super(params, messageBus, toolName, displayName); + if (!path.isAbsolute(this.params.file_path)) { + const result = correctPath(this.params.file_path, this.config); + if (result.success) { + this.resolvedPath = result.correctedPath; + } else { + this.resolvedPath = path.resolve( + this.config.getTargetDir(), + this.params.file_path, + ); + } + } else { + this.resolvedPath = this.params.file_path; + } } override toolLocations(): ToolLocation[] { - return [{ path: this.params.file_path }]; + return [{ path: this.resolvedPath }]; + } + + override getPolicyUpdateOptions( + _outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined { + return { + argsPattern: buildFilePathArgsPattern(this.params.file_path), + }; } private async attemptSelfCorrection( @@ -471,7 +496,7 @@ class EditToolInvocation const initialContentHash = hashContent(currentContent); const onDiskContent = await this.config .getFileSystemService() - .readTextFile(params.file_path); + .readTextFile(this.resolvedPath); const onDiskContentHash = hashContent(onDiskContent.replace(/\r\n/g, '\n')); if (initialContentHash !== onDiskContentHash) { @@ -582,7 +607,7 @@ class EditToolInvocation try { currentContent = await this.config .getFileSystemService() - .readTextFile(params.file_path); + .readTextFile(this.resolvedPath); originalLineEnding = detectLineEnding(currentContent); currentContent = currentContent.replace(/\r\n/g, '\n'); fileExists = true; @@ -615,7 +640,7 @@ class EditToolInvocation isNewFile: false, error: { display: `File not found. Cannot apply edit. Use an empty old_string to create a new file.`, - raw: `File not found: ${params.file_path}`, + raw: `File not found: ${this.resolvedPath}`, type: ToolErrorType.FILE_NOT_FOUND, }, originalLineEnding, @@ -630,7 +655,7 @@ class EditToolInvocation isNewFile: false, error: { display: `Failed to read content of file.`, - raw: `Failed to read content of existing file: ${params.file_path}`, + raw: `Failed to read content of existing file: ${this.resolvedPath}`, type: ToolErrorType.READ_CONTENT_FAILURE, }, originalLineEnding, @@ -645,7 +670,7 @@ class EditToolInvocation isNewFile: false, error: { display: `Failed to edit. Attempted to create a file that already exists.`, - raw: `File already exists, cannot create: ${params.file_path}`, + raw: `File already exists, cannot create: ${this.resolvedPath}`, type: ToolErrorType.ATTEMPT_TO_CREATE_EXISTING_FILE, }, originalLineEnding, @@ -727,7 +752,7 @@ class EditToolInvocation return false; } - const fileName = path.basename(this.params.file_path); + const fileName = path.basename(this.resolvedPath); const fileDiff = Diff.createPatch( fileName, editData.currentContent ?? '', @@ -739,14 +764,14 @@ class EditToolInvocation const ideClient = await IdeClient.getInstance(); const ideConfirmation = this.config.getIdeMode() && ideClient.isDiffingEnabled() - ? ideClient.openDiff(this.params.file_path, editData.newContent) + ? ideClient.openDiff(this.resolvedPath, editData.newContent) : undefined; const confirmationDetails: ToolEditConfirmationDetails = { type: 'edit', - title: `Confirm Edit: ${shortenPath(makeRelative(this.params.file_path, this.config.getTargetDir()))}`, + title: `Confirm Edit: ${shortenPath(makeRelative(this.resolvedPath, this.config.getTargetDir()))}`, fileName, - filePath: this.params.file_path, + filePath: this.resolvedPath, fileDiff, originalContent: editData.currentContent, newContent: editData.newContent, @@ -771,7 +796,7 @@ class EditToolInvocation getDescription(): string { const relativePath = makeRelative( - this.params.file_path, + this.resolvedPath, this.config.getTargetDir(), ); if (this.params.old_string === '') { @@ -797,11 +822,7 @@ class EditToolInvocation * @returns Result of the edit operation */ async execute(signal: AbortSignal): Promise { - const resolvedPath = path.resolve( - this.config.getTargetDir(), - this.params.file_path, - ); - const validationError = this.config.validatePathAccess(resolvedPath); + const validationError = this.config.validatePathAccess(this.resolvedPath); if (validationError) { return { llmContent: validationError, @@ -843,7 +864,7 @@ class EditToolInvocation } try { - await this.ensureParentDirectoriesExistAsync(this.params.file_path); + await this.ensureParentDirectoriesExistAsync(this.resolvedPath); let finalContent = editData.newContent; // Restore original line endings if they were CRLF, or use OS default for new files @@ -856,15 +877,15 @@ class EditToolInvocation } await this.config .getFileSystemService() - .writeTextFile(this.params.file_path, finalContent); + .writeTextFile(this.resolvedPath, finalContent); let displayResult: ToolResultDisplay; if (editData.isNewFile) { - displayResult = `Created ${shortenPath(makeRelative(this.params.file_path, this.config.getTargetDir()))}`; + displayResult = `Created ${shortenPath(makeRelative(this.resolvedPath, this.config.getTargetDir()))}`; } else { // Generate diff for display, even though core logic doesn't technically need it // The CLI wrapper will use this part of the ToolResult - const fileName = path.basename(this.params.file_path); + const fileName = path.basename(this.resolvedPath); const fileDiff = Diff.createPatch( fileName, editData.currentContent ?? '', // Should not be null here if not isNewFile @@ -883,7 +904,7 @@ class EditToolInvocation displayResult = { fileDiff, fileName, - filePath: this.params.file_path, + filePath: this.resolvedPath, originalContent: editData.currentContent, newContent: editData.newContent, diffStat, @@ -893,8 +914,8 @@ class EditToolInvocation const llmSuccessMessageParts = [ editData.isNewFile - ? `Created new file: ${this.params.file_path} with provided content.` - : `Successfully modified file: ${this.params.file_path} (${editData.occurrences} replacements).`, + ? `Created new file: ${this.resolvedPath} with provided content.` + : `Successfully modified file: ${this.resolvedPath} (${editData.occurrences} replacements).`, ]; // Return a diff of the file before and after the write so that the agent @@ -985,16 +1006,20 @@ export class EditTool return "The 'file_path' parameter must be non-empty."; } - let filePath = params.file_path; - if (!path.isAbsolute(filePath)) { - // Attempt to auto-correct to an absolute path - const result = correctPath(filePath, this.config); - if (!result.success) { - return result.error; + let resolvedPath: string; + if (!path.isAbsolute(params.file_path)) { + const result = correctPath(params.file_path, this.config); + if (result.success) { + resolvedPath = result.correctedPath; + } else { + resolvedPath = path.resolve( + this.config.getTargetDir(), + params.file_path, + ); } - filePath = result.correctedPath; + } else { + resolvedPath = params.file_path; } - params.file_path = filePath; const newPlaceholders = detectOmissionPlaceholders(params.new_string); if (newPlaceholders.length > 0) { @@ -1009,7 +1034,7 @@ export class EditTool } } - return this.config.validatePathAccess(params.file_path); + return this.config.validatePathAccess(resolvedPath); } protected createInvocation( diff --git a/packages/core/src/tools/glob.ts b/packages/core/src/tools/glob.ts index c2f3c4ab54..9cef63759d 100644 --- a/packages/core/src/tools/glob.ts +++ b/packages/core/src/tools/glob.ts @@ -14,12 +14,15 @@ import { Kind, type ToolInvocation, type ToolResult, + type PolicyUpdateOptions, + type ToolConfirmationOutcome, } from './tools.js'; import { shortenPath, makeRelative } from '../utils/paths.js'; import { type Config } from '../config/config.js'; import { DEFAULT_FILE_FILTERING_OPTIONS } from '../config/constants.js'; import { ToolErrorType } from './tool-error.js'; import { GLOB_TOOL_NAME, GLOB_DISPLAY_NAME } from './tool-names.js'; +import { buildPatternArgsPattern } from '../policy/utils.js'; import { getErrorMessage } from '../utils/errors.js'; import { debugLogger } from '../utils/debugLogger.js'; import { GLOB_DEFINITION } from './definitions/coreTools.js'; @@ -118,6 +121,14 @@ class GlobToolInvocation extends BaseToolInvocation< return description; } + override getPolicyUpdateOptions( + _outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined { + return { + argsPattern: buildPatternArgsPattern(this.params.pattern), + }; + } + async execute(signal: AbortSignal): Promise { try { const workspaceContext = this.config.getWorkspaceContext(); diff --git a/packages/core/src/tools/grep.ts b/packages/core/src/tools/grep.ts index c7e676951a..f0d7aaa4aa 100644 --- a/packages/core/src/tools/grep.ts +++ b/packages/core/src/tools/grep.ts @@ -21,6 +21,8 @@ import { Kind, type ToolInvocation, type ToolResult, + type PolicyUpdateOptions, + type ToolConfirmationOutcome, } from './tools.js'; import { makeRelative, shortenPath } from '../utils/paths.js'; import { getErrorMessage, isNodeError } from '../utils/errors.js'; @@ -29,6 +31,7 @@ import type { Config } from '../config/config.js'; import type { FileExclusions } from '../utils/ignorePatterns.js'; import { ToolErrorType } from './tool-error.js'; import { GREP_TOOL_NAME } from './tool-names.js'; +import { buildPatternArgsPattern } from '../policy/utils.js'; import { debugLogger } from '../utils/debugLogger.js'; import { GREP_DEFINITION } from './definitions/coreTools.js'; import { resolveToolDeclaration } from './definitions/resolver.js'; @@ -285,6 +288,14 @@ class GrepToolInvocation extends BaseToolInvocation< } } + override getPolicyUpdateOptions( + _outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined { + return { + argsPattern: buildPatternArgsPattern(this.params.pattern), + }; + } + /** * Checks if a command is available in the system's PATH. * @param {string} command The command name (e.g., 'git', 'grep'). diff --git a/packages/core/src/tools/ls.ts b/packages/core/src/tools/ls.ts index 9456f8ffc9..1e2d1cccf8 100644 --- a/packages/core/src/tools/ls.ts +++ b/packages/core/src/tools/ls.ts @@ -13,12 +13,15 @@ import { Kind, type ToolInvocation, type ToolResult, + type PolicyUpdateOptions, + type ToolConfirmationOutcome, } from './tools.js'; import { makeRelative, shortenPath } from '../utils/paths.js'; import type { Config } from '../config/config.js'; import { DEFAULT_FILE_FILTERING_OPTIONS } from '../config/constants.js'; import { ToolErrorType } from './tool-error.js'; import { LS_TOOL_NAME } from './tool-names.js'; +import { buildFilePathArgsPattern } from '../policy/utils.js'; import { debugLogger } from '../utils/debugLogger.js'; import { LS_DEFINITION } from './definitions/coreTools.js'; import { resolveToolDeclaration } from './definitions/resolver.js'; @@ -123,6 +126,14 @@ class LSToolInvocation extends BaseToolInvocation { return shortenPath(relativePath); } + override getPolicyUpdateOptions( + _outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined { + return { + argsPattern: buildFilePathArgsPattern(this.params.dir_path), + }; + } + // Helper for consistent error formatting private errorResult( llmContent: string, diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index f67d1f9bea..523eac62ad 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -184,7 +184,7 @@ export class DiscoveredMCPToolInvocation extends BaseToolInvocation< ); } - protected override getPolicyUpdateOptions( + override getPolicyUpdateOptions( _outcome: ToolConfirmationOutcome, ): PolicyUpdateOptions | undefined { return { mcpName: this.serverName }; diff --git a/packages/core/src/tools/read-file.ts b/packages/core/src/tools/read-file.ts index 0f044a4998..a5145c399d 100644 --- a/packages/core/src/tools/read-file.ts +++ b/packages/core/src/tools/read-file.ts @@ -14,8 +14,11 @@ import { type ToolInvocation, type ToolLocation, type ToolResult, + type PolicyUpdateOptions, + type ToolConfirmationOutcome, } from './tools.js'; import { ToolErrorType } from './tool-error.js'; +import { buildFilePathArgsPattern } from '../policy/utils.js'; import type { PartUnion } from '@google/genai'; import { @@ -88,6 +91,14 @@ class ReadFileToolInvocation extends BaseToolInvocation< ]; } + override getPolicyUpdateOptions( + _outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined { + return { + argsPattern: buildFilePathArgsPattern(this.params.file_path), + }; + } + async execute(): Promise { const validationError = this.config.validatePathAccess( this.resolvedPath, diff --git a/packages/core/src/tools/read-many-files.ts b/packages/core/src/tools/read-many-files.ts index c9c4e230e6..4a2ae9a4c0 100644 --- a/packages/core/src/tools/read-many-files.ts +++ b/packages/core/src/tools/read-many-files.ts @@ -11,11 +11,14 @@ import { Kind, type ToolInvocation, type ToolResult, + type PolicyUpdateOptions, + type ToolConfirmationOutcome, } from './tools.js'; import { getErrorMessage } from '../utils/errors.js'; import * as fsPromises from 'node:fs/promises'; import * as path from 'node:path'; import { glob, escape } from 'glob'; +import { buildPatternArgsPattern } from '../policy/utils.js'; import { detectFileType, processSingleFileContent, @@ -155,6 +158,16 @@ ${finalExclusionPatternsForDescription )}".`; } + override getPolicyUpdateOptions( + _outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined { + // We join the include patterns to match the JSON stringified arguments. + // buildPatternArgsPattern handles JSON stringification. + return { + argsPattern: buildPatternArgsPattern(JSON.stringify(this.params.include)), + }; + } + async execute(signal: AbortSignal): Promise { const { include, exclude = [], useDefaultExcludes = true } = this.params; diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 4ea83b0af4..a1bef189b5 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -90,7 +90,7 @@ export class ShellToolInvocation extends BaseToolInvocation< return description; } - protected override getPolicyUpdateOptions( + override getPolicyUpdateOptions( outcome: ToolConfirmationOutcome, ): PolicyUpdateOptions | undefined { if ( diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index fcdcbd6df6..38a868d665 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -154,12 +154,22 @@ export const LS_TOOL_NAME_LEGACY = 'list_directory'; // Just to be safe if anyth export const EDIT_TOOL_NAMES = new Set([EDIT_TOOL_NAME, WRITE_FILE_TOOL_NAME]); -// Tool Display Names -export const WRITE_FILE_DISPLAY_NAME = 'WriteFile'; -export const EDIT_DISPLAY_NAME = 'Edit'; -export const ASK_USER_DISPLAY_NAME = 'Ask User'; -export const READ_FILE_DISPLAY_NAME = 'ReadFile'; -export const GLOB_DISPLAY_NAME = 'FindFiles'; +/** + * Tools that can access local files or remote resources and should be + * treated with extra caution when updating policies. + */ +export const SENSITIVE_TOOLS = new Set([ + GLOB_TOOL_NAME, + GREP_TOOL_NAME, + READ_MANY_FILES_TOOL_NAME, + WEB_FETCH_TOOL_NAME, + READ_FILE_TOOL_NAME, + LS_TOOL_NAME, + WRITE_FILE_TOOL_NAME, + EDIT_TOOL_NAME, + SHELL_TOOL_NAME, +]); + export const TRACKER_CREATE_TASK_TOOL_NAME = 'tracker_create_task'; export const TRACKER_UPDATE_TASK_TOOL_NAME = 'tracker_update_task'; export const TRACKER_GET_TASK_TOOL_NAME = 'tracker_get_task'; @@ -167,6 +177,13 @@ export const TRACKER_LIST_TASKS_TOOL_NAME = 'tracker_list_tasks'; export const TRACKER_ADD_DEPENDENCY_TOOL_NAME = 'tracker_add_dependency'; export const TRACKER_VISUALIZE_TOOL_NAME = 'tracker_visualize'; +// Tool Display Names +export const WRITE_FILE_DISPLAY_NAME = 'WriteFile'; +export const EDIT_DISPLAY_NAME = 'Edit'; +export const ASK_USER_DISPLAY_NAME = 'Ask User'; +export const READ_FILE_DISPLAY_NAME = 'ReadFile'; +export const GLOB_DISPLAY_NAME = 'FindFiles'; + /** * Mapping of legacy tool names to their current names. * This ensures backward compatibility for user-defined policies, skills, and hooks. diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index 0a82cc1510..828461ea65 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -68,12 +68,21 @@ export interface ToolInvocation< updateOutput?: (output: ToolLiveOutput) => void, shellExecutionConfig?: ShellExecutionConfig, ): Promise; + + /** + * Returns tool-specific options for policy updates. + * This is used by the scheduler to narrow policy rules when a tool is approved. + */ + getPolicyUpdateOptions?( + outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined; } /** * Options for policy updates that can be customized by tool invocations. */ export interface PolicyUpdateOptions { + argsPattern?: string; commandPrefix?: string | string[]; mcpName?: string; } @@ -130,7 +139,7 @@ export abstract class BaseToolInvocation< * Subclasses can override this to provide additional options like * commandPrefix (for shell) or mcpName (for MCP tools). */ - protected getPolicyUpdateOptions( + getPolicyUpdateOptions( _outcome: ToolConfirmationOutcome, ): PolicyUpdateOptions | undefined { return undefined; diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts index 3170227188..50960a9f7f 100644 --- a/packages/core/src/tools/web-fetch.ts +++ b/packages/core/src/tools/web-fetch.ts @@ -12,7 +12,9 @@ import { type ToolInvocation, type ToolResult, type ToolConfirmationOutcome, + type PolicyUpdateOptions, } from './tools.js'; +import { buildPatternArgsPattern } from '../policy/utils.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { ToolErrorType } from './tool-error.js'; import { getErrorMessage } from '../utils/errors.js'; @@ -291,6 +293,22 @@ ${textContent} return `Processing URLs and instructions from prompt: "${displayPrompt}"`; } + override getPolicyUpdateOptions( + _outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined { + if (this.params.url) { + return { + argsPattern: buildPatternArgsPattern(this.params.url), + }; + } + if (this.params.prompt) { + return { + argsPattern: buildPatternArgsPattern(this.params.prompt), + }; + } + return undefined; + } + protected override async getConfirmationDetails( _abortSignal: AbortSignal, ): Promise { diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index 8ec660b661..4c0a533689 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -24,7 +24,9 @@ import { type ToolLocation, type ToolResult, type ToolConfirmationOutcome, + type PolicyUpdateOptions, } from './tools.js'; +import { buildFilePathArgsPattern } from '../policy/utils.js'; import { ToolErrorType } from './tool-error.js'; import { makeRelative, shortenPath } from '../utils/paths.js'; import { getErrorMessage, isNodeError } from '../utils/errors.js'; @@ -164,6 +166,14 @@ class WriteFileToolInvocation extends BaseToolInvocation< return [{ path: this.resolvedPath }]; } + override getPolicyUpdateOptions( + _outcome: ToolConfirmationOutcome, + ): PolicyUpdateOptions | undefined { + return { + argsPattern: buildFilePathArgsPattern(this.params.file_path), + }; + } + override getDescription(): string { const relativePath = makeRelative( this.resolvedPath, diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index 280ad18db5..adfb1044b6 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -1461,6 +1461,13 @@ "default": false, "type": "boolean" }, + "autoAddToPolicyByDefault": { + "title": "Auto-add to Policy by Default", + "description": "When enabled, the \"Allow for all future sessions\" option becomes the default choice for low-risk tools in trusted workspaces.", + "markdownDescription": "When enabled, the \"Allow for all future sessions\" option becomes the default choice for low-risk tools in trusted workspaces.\n\n- Category: `Security`\n- Requires restart: `no`\n- Default: `false`", + "default": false, + "type": "boolean" + }, "blockGitExtensions": { "title": "Blocks extensions from Git", "description": "Blocks installing and loading extensions from Git.",