diff --git a/packages/cli/src/config/policy-engine.integration.test.ts b/packages/cli/src/config/policy-engine.integration.test.ts index 6847865434..1d7573337e 100644 --- a/packages/cli/src/config/policy-engine.integration.test.ts +++ b/packages/cli/src/config/policy-engine.integration.test.ts @@ -352,6 +352,38 @@ describe('Policy Engine Integration Tests', () => { ).toBe(PolicyDecision.DENY); }); + it('should correctly match tool annotations', async () => { + const settings: Settings = {}; + + const config = await createPolicyEngineConfig( + settings, + ApprovalMode.DEFAULT, + ); + + // Add a manual rule with annotations to the config + config.rules = config.rules || []; + config.rules.push({ + toolAnnotations: { readOnlyHint: true }, + decision: PolicyDecision.ALLOW, + priority: 10, + }); + + const engine = new PolicyEngine(config); + + // A tool with readOnlyHint=true should be ALLOWED + const roCall = { name: 'some_tool', args: {} }; + const roMeta = { readOnlyHint: true }; + expect((await engine.check(roCall, undefined, roMeta)).decision).toBe( + PolicyDecision.ALLOW, + ); + + // A tool without the hint (or with false) should follow default decision (ASK_USER) + const rwMeta = { readOnlyHint: false }; + expect((await engine.check(roCall, undefined, rwMeta)).decision).toBe( + PolicyDecision.ASK_USER, + ); + }); + describe.each(['write_file', 'replace'])( 'Plan Mode policy for %s', (toolName) => { diff --git a/packages/core/src/policy/policy-engine.test.ts b/packages/core/src/policy/policy-engine.test.ts index 121dfb7c0c..b972ce0e8f 100644 --- a/packages/core/src/policy/policy-engine.test.ts +++ b/packages/core/src/policy/policy-engine.test.ts @@ -2557,4 +2557,68 @@ describe('PolicyEngine', () => { expect(checkers[0].priority).toBe(2.5); }); }); + + describe('Tool Annotations', () => { + it('should match tools by semantic annotations', async () => { + engine = new PolicyEngine({ + rules: [ + { + toolAnnotations: { readOnlyHint: true }, + decision: PolicyDecision.ALLOW, + priority: 10, + }, + ], + defaultDecision: PolicyDecision.DENY, + }); + + const readOnlyTool = { name: 'read', args: {} }; + const readOnlyMeta = { readOnlyHint: true, extra: 'info' }; + + const writeTool = { name: 'write', args: {} }; + const writeMeta = { readOnlyHint: false }; + + expect( + (await engine.check(readOnlyTool, undefined, readOnlyMeta)).decision, + ).toBe(PolicyDecision.ALLOW); + expect( + (await engine.check(writeTool, undefined, writeMeta)).decision, + ).toBe(PolicyDecision.DENY); + expect((await engine.check(writeTool, undefined, {})).decision).toBe( + PolicyDecision.DENY, + ); + }); + + it('should support scoped annotation rules', async () => { + engine = new PolicyEngine({ + rules: [ + { + toolName: '*__*', + toolAnnotations: { experimental: true }, + decision: PolicyDecision.DENY, + priority: 20, + }, + { + toolName: '*__*', + decision: PolicyDecision.ALLOW, + priority: 10, + }, + ], + }); + + expect( + ( + await engine.check({ name: 'mcp__test' }, 'mcp', { + experimental: true, + }) + ).decision, + ).toBe(PolicyDecision.DENY); + expect( + ( + await engine.check({ name: 'mcp__stable' }, 'mcp', { + experimental: false, + }) + ).decision, + ).toBe(PolicyDecision.ALLOW); + }); + }); }); diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index 0998ccb2b5..1b4c976f89 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -102,6 +102,7 @@ function ruleMatches( stringifiedArgs: string | undefined, serverName: string | undefined, currentApprovalMode: ApprovalMode, + toolAnnotations?: Record, ): boolean { // Check if rule applies to current approval mode if (rule.modes && rule.modes.length > 0) { @@ -124,6 +125,18 @@ function ruleMatches( } } + // Check annotations if specified + if (rule.toolAnnotations) { + if (!toolAnnotations) { + return false; + } + for (const [key, value] of Object.entries(rule.toolAnnotations)) { + if (toolAnnotations[key] !== value) { + return false; + } + } + } + // Check args pattern if specified if (rule.argsPattern) { // If rule has an args pattern but tool has no args, no match @@ -204,6 +217,7 @@ export class PolicyEngine { dir_path: string | undefined, allowRedirection?: boolean, rule?: PolicyRule, + toolAnnotations?: Record, ): Promise { if (!command) { return { @@ -294,6 +308,7 @@ export class PolicyEngine { const subResult = await this.check( { name: toolName, args: { command: subCmd, dir_path } }, serverName, + toolAnnotations, ); // subResult.decision is already filtered through applyNonInteractiveMode by this.check() @@ -351,6 +366,7 @@ export class PolicyEngine { async check( toolCall: FunctionCall, serverName: string | undefined, + toolAnnotations?: Record, ): Promise { let stringifiedArgs: string | undefined; // Compute stringified args once before the loop @@ -403,7 +419,14 @@ export class PolicyEngine { for (const rule of this.rules) { const match = toolCallsToTry.some((tc) => - ruleMatches(rule, tc, stringifiedArgs, serverName, this.approvalMode), + ruleMatches( + rule, + tc, + stringifiedArgs, + serverName, + this.approvalMode, + toolAnnotations, + ), ); if (match) { @@ -420,6 +443,7 @@ export class PolicyEngine { shellDirPath, rule.allowRedirection, rule, + toolAnnotations, ); decision = shellResult.decision; if (shellResult.rule) { @@ -446,6 +470,9 @@ export class PolicyEngine { this.defaultDecision, serverName, shellDirPath, + undefined, + undefined, + toolAnnotations, ); decision = shellResult.decision; matchedRule = shellResult.rule; @@ -464,6 +491,7 @@ export class PolicyEngine { stringifiedArgs, serverName, this.approvalMode, + toolAnnotations, ) ) { debugLogger.debug( diff --git a/packages/core/src/policy/toml-loader.test.ts b/packages/core/src/policy/toml-loader.test.ts index 1ea83775fe..785d56cf3e 100644 --- a/packages/core/src/policy/toml-loader.test.ts +++ b/packages/core/src/policy/toml-loader.test.ts @@ -89,6 +89,24 @@ priority = 100 expect(result.errors).toHaveLength(0); }); + it('should parse toolAnnotations from TOML', async () => { + const result = await runLoadPoliciesFromToml(` +[[rule]] +toolName = "annotated-tool" +toolAnnotations = { readOnlyHint = true, custom = "value" } +decision = "allow" +priority = 70 +`); + + expect(result.rules).toHaveLength(1); + expect(result.rules[0].toolName).toBe('annotated-tool'); + expect(result.rules[0].toolAnnotations).toEqual({ + readOnlyHint: true, + custom: 'value', + }); + expect(result.errors).toHaveLength(0); + }); + it('should transform mcpName = "*" to wildcard toolName', async () => { const result = await runLoadPoliciesFromToml(` [[rule]] diff --git a/packages/core/src/policy/toml-loader.ts b/packages/core/src/policy/toml-loader.ts index 7be3fe27dc..6b164d59b8 100644 --- a/packages/core/src/policy/toml-loader.ts +++ b/packages/core/src/policy/toml-loader.ts @@ -46,6 +46,7 @@ const PolicyRuleSchema = z.object({ 'priority must be <= 999 to prevent tier overflow. Priorities >= 1000 would jump to the next tier.', }), modes: z.array(z.nativeEnum(ApprovalMode)).optional(), + toolAnnotations: z.record(z.any()).optional(), allow_redirection: z.boolean().optional(), deny_message: z.string().optional(), }); @@ -61,6 +62,7 @@ const SafetyCheckerRuleSchema = z.object({ commandRegex: z.string().optional(), priority: z.number().int().default(0), modes: z.array(z.nativeEnum(ApprovalMode)).optional(), + toolAnnotations: z.record(z.any()).optional(), checker: z.discriminatedUnion('type', [ z.object({ type: z.literal('in-process'), @@ -383,6 +385,7 @@ export async function loadPoliciesFromToml( decision: rule.decision, priority: transformPriority(rule.priority, tier), modes: rule.modes, + toolAnnotations: rule.toolAnnotations, allowRedirection: rule.allow_redirection, source: `${tierName.charAt(0).toUpperCase() + tierName.slice(1)}: ${file}`, denyMessage: rule.deny_message, @@ -467,6 +470,7 @@ export async function loadPoliciesFromToml( // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion checker: checker.checker as SafetyCheckerConfig, modes: checker.modes, + toolAnnotations: checker.toolAnnotations, source: `${tierName.charAt(0).toUpperCase() + tierName.slice(1)}: ${file}`, }; diff --git a/packages/core/src/policy/types.ts b/packages/core/src/policy/types.ts index e8aa0e6dd1..79bf74f408 100644 --- a/packages/core/src/policy/types.ts +++ b/packages/core/src/policy/types.ts @@ -115,6 +115,12 @@ export interface PolicyRule { */ argsPattern?: RegExp; + /** + * Metadata annotations provided by the tool (e.g. readOnlyHint). + * All keys and values in this record must match the tool's annotations. + */ + toolAnnotations?: Record; + /** * The decision to make when this rule matches. */ @@ -165,6 +171,12 @@ export interface SafetyCheckerRule { */ argsPattern?: RegExp; + /** + * Metadata annotations provided by the tool (e.g. readOnlyHint). + * All keys and values in this record must match the tool's annotations. + */ + toolAnnotations?: Record; + /** * Priority of this checker. Higher numbers run first. * Default is 0.