feat(policy): support auto-add to policy by default and scoped persistence

This commit is contained in:
Spencer
2026-03-04 04:01:23 +00:00
parent ab64b15d51
commit 211d996289
17 changed files with 732 additions and 270 deletions
+33 -26
View File
@@ -549,6 +549,7 @@ export interface ConfigParameters {
truncateToolOutputThreshold?: number;
eventEmitter?: EventEmitter;
useWriteTodos?: boolean;
workspacePoliciesDir?: string;
policyEngineConfig?: PolicyEngineConfig;
directWebFetch?: boolean;
policyUpdateConfirmationRequest?: PolicyUpdateConfirmationRequest;
@@ -742,6 +743,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:
@@ -952,6 +954,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 ?? [];
@@ -1183,7 +1186,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),
);
}
}
@@ -1995,6 +1998,10 @@ export class Config implements McpContext, AgentLoopContext {
return this.geminiMdFilePaths;
}
getWorkspacePoliciesDir(): string | undefined {
return this.workspacePoliciesDir;
}
setGeminiMdFilePaths(paths: string[]): void {
this.geminiMdFilePaths = paths;
}
@@ -2597,7 +2604,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);
@@ -2781,7 +2788,7 @@ export class Config implements McpContext, AgentLoopContext {
}
async createToolRegistry(): Promise<ToolRegistry> {
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 = (
@@ -2811,10 +2818,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()) {
@@ -2827,85 +2834,85 @@ 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),
new TrackerCreateTaskTool(this, this.messageBus),
),
);
maybeRegister(TrackerUpdateTaskTool, () =>
registry.registerTool(
new TrackerUpdateTaskTool(this, this._messageBus),
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)),
);
}
+7
View File
@@ -168,6 +168,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);
}
@@ -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;
+13 -1
View File
@@ -530,9 +530,21 @@ 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 });
// Backup existing file if it exists
try {
await fs.copyFile(policyFile, `${policyFile}.bak`);
} catch (error) {
if (!isNodeError(error) || error.code !== 'ENOENT') {
debugLogger.warn(`Failed to backup ${policyFile}`, error);
}
}
// Read existing file
let existingData: { rule?: TomlRule[] } = {};
try {
+76 -4
View File
@@ -230,15 +230,87 @@ describe('createPolicyUpdater', () => {
// 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,
);
(fs.mkdir as unknown as Mock).mockResolvedValue(undefined);
(fs.readFile as unknown as Mock).mockRejectedValue(
new Error('File not found'),
);
(fs.copyFile as unknown as Mock).mockResolvedValue(undefined);
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);
await messageBus.publish({
type: MessageBusType.UPDATE_POLICY,
toolName: 'test_tool',
persist: true,
persistScope: 'workspace',
});
await new Promise((resolve) => setTimeout(resolve, 0));
expect(mockStorage.getWorkspaceAutoSavedPolicyPath).toHaveBeenCalled();
expect(fs.mkdir).toHaveBeenCalledWith(workspacePoliciesDir, {
recursive: true,
});
expect(fs.rename).toHaveBeenCalledWith(
expect.stringMatching(/\.tmp$/),
policyFile,
);
});
it('should backup existing policy file before writing', async () => {
createPolicyUpdater(policyEngine, messageBus, mockStorage);
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).mockResolvedValue(
'[[rule]]\ntoolName = "existing"',
);
(fs.copyFile as unknown as Mock).mockResolvedValue(undefined);
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);
await messageBus.publish({
type: MessageBusType.UPDATE_POLICY,
toolName: 'new_tool',
persist: true,
});
await new Promise((resolve) => setTimeout(resolve, 0));
expect(fs.copyFile).toHaveBeenCalledWith(policyFile, `${policyFile}.bak`);
});
});
+12
View File
@@ -82,3 +82,15 @@ 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":"<path>" in a JSON string.
*/
export function buildFilePathArgsPattern(filePath: string): string {
const jsonPath = JSON.stringify(filePath).slice(1, -1);
return `"file_path":"${escapeRegex(jsonPath)}"`;
}
+94 -1
View File
@@ -16,7 +16,10 @@ 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 {
ToolConfirmationOutcome,
@@ -219,6 +222,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<Config>;
@@ -453,6 +458,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<Config>;
@@ -487,6 +494,92 @@ 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<Config>;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked<MessageBus>;
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<Config>;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked<MessageBus>;
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),
setApprovalMode: vi.fn(),
} as unknown as Mocked<Config>;
const mockMessageBus = {
publish: vi.fn(),
} as unknown as Mocked<MessageBus>;
const tool = { name: 'write_file' } as AnyDeclarativeTool;
const details = {
type: 'edit',
filePath: 'src/foo.ts',
title: 'Edit',
onConfirm: vi.fn(),
} as unknown as SerializableConfirmationDetails;
await updatePolicy(
tool,
ToolConfirmationOutcome.ProceedAlwaysAndSave,
details,
{ config: mockConfig, messageBus: mockMessageBus },
);
expect(mockMessageBus.publish).toHaveBeenCalledWith(
expect.objectContaining({
toolName: 'write_file',
argsPattern: '"file_path":"src/foo\\.ts"',
}),
);
});
});
describe('getPolicyDenialError', () => {
+25
View File
@@ -22,6 +22,7 @@ import {
type AnyDeclarativeTool,
type PolicyUpdateOptions,
} from '../tools/tools.js';
import { buildFilePathArgsPattern } from '../policy/utils.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { EDIT_TOOL_NAMES } from '../tools/tool-names.js';
import type { ValidatingToolCall } from './types.js';
@@ -102,6 +103,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()
) {
persistScope = 'workspace';
} else {
persistScope = 'user';
}
}
// Specialized Tools (MCP)
if (confirmationDetails?.type === 'mcp') {
await handleMcpPolicyUpdate(
@@ -109,6 +124,7 @@ export async function updatePolicy(
outcome,
confirmationDetails,
deps.messageBus,
persistScope,
);
return;
}
@@ -119,6 +135,7 @@ export async function updatePolicy(
outcome,
confirmationDetails,
deps.messageBus,
persistScope,
);
}
@@ -148,6 +165,7 @@ async function handleStandardPolicyUpdate(
outcome: ToolConfirmationOutcome,
confirmationDetails: SerializableConfirmationDetails | undefined,
messageBus: MessageBus,
persistScope?: 'workspace' | 'user',
): Promise<void> {
if (
outcome === ToolConfirmationOutcome.ProceedAlways ||
@@ -157,12 +175,17 @@ async function handleStandardPolicyUpdate(
if (confirmationDetails?.type === 'exec') {
options.commandPrefix = confirmationDetails.rootCommands;
} else if (confirmationDetails?.type === 'edit') {
options.argsPattern = buildFilePathArgsPattern(
confirmationDetails.filePath,
);
}
await messageBus.publish({
type: MessageBusType.UPDATE_POLICY,
toolName: tool.name,
persist: outcome === ToolConfirmationOutcome.ProceedAlwaysAndSave,
persistScope,
...options,
});
}
@@ -180,6 +203,7 @@ async function handleMcpPolicyUpdate(
{ type: 'mcp' }
>,
messageBus: MessageBus,
persistScope?: 'workspace' | 'user',
): Promise<void> {
const isMcpAlways =
outcome === ToolConfirmationOutcome.ProceedAlways ||
@@ -204,5 +228,6 @@ async function handleMcpPolicyUpdate(
toolName,
mcpName: confirmationDetails.serverName,
persist,
persistScope,
});
}
+10
View File
@@ -20,7 +20,9 @@ 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';
@@ -456,6 +458,14 @@ class EditToolInvocation
return [{ path: this.params.file_path }];
}
protected override getPolicyUpdateOptions(
_outcome: ToolConfirmationOutcome,
): PolicyUpdateOptions | undefined {
return {
argsPattern: buildFilePathArgsPattern(this.params.file_path),
};
}
private async attemptSelfCorrection(
params: EditToolParams,
currentContent: string,
+1
View File
@@ -74,6 +74,7 @@ export interface ToolInvocation<
* Options for policy updates that can be customized by tool invocations.
*/
export interface PolicyUpdateOptions {
argsPattern?: string;
commandPrefix?: string | string[];
mcpName?: string;
}
+10
View File
@@ -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 }];
}
protected override getPolicyUpdateOptions(
_outcome: ToolConfirmationOutcome,
): PolicyUpdateOptions | undefined {
return {
argsPattern: buildFilePathArgsPattern(this.params.file_path),
};
}
override getDescription(): string {
const relativePath = makeRelative(
this.resolvedPath,