diff --git a/packages/cli/src/test-utils/render.tsx b/packages/cli/src/test-utils/render.tsx index bbc9576ff2..a9f786f11c 100644 --- a/packages/cli/src/test-utils/render.tsx +++ b/packages/cli/src/test-utils/render.tsx @@ -602,6 +602,7 @@ const mockUIActions: UIActions = { import { type TextBuffer } from '../ui/components/shared/text-buffer.js'; import { InputContext, type InputState } from '../ui/contexts/InputContext.js'; +import { QuotaContext, type QuotaState } from '../ui/contexts/QuotaContext.js'; let capturedOverflowState: OverflowState | undefined; let capturedOverflowActions: OverflowActions | undefined; @@ -619,6 +620,7 @@ export const renderWithProviders = async ( shellFocus = true, settings = mockSettings, uiState: providedUiState, + quotaState: providedQuotaState, inputState: providedInputState, width, mouseEventsEnabled = false, @@ -631,6 +633,7 @@ export const renderWithProviders = async ( shellFocus?: boolean; settings?: LoadedSettings; uiState?: Partial; + quotaState?: Partial; inputState?: Partial; width?: number; mouseEventsEnabled?: boolean; @@ -666,6 +669,16 @@ export const renderWithProviders = async ( }, ) as UIState; + const quotaState: QuotaState = { + userTier: undefined, + stats: undefined, + proQuotaRequest: null, + validationRequest: null, + overageMenuRequest: null, + emptyWalletRequest: null, + ...providedQuotaState, + }; + const inputState = { buffer: { text: '' } as unknown as TextBuffer, userMessages: [], @@ -727,65 +740,67 @@ export const renderWithProviders = async ( - - - - - - - - - - + + + + + + + + + - - - - - - - {comp} - - - - - - - - - - - - - - - - + + + + + + + + {comp} + + + + + + + + + + + + + + + + + diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index d78b56e11d..8f05b996dc 100644 --- a/packages/cli/src/ui/AppContainer.test.tsx +++ b/packages/cli/src/ui/AppContainer.test.tsx @@ -123,16 +123,19 @@ vi.mock('ink', async (importOriginal) => { }); import { InputContext, type InputState } from './contexts/InputContext.js'; +import { QuotaContext, type QuotaState } from './contexts/QuotaContext.js'; // Helper component will read the context values provided by AppContainer // so we can assert against them in our tests. let capturedUIState: UIState; let capturedInputState: InputState; +let capturedQuotaState: QuotaState; let capturedUIActions: UIActions; let capturedOverflowActions: OverflowActions; function TestContextConsumer() { capturedUIState = useContext(UIStateContext)!; capturedInputState = useContext(InputContext)!; + capturedQuotaState = useContext(QuotaContext)!; capturedUIActions = useContext(UIActionsContext)!; capturedOverflowActions = useOverflowActions()!; return null; @@ -1309,15 +1312,15 @@ describe('AppContainer State Management', () => { }); describe('Quota and Fallback Integration', () => { - it('passes a null proQuotaRequest to UIStateContext by default', async () => { + it('passes a null proQuotaRequest to QuotaContext by default', async () => { // The default mock from beforeEach already sets proQuotaRequest to null const { unmount } = await act(async () => renderAppContainer()); // Assert that the context value is as expected - expect(capturedUIState.quota.proQuotaRequest).toBeNull(); + expect(capturedQuotaState.proQuotaRequest).toBeNull(); unmount(); }); - it('passes a valid proQuotaRequest to UIStateContext when provided by the hook', async () => { + it('passes a valid proQuotaRequest to QuotaContext when provided by the hook', async () => { // Arrange: Create a mock request object that a UI dialog would receive const mockRequest = { failedModel: 'gemini-pro', @@ -1332,7 +1335,7 @@ describe('AppContainer State Management', () => { // Act: Render the container const { unmount } = await act(async () => renderAppContainer()); // Assert: The mock request is correctly passed through the context - expect(capturedUIState.quota.proQuotaRequest).toEqual(mockRequest); + expect(capturedQuotaState.proQuotaRequest).toEqual(mockRequest); unmount(); }); diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index eaf6fc3e75..f17ac0d756 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -25,6 +25,7 @@ import { import { App } from './App.js'; import { AppContext } from './contexts/AppContext.js'; import { UIStateContext, type UIState } from './contexts/UIStateContext.js'; +import { QuotaContext } from './contexts/QuotaContext.js'; import { UIActionsContext, type UIActions, @@ -2401,6 +2402,26 @@ Logging in with Google... Restarting Gemini CLI to continue. ], ); + const quotaState = useMemo( + () => ({ + userTier, + stats: quotaStats, + proQuotaRequest, + validationRequest, + // G1 AI Credits dialog state + overageMenuRequest, + emptyWalletRequest, + }), + [ + userTier, + quotaStats, + proQuotaRequest, + validationRequest, + overageMenuRequest, + emptyWalletRequest, + ], + ); + const uiState: UIState = useMemo( () => ({ history: historyManager.history, @@ -2473,15 +2494,6 @@ Logging in with Google... Restarting Gemini CLI to continue. showApprovalModeIndicator, allowPlanMode, currentModel, - quota: { - userTier, - stats: quotaStats, - proQuotaRequest, - validationRequest, - // G1 AI Credits dialog state - overageMenuRequest, - emptyWalletRequest, - }, contextFileNames, errorCount, availableTerminalHeight, @@ -2592,12 +2604,6 @@ Logging in with Google... Restarting Gemini CLI to continue. queueErrorMessage, showApprovalModeIndicator, allowPlanMode, - userTier, - quotaStats, - proQuotaRequest, - validationRequest, - overageMenuRequest, - emptyWalletRequest, contextFileNames, errorCount, availableTerminalHeight, @@ -2816,34 +2822,36 @@ Logging in with Google... Restarting Gemini CLI to continue. return ( - - - - - + + + + - - - - - - - - - - - - + + + + + + + + + + + + + + ); }; diff --git a/packages/cli/src/ui/ToolConfirmationFullFrame.test.tsx b/packages/cli/src/ui/ToolConfirmationFullFrame.test.tsx index c8456fb237..5fde51c429 100644 --- a/packages/cli/src/ui/ToolConfirmationFullFrame.test.tsx +++ b/packages/cli/src/ui/ToolConfirmationFullFrame.test.tsx @@ -11,9 +11,9 @@ import { CoreToolCallStatus, ApprovalMode, makeFakeConfig, + type SerializableConfirmationDetails, } from '@google/gemini-cli-core'; import { type UIState } from './contexts/UIStateContext.js'; -import type { SerializableConfirmationDetails } from '@google/gemini-cli-core'; import { act } from 'react'; import { StreamingState } from './types.js'; @@ -107,15 +107,6 @@ describe('Full Terminal Tool Confirmation Snapshot', () => { constrainHeight: true, isConfigInitialized: true, cleanUiDetailsVisible: true, - quota: { - userTier: 'PRO', - stats: { - limits: {}, - usage: {}, - }, - proQuotaRequest: null, - validationRequest: null, - }, pendingHistoryItems: [ { id: 2, @@ -145,6 +136,13 @@ describe('Full Terminal Tool Confirmation Snapshot', () => { const { waitUntilReady, lastFrame, generateSvg, unmount } = await renderWithProviders(, { uiState: mockUIState, + quotaState: { + userTier: 'PRO', + stats: { + remaining: 100, + limit: 1000, + }, + }, config: mockConfig, settings: createMockSettings({ merged: { diff --git a/packages/cli/src/ui/components/Composer.test.tsx b/packages/cli/src/ui/components/Composer.test.tsx index 316b9a1780..8a7ca134a8 100644 --- a/packages/cli/src/ui/components/Composer.test.tsx +++ b/packages/cli/src/ui/components/Composer.test.tsx @@ -201,12 +201,6 @@ const createMockUIState = (overrides: Partial = {}): UIState => isBackgroundTaskVisible: false, embeddedShellFocused: false, showIsExpandableHint: false, - quota: { - userTier: undefined, - stats: undefined, - proQuotaRequest: null, - validationRequest: null, - }, ...overrides, }) as UIState; @@ -245,6 +239,7 @@ const createMockConfig = (overrides = {}): Config => ...overrides, }) as unknown as Config; +import { QuotaContext, type QuotaState } from '../contexts/QuotaContext.js'; import { InputContext, type InputState } from '../contexts/InputContext.js'; const renderComposer = async ( @@ -253,6 +248,7 @@ const renderComposer = async ( config = createMockConfig(), uiActions = createMockUIActions(), inputStateOverrides: Partial = {}, + quotaStateOverrides: Partial = {}, ) => { const inputState = { buffer: { text: '' } as unknown as TextBuffer, @@ -266,16 +262,28 @@ const renderComposer = async ( ...inputStateOverrides, }; + const quotaState: QuotaState = { + userTier: undefined, + stats: undefined, + proQuotaRequest: null, + validationRequest: null, + overageMenuRequest: null, + emptyWalletRequest: null, + ...quotaStateOverrides, + }; + const result = await render( - - - - - - - + + + + + + + + + , ); diff --git a/packages/cli/src/ui/components/DialogManager.test.tsx b/packages/cli/src/ui/components/DialogManager.test.tsx index 31b28f5223..6acc76303c 100644 --- a/packages/cli/src/ui/components/DialogManager.test.tsx +++ b/packages/cli/src/ui/components/DialogManager.test.tsx @@ -9,6 +9,7 @@ import { DialogManager } from './DialogManager.js'; import { describe, it, expect, vi } from 'vitest'; import { Text } from 'ink'; import { type UIState } from '../contexts/UIStateContext.js'; +import { type QuotaState } from '../contexts/QuotaContext.js'; import { type RestartReason } from '../hooks/useIdeTrustListener.js'; import { type IdeInfo } from '@google/gemini-cli-core'; @@ -75,14 +76,6 @@ describe('DialogManager', () => { terminalWidth: 80, confirmUpdateExtensionRequests: [], showIdeRestartPrompt: false, - quota: { - userTier: undefined, - stats: undefined, - proQuotaRequest: null, - validationRequest: null, - overageMenuRequest: null, - emptyWalletRequest: null, - }, shouldShowIdePrompt: false, isFolderTrustDialogOpen: false, loopDetectionConfirmationRequest: null, @@ -112,7 +105,7 @@ describe('DialogManager', () => { unmount(); }); - const testCases: Array<[Partial, string]> = [ + const testCases: Array<[Partial, string, Partial?]> = [ [ { showIdeRestartPrompt: true, @@ -121,23 +114,17 @@ describe('DialogManager', () => { 'IdeTrustChangeDialog', ], [ + {}, + 'ProQuotaDialog', { - quota: { - userTier: undefined, - stats: undefined, - proQuotaRequest: { - failedModel: 'a', - fallbackModel: 'b', - message: 'c', - isTerminalQuotaError: false, - resolve: vi.fn(), - }, - validationRequest: null, - overageMenuRequest: null, - emptyWalletRequest: null, + proQuotaRequest: { + failedModel: 'a', + fallbackModel: 'b', + message: 'c', + isTerminalQuotaError: false, + resolve: vi.fn(), }, }, - 'ProQuotaDialog', ], [ { @@ -195,7 +182,11 @@ describe('DialogManager', () => { it.each(testCases)( 'renders %s when state is %o', - async (uiStateOverride, expectedComponent) => { + async ( + uiStateOverride: Partial, + expectedComponent: string, + quotaStateOverride?: Partial, + ) => { const { lastFrame, unmount } = await renderWithProviders( , { @@ -203,6 +194,7 @@ describe('DialogManager', () => { ...baseUiState, ...uiStateOverride, } as Partial as UIState, + quotaState: quotaStateOverride, }, ); expect(lastFrame()).toContain(expectedComponent); diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index e7e23c834d..b231a62db5 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -27,6 +27,7 @@ import { PermissionsModifyTrustDialog } from './PermissionsModifyTrustDialog.js' import { ModelDialog } from './ModelDialog.js'; import { theme } from '../semantic-colors.js'; import { useUIState } from '../contexts/UIStateContext.js'; +import { useQuotaState } from '../contexts/QuotaContext.js'; import { useUIActions } from '../contexts/UIActionsContext.js'; import { useConfig } from '../contexts/ConfigContext.js'; import { useSettings } from '../contexts/SettingsContext.js'; @@ -52,6 +53,7 @@ export const DialogManager = ({ const settings = useSettings(); const uiState = useUIState(); + const quotaState = useQuotaState(); const uiActions = useUIActions(); const { constrainHeight, @@ -74,54 +76,50 @@ export const DialogManager = ({ /> ); } - if (uiState.quota.proQuotaRequest) { + if (quotaState.proQuotaRequest) { return ( ); } - if (uiState.quota.validationRequest) { + if (quotaState.validationRequest) { return ( ); } - if (uiState.quota.overageMenuRequest) { + if (quotaState.overageMenuRequest) { return ( ); } - if (uiState.quota.emptyWalletRequest) { + if (quotaState.emptyWalletRequest) { return ( ); diff --git a/packages/cli/src/ui/components/Footer.test.tsx b/packages/cli/src/ui/components/Footer.test.tsx index 277a8d3b13..ab242928aa 100644 --- a/packages/cli/src/ui/components/Footer.test.tsx +++ b/packages/cli/src/ui/components/Footer.test.tsx @@ -267,17 +267,12 @@ describe('