From 3e50be1658a8f2be9c913211cb945aeb38a112d5 Mon Sep 17 00:00:00 2001 From: Adib234 <30782825+Adib234@users.noreply.github.com> Date: Thu, 20 Nov 2025 13:57:06 -0800 Subject: [PATCH] Update persistence state to track counts of messages instead of times banner has been displayed (#13428) --- .../cli/src/ui/components/AppHeader.test.tsx | 14 +- packages/cli/src/ui/components/AppHeader.tsx | 27 +--- packages/cli/src/ui/hooks/useBanner.test.ts | 147 ++++++++++++++++++ packages/cli/src/ui/hooks/useBanner.ts | 73 +++++++++ packages/cli/src/utils/persistentState.ts | 2 +- 5 files changed, 234 insertions(+), 29 deletions(-) create mode 100644 packages/cli/src/ui/hooks/useBanner.test.ts create mode 100644 packages/cli/src/ui/hooks/useBanner.ts diff --git a/packages/cli/src/ui/components/AppHeader.test.tsx b/packages/cli/src/ui/components/AppHeader.test.tsx index a6f9103f18..478aaceeeb 100644 --- a/packages/cli/src/ui/components/AppHeader.test.tsx +++ b/packages/cli/src/ui/components/AppHeader.test.tsx @@ -8,6 +8,7 @@ import { renderWithProviders } from '../../test-utils/render.js'; import { AppHeader } from './AppHeader.js'; import { describe, it, expect, vi, beforeEach } from 'vitest'; import { makeFakeConfig } from '@google/gemini-cli-core'; +import crypto from 'node:crypto'; const persistentStateMock = vi.hoisted(() => ({ get: vi.fn(), @@ -25,7 +26,7 @@ vi.mock('../utils/terminalSetup.js', () => ({ describe('', () => { beforeEach(() => { vi.clearAllMocks(); - persistentStateMock.get.mockReturnValue(0); + persistentStateMock.get.mockReturnValue({}); }); it('should render the banner with default text', () => { @@ -146,8 +147,8 @@ describe('', () => { unmount(); }); - it('should increment the shown count when default banner is displayed', () => { - persistentStateMock.get.mockReturnValue(0); + it('should increment the version count when default banner is displayed', () => { + persistentStateMock.get.mockReturnValue({}); const mockConfig = makeFakeConfig(); const uiState = { bannerData: { @@ -163,7 +164,12 @@ describe('', () => { expect(persistentStateMock.set).toHaveBeenCalledWith( 'defaultBannerShownCount', - 1, + { + [crypto + .createHash('sha256') + .update(uiState.bannerData.defaultText) + .digest('hex')]: 1, + }, ); unmount(); }); diff --git a/packages/cli/src/ui/components/AppHeader.tsx b/packages/cli/src/ui/components/AppHeader.tsx index 05e528f6c7..c404c0e9f9 100644 --- a/packages/cli/src/ui/components/AppHeader.tsx +++ b/packages/cli/src/ui/components/AppHeader.tsx @@ -10,9 +10,8 @@ import { Tips } from './Tips.js'; import { useSettings } from '../contexts/SettingsContext.js'; import { useConfig } from '../contexts/ConfigContext.js'; import { useUIState } from '../contexts/UIStateContext.js'; -import { persistentState } from '../../utils/persistentState.js'; -import { useEffect, useRef, useState } from 'react'; import { Banner } from './Banner.js'; +import { useBanner } from '../hooks/useBanner.js'; interface AppHeaderProps { version: string; @@ -23,27 +22,7 @@ export const AppHeader = ({ version }: AppHeaderProps) => { const config = useConfig(); const { nightly, mainAreaWidth, bannerData, bannerVisible } = useUIState(); - const [defaultBannerShownCount] = useState( - () => persistentState.get('defaultBannerShownCount') || 0, - ); - - const { defaultText, warningText } = bannerData; - - const showDefaultBanner = - warningText === '' && - !config.getPreviewFeatures() && - defaultBannerShownCount < 5; - - const bannerText = showDefaultBanner ? defaultText : warningText; - - const hasIncrementedRef = useRef(false); - useEffect(() => { - if (showDefaultBanner && defaultText && !hasIncrementedRef.current) { - hasIncrementedRef.current = true; - const current = persistentState.get('defaultBannerShownCount') || 0; - persistentState.set('defaultBannerShownCount', current + 1); - } - }, [showDefaultBanner, defaultText]); + const { bannerText } = useBanner(bannerData, config); return ( @@ -54,7 +33,7 @@ export const AppHeader = ({ version }: AppHeaderProps) => { )} diff --git a/packages/cli/src/ui/hooks/useBanner.test.ts b/packages/cli/src/ui/hooks/useBanner.test.ts new file mode 100644 index 0000000000..27909fae27 --- /dev/null +++ b/packages/cli/src/ui/hooks/useBanner.test.ts @@ -0,0 +1,147 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + describe, + it, + expect, + vi, + beforeEach, + type MockedFunction, +} from 'vitest'; +import { renderHook } from '../../test-utils/render.js'; +import { useBanner } from './useBanner.js'; +import { persistentState } from '../../utils/persistentState.js'; +import type { Config } from '@google/gemini-cli-core'; +import crypto from 'node:crypto'; + +vi.mock('../../utils/persistentState.js', () => ({ + persistentState: { + get: vi.fn(), + set: vi.fn(), + }, +})); + +vi.mock('../semantic-colors.js', () => ({ + theme: { + status: { + warning: 'mock-warning-color', + }, + }, +})); + +vi.mock('../colors.js', () => ({ + Colors: { + AccentBlue: 'mock-accent-blue', + }, +})); + +// Define the shape of the config methods used by this hook +interface MockConfigShape { + getPreviewFeatures: MockedFunction<() => boolean>; +} + +describe('useBanner', () => { + let mockConfig: MockConfigShape; + const mockedPersistentStateGet = persistentState.get as MockedFunction< + typeof persistentState.get + >; + const mockedPersistentStateSet = persistentState.set as MockedFunction< + typeof persistentState.set + >; + + const defaultBannerData = { + defaultText: 'Standard Banner', + warningText: '', + }; + + beforeEach(() => { + vi.resetAllMocks(); + + // Initialize the mock config with default behavior + mockConfig = { + getPreviewFeatures: vi.fn().mockReturnValue(false), + }; + + // Default persistentState behavior: return empty object (no counts) + mockedPersistentStateGet.mockReturnValue({}); + }); + + it('should return warning text and warning color if warningText is present', () => { + const data = { defaultText: 'Standard', warningText: 'Critical Error' }; + + const { result } = renderHook(() => + useBanner(data, mockConfig as unknown as Config), + ); + + expect(result.current.bannerText).toBe('Critical Error'); + }); + + it('should NOT show default banner if preview features are enabled in config', () => { + // Simulate Preview Features Enabled + mockConfig.getPreviewFeatures.mockReturnValue(true); + + const { result } = renderHook(() => + useBanner(defaultBannerData, mockConfig as unknown as Config), + ); + + // Should fall back to warningText (which is empty) + expect(result.current.bannerText).toBe(''); + }); + + it('should hide banner if show count exceeds max limit (Legacy format)', () => { + mockedPersistentStateGet.mockReturnValue({ + [crypto + .createHash('sha256') + .update(defaultBannerData.defaultText) + .digest('hex')]: 5, + }); + + const { result } = renderHook(() => + useBanner(defaultBannerData, mockConfig as unknown as Config), + ); + + expect(result.current.bannerText).toBe(''); + }); + + it('should increment the persistent count when banner is shown', () => { + const data = { defaultText: 'Tracker', warningText: '' }; + + // Current count is 1 + mockedPersistentStateGet.mockReturnValue({ + [crypto.createHash('sha256').update(data.defaultText).digest('hex')]: 1, + }); + + renderHook(() => useBanner(data, mockConfig as unknown as Config)); + + // Expect set to be called with incremented count + expect(mockedPersistentStateSet).toHaveBeenCalledWith( + 'defaultBannerShownCount', + { + [crypto.createHash('sha256').update(data.defaultText).digest('hex')]: 2, + }, + ); + }); + + it('should NOT increment count if warning text is shown instead', () => { + const data = { defaultText: 'Standard', warningText: 'Warning' }; + + renderHook(() => useBanner(data, mockConfig as unknown as Config)); + + // Since warning text takes precedence, default banner logic (and increment) is skipped + expect(mockedPersistentStateSet).not.toHaveBeenCalled(); + }); + + it('should handle newline replacements', () => { + const data = { defaultText: 'Line1\\nLine2', warningText: '' }; + + const { result } = renderHook(() => + useBanner(data, mockConfig as unknown as Config), + ); + + expect(result.current.bannerText).toBe('Line1\nLine2'); + }); +}); diff --git a/packages/cli/src/ui/hooks/useBanner.ts b/packages/cli/src/ui/hooks/useBanner.ts new file mode 100644 index 0000000000..faca37ca02 --- /dev/null +++ b/packages/cli/src/ui/hooks/useBanner.ts @@ -0,0 +1,73 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useState, useEffect, useRef } from 'react'; +import { persistentState } from '../../utils/persistentState.js'; +import type { Config } from '@google/gemini-cli-core'; +import crypto from 'node:crypto'; + +const DEFAULT_MAX_BANNER_SHOWN_COUNT = 5; + +interface BannerData { + defaultText: string; + warningText: string; +} + +export function useBanner(bannerData: BannerData, config: Config) { + const { defaultText, warningText } = bannerData; + + const [previewEnabled, setPreviewEnabled] = useState( + config.getPreviewFeatures(), + ); + + useEffect(() => { + const isEnabled = config.getPreviewFeatures(); + if (isEnabled !== previewEnabled) { + setPreviewEnabled(isEnabled); + } + }, [config, previewEnabled]); + + const [bannerCounts] = useState( + () => persistentState.get('defaultBannerShownCount') || {}, + ); + + const hashedText = crypto + .createHash('sha256') + .update(defaultText) + .digest('hex'); + + const currentBannerCount = bannerCounts[hashedText] || 0; + + const showDefaultBanner = + warningText === '' && + !previewEnabled && + currentBannerCount < DEFAULT_MAX_BANNER_SHOWN_COUNT; + + const rawBannerText = showDefaultBanner ? defaultText : warningText; + const bannerText = rawBannerText.replace(/\\n/g, '\n'); + + const lastIncrementedKey = useRef(null); + + useEffect(() => { + if (showDefaultBanner && defaultText) { + if (lastIncrementedKey.current !== defaultText) { + lastIncrementedKey.current = defaultText; + + const allCounts = persistentState.get('defaultBannerShownCount') || {}; + const current = allCounts[hashedText] || 0; + + persistentState.set('defaultBannerShownCount', { + ...allCounts, + [hashedText]: current + 1, + }); + } + } + }, [showDefaultBanner, defaultText, hashedText]); + + return { + bannerText, + }; +} diff --git a/packages/cli/src/utils/persistentState.ts b/packages/cli/src/utils/persistentState.ts index 6e4a8810b0..f5fc4d4b29 100644 --- a/packages/cli/src/utils/persistentState.ts +++ b/packages/cli/src/utils/persistentState.ts @@ -11,7 +11,7 @@ import * as path from 'node:path'; const STATE_FILENAME = 'state.json'; interface PersistentStateData { - defaultBannerShownCount?: number; + defaultBannerShownCount?: Record; // Add other persistent state keys here as needed }