workspace-policy-cli.test.ts

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

import { describe, it, expect, vi, beforeEach } from 'vitest';
import * as path from 'node:path';
import { loadCliConfig, type CliArgs } from './config.js';
import { createTestMergedSettings } from './settings.js';
import * as ServerConfig from '@google/gemini-cli-core';
import { isWorkspaceTrusted } from './trustedFolders.js';
import * as Policy from './policy.js';

// Mock dependencies
vi.mock('./trustedFolders.js', () => ({
  isWorkspaceTrusted: vi.fn(),
}));

const mockCheckIntegrity = vi.fn();
const mockAcceptIntegrity = vi.fn();

vi.mock('@google/gemini-cli-core', async () => {
  const actual = await vi.importActual<typeof ServerConfig>(
    '@google/gemini-cli-core',
  );
  return {
    ...actual,
    loadServerHierarchicalMemory: vi.fn().mockResolvedValue({
      memoryContent: '',
      fileCount: 0,
      filePaths: [],
    }),
    createPolicyEngineConfig: vi.fn().mockResolvedValue({
      rules: [],
      checkers: [],
    }),
    getVersion: vi.fn().mockResolvedValue('test-version'),
    PolicyIntegrityManager: vi.fn().mockImplementation(() => ({
      checkIntegrity: mockCheckIntegrity,
      acceptIntegrity: mockAcceptIntegrity,
    })),
    IntegrityStatus: { MATCH: 'match', NEW: 'new', MISMATCH: 'mismatch' },
    debugLogger: {
      warn: vi.fn(),
      error: vi.fn(),
    },
    isHeadlessMode: vi.fn().mockReturnValue(false), // Default to interactive
  };
});

describe('Workspace-Level Policy CLI Integration', () => {
  const MOCK_CWD = process.cwd();

  beforeEach(() => {
    vi.clearAllMocks();
    Policy.setDisableWorkspacePolicies(false);
    // Default to MATCH for existing tests
    mockCheckIntegrity.mockResolvedValue({
      status: 'match',
      hash: 'test-hash',
      fileCount: 1,
    });
    vi.mocked(ServerConfig.isHeadlessMode).mockReturnValue(false);
  });

  it('should have getWorkspacePoliciesDir on Storage class', () => {
    const storage = new ServerConfig.Storage(MOCK_CWD);
    expect(storage.getWorkspacePoliciesDir).toBeDefined();
    expect(typeof storage.getWorkspacePoliciesDir).toBe('function');
  });

  it('should pass workspacePoliciesDir to createPolicyEngineConfig when folder is trusted', async () => {
    vi.mocked(isWorkspaceTrusted).mockReturnValue({
      isTrusted: true,
      source: 'file',
    });

    const settings = createTestMergedSettings();
    const argv = { query: 'test' } as unknown as CliArgs;

    await loadCliConfig(settings, 'test-session', argv, { cwd: MOCK_CWD });

    expect(ServerConfig.createPolicyEngineConfig).toHaveBeenCalledWith(
      expect.objectContaining({
        workspacePoliciesDir: expect.stringContaining(
          path.join('.gemini', 'policies'),
        ),
      }),
      expect.anything(),
    );
  });

  it('should NOT pass workspacePoliciesDir to createPolicyEngineConfig when folder is NOT trusted', async () => {
    vi.mocked(isWorkspaceTrusted).mockReturnValue({
      isTrusted: false,
      source: 'file',
    });

    const settings = createTestMergedSettings();
    const argv = { query: 'test' } as unknown as CliArgs;

    await loadCliConfig(settings, 'test-session', argv, { cwd: MOCK_CWD });

    expect(ServerConfig.createPolicyEngineConfig).toHaveBeenCalledWith(
      expect.objectContaining({
        workspacePoliciesDir: undefined,
      }),
      expect.anything(),
    );
  });

  it('should NOT pass workspacePoliciesDir if integrity is NEW but fileCount is 0', async () => {
    vi.mocked(isWorkspaceTrusted).mockReturnValue({
      isTrusted: true,
      source: 'file',
    });
    mockCheckIntegrity.mockResolvedValue({
      status: 'new',
      hash: 'hash',
      fileCount: 0,
    });

    const settings = createTestMergedSettings();
    const argv = { query: 'test' } as unknown as CliArgs;

    await loadCliConfig(settings, 'test-session', argv, { cwd: MOCK_CWD });

    expect(ServerConfig.createPolicyEngineConfig).toHaveBeenCalledWith(
      expect.objectContaining({
        workspacePoliciesDir: undefined,
      }),
      expect.anything(),
    );
  });

  it('should automatically accept and load workspacePoliciesDir if integrity MISMATCH in non-interactive mode', async () => {
    vi.mocked(isWorkspaceTrusted).mockReturnValue({
      isTrusted: true,
      source: 'file',
    });
    mockCheckIntegrity.mockResolvedValue({
      status: 'mismatch',
      hash: 'new-hash',
      fileCount: 1,
    });
    vi.mocked(ServerConfig.isHeadlessMode).mockReturnValue(true); // Non-interactive

    const settings = createTestMergedSettings();
    const argv = { prompt: 'do something' } as unknown as CliArgs;

    await loadCliConfig(settings, 'test-session', argv, { cwd: MOCK_CWD });

    expect(mockAcceptIntegrity).toHaveBeenCalledWith(
      'workspace',
      MOCK_CWD,
      'new-hash',
    );
    expect(ServerConfig.createPolicyEngineConfig).toHaveBeenCalledWith(
      expect.objectContaining({
        workspacePoliciesDir: expect.stringContaining(
          path.join('.gemini', 'policies'),
        ),
      }),
      expect.anything(),
    );
  });

  it('should automatically accept and load workspacePoliciesDir if integrity MISMATCH in interactive mode when AUTO_ACCEPT is true', async () => {
    vi.mocked(isWorkspaceTrusted).mockReturnValue({
      isTrusted: true,
      source: 'file',
    });
    mockCheckIntegrity.mockResolvedValue({
      status: 'mismatch',
      hash: 'new-hash',
      fileCount: 1,
    });
    vi.mocked(ServerConfig.isHeadlessMode).mockReturnValue(false); // Interactive

    const settings = createTestMergedSettings();
    const argv = {
      query: 'test',
      promptInteractive: 'test',
    } as unknown as CliArgs;

    const config = await loadCliConfig(settings, 'test-session', argv, {
      cwd: MOCK_CWD,
    });

    expect(config.getPolicyUpdateConfirmationRequest()).toBeUndefined();
    expect(mockAcceptIntegrity).toHaveBeenCalledWith(
      'workspace',
      MOCK_CWD,
      'new-hash',
    );
    expect(ServerConfig.createPolicyEngineConfig).toHaveBeenCalledWith(
      expect.objectContaining({
        workspacePoliciesDir: expect.stringContaining(
          path.join('.gemini', 'policies'),
        ),
      }),
      expect.anything(),
    );
  });

  it('should automatically accept and load workspacePoliciesDir if integrity is NEW in interactive mode when AUTO_ACCEPT is true', async () => {
    vi.mocked(isWorkspaceTrusted).mockReturnValue({
      isTrusted: true,
      source: 'file',
    });
    mockCheckIntegrity.mockResolvedValue({
      status: 'new',
      hash: 'new-hash',
      fileCount: 5,
    });
    vi.mocked(ServerConfig.isHeadlessMode).mockReturnValue(false); // Interactive

    const settings = createTestMergedSettings();
    const argv = { query: 'test' } as unknown as CliArgs;

    const config = await loadCliConfig(settings, 'test-session', argv, {
      cwd: MOCK_CWD,
    });

    expect(config.getPolicyUpdateConfirmationRequest()).toBeUndefined();
    expect(mockAcceptIntegrity).toHaveBeenCalledWith(
      'workspace',
      MOCK_CWD,
      'new-hash',
    );

    expect(ServerConfig.createPolicyEngineConfig).toHaveBeenCalledWith(
      expect.objectContaining({
        workspacePoliciesDir: expect.stringContaining(
          path.join('.gemini', 'policies'),
        ),
      }),
      expect.anything(),
    );
  });

  it('should set policyUpdateConfirmationRequest if integrity MISMATCH in interactive mode when AUTO_ACCEPT is false', async () => {
    // Monkey patch autoAcceptWorkspacePolicies using setter
    const originalValue = Policy.autoAcceptWorkspacePolicies;
    Policy.setAutoAcceptWorkspacePolicies(false);

    try {
      vi.mocked(isWorkspaceTrusted).mockReturnValue({
        isTrusted: true,
        source: 'file',
      });
      mockCheckIntegrity.mockResolvedValue({
        status: 'mismatch',
        hash: 'new-hash',
        fileCount: 1,
      });
      vi.mocked(ServerConfig.isHeadlessMode).mockReturnValue(false); // Interactive

      const settings = createTestMergedSettings();
      const argv = {
        query: 'test',
        promptInteractive: 'test',
      } as unknown as CliArgs;

      const config = await loadCliConfig(settings, 'test-session', argv, {
        cwd: MOCK_CWD,
      });

      expect(config.getPolicyUpdateConfirmationRequest()).toEqual({
        scope: 'workspace',
        identifier: MOCK_CWD,
        policyDir: expect.stringContaining(path.join('.gemini', 'policies')),
        newHash: 'new-hash',
      });
      expect(ServerConfig.createPolicyEngineConfig).toHaveBeenCalledWith(
        expect.objectContaining({
          workspacePoliciesDir: undefined,
        }),
        expect.anything(),
      );
    } finally {
      // Restore for other tests
      Policy.setAutoAcceptWorkspacePolicies(originalValue);
    }
  });
});