From 662654c5d203eac993f02562c78ecc2f4eaf1b24 Mon Sep 17 00:00:00 2001 From: Abhijit Balaji Date: Wed, 18 Feb 2026 12:54:47 -0800 Subject: [PATCH] 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. --- packages/cli/src/test-utils/render.tsx | 2 +- packages/cli/src/ui/AppContainer.tsx | 35 +-- .../cli/src/ui/components/DialogManager.tsx | 12 +- .../ui/components/PolicyUpdateDialog.test.tsx | 115 +++++---- .../src/ui/components/PolicyUpdateDialog.tsx | 60 +++-- .../PolicyUpdateDialog.test.tsx.snap | 14 + .../cli/src/ui/contexts/UIActionsContext.tsx | 5 +- .../cli/src/ui/contexts/UIStateContext.tsx | 1 - packages/core/src/config/config.ts | 28 ++ packages/core/src/policy/config.ts | 2 +- packages/core/src/policy/integrity.test.ts | 239 +++++++----------- packages/core/src/policy/integrity.ts | 2 +- packages/core/src/policy/policy-engine.ts | 2 +- 13 files changed, 248 insertions(+), 269 deletions(-) create mode 100644 packages/cli/src/ui/components/__snapshots__/PolicyUpdateDialog.test.tsx.snap diff --git a/packages/cli/src/test-utils/render.tsx b/packages/cli/src/test-utils/render.tsx index 07ebf9246f..c65e5ffe94 100644 --- 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(), diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 7a270b8b74..a372125f7e 100644 --- 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, diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index 11119c12b0..9fdd4718a6 100644 --- 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 ( uiActions.setIsPolicyUpdateDialogOpen(false)} /> ); } diff --git a/packages/cli/src/ui/components/PolicyUpdateDialog.test.tsx b/packages/cli/src/ui/components/PolicyUpdateDialog.test.tsx index 5175564886..d54b610638 100644 --- 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( , ); 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( , ); @@ -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( , ); @@ -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( , ); 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( - , - ); - - const output = lastFrame(); - expect(output).toContain( - 'Gemini CLI is restarting to apply the policy changes...', - ); - }); }); diff --git a/packages/cli/src/ui/components/PolicyUpdateDialog.tsx b/packages/cli/src/ui/components/PolicyUpdateDialog.tsx index c05044bd9c..395a2dd3cc 100644 --- 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 = ({ - 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> = [ { - label: 'Accept and Load (Requires Restart)', + label: 'Accept and Load', value: PolicyUpdateChoice.ACCEPT, key: 'accept', }, @@ -65,9 +86,9 @@ export const PolicyUpdateDialog: React.FC = ({ > - New or changed {scope} policies detected + New or changed {request.scope} policies detected - Location: {identifier} + Location: {request.identifier} Do you want to accept and load these policies? @@ -75,17 +96,10 @@ export const PolicyUpdateDialog: React.FC = ({ - {isRestarting && ( - - - Gemini CLI is restarting to apply the policy changes... - - - )} ); }; diff --git a/packages/cli/src/ui/components/__snapshots__/PolicyUpdateDialog.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/PolicyUpdateDialog.test.tsx.snap new file mode 100644 index 0000000000..a8bd583cb2 --- /dev/null +++ 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) │ + │ │ + ╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯" +`; diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index afd49f4f4e..03780c5068 100644 --- 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; + setIsPolicyUpdateDialogOpen: (value: boolean) => void; setConstrainHeight: (value: boolean) => void; onEscapePromptChange: (show: boolean) => void; refreshStatic: () => void; diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 82b43d3616..56d4b83c09 100644 --- 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; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index cca6c60afb..ac7028b0be 100644 --- 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 { + 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( diff --git a/packages/core/src/policy/config.ts b/packages/core/src/policy/config.ts index a9414d65b6..413ef81ae7 100644 --- 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 */ diff --git a/packages/core/src/policy/integrity.test.ts b/packages/core/src/policy/integrity.test.ts index a289c513d0..32ebf56058 100644 --- 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'); }); }); }); diff --git a/packages/core/src/policy/integrity.ts b/packages/core/src/policy/integrity.ts index d9661853ae..77eb49f7e4 100644 --- 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 */ diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index 3f386edd8f..25d1983278 100644 --- 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 */