Skip to content

Commit 3e50be1

Browse files
Adib234Adib234
authored
Update persistence state to track counts of messages instead of times banner has been displayed (#13428)
1 parent 9937fb2 commit 3e50be1

File tree

5 files changed

+234
-29
lines changed

5 files changed

+234
-29
lines changed

‎packages/cli/src/ui/components/AppHeader.test.tsx‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { renderWithProviders } from '../../test-utils/render.js';
88
import { AppHeader } from './AppHeader.js';
99
import { describe, it, expect, vi, beforeEach } from 'vitest';
1010
import { makeFakeConfig } from '@google/gemini-cli-core';
11+
import crypto from 'node:crypto';
1112

1213
const persistentStateMock = vi.hoisted(() => ({
1314
get: vi.fn(),
@@ -25,7 +26,7 @@ vi.mock('../utils/terminalSetup.js', () => ({
2526
describe('<AppHeader />', () => {
2627
beforeEach(() => {
2728
vi.clearAllMocks();
28-
persistentStateMock.get.mockReturnValue(0);
29+
persistentStateMock.get.mockReturnValue({});
2930
});
3031

3132
it('should render the banner with default text', () => {
@@ -146,8 +147,8 @@ describe('<AppHeader />', () => {
146147
unmount();
147148
});
148149

149-
it('should increment the shown count when default banner is displayed', () => {
150-
persistentStateMock.get.mockReturnValue(0);
150+
it('should increment the version count when default banner is displayed', () => {
151+
persistentStateMock.get.mockReturnValue({});
151152
const mockConfig = makeFakeConfig();
152153
const uiState = {
153154
bannerData: {
@@ -163,7 +164,12 @@ describe('<AppHeader />', () => {
163164

164165
expect(persistentStateMock.set).toHaveBeenCalledWith(
165166
'defaultBannerShownCount',
166-
1,
167+
{
168+
[crypto
169+
.createHash('sha256')
170+
.update(uiState.bannerData.defaultText)
171+
.digest('hex')]: 1,
172+
},
167173
);
168174
unmount();
169175
});

‎packages/cli/src/ui/components/AppHeader.tsx‎

Lines changed: 3 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,8 @@ import { Tips } from './Tips.js';
1010
import { useSettings } from '../contexts/SettingsContext.js';
1111
import { useConfig } from '../contexts/ConfigContext.js';
1212
import { useUIState } from '../contexts/UIStateContext.js';
13-
import { persistentState } from '../../utils/persistentState.js';
14-
import { useEffect, useRef, useState } from 'react';
1513
import { Banner } from './Banner.js';
14+
import { useBanner } from '../hooks/useBanner.js';
1615

1716
interface AppHeaderProps {
1817
version: string;
@@ -23,27 +22,7 @@ export const AppHeader = ({ version }: AppHeaderProps) => {
2322
const config = useConfig();
2423
const { nightly, mainAreaWidth, bannerData, bannerVisible } = useUIState();
2524

26-
const [defaultBannerShownCount] = useState(
27-
() => persistentState.get('defaultBannerShownCount') || 0,
28-
);
29-
30-
const { defaultText, warningText } = bannerData;
31-
32-
const showDefaultBanner =
33-
warningText === '' &&
34-
!config.getPreviewFeatures() &&
35-
defaultBannerShownCount < 5;
36-
37-
const bannerText = showDefaultBanner ? defaultText : warningText;
38-
39-
const hasIncrementedRef = useRef(false);
40-
useEffect(() => {
41-
if (showDefaultBanner && defaultText && !hasIncrementedRef.current) {
42-
hasIncrementedRef.current = true;
43-
const current = persistentState.get('defaultBannerShownCount') || 0;
44-
persistentState.set('defaultBannerShownCount', current + 1);
45-
}
46-
}, [showDefaultBanner, defaultText]);
25+
const { bannerText } = useBanner(bannerData, config);
4726

4827
return (
4928
<Box flexDirection="column">
@@ -54,7 +33,7 @@ export const AppHeader = ({ version }: AppHeaderProps) => {
5433
<Banner
5534
width={mainAreaWidth}
5635
bannerText={bannerText}
57-
isWarning={warningText !== ''}
36+
isWarning={bannerData.warningText !== ''}
5837
/>
5938
)}
6039
</>

‎packages/cli/src/ui/hooks/useBanner.test.ts‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import {
8+
describe,
9+
it,
10+
expect,
11+
vi,
12+
beforeEach,
13+
type MockedFunction,
14+
} from 'vitest';
15+
import { renderHook } from '../../test-utils/render.js';
16+
import { useBanner } from './useBanner.js';
17+
import { persistentState } from '../../utils/persistentState.js';
18+
import type { Config } from '@google/gemini-cli-core';
19+
import crypto from 'node:crypto';
20+
21+
vi.mock('../../utils/persistentState.js', () => ({
22+
persistentState: {
23+
get: vi.fn(),
24+
set: vi.fn(),
25+
},
26+
}));
27+
28+
vi.mock('../semantic-colors.js', () => ({
29+
theme: {
30+
status: {
31+
warning: 'mock-warning-color',
32+
},
33+
},
34+
}));
35+
36+
vi.mock('../colors.js', () => ({
37+
Colors: {
38+
AccentBlue: 'mock-accent-blue',
39+
},
40+
}));
41+
42+
// Define the shape of the config methods used by this hook
43+
interface MockConfigShape {
44+
getPreviewFeatures: MockedFunction<() => boolean>;
45+
}
46+
47+
describe('useBanner', () => {
48+
let mockConfig: MockConfigShape;
49+
const mockedPersistentStateGet = persistentState.get as MockedFunction<
50+
typeof persistentState.get
51+
>;
52+
const mockedPersistentStateSet = persistentState.set as MockedFunction<
53+
typeof persistentState.set
54+
>;
55+
56+
const defaultBannerData = {
57+
defaultText: 'Standard Banner',
58+
warningText: '',
59+
};
60+
61+
beforeEach(() => {
62+
vi.resetAllMocks();
63+
64+
// Initialize the mock config with default behavior
65+
mockConfig = {
66+
getPreviewFeatures: vi.fn().mockReturnValue(false),
67+
};
68+
69+
// Default persistentState behavior: return empty object (no counts)
70+
mockedPersistentStateGet.mockReturnValue({});
71+
});
72+
73+
it('should return warning text and warning color if warningText is present', () => {
74+
const data = { defaultText: 'Standard', warningText: 'Critical Error' };
75+
76+
const { result } = renderHook(() =>
77+
useBanner(data, mockConfig as unknown as Config),
78+
);
79+
80+
expect(result.current.bannerText).toBe('Critical Error');
81+
});
82+
83+
it('should NOT show default banner if preview features are enabled in config', () => {
84+
// Simulate Preview Features Enabled
85+
mockConfig.getPreviewFeatures.mockReturnValue(true);
86+
87+
const { result } = renderHook(() =>
88+
useBanner(defaultBannerData, mockConfig as unknown as Config),
89+
);
90+
91+
// Should fall back to warningText (which is empty)
92+
expect(result.current.bannerText).toBe('');
93+
});
94+
95+
it('should hide banner if show count exceeds max limit (Legacy format)', () => {
96+
mockedPersistentStateGet.mockReturnValue({
97+
[crypto
98+
.createHash('sha256')
99+
.update(defaultBannerData.defaultText)
100+
.digest('hex')]: 5,
101+
});
102+
103+
const { result } = renderHook(() =>
104+
useBanner(defaultBannerData, mockConfig as unknown as Config),
105+
);
106+
107+
expect(result.current.bannerText).toBe('');
108+
});
109+
110+
it('should increment the persistent count when banner is shown', () => {
111+
const data = { defaultText: 'Tracker', warningText: '' };
112+
113+
// Current count is 1
114+
mockedPersistentStateGet.mockReturnValue({
115+
[crypto.createHash('sha256').update(data.defaultText).digest('hex')]: 1,
116+
});
117+
118+
renderHook(() => useBanner(data, mockConfig as unknown as Config));
119+
120+
// Expect set to be called with incremented count
121+
expect(mockedPersistentStateSet).toHaveBeenCalledWith(
122+
'defaultBannerShownCount',
123+
{
124+
[crypto.createHash('sha256').update(data.defaultText).digest('hex')]: 2,
125+
},
126+
);
127+
});
128+
129+
it('should NOT increment count if warning text is shown instead', () => {
130+
const data = { defaultText: 'Standard', warningText: 'Warning' };
131+
132+
renderHook(() => useBanner(data, mockConfig as unknown as Config));
133+
134+
// Since warning text takes precedence, default banner logic (and increment) is skipped
135+
expect(mockedPersistentStateSet).not.toHaveBeenCalled();
136+
});
137+
138+
it('should handle newline replacements', () => {
139+
const data = { defaultText: 'Line1\\nLine2', warningText: '' };
140+
141+
const { result } = renderHook(() =>
142+
useBanner(data, mockConfig as unknown as Config),
143+
);
144+
145+
expect(result.current.bannerText).toBe('Line1\nLine2');
146+
});
147+
});

‎packages/cli/src/ui/hooks/useBanner.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { useState, useEffect, useRef } from 'react';
8+
import { persistentState } from '../../utils/persistentState.js';
9+
import type { Config } from '@google/gemini-cli-core';
10+
import crypto from 'node:crypto';
11+
12+
const DEFAULT_MAX_BANNER_SHOWN_COUNT = 5;
13+
14+
interface BannerData {
15+
defaultText: string;
16+
warningText: string;
17+
}
18+
19+
export function useBanner(bannerData: BannerData, config: Config) {
20+
const { defaultText, warningText } = bannerData;
21+
22+
const [previewEnabled, setPreviewEnabled] = useState(
23+
config.getPreviewFeatures(),
24+
);
25+
26+
useEffect(() => {
27+
const isEnabled = config.getPreviewFeatures();
28+
if (isEnabled !== previewEnabled) {
29+
setPreviewEnabled(isEnabled);
30+
}
31+
}, [config, previewEnabled]);
32+
33+
const [bannerCounts] = useState(
34+
() => persistentState.get('defaultBannerShownCount') || {},
35+
);
36+
37+
const hashedText = crypto
38+
.createHash('sha256')
39+
.update(defaultText)
40+
.digest('hex');
41+
42+
const currentBannerCount = bannerCounts[hashedText] || 0;
43+
44+
const showDefaultBanner =
45+
warningText === '' &&
46+
!previewEnabled &&
47+
currentBannerCount < DEFAULT_MAX_BANNER_SHOWN_COUNT;
48+
49+
const rawBannerText = showDefaultBanner ? defaultText : warningText;
50+
const bannerText = rawBannerText.replace(/\\n/g, '\n');
51+
52+
const lastIncrementedKey = useRef<string | null>(null);
53+
54+
useEffect(() => {
55+
if (showDefaultBanner && defaultText) {
56+
if (lastIncrementedKey.current !== defaultText) {
57+
lastIncrementedKey.current = defaultText;
58+
59+
const allCounts = persistentState.get('defaultBannerShownCount') || {};
60+
const current = allCounts[hashedText] || 0;
61+
62+
persistentState.set('defaultBannerShownCount', {
63+
...allCounts,
64+
[hashedText]: current + 1,
65+
});
66+
}
67+
}
68+
}, [showDefaultBanner, defaultText, hashedText]);
69+
70+
return {
71+
bannerText,
72+
};
73+
}

‎packages/cli/src/utils/persistentState.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import * as path from 'node:path';
1111
const STATE_FILENAME = 'state.json';
1212

1313
interface PersistentStateData {
14-
defaultBannerShownCount?: number;
14+
defaultBannerShownCount?: Record<string, number>;
1515
// Add other persistent state keys here as needed
1616
}
1717

0 commit comments

Comments
 (0)