gemini-cli/packages/core/src/scheduler/policy.test.ts

/**
 * @license
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import {
  describe,
  it,
  expect,
  vi,
  type Mocked,
  beforeEach,
  afterEach,
} from 'vitest';
import { checkPolicy, updatePolicy, getPolicyDenialError } from './policy.js';
import type { Config } from '../config/config.js';
import type { AgentLoopContext } from '../config/agent-loop-context.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import {
  MessageBusType,
  type SerializableConfirmationDetails,
} from '../confirmation-bus/types.js';
import { ApprovalMode, PolicyDecision } from '../policy/types.js';
import { escapeRegex } from '../policy/utils.js';
import {
  ToolConfirmationOutcome,
  type AnyDeclarativeTool,
  type ToolMcpConfirmationDetails,
  type ToolExecuteConfirmationDetails,
  type AnyToolInvocation,
} from '../tools/tools.js';
import {
  ROOT_SCHEDULER_ID,
  type ValidatingToolCall,
  type ToolCallRequestInfo,
  type CompletedToolCall,
} from './types.js';
import type { PolicyEngine } from '../policy/policy-engine.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { CoreToolScheduler } from '../core/coreToolScheduler.js';
import { Scheduler } from './scheduler.js';
import { ToolErrorType } from '../tools/tool-error.js';
import type { ToolRegistry } from '../tools/tool-registry.js';

describe('policy.ts', () => {
  describe('checkPolicy', () => {
    it('should return the decision from the policy engine', async () => {
      const mockPolicyEngine = {
        check: vi.fn().mockResolvedValue({ decision: PolicyDecision.ALLOW }),
      } as unknown as Mocked<PolicyEngine>;

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;

      const toolCall = {
        request: { name: 'test-tool', args: {} },
        tool: { name: 'test-tool' },
      } as ValidatingToolCall;

      const result = await checkPolicy(toolCall, mockConfig);
      expect(result.decision).toBe(PolicyDecision.ALLOW);
      expect(mockPolicyEngine.check).toHaveBeenCalledWith(
        { name: 'test-tool', args: {} },
        undefined,
        undefined,
        undefined,
      );
    });

    it('should pass serverName and toolAnnotations for MCP tools', async () => {
      const mockPolicyEngine = {
        check: vi.fn().mockResolvedValue({ decision: PolicyDecision.ALLOW }),
      } as unknown as Mocked<PolicyEngine>;

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;

      const mcpTool = Object.create(DiscoveredMCPTool.prototype);
      mcpTool.serverName = 'my-server';
      mcpTool._toolAnnotations = { readOnlyHint: true };

      const toolCall = {
        request: { name: 'mcp-tool', args: {} },
        tool: mcpTool,
      } as ValidatingToolCall;

      await checkPolicy(toolCall, mockConfig);
      expect(mockPolicyEngine.check).toHaveBeenCalledWith(
        { name: 'mcp-tool', args: {} },
        'my-server',
        { readOnlyHint: true },
        undefined,
      );
    });

    it('should respect disableAlwaysAllow from config', async () => {
      const mockPolicyEngine = {
        check: vi.fn().mockResolvedValue({ decision: PolicyDecision.ALLOW }),
      } as unknown as Mocked<PolicyEngine>;

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
        getDisableAlwaysAllow: vi.fn().mockReturnValue(true),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;

      const toolCall = {
        request: { name: 'test-tool', args: {} },
        tool: { name: 'test-tool' },
      } as ValidatingToolCall;

      // Note: checkPolicy calls config.getPolicyEngine().check()
      // The PolicyEngine itself is already configured with disableAlwaysAllow
      // when created in Config. Here we are just verifying that checkPolicy
      // doesn't somehow bypass it.
      await checkPolicy(toolCall, mockConfig);
      expect(mockPolicyEngine.check).toHaveBeenCalled();
    });

    it('should throw if ASK_USER is returned in non-interactive mode', async () => {
      const mockPolicyEngine = {
        check: vi.fn().mockResolvedValue({ decision: PolicyDecision.ASK_USER }),
      } as unknown as Mocked<PolicyEngine>;

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
        isInteractive: vi.fn().mockReturnValue(false),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;

      const toolCall = {
        request: { name: 'test-tool', args: {} },
        tool: { name: 'test-tool' },
      } as ValidatingToolCall;

      await expect(checkPolicy(toolCall, mockConfig)).rejects.toThrow(
        /not supported in non-interactive mode/,
      );
    });

    it('should return DENY without throwing', async () => {
      const mockPolicyEngine = {
        check: vi.fn().mockResolvedValue({ decision: PolicyDecision.DENY }),
      } as unknown as Mocked<PolicyEngine>;

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;

      const toolCall = {
        request: { name: 'test-tool', args: {} },
        tool: { name: 'test-tool' },
      } as ValidatingToolCall;

      const result = await checkPolicy(toolCall, mockConfig);
      expect(result.decision).toBe(PolicyDecision.DENY);
    });

    it('should return ASK_USER without throwing in interactive mode', async () => {
      const mockPolicyEngine = {
        check: vi.fn().mockResolvedValue({ decision: PolicyDecision.ASK_USER }),
      } as unknown as Mocked<PolicyEngine>;

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
        isInteractive: vi.fn().mockReturnValue(true),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;

      const toolCall = {
        request: { name: 'test-tool', args: {} },
        tool: { name: 'test-tool' },
      } as ValidatingToolCall;

      const result = await checkPolicy(toolCall, mockConfig);
      expect(result.decision).toBe(PolicyDecision.ASK_USER);
    });

    it('should return ALLOW if decision is ASK_USER and request is client-initiated', async () => {
      const mockPolicyEngine = {
        check: vi.fn().mockResolvedValue({ decision: PolicyDecision.ASK_USER }),
      } as unknown as Mocked<PolicyEngine>;

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
        isInteractive: vi.fn().mockReturnValue(true),
      } as unknown as Mocked<Config>;

      const toolCall = {
        request: { name: 'test-tool', args: {}, isClientInitiated: true },
        tool: { name: 'test-tool' },
      } as ValidatingToolCall;

      const result = await checkPolicy(toolCall, mockConfig);
      expect(result.decision).toBe(PolicyDecision.ALLOW);
    });

    it('should still return DENY if request is client-initiated but policy says DENY', async () => {
      const mockPolicyEngine = {
        check: vi.fn().mockResolvedValue({ decision: PolicyDecision.DENY }),
      } as unknown as Mocked<PolicyEngine>;

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
      } as unknown as Mocked<Config>;

      const toolCall = {
        request: { name: 'test-tool', args: {}, isClientInitiated: true },
        tool: { name: 'test-tool' },
      } as ValidatingToolCall;

      const result = await checkPolicy(toolCall, mockConfig);
      expect(result.decision).toBe(PolicyDecision.DENY);
    });
  });

  describe('updatePolicy', () => {
    it('should set AUTO_EDIT mode for auto-edit transition tools', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;

      const tool = { name: 'replace' } as AnyDeclarativeTool; // 'replace' is in EDIT_TOOL_NAMES

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlways,
        undefined,
        mockConfig,
        mockMessageBus,
      );

      expect(mockConfig.setApprovalMode).toHaveBeenCalledWith(
        ApprovalMode.AUTO_EDIT,
      );
      expect(mockMessageBus.publish).not.toHaveBeenCalled();
    });

    it('should handle standard policy updates (persist=false)', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'test-tool' } as AnyDeclarativeTool;

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlways,
        undefined,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'test-tool',
          persist: false,
        }),
      );
    });

    it('should handle standard policy updates with persistence', async () => {
      const mockConfig = {
        isTrustedFolder: vi.fn().mockReturnValue(false),
        getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'test-tool' } as AnyDeclarativeTool;

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlwaysAndSave,
        undefined,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'test-tool',
          persist: true,
        }),
      );
    });

    it('should handle shell command prefixes', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'run_shell_command' } as AnyDeclarativeTool;
      const details: ToolExecuteConfirmationDetails = {
        type: 'exec',
        command: 'ls -la',
        rootCommand: 'ls',
        rootCommands: ['ls'],
        title: 'Shell',
        onConfirm: vi.fn(),
      };

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlways,
        details,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'run_shell_command',
          commandPrefix: ['ls'],
        }),
      );
    });

    it('should handle MCP policy updates (server scope)', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'mcp-tool' } as AnyDeclarativeTool;
      const details: ToolMcpConfirmationDetails = {
        type: 'mcp',
        serverName: 'my-server',
        toolName: 'mcp-tool',
        toolDisplayName: 'My Tool',
        title: 'MCP',
        onConfirm: vi.fn(),
      };

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlwaysServer,
        details,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'mcp_my-server_*',
          mcpName: 'my-server',
          persist: false,
        }),
      );
    });

    it('should NOT publish update for ProceedOnce', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'test-tool' } as AnyDeclarativeTool;

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedOnce,
        undefined,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).not.toHaveBeenCalled();
      expect(mockConfig.setApprovalMode).not.toHaveBeenCalled();
    });

    it('should NOT publish update for Cancel', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'test-tool' } as AnyDeclarativeTool;

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.Cancel,
        undefined,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).not.toHaveBeenCalled();
    });

    it('should NOT publish update for ModifyWithEditor', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'test-tool' } as AnyDeclarativeTool;

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ModifyWithEditor,
        undefined,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).not.toHaveBeenCalled();
    });

    it('should handle MCP ProceedAlwaysTool (specific tool name)', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'mcp-tool' } as AnyDeclarativeTool;
      const details: ToolMcpConfirmationDetails = {
        type: 'mcp',
        serverName: 'my-server',
        toolName: 'mcp-tool',
        toolDisplayName: 'My Tool',
        title: 'MCP',
        onConfirm: vi.fn(),
      };

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlwaysTool,
        details,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'mcp-tool', // Specific name, not wildcard
          mcpName: 'my-server',
          persist: false,
        }),
      );
    });

    it('should handle MCP ProceedAlways (persist: false)', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'mcp-tool' } as AnyDeclarativeTool;
      const details: ToolMcpConfirmationDetails = {
        type: 'mcp',
        serverName: 'my-server',
        toolName: 'mcp-tool',
        toolDisplayName: 'My Tool',
        title: 'MCP',
        onConfirm: vi.fn(),
      };

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlways,
        details,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'mcp-tool',
          mcpName: 'my-server',
          persist: false,
        }),
      );
    });

    it('should handle MCP ProceedAlwaysAndSave (persist: true)', async () => {
      const mockConfig = {
        isTrustedFolder: vi.fn().mockReturnValue(false),
        getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;

      (mockConfig as unknown as { config: Config }).config =
        mockConfig as Config;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
        mockMessageBus;
      const tool = { name: 'mcp-tool' } as AnyDeclarativeTool;
      const details: ToolMcpConfirmationDetails = {
        type: 'mcp',
        serverName: 'my-server',
        toolName: 'mcp-tool',
        toolDisplayName: 'My Tool',
        title: 'MCP',
        onConfirm: vi.fn(),
      };

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlwaysAndSave,
        details,
        mockConfig,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'mcp-tool',
          mcpName: 'my-server',
          persist: true,
        }),
      );
    });

    it('should determine persistScope: workspace in trusted folders', async () => {
      const mockConfig = {
        isTrustedFolder: vi.fn().mockReturnValue(true),
        getWorkspacePoliciesDir: vi
          .fn()
          .mockReturnValue('/mock/project/policies'),
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      const tool = { name: 'test-tool' } as AnyDeclarativeTool;

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlwaysAndSave,
        undefined,
        {
          config: mockConfig,
        } as unknown as AgentLoopContext,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          persistScope: 'workspace',
        }),
      );
    });

    it('should determine persistScope: user in untrusted folders', async () => {
      const mockConfig = {
        isTrustedFolder: vi.fn().mockReturnValue(false),
        getWorkspacePoliciesDir: vi
          .fn()
          .mockReturnValue('/mock/project/policies'),
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      const tool = { name: 'test-tool' } as AnyDeclarativeTool;

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlwaysAndSave,
        undefined,
        {
          config: mockConfig,
        } as unknown as AgentLoopContext,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          persistScope: 'user',
        }),
      );
    });

    it('should narrow edit tools with argsPattern', async () => {
      const mockConfig = {
        isTrustedFolder: vi.fn().mockReturnValue(false),
        getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
        getTargetDir: vi.fn().mockReturnValue('/mock/dir'),
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;
      const tool = { name: 'write_file' } as AnyDeclarativeTool;
      const details: SerializableConfirmationDetails = {
        type: 'edit',
        title: 'Edit',
        filePath: 'src/foo.ts',
        fileName: 'foo.ts',
        fileDiff: '--- foo.ts\n+++ foo.ts\n@@ -1 +1 @@\n-old\n+new',
        originalContent: 'old',
        newContent: 'new',
      };

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlwaysAndSave,
        details,
        {
          config: mockConfig,
        } as unknown as AgentLoopContext,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          toolName: 'write_file',
          argsPattern:
            '\\\\0' + escapeRegex('"file_path":"src/foo.ts"') + '\\\\0',
        }),
      );
    });

    it('should work when context is created via Object.create (prototype chain)', async () => {
      const mockConfig = {
        setApprovalMode: vi.fn(),
      } as unknown as Mocked<Config>;
      const mockMessageBus = {
        publish: vi.fn(),
      } as unknown as Mocked<MessageBus>;

      const baseContext = {
        config: mockConfig,
        messageBus: mockMessageBus,
      };
      const protoContext: AgentLoopContext = Object.create(baseContext);

      expect(Object.keys(protoContext)).toHaveLength(0);
      expect(protoContext.config).toBe(mockConfig);
      expect(protoContext.messageBus).toBe(mockMessageBus);

      const tool = { name: 'test-tool' } as AnyDeclarativeTool;

      await updatePolicy(
        tool,
        ToolConfirmationOutcome.ProceedAlways,
        undefined,
        protoContext,
        mockMessageBus,
      );

      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'test-tool',
          persist: false,
        }),
      );
    });
  });

  describe('getPolicyDenialError', () => {
    it('should return default denial message when no rule provided', () => {
      const mockConfig = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
      } as unknown as Config;

      (mockConfig as unknown as { config: Config }).config = mockConfig;

      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;

      (mockConfig as unknown as { config: Config }).config = mockConfig;
      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', () => {
  let mockConfig: Mocked<Config>;
  let mockMessageBus: Mocked<MessageBus>;
  let mockPolicyEngine: Mocked<PolicyEngine>;
  let mockToolRegistry: Mocked<ToolRegistry>;
  let mockTool: AnyDeclarativeTool;
  let mockInvocation: AnyToolInvocation;

  const req: ToolCallRequestInfo = {
    callId: 'call-1',
    name: 'test-tool',
    args: { foo: 'bar' },
    isClientInitiated: false,
    prompt_id: 'prompt-1',
    schedulerId: ROOT_SCHEDULER_ID,
  };

  beforeEach(() => {
    mockTool = {
      name: 'test-tool',
      build: vi.fn(),
    } as unknown as AnyDeclarativeTool;

    mockInvocation = {
      shouldConfirmExecute: vi.fn(),
    } as unknown as AnyToolInvocation;
    vi.mocked(mockTool.build).mockReturnValue(mockInvocation);

    mockPolicyEngine = {
      check: vi.fn().mockResolvedValue({ decision: PolicyDecision.DENY }), // Default to DENY for this test
    } as unknown as Mocked<PolicyEngine>;

    mockToolRegistry = {
      getTool: vi.fn().mockReturnValue(mockTool),
      getAllToolNames: vi.fn().mockReturnValue(['test-tool']),
    } as unknown as Mocked<ToolRegistry>;

    mockMessageBus = {
      publish: vi.fn(),
      subscribe: vi.fn(),
    } as unknown as Mocked<MessageBus>;

    mockConfig = {
      getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
      toolRegistry: mockToolRegistry,
      getToolRegistry: () => mockToolRegistry,
      getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
      isInteractive: vi.fn().mockReturnValue(true),
      getEnableHooks: vi.fn().mockReturnValue(false),
      getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.PLAN), // Key: Plan Mode
      setApprovalMode: vi.fn(),
      getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
    } as unknown as Mocked<Config>;
    (mockConfig as unknown as { config: Config }).config = mockConfig as Config;
    (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
      mockMessageBus;
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  describe.each([
    { enableEventDrivenScheduler: false, name: 'Legacy CoreToolScheduler' },
    { enableEventDrivenScheduler: true, name: 'Event-Driven Scheduler' },
  ])('$name', ({ enableEventDrivenScheduler }) => {
    it('should return the correct Plan Mode denial message when policy denies execution', async () => {
      let resultMessage: string | undefined;
      let resultErrorType: ToolErrorType | undefined;

      const signal = new AbortController().signal;

      if (enableEventDrivenScheduler) {
        const scheduler = new Scheduler({
          context: {
            config: mockConfig,
            messageBus: mockMessageBus,
            toolRegistry: mockToolRegistry,
          } as unknown as AgentLoopContext,
          getPreferredEditor: () => undefined,
          schedulerId: ROOT_SCHEDULER_ID,
        });

        const results = await scheduler.schedule(req, signal);
        const result = results[0];

        expect(result.status).toBe('error');
        if (result.status === 'error') {
          resultMessage = result.response.error?.message;
          resultErrorType = result.response.errorType;
        }
      } else {
        let capturedCalls: CompletedToolCall[] = [];
        const scheduler = new CoreToolScheduler({
          context: {
            config: mockConfig,
            messageBus: mockMessageBus,
            toolRegistry: mockToolRegistry,
          } as unknown as AgentLoopContext,
          getPreferredEditor: () => undefined,
          onAllToolCallsComplete: async (calls) => {
            capturedCalls = calls;
          },
        });

        await scheduler.schedule(req, signal);

        expect(capturedCalls.length).toBeGreaterThan(0);
        const call = capturedCalls[0];
        if (call.status === 'error') {
          resultMessage = call.response.error?.message;
          resultErrorType = call.response.errorType;
        }
      }

      expect(resultMessage).toBe('Tool execution denied by policy.');
      expect(resultErrorType).toBe(ToolErrorType.POLICY_VIOLATION);
    });
  });
});