/**
 * @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,
} from './types.js';
import type { PolicyEngine } from '../policy/policy-engine.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 and publish policy update', async () => {
      const mockConfig = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'replace',
          persist: false,
        }),
      );
    });

    it('should preserve the original mode set when a session allow triggers AUTO_EDIT', async () => {
      let currentMode = ApprovalMode.DEFAULT;
      const mockConfig = {
        getApprovalMode: vi.fn(() => currentMode),
        setApprovalMode: vi.fn((mode: ApprovalMode) => {
          currentMode = mode;
        }),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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>;
      const tool = { name: 'replace' } as AnyDeclarativeTool;

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

      expect(mockConfig.setApprovalMode).toHaveBeenCalledWith(
        ApprovalMode.AUTO_EDIT,
      );
      expect(mockMessageBus.publish).toHaveBeenCalledWith(
        expect.objectContaining({
          type: MessageBusType.UPDATE_POLICY,
          toolName: 'replace',
          persist: false,
          modes: [
            ApprovalMode.DEFAULT,
            ApprovalMode.AUTO_EDIT,
            ApprovalMode.YOLO,
          ],
        }),
      );
    });

    it('should handle standard policy updates (persist=false)', async () => {
      const mockConfig = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        isTrustedFolder: vi.fn().mockReturnValue(false),
        getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        isTrustedFolder: vi.fn().mockReturnValue(false),
        getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        isTrustedFolder: vi.fn().mockReturnValue(true),
        getWorkspacePoliciesDir: vi
          .fn()
          .mockReturnValue('/mock/project/policies'),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        isTrustedFolder: vi.fn().mockReturnValue(false),
        getWorkspacePoliciesDir: vi
          .fn()
          .mockReturnValue('/mock/project/policies'),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        isTrustedFolder: vi.fn().mockReturnValue(false),
        getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
        getTargetDir: vi.fn().mockReturnValue('/mock/dir'),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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 = {
        getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
        setApprovalMode: vi.fn(),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } 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),
        getSessionId: vi.fn().mockReturnValue('test-session-id'),
      } as unknown as Config;

      (mockConfig as unknown as { config: Config }).config = mockConfig;
      const rule = {
        toolName: '*',
        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),
      getHookSystem: vi.fn().mockReturnValue(undefined),
      isInteractive: vi.fn().mockReturnValue(true),
      getEnableHooks: vi.fn().mockReturnValue(false),
      getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.PLAN), // Key: Plan Mode
      getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(false),
      getTelemetryTracesEnabled: vi.fn().mockReturnValue(false),
      setApprovalMode: vi.fn(),
      getSessionId: vi.fn().mockReturnValue('test-session-id'),
      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();
  });

  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;

    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;
    }

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

  describe('updatePolicy - context-aware modes', () => {
    const testCases = [
      {
        currentMode: ApprovalMode.DEFAULT,
        expectedModes: [
          ApprovalMode.DEFAULT,
          ApprovalMode.AUTO_EDIT,
          ApprovalMode.YOLO,
        ],
        description:
          'include current and more permissive modes in DEFAULT mode',
      },
      {
        currentMode: ApprovalMode.AUTO_EDIT,
        expectedModes: [ApprovalMode.AUTO_EDIT, ApprovalMode.YOLO],
        description:
          'include current and more permissive modes in AUTO_EDIT mode',
      },
      {
        currentMode: ApprovalMode.YOLO,
        expectedModes: [ApprovalMode.YOLO],
        description: 'include current and more permissive modes in YOLO mode',
      },
      {
        currentMode: ApprovalMode.PLAN,
        expectedModes: [
          ApprovalMode.PLAN,
          ApprovalMode.DEFAULT,
          ApprovalMode.AUTO_EDIT,
          ApprovalMode.YOLO,
        ],
        description: 'include all modes explicitly when granted in PLAN mode',
      },
    ];

    testCases.forEach(({ currentMode, expectedModes, description }) => {
      it(`should ${description}`, async () => {
        const mockConfig = {
          getApprovalMode: vi.fn().mockReturnValue(currentMode),
          isTrustedFolder: vi.fn().mockReturnValue(false),
          getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
          getSessionId: vi.fn().mockReturnValue('test-session-id'),
        } as unknown as Mocked<Config>;

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

        const context = {
          config: mockConfig,
          messageBus: mockMessageBus,
        } as unknown as AgentLoopContext;

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

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

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