Update persistence state to track counts of messages instead of times banner has been displayed (#13428)

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('<AppHeader />', () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    persistentStateMock.get.mockReturnValue(0);
+    persistentStateMock.get.mockReturnValue({});
   });
 
   it('should render the banner with default text', () => {
@@ -146,8 +147,8 @@ describe('<AppHeader />', () => {
     unmount();
   });
 
-  it('should increment the shown count when default banner is displayed', () => {
+  it('should increment the version count when default banner is displayed', () => {
-    persistentStateMock.get.mockReturnValue(0);
+    persistentStateMock.get.mockReturnValue({});
     const mockConfig = makeFakeConfig();
     const uiState = {
       bannerData: {
@@ -163,7 +164,12 @@ describe('<AppHeader />', () => {
 
     expect(persistentStateMock.set).toHaveBeenCalledWith(
       'defaultBannerShownCount',
-      1,
+      {
+        [crypto
+          .createHash('sha256')
+          .update(uiState.bannerData.defaultText)
+          .digest('hex')]: 1,
+      },
     );
     unmount();
   });

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(
+  const { bannerText } = useBanner(bannerData, config);
-    () => 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]);
 
   return (
     <Box flexDirection="column">
@@ -54,7 +33,7 @@ export const AppHeader = ({ version }: AppHeaderProps) => {
             <Banner
               width={mainAreaWidth}
               bannerText={bannerText}
-              isWarning={warningText !== ''}
+              isWarning={bannerData.warningText !== ''}
             />
           )}
         </>

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');
+  });
+});

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<string | null>(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,
+  };
+}

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<string, number>;
   // Add other persistent state keys here as needed
 }