diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 154b975638..a250eb0f81 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -7,10 +7,8 @@ import { describe, it, expect, vi } from 'vitest'; import type { Mock } from 'vitest'; import type { CallableTool } from '@google/genai'; -import { - CoreToolScheduler, - PLAN_MODE_DENIAL_MESSAGE, -} from './coreToolScheduler.js'; +import { CoreToolScheduler } from './coreToolScheduler.js'; +import { PLAN_MODE_DENIAL_MESSAGE } from '../scheduler/policy.js'; import type { ToolCall, WaitingToolCall, diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 124cef32b9..c8647ffc99 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -14,7 +14,7 @@ import { } from '../tools/tools.js'; import type { EditorType } from '../utils/editor.js'; import type { Config } from '../config/config.js'; -import { PolicyDecision, ApprovalMode } from '../policy/types.js'; +import { PolicyDecision } from '../policy/types.js'; import { logToolCall } from '../telemetry/loggers.js'; import { ToolErrorType } from '../tools/tool-error.js'; import { ToolCallEvent } from '../telemetry/types.js'; @@ -44,6 +44,7 @@ import { } from '../scheduler/types.js'; import { ToolExecutor } from '../scheduler/tool-executor.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; +import { getPolicyDenialError } from '../scheduler/policy.js'; export type { ToolCall, @@ -64,9 +65,6 @@ export type { ToolCallResponseInfo, }; -export const PLAN_MODE_DENIAL_MESSAGE = - 'You are in Plan Mode - adjust your prompt to only use read and search tools.'; - const createErrorResponse = ( request: ToolCallRequestInfo, error: Error, @@ -599,18 +597,15 @@ export class CoreToolScheduler { ? toolCall.tool.serverName : undefined; - const { decision } = await this.config + const { decision, rule } = await this.config .getPolicyEngine() .check(toolCallForPolicy, serverName); if (decision === PolicyDecision.DENY) { - let errorMessage = `Tool execution denied by policy.`; - let errorType = ToolErrorType.POLICY_VIOLATION; - - if (this.config.getApprovalMode() === ApprovalMode.PLAN) { - errorMessage = PLAN_MODE_DENIAL_MESSAGE; - errorType = ToolErrorType.STOP_EXECUTION; - } + const { errorMessage, errorType } = getPolicyDenialError( + this.config, + rule, + ); this.setStatusInternal( reqInfo.callId, 'error', diff --git a/packages/core/src/scheduler/policy.ts b/packages/core/src/scheduler/policy.ts index d28ca6dad6..279dea85c7 100644 --- a/packages/core/src/scheduler/policy.ts +++ b/packages/core/src/scheduler/policy.ts @@ -4,10 +4,12 @@ * SPDX-License-Identifier: Apache-2.0 */ +import { ToolErrorType } from '../tools/tool-error.js'; import { ApprovalMode, PolicyDecision, type CheckResult, + type PolicyRule, } from '../policy/types.js'; import type { Config } from '../config/config.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; @@ -24,6 +26,30 @@ 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. + */ +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}`, + errorType: ToolErrorType.POLICY_VIOLATION, + }; +} + /** * Queries the system PolicyEngine to determine tool allowance. * @returns The PolicyDecision. diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index 4ae3e84c8c..7fd815a597 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -46,7 +46,14 @@ import { ToolModificationHandler } from './tool-modifier.js'; vi.mock('./state-manager.js'); vi.mock('./confirmation.js'); -vi.mock('./policy.js'); +vi.mock('./policy.js', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + checkPolicy: vi.fn(), + updatePolicy: vi.fn(), + }; +}); vi.mock('./tool-executor.js'); vi.mock('./tool-modifier.js'); @@ -55,7 +62,7 @@ import type { Config } from '../config/config.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; import type { PolicyEngine } from '../policy/policy-engine.js'; import type { ToolRegistry } from '../tools/tool-registry.js'; -import { PolicyDecision } from '../policy/types.js'; +import { PolicyDecision, ApprovalMode } from '../policy/types.js'; import { ToolConfirmationOutcome, type AnyDeclarativeTool, @@ -149,6 +156,7 @@ describe('Scheduler (Orchestrator)', () => { isInteractive: vi.fn().mockReturnValue(true), getEnableHooks: vi.fn().mockReturnValue(true), setApprovalMode: vi.fn(), + getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT), } as unknown as Mocked; mockMessageBus = { diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 0589c50a72..71729923d0 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -8,7 +8,7 @@ import type { Config } from '../config/config.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { SchedulerStateManager } from './state-manager.js'; import { resolveConfirmation } from './confirmation.js'; -import { checkPolicy, updatePolicy } from './policy.js'; +import { checkPolicy, updatePolicy, getPolicyDenialError } from './policy.js'; import { ToolExecutor } from './tool-executor.js'; import { ToolModificationHandler } from './tool-modifier.js'; import { @@ -407,14 +407,18 @@ export class Scheduler { const { decision, rule } = await checkPolicy(toolCall, this.config); if (decision === PolicyDecision.DENY) { - const denyMessage = rule?.denyMessage ? ` ${rule.denyMessage}` : ''; + const { errorMessage, errorType } = getPolicyDenialError( + this.config, + rule, + ); + this.state.updateStatus( callId, 'error', createErrorResponse( toolCall.request, - new Error(`Tool execution denied by policy.${denyMessage}`), - ToolErrorType.POLICY_VIOLATION, + new Error(errorMessage), + errorType, ), ); this.state.finalizeCall(callId);