diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 3b582abe89..6a5e3524a0 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -8,7 +8,6 @@ import { describe, it, expect, vi } from 'vitest'; import type { Mock } from 'vitest'; import type { CallableTool } from '@google/genai'; import { CoreToolScheduler } from './coreToolScheduler.js'; -import { PLAN_MODE_DENIAL_MESSAGE } from '../scheduler/policy.js'; import type { ToolCall, WaitingToolCall, @@ -2161,7 +2160,7 @@ describe('CoreToolScheduler Sequential Execution', () => { }); describe('Policy Decisions in Plan Mode', () => { - it('should return STOP_EXECUTION error type and informative message when denied in Plan Mode', async () => { + it('should return POLICY_VIOLATION error type and informative message when denied in Plan Mode', async () => { const mockTool = new MockTool({ name: 'dangerous_tool', displayName: 'Dangerous Tool', @@ -2205,8 +2204,64 @@ describe('CoreToolScheduler Sequential Execution', () => { const result = reportedTools[0]; expect(result.status).toBe('error'); - expect(result.response.errorType).toBe(ToolErrorType.STOP_EXECUTION); - expect(result.response.error.message).toBe(PLAN_MODE_DENIAL_MESSAGE); + expect(result.response.errorType).toBe(ToolErrorType.POLICY_VIOLATION); + expect(result.response.error.message).toBe( + 'Tool execution denied by policy.', + ); + }); + + it('should return custom deny message when denied in Plan Mode with a specific rule message', async () => { + const mockTool = new MockTool({ + name: 'dangerous_tool', + displayName: 'Dangerous Tool', + description: 'Does risky stuff', + }); + const mockToolRegistry = { + getTool: () => mockTool, + getAllToolNames: () => ['dangerous_tool'], + } as unknown as ToolRegistry; + + const onAllToolCallsComplete = vi.fn(); + const customDenyMessage = 'Custom denial message for testing'; + + const mockConfig = createMockConfig({ + getToolRegistry: () => mockToolRegistry, + getApprovalMode: () => ApprovalMode.PLAN, + getPolicyEngine: () => + ({ + check: async () => ({ + decision: PolicyDecision.DENY, + rule: { denyMessage: customDenyMessage }, + }), + }) as unknown as PolicyEngine, + }); + mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined); + + const scheduler = new CoreToolScheduler({ + config: mockConfig, + onAllToolCallsComplete, + getPreferredEditor: () => 'vscode', + }); + + const request = { + callId: 'call-1', + name: 'dangerous_tool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-1', + }; + + await scheduler.schedule(request, new AbortController().signal); + + expect(onAllToolCallsComplete).toHaveBeenCalledTimes(1); + const reportedTools = onAllToolCallsComplete.mock.calls[0][0]; + const result = reportedTools[0]; + + expect(result.status).toBe('error'); + expect(result.response.errorType).toBe(ToolErrorType.POLICY_VIOLATION); + expect(result.response.error.message).toBe( + `Tool execution denied by policy. ${customDenyMessage}`, + ); }); }); }); diff --git a/packages/core/src/policy/policies/plan.toml b/packages/core/src/policy/policies/plan.toml index 4fbcb6c376..4bcecab29f 100644 --- a/packages/core/src/policy/policies/plan.toml +++ b/packages/core/src/policy/policies/plan.toml @@ -31,6 +31,7 @@ decision = "deny" priority = 20 modes = ["plan"] +deny_message = "You are in Plan Mode - adjust your prompt to only use read and search tools." # Explicitly Allow Read-Only Tools in Plan mode. diff --git a/packages/core/src/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts index ad32b93f93..a076e4c44f 100644 --- a/packages/core/src/scheduler/policy.test.ts +++ b/packages/core/src/scheduler/policy.test.ts @@ -13,11 +13,7 @@ import { beforeEach, afterEach, } from 'vitest'; -import { - checkPolicy, - updatePolicy, - PLAN_MODE_DENIAL_MESSAGE, -} from './policy.js'; +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'; @@ -441,6 +437,37 @@ describe('policy.ts', () => { ); }); }); + + describe('getPolicyDenialError', () => { + it('should return default denial message when no rule provided', () => { + const mockConfig = { + getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT), + } as unknown as Config; + + const { errorMessage, errorType } = getPolicyDenialError(mockConfig); + + expect(errorMessage).toBe('Tool execution denied by policy.'); + expect(errorType).toBe(ToolErrorType.POLICY_VIOLATION); + }); + + it('should return custom deny message if provided', () => { + const mockConfig = { + getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT), + } as unknown as Config; + const rule = { + decision: PolicyDecision.DENY, + denyMessage: 'Custom Deny', + }; + + const { errorMessage, errorType } = getPolicyDenialError( + mockConfig, + rule, + ); + + expect(errorMessage).toBe('Tool execution denied by policy. Custom Deny'); + expect(errorType).toBe(ToolErrorType.POLICY_VIOLATION); + }); + }); }); describe('Plan Mode Denial Consistency', () => { @@ -547,8 +574,8 @@ describe('Plan Mode Denial Consistency', () => { } } - expect(resultMessage).toBe(PLAN_MODE_DENIAL_MESSAGE); - expect(resultErrorType).toBe(ToolErrorType.STOP_EXECUTION); + expect(resultMessage).toBe('Tool execution denied by policy.'); + expect(resultErrorType).toBe(ToolErrorType.POLICY_VIOLATION); }); }); }); diff --git a/packages/core/src/scheduler/policy.ts b/packages/core/src/scheduler/policy.ts index 279dea85c7..247b696f22 100644 --- a/packages/core/src/scheduler/policy.ts +++ b/packages/core/src/scheduler/policy.ts @@ -26,23 +26,13 @@ import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; import { EDIT_TOOL_NAMES } from '../tools/tool-names.js'; import type { ValidatingToolCall } from './types.js'; -export const PLAN_MODE_DENIAL_MESSAGE = - 'You are in Plan Mode - adjust your prompt to only use read and search tools.'; - /** - * Helper to determine the error message and type for a policy denial. + * Helper to format the policy denial error. */ export function getPolicyDenialError( config: Config, rule?: PolicyRule, ): { errorMessage: string; errorType: ToolErrorType } { - if (config.getApprovalMode() === ApprovalMode.PLAN) { - return { - errorMessage: PLAN_MODE_DENIAL_MESSAGE, - errorType: ToolErrorType.STOP_EXECUTION, - }; - } - const denyMessage = rule?.denyMessage ? ` ${rule.denyMessage}` : ''; return { errorMessage: `Tool execution denied by policy.${denyMessage}`, diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index 7fd815a597..a3979f43a6 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -745,6 +745,63 @@ describe('Scheduler (Orchestrator)', () => { ); }); + it('should return POLICY_VIOLATION error type when denied in Plan Mode', async () => { + vi.mocked(checkPolicy).mockResolvedValue({ + decision: PolicyDecision.DENY, + rule: { decision: PolicyDecision.DENY }, + }); + + mockConfig.getApprovalMode.mockReturnValue(ApprovalMode.PLAN); + + await scheduler.schedule(req1, signal); + + expect(mockStateManager.updateStatus).toHaveBeenCalledWith( + 'call-1', + 'error', + expect.objectContaining({ + errorType: ToolErrorType.POLICY_VIOLATION, + responseParts: expect.arrayContaining([ + expect.objectContaining({ + functionResponse: expect.objectContaining({ + response: { + error: 'Tool execution denied by policy.', + }, + }), + }), + ]), + }), + ); + }); + + it('should return POLICY_VIOLATION and custom deny message when denied in Plan Mode with rule message', async () => { + const customMessage = 'Custom Plan Mode Deny'; + vi.mocked(checkPolicy).mockResolvedValue({ + decision: PolicyDecision.DENY, + rule: { decision: PolicyDecision.DENY, denyMessage: customMessage }, + }); + + mockConfig.getApprovalMode.mockReturnValue(ApprovalMode.PLAN); + + await scheduler.schedule(req1, signal); + + expect(mockStateManager.updateStatus).toHaveBeenCalledWith( + 'call-1', + 'error', + expect.objectContaining({ + errorType: ToolErrorType.POLICY_VIOLATION, + responseParts: expect.arrayContaining([ + expect.objectContaining({ + functionResponse: expect.objectContaining({ + response: { + error: `Tool execution denied by policy. ${customMessage}`, + }, + }), + }), + ]), + }), + ); + }); + it('should bypass confirmation and ProceedOnce if Policy returns ALLOW (YOLO/AllowedTools)', async () => { vi.mocked(checkPolicy).mockResolvedValue({ decision: PolicyDecision.ALLOW,