From e35ae5425547abb492415f604378692795c89569 Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Fri, 5 Sep 2025 15:36:05 -0700 Subject: [PATCH 01/10] feat(core): implement Tool Confirmation Message Bus foundation (#7231) Introduces a pub/sub message bus architecture to decouple tool confirmation logic from core tool implementation. This is the first PR in a series to implement the Tool Confirmation Message Bus as described in issue #7231. ## Changes - Add MessageBus class with event-driven publish/subscribe pattern - Implement PolicyEngine with configurable rules for tool execution control - Integrate MessageBus and PolicyEngine into Config class - Add comprehensive unit tests for both components ## Features - Priority-based policy rules with tool name and args pattern matching - Non-interactive mode support (ASK_USER becomes DENY) - Error handling and message validation - Support for multiple message types (confirmation, rejection, execution) ## Testing - 25 new unit tests covering all functionality - Full test suite passes (1964 total tests) - TypeScript compilation and linting checks pass --- packages/core/src/config/config.ts | 16 + packages/core/src/confirmation-bus/index.ts | 8 + .../src/confirmation-bus/message-bus.test.ts | 226 +++++++++++++++ .../core/src/confirmation-bus/message-bus.ts | 97 +++++++ packages/core/src/confirmation-bus/types.ts | 51 ++++ packages/core/src/policy/index.ts | 8 + .../core/src/policy/policy-engine.test.ts | 274 ++++++++++++++++++ packages/core/src/policy/policy-engine.ts | 98 +++++++ packages/core/src/policy/types.ts | 55 ++++ 9 files changed, 833 insertions(+) create mode 100644 packages/core/src/confirmation-bus/index.ts create mode 100644 packages/core/src/confirmation-bus/message-bus.test.ts create mode 100644 packages/core/src/confirmation-bus/message-bus.ts create mode 100644 packages/core/src/confirmation-bus/types.ts create mode 100644 packages/core/src/policy/index.ts create mode 100644 packages/core/src/policy/policy-engine.test.ts create mode 100644 packages/core/src/policy/policy-engine.ts create mode 100644 packages/core/src/policy/types.ts diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 543199274fb..2461ca2df12 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -57,6 +57,9 @@ import { WorkspaceContext } from '../utils/workspaceContext.js'; import { Storage } from './storage.js'; import { FileExclusions } from '../utils/ignorePatterns.js'; import type { EventEmitter } from 'node:events'; +import { MessageBus } from '../confirmation-bus/message-bus.js'; +import { PolicyEngine } from '../policy/policy-engine.js'; +import type { PolicyEngineConfig } from '../policy/types.js'; export enum ApprovalMode { DEFAULT = 'default', @@ -212,6 +215,7 @@ export interface ConfigParameters { enablePromptCompletion?: boolean; eventEmitter?: EventEmitter; useSmartEdit?: boolean; + policyEngineConfig?: PolicyEngineConfig; } export class Config { @@ -289,6 +293,8 @@ export class Config { private readonly fileExclusions: FileExclusions; private readonly eventEmitter?: EventEmitter; private readonly useSmartEdit: boolean; + private readonly messageBus: MessageBus; + private readonly policyEngine: PolicyEngine; constructor(params: ConfigParameters) { this.sessionId = params.sessionId; @@ -365,6 +371,8 @@ export class Config { this.enablePromptCompletion = params.enablePromptCompletion ?? false; this.fileExclusions = new FileExclusions(this); this.eventEmitter = params.eventEmitter; + this.policyEngine = new PolicyEngine(params.policyEngineConfig); + this.messageBus = new MessageBus(this.policyEngine); if (params.contextFileName) { setGeminiMdFilename(params.contextFileName); @@ -834,6 +842,14 @@ export class Config { return this.fileExclusions; } + getMessageBus(): MessageBus { + return this.messageBus; + } + + getPolicyEngine(): PolicyEngine { + return this.policyEngine; + } + async createToolRegistry(): Promise { const registry = new ToolRegistry(this, this.eventEmitter); diff --git a/packages/core/src/confirmation-bus/index.ts b/packages/core/src/confirmation-bus/index.ts new file mode 100644 index 00000000000..379d9aa4d84 --- /dev/null +++ b/packages/core/src/confirmation-bus/index.ts @@ -0,0 +1,8 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export * from './message-bus.js'; +export * from './types.js'; diff --git a/packages/core/src/confirmation-bus/message-bus.test.ts b/packages/core/src/confirmation-bus/message-bus.test.ts new file mode 100644 index 00000000000..e2d0bd2fe09 --- /dev/null +++ b/packages/core/src/confirmation-bus/message-bus.test.ts @@ -0,0 +1,226 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { MessageBus } from './message-bus.js'; +import { PolicyEngine } from '../policy/policy-engine.js'; +import { PolicyDecision } from '../policy/types.js'; +import { MessageBusType, type ToolConfirmationRequest } from './types.js'; + +describe('MessageBus', () => { + let messageBus: MessageBus; + let policyEngine: PolicyEngine; + + beforeEach(() => { + policyEngine = new PolicyEngine(); + messageBus = new MessageBus(policyEngine); + }); + + describe('publish', () => { + it('should emit error for invalid message', () => { + const errorHandler = vi.fn(); + messageBus.on('error', errorHandler); + + // @ts-expect-error - Testing invalid message + messageBus.publish({ invalid: 'message' }); + + expect(errorHandler).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining('Invalid message structure'), + }), + ); + }); + + it('should validate tool confirmation requests have correlationId', () => { + const errorHandler = vi.fn(); + messageBus.on('error', errorHandler); + + // @ts-expect-error - Testing missing correlationId + messageBus.publish({ + type: MessageBusType.TOOL_CONFIRMATION_REQUEST, + toolCall: { name: 'test' }, + }); + + expect(errorHandler).toHaveBeenCalled(); + }); + + it('should emit confirmation response when policy allows', () => { + vi.spyOn(policyEngine, 'check').mockReturnValue(PolicyDecision.ALLOW); + + const responseHandler = vi.fn(); + messageBus.subscribe( + MessageBusType.TOOL_CONFIRMATION_RESPONSE, + responseHandler, + ); + + const request: ToolConfirmationRequest = { + type: MessageBusType.TOOL_CONFIRMATION_REQUEST, + toolCall: { name: 'test-tool', args: {} }, + correlationId: '123', + }; + + messageBus.publish(request); + + expect(responseHandler).toHaveBeenCalledWith({ + type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, + correlationId: '123', + confirmed: true, + }); + }); + + it('should emit rejection and response when policy denies', () => { + vi.spyOn(policyEngine, 'check').mockReturnValue(PolicyDecision.DENY); + + const responseHandler = vi.fn(); + const rejectionHandler = vi.fn(); + messageBus.subscribe( + MessageBusType.TOOL_CONFIRMATION_RESPONSE, + responseHandler, + ); + messageBus.subscribe( + MessageBusType.TOOL_POLICY_REJECTION, + rejectionHandler, + ); + + const request: ToolConfirmationRequest = { + type: MessageBusType.TOOL_CONFIRMATION_REQUEST, + toolCall: { name: 'test-tool', args: {} }, + correlationId: '123', + }; + + messageBus.publish(request); + + expect(rejectionHandler).toHaveBeenCalledWith({ + type: MessageBusType.TOOL_POLICY_REJECTION, + toolCall: { name: 'test-tool', args: {} }, + }); + + expect(responseHandler).toHaveBeenCalledWith({ + type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, + correlationId: '123', + confirmed: false, + }); + }); + + it('should pass through to UI when policy says ASK_USER', () => { + vi.spyOn(policyEngine, 'check').mockReturnValue(PolicyDecision.ASK_USER); + + const requestHandler = vi.fn(); + messageBus.subscribe( + MessageBusType.TOOL_CONFIRMATION_REQUEST, + requestHandler, + ); + + const request: ToolConfirmationRequest = { + type: MessageBusType.TOOL_CONFIRMATION_REQUEST, + toolCall: { name: 'test-tool', args: {} }, + correlationId: '123', + }; + + messageBus.publish(request); + + expect(requestHandler).toHaveBeenCalledWith(request); + }); + + it('should emit other message types directly', () => { + const successHandler = vi.fn(); + messageBus.subscribe( + MessageBusType.TOOL_EXECUTION_SUCCESS, + successHandler, + ); + + const message = { + type: MessageBusType.TOOL_EXECUTION_SUCCESS as const, + toolCall: { name: 'test-tool' }, + result: 'success', + }; + + messageBus.publish(message); + + expect(successHandler).toHaveBeenCalledWith(message); + }); + }); + + describe('subscribe/unsubscribe', () => { + it('should allow subscribing to specific message types', () => { + const handler = vi.fn(); + messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler); + + const message = { + type: MessageBusType.TOOL_EXECUTION_SUCCESS as const, + toolCall: { name: 'test' }, + result: 'test', + }; + + messageBus.publish(message); + + expect(handler).toHaveBeenCalledWith(message); + }); + + it('should allow unsubscribing from message types', () => { + const handler = vi.fn(); + messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler); + messageBus.unsubscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler); + + const message = { + type: MessageBusType.TOOL_EXECUTION_SUCCESS as const, + toolCall: { name: 'test' }, + result: 'test', + }; + + messageBus.publish(message); + + expect(handler).not.toHaveBeenCalled(); + }); + + it('should support multiple subscribers for the same message type', () => { + const handler1 = vi.fn(); + const handler2 = vi.fn(); + + messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler1); + messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler2); + + const message = { + type: MessageBusType.TOOL_EXECUTION_SUCCESS as const, + toolCall: { name: 'test' }, + result: 'test', + }; + + messageBus.publish(message); + + expect(handler1).toHaveBeenCalledWith(message); + expect(handler2).toHaveBeenCalledWith(message); + }); + }); + + describe('error handling', () => { + it('should not crash on errors during message processing', () => { + const errorHandler = vi.fn(); + messageBus.on('error', errorHandler); + + // Mock policyEngine to throw an error + vi.spyOn(policyEngine, 'check').mockImplementation(() => { + throw new Error('Policy check failed'); + }); + + const request: ToolConfirmationRequest = { + type: MessageBusType.TOOL_CONFIRMATION_REQUEST, + toolCall: { name: 'test-tool' }, + correlationId: '123', + }; + + // Should not throw + expect(() => messageBus.publish(request)).not.toThrow(); + + // Should emit error + expect(errorHandler).toHaveBeenCalledWith( + expect.objectContaining({ + message: 'Policy check failed', + }), + ); + }); + }); +}); diff --git a/packages/core/src/confirmation-bus/message-bus.ts b/packages/core/src/confirmation-bus/message-bus.ts new file mode 100644 index 00000000000..40fd263d7d9 --- /dev/null +++ b/packages/core/src/confirmation-bus/message-bus.ts @@ -0,0 +1,97 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EventEmitter } from 'node:events'; +import type { PolicyEngine } from '../policy/policy-engine.js'; +import { PolicyDecision } from '../policy/types.js'; +import { MessageBusType, type Message } from './types.js'; + +export class MessageBus extends EventEmitter { + constructor(private readonly policyEngine: PolicyEngine) { + super(); + } + + private isValidMessage(message: Message): boolean { + if (!message || !message.type) { + return false; + } + + if ( + message.type === MessageBusType.TOOL_CONFIRMATION_REQUEST && + !('correlationId' in message) + ) { + return false; + } + + return true; + } + + private emitMessage(message: Message): void { + this.emit(message.type, message); + } + + publish(message: Message): void { + try { + if (!this.isValidMessage(message)) { + throw new Error( + `Invalid message structure: ${JSON.stringify(message)}`, + ); + } + + if (message.type === MessageBusType.TOOL_CONFIRMATION_REQUEST) { + const decision = this.policyEngine.check(message.toolCall); + + switch (decision) { + case PolicyDecision.ALLOW: + // Directly emit the response instead of recursive publish + this.emitMessage({ + type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, + correlationId: message.correlationId, + confirmed: true, + }); + break; + case PolicyDecision.DENY: + // Emit both rejection and response messages + this.emitMessage({ + type: MessageBusType.TOOL_POLICY_REJECTION, + toolCall: message.toolCall, + }); + this.emitMessage({ + type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, + correlationId: message.correlationId, + confirmed: false, + }); + break; + case PolicyDecision.ASK_USER: + // Pass through to UI for user confirmation + this.emitMessage(message); + break; + default: + throw new Error(`Unknown policy decision: ${decision}`); + } + } else { + // For all other message types, just emit them + this.emitMessage(message); + } + } catch (error) { + this.emit('error', error); + } + } + + subscribe( + type: T['type'], + listener: (message: T) => void, + ): void { + this.on(type, listener); + } + + unsubscribe( + type: T['type'], + listener: (message: T) => void, + ): void { + this.off(type, listener); + } +} diff --git a/packages/core/src/confirmation-bus/types.ts b/packages/core/src/confirmation-bus/types.ts new file mode 100644 index 00000000000..ea75a543d8c --- /dev/null +++ b/packages/core/src/confirmation-bus/types.ts @@ -0,0 +1,51 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { type FunctionCall } from '@google/genai'; + +export enum MessageBusType { + TOOL_CONFIRMATION_REQUEST = 'tool-confirmation-request', + TOOL_CONFIRMATION_RESPONSE = 'tool-confirmation-response', + TOOL_POLICY_REJECTION = 'tool-policy-rejection', + TOOL_EXECUTION_SUCCESS = 'tool-execution-success', + TOOL_EXECUTION_FAILURE = 'tool-execution-failure', +} + +export interface ToolConfirmationRequest { + type: MessageBusType.TOOL_CONFIRMATION_REQUEST; + toolCall: FunctionCall; + correlationId: string; +} + +export interface ToolConfirmationResponse { + type: MessageBusType.TOOL_CONFIRMATION_RESPONSE; + correlationId: string; + confirmed: boolean; +} + +export interface ToolPolicyRejection { + type: MessageBusType.TOOL_POLICY_REJECTION; + toolCall: FunctionCall; +} + +export interface ToolExecutionSuccess { + type: MessageBusType.TOOL_EXECUTION_SUCCESS; + toolCall: FunctionCall; + result: unknown; +} + +export interface ToolExecutionFailure { + type: MessageBusType.TOOL_EXECUTION_FAILURE; + toolCall: FunctionCall; + error: Error; +} + +export type Message = + | ToolConfirmationRequest + | ToolConfirmationResponse + | ToolPolicyRejection + | ToolExecutionSuccess + | ToolExecutionFailure; diff --git a/packages/core/src/policy/index.ts b/packages/core/src/policy/index.ts new file mode 100644 index 00000000000..e15309ca69b --- /dev/null +++ b/packages/core/src/policy/index.ts @@ -0,0 +1,8 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export * from './policy-engine.js'; +export * from './types.js'; diff --git a/packages/core/src/policy/policy-engine.test.ts b/packages/core/src/policy/policy-engine.test.ts new file mode 100644 index 00000000000..c734394f712 --- /dev/null +++ b/packages/core/src/policy/policy-engine.test.ts @@ -0,0 +1,274 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach } from 'vitest'; +import { PolicyEngine } from './policy-engine.js'; +import { + PolicyDecision, + type PolicyRule, + type PolicyEngineConfig, +} from './types.js'; +import type { FunctionCall } from '@google/genai'; + +describe('PolicyEngine', () => { + let engine: PolicyEngine; + + beforeEach(() => { + engine = new PolicyEngine(); + }); + + describe('constructor', () => { + it('should use default config when none provided', () => { + const decision = engine.check({ name: 'test' }); + expect(decision).toBe(PolicyDecision.ASK_USER); + }); + + it('should respect custom default decision', () => { + engine = new PolicyEngine({ defaultDecision: PolicyDecision.DENY }); + const decision = engine.check({ name: 'test' }); + expect(decision).toBe(PolicyDecision.DENY); + }); + + it('should sort rules by priority', () => { + const rules: PolicyRule[] = [ + { toolName: 'tool1', decision: PolicyDecision.DENY, priority: 1 }, + { toolName: 'tool2', decision: PolicyDecision.ALLOW, priority: 10 }, + { toolName: 'tool3', decision: PolicyDecision.ASK_USER, priority: 5 }, + ]; + + engine = new PolicyEngine({ rules }); + const sortedRules = engine.getRules(); + + expect(sortedRules[0].priority).toBe(10); + expect(sortedRules[1].priority).toBe(5); + expect(sortedRules[2].priority).toBe(1); + }); + }); + + describe('check', () => { + it('should match tool by name', () => { + const rules: PolicyRule[] = [ + { toolName: 'shell', decision: PolicyDecision.ALLOW }, + { toolName: 'edit', decision: PolicyDecision.DENY }, + ]; + + engine = new PolicyEngine({ rules }); + + expect(engine.check({ name: 'shell' })).toBe(PolicyDecision.ALLOW); + expect(engine.check({ name: 'edit' })).toBe(PolicyDecision.DENY); + expect(engine.check({ name: 'other' })).toBe(PolicyDecision.ASK_USER); + }); + + it('should match by args pattern', () => { + const rules: PolicyRule[] = [ + { + toolName: 'shell', + argsPattern: /rm -rf/, + decision: PolicyDecision.DENY, + }, + { + toolName: 'shell', + decision: PolicyDecision.ALLOW, + }, + ]; + + engine = new PolicyEngine({ rules }); + + const dangerousCall: FunctionCall = { + name: 'shell', + args: { command: 'rm -rf /' }, + }; + + const safeCall: FunctionCall = { + name: 'shell', + args: { command: 'ls -la' }, + }; + + expect(engine.check(dangerousCall)).toBe(PolicyDecision.DENY); + expect(engine.check(safeCall)).toBe(PolicyDecision.ALLOW); + }); + + it('should apply rules by priority', () => { + const rules: PolicyRule[] = [ + { toolName: 'shell', decision: PolicyDecision.DENY, priority: 1 }, + { toolName: 'shell', decision: PolicyDecision.ALLOW, priority: 10 }, + ]; + + engine = new PolicyEngine({ rules }); + + // Higher priority rule (ALLOW) should win + expect(engine.check({ name: 'shell' })).toBe(PolicyDecision.ALLOW); + }); + + it('should apply wildcard rules (no toolName)', () => { + const rules: PolicyRule[] = [ + { decision: PolicyDecision.DENY }, // Applies to all tools + { toolName: 'safe-tool', decision: PolicyDecision.ALLOW, priority: 10 }, + ]; + + engine = new PolicyEngine({ rules }); + + expect(engine.check({ name: 'safe-tool' })).toBe(PolicyDecision.ALLOW); + expect(engine.check({ name: 'any-other-tool' })).toBe( + PolicyDecision.DENY, + ); + }); + + it('should handle non-interactive mode', () => { + const config: PolicyEngineConfig = { + nonInteractive: true, + rules: [ + { toolName: 'interactive-tool', decision: PolicyDecision.ASK_USER }, + { toolName: 'allowed-tool', decision: PolicyDecision.ALLOW }, + ], + }; + + engine = new PolicyEngine(config); + + // ASK_USER should become DENY in non-interactive mode + expect(engine.check({ name: 'interactive-tool' })).toBe( + PolicyDecision.DENY, + ); + // ALLOW should remain ALLOW + expect(engine.check({ name: 'allowed-tool' })).toBe(PolicyDecision.ALLOW); + // Default ASK_USER should also become DENY + expect(engine.check({ name: 'unknown-tool' })).toBe(PolicyDecision.DENY); + }); + }); + + describe('addRule', () => { + it('should add a new rule and maintain priority order', () => { + engine.addRule({ + toolName: 'tool1', + decision: PolicyDecision.ALLOW, + priority: 5, + }); + engine.addRule({ + toolName: 'tool2', + decision: PolicyDecision.DENY, + priority: 10, + }); + engine.addRule({ + toolName: 'tool3', + decision: PolicyDecision.ASK_USER, + priority: 1, + }); + + const rules = engine.getRules(); + expect(rules).toHaveLength(3); + expect(rules[0].priority).toBe(10); + expect(rules[1].priority).toBe(5); + expect(rules[2].priority).toBe(1); + }); + + it('should apply newly added rules', () => { + expect(engine.check({ name: 'new-tool' })).toBe(PolicyDecision.ASK_USER); + + engine.addRule({ toolName: 'new-tool', decision: PolicyDecision.ALLOW }); + + expect(engine.check({ name: 'new-tool' })).toBe(PolicyDecision.ALLOW); + }); + }); + + describe('removeRulesForTool', () => { + it('should remove rules for specific tool', () => { + engine.addRule({ toolName: 'tool1', decision: PolicyDecision.ALLOW }); + engine.addRule({ toolName: 'tool2', decision: PolicyDecision.DENY }); + engine.addRule({ + toolName: 'tool1', + decision: PolicyDecision.ASK_USER, + priority: 10, + }); + + expect(engine.getRules()).toHaveLength(3); + + engine.removeRulesForTool('tool1'); + + const remainingRules = engine.getRules(); + expect(remainingRules).toHaveLength(1); + expect(remainingRules.some((r) => r.toolName === 'tool1')).toBe(false); + expect(remainingRules.some((r) => r.toolName === 'tool2')).toBe(true); + }); + + it('should handle removing non-existent tool', () => { + engine.addRule({ toolName: 'existing', decision: PolicyDecision.ALLOW }); + + expect(() => engine.removeRulesForTool('non-existent')).not.toThrow(); + expect(engine.getRules()).toHaveLength(1); + }); + }); + + describe('getRules', () => { + it('should return readonly array of rules', () => { + const rules: PolicyRule[] = [ + { toolName: 'tool1', decision: PolicyDecision.ALLOW }, + { toolName: 'tool2', decision: PolicyDecision.DENY }, + ]; + + engine = new PolicyEngine({ rules }); + + const retrievedRules = engine.getRules(); + expect(retrievedRules).toHaveLength(2); + expect(retrievedRules[0].toolName).toBe('tool1'); + expect(retrievedRules[1].toolName).toBe('tool2'); + }); + }); + + describe('complex scenarios', () => { + it('should handle multiple matching rules with different priorities', () => { + const rules: PolicyRule[] = [ + { decision: PolicyDecision.DENY, priority: 0 }, // Default deny all + { toolName: 'shell', decision: PolicyDecision.ASK_USER, priority: 5 }, + { + toolName: 'shell', + argsPattern: /"command":"ls/, + decision: PolicyDecision.ALLOW, + priority: 10, + }, + ]; + + engine = new PolicyEngine({ rules }); + + // Matches highest priority rule (ls command) + expect(engine.check({ name: 'shell', args: { command: 'ls -la' } })).toBe( + PolicyDecision.ALLOW, + ); + + // Matches middle priority rule (shell without ls) + expect(engine.check({ name: 'shell', args: { command: 'pwd' } })).toBe( + PolicyDecision.ASK_USER, + ); + + // Matches lowest priority rule (not shell) + expect(engine.check({ name: 'edit' })).toBe(PolicyDecision.DENY); + }); + + it('should handle tools with no args', () => { + const rules: PolicyRule[] = [ + { + toolName: 'read', + argsPattern: /secret/, + decision: PolicyDecision.DENY, + }, + ]; + + engine = new PolicyEngine({ rules }); + + // Tool call without args should not match pattern + expect(engine.check({ name: 'read' })).toBe(PolicyDecision.ASK_USER); + + // Tool call with args not matching pattern + expect(engine.check({ name: 'read', args: { file: 'public.txt' } })).toBe( + PolicyDecision.ASK_USER, + ); + + // Tool call with args matching pattern + expect(engine.check({ name: 'read', args: { file: 'secret.txt' } })).toBe( + PolicyDecision.DENY, + ); + }); + }); +}); diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts new file mode 100644 index 00000000000..027fb95ad05 --- /dev/null +++ b/packages/core/src/policy/policy-engine.ts @@ -0,0 +1,98 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { type FunctionCall } from '@google/genai'; +import { + PolicyDecision, + type PolicyEngineConfig, + type PolicyRule, +} from './types.js'; + +export class PolicyEngine { + private readonly rules: PolicyRule[]; + private readonly defaultDecision: PolicyDecision; + private readonly nonInteractive: boolean; + + constructor(config: PolicyEngineConfig = {}) { + this.rules = (config.rules ?? []).sort( + (a, b) => (b.priority ?? 0) - (a.priority ?? 0), + ); + this.defaultDecision = config.defaultDecision ?? PolicyDecision.ASK_USER; + this.nonInteractive = config.nonInteractive ?? false; + } + + /** + * Check if a tool call is allowed based on the configured policies. + */ + check(toolCall: FunctionCall): PolicyDecision { + // Find the first matching rule (already sorted by priority) + for (const rule of this.rules) { + if (this.ruleMatches(rule, toolCall)) { + return this.applyNonInteractiveMode(rule.decision); + } + } + + // No matching rule found, use default decision + return this.applyNonInteractiveMode(this.defaultDecision); + } + + /** + * Add a new rule to the policy engine. + */ + addRule(rule: PolicyRule): void { + this.rules.push(rule); + // Re-sort rules by priority + this.rules.sort((a, b) => (b.priority ?? 0) - (a.priority ?? 0)); + } + + /** + * Remove rules for a specific tool. + */ + removeRulesForTool(toolName: string): void { + // Remove all rules for the specified tool, not just the first one + for (let i = this.rules.length - 1; i >= 0; i--) { + if (this.rules[i].toolName === toolName) { + this.rules.splice(i, 1); + } + } + } + + /** + * Get all current rules. + */ + getRules(): readonly PolicyRule[] { + return this.rules; + } + + private ruleMatches(rule: PolicyRule, toolCall: FunctionCall): boolean { + // Check tool name if specified + if (rule.toolName && toolCall.name !== rule.toolName) { + return false; + } + + // Check args pattern if specified + if (rule.argsPattern) { + // If rule has an args pattern but tool has no args, no match + if (!toolCall.args) { + return false; + } + const argsString = JSON.stringify(toolCall.args); + if (!rule.argsPattern.test(argsString)) { + return false; + } + } + + return true; + } + + private applyNonInteractiveMode(decision: PolicyDecision): PolicyDecision { + // In non-interactive mode, ASK_USER becomes DENY + if (this.nonInteractive && decision === PolicyDecision.ASK_USER) { + return PolicyDecision.DENY; + } + return decision; + } +} diff --git a/packages/core/src/policy/types.ts b/packages/core/src/policy/types.ts new file mode 100644 index 00000000000..f20a88e70c6 --- /dev/null +++ b/packages/core/src/policy/types.ts @@ -0,0 +1,55 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export enum PolicyDecision { + ALLOW = 'allow', + DENY = 'deny', + ASK_USER = 'ask_user', +} + +export interface PolicyRule { + /** + * The name of the tool this rule applies to. + * If undefined, the rule applies to all tools. + */ + toolName?: string; + + /** + * Pattern to match against tool arguments. + * Can be used for more fine-grained control. + */ + argsPattern?: RegExp; + + /** + * The decision to make when this rule matches. + */ + decision: PolicyDecision; + + /** + * Priority of this rule. Higher numbers take precedence. + * Default is 0. + */ + priority?: number; +} + +export interface PolicyEngineConfig { + /** + * List of policy rules to apply. + */ + rules?: PolicyRule[]; + + /** + * Default decision when no rules match. + * Defaults to ASK_USER. + */ + defaultDecision?: PolicyDecision; + + /** + * Whether to allow tools in non-interactive mode. + * When true, ASK_USER decisions become DENY. + */ + nonInteractive?: boolean; +} From dccd03a6d97c02a004359a52f80c0fada5318625 Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Fri, 5 Sep 2025 18:34:15 -0700 Subject: [PATCH 02/10] fix(policy): address security issue in PolicyEngine argument matching Fixes the security vulnerability identified by Gemini Code Assist where JSON.stringify() could produce different output for the same arguments due to non-deterministic property ordering. This could allow security policies to be bypassed. Changes: - Implement stableStringify() method with sorted keys for consistent output - Ensure pattern matching works regardless of argument property order - Add comprehensive tests for order-independent matching - Test nested objects to ensure deep stability This ensures PolicyEngine rules cannot be bypassed by reordering object properties in tool arguments. --- .../core/src/policy/policy-engine.test.ts | 59 +++++++++++++++++++ packages/core/src/policy/policy-engine.ts | 29 ++++++++- 2 files changed, 87 insertions(+), 1 deletion(-) diff --git a/packages/core/src/policy/policy-engine.test.ts b/packages/core/src/policy/policy-engine.test.ts index c734394f712..da2c9eded5d 100644 --- a/packages/core/src/policy/policy-engine.test.ts +++ b/packages/core/src/policy/policy-engine.test.ts @@ -270,5 +270,64 @@ describe('PolicyEngine', () => { PolicyDecision.DENY, ); }); + + it('should match args pattern regardless of property order', () => { + const rules: PolicyRule[] = [ + { + toolName: 'shell', + // Pattern matches the stable stringified format + argsPattern: /"command":"rm[^"]*-rf/, + decision: PolicyDecision.DENY, + }, + ]; + + engine = new PolicyEngine({ rules }); + + // Same args with different property order should both match + const args1 = { command: 'rm -rf /', path: '/home' }; + const args2 = { path: '/home', command: 'rm -rf /' }; + + expect(engine.check({ name: 'shell', args: args1 })).toBe( + PolicyDecision.DENY, + ); + expect(engine.check({ name: 'shell', args: args2 })).toBe( + PolicyDecision.DENY, + ); + + // Verify safe command doesn't match + const safeArgs = { command: 'ls -la', path: '/home' }; + expect(engine.check({ name: 'shell', args: safeArgs })).toBe( + PolicyDecision.ASK_USER, + ); + }); + + it('should handle nested objects in args with stable stringification', () => { + const rules: PolicyRule[] = [ + { + toolName: 'api', + argsPattern: /"sensitive":true/, + decision: PolicyDecision.DENY, + }, + ]; + + engine = new PolicyEngine({ rules }); + + // Nested objects with different key orders should match consistently + const args1 = { + data: { sensitive: true, value: 'secret' }, + method: 'POST', + }; + const args2 = { + method: 'POST', + data: { value: 'secret', sensitive: true }, + }; + + expect(engine.check({ name: 'api', args: args1 })).toBe( + PolicyDecision.DENY, + ); + expect(engine.check({ name: 'api', args: args2 })).toBe( + PolicyDecision.DENY, + ); + }); }); }); diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index 027fb95ad05..a2d0225d88c 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -79,7 +79,8 @@ export class PolicyEngine { if (!toolCall.args) { return false; } - const argsString = JSON.stringify(toolCall.args); + // Use stable JSON stringification with sorted keys to ensure consistent matching + const argsString = this.stableStringify(toolCall.args); if (!rule.argsPattern.test(argsString)) { return false; } @@ -88,6 +89,32 @@ export class PolicyEngine { return true; } + /** + * Produces a stable JSON string representation with sorted keys. + * This ensures consistent pattern matching regardless of property order. + */ + private stableStringify(obj: unknown): string { + if (obj === null || obj === undefined) { + return JSON.stringify(obj); + } + + if (typeof obj !== 'object') { + return JSON.stringify(obj); + } + + if (Array.isArray(obj)) { + return '[' + obj.map(item => this.stableStringify(item)).join(',') + ']'; + } + + // Sort object keys and recursively stringify + const sortedKeys = Object.keys(obj).sort(); + const pairs = sortedKeys.map(key => { + const value = (obj as Record)[key]; + return JSON.stringify(key) + ':' + this.stableStringify(value); + }); + return '{' + pairs.join(',') + '}'; + } + private applyNonInteractiveMode(decision: PolicyDecision): PolicyDecision { // In non-interactive mode, ASK_USER becomes DENY if (this.nonInteractive && decision === PolicyDecision.ASK_USER) { From 805270bb1f9cb4ff7d70a5a8d639fac949dd0f5b Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Fri, 5 Sep 2025 19:42:16 -0700 Subject: [PATCH 03/10] fix(policy): prevent stack overflow from circular references in stableStringify Addresses the denial-of-service vulnerability identified in code review where circular references in tool arguments could cause infinite recursion and stack overflow in the stableStringify method. Changes: - Implement WeakSet tracking to detect and handle circular references - Replace recursive calls with inner stringify function that tracks visited objects - Return "[Circular]" placeholder when circular reference detected - Add comprehensive tests for both shallow and deep circular references - Use proper TypeScript types instead of 'any' in tests Security impact: - Prevents potential DoS attacks via crafted tool arguments - Ensures PolicyEngine remains robust against malformed input - Maintains stable stringification while preventing infinite recursion --- .../core/src/policy/policy-engine.test.ts | 81 +++++++++++++++++++ packages/core/src/policy/policy-engine.ts | 55 ++++++++----- 2 files changed, 117 insertions(+), 19 deletions(-) diff --git a/packages/core/src/policy/policy-engine.test.ts b/packages/core/src/policy/policy-engine.test.ts index da2c9eded5d..9d490f95d29 100644 --- a/packages/core/src/policy/policy-engine.test.ts +++ b/packages/core/src/policy/policy-engine.test.ts @@ -329,5 +329,86 @@ describe('PolicyEngine', () => { PolicyDecision.DENY, ); }); + + it('should handle circular references without stack overflow', () => { + const rules: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /\[Circular\]/, + decision: PolicyDecision.DENY, + }, + ]; + + engine = new PolicyEngine({ rules }); + + // Create an object with a circular reference + type CircularArgs = Record & { + data?: Record; + }; + const circularArgs: CircularArgs = { + name: 'test', + data: {}, + }; + // Create circular reference - TypeScript allows this since data is Record + (circularArgs.data as Record)['self'] = + circularArgs.data; + + // Should not throw stack overflow error + expect(() => + engine.check({ name: 'test', args: circularArgs }), + ).not.toThrow(); + + // Should detect the circular reference pattern + expect(engine.check({ name: 'test', args: circularArgs })).toBe( + PolicyDecision.DENY, + ); + + // Non-circular object should not match + const normalArgs = { name: 'test', data: { value: 'normal' } }; + expect(engine.check({ name: 'test', args: normalArgs })).toBe( + PolicyDecision.ASK_USER, + ); + }); + + it('should handle deep circular references', () => { + const rules: PolicyRule[] = [ + { + toolName: 'deep', + argsPattern: /\[Circular\]/, + decision: PolicyDecision.DENY, + }, + ]; + + engine = new PolicyEngine({ rules }); + + // Create a deep circular reference + type DeepCircular = Record & { + level1?: { + level2?: { + level3?: Record; + }; + }; + }; + const deepCircular: DeepCircular = { + level1: { + level2: { + level3: {}, + }, + }, + }; + // Create circular reference with proper type assertions + const level3 = deepCircular.level1!.level2!.level3!; + level3['back'] = deepCircular.level1; + + // Should handle without stack overflow + expect(() => + engine.check({ name: 'deep', args: deepCircular }), + ).not.toThrow(); + + // Should detect the circular reference + expect(engine.check({ name: 'deep', args: deepCircular })).toBe( + PolicyDecision.DENY, + ); + }); }); }); diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index a2d0225d88c..7dc385891cf 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -92,27 +92,44 @@ export class PolicyEngine { /** * Produces a stable JSON string representation with sorted keys. * This ensures consistent pattern matching regardless of property order. + * Handles circular references to prevent stack overflow attacks. */ private stableStringify(obj: unknown): string { - if (obj === null || obj === undefined) { - return JSON.stringify(obj); - } - - if (typeof obj !== 'object') { - return JSON.stringify(obj); - } - - if (Array.isArray(obj)) { - return '[' + obj.map(item => this.stableStringify(item)).join(',') + ']'; - } - - // Sort object keys and recursively stringify - const sortedKeys = Object.keys(obj).sort(); - const pairs = sortedKeys.map(key => { - const value = (obj as Record)[key]; - return JSON.stringify(key) + ':' + this.stableStringify(value); - }); - return '{' + pairs.join(',') + '}'; + const stringify = (currentObj: unknown, seen: WeakSet): string => { + if (currentObj === null || currentObj === undefined) { + return JSON.stringify(currentObj); + } + + if (typeof currentObj !== 'object') { + return JSON.stringify(currentObj); + } + + // Check for circular reference + if (seen.has(currentObj)) { + return '"[Circular]"'; + } + seen.add(currentObj); + + if (Array.isArray(currentObj)) { + return ( + '[' + + (currentObj as unknown[]) + .map((item) => stringify(item, seen)) + .join(',') + + ']' + ); + } + + // Sort object keys and recursively stringify + const sortedKeys = Object.keys(currentObj).sort(); + const pairs = sortedKeys.map((key) => { + const value = (currentObj as Record)[key]; + return JSON.stringify(key) + ':' + stringify(value, seen); + }); + return '{' + pairs.join(',') + '}'; + }; + + return stringify(obj, new WeakSet()); } private applyNonInteractiveMode(decision: PolicyDecision): PolicyDecision { From f2ea10a46adc8ae79fd47750e4cefe94bfcdc21d Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Fri, 5 Sep 2025 20:13:51 -0700 Subject: [PATCH 04/10] fix(policy-engine): address high-severity security issues in stableStringify Fixed multiple security vulnerabilities identified in code review: - Correctly handle repeated non-circular objects (no false positives) - Properly handle undefined and function values per JSON spec - Use Set with ancestor chain tracking instead of WeakSet - Ensure generated JSON is always valid - Add comprehensive test coverage for edge cases The stableStringify method now: - Tracks ancestor chain per branch to detect true circular references - Omits undefined/function values from objects (per JSON spec) - Converts undefined/functions to null in arrays (per JSON spec) - Produces deterministic output with sorted keys for consistent pattern matching --- .../core/src/policy/policy-engine.test.ts | 134 ++++++++++++++++++ packages/core/src/policy/policy-engine.ts | 60 +++++--- 2 files changed, 174 insertions(+), 20 deletions(-) diff --git a/packages/core/src/policy/policy-engine.test.ts b/packages/core/src/policy/policy-engine.test.ts index 9d490f95d29..241feab978f 100644 --- a/packages/core/src/policy/policy-engine.test.ts +++ b/packages/core/src/policy/policy-engine.test.ts @@ -410,5 +410,139 @@ describe('PolicyEngine', () => { PolicyDecision.DENY, ); }); + + it('should handle repeated non-circular objects correctly', () => { + const rules: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /\[Circular\]/, + decision: PolicyDecision.DENY, + }, + { + toolName: 'test', + argsPattern: /"value":"shared"/, + decision: PolicyDecision.ALLOW, + priority: 10, + }, + ]; + + engine = new PolicyEngine({ rules }); + + // Create an object with repeated references but no cycles + const sharedObj = { value: 'shared' }; + const args = { + first: sharedObj, + second: sharedObj, + third: { nested: sharedObj }, + }; + + // Should NOT mark repeated objects as circular, and should match the shared value pattern + expect(engine.check({ name: 'test', args })).toBe(PolicyDecision.ALLOW); + }); + + it('should omit undefined and function values from objects', () => { + const rules: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /"definedValue":"test"/, + decision: PolicyDecision.ALLOW, + }, + ]; + + engine = new PolicyEngine({ rules }); + + const args = { + definedValue: 'test', + undefinedValue: undefined, + functionValue: () => 'hello', + nullValue: null, + }; + + // Should match pattern with defined value, undefined and functions omitted + expect(engine.check({ name: 'test', args })).toBe(PolicyDecision.ALLOW); + + // Check that the pattern would NOT match if undefined was included + const rulesWithUndefined: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /undefinedValue/, + decision: PolicyDecision.DENY, + }, + ]; + engine = new PolicyEngine({ rules: rulesWithUndefined }); + expect(engine.check({ name: 'test', args })).toBe( + PolicyDecision.ASK_USER, + ); + + // Check that the pattern would NOT match if function was included + const rulesWithFunction: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /functionValue/, + decision: PolicyDecision.DENY, + }, + ]; + engine = new PolicyEngine({ rules: rulesWithFunction }); + expect(engine.check({ name: 'test', args })).toBe( + PolicyDecision.ASK_USER, + ); + }); + + it('should convert undefined and functions to null in arrays', () => { + const rules: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /\["value",null,null,null\]/, + decision: PolicyDecision.ALLOW, + }, + ]; + + engine = new PolicyEngine({ rules }); + + const args = { + array: ['value', undefined, () => 'hello', null], + }; + + // Should match pattern with undefined and functions converted to null + expect(engine.check({ name: 'test', args })).toBe(PolicyDecision.ALLOW); + }); + + it('should produce valid JSON for all inputs', () => { + const testCases = [ + { input: { simple: 'string' }, desc: 'simple object' }, + { input: { nested: { deep: { value: 123 } } }, desc: 'nested object' }, + { input: [1, 2, 3], desc: 'simple array' }, + { input: { mixed: [1, { a: 'b' }, null] }, desc: 'mixed array' }, + { + input: { undef: undefined, func: () => {}, normal: 'value' }, + desc: 'object with undefined and function', + }, + { + input: ['a', undefined, () => {}, null], + desc: 'array with undefined and function', + }, + ]; + + for (const { input } of testCases) { + const rules: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /.*/, + decision: PolicyDecision.ALLOW, + }, + ]; + engine = new PolicyEngine({ rules }); + + // Should not throw when checking (which internally uses stableStringify) + expect(() => + engine.check({ name: 'test', args: input }), + ).not.toThrow(); + + // The check should succeed + expect(engine.check({ name: 'test', args: input })).toBe( + PolicyDecision.ALLOW, + ); + } + }); }); }); diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index 7dc385891cf..76a30c28e90 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -93,43 +93,63 @@ export class PolicyEngine { * Produces a stable JSON string representation with sorted keys. * This ensures consistent pattern matching regardless of property order. * Handles circular references to prevent stack overflow attacks. + * Properly handles undefined and function values per JSON spec. */ private stableStringify(obj: unknown): string { - const stringify = (currentObj: unknown, seen: WeakSet): string => { - if (currentObj === null || currentObj === undefined) { - return JSON.stringify(currentObj); + const stringify = ( + currentObj: unknown, + ancestors: Set + ): string => { + // Handle primitives and null + if (currentObj === undefined) { + return 'null'; // undefined in arrays becomes null in JSON + } + if (currentObj === null) { + return 'null'; + } + if (typeof currentObj === 'function') { + return 'null'; // functions in arrays become null in JSON } - if (typeof currentObj !== 'object') { return JSON.stringify(currentObj); } - // Check for circular reference - if (seen.has(currentObj)) { + // Check for circular reference (object is in ancestor chain) + if (ancestors.has(currentObj)) { return '"[Circular]"'; } - seen.add(currentObj); + + // Create new ancestors set for this branch + const newAncestors = new Set(ancestors); + newAncestors.add(currentObj); if (Array.isArray(currentObj)) { - return ( - '[' + - (currentObj as unknown[]) - .map((item) => stringify(item, seen)) - .join(',') + - ']' - ); + const items = currentObj.map((item) => { + // undefined and functions in arrays become null + if (item === undefined || typeof item === 'function') { + return 'null'; + } + return stringify(item, newAncestors); + }); + return '[' + items.join(',') + ']'; } - // Sort object keys and recursively stringify + // Handle objects - sort keys and filter out undefined/function values const sortedKeys = Object.keys(currentObj).sort(); - const pairs = sortedKeys.map((key) => { + const pairs: string[] = []; + + for (const key of sortedKeys) { const value = (currentObj as Record)[key]; - return JSON.stringify(key) + ':' + stringify(value, seen); - }); + // Skip undefined and function values in objects (per JSON spec) + if (value !== undefined && typeof value !== 'function') { + pairs.push(JSON.stringify(key) + ':' + stringify(value, newAncestors)); + } + } + return '{' + pairs.join(',') + '}'; }; - return stringify(obj, new WeakSet()); + return stringify(obj, new Set()); } private applyNonInteractiveMode(decision: PolicyDecision): PolicyDecision { @@ -139,4 +159,4 @@ export class PolicyEngine { } return decision; } -} +} \ No newline at end of file From 679f05eb336097b34d5c3881c5925349f33a5175 Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Fri, 5 Sep 2025 20:27:26 -0700 Subject: [PATCH 05/10] fix(tests): resolve TypeScript build errors in policy-engine tests - Wrap array test inputs in objects to match Record type - Add explicit type annotation for testCases array - Ensure all test inputs have correct types for args parameter --- .../core/src/policy/policy-engine.test.ts | 36 ++++++++++--------- packages/core/src/policy/policy-engine.ts | 12 ++++--- 2 files changed, 26 insertions(+), 22 deletions(-) diff --git a/packages/core/src/policy/policy-engine.test.ts b/packages/core/src/policy/policy-engine.test.ts index 241feab978f..865c12cd636 100644 --- a/packages/core/src/policy/policy-engine.test.ts +++ b/packages/core/src/policy/policy-engine.test.ts @@ -508,20 +508,24 @@ describe('PolicyEngine', () => { }); it('should produce valid JSON for all inputs', () => { - const testCases = [ - { input: { simple: 'string' }, desc: 'simple object' }, - { input: { nested: { deep: { value: 123 } } }, desc: 'nested object' }, - { input: [1, 2, 3], desc: 'simple array' }, - { input: { mixed: [1, { a: 'b' }, null] }, desc: 'mixed array' }, - { - input: { undef: undefined, func: () => {}, normal: 'value' }, - desc: 'object with undefined and function', - }, - { - input: ['a', undefined, () => {}, null], - desc: 'array with undefined and function', - }, - ]; + const testCases: Array<{ input: Record; desc: string }> = + [ + { input: { simple: 'string' }, desc: 'simple object' }, + { + input: { nested: { deep: { value: 123 } } }, + desc: 'nested object', + }, + { input: { data: [1, 2, 3] }, desc: 'simple array' }, + { input: { mixed: [1, { a: 'b' }, null] }, desc: 'mixed array' }, + { + input: { undef: undefined, func: () => {}, normal: 'value' }, + desc: 'object with undefined and function', + }, + { + input: { data: ['a', undefined, () => {}, null] }, + desc: 'array with undefined and function', + }, + ]; for (const { input } of testCases) { const rules: PolicyRule[] = [ @@ -534,9 +538,7 @@ describe('PolicyEngine', () => { engine = new PolicyEngine({ rules }); // Should not throw when checking (which internally uses stableStringify) - expect(() => - engine.check({ name: 'test', args: input }), - ).not.toThrow(); + expect(() => engine.check({ name: 'test', args: input })).not.toThrow(); // The check should succeed expect(engine.check({ name: 'test', args: input })).toBe( diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index 76a30c28e90..a079ea45d37 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -98,7 +98,7 @@ export class PolicyEngine { private stableStringify(obj: unknown): string { const stringify = ( currentObj: unknown, - ancestors: Set + ancestors: Set, ): string => { // Handle primitives and null if (currentObj === undefined) { @@ -137,15 +137,17 @@ export class PolicyEngine { // Handle objects - sort keys and filter out undefined/function values const sortedKeys = Object.keys(currentObj).sort(); const pairs: string[] = []; - + for (const key of sortedKeys) { const value = (currentObj as Record)[key]; // Skip undefined and function values in objects (per JSON spec) if (value !== undefined && typeof value !== 'function') { - pairs.push(JSON.stringify(key) + ':' + stringify(value, newAncestors)); + pairs.push( + JSON.stringify(key) + ':' + stringify(value, newAncestors), + ); } } - + return '{' + pairs.join(',') + '}'; }; @@ -159,4 +161,4 @@ export class PolicyEngine { } return decision; } -} \ No newline at end of file +} From d693fbf685ac7a3332015c24d9cbe30c9ba5ac1a Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Fri, 5 Sep 2025 20:35:04 -0700 Subject: [PATCH 06/10] fix(policy-engine): address critical security issues and improve documentation Fixes based on code review: - Remove readonly modifier from rules array since it's mutated by addRule/removeRulesForTool - Add support for toJSON methods per JSON.stringify specification - Add comprehensive documentation for stableStringify method - Add tests for toJSON handling including edge cases The stableStringify method now fully respects the JSON specification while maintaining deterministic output for security policy matching. --- .../core/src/policy/policy-engine.test.ts | 74 +++++++++++++++++++ packages/core/src/policy/policy-engine.ts | 68 +++++++++++++++-- 2 files changed, 137 insertions(+), 5 deletions(-) diff --git a/packages/core/src/policy/policy-engine.test.ts b/packages/core/src/policy/policy-engine.test.ts index 865c12cd636..51f222b2e42 100644 --- a/packages/core/src/policy/policy-engine.test.ts +++ b/packages/core/src/policy/policy-engine.test.ts @@ -546,5 +546,79 @@ describe('PolicyEngine', () => { ); } }); + + it('should respect toJSON methods on objects', () => { + const rules: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /"sanitized":"safe"/, + decision: PolicyDecision.ALLOW, + }, + { + toolName: 'test', + argsPattern: /"dangerous":"data"/, + decision: PolicyDecision.DENY, + }, + ]; + + engine = new PolicyEngine({ rules }); + + // Object with toJSON that sanitizes output + const args = { + data: { + dangerous: 'data', + toJSON: () => ({ sanitized: 'safe' }), + }, + }; + + // Should match the sanitized pattern, not the dangerous one + expect(engine.check({ name: 'test', args })).toBe(PolicyDecision.ALLOW); + }); + + it('should handle toJSON that returns primitives', () => { + const rules: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /"value":"string-value"/, + decision: PolicyDecision.ALLOW, + }, + ]; + + engine = new PolicyEngine({ rules }); + + const args = { + value: { + complex: 'object', + toJSON: () => 'string-value', + }, + }; + + // toJSON returns a string, which should be properly stringified + expect(engine.check({ name: 'test', args })).toBe(PolicyDecision.ALLOW); + }); + + it('should handle toJSON that throws an error', () => { + const rules: PolicyRule[] = [ + { + toolName: 'test', + argsPattern: /"fallback":"value"/, + decision: PolicyDecision.ALLOW, + }, + ]; + + engine = new PolicyEngine({ rules }); + + const args = { + data: { + fallback: 'value', + toJSON: () => { + throw new Error('toJSON error'); + }, + }, + }; + + // Should fall back to regular object serialization when toJSON throws + expect(engine.check({ name: 'test', args })).toBe(PolicyDecision.ALLOW); + }); }); }); diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index a079ea45d37..a01f014c6b0 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -12,7 +12,7 @@ import { } from './types.js'; export class PolicyEngine { - private readonly rules: PolicyRule[]; + private rules: PolicyRule[]; private readonly defaultDecision: PolicyDecision; private readonly nonInteractive: boolean; @@ -90,10 +90,56 @@ export class PolicyEngine { } /** - * Produces a stable JSON string representation with sorted keys. - * This ensures consistent pattern matching regardless of property order. - * Handles circular references to prevent stack overflow attacks. - * Properly handles undefined and function values per JSON spec. + * Produces a stable, deterministic JSON string representation with sorted keys. + * + * This method is critical for security policy matching. It ensures that the same + * object always produces the same string representation, regardless of property + * insertion order, which could vary across different JavaScript engines or + * runtime conditions. + * + * Key behaviors: + * 1. **Sorted Keys**: Object properties are always serialized in alphabetical order, + * ensuring deterministic output for pattern matching. + * + * 2. **Circular Reference Protection**: Uses ancestor chain tracking (not just + * object identity) to detect true circular references while correctly handling + * repeated non-circular object references. Circular references are replaced + * with "[Circular]" to prevent stack overflow attacks. + * + * 3. **JSON Spec Compliance**: + * - undefined values: Omitted from objects, converted to null in arrays + * - Functions: Omitted from objects, converted to null in arrays + * - toJSON methods: Respected and called when present (per JSON.stringify spec) + * + * 4. **Security Considerations**: + * - Prevents DoS via circular references that would cause infinite recursion + * - Ensures consistent policy rule matching by normalizing property order + * - Respects toJSON for objects that sanitize their output + * - Handles toJSON methods that throw errors gracefully + * + * @param obj - The object to stringify (typically toolCall.args) + * @returns A deterministic JSON string representation + * + * @example + * // Different property orders produce the same output: + * stableStringify({b: 2, a: 1}) === stableStringify({a: 1, b: 2}) + * // Returns: '{"a":1,"b":2}' + * + * @example + * // Circular references are handled safely: + * const obj = {a: 1}; + * obj.self = obj; + * stableStringify(obj) + * // Returns: '{"a":1,"self":"[Circular]"}' + * + * @example + * // toJSON methods are respected: + * const obj = { + * sensitive: 'secret', + * toJSON: () => ({ safe: 'data' }) + * }; + * stableStringify(obj) + * // Returns: '{"safe":"data"}' */ private stableStringify(obj: unknown): string { const stringify = ( @@ -123,6 +169,18 @@ export class PolicyEngine { const newAncestors = new Set(ancestors); newAncestors.add(currentObj); + // Check for toJSON method and use it if present + const objWithToJSON = currentObj as { toJSON?: () => unknown }; + if (typeof objWithToJSON.toJSON === 'function') { + try { + const jsonValue = objWithToJSON.toJSON(); + // The result of toJSON needs to be stringified recursively + return stringify(jsonValue, newAncestors); + } catch { + // If toJSON throws, treat as a regular object + } + } + if (Array.isArray(currentObj)) { const items = currentObj.map((item) => { // undefined and functions in arrays become null From c35d83cedb102a64693e0f073921ef06c8ec10bc Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Fri, 5 Sep 2025 20:36:57 -0700 Subject: [PATCH 07/10] Fix lint --- packages/core/src/policy/policy-engine.ts | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index a01f014c6b0..0472a8eedb3 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -91,47 +91,47 @@ export class PolicyEngine { /** * Produces a stable, deterministic JSON string representation with sorted keys. - * + * * This method is critical for security policy matching. It ensures that the same * object always produces the same string representation, regardless of property * insertion order, which could vary across different JavaScript engines or * runtime conditions. - * + * * Key behaviors: * 1. **Sorted Keys**: Object properties are always serialized in alphabetical order, * ensuring deterministic output for pattern matching. - * + * * 2. **Circular Reference Protection**: Uses ancestor chain tracking (not just * object identity) to detect true circular references while correctly handling * repeated non-circular object references. Circular references are replaced * with "[Circular]" to prevent stack overflow attacks. - * + * * 3. **JSON Spec Compliance**: * - undefined values: Omitted from objects, converted to null in arrays * - Functions: Omitted from objects, converted to null in arrays * - toJSON methods: Respected and called when present (per JSON.stringify spec) - * + * * 4. **Security Considerations**: * - Prevents DoS via circular references that would cause infinite recursion * - Ensures consistent policy rule matching by normalizing property order * - Respects toJSON for objects that sanitize their output * - Handles toJSON methods that throw errors gracefully - * + * * @param obj - The object to stringify (typically toolCall.args) * @returns A deterministic JSON string representation - * + * * @example * // Different property orders produce the same output: * stableStringify({b: 2, a: 1}) === stableStringify({a: 1, b: 2}) * // Returns: '{"a":1,"b":2}' - * + * * @example * // Circular references are handled safely: * const obj = {a: 1}; * obj.self = obj; * stableStringify(obj) * // Returns: '{"a":1,"self":"[Circular]"}' - * + * * @example * // toJSON methods are respected: * const obj = { From c3c8de81c03ecbf2f9790ca97c74c17b17d4e4d3 Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Fri, 5 Sep 2025 20:51:28 -0700 Subject: [PATCH 08/10] fix(message-bus): use safeJsonStringify for error messages Replace JSON.stringify with safeJsonStringify to handle potential circular references in message objects. This prevents TypeError exceptions when logging invalid messages with circular references in toolCall.args. --- packages/core/src/confirmation-bus/message-bus.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/core/src/confirmation-bus/message-bus.ts b/packages/core/src/confirmation-bus/message-bus.ts index 40fd263d7d9..b9d66eff6ab 100644 --- a/packages/core/src/confirmation-bus/message-bus.ts +++ b/packages/core/src/confirmation-bus/message-bus.ts @@ -8,6 +8,7 @@ import { EventEmitter } from 'node:events'; import type { PolicyEngine } from '../policy/policy-engine.js'; import { PolicyDecision } from '../policy/types.js'; import { MessageBusType, type Message } from './types.js'; +import { safeJsonStringify } from '../utils/safeJsonStringify.js'; export class MessageBus extends EventEmitter { constructor(private readonly policyEngine: PolicyEngine) { @@ -37,7 +38,7 @@ export class MessageBus extends EventEmitter { try { if (!this.isValidMessage(message)) { throw new Error( - `Invalid message structure: ${JSON.stringify(message)}`, + `Invalid message structure: ${safeJsonStringify(message)}`, ); } From 351cd63694b166a538b00fa2464171df384f0e27 Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Tue, 9 Sep 2025 20:13:35 -0400 Subject: [PATCH 09/10] feat(policy): Refactor policy engine and message bus This commit addresses several comments from PR #7835. - Refactors `removeRulesForTool` in `PolicyEngine` to use `.filter()` for better readability. - Extracts `stableStringify` into its own utility file and makes `ruleMatches` a file-level function. - Adds structural types to message objects in `message-bus.test.ts` to improve type safety. - Makes `ToolExecutionSuccess` and `ToolExecutionFailure` interfaces generic for more flexible typing. - Fixes failing tests in `message-bus.test.ts`. --- .../src/confirmation-bus/message-bus.test.ts | 31 +-- packages/core/src/confirmation-bus/types.ts | 8 +- packages/core/src/policy/policy-engine.ts | 177 +++--------------- packages/core/src/policy/stable-stringify.ts | 123 ++++++++++++ 4 files changed, 172 insertions(+), 167 deletions(-) create mode 100644 packages/core/src/policy/stable-stringify.ts diff --git a/packages/core/src/confirmation-bus/message-bus.test.ts b/packages/core/src/confirmation-bus/message-bus.test.ts index e2d0bd2fe09..8156671c9b4 100644 --- a/packages/core/src/confirmation-bus/message-bus.test.ts +++ b/packages/core/src/confirmation-bus/message-bus.test.ts @@ -8,7 +8,13 @@ import { describe, it, expect, beforeEach, vi } from 'vitest'; import { MessageBus } from './message-bus.js'; import { PolicyEngine } from '../policy/policy-engine.js'; import { PolicyDecision } from '../policy/types.js'; -import { MessageBusType, type ToolConfirmationRequest } from './types.js'; +import { + MessageBusType, + type ToolConfirmationRequest, + type ToolConfirmationResponse, + type ToolPolicyRejection, + type ToolExecutionSuccess, +} from './types.js'; describe('MessageBus', () => { let messageBus: MessageBus; @@ -64,11 +70,12 @@ describe('MessageBus', () => { messageBus.publish(request); - expect(responseHandler).toHaveBeenCalledWith({ + const expectedResponse: ToolConfirmationResponse = { type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, correlationId: '123', confirmed: true, - }); + }; + expect(responseHandler).toHaveBeenCalledWith(expectedResponse); }); it('should emit rejection and response when policy denies', () => { @@ -93,16 +100,18 @@ describe('MessageBus', () => { messageBus.publish(request); - expect(rejectionHandler).toHaveBeenCalledWith({ + const expectedRejection: ToolPolicyRejection = { type: MessageBusType.TOOL_POLICY_REJECTION, toolCall: { name: 'test-tool', args: {} }, - }); + }; + expect(rejectionHandler).toHaveBeenCalledWith(expectedRejection); - expect(responseHandler).toHaveBeenCalledWith({ + const expectedResponse: ToolConfirmationResponse = { type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, correlationId: '123', confirmed: false, - }); + }; + expect(responseHandler).toHaveBeenCalledWith(expectedResponse); }); it('should pass through to UI when policy says ASK_USER', () => { @@ -132,7 +141,7 @@ describe('MessageBus', () => { successHandler, ); - const message = { + const message: ToolExecutionSuccess = { type: MessageBusType.TOOL_EXECUTION_SUCCESS as const, toolCall: { name: 'test-tool' }, result: 'success', @@ -149,7 +158,7 @@ describe('MessageBus', () => { const handler = vi.fn(); messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler); - const message = { + const message: ToolExecutionSuccess = { type: MessageBusType.TOOL_EXECUTION_SUCCESS as const, toolCall: { name: 'test' }, result: 'test', @@ -165,7 +174,7 @@ describe('MessageBus', () => { messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler); messageBus.unsubscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler); - const message = { + const message: ToolExecutionSuccess = { type: MessageBusType.TOOL_EXECUTION_SUCCESS as const, toolCall: { name: 'test' }, result: 'test', @@ -183,7 +192,7 @@ describe('MessageBus', () => { messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler1); messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler2); - const message = { + const message: ToolExecutionSuccess = { type: MessageBusType.TOOL_EXECUTION_SUCCESS as const, toolCall: { name: 'test' }, result: 'test', diff --git a/packages/core/src/confirmation-bus/types.ts b/packages/core/src/confirmation-bus/types.ts index ea75a543d8c..cb86595be9a 100644 --- a/packages/core/src/confirmation-bus/types.ts +++ b/packages/core/src/confirmation-bus/types.ts @@ -31,16 +31,16 @@ export interface ToolPolicyRejection { toolCall: FunctionCall; } -export interface ToolExecutionSuccess { +export interface ToolExecutionSuccess { type: MessageBusType.TOOL_EXECUTION_SUCCESS; toolCall: FunctionCall; - result: unknown; + result: T; } -export interface ToolExecutionFailure { +export interface ToolExecutionFailure { type: MessageBusType.TOOL_EXECUTION_FAILURE; toolCall: FunctionCall; - error: Error; + error: E; } export type Message = diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index 0472a8eedb3..bc7f0f265c4 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -10,6 +10,29 @@ import { type PolicyEngineConfig, type PolicyRule, } from './types.js'; +import { stableStringify } from './stable-stringify.js'; + +function ruleMatches(rule: PolicyRule, toolCall: FunctionCall): boolean { + // Check tool name if specified + if (rule.toolName && toolCall.name !== rule.toolName) { + return false; + } + + // Check args pattern if specified + if (rule.argsPattern) { + // If rule has an args pattern but tool has no args, no match + if (!toolCall.args) { + return false; + } + // Use stable JSON stringification with sorted keys to ensure consistent matching + const argsString = stableStringify(toolCall.args); + if (!rule.argsPattern.test(argsString)) { + return false; + } + } + + return true; +} export class PolicyEngine { private rules: PolicyRule[]; @@ -30,7 +53,7 @@ export class PolicyEngine { check(toolCall: FunctionCall): PolicyDecision { // Find the first matching rule (already sorted by priority) for (const rule of this.rules) { - if (this.ruleMatches(rule, toolCall)) { + if (ruleMatches(rule, toolCall)) { return this.applyNonInteractiveMode(rule.decision); } } @@ -52,12 +75,7 @@ export class PolicyEngine { * Remove rules for a specific tool. */ removeRulesForTool(toolName: string): void { - // Remove all rules for the specified tool, not just the first one - for (let i = this.rules.length - 1; i >= 0; i--) { - if (this.rules[i].toolName === toolName) { - this.rules.splice(i, 1); - } - } + this.rules = this.rules.filter((rule) => rule.toolName !== toolName); } /** @@ -67,151 +85,6 @@ export class PolicyEngine { return this.rules; } - private ruleMatches(rule: PolicyRule, toolCall: FunctionCall): boolean { - // Check tool name if specified - if (rule.toolName && toolCall.name !== rule.toolName) { - return false; - } - - // Check args pattern if specified - if (rule.argsPattern) { - // If rule has an args pattern but tool has no args, no match - if (!toolCall.args) { - return false; - } - // Use stable JSON stringification with sorted keys to ensure consistent matching - const argsString = this.stableStringify(toolCall.args); - if (!rule.argsPattern.test(argsString)) { - return false; - } - } - - return true; - } - - /** - * Produces a stable, deterministic JSON string representation with sorted keys. - * - * This method is critical for security policy matching. It ensures that the same - * object always produces the same string representation, regardless of property - * insertion order, which could vary across different JavaScript engines or - * runtime conditions. - * - * Key behaviors: - * 1. **Sorted Keys**: Object properties are always serialized in alphabetical order, - * ensuring deterministic output for pattern matching. - * - * 2. **Circular Reference Protection**: Uses ancestor chain tracking (not just - * object identity) to detect true circular references while correctly handling - * repeated non-circular object references. Circular references are replaced - * with "[Circular]" to prevent stack overflow attacks. - * - * 3. **JSON Spec Compliance**: - * - undefined values: Omitted from objects, converted to null in arrays - * - Functions: Omitted from objects, converted to null in arrays - * - toJSON methods: Respected and called when present (per JSON.stringify spec) - * - * 4. **Security Considerations**: - * - Prevents DoS via circular references that would cause infinite recursion - * - Ensures consistent policy rule matching by normalizing property order - * - Respects toJSON for objects that sanitize their output - * - Handles toJSON methods that throw errors gracefully - * - * @param obj - The object to stringify (typically toolCall.args) - * @returns A deterministic JSON string representation - * - * @example - * // Different property orders produce the same output: - * stableStringify({b: 2, a: 1}) === stableStringify({a: 1, b: 2}) - * // Returns: '{"a":1,"b":2}' - * - * @example - * // Circular references are handled safely: - * const obj = {a: 1}; - * obj.self = obj; - * stableStringify(obj) - * // Returns: '{"a":1,"self":"[Circular]"}' - * - * @example - * // toJSON methods are respected: - * const obj = { - * sensitive: 'secret', - * toJSON: () => ({ safe: 'data' }) - * }; - * stableStringify(obj) - * // Returns: '{"safe":"data"}' - */ - private stableStringify(obj: unknown): string { - const stringify = ( - currentObj: unknown, - ancestors: Set, - ): string => { - // Handle primitives and null - if (currentObj === undefined) { - return 'null'; // undefined in arrays becomes null in JSON - } - if (currentObj === null) { - return 'null'; - } - if (typeof currentObj === 'function') { - return 'null'; // functions in arrays become null in JSON - } - if (typeof currentObj !== 'object') { - return JSON.stringify(currentObj); - } - - // Check for circular reference (object is in ancestor chain) - if (ancestors.has(currentObj)) { - return '"[Circular]"'; - } - - // Create new ancestors set for this branch - const newAncestors = new Set(ancestors); - newAncestors.add(currentObj); - - // Check for toJSON method and use it if present - const objWithToJSON = currentObj as { toJSON?: () => unknown }; - if (typeof objWithToJSON.toJSON === 'function') { - try { - const jsonValue = objWithToJSON.toJSON(); - // The result of toJSON needs to be stringified recursively - return stringify(jsonValue, newAncestors); - } catch { - // If toJSON throws, treat as a regular object - } - } - - if (Array.isArray(currentObj)) { - const items = currentObj.map((item) => { - // undefined and functions in arrays become null - if (item === undefined || typeof item === 'function') { - return 'null'; - } - return stringify(item, newAncestors); - }); - return '[' + items.join(',') + ']'; - } - - // Handle objects - sort keys and filter out undefined/function values - const sortedKeys = Object.keys(currentObj).sort(); - const pairs: string[] = []; - - for (const key of sortedKeys) { - const value = (currentObj as Record)[key]; - // Skip undefined and function values in objects (per JSON spec) - if (value !== undefined && typeof value !== 'function') { - pairs.push( - JSON.stringify(key) + ':' + stringify(value, newAncestors), - ); - } - } - - return '{' + pairs.join(',') + '}'; - }; - - return stringify(obj, new Set()); - } - private applyNonInteractiveMode(decision: PolicyDecision): PolicyDecision { // In non-interactive mode, ASK_USER becomes DENY if (this.nonInteractive && decision === PolicyDecision.ASK_USER) { diff --git a/packages/core/src/policy/stable-stringify.ts b/packages/core/src/policy/stable-stringify.ts new file mode 100644 index 00000000000..ae74b966d61 --- /dev/null +++ b/packages/core/src/policy/stable-stringify.ts @@ -0,0 +1,123 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Produces a stable, deterministic JSON string representation with sorted keys. + * + * This method is critical for security policy matching. It ensures that the same + * object always produces the same string representation, regardless of property + * insertion order, which could vary across different JavaScript engines or + * runtime conditions. + * + * Key behaviors: + * 1. **Sorted Keys**: Object properties are always serialized in alphabetical order, + * ensuring deterministic output for pattern matching. + * + * 2. **Circular Reference Protection**: Uses ancestor chain tracking (not just + * object identity) to detect true circular references while correctly handling + * repeated non-circular object references. Circular references are replaced + * with "[Circular]" to prevent stack overflow attacks. + * + * 3. **JSON Spec Compliance**: + * - undefined values: Omitted from objects, converted to null in arrays + * - Functions: Omitted from objects, converted to null in arrays + * - toJSON methods: Respected and called when present (per JSON.stringify spec) + * + * 4. **Security Considerations**: + * - Prevents DoS via circular references that would cause infinite recursion + * - Ensures consistent policy rule matching by normalizing property order + * - Respects toJSON for objects that sanitize their output + * - Handles toJSON methods that throw errors gracefully + * + * @param obj - The object to stringify (typically toolCall.args) + * @returns A deterministic JSON string representation + * + * @example + * // Different property orders produce the same output: + * stableStringify({b: 2, a: 1}) === stableStringify({a: 1, b: 2}) + * // Returns: '{"a":1,"b":2}' + * + * @example + * // Circular references are handled safely: + * const obj = {a: 1}; + * obj.self = obj; + * stableStringify(obj) + * // Returns: '{"a":1,"self":"[Circular]"}' + * + * @example + * // toJSON methods are respected: + * const obj = { + * sensitive: 'secret', + * toJSON: () => ({ safe: 'data' }) + * }; + * stableStringify(obj) + * // Returns: '{"safe":"data"}' + */ +export function stableStringify(obj: unknown): string { + const stringify = (currentObj: unknown, ancestors: Set): string => { + // Handle primitives and null + if (currentObj === undefined) { + return 'null'; // undefined in arrays becomes null in JSON + } + if (currentObj === null) { + return 'null'; + } + if (typeof currentObj === 'function') { + return 'null'; // functions in arrays become null in JSON + } + if (typeof currentObj !== 'object') { + return JSON.stringify(currentObj); + } + + // Check for circular reference (object is in ancestor chain) + if (ancestors.has(currentObj)) { + return '"[Circular]"'; + } + + // Create new ancestors set for this branch + const newAncestors = new Set(ancestors); + newAncestors.add(currentObj); + + // Check for toJSON method and use it if present + const objWithToJSON = currentObj as { toJSON?: () => unknown }; + if (typeof objWithToJSON.toJSON === 'function') { + try { + const jsonValue = objWithToJSON.toJSON(); + // The result of toJSON needs to be stringified recursively + return stringify(jsonValue, newAncestors); + } catch { + // If toJSON throws, treat as a regular object + } + } + + if (Array.isArray(currentObj)) { + const items = currentObj.map((item) => { + // undefined and functions in arrays become null + if (item === undefined || typeof item === 'function') { + return 'null'; + } + return stringify(item, newAncestors); + }); + return '[' + items.join(',') + ']'; + } + + // Handle objects - sort keys and filter out undefined/function values + const sortedKeys = Object.keys(currentObj).sort(); + const pairs: string[] = []; + + for (const key of sortedKeys) { + const value = (currentObj as Record)[key]; + // Skip undefined and function values in objects (per JSON spec) + if (value !== undefined && typeof value !== 'function') { + pairs.push(JSON.stringify(key) + ':' + stringify(value, newAncestors)); + } + } + + return '{' + pairs.join(',') + '}'; + }; + + return stringify(obj, new Set()); +} From f086b5da516eb13d0ad918da1289e01083977709 Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Wed, 10 Sep 2025 16:28:15 -0400 Subject: [PATCH 10/10] feat(policy): Optimize stableStringify and cache result This commit addresses performance feedback on the `stableStringify` function and its usage within the `PolicyEngine`. - The `stableStringify` function is optimized to avoid creating new sets for ancestor tracking on each recursion, and to handle nulls from `toJSON` more efficiently. - The `PolicyEngine` now caches the result of `stableStringify` during a `check` operation, preventing it from being called repeatedly for each policy rule with an argument pattern. This significantly reduces the performance overhead when multiple such rules exist. --- packages/core/src/policy/policy-engine.ts | 20 ++++-- packages/core/src/policy/stable-stringify.ts | 71 +++++++++++--------- 2 files changed, 54 insertions(+), 37 deletions(-) diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index bc7f0f265c4..e1006ffdef1 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -12,7 +12,11 @@ import { } from './types.js'; import { stableStringify } from './stable-stringify.js'; -function ruleMatches(rule: PolicyRule, toolCall: FunctionCall): boolean { +function ruleMatches( + rule: PolicyRule, + toolCall: FunctionCall, + stringifiedArgs: string | undefined, +): boolean { // Check tool name if specified if (rule.toolName && toolCall.name !== rule.toolName) { return false; @@ -25,8 +29,10 @@ function ruleMatches(rule: PolicyRule, toolCall: FunctionCall): boolean { return false; } // Use stable JSON stringification with sorted keys to ensure consistent matching - const argsString = stableStringify(toolCall.args); - if (!rule.argsPattern.test(argsString)) { + if ( + stringifiedArgs === undefined || + !rule.argsPattern.test(stringifiedArgs) + ) { return false; } } @@ -51,9 +57,15 @@ export class PolicyEngine { * Check if a tool call is allowed based on the configured policies. */ check(toolCall: FunctionCall): PolicyDecision { + let stringifiedArgs: string | undefined; + // Compute stringified args once before the loop + if (toolCall.args && this.rules.some((rule) => rule.argsPattern)) { + stringifiedArgs = stableStringify(toolCall.args); + } + // Find the first matching rule (already sorted by priority) for (const rule of this.rules) { - if (ruleMatches(rule, toolCall)) { + if (ruleMatches(rule, toolCall, stringifiedArgs)) { return this.applyNonInteractiveMode(rule.decision); } } diff --git a/packages/core/src/policy/stable-stringify.ts b/packages/core/src/policy/stable-stringify.ts index ae74b966d61..78db692eab7 100644 --- a/packages/core/src/policy/stable-stringify.ts +++ b/packages/core/src/policy/stable-stringify.ts @@ -77,46 +77,51 @@ export function stableStringify(obj: unknown): string { return '"[Circular]"'; } - // Create new ancestors set for this branch - const newAncestors = new Set(ancestors); - newAncestors.add(currentObj); + ancestors.add(currentObj); - // Check for toJSON method and use it if present - const objWithToJSON = currentObj as { toJSON?: () => unknown }; - if (typeof objWithToJSON.toJSON === 'function') { - try { - const jsonValue = objWithToJSON.toJSON(); - // The result of toJSON needs to be stringified recursively - return stringify(jsonValue, newAncestors); - } catch { - // If toJSON throws, treat as a regular object + try { + // Check for toJSON method and use it if present + const objWithToJSON = currentObj as { toJSON?: () => unknown }; + if (typeof objWithToJSON.toJSON === 'function') { + try { + const jsonValue = objWithToJSON.toJSON(); + // The result of toJSON needs to be stringified recursively + if (jsonValue === null) { + return 'null'; + } + return stringify(jsonValue, ancestors); + } catch { + // If toJSON throws, treat as a regular object + } } - } - if (Array.isArray(currentObj)) { - const items = currentObj.map((item) => { - // undefined and functions in arrays become null - if (item === undefined || typeof item === 'function') { - return 'null'; - } - return stringify(item, newAncestors); - }); - return '[' + items.join(',') + ']'; - } + if (Array.isArray(currentObj)) { + const items = currentObj.map((item) => { + // undefined and functions in arrays become null + if (item === undefined || typeof item === 'function') { + return 'null'; + } + return stringify(item, ancestors); + }); + return '[' + items.join(',') + ']'; + } - // Handle objects - sort keys and filter out undefined/function values - const sortedKeys = Object.keys(currentObj).sort(); - const pairs: string[] = []; + // Handle objects - sort keys and filter out undefined/function values + const sortedKeys = Object.keys(currentObj).sort(); + const pairs: string[] = []; - for (const key of sortedKeys) { - const value = (currentObj as Record)[key]; - // Skip undefined and function values in objects (per JSON spec) - if (value !== undefined && typeof value !== 'function') { - pairs.push(JSON.stringify(key) + ':' + stringify(value, newAncestors)); + for (const key of sortedKeys) { + const value = (currentObj as Record)[key]; + // Skip undefined and function values in objects (per JSON spec) + if (value !== undefined && typeof value !== 'function') { + pairs.push(JSON.stringify(key) + ':' + stringify(value, ancestors)); + } } - } - return '{' + pairs.join(',') + '}'; + return '{' + pairs.join(',') + '}'; + } finally { + ancestors.delete(currentObj); + } }; return stringify(obj, new Set());