From 1625e4f530a05091a7546980d9b2824708b94671 Mon Sep 17 00:00:00 2001 From: Spencer Date: Wed, 4 Mar 2026 20:36:08 +0000 Subject: [PATCH] fix(policy): address review feedback for robust path correction, priority consistency, and narrowing enforcement --- packages/core/src/policy/config.ts | 32 +++++++++++++++-- packages/core/src/policy/persistence.test.ts | 6 ++-- packages/core/src/policy/utils.ts | 23 +++++++------ packages/core/src/scheduler/policy.test.ts | 1 + packages/core/src/scheduler/policy.ts | 12 ++++--- packages/core/src/tools/edit.ts | 36 +++++++++++++++----- 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/read-file.ts | 11 ++++++ packages/core/src/tools/read-many-files.ts | 13 +++++++ packages/core/src/tools/tool-names.ts | 29 ++++++++++++---- packages/core/src/tools/web-fetch.ts | 18 ++++++++++ 13 files changed, 178 insertions(+), 36 deletions(-) diff --git a/packages/core/src/policy/config.ts b/packages/core/src/policy/config.ts index 83340c4261..fc52ce8427 100644 --- a/packages/core/src/policy/config.ts +++ b/packages/core/src/policy/config.ts @@ -34,7 +34,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'; @@ -491,6 +491,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 + (ALWAYS_ALLOW_PRIORITY % 1); + + 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, @@ -498,7 +511,7 @@ export function createPolicyUpdater( policyEngine.addRule({ toolName, decision: PolicyDecision.ALLOW, - priority: ALWAYS_ALLOW_PRIORITY, + priority, argsPattern: new RegExp(pattern), source: 'Dynamic (Confirmed)', }); @@ -517,10 +530,23 @@ export function createPolicyUpdater( ? new RegExp(message.argsPattern) : undefined; + const tier = + message.persistScope === 'user' + ? USER_POLICY_TIER + : WORKSPACE_POLICY_TIER; + const priority = tier + (ALWAYS_ALLOW_PRIORITY % 1); + + 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)', }); diff --git a/packages/core/src/policy/persistence.test.ts b/packages/core/src/policy/persistence.test.ts index 8b3667c5a2..fd445afc40 100644 --- a/packages/core/src/policy/persistence.test.ts +++ b/packages/core/src/policy/persistence.test.ts @@ -7,7 +7,7 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import * as path from 'node:path'; import { createPolicyUpdater } from './config.js'; -import { ALWAYS_ALLOW_PRIORITY } from './utils.js'; +import { getAlwaysAllowPriorityFraction } from './utils.js'; import { PolicyEngine } from './policy-engine.js'; import { MessageBus } from '../confirmation-bus/message-bus.js'; import { MessageBusType } from '../confirmation-bus/types.js'; @@ -62,9 +62,7 @@ describe('createPolicyUpdater', () => { const content = memfs.readFileSync(policyFile, 'utf-8') as string; expect(content).toContain('toolName = "test_tool"'); expect(content).toContain('decision = "allow"'); - const expectedPriority = Math.round( - (ALWAYS_ALLOW_PRIORITY - Math.floor(ALWAYS_ALLOW_PRIORITY)) * 1000, - ); + const expectedPriority = getAlwaysAllowPriorityFraction(); expect(content).toContain(`priority = ${expectedPriority}`); }); diff --git a/packages/core/src/policy/utils.ts b/packages/core/src/policy/utils.ts index 6dd4f19031..f79e865700 100644 --- a/packages/core/src/policy/utils.ts +++ b/packages/core/src/policy/utils.ts @@ -10,17 +10,6 @@ */ export const ALWAYS_ALLOW_PRIORITY = 3.95; -/** - * Calculates a unique priority within the ALWAYS_ALLOW_PRIORITY tier. - * It uses the fractional part as a base and adds a small offset. - */ -export function getAlwaysAllowPriority(offset: number): number { - const base = Math.floor(ALWAYS_ALLOW_PRIORITY); - const fraction = ALWAYS_ALLOW_PRIORITY - base; - // Use a precision of 3 decimal places for the offset - return base + fraction + offset / 1000; -} - /** * Returns the fractional priority of ALWAYS_ALLOW_PRIORITY scaled to 1000. */ @@ -120,3 +109,15 @@ export function buildFilePathArgsPattern(filePath: string): string { const jsonPath = JSON.stringify(filePath).slice(1, -1); return `"file_path":"${escapeRegex(jsonPath)}"`; } + +/** + * 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 jsonPattern = JSON.stringify(pattern).slice(1, -1); + return `"pattern":"${escapeRegex(jsonPattern)}"`; +} diff --git a/packages/core/src/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts index af5883bd0d..82691c6c37 100644 --- a/packages/core/src/scheduler/policy.test.ts +++ b/packages/core/src/scheduler/policy.test.ts @@ -553,6 +553,7 @@ describe('policy.ts', () => { 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 = { diff --git a/packages/core/src/scheduler/policy.ts b/packages/core/src/scheduler/policy.ts index 59459b6214..1ac70a108b 100644 --- a/packages/core/src/scheduler/policy.ts +++ b/packages/core/src/scheduler/policy.ts @@ -24,6 +24,7 @@ import { 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'; @@ -114,7 +115,7 @@ export async function updatePolicy( // If folder is trusted and workspace policies are enabled, we prefer workspace scope. if ( deps.config.isTrustedFolder() && - deps.config.getWorkspacePoliciesDir() + deps.config.getWorkspacePoliciesDir() !== undefined ) { persistScope = 'workspace'; } else { @@ -142,6 +143,7 @@ export async function updatePolicy( deps.messageBus, persistScope, deps.toolInvocation, + deps.config, ); } @@ -173,6 +175,7 @@ async function handleStandardPolicyUpdate( messageBus: MessageBus, persistScope?: 'workspace' | 'user', toolInvocation?: AnyToolInvocation, + config?: Config, ): Promise { if ( outcome === ToolConfirmationOutcome.ProceedAlways || @@ -184,9 +187,10 @@ async function handleStandardPolicyUpdate( if (!options.commandPrefix && confirmationDetails?.type === 'exec') { options.commandPrefix = confirmationDetails.rootCommands; } else if (!options.argsPattern && confirmationDetails?.type === 'edit') { - options.argsPattern = buildFilePathArgsPattern( - confirmationDetails.filePath, - ); + const filePath = config + ? makeRelative(confirmationDetails.filePath, config.getTargetDir()) + : confirmationDetails.filePath; + options.argsPattern = buildFilePathArgsPattern(filePath); } await messageBus.publish({ diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index 75c0709e72..06f9657745 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -27,6 +27,7 @@ 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'; @@ -453,10 +454,19 @@ class EditToolInvocation displayName?: string, ) { super(params, messageBus, toolName, displayName); - this.resolvedPath = path.resolve( - this.config.getTargetDir(), - this.params.file_path, - ); + 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[] { @@ -996,10 +1006,20 @@ export class EditTool return "The 'file_path' parameter must be non-empty."; } - const resolvedPath = path.resolve( - this.config.getTargetDir(), - params.file_path, - ); + 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, + ); + } + } else { + resolvedPath = params.file_path; + } const newPlaceholders = detectOmissionPlaceholders(params.new_string); if (newPlaceholders.length > 0) { 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/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/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/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 {