From 6dae3a54024d01e95a75bb6cecb2467dccd54067 Mon Sep 17 00:00:00 2001 From: Spencer Date: Mon, 9 Feb 2026 21:53:10 -0500 Subject: [PATCH] Feature/quota visibility 16795 (#18203) --- packages/cli/src/test-utils/render.tsx | 8 +- packages/cli/src/ui/App.test.tsx | 2 +- packages/cli/src/ui/AppContainer.test.tsx | 6 +- packages/cli/src/ui/AppContainer.tsx | 37 ++- .../cli/src/ui/commands/statsCommand.test.ts | 21 +- packages/cli/src/ui/commands/statsCommand.ts | 15 +- .../cli/src/ui/components/AppHeader.test.tsx | 2 +- .../cli/src/ui/components/Composer.test.tsx | 82 +++--- packages/cli/src/ui/components/Composer.tsx | 6 +- .../cli/src/ui/components/ConsentPrompt.tsx | 4 +- .../src/ui/components/DialogManager.test.tsx | 35 ++- .../cli/src/ui/components/DialogManager.tsx | 28 +- .../cli/src/ui/components/Footer.test.tsx | 67 ++++- packages/cli/src/ui/components/Footer.tsx | 16 +- .../src/ui/components/HistoryItemDisplay.tsx | 26 +- .../ui/components/ModelStatsDisplay.test.tsx | 11 +- .../src/ui/components/ModelStatsDisplay.tsx | 37 ++- .../src/ui/components/QuotaDisplay.test.tsx | 73 +++++ .../cli/src/ui/components/QuotaDisplay.tsx | 64 +++++ .../cli/src/ui/components/QuotaStatsInfo.tsx | 65 +++++ .../src/ui/components/StatsDisplay.test.tsx | 65 ++++- .../cli/src/ui/components/StatsDisplay.tsx | 115 +++++--- .../src/ui/components/StatusDisplay.test.tsx | 10 +- .../src/ui/components/ToolStatsDisplay.tsx | 6 +- .../__snapshots__/Footer.test.tsx.snap | 6 + .../ModelStatsDisplay.test.tsx.snap | 15 +- .../__snapshots__/QuotaDisplay.test.tsx.snap | 11 + .../SessionSummaryDisplay.test.tsx.snap | 4 +- .../__snapshots__/StatsDisplay.test.tsx.snap | 73 ++--- .../ToolStatsDisplay.test.tsx.snap | 5 - .../cli/src/ui/contexts/UIStateContext.tsx | 14 +- .../src/ui/hooks/useQuotaAndFallback.test.ts | 4 +- .../cli/src/ui/hooks/useQuotaAndFallback.ts | 4 +- packages/cli/src/ui/types.ts | 28 +- packages/cli/src/ui/utils/displayUtils.ts | 5 +- packages/cli/src/ui/utils/formatters.ts | 26 +- .../core/src/code_assist/codeAssist.test.ts | 7 +- packages/core/src/config/config.test.ts | 272 ++++++++++++++---- packages/core/src/config/config.ts | 269 +++++++++++++++-- .../core/src/core/contentGenerator.test.ts | 51 ++-- .../src/core/loggingContentGenerator.test.ts | 3 +- .../core/src/core/loggingContentGenerator.ts | 9 +- packages/core/src/utils/events.ts | 25 +- 43 files changed, 1315 insertions(+), 317 deletions(-) create mode 100644 packages/cli/src/ui/components/QuotaDisplay.test.tsx create mode 100644 packages/cli/src/ui/components/QuotaDisplay.tsx create mode 100644 packages/cli/src/ui/components/QuotaStatsInfo.tsx create mode 100644 packages/cli/src/ui/components/__snapshots__/QuotaDisplay.test.tsx.snap diff --git a/packages/cli/src/test-utils/render.tsx b/packages/cli/src/test-utils/render.tsx index 2ac08ee977..6b013c16fb 100644 --- a/packages/cli/src/test-utils/render.tsx +++ b/packages/cli/src/test-utils/render.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -151,6 +151,12 @@ const baseMockUiState = { activePtyId: undefined, backgroundShells: new Map(), backgroundShellHeight: 0, + quota: { + userTier: undefined, + stats: undefined, + proQuotaRequest: null, + validationRequest: null, + }, }; export const mockAppState: AppState = { diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index bd663ba195..6a19d80184 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index 1cddd7c094..385185d0d3 100644 --- a/packages/cli/src/ui/AppContainer.test.tsx +++ b/packages/cli/src/ui/AppContainer.test.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -951,7 +951,7 @@ describe('AppContainer State Management', () => { }); await waitFor(() => { // Assert that the context value is as expected - expect(capturedUIState.proQuotaRequest).toBeNull(); + expect(capturedUIState.quota.proQuotaRequest).toBeNull(); }); unmount!(); }); @@ -976,7 +976,7 @@ describe('AppContainer State Management', () => { }); await waitFor(() => { // Assert: The mock request is correctly passed through the context - expect(capturedUIState.proQuotaRequest).toEqual(mockRequest); + expect(capturedUIState.quota.proQuotaRequest).toEqual(mockRequest); }); unmount!(); }); diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index a02512f189..49ca8e1a92 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -29,6 +29,7 @@ import { AuthState, type ConfirmationRequest, type PermissionConfirmationRequest, + type QuotaStats, } from './types.js'; import { checkPermissions } from './hooks/atCommandProcessor.js'; import { MessageType, StreamingState } from './types.js'; @@ -323,6 +324,16 @@ export const AppContainer = (props: AppContainerProps) => { const [currentModel, setCurrentModel] = useState(config.getModel()); const [userTier, setUserTier] = useState(undefined); + const [quotaStats, setQuotaStats] = useState(() => { + const remaining = config.getQuotaRemaining(); + const limit = config.getQuotaLimit(); + const resetTime = config.getQuotaResetTime(); + return remaining !== undefined || + limit !== undefined || + resetTime !== undefined + ? { remaining, limit, resetTime } + : undefined; + }); const [isConfigInitialized, setConfigInitialized] = useState(false); @@ -425,9 +436,23 @@ export const AppContainer = (props: AppContainerProps) => { setCurrentModel(config.getModel()); }; + const handleQuotaChanged = (payload: { + remaining: number | undefined; + limit: number | undefined; + resetTime?: string; + }) => { + setQuotaStats({ + remaining: payload.remaining, + limit: payload.limit, + resetTime: payload.resetTime, + }); + }; + coreEvents.on(CoreEvent.ModelChanged, handleModelChanged); + coreEvents.on(CoreEvent.QuotaChanged, handleQuotaChanged); return () => { coreEvents.off(CoreEvent.ModelChanged, handleModelChanged); + coreEvents.off(CoreEvent.QuotaChanged, handleQuotaChanged); }; }, [config]); @@ -1887,9 +1912,12 @@ Logging in with Google... Restarting Gemini CLI to continue. queueErrorMessage, showApprovalModeIndicator, currentModel, - userTier, - proQuotaRequest, - validationRequest, + quota: { + userTier, + stats: quotaStats, + proQuotaRequest, + validationRequest, + }, contextFileNames, errorCount, availableTerminalHeight, @@ -1994,6 +2022,7 @@ Logging in with Google... Restarting Gemini CLI to continue. queueErrorMessage, showApprovalModeIndicator, userTier, + quotaStats, proQuotaRequest, validationRequest, contextFileNames, diff --git a/packages/cli/src/ui/commands/statsCommand.test.ts b/packages/cli/src/ui/commands/statsCommand.test.ts index f89c76caac..63fe3eb9e5 100644 --- a/packages/cli/src/ui/commands/statsCommand.test.ts +++ b/packages/cli/src/ui/commands/statsCommand.test.ts @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -54,6 +54,7 @@ describe('statsCommand', () => { selectedAuthType: '', tier: undefined, userEmail: 'mock@example.com', + currentModel: undefined, }); }); @@ -63,9 +64,20 @@ describe('statsCommand', () => { const mockQuota = { buckets: [] }; const mockRefreshUserQuota = vi.fn().mockResolvedValue(mockQuota); const mockGetUserTierName = vi.fn().mockReturnValue('Basic'); + const mockGetModel = vi.fn().mockReturnValue('gemini-pro'); + const mockGetQuotaRemaining = vi.fn().mockReturnValue(85); + const mockGetQuotaLimit = vi.fn().mockReturnValue(100); + const mockGetQuotaResetTime = vi + .fn() + .mockReturnValue('2025-01-01T12:00:00Z'); + mockContext.services.config = { refreshUserQuota: mockRefreshUserQuota, getUserTierName: mockGetUserTierName, + getModel: mockGetModel, + getQuotaRemaining: mockGetQuotaRemaining, + getQuotaLimit: mockGetQuotaLimit, + getQuotaResetTime: mockGetQuotaResetTime, } as unknown as Config; await statsCommand.action(mockContext, ''); @@ -75,6 +87,10 @@ describe('statsCommand', () => { expect.objectContaining({ quotas: mockQuota, tier: 'Basic', + currentModel: 'gemini-pro', + pooledRemaining: 85, + pooledLimit: 100, + pooledResetTime: '2025-01-01T12:00:00Z', }), ); }); @@ -93,6 +109,9 @@ describe('statsCommand', () => { selectedAuthType: '', tier: undefined, userEmail: 'mock@example.com', + currentModel: undefined, + pooledRemaining: undefined, + pooledLimit: undefined, }); }); diff --git a/packages/cli/src/ui/commands/statsCommand.ts b/packages/cli/src/ui/commands/statsCommand.ts index 8d4466ba86..b90e7309e1 100644 --- a/packages/cli/src/ui/commands/statsCommand.ts +++ b/packages/cli/src/ui/commands/statsCommand.ts @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -44,6 +44,7 @@ async function defaultSessionView(context: CommandContext) { const wallDuration = now.getTime() - sessionStartTime.getTime(); const { selectedAuthType, userEmail, tier } = getUserIdentity(context); + const currentModel = context.services.config?.getModel(); const statsItem: HistoryItemStats = { type: MessageType.STATS, @@ -51,12 +52,16 @@ async function defaultSessionView(context: CommandContext) { selectedAuthType, userEmail, tier, + currentModel, }; if (context.services.config) { const quota = await context.services.config.refreshUserQuota(); if (quota) { statsItem.quotas = quota; + statsItem.pooledRemaining = context.services.config.getQuotaRemaining(); + statsItem.pooledLimit = context.services.config.getQuotaLimit(); + statsItem.pooledResetTime = context.services.config.getQuotaResetTime(); } } @@ -89,11 +94,19 @@ export const statsCommand: SlashCommand = { autoExecute: true, action: (context: CommandContext) => { const { selectedAuthType, userEmail, tier } = getUserIdentity(context); + const currentModel = context.services.config?.getModel(); + const pooledRemaining = context.services.config?.getQuotaRemaining(); + const pooledLimit = context.services.config?.getQuotaLimit(); + const pooledResetTime = context.services.config?.getQuotaResetTime(); context.ui.addItem({ type: MessageType.MODEL_STATS, selectedAuthType, userEmail, tier, + currentModel, + pooledRemaining, + pooledLimit, + pooledResetTime, } as HistoryItemModelStats); }, }, diff --git a/packages/cli/src/ui/components/AppHeader.test.tsx b/packages/cli/src/ui/components/AppHeader.test.tsx index 13f7b13e77..b827de6dc9 100644 --- a/packages/cli/src/ui/components/AppHeader.test.tsx +++ b/packages/cli/src/ui/components/AppHeader.test.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ diff --git a/packages/cli/src/ui/components/Composer.test.tsx b/packages/cli/src/ui/components/Composer.test.tsx index 73765dcf04..2e59d78772 100644 --- a/packages/cli/src/ui/components/Composer.test.tsx +++ b/packages/cli/src/ui/components/Composer.test.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -24,7 +24,10 @@ vi.mock('../contexts/VimModeContext.js', () => ({ })), })); import { ApprovalMode } from '@google/gemini-cli-core'; +import type { Config } from '@google/gemini-cli-core'; import { StreamingState, ToolCallStatus } from '../types.js'; +import type { LoadedSettings } from '../../config/settings.js'; +import type { SessionMetrics } from '../contexts/SessionContext.js'; // Mock child components vi.mock('./LoadingIndicator.js', () => ({ @@ -145,6 +148,12 @@ const createMockUIState = (overrides: Partial = {}): UIState => activeHooks: [], isBackgroundShellVisible: false, embeddedShellFocused: false, + quota: { + userTier: undefined, + stats: undefined, + proQuotaRequest: null, + validationRequest: null, + }, ...overrides, }) as UIState; @@ -155,31 +164,30 @@ const createMockUIActions = (): UIActions => setShellModeActive: vi.fn(), onEscapePromptChange: vi.fn(), vimHandleInput: vi.fn(), - // eslint-disable-next-line @typescript-eslint/no-explicit-any - }) as any; + }) as Partial as UIActions; -const createMockConfig = (overrides = {}) => ({ - getModel: vi.fn(() => 'gemini-1.5-pro'), - getTargetDir: vi.fn(() => '/test/dir'), - getDebugMode: vi.fn(() => false), - getAccessibility: vi.fn(() => ({})), - getMcpServers: vi.fn(() => ({})), - isPlanEnabled: vi.fn(() => false), - getToolRegistry: () => ({ - getTool: vi.fn(), - }), - getSkillManager: () => ({ - getSkills: () => [], - getDisplayableSkills: () => [], - }), - getMcpClientManager: () => ({ - getMcpServers: () => ({}), - getBlockedMcpServers: () => [], - }), - ...overrides, -}); +const createMockConfig = (overrides = {}): Config => + ({ + getModel: vi.fn(() => 'gemini-1.5-pro'), + getTargetDir: vi.fn(() => '/test/dir'), + getDebugMode: vi.fn(() => false), + getAccessibility: vi.fn(() => ({})), + getMcpServers: vi.fn(() => ({})), + isPlanEnabled: vi.fn(() => false), + getToolRegistry: () => ({ + getTool: vi.fn(), + }), + getSkillManager: () => ({ + getSkills: () => [], + getDisplayableSkills: () => [], + }), + getMcpClientManager: () => ({ + getMcpServers: () => ({}), + getBlockedMcpServers: () => [], + }), + ...overrides, + }) as unknown as Config; -/* eslint-disable @typescript-eslint/no-explicit-any */ const renderComposer = ( uiState: UIState, settings = createMockSettings(), @@ -187,8 +195,8 @@ const renderComposer = ( uiActions = createMockUIActions(), ) => render( - - + + @@ -197,7 +205,6 @@ const renderComposer = ( , ); -/* eslint-enable @typescript-eslint/no-explicit-any */ describe('Composer', () => { describe('Footer Display Settings', () => { @@ -229,8 +236,11 @@ describe('Composer', () => { sessionStats: { sessionId: 'test-session', sessionStartTime: new Date(), - // eslint-disable-next-line @typescript-eslint/no-explicit-any - metrics: {} as any, + metrics: { + models: {}, + tools: {}, + files: {}, + } as SessionMetrics, lastPromptTokenCount: 150, promptCount: 5, }, @@ -251,8 +261,9 @@ describe('Composer', () => { vi.mocked(useVimMode).mockReturnValueOnce({ vimEnabled: true, vimMode: 'INSERT', - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as any); + toggleVimEnabled: vi.fn(), + setVimMode: vi.fn(), + } as unknown as ReturnType); const { lastFrame } = renderComposer(uiState, settings, config); @@ -541,9 +552,12 @@ describe('Composer', () => { const uiState = createMockUIState({ showErrorDetails: true, filteredConsoleMessages: [ - { level: 'error', message: 'Test error', timestamp: new Date() }, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - ] as any, + { + type: 'error', + content: 'Test error', + count: 1, + }, + ], }); const { lastFrame } = renderComposer(uiState); diff --git a/packages/cli/src/ui/components/Composer.tsx b/packages/cli/src/ui/components/Composer.tsx index 2b515fa675..4ccca33e4f 100644 --- a/packages/cli/src/ui/components/Composer.tsx +++ b/packages/cli/src/ui/components/Composer.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -59,8 +59,8 @@ export const Composer = ({ isFocused = true }: { isFocused?: boolean }) => { Boolean(uiState.authConsentRequest) || (uiState.confirmUpdateExtensionRequests?.length ?? 0) > 0 || Boolean(uiState.loopDetectionConfirmationRequest) || - Boolean(uiState.proQuotaRequest) || - Boolean(uiState.validationRequest) || + Boolean(uiState.quota.proQuotaRequest) || + Boolean(uiState.quota.validationRequest) || Boolean(uiState.customDialog); const showLoadingIndicator = (!uiState.embeddedShellFocused || uiState.isBackgroundShellVisible) && diff --git a/packages/cli/src/ui/components/ConsentPrompt.tsx b/packages/cli/src/ui/components/ConsentPrompt.tsx index efa6b136a3..3f255d2606 100644 --- a/packages/cli/src/ui/components/ConsentPrompt.tsx +++ b/packages/cli/src/ui/components/ConsentPrompt.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -25,7 +25,7 @@ export const ConsentPrompt = (props: ConsentPromptProps) => { borderStyle="round" borderColor={theme.border.default} flexDirection="column" - paddingY={1} + paddingTop={1} paddingX={2} > {typeof prompt === 'string' ? ( diff --git a/packages/cli/src/ui/components/DialogManager.test.tsx b/packages/cli/src/ui/components/DialogManager.test.tsx index 78e292e344..da10e97d50 100644 --- a/packages/cli/src/ui/components/DialogManager.test.tsx +++ b/packages/cli/src/ui/components/DialogManager.test.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -75,7 +75,12 @@ describe('DialogManager', () => { terminalWidth: 80, confirmUpdateExtensionRequests: [], showIdeRestartPrompt: false, - proQuotaRequest: null, + quota: { + userTier: undefined, + stats: undefined, + proQuotaRequest: null, + validationRequest: null, + }, shouldShowIdePrompt: false, isFolderTrustDialogOpen: false, loopDetectionConfirmationRequest: null, @@ -99,8 +104,7 @@ describe('DialogManager', () => { it('renders nothing by default', () => { const { lastFrame } = renderWithProviders( , - // eslint-disable-next-line @typescript-eslint/no-explicit-any - { uiState: baseUiState as any }, + { uiState: baseUiState as Partial as UIState }, ); expect(lastFrame()).toBe(''); }); @@ -115,12 +119,17 @@ describe('DialogManager', () => { ], [ { - proQuotaRequest: { - failedModel: 'a', - fallbackModel: 'b', - message: 'c', - isTerminalQuotaError: false, - resolve: vi.fn(), + quota: { + userTier: undefined, + stats: undefined, + proQuotaRequest: { + failedModel: 'a', + fallbackModel: 'b', + message: 'c', + isTerminalQuotaError: false, + resolve: vi.fn(), + }, + validationRequest: null, }, }, 'ProQuotaDialog', @@ -185,8 +194,10 @@ describe('DialogManager', () => { const { lastFrame } = renderWithProviders( , { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - uiState: { ...baseUiState, ...uiStateOverride } as any, + uiState: { + ...baseUiState, + ...uiStateOverride, + } as Partial as UIState, }, ); expect(lastFrame()).toContain(expectedComponent); diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index a502a39030..e4e2f4a6e6 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -71,24 +71,30 @@ export const DialogManager = ({ /> ); } - if (uiState.proQuotaRequest) { + if (uiState.quota.proQuotaRequest) { return ( ); } - if (uiState.validationRequest) { + if (uiState.quota.validationRequest) { return ( ); diff --git a/packages/cli/src/ui/components/Footer.test.tsx b/packages/cli/src/ui/components/Footer.test.tsx index 4113060081..102ddfb1b7 100644 --- a/packages/cli/src/ui/components/Footer.test.tsx +++ b/packages/cli/src/ui/components/Footer.test.tsx @@ -1,10 +1,10 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { renderWithProviders } from '../../test-utils/render.js'; import { createMockSettings } from '../../test-utils/settings.js'; import { Footer } from './Footer.js'; @@ -131,6 +131,69 @@ describe('