feat(policy): implement hot-reloading for workspace policies

This change eliminates the need for a CLI restart when a user accepts new or changed project-level policies. Workspace rules are now dynamically injected into the active PolicyEngine instance. Key improvements: - Added Config.loadWorkspacePolicies() to handle mid-session rule injection. - Fully encapsulated acceptance and integrity logic within PolicyUpdateDialog. - Integrated centralized keybindings (Command.ESCAPE) for dialog dismissal. - Refactored PolicyIntegrityManager tests to use a real temporary directory instead of filesystem mocks for improved reliability. - Updated copyright headers to 2026 across affected files. - Added UI snapshot tests for the policy update dialog. Addresses review feedback from PR #18682.
2026-03-20 19:11:23 -07:00 · 2026-02-18 12:54:47 -08:00
parent 8feff1cc9b
commit 662654c5d2
13 changed files with 248 additions and 269 deletions
--- a/packages/cli/src/test-utils/render.tsx
+++ b/packages/cli/src/test-utils/render.tsx
@@ -199,7 +199,7 @@ const mockUIActions: UIActions = {
  vimHandleInput: vi.fn(),
  handleIdePromptComplete: vi.fn(),
  handleFolderTrustSelect: vi.fn(),
-  handlePolicyUpdateSelect: vi.fn(),
+  setIsPolicyUpdateDialogOpen: vi.fn(),
  setConstrainHeight: vi.fn(),
  onEscapePromptChange: vi.fn(),
  refreshStatic: vi.fn(),
--- a/packages/cli/src/ui/AppContainer.tsx
+++ b/packages/cli/src/ui/AppContainer.tsx
@@ -122,7 +122,7 @@ import { appEvents, AppEvent, TransientMessageType } from '../utils/events.js';
 import { type UpdateObject } from './utils/updateCheck.js';
 import { setUpdateHandler } from '../utils/handleAutoUpdate.js';
 import { registerCleanup, runExitCleanup } from '../utils/cleanup.js';
-import { RELAUNCH_EXIT_CODE, relaunchApp } from '../utils/processUtils.js';
+import { RELAUNCH_EXIT_CODE } from '../utils/processUtils.js';
 import type { SessionInfo } from '../utils/sessionUtils.js';
 import { useMessageQueue } from './hooks/useMessageQueue.js';
 import { useMcpStatus } from './hooks/useMcpStatus.js';
@@ -154,7 +154,6 @@ import {
 } from './constants.js';
 import { LoginWithGoogleRestartDialog } from './auth/LoginWithGoogleRestartDialog.js';
 import { NewAgentsChoice } from './components/NewAgentsNotification.js';
-import { PolicyUpdateChoice } from './components/PolicyUpdateDialog.js';
 import { isSlashCommand } from './utils/commandUtils.js';
 import { useTerminalTheme } from './hooks/useTerminalTheme.js';
 import { useTimedMessage } from './hooks/useTimedMessage.js';
@@ -1446,32 +1445,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
  const [isPolicyUpdateDialogOpen, setIsPolicyUpdateDialogOpen] = useState(
    !!policyUpdateConfirmationRequest,
  );
-  const [isRestartingPolicyUpdate, setIsRestartingPolicyUpdate] =
-    useState(false);
-
-  const handlePolicyUpdateSelect = useCallback(
-    async (choice: PolicyUpdateChoice) => {
-      if (
-        choice === PolicyUpdateChoice.ACCEPT &&
-        policyUpdateConfirmationRequest
-      ) {
-        const integrityManager = new PolicyIntegrityManager();
-        await integrityManager.acceptIntegrity(
-          policyUpdateConfirmationRequest.scope,
-          policyUpdateConfirmationRequest.identifier,
-          policyUpdateConfirmationRequest.newHash,
-        );
-        setIsRestartingPolicyUpdate(true);
-        // Give time for the UI to render the restarting message
-        setTimeout(async () => {
-          await relaunchApp();
-        }, 250);
-      } else {
-        setIsPolicyUpdateDialogOpen(false);
-      }
-    },
-    [policyUpdateConfirmationRequest],
-  );

  const {
    needsRestart: ideNeedsRestart,
@@ -2173,7 +2146,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
      isFolderTrustDialogOpen: isFolderTrustDialogOpen ?? false,
      isPolicyUpdateDialogOpen,
      policyUpdateConfirmationRequest,
-      isRestartingPolicyUpdate,
      isTrustedFolder,
      constrainHeight,
      showErrorDetails,
@@ -2298,7 +2270,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
      isFolderTrustDialogOpen,
      isPolicyUpdateDialogOpen,
      policyUpdateConfirmationRequest,
-      isRestartingPolicyUpdate,
      isTrustedFolder,
      constrainHeight,
      showErrorDetails,
@@ -2396,7 +2367,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
      vimHandleInput,
      handleIdePromptComplete,
      handleFolderTrustSelect,
-      handlePolicyUpdateSelect,
+      setIsPolicyUpdateDialogOpen,
      setConstrainHeight,
      onEscapePromptChange: handleEscapePromptChange,
      refreshStatic,
@@ -2481,7 +2452,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
      vimHandleInput,
      handleIdePromptComplete,
      handleFolderTrustSelect,
-      handlePolicyUpdateSelect,
+      setIsPolicyUpdateDialogOpen,
      setConstrainHeight,
      handleEscapePromptChange,
      refreshStatic,
--- a/packages/cli/src/ui/components/DialogManager.tsx
+++ b/packages/cli/src/ui/components/DialogManager.tsx
@@ -167,16 +167,12 @@ export const DialogManager = ({
      />
    );
  }
-  if (
-    uiState.isPolicyUpdateDialogOpen &&
-    uiState.policyUpdateConfirmationRequest
-  ) {
+  if (uiState.isPolicyUpdateDialogOpen) {
    return (
      <PolicyUpdateDialog
-        onSelect={uiActions.handlePolicyUpdateSelect}
-        scope={uiState.policyUpdateConfirmationRequest.scope}
-        identifier={uiState.policyUpdateConfirmationRequest.policyDir}
-        isRestarting={uiState.isRestartingPolicyUpdate}
+        config={config}
+        request={uiState.policyUpdateConfirmationRequest!}
+        onClose={() => uiActions.setIsPolicyUpdateDialogOpen(false)}
      />
    );
  }
--- a/packages/cli/src/ui/components/PolicyUpdateDialog.test.tsx
+++ b/packages/cli/src/ui/components/PolicyUpdateDialog.test.tsx
@@ -4,46 +4,76 @@
 * SPDX-License-Identifier: Apache-2.0
 */

-import { describe, it, expect, vi, afterEach } from 'vitest';
+import { describe, it, expect, vi, afterEach, beforeEach } from 'vitest';
 import { act } from 'react';
 import { renderWithProviders } from '../../test-utils/render.js';
 import { waitFor } from '../../test-utils/async.js';
+import { PolicyUpdateDialog } from './PolicyUpdateDialog.js';
 import {
-  PolicyUpdateDialog,
-  PolicyUpdateChoice,
-} from './PolicyUpdateDialog.js';
+  type Config,
+  type PolicyUpdateConfirmationRequest,
+  PolicyIntegrityManager,
+} from '@google/gemini-cli-core';
+
+// Mock PolicyIntegrityManager
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const original = (await importOriginal()) as any;
+  return {
+    ...original,
+    PolicyIntegrityManager: vi.fn().mockImplementation(() => ({
+      acceptIntegrity: vi.fn().mockResolvedValue(undefined),
+    })),
+  };
+});

 describe('PolicyUpdateDialog', () => {
+  let mockConfig: Config;
+  let mockRequest: PolicyUpdateConfirmationRequest;
+  let onClose: () => void;
+
+  beforeEach(() => {
+    mockConfig = {
+      loadWorkspacePolicies: vi.fn().mockResolvedValue(undefined),
+    } as unknown as Config;
+
+    mockRequest = {
+      scope: 'workspace',
+      identifier: '/test/workspace/.gemini/policies',
+      policyDir: '/test/workspace/.gemini/policies',
+      newHash: 'test-hash',
+    } as PolicyUpdateConfirmationRequest;
+
+    onClose = vi.fn();
+  });
+
  afterEach(() => {
    vi.clearAllMocks();
  });

-  it('renders correctly with default props', () => {
-    const onSelect = vi.fn();
+  it('renders correctly and matches snapshot', () => {
    const { lastFrame } = renderWithProviders(
      <PolicyUpdateDialog
-        onSelect={onSelect}
-        scope="workspace"
-        identifier="/test/path"
-        isRestarting={false}
+        config={mockConfig}
+        request={mockRequest}
+        onClose={onClose}
      />,
    );

    const output = lastFrame();
+    expect(output).toMatchSnapshot();
    expect(output).toContain('New or changed workspace policies detected');
-    expect(output).toContain('Location: /test/path');
+    expect(output).toContain('Location: /test/workspace/.gemini/policies');
    expect(output).toContain('Accept and Load');
    expect(output).toContain('Ignore');
  });

-  it('calls onSelect with ACCEPT when accept option is chosen', async () => {
-    const onSelect = vi.fn();
+  it('handles ACCEPT correctly', async () => {
    const { stdin } = renderWithProviders(
      <PolicyUpdateDialog
-        onSelect={onSelect}
-        scope="workspace"
-        identifier="/test/path"
-        isRestarting={false}
+        config={mockConfig}
+        request={mockRequest}
+        onClose={onClose}
      />,
    );

@@ -53,18 +83,20 @@ describe('PolicyUpdateDialog', () => {
    });

    await waitFor(() => {
-      expect(onSelect).toHaveBeenCalledWith(PolicyUpdateChoice.ACCEPT);
+      expect(PolicyIntegrityManager).toHaveBeenCalled();
+      expect(mockConfig.loadWorkspacePolicies).toHaveBeenCalledWith(
+        mockRequest.policyDir,
+      );
+      expect(onClose).toHaveBeenCalled();
    });
  });

-  it('calls onSelect with IGNORE when ignore option is chosen', async () => {
-    const onSelect = vi.fn();
+  it('handles IGNORE correctly', async () => {
    const { stdin } = renderWithProviders(
      <PolicyUpdateDialog
-        onSelect={onSelect}
-        scope="workspace"
-        identifier="/test/path"
-        isRestarting={false}
+        config={mockConfig}
+        request={mockRequest}
+        onClose={onClose}
      />,
    );

@@ -77,44 +109,27 @@ describe('PolicyUpdateDialog', () => {
    });

    await waitFor(() => {
-      expect(onSelect).toHaveBeenCalledWith(PolicyUpdateChoice.IGNORE);
+      expect(PolicyIntegrityManager).not.toHaveBeenCalled();
+      expect(mockConfig.loadWorkspacePolicies).not.toHaveBeenCalled();
+      expect(onClose).toHaveBeenCalled();
    });
  });

-  it('calls onSelect with IGNORE when Escape is pressed', async () => {
-    const onSelect = vi.fn();
+  it('calls onClose when Escape key is pressed', async () => {
    const { stdin } = renderWithProviders(
      <PolicyUpdateDialog
-        onSelect={onSelect}
-        scope="workspace"
-        identifier="/test/path"
-        isRestarting={false}
+        config={mockConfig}
+        request={mockRequest}
+        onClose={onClose}
      />,
    );

    await act(async () => {
-      stdin.write('\x1B'); // Escape key
+      stdin.write('\x1B'); // Escape key (matches Command.ESCAPE default)
    });

    await waitFor(() => {
-      expect(onSelect).toHaveBeenCalledWith(PolicyUpdateChoice.IGNORE);
+      expect(onClose).toHaveBeenCalled();
    });
  });
-
-  it('displays restarting message when isRestarting is true', () => {
-    const onSelect = vi.fn();
-    const { lastFrame } = renderWithProviders(
-      <PolicyUpdateDialog
-        onSelect={onSelect}
-        scope="workspace"
-        identifier="/test/path"
-        isRestarting={true}
-      />,
-    );
-
-    const output = lastFrame();
-    expect(output).toContain(
-      'Gemini CLI is restarting to apply the policy changes...',
-    );
-  });
 });
--- a/packages/cli/src/ui/components/PolicyUpdateDialog.tsx
+++ b/packages/cli/src/ui/components/PolicyUpdateDialog.tsx
@@ -5,11 +5,18 @@
 */

 import { Box, Text } from 'ink';
+import { useCallback } from 'react';
 import type React from 'react';
+import {
+  type Config,
+  type PolicyUpdateConfirmationRequest,
+  PolicyIntegrityManager,
+} from '@google/gemini-cli-core';
 import { theme } from '../semantic-colors.js';
 import type { RadioSelectItem } from './shared/RadioButtonSelect.js';
 import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
 import { useKeypress } from '../hooks/useKeypress.js';
+import { keyMatchers, Command } from '../keyMatchers.js';

 export enum PolicyUpdateChoice {
  ACCEPT = 'accept',
@@ -17,32 +24,46 @@ export enum PolicyUpdateChoice {
 }

 interface PolicyUpdateDialogProps {
-  onSelect: (choice: PolicyUpdateChoice) => void;
-  scope: string;
-  identifier: string;
-  isRestarting?: boolean;
+  config: Config;
+  request: PolicyUpdateConfirmationRequest;
+  onClose: () => void;
 }

 export const PolicyUpdateDialog: React.FC<PolicyUpdateDialogProps> = ({
-  onSelect,
-  scope,
-  identifier,
-  isRestarting,
+  config,
+  request,
+  onClose,
 }) => {
+  const handleSelect = useCallback(
+    async (choice: PolicyUpdateChoice) => {
+      if (choice === PolicyUpdateChoice.ACCEPT) {
+        const integrityManager = new PolicyIntegrityManager();
+        await integrityManager.acceptIntegrity(
+          request.scope,
+          request.identifier,
+          request.newHash,
+        );
+        await config.loadWorkspacePolicies(request.policyDir);
+      }
+      onClose();
+    },
+    [config, request, onClose],
+  );
+
  useKeypress(
    (key) => {
-      if (key.name === 'escape') {
-        onSelect(PolicyUpdateChoice.IGNORE);
+      if (keyMatchers[Command.ESCAPE](key)) {
+        onClose();
        return true;
      }
      return false;
    },
-    { isActive: !isRestarting },
+    { isActive: true },
  );

  const options: Array<RadioSelectItem<PolicyUpdateChoice>> = [
    {
-      label: 'Accept and Load (Requires Restart)',
+      label: 'Accept and Load',
      value: PolicyUpdateChoice.ACCEPT,
      key: 'accept',
    },
@@ -65,9 +86,9 @@ export const PolicyUpdateDialog: React.FC<PolicyUpdateDialogProps> = ({
      >
        <Box flexDirection="column" marginBottom={1}>
          <Text bold color={theme.text.primary}>
-            New or changed {scope} policies detected
+            New or changed {request.scope} policies detected
          </Text>
-          <Text color={theme.text.primary}>Location: {identifier}</Text>
+          <Text color={theme.text.primary}>Location: {request.identifier}</Text>
          <Text color={theme.text.primary}>
            Do you want to accept and load these policies?
          </Text>
@@ -75,17 +96,10 @@ export const PolicyUpdateDialog: React.FC<PolicyUpdateDialogProps> = ({

        <RadioButtonSelect
          items={options}
-          onSelect={onSelect}
-          isFocused={!isRestarting}
+          onSelect={handleSelect}
+          isFocused={true}
        />
      </Box>
-      {isRestarting && (
-        <Box marginLeft={1} marginTop={1}>
-          <Text color={theme.status.warning}>
-            Gemini CLI is restarting to apply the policy changes...
-          </Text>
-        </Box>
-      )}
    </Box>
  );
 };
--- a/packages/cli/src/ui/components/snapshots/PolicyUpdateDialog.test.tsx.snap
+++ b/packages/cli/src/ui/components/snapshots/PolicyUpdateDialog.test.tsx.snap
@@ -0,0 +1,14 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`PolicyUpdateDialog > renders correctly and matches snapshot 1`] = `
+" ╭────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
+ │                                                                                                                    │
+ │ New or changed workspace policies detected                                                                         │
+ │ Location: /test/workspace/.gemini/policies                                                                         │
+ │ Do you want to accept and load these policies?                                                                     │
+ │                                                                                                                    │
+ │ ● 1. Accept and Load                                                                                               │
+ │   2. Ignore (Use Default Policies)                                                                                 │
+ │                                                                                                                    │
+ ╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯"
+`;
--- a/packages/cli/src/ui/contexts/UIActionsContext.tsx
+++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx
@@ -1,6 +1,6 @@
 /**
 * @license
- * Copyright 2025 Google LLC
+ * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

@@ -8,7 +8,6 @@ import { createContext, useContext } from 'react';
 import { type Key } from '../hooks/useKeypress.js';
 import { type IdeIntegrationNudgeResult } from '../IdeIntegrationNudge.js';
 import { type FolderTrustChoice } from '../components/FolderTrustDialog.js';
-import { type PolicyUpdateChoice } from '../components/PolicyUpdateDialog.js';
 import {
  type AuthType,
  type EditorType,
@@ -53,7 +52,7 @@ export interface UIActions {
  vimHandleInput: (key: Key) => boolean;
  handleIdePromptComplete: (result: IdeIntegrationNudgeResult) => void;
  handleFolderTrustSelect: (choice: FolderTrustChoice) => void;
-  handlePolicyUpdateSelect: (choice: PolicyUpdateChoice) => Promise<void>;
+  setIsPolicyUpdateDialogOpen: (value: boolean) => void;
  setConstrainHeight: (value: boolean) => void;
  onEscapePromptChange: (show: boolean) => void;
  refreshStatic: () => void;
--- a/packages/cli/src/ui/contexts/UIStateContext.tsx
+++ b/packages/cli/src/ui/contexts/UIStateContext.tsx
@@ -115,7 +115,6 @@ export interface UIState {
  isFolderTrustDialogOpen: boolean;
  isPolicyUpdateDialogOpen: boolean;
  policyUpdateConfirmationRequest: PolicyUpdateConfirmationRequest | undefined;
-  isRestartingPolicyUpdate: boolean;
  isTrustedFolder: boolean | undefined;
  constrainHeight: boolean;
  showErrorDetails: boolean;
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -126,6 +126,8 @@ import {
 import { fetchAdminControls } from '../code_assist/admin/admin_controls.js';
 import { isSubpath } from '../utils/paths.js';
 import { UserHintService } from './userHintService.js';
+import { WORKSPACE_POLICY_TIER } from '../policy/config.js';
+import { loadPoliciesFromToml } from '../policy/toml-loader.js';

 export interface AccessibilitySettings {
  enableLoadingPhrases?: boolean;
@@ -1733,6 +1735,32 @@ export class Config {
    return this.policyUpdateConfirmationRequest;
  }

+  /**
+   * Hot-loads workspace policies from the specified directory into the active policy engine.
+   * This allows applying newly accepted policies without requiring an application restart.
+   *
+   * @param policyDir The directory containing the workspace policy TOML files.
+   */
+  async loadWorkspacePolicies(policyDir: string): Promise<void> {
+    const { rules, checkers } = await loadPoliciesFromToml(
+      [policyDir],
+      () => WORKSPACE_POLICY_TIER,
+    );
+
+    for (const rule of rules) {
+      this.policyEngine.addRule(rule);
+    }
+
+    for (const checker of checkers) {
+      this.policyEngine.addChecker(checker);
+    }
+
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any
+    (this as any).policyUpdateConfirmationRequest = undefined;
+
+    debugLogger.debug(`Workspace policies loaded from: ${policyDir}`);
+  }
+
  setApprovalMode(mode: ApprovalMode): void {
    if (!this.isTrustedFolder() && mode !== ApprovalMode.DEFAULT) {
      throw new Error(
--- a/packages/core/src/policy/config.ts
+++ b/packages/core/src/policy/config.ts
@@ -1,6 +1,6 @@
 /**
 * @license
- * Copyright 2025 Google LLC
+ * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

--- a/packages/core/src/policy/integrity.test.ts
+++ b/packages/core/src/policy/integrity.test.ts
@@ -1,73 +1,47 @@
 /**
 * @license
- * Copyright 2025 Google LLC
+ * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

-import {
-  describe,
-  it,
-  expect,
-  vi,
-  afterEach,
-  beforeEach,
-  type Mock,
-} from 'vitest';
+import { describe, it, expect, vi, afterEach, beforeEach } from 'vitest';
 import { PolicyIntegrityManager, IntegrityStatus } from './integrity.js';
-
-// Mock dependencies
-vi.mock('../config/storage.js', () => ({
-  Storage: {
-    getPolicyIntegrityStoragePath: vi
-      .fn()
-      .mockReturnValue('/mock/storage/policy_integrity.json'),
-  },
-}));
-
-vi.mock('./toml-loader.js', () => ({
-  readPolicyFiles: vi.fn(),
-}));
-
-// Mock FS
-const mockFs = vi.hoisted(() => ({
-  readFile: vi.fn(),
-  writeFile: vi.fn(),
-  mkdir: vi.fn(),
-}));
-
-vi.mock('node:fs/promises', () => ({
-  default: mockFs,
-  readFile: mockFs.readFile,
-  writeFile: mockFs.writeFile,
-  mkdir: mockFs.mkdir,
-}));
+import * as fs from 'node:fs/promises';
+import * as path from 'node:path';
+import * as os from 'node:os';
+import { Storage } from '../config/storage.js';

 describe('PolicyIntegrityManager', () => {
  let integrityManager: PolicyIntegrityManager;
-  let readPolicyFilesMock: Mock;
+  let tempDir: string;
+  let integrityStoragePath: string;

  beforeEach(async () => {
-    vi.resetModules();
-    const { readPolicyFiles } = await import('./toml-loader.js');
-    readPolicyFilesMock = readPolicyFiles as Mock;
+    tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'gemini-cli-test-'));
+    integrityStoragePath = path.join(tempDir, 'policy_integrity.json');
+
+    vi.spyOn(Storage, 'getPolicyIntegrityStoragePath').mockReturnValue(
+      integrityStoragePath,
+    );
+
    integrityManager = new PolicyIntegrityManager();
  });

-  afterEach(() => {
-    vi.clearAllMocks();
+  afterEach(async () => {
+    await fs.rm(tempDir, { recursive: true, force: true });
+    vi.restoreAllMocks();
  });

  describe('checkIntegrity', () => {
    it('should return NEW if no stored hash', async () => {
-      mockFs.readFile.mockRejectedValue({ code: 'ENOENT' }); // No stored file
-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/a.toml', content: 'contentA' },
-      ]);
+      const policyDir = path.join(tempDir, 'policies');
+      await fs.mkdir(policyDir);
+      await fs.writeFile(path.join(policyDir, 'a.toml'), 'contentA');

      const result = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/dir',
+        policyDir,
      );
      expect(result.status).toBe(IntegrityStatus.NEW);
      expect(result.hash).toBeDefined();
@@ -76,187 +50,171 @@ describe('PolicyIntegrityManager', () => {
    });

    it('should return MATCH if stored hash matches', async () => {
-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/a.toml', content: 'contentA' },
-      ]);
-      // We can't easily get the expected hash without calling private method or re-implementing logic.
-      // But we can run checkIntegrity once (NEW) to get the hash, then mock FS with that hash.
-      mockFs.readFile.mockRejectedValue({ code: 'ENOENT' });
+      const policyDir = path.join(tempDir, 'policies');
+      await fs.mkdir(policyDir);
+      await fs.writeFile(path.join(policyDir, 'a.toml'), 'contentA');
+
+      // First run to get the hash
      const resultNew = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/dir',
+        policyDir,
      );
      const currentHash = resultNew.hash;

-      mockFs.readFile.mockResolvedValue(
-        JSON.stringify({
-          'workspace:id': currentHash,
-        }),
+      // Save the hash to mock storage
+      await fs.writeFile(
+        integrityStoragePath,
+        JSON.stringify({ 'workspace:id': currentHash }),
      );

      const result = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/dir',
+        policyDir,
      );
      expect(result.status).toBe(IntegrityStatus.MATCH);
      expect(result.hash).toBe(currentHash);
-      expect(result.fileCount).toBe(1);
    });

    it('should return MISMATCH if stored hash differs', async () => {
-      mockFs.readFile.mockRejectedValue({ code: 'ENOENT' });
-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/a.toml', content: 'contentA' },
-      ]);
+      const policyDir = path.join(tempDir, 'policies');
+      await fs.mkdir(policyDir);
+      await fs.writeFile(path.join(policyDir, 'a.toml'), 'contentA');
+
      const resultNew = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/dir',
+        policyDir,
      );
      const currentHash = resultNew.hash;

-      mockFs.readFile.mockResolvedValue(
-        JSON.stringify({
-          'workspace:id': 'different_hash',
-        }),
+      // Save a different hash
+      await fs.writeFile(
+        integrityStoragePath,
+        JSON.stringify({ 'workspace:id': 'different_hash' }),
      );

      const result = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/dir',
+        policyDir,
      );
      expect(result.status).toBe(IntegrityStatus.MISMATCH);
      expect(result.hash).toBe(currentHash);
-      expect(result.fileCount).toBe(1);
    });

    it('should result in different hash if filename changes', async () => {
-      mockFs.readFile.mockRejectedValue({ code: 'ENOENT' });
+      const policyDir1 = path.join(tempDir, 'policies1');
+      await fs.mkdir(policyDir1);
+      await fs.writeFile(path.join(policyDir1, 'a.toml'), 'contentA');

-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/a.toml', content: 'contentA' },
-      ]);
      const result1 = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/workspace/policies',
+        policyDir1,
      );

-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/b.toml', content: 'contentA' },
-      ]);
+      const policyDir2 = path.join(tempDir, 'policies2');
+      await fs.mkdir(policyDir2);
+      await fs.writeFile(path.join(policyDir2, 'b.toml'), 'contentA');
+
      const result2 = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/workspace/policies',
+        policyDir2,
      );

      expect(result1.hash).not.toBe(result2.hash);
    });

    it('should result in different hash if content changes', async () => {
-      mockFs.readFile.mockRejectedValue({ code: 'ENOENT' });
+      const policyDir = path.join(tempDir, 'policies');
+      await fs.mkdir(policyDir);

-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/a.toml', content: 'contentA' },
-      ]);
+      await fs.writeFile(path.join(policyDir, 'a.toml'), 'contentA');
      const result1 = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/workspace/policies',
+        policyDir,
      );

-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/a.toml', content: 'contentB' },
-      ]);
+      await fs.writeFile(path.join(policyDir, 'a.toml'), 'contentB');
      const result2 = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/workspace/policies',
+        policyDir,
      );

      expect(result1.hash).not.toBe(result2.hash);
    });

    it('should be deterministic (sort order)', async () => {
-      mockFs.readFile.mockRejectedValue({ code: 'ENOENT' });
+      const policyDir1 = path.join(tempDir, 'policies1');
+      await fs.mkdir(policyDir1);
+      await fs.writeFile(path.join(policyDir1, 'a.toml'), 'contentA');
+      await fs.writeFile(path.join(policyDir1, 'b.toml'), 'contentB');

-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/a.toml', content: 'contentA' },
-        { path: '/workspace/policies/b.toml', content: 'contentB' },
-      ]);
      const result1 = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/workspace/policies',
+        policyDir1,
      );

-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/workspace/policies/b.toml', content: 'contentB' },
-        { path: '/workspace/policies/a.toml', content: 'contentA' },
-      ]);
+      // Re-read with same files but they might be in different order in readdir
+      // PolicyIntegrityManager should sort them.
      const result2 = await integrityManager.checkIntegrity(
        'workspace',
        'id',
-        '/workspace/policies',
+        policyDir1,
      );

      expect(result1.hash).toBe(result2.hash);
    });

    it('should handle multiple projects correctly', async () => {
-      mockFs.readFile.mockRejectedValue({ code: 'ENOENT' });
+      const dirA = path.join(tempDir, 'dirA');
+      await fs.mkdir(dirA);
+      await fs.writeFile(path.join(dirA, 'p.toml'), 'contentA');
+
+      const dirB = path.join(tempDir, 'dirB');
+      await fs.mkdir(dirB);
+      await fs.writeFile(path.join(dirB, 'p.toml'), 'contentB');

-      // First, get hashes for two different projects
-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/dirA/p.toml', content: 'contentA' },
-      ]);
      const { hash: hashA } = await integrityManager.checkIntegrity(
        'workspace',
        'idA',
-        '/dirA',
+        dirA,
      );
-
-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/dirB/p.toml', content: 'contentB' },
-      ]);
      const { hash: hashB } = await integrityManager.checkIntegrity(
        'workspace',
        'idB',
-        '/dirB',
+        dirB,
      );

-      // Now mock storage with both
-      mockFs.readFile.mockResolvedValue(
+      // Save to storage
+      await fs.writeFile(
+        integrityStoragePath,
        JSON.stringify({
          'workspace:idA': hashA,
-          'workspace:idB': 'oldHashB', // Different from hashB
+          'workspace:idB': 'oldHashB',
        }),
      );

      // Project A should match
-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/dirA/p.toml', content: 'contentA' },
-      ]);
      const resultA = await integrityManager.checkIntegrity(
        'workspace',
        'idA',
-        '/dirA',
+        dirA,
      );
      expect(resultA.status).toBe(IntegrityStatus.MATCH);
      expect(resultA.hash).toBe(hashA);

      // Project B should mismatch
-      readPolicyFilesMock.mockResolvedValue([
-        { path: '/dirB/p.toml', content: 'contentB' },
-      ]);
      const resultB = await integrityManager.checkIntegrity(
        'workspace',
        'idB',
-        '/dirB',
+        dirB,
      );
      expect(resultB.status).toBe(IntegrityStatus.MISMATCH);
      expect(resultB.hash).toBe(hashB);
@@ -265,42 +223,27 @@ describe('PolicyIntegrityManager', () => {

  describe('acceptIntegrity', () => {
    it('should save the hash to storage', async () => {
-      mockFs.readFile.mockRejectedValue({ code: 'ENOENT' }); // Start empty
-      mockFs.mkdir.mockResolvedValue(undefined);
-      mockFs.writeFile.mockResolvedValue(undefined);
-
      await integrityManager.acceptIntegrity('workspace', 'id', 'hash123');

-      expect(mockFs.writeFile).toHaveBeenCalledWith(
-        '/mock/storage/policy_integrity.json',
-        JSON.stringify({ 'workspace:id': 'hash123' }, null, 2),
-        'utf-8',
+      const stored = JSON.parse(
+        await fs.readFile(integrityStoragePath, 'utf-8'),
      );
+      expect(stored['workspace:id']).toBe('hash123');
    });

    it('should update existing hash', async () => {
-      mockFs.readFile.mockResolvedValue(
-        JSON.stringify({
-          'other:id': 'otherhash',
-        }),
+      await fs.writeFile(
+        integrityStoragePath,
+        JSON.stringify({ 'other:id': 'otherhash' }),
      );
-      mockFs.mkdir.mockResolvedValue(undefined);
-      mockFs.writeFile.mockResolvedValue(undefined);

      await integrityManager.acceptIntegrity('workspace', 'id', 'hash123');

-      expect(mockFs.writeFile).toHaveBeenCalledWith(
-        '/mock/storage/policy_integrity.json',
-        JSON.stringify(
-          {
-            'other:id': 'otherhash',
-            'workspace:id': 'hash123',
-          },
-          null,
-          2,
-        ),
-        'utf-8',
+      const stored = JSON.parse(
+        await fs.readFile(integrityStoragePath, 'utf-8'),
      );
+      expect(stored['other:id']).toBe('otherhash');
+      expect(stored['workspace:id']).toBe('hash123');
    });
  });
 });
--- a/packages/core/src/policy/integrity.ts
+++ b/packages/core/src/policy/integrity.ts
@@ -1,6 +1,6 @@
 /**
 * @license
- * Copyright 2025 Google LLC
+ * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

--- a/packages/core/src/policy/policy-engine.ts
+++ b/packages/core/src/policy/policy-engine.ts
@@ -1,6 +1,6 @@
 /**
 * @license
- * Copyright 2025 Google LLC
+ * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */