Feature/quota visibility 16795 (#18203)

This commit is contained in:
Spencer
2026-02-09 21:53:10 -05:00
committed by GitHub
parent 0a3ecf3a75
commit 6dae3a5402
43 changed files with 1315 additions and 317 deletions
+7 -1
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -151,6 +151,12 @@ const baseMockUiState = {
activePtyId: undefined, activePtyId: undefined,
backgroundShells: new Map(), backgroundShells: new Map(),
backgroundShellHeight: 0, backgroundShellHeight: 0,
quota: {
userTier: undefined,
stats: undefined,
proQuotaRequest: null,
validationRequest: null,
},
}; };
export const mockAppState: AppState = { export const mockAppState: AppState = {
+1 -1
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
+3 -3
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -951,7 +951,7 @@ describe('AppContainer State Management', () => {
}); });
await waitFor(() => { await waitFor(() => {
// Assert that the context value is as expected // Assert that the context value is as expected
expect(capturedUIState.proQuotaRequest).toBeNull(); expect(capturedUIState.quota.proQuotaRequest).toBeNull();
}); });
unmount!(); unmount!();
}); });
@@ -976,7 +976,7 @@ describe('AppContainer State Management', () => {
}); });
await waitFor(() => { await waitFor(() => {
// Assert: The mock request is correctly passed through the context // Assert: The mock request is correctly passed through the context
expect(capturedUIState.proQuotaRequest).toEqual(mockRequest); expect(capturedUIState.quota.proQuotaRequest).toEqual(mockRequest);
}); });
unmount!(); unmount!();
}); });
+33 -4
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -29,6 +29,7 @@ import {
AuthState, AuthState,
type ConfirmationRequest, type ConfirmationRequest,
type PermissionConfirmationRequest, type PermissionConfirmationRequest,
type QuotaStats,
} from './types.js'; } from './types.js';
import { checkPermissions } from './hooks/atCommandProcessor.js'; import { checkPermissions } from './hooks/atCommandProcessor.js';
import { MessageType, StreamingState } from './types.js'; import { MessageType, StreamingState } from './types.js';
@@ -323,6 +324,16 @@ export const AppContainer = (props: AppContainerProps) => {
const [currentModel, setCurrentModel] = useState(config.getModel()); const [currentModel, setCurrentModel] = useState(config.getModel());
const [userTier, setUserTier] = useState<UserTierId | undefined>(undefined); const [userTier, setUserTier] = useState<UserTierId | undefined>(undefined);
const [quotaStats, setQuotaStats] = useState<QuotaStats | undefined>(() => {
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); const [isConfigInitialized, setConfigInitialized] = useState(false);
@@ -425,9 +436,23 @@ export const AppContainer = (props: AppContainerProps) => {
setCurrentModel(config.getModel()); 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.ModelChanged, handleModelChanged);
coreEvents.on(CoreEvent.QuotaChanged, handleQuotaChanged);
return () => { return () => {
coreEvents.off(CoreEvent.ModelChanged, handleModelChanged); coreEvents.off(CoreEvent.ModelChanged, handleModelChanged);
coreEvents.off(CoreEvent.QuotaChanged, handleQuotaChanged);
}; };
}, [config]); }, [config]);
@@ -1887,9 +1912,12 @@ Logging in with Google... Restarting Gemini CLI to continue.
queueErrorMessage, queueErrorMessage,
showApprovalModeIndicator, showApprovalModeIndicator,
currentModel, currentModel,
userTier, quota: {
proQuotaRequest, userTier,
validationRequest, stats: quotaStats,
proQuotaRequest,
validationRequest,
},
contextFileNames, contextFileNames,
errorCount, errorCount,
availableTerminalHeight, availableTerminalHeight,
@@ -1994,6 +2022,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
queueErrorMessage, queueErrorMessage,
showApprovalModeIndicator, showApprovalModeIndicator,
userTier, userTier,
quotaStats,
proQuotaRequest, proQuotaRequest,
validationRequest, validationRequest,
contextFileNames, contextFileNames,
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -54,6 +54,7 @@ describe('statsCommand', () => {
selectedAuthType: '', selectedAuthType: '',
tier: undefined, tier: undefined,
userEmail: 'mock@example.com', userEmail: 'mock@example.com',
currentModel: undefined,
}); });
}); });
@@ -63,9 +64,20 @@ describe('statsCommand', () => {
const mockQuota = { buckets: [] }; const mockQuota = { buckets: [] };
const mockRefreshUserQuota = vi.fn().mockResolvedValue(mockQuota); const mockRefreshUserQuota = vi.fn().mockResolvedValue(mockQuota);
const mockGetUserTierName = vi.fn().mockReturnValue('Basic'); 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 = { mockContext.services.config = {
refreshUserQuota: mockRefreshUserQuota, refreshUserQuota: mockRefreshUserQuota,
getUserTierName: mockGetUserTierName, getUserTierName: mockGetUserTierName,
getModel: mockGetModel,
getQuotaRemaining: mockGetQuotaRemaining,
getQuotaLimit: mockGetQuotaLimit,
getQuotaResetTime: mockGetQuotaResetTime,
} as unknown as Config; } as unknown as Config;
await statsCommand.action(mockContext, ''); await statsCommand.action(mockContext, '');
@@ -75,6 +87,10 @@ describe('statsCommand', () => {
expect.objectContaining({ expect.objectContaining({
quotas: mockQuota, quotas: mockQuota,
tier: 'Basic', tier: 'Basic',
currentModel: 'gemini-pro',
pooledRemaining: 85,
pooledLimit: 100,
pooledResetTime: '2025-01-01T12:00:00Z',
}), }),
); );
}); });
@@ -93,6 +109,9 @@ describe('statsCommand', () => {
selectedAuthType: '', selectedAuthType: '',
tier: undefined, tier: undefined,
userEmail: 'mock@example.com', userEmail: 'mock@example.com',
currentModel: undefined,
pooledRemaining: undefined,
pooledLimit: undefined,
}); });
}); });
+14 -1
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -44,6 +44,7 @@ async function defaultSessionView(context: CommandContext) {
const wallDuration = now.getTime() - sessionStartTime.getTime(); const wallDuration = now.getTime() - sessionStartTime.getTime();
const { selectedAuthType, userEmail, tier } = getUserIdentity(context); const { selectedAuthType, userEmail, tier } = getUserIdentity(context);
const currentModel = context.services.config?.getModel();
const statsItem: HistoryItemStats = { const statsItem: HistoryItemStats = {
type: MessageType.STATS, type: MessageType.STATS,
@@ -51,12 +52,16 @@ async function defaultSessionView(context: CommandContext) {
selectedAuthType, selectedAuthType,
userEmail, userEmail,
tier, tier,
currentModel,
}; };
if (context.services.config) { if (context.services.config) {
const quota = await context.services.config.refreshUserQuota(); const quota = await context.services.config.refreshUserQuota();
if (quota) { if (quota) {
statsItem.quotas = 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, autoExecute: true,
action: (context: CommandContext) => { action: (context: CommandContext) => {
const { selectedAuthType, userEmail, tier } = getUserIdentity(context); 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({ context.ui.addItem({
type: MessageType.MODEL_STATS, type: MessageType.MODEL_STATS,
selectedAuthType, selectedAuthType,
userEmail, userEmail,
tier, tier,
currentModel,
pooledRemaining,
pooledLimit,
pooledResetTime,
} as HistoryItemModelStats); } as HistoryItemModelStats);
}, },
}, },
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -24,7 +24,10 @@ vi.mock('../contexts/VimModeContext.js', () => ({
})), })),
})); }));
import { ApprovalMode } from '@google/gemini-cli-core'; import { ApprovalMode } from '@google/gemini-cli-core';
import type { Config } from '@google/gemini-cli-core';
import { StreamingState, ToolCallStatus } from '../types.js'; import { StreamingState, ToolCallStatus } from '../types.js';
import type { LoadedSettings } from '../../config/settings.js';
import type { SessionMetrics } from '../contexts/SessionContext.js';
// Mock child components // Mock child components
vi.mock('./LoadingIndicator.js', () => ({ vi.mock('./LoadingIndicator.js', () => ({
@@ -145,6 +148,12 @@ const createMockUIState = (overrides: Partial<UIState> = {}): UIState =>
activeHooks: [], activeHooks: [],
isBackgroundShellVisible: false, isBackgroundShellVisible: false,
embeddedShellFocused: false, embeddedShellFocused: false,
quota: {
userTier: undefined,
stats: undefined,
proQuotaRequest: null,
validationRequest: null,
},
...overrides, ...overrides,
}) as UIState; }) as UIState;
@@ -155,31 +164,30 @@ const createMockUIActions = (): UIActions =>
setShellModeActive: vi.fn(), setShellModeActive: vi.fn(),
onEscapePromptChange: vi.fn(), onEscapePromptChange: vi.fn(),
vimHandleInput: vi.fn(), vimHandleInput: vi.fn(),
// eslint-disable-next-line @typescript-eslint/no-explicit-any }) as Partial<UIActions> as UIActions;
}) as any;
const createMockConfig = (overrides = {}) => ({ const createMockConfig = (overrides = {}): Config =>
getModel: vi.fn(() => 'gemini-1.5-pro'), ({
getTargetDir: vi.fn(() => '/test/dir'), getModel: vi.fn(() => 'gemini-1.5-pro'),
getDebugMode: vi.fn(() => false), getTargetDir: vi.fn(() => '/test/dir'),
getAccessibility: vi.fn(() => ({})), getDebugMode: vi.fn(() => false),
getMcpServers: vi.fn(() => ({})), getAccessibility: vi.fn(() => ({})),
isPlanEnabled: vi.fn(() => false), getMcpServers: vi.fn(() => ({})),
getToolRegistry: () => ({ isPlanEnabled: vi.fn(() => false),
getTool: vi.fn(), getToolRegistry: () => ({
}), getTool: vi.fn(),
getSkillManager: () => ({ }),
getSkills: () => [], getSkillManager: () => ({
getDisplayableSkills: () => [], getSkills: () => [],
}), getDisplayableSkills: () => [],
getMcpClientManager: () => ({ }),
getMcpServers: () => ({}), getMcpClientManager: () => ({
getBlockedMcpServers: () => [], getMcpServers: () => ({}),
}), getBlockedMcpServers: () => [],
...overrides, }),
}); ...overrides,
}) as unknown as Config;
/* eslint-disable @typescript-eslint/no-explicit-any */
const renderComposer = ( const renderComposer = (
uiState: UIState, uiState: UIState,
settings = createMockSettings(), settings = createMockSettings(),
@@ -187,8 +195,8 @@ const renderComposer = (
uiActions = createMockUIActions(), uiActions = createMockUIActions(),
) => ) =>
render( render(
<ConfigContext.Provider value={config as any}> <ConfigContext.Provider value={config as unknown as Config}>
<SettingsContext.Provider value={settings as any}> <SettingsContext.Provider value={settings as unknown as LoadedSettings}>
<UIStateContext.Provider value={uiState}> <UIStateContext.Provider value={uiState}>
<UIActionsContext.Provider value={uiActions}> <UIActionsContext.Provider value={uiActions}>
<Composer /> <Composer />
@@ -197,7 +205,6 @@ const renderComposer = (
</SettingsContext.Provider> </SettingsContext.Provider>
</ConfigContext.Provider>, </ConfigContext.Provider>,
); );
/* eslint-enable @typescript-eslint/no-explicit-any */
describe('Composer', () => { describe('Composer', () => {
describe('Footer Display Settings', () => { describe('Footer Display Settings', () => {
@@ -229,8 +236,11 @@ describe('Composer', () => {
sessionStats: { sessionStats: {
sessionId: 'test-session', sessionId: 'test-session',
sessionStartTime: new Date(), sessionStartTime: new Date(),
// eslint-disable-next-line @typescript-eslint/no-explicit-any metrics: {
metrics: {} as any, models: {},
tools: {},
files: {},
} as SessionMetrics,
lastPromptTokenCount: 150, lastPromptTokenCount: 150,
promptCount: 5, promptCount: 5,
}, },
@@ -251,8 +261,9 @@ describe('Composer', () => {
vi.mocked(useVimMode).mockReturnValueOnce({ vi.mocked(useVimMode).mockReturnValueOnce({
vimEnabled: true, vimEnabled: true,
vimMode: 'INSERT', vimMode: 'INSERT',
// eslint-disable-next-line @typescript-eslint/no-explicit-any toggleVimEnabled: vi.fn(),
} as any); setVimMode: vi.fn(),
} as unknown as ReturnType<typeof useVimMode>);
const { lastFrame } = renderComposer(uiState, settings, config); const { lastFrame } = renderComposer(uiState, settings, config);
@@ -541,9 +552,12 @@ describe('Composer', () => {
const uiState = createMockUIState({ const uiState = createMockUIState({
showErrorDetails: true, showErrorDetails: true,
filteredConsoleMessages: [ filteredConsoleMessages: [
{ level: 'error', message: 'Test error', timestamp: new Date() }, {
// eslint-disable-next-line @typescript-eslint/no-explicit-any type: 'error',
] as any, content: 'Test error',
count: 1,
},
],
}); });
const { lastFrame } = renderComposer(uiState); const { lastFrame } = renderComposer(uiState);
+3 -3
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -59,8 +59,8 @@ export const Composer = ({ isFocused = true }: { isFocused?: boolean }) => {
Boolean(uiState.authConsentRequest) || Boolean(uiState.authConsentRequest) ||
(uiState.confirmUpdateExtensionRequests?.length ?? 0) > 0 || (uiState.confirmUpdateExtensionRequests?.length ?? 0) > 0 ||
Boolean(uiState.loopDetectionConfirmationRequest) || Boolean(uiState.loopDetectionConfirmationRequest) ||
Boolean(uiState.proQuotaRequest) || Boolean(uiState.quota.proQuotaRequest) ||
Boolean(uiState.validationRequest) || Boolean(uiState.quota.validationRequest) ||
Boolean(uiState.customDialog); Boolean(uiState.customDialog);
const showLoadingIndicator = const showLoadingIndicator =
(!uiState.embeddedShellFocused || uiState.isBackgroundShellVisible) && (!uiState.embeddedShellFocused || uiState.isBackgroundShellVisible) &&
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -25,7 +25,7 @@ export const ConsentPrompt = (props: ConsentPromptProps) => {
borderStyle="round" borderStyle="round"
borderColor={theme.border.default} borderColor={theme.border.default}
flexDirection="column" flexDirection="column"
paddingY={1} paddingTop={1}
paddingX={2} paddingX={2}
> >
{typeof prompt === 'string' ? ( {typeof prompt === 'string' ? (
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -75,7 +75,12 @@ describe('DialogManager', () => {
terminalWidth: 80, terminalWidth: 80,
confirmUpdateExtensionRequests: [], confirmUpdateExtensionRequests: [],
showIdeRestartPrompt: false, showIdeRestartPrompt: false,
proQuotaRequest: null, quota: {
userTier: undefined,
stats: undefined,
proQuotaRequest: null,
validationRequest: null,
},
shouldShowIdePrompt: false, shouldShowIdePrompt: false,
isFolderTrustDialogOpen: false, isFolderTrustDialogOpen: false,
loopDetectionConfirmationRequest: null, loopDetectionConfirmationRequest: null,
@@ -99,8 +104,7 @@ describe('DialogManager', () => {
it('renders nothing by default', () => { it('renders nothing by default', () => {
const { lastFrame } = renderWithProviders( const { lastFrame } = renderWithProviders(
<DialogManager {...defaultProps} />, <DialogManager {...defaultProps} />,
// eslint-disable-next-line @typescript-eslint/no-explicit-any { uiState: baseUiState as Partial<UIState> as UIState },
{ uiState: baseUiState as any },
); );
expect(lastFrame()).toBe(''); expect(lastFrame()).toBe('');
}); });
@@ -115,12 +119,17 @@ describe('DialogManager', () => {
], ],
[ [
{ {
proQuotaRequest: { quota: {
failedModel: 'a', userTier: undefined,
fallbackModel: 'b', stats: undefined,
message: 'c', proQuotaRequest: {
isTerminalQuotaError: false, failedModel: 'a',
resolve: vi.fn(), fallbackModel: 'b',
message: 'c',
isTerminalQuotaError: false,
resolve: vi.fn(),
},
validationRequest: null,
}, },
}, },
'ProQuotaDialog', 'ProQuotaDialog',
@@ -185,8 +194,10 @@ describe('DialogManager', () => {
const { lastFrame } = renderWithProviders( const { lastFrame } = renderWithProviders(
<DialogManager {...defaultProps} />, <DialogManager {...defaultProps} />,
{ {
// eslint-disable-next-line @typescript-eslint/no-explicit-any uiState: {
uiState: { ...baseUiState, ...uiStateOverride } as any, ...baseUiState,
...uiStateOverride,
} as Partial<UIState> as UIState,
}, },
); );
expect(lastFrame()).toContain(expectedComponent); expect(lastFrame()).toContain(expectedComponent);
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -71,24 +71,30 @@ export const DialogManager = ({
/> />
); );
} }
if (uiState.proQuotaRequest) { if (uiState.quota.proQuotaRequest) {
return ( return (
<ProQuotaDialog <ProQuotaDialog
failedModel={uiState.proQuotaRequest.failedModel} failedModel={uiState.quota.proQuotaRequest.failedModel}
fallbackModel={uiState.proQuotaRequest.fallbackModel} fallbackModel={uiState.quota.proQuotaRequest.fallbackModel}
message={uiState.proQuotaRequest.message} message={uiState.quota.proQuotaRequest.message}
isTerminalQuotaError={uiState.proQuotaRequest.isTerminalQuotaError} isTerminalQuotaError={
isModelNotFoundError={!!uiState.proQuotaRequest.isModelNotFoundError} uiState.quota.proQuotaRequest.isTerminalQuotaError
}
isModelNotFoundError={
!!uiState.quota.proQuotaRequest.isModelNotFoundError
}
onChoice={uiActions.handleProQuotaChoice} onChoice={uiActions.handleProQuotaChoice}
/> />
); );
} }
if (uiState.validationRequest) { if (uiState.quota.validationRequest) {
return ( return (
<ValidationDialog <ValidationDialog
validationLink={uiState.validationRequest.validationLink} validationLink={uiState.quota.validationRequest.validationLink}
validationDescription={uiState.validationRequest.validationDescription} validationDescription={
learnMoreUrl={uiState.validationRequest.learnMoreUrl} uiState.quota.validationRequest.validationDescription
}
learnMoreUrl={uiState.quota.validationRequest.learnMoreUrl}
onChoice={uiActions.handleValidationChoice} onChoice={uiActions.handleValidationChoice}
/> />
); );
+65 -2
View File
@@ -1,10 +1,10 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * 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 { renderWithProviders } from '../../test-utils/render.js';
import { createMockSettings } from '../../test-utils/settings.js'; import { createMockSettings } from '../../test-utils/settings.js';
import { Footer } from './Footer.js'; import { Footer } from './Footer.js';
@@ -131,6 +131,69 @@ describe('<Footer />', () => {
expect(lastFrame()).toMatch(/\(\d+% context left\)/); expect(lastFrame()).toMatch(/\(\d+% context left\)/);
}); });
it('displays the usage indicator when usage is low', () => {
const { lastFrame } = renderWithProviders(<Footer />, {
width: 120,
uiState: {
sessionStats: mockSessionStats,
quota: {
userTier: undefined,
stats: {
remaining: 15,
limit: 100,
resetTime: undefined,
},
proQuotaRequest: null,
validationRequest: null,
},
},
});
expect(lastFrame()).toContain('15%');
expect(lastFrame()).toMatchSnapshot();
});
it('hides the usage indicator when usage is not near limit', () => {
const { lastFrame } = renderWithProviders(<Footer />, {
width: 120,
uiState: {
sessionStats: mockSessionStats,
quota: {
userTier: undefined,
stats: {
remaining: 85,
limit: 100,
resetTime: undefined,
},
proQuotaRequest: null,
validationRequest: null,
},
},
});
expect(lastFrame()).not.toContain('Usage remaining');
expect(lastFrame()).toMatchSnapshot();
});
it('displays "Limit reached" message when remaining is 0', () => {
const { lastFrame } = renderWithProviders(<Footer />, {
width: 120,
uiState: {
sessionStats: mockSessionStats,
quota: {
userTier: undefined,
stats: {
remaining: 0,
limit: 100,
resetTime: undefined,
},
proQuotaRequest: null,
validationRequest: null,
},
},
});
expect(lastFrame()).toContain('Limit reached');
expect(lastFrame()).toMatchSnapshot();
});
it('displays the model name and abbreviated context percentage', () => { it('displays the model name and abbreviated context percentage', () => {
const { lastFrame } = renderWithProviders(<Footer />, { const { lastFrame } = renderWithProviders(<Footer />, {
width: 99, width: 99,
+15 -1
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -17,6 +17,7 @@ import process from 'node:process';
import { ThemedGradient } from './ThemedGradient.js'; import { ThemedGradient } from './ThemedGradient.js';
import { MemoryUsageDisplay } from './MemoryUsageDisplay.js'; import { MemoryUsageDisplay } from './MemoryUsageDisplay.js';
import { ContextUsageDisplay } from './ContextUsageDisplay.js'; import { ContextUsageDisplay } from './ContextUsageDisplay.js';
import { QuotaDisplay } from './QuotaDisplay.js';
import { DebugProfiler } from './DebugProfiler.js'; import { DebugProfiler } from './DebugProfiler.js';
import { isDevelopment } from '../../utils/installationInfo.js'; import { isDevelopment } from '../../utils/installationInfo.js';
import { useUIState } from '../contexts/UIStateContext.js'; import { useUIState } from '../contexts/UIStateContext.js';
@@ -43,6 +44,7 @@ export const Footer: React.FC = () => {
nightly, nightly,
isTrustedFolder, isTrustedFolder,
terminalWidth, terminalWidth,
quotaStats,
} = { } = {
model: uiState.currentModel, model: uiState.currentModel,
targetDir: config.getTargetDir(), targetDir: config.getTargetDir(),
@@ -56,6 +58,7 @@ export const Footer: React.FC = () => {
nightly: uiState.nightly, nightly: uiState.nightly,
isTrustedFolder: uiState.isTrustedFolder, isTrustedFolder: uiState.isTrustedFolder,
terminalWidth: uiState.terminalWidth, terminalWidth: uiState.terminalWidth,
quotaStats: uiState.quota.stats,
}; };
const showMemoryUsage = const showMemoryUsage =
@@ -159,6 +162,17 @@ export const Footer: React.FC = () => {
/> />
</> </>
)} )}
{quotaStats && (
<>
{' '}
<QuotaDisplay
remaining={quotaStats.remaining}
limit={quotaStats.limit}
resetTime={quotaStats.resetTime}
terse={true}
/>
</>
)}
</Text> </Text>
{showMemoryUsage && <MemoryUsageDisplay />} {showMemoryUsage && <MemoryUsageDisplay />}
</Box> </Box>
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -125,6 +125,18 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
selectedAuthType={itemForDisplay.selectedAuthType} selectedAuthType={itemForDisplay.selectedAuthType}
userEmail={itemForDisplay.userEmail} userEmail={itemForDisplay.userEmail}
tier={itemForDisplay.tier} tier={itemForDisplay.tier}
currentModel={itemForDisplay.currentModel}
quotaStats={
itemForDisplay.pooledRemaining !== undefined ||
itemForDisplay.pooledLimit !== undefined ||
itemForDisplay.pooledResetTime !== undefined
? {
remaining: itemForDisplay.pooledRemaining,
limit: itemForDisplay.pooledLimit,
resetTime: itemForDisplay.pooledResetTime,
}
: undefined
}
/> />
)} )}
{itemForDisplay.type === 'model_stats' && ( {itemForDisplay.type === 'model_stats' && (
@@ -132,6 +144,18 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
selectedAuthType={itemForDisplay.selectedAuthType} selectedAuthType={itemForDisplay.selectedAuthType}
userEmail={itemForDisplay.userEmail} userEmail={itemForDisplay.userEmail}
tier={itemForDisplay.tier} tier={itemForDisplay.tier}
currentModel={itemForDisplay.currentModel}
quotaStats={
itemForDisplay.pooledRemaining !== undefined ||
itemForDisplay.pooledLimit !== undefined ||
itemForDisplay.pooledResetTime !== undefined
? {
remaining: itemForDisplay.pooledRemaining,
limit: itemForDisplay.pooledLimit,
resetTime: itemForDisplay.pooledResetTime,
}
: undefined
}
/> />
)} )}
{itemForDisplay.type === 'tool_stats' && <ToolStatsDisplay />} {itemForDisplay.type === 'tool_stats' && <ToolStatsDisplay />}
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -33,7 +33,11 @@ vi.mock('../contexts/SettingsContext.js', async (importOriginal) => {
const useSessionStatsMock = vi.mocked(SessionContext.useSessionStats); const useSessionStatsMock = vi.mocked(SessionContext.useSessionStats);
const useSettingsMock = vi.mocked(SettingsContext.useSettings); const useSettingsMock = vi.mocked(SettingsContext.useSettings);
const renderWithMockedStats = (metrics: SessionMetrics, width?: number) => { const renderWithMockedStats = (
metrics: SessionMetrics,
width?: number,
currentModel: string = 'gemini-2.5-pro',
) => {
useSessionStatsMock.mockReturnValue({ useSessionStatsMock.mockReturnValue({
stats: { stats: {
sessionId: 'test-session', sessionId: 'test-session',
@@ -55,7 +59,7 @@ const renderWithMockedStats = (metrics: SessionMetrics, width?: number) => {
}, },
} as unknown as LoadedSettings); } as unknown as LoadedSettings);
return render(<ModelStatsDisplay />, width); return render(<ModelStatsDisplay currentModel={currentModel} />, width);
}; };
describe('<ModelStatsDisplay />', () => { describe('<ModelStatsDisplay />', () => {
@@ -380,6 +384,7 @@ describe('<ModelStatsDisplay />', () => {
}, },
}, },
80, 80,
'auto-gemini-3',
); );
const output = lastFrame(); const output = lastFrame();
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -16,6 +16,9 @@ import {
import { useSessionStats } from '../contexts/SessionContext.js'; import { useSessionStats } from '../contexts/SessionContext.js';
import { Table, type Column } from './Table.js'; import { Table, type Column } from './Table.js';
import { useSettings } from '../contexts/SettingsContext.js'; import { useSettings } from '../contexts/SettingsContext.js';
import { getDisplayString, isAutoModel } from '@google/gemini-cli-core';
import type { QuotaStats } from '../types.js';
import { QuotaStatsInfo } from './QuotaStatsInfo.js';
interface StatRowData { interface StatRowData {
metric: string; metric: string;
@@ -29,14 +32,23 @@ interface ModelStatsDisplayProps {
selectedAuthType?: string; selectedAuthType?: string;
userEmail?: string; userEmail?: string;
tier?: string; tier?: string;
currentModel?: string;
quotaStats?: QuotaStats;
} }
export const ModelStatsDisplay: React.FC<ModelStatsDisplayProps> = ({ export const ModelStatsDisplay: React.FC<ModelStatsDisplayProps> = ({
selectedAuthType, selectedAuthType,
userEmail, userEmail,
tier, tier,
currentModel,
quotaStats,
}) => { }) => {
const { stats } = useSessionStats(); const { stats } = useSessionStats();
const pooledRemaining = quotaStats?.remaining;
const pooledLimit = quotaStats?.limit;
const pooledResetTime = quotaStats?.resetTime;
const { models } = stats.metrics; const { models } = stats.metrics;
const settings = useSettings(); const settings = useSettings();
const showUserIdentity = settings.merged.ui.showUserIdentity; const showUserIdentity = settings.merged.ui.showUserIdentity;
@@ -49,7 +61,7 @@ export const ModelStatsDisplay: React.FC<ModelStatsDisplayProps> = ({
<Box <Box
borderStyle="round" borderStyle="round"
borderColor={theme.border.default} borderColor={theme.border.default}
paddingY={1} paddingTop={1}
paddingX={2} paddingX={2}
> >
<Text color={theme.text.primary}> <Text color={theme.text.primary}>
@@ -223,16 +235,21 @@ export const ModelStatsDisplay: React.FC<ModelStatsDisplayProps> = ({
})), })),
]; ];
const isAuto = currentModel && isAutoModel(currentModel);
const statsTitle = isAuto
? `${getDisplayString(currentModel)} Stats For Nerds`
: 'Model Stats For Nerds';
return ( return (
<Box <Box
borderStyle="round" borderStyle="round"
borderColor={theme.border.default} borderColor={theme.border.default}
flexDirection="column" flexDirection="column"
paddingY={1} paddingTop={1}
paddingX={2} paddingX={2}
> >
<Text bold color={theme.text.accent}> <Text bold color={theme.text.accent}>
Model Stats For Nerds {statsTitle}
</Text> </Text>
<Box height={1} /> <Box height={1} />
@@ -258,7 +275,17 @@ export const ModelStatsDisplay: React.FC<ModelStatsDisplayProps> = ({
<Text color={theme.text.primary}>{tier}</Text> <Text color={theme.text.primary}>{tier}</Text>
</Box> </Box>
)} )}
{showUserIdentity && (selectedAuthType || tier) && <Box height={1} />} {isAuto &&
pooledRemaining !== undefined &&
pooledLimit !== undefined &&
pooledLimit > 0 && (
<QuotaStatsInfo
remaining={pooledRemaining}
limit={pooledLimit}
resetTime={pooledResetTime}
/>
)}
{(showUserIdentity || isAuto) && <Box height={1} />}
<Table data={rows} columns={columns} /> <Table data={rows} columns={columns} />
</Box> </Box>
@@ -0,0 +1,73 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { render } from '../../test-utils/render.js';
import { describe, it, expect } from 'vitest';
import { QuotaDisplay } from './QuotaDisplay.js';
describe('QuotaDisplay', () => {
it('should not render when remaining is undefined', () => {
const { lastFrame } = render(
<QuotaDisplay remaining={undefined} limit={100} />,
);
expect(lastFrame()).toBe('');
});
it('should not render when limit is undefined', () => {
const { lastFrame } = render(
<QuotaDisplay remaining={100} limit={undefined} />,
);
expect(lastFrame()).toBe('');
});
it('should not render when limit is 0', () => {
const { lastFrame } = render(<QuotaDisplay remaining={100} limit={0} />);
expect(lastFrame()).toBe('');
});
it('should not render when usage > 20%', () => {
const { lastFrame } = render(<QuotaDisplay remaining={85} limit={100} />);
expect(lastFrame()).toBe('');
});
it('should render yellow when usage < 20%', () => {
const { lastFrame } = render(<QuotaDisplay remaining={15} limit={100} />);
expect(lastFrame()).toMatchSnapshot();
});
it('should render red when usage < 5%', () => {
const { lastFrame } = render(<QuotaDisplay remaining={4} limit={100} />);
expect(lastFrame()).toMatchSnapshot();
});
it('should render with reset time when provided', () => {
const resetTime = new Date(Date.now() + 3600000).toISOString(); // 1 hour from now
const { lastFrame } = render(
<QuotaDisplay remaining={15} limit={100} resetTime={resetTime} />,
);
expect(lastFrame()).toMatchSnapshot();
});
it('should NOT render reset time when terse is true', () => {
const resetTime = new Date(Date.now() + 3600000).toISOString();
const { lastFrame } = render(
<QuotaDisplay
remaining={15}
limit={100}
resetTime={resetTime}
terse={true}
/>,
);
expect(lastFrame()).toMatchSnapshot();
});
it('should render terse limit reached message', () => {
const { lastFrame } = render(
<QuotaDisplay remaining={0} limit={100} terse={true} />,
);
expect(lastFrame()).toMatchSnapshot();
});
});
@@ -0,0 +1,64 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type React from 'react';
import { Text } from 'ink';
import {
getStatusColor,
QUOTA_THRESHOLD_HIGH,
QUOTA_THRESHOLD_MEDIUM,
} from '../utils/displayUtils.js';
import { formatResetTime } from '../utils/formatters.js';
interface QuotaDisplayProps {
remaining: number | undefined;
limit: number | undefined;
resetTime?: string;
terse?: boolean;
}
export const QuotaDisplay: React.FC<QuotaDisplayProps> = ({
remaining,
limit,
resetTime,
terse = false,
}) => {
if (remaining === undefined || limit === undefined || limit === 0) {
return null;
}
const percentage = (remaining / limit) * 100;
if (percentage > QUOTA_THRESHOLD_HIGH) {
return null;
}
const color = getStatusColor(percentage, {
green: QUOTA_THRESHOLD_HIGH,
yellow: QUOTA_THRESHOLD_MEDIUM,
});
const resetInfo =
!terse && resetTime ? `, ${formatResetTime(resetTime)}` : '';
if (remaining === 0) {
return (
<Text color={color}>
{terse
? 'Limit reached'
: `/stats Limit reached${resetInfo}${!terse && '. /auth to continue.'}`}
</Text>
);
}
return (
<Text color={color}>
{terse
? `${percentage.toFixed(0)}%`
: `/stats ${percentage.toFixed(0)}% usage remaining${resetInfo}`}
</Text>
);
};
@@ -0,0 +1,65 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type React from 'react';
import { Box, Text } from 'ink';
import { theme } from '../semantic-colors.js';
import { formatResetTime } from '../utils/formatters.js';
import {
getStatusColor,
QUOTA_THRESHOLD_HIGH,
QUOTA_THRESHOLD_MEDIUM,
} from '../utils/displayUtils.js';
interface QuotaStatsInfoProps {
remaining: number | undefined;
limit: number | undefined;
resetTime?: string;
showDetails?: boolean;
}
export const QuotaStatsInfo: React.FC<QuotaStatsInfoProps> = ({
remaining,
limit,
resetTime,
showDetails = true,
}) => {
if (remaining === undefined || limit === undefined || limit === 0) {
return null;
}
const percentage = (remaining / limit) * 100;
const color = getStatusColor(percentage, {
green: QUOTA_THRESHOLD_HIGH,
yellow: QUOTA_THRESHOLD_MEDIUM,
});
return (
<Box flexDirection="column" marginTop={0} marginBottom={0}>
<Text color={color}>
{remaining === 0
? `Limit reached`
: `${percentage.toFixed(0)}% usage remaining`}
{resetTime && `, ${formatResetTime(resetTime)}`}
</Text>
{showDetails && (
<>
<Text color={theme.text.primary}>
Usage limit: {limit.toLocaleString()}
</Text>
<Text color={theme.text.primary}>
Usage limits span all sessions and reset daily.
</Text>
{remaining === 0 && (
<Text color={theme.text.primary}>
Please /auth to upgrade or switch to an API key to continue.
</Text>
)}
</>
)}
</Box>
);
};
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -421,6 +421,7 @@ describe('<StatsDisplay />', () => {
buckets: [ buckets: [
{ {
modelId: 'gemini-2.5-pro', modelId: 'gemini-2.5-pro',
remainingAmount: '75',
remainingFraction: 0.75, remainingFraction: 0.75,
resetTime, resetTime,
}, },
@@ -446,9 +447,64 @@ describe('<StatsDisplay />', () => {
); );
const output = lastFrame(); const output = lastFrame();
expect(output).toContain('Usage left'); expect(output).toContain('Usage remaining');
expect(output).toContain('75.0%'); expect(output).toContain('75.0%');
expect(output).toContain('(Resets in 1h 30m)'); expect(output).toContain('resets in 1h 30m');
expect(output).toMatchSnapshot();
vi.useRealTimers();
});
it('renders pooled quota information for auto mode', () => {
const now = new Date('2025-01-01T12:00:00Z');
vi.useFakeTimers();
vi.setSystemTime(now);
const metrics = createTestMetrics();
const quotas: RetrieveUserQuotaResponse = {
buckets: [
{
modelId: 'gemini-2.5-pro',
remainingAmount: '10',
remainingFraction: 0.1, // limit = 100
},
{
modelId: 'gemini-2.5-flash',
remainingAmount: '700',
remainingFraction: 0.7, // limit = 1000
},
],
};
useSessionStatsMock.mockReturnValue({
stats: {
sessionId: 'test-session-id',
sessionStartTime: new Date(),
metrics,
lastPromptTokenCount: 0,
promptCount: 5,
},
getPromptCount: () => 5,
startNewPrompt: vi.fn(),
});
const { lastFrame } = renderWithProviders(
<StatsDisplay
duration="1s"
quotas={quotas}
currentModel="auto"
quotaStats={{
remaining: 710,
limit: 1100,
}}
/>,
{ width: 100 },
);
const output = lastFrame();
// (10 + 700) / (100 + 1000) = 710 / 1100 = 64.5%
expect(output).toContain('65% usage remaining');
expect(output).toContain('Usage limit: 1,100');
expect(output).toMatchSnapshot(); expect(output).toMatchSnapshot();
vi.useRealTimers(); vi.useRealTimers();
@@ -468,6 +524,7 @@ describe('<StatsDisplay />', () => {
buckets: [ buckets: [
{ {
modelId: 'gemini-2.5-flash', modelId: 'gemini-2.5-flash',
remainingAmount: '50',
remainingFraction: 0.5, remainingFraction: 0.5,
resetTime, resetTime,
}, },
@@ -495,7 +552,7 @@ describe('<StatsDisplay />', () => {
expect(output).toContain('gemini-2.5-flash'); expect(output).toContain('gemini-2.5-flash');
expect(output).toContain('-'); // for requests expect(output).toContain('-'); // for requests
expect(output).toContain('50.0%'); expect(output).toContain('50.0%');
expect(output).toContain('(Resets in 2h)'); expect(output).toContain('resets in 2h');
expect(output).toMatchSnapshot(); expect(output).toMatchSnapshot();
vi.useRealTimers(); vi.useRealTimers();
+71 -44
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -8,7 +8,7 @@ import type React from 'react';
import { Box, Text } from 'ink'; import { Box, Text } from 'ink';
import { ThemedGradient } from './ThemedGradient.js'; import { ThemedGradient } from './ThemedGradient.js';
import { theme } from '../semantic-colors.js'; import { theme } from '../semantic-colors.js';
import { formatDuration } from '../utils/formatters.js'; import { formatDuration, formatResetTime } from '../utils/formatters.js';
import type { ModelMetrics } from '../contexts/SessionContext.js'; import type { ModelMetrics } from '../contexts/SessionContext.js';
import { useSessionStats } from '../contexts/SessionContext.js'; import { useSessionStats } from '../contexts/SessionContext.js';
import { import {
@@ -24,8 +24,12 @@ import { computeSessionStats } from '../utils/computeStats.js';
import { import {
type RetrieveUserQuotaResponse, type RetrieveUserQuotaResponse,
VALID_GEMINI_MODELS, VALID_GEMINI_MODELS,
getDisplayString,
isAutoModel,
} from '@google/gemini-cli-core'; } from '@google/gemini-cli-core';
import { useSettings } from '../contexts/SettingsContext.js'; import { useSettings } from '../contexts/SettingsContext.js';
import type { QuotaStats } from '../types.js';
import { QuotaStatsInfo } from './QuotaStatsInfo.js';
// A more flexible and powerful StatRow component // A more flexible and powerful StatRow component
interface StatRowProps { interface StatRowProps {
@@ -122,36 +126,25 @@ const buildModelRows = (
return [...activeRows, ...quotaRows]; return [...activeRows, ...quotaRows];
}; };
const formatResetTime = (resetTime: string): string => {
const diff = new Date(resetTime).getTime() - Date.now();
if (diff <= 0) return '';
const totalMinutes = Math.ceil(diff / (1000 * 60));
const hours = Math.floor(totalMinutes / 60);
const minutes = totalMinutes % 60;
const fmt = (val: number, unit: 'hour' | 'minute') =>
new Intl.NumberFormat('en', {
style: 'unit',
unit,
unitDisplay: 'narrow',
}).format(val);
if (hours > 0 && minutes > 0) {
return `(Resets in ${fmt(hours, 'hour')} ${fmt(minutes, 'minute')})`;
} else if (hours > 0) {
return `(Resets in ${fmt(hours, 'hour')})`;
}
return `(Resets in ${fmt(minutes, 'minute')})`;
};
const ModelUsageTable: React.FC<{ const ModelUsageTable: React.FC<{
models: Record<string, ModelMetrics>; models: Record<string, ModelMetrics>;
quotas?: RetrieveUserQuotaResponse; quotas?: RetrieveUserQuotaResponse;
cacheEfficiency: number; cacheEfficiency: number;
totalCachedTokens: number; totalCachedTokens: number;
}> = ({ models, quotas, cacheEfficiency, totalCachedTokens }) => { currentModel?: string;
pooledRemaining?: number;
pooledLimit?: number;
pooledResetTime?: string;
}> = ({
models,
quotas,
cacheEfficiency,
totalCachedTokens,
currentModel,
pooledRemaining,
pooledLimit,
pooledResetTime,
}) => {
const rows = buildModelRows(models, quotas); const rows = buildModelRows(models, quotas);
if (rows.length === 0) { if (rows.length === 0) {
@@ -179,13 +172,43 @@ const ModelUsageTable: React.FC<{
? usageLimitWidth ? usageLimitWidth
: uncachedWidth + cachedWidth + outputTokensWidth); : uncachedWidth + cachedWidth + outputTokensWidth);
const isAuto = currentModel && isAutoModel(currentModel);
const modelUsageTitle = isAuto
? `${getDisplayString(currentModel)} Usage`
: `Model Usage`;
return ( return (
<Box flexDirection="column" marginTop={1}> <Box flexDirection="column" marginTop={1}>
{/* Header */} {/* Header */}
<Box alignItems="flex-end"> <Box alignItems="flex-end">
<Box width={nameWidth}> <Box width={nameWidth}>
<Text bold color={theme.text.primary} wrap="truncate-end"> <Text bold color={theme.text.primary} wrap="truncate-end">
Model Usage {modelUsageTitle}
</Text>
</Box>
</Box>
{isAuto &&
showQuotaColumn &&
pooledRemaining !== undefined &&
pooledLimit !== undefined &&
pooledLimit > 0 && (
<Box flexDirection="column" marginTop={0} marginBottom={1}>
<QuotaStatsInfo
remaining={pooledRemaining}
limit={pooledLimit}
resetTime={pooledResetTime}
/>
<Text color={theme.text.primary}>
For a full token breakdown, run `/stats model`.
</Text>
</Box>
)}
<Box alignItems="flex-end">
<Box width={nameWidth}>
<Text bold color={theme.text.primary}>
Model
</Text> </Text>
</Box> </Box>
<Box <Box
@@ -198,6 +221,7 @@ const ModelUsageTable: React.FC<{
Reqs Reqs
</Text> </Text>
</Box> </Box>
{!showQuotaColumn && ( {!showQuotaColumn && (
<> <>
<Box <Box
@@ -239,7 +263,7 @@ const ModelUsageTable: React.FC<{
alignItems="flex-end" alignItems="flex-end"
> >
<Text bold color={theme.text.primary}> <Text bold color={theme.text.primary}>
Usage left Usage remaining
</Text> </Text>
</Box> </Box>
)} )}
@@ -259,7 +283,10 @@ const ModelUsageTable: React.FC<{
{rows.map((row) => ( {rows.map((row) => (
<Box key={row.key}> <Box key={row.key}>
<Box width={nameWidth}> <Box width={nameWidth}>
<Text color={theme.text.primary} wrap="truncate-end"> <Text
color={row.isActive ? theme.text.primary : theme.text.secondary}
wrap="truncate-end"
>
{row.modelName} {row.modelName}
</Text> </Text>
</Box> </Box>
@@ -344,19 +371,6 @@ const ModelUsageTable: React.FC<{
</Text> </Text>
</Box> </Box>
)} )}
{showQuotaColumn && (
<>
<Box marginTop={1} marginBottom={2}>
<Text color={theme.text.primary}>
{`Usage limits span all sessions and reset daily.\n/auth to upgrade or switch to API key.`}
</Text>
</Box>
<Text color={theme.text.secondary}>
» Tip: For a full token breakdown, run `/stats model`.
</Text>
</>
)}
</Box> </Box>
); );
}; };
@@ -368,6 +382,8 @@ interface StatsDisplayProps {
selectedAuthType?: string; selectedAuthType?: string;
userEmail?: string; userEmail?: string;
tier?: string; tier?: string;
currentModel?: string;
quotaStats?: QuotaStats;
} }
export const StatsDisplay: React.FC<StatsDisplayProps> = ({ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
@@ -377,12 +393,19 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
selectedAuthType, selectedAuthType,
userEmail, userEmail,
tier, tier,
currentModel,
quotaStats,
}) => { }) => {
const { stats } = useSessionStats(); const { stats } = useSessionStats();
const { metrics } = stats; const { metrics } = stats;
const { models, tools, files } = metrics; const { models, tools, files } = metrics;
const computed = computeSessionStats(metrics); const computed = computeSessionStats(metrics);
const settings = useSettings(); const settings = useSettings();
const pooledRemaining = quotaStats?.remaining;
const pooledLimit = quotaStats?.limit;
const pooledResetTime = quotaStats?.resetTime;
const showUserIdentity = settings.merged.ui.showUserIdentity; const showUserIdentity = settings.merged.ui.showUserIdentity;
const successThresholds = { const successThresholds = {
@@ -415,7 +438,7 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
borderStyle="round" borderStyle="round"
borderColor={theme.border.default} borderColor={theme.border.default}
flexDirection="column" flexDirection="column"
paddingY={1} paddingTop={1}
paddingX={2} paddingX={2}
overflow="hidden" overflow="hidden"
> >
@@ -508,6 +531,10 @@ export const StatsDisplay: React.FC<StatsDisplayProps> = ({
quotas={quotas} quotas={quotas}
cacheEfficiency={computed.cacheEfficiency} cacheEfficiency={computed.cacheEfficiency}
totalCachedTokens={computed.totalCachedTokens} totalCachedTokens={computed.totalCachedTokens}
currentModel={currentModel}
pooledRemaining={pooledRemaining}
pooledLimit={pooledLimit}
pooledResetTime={pooledResetTime}
/> />
</Box> </Box>
); );
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -12,6 +12,8 @@ import { UIStateContext, type UIState } from '../contexts/UIStateContext.js';
import { TransientMessageType } from '../../utils/events.js'; import { TransientMessageType } from '../../utils/events.js';
import { ConfigContext } from '../contexts/ConfigContext.js'; import { ConfigContext } from '../contexts/ConfigContext.js';
import { SettingsContext } from '../contexts/SettingsContext.js'; import { SettingsContext } from '../contexts/SettingsContext.js';
import type { Config } from '@google/gemini-cli-core';
import type { LoadedSettings } from '../../config/settings.js';
import { createMockSettings } from '../../test-utils/settings.js'; import { createMockSettings } from '../../test-utils/settings.js';
import type { TextBuffer } from './shared/text-buffer.js'; import type { TextBuffer } from './shared/text-buffer.js';
@@ -68,7 +70,6 @@ const createMockConfig = (overrides = {}) => ({
...overrides, ...overrides,
}); });
/* eslint-disable @typescript-eslint/no-explicit-any */
const renderStatusDisplay = ( const renderStatusDisplay = (
props: { hideContextSummary: boolean } = { hideContextSummary: false }, props: { hideContextSummary: boolean } = { hideContextSummary: false },
uiState: UIState = createMockUIState(), uiState: UIState = createMockUIState(),
@@ -76,15 +77,14 @@ const renderStatusDisplay = (
config = createMockConfig(), config = createMockConfig(),
) => ) =>
render( render(
<ConfigContext.Provider value={config as any}> <ConfigContext.Provider value={config as unknown as Config}>
<SettingsContext.Provider value={settings as any}> <SettingsContext.Provider value={settings as unknown as LoadedSettings}>
<UIStateContext.Provider value={uiState}> <UIStateContext.Provider value={uiState}>
<StatusDisplay {...props} /> <StatusDisplay {...props} />
</UIStateContext.Provider> </UIStateContext.Provider>
</SettingsContext.Provider> </SettingsContext.Provider>
</ConfigContext.Provider>, </ConfigContext.Provider>,
); );
/* eslint-enable @typescript-eslint/no-explicit-any */
describe('StatusDisplay', () => { describe('StatusDisplay', () => {
const originalEnv = process.env; const originalEnv = process.env;
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -64,7 +64,7 @@ export const ToolStatsDisplay: React.FC = () => {
<Box <Box
borderStyle="round" borderStyle="round"
borderColor={theme.border.default} borderColor={theme.border.default}
paddingY={1} paddingTop={1}
paddingX={2} paddingX={2}
> >
<Text color={theme.text.primary}> <Text color={theme.text.primary}>
@@ -98,7 +98,7 @@ export const ToolStatsDisplay: React.FC = () => {
borderStyle="round" borderStyle="round"
borderColor={theme.border.default} borderColor={theme.border.default}
flexDirection="column" flexDirection="column"
paddingY={1} paddingTop={1}
paddingX={2} paddingX={2}
width={70} width={70}
> >
@@ -1,5 +1,9 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`<Footer /> > displays "Limit reached" message when remaining is 0 1`] = `" ...directories/to/make/it/long no sandbox (see /docs) gemini-pro /model Limit reached"`;
exports[`<Footer /> > displays the usage indicator when usage is low 1`] = `" ...directories/to/make/it/long no sandbox (see /docs) gemini-pro /model 15%"`;
exports[`<Footer /> > footer configuration filtering (golden snapshots) > renders complete footer in narrow terminal (baseline narrow) > complete-footer-narrow 1`] = `" ...s/to/make/it/long no sandbox gemini-pro /model (100%)"`; exports[`<Footer /> > footer configuration filtering (golden snapshots) > renders complete footer in narrow terminal (baseline narrow) > complete-footer-narrow 1`] = `" ...s/to/make/it/long no sandbox gemini-pro /model (100%)"`;
exports[`<Footer /> > footer configuration filtering (golden snapshots) > renders complete footer with all sections visible (baseline) > complete-footer-wide 1`] = `" ...directories/to/make/it/long no sandbox (see /docs) gemini-pro /model (100% context left)"`; exports[`<Footer /> > footer configuration filtering (golden snapshots) > renders complete footer with all sections visible (baseline) > complete-footer-wide 1`] = `" ...directories/to/make/it/long no sandbox (see /docs) gemini-pro /model (100% context left)"`;
@@ -9,3 +13,5 @@ exports[`<Footer /> > footer configuration filtering (golden snapshots) > render
exports[`<Footer /> > footer configuration filtering (golden snapshots) > renders footer with all optional sections hidden (minimal footer) > footer-minimal 1`] = `""`; exports[`<Footer /> > footer configuration filtering (golden snapshots) > renders footer with all optional sections hidden (minimal footer) > footer-minimal 1`] = `""`;
exports[`<Footer /> > footer configuration filtering (golden snapshots) > renders footer with only model info hidden (partial filtering) > footer-no-model 1`] = `" ...directories/to/make/it/long no sandbox (see /docs)"`; exports[`<Footer /> > footer configuration filtering (golden snapshots) > renders footer with only model info hidden (partial filtering) > footer-no-model 1`] = `" ...directories/to/make/it/long no sandbox (see /docs)"`;
exports[`<Footer /> > hides the usage indicator when usage is not near limit 1`] = `" ...directories/to/make/it/long no sandbox (see /docs) gemini-pro /model"`;
@@ -5,6 +5,7 @@ exports[`<ModelStatsDisplay /> > should display a single model correctly 1`] = `
│ │ │ │
│ Model Stats For Nerds │ │ Model Stats For Nerds │
│ │ │ │
│ │
│ Metric gemini-2.5-pro │ │ Metric gemini-2.5-pro │
│ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │
│ API │ │ API │
@@ -18,7 +19,6 @@ exports[`<ModelStatsDisplay /> > should display a single model correctly 1`] = `
│ ↳ Thoughts 2 │ │ ↳ Thoughts 2 │
│ ↳ Tool 1 │ │ ↳ Tool 1 │
│ ↳ Output 20 │ │ ↳ Output 20 │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -27,6 +27,7 @@ exports[`<ModelStatsDisplay /> > should display conditional rows if at least one
│ │ │ │
│ Model Stats For Nerds │ │ Model Stats For Nerds │
│ │ │ │
│ │
│ Metric gemini-2.5-pro gemini-2.5-flash │ │ Metric gemini-2.5-pro gemini-2.5-flash │
│ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │
│ API │ │ API │
@@ -40,7 +41,6 @@ exports[`<ModelStatsDisplay /> > should display conditional rows if at least one
│ ↳ Thoughts 2 0 │ │ ↳ Thoughts 2 0 │
│ ↳ Tool 0 3 │ │ ↳ Tool 0 3 │
│ ↳ Output 20 10 │ │ ↳ Output 20 10 │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -49,6 +49,7 @@ exports[`<ModelStatsDisplay /> > should display stats for multiple models correc
│ │ │ │
│ Model Stats For Nerds │ │ Model Stats For Nerds │
│ │ │ │
│ │
│ Metric gemini-2.5-pro gemini-2.5-flash │ │ Metric gemini-2.5-pro gemini-2.5-flash │
│ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │
│ API │ │ API │
@@ -62,7 +63,6 @@ exports[`<ModelStatsDisplay /> > should display stats for multiple models correc
│ ↳ Thoughts 10 20 │ │ ↳ Thoughts 10 20 │
│ ↳ Tool 5 10 │ │ ↳ Tool 5 10 │
│ ↳ Output 200 400 │ │ ↳ Output 200 400 │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -71,6 +71,7 @@ exports[`<ModelStatsDisplay /> > should handle large values without wrapping or
│ │ │ │
│ Model Stats For Nerds │ │ Model Stats For Nerds │
│ │ │ │
│ │
│ Metric gemini-2.5-pro │ │ Metric gemini-2.5-pro │
│ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │
│ API │ │ API │
@@ -84,14 +85,14 @@ exports[`<ModelStatsDisplay /> > should handle large values without wrapping or
│ ↳ Thoughts 111,111,111 │ │ ↳ Thoughts 111,111,111 │
│ ↳ Tool 222,222,222 │ │ ↳ Tool 222,222,222 │
│ ↳ Output 123,456,789 │ │ ↳ Output 123,456,789 │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
exports[`<ModelStatsDisplay /> > should handle models with long names (gemini-3-*-preview) without layout breaking 1`] = ` exports[`<ModelStatsDisplay /> > should handle models with long names (gemini-3-*-preview) without layout breaking 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────╮ "╭──────────────────────────────────────────────────────────────────────────────╮
│ │ │ │
Model Stats For Nerds Auto (Gemini 3) Stats For Nerds │
│ │
│ │ │ │
│ Metric gemini-3-pro-preview gemini-3-flash-preview │ │ Metric gemini-3-pro-preview gemini-3-flash-preview │
│ ────────────────────────────────────────────────────────────────────────── │ │ ────────────────────────────────────────────────────────────────────────── │
@@ -106,7 +107,6 @@ exports[`<ModelStatsDisplay /> > should handle models with long names (gemini-3-
│ ↳ Thoughts 100 200 │ │ ↳ Thoughts 100 200 │
│ ↳ Tool 50 100 │ │ ↳ Tool 50 100 │
│ ↳ Output 4,000 8,000 │ │ ↳ Output 4,000 8,000 │
│ │
╰──────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -115,6 +115,7 @@ exports[`<ModelStatsDisplay /> > should not display conditional rows if no model
│ │ │ │
│ Model Stats For Nerds │ │ Model Stats For Nerds │
│ │ │ │
│ │
│ Metric gemini-2.5-pro │ │ Metric gemini-2.5-pro │
│ ────────────────────────────────────────────────────────────────────────────────────────────── │ │ ────────────────────────────────────────────────────────────────────────────────────────────── │
│ API │ │ API │
@@ -125,7 +126,6 @@ exports[`<ModelStatsDisplay /> > should not display conditional rows if no model
│ Total 30 │ │ Total 30 │
│ ↳ Input 10 │ │ ↳ Input 10 │
│ ↳ Output 20 │ │ ↳ Output 20 │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -133,6 +133,5 @@ exports[`<ModelStatsDisplay /> > should render "no API calls" message when there
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ "╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │ │ │
│ No API calls have been made in this session. │ │ No API calls have been made in this session. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -0,0 +1,11 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`QuotaDisplay > should NOT render reset time when terse is true 1`] = `"15%"`;
exports[`QuotaDisplay > should render red when usage < 5% 1`] = `"/stats 4% usage remaining"`;
exports[`QuotaDisplay > should render terse limit reached message 1`] = `"Limit reached"`;
exports[`QuotaDisplay > should render with reset time when provided 1`] = `"/stats 15% usage remaining, resets in 1h"`;
exports[`QuotaDisplay > should render yellow when usage < 20% 1`] = `"/stats 15% usage remaining"`;
@@ -18,11 +18,11 @@ exports[`<SessionSummaryDisplay /> > renders the summary display with a title 1`
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │ │ │
│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ Model Usage
│ Model Reqs Input Tokens Cache Reads Output Tokens │
│ ──────────────────────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────────────────────── │
│ gemini-2.5-pro 10 500 500 2,000 │ │ gemini-2.5-pro 10 500 500 2,000 │
│ │ │ │
│ Savings Highlight: 500 (50.0%) of input tokens were served from the cache, reducing costs. │ │ Savings Highlight: 500 (50.0%) of input tokens were served from the cache, reducing costs. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -17,7 +17,6 @@ exports[`<StatsDisplay /> > Code Changes Display > displays Code Changes when li
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 100ms (100.0%) │ │ » Tool Time: 100ms (100.0%) │
│ │ │ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -37,7 +36,6 @@ exports[`<StatsDisplay /> > Code Changes Display > hides Code Changes when no li
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 100ms (100.0%) │ │ » Tool Time: 100ms (100.0%) │
│ │ │ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -57,7 +55,6 @@ exports[`<StatsDisplay /> > Conditional Color Tests > renders success rate in gr
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -77,7 +74,6 @@ exports[`<StatsDisplay /> > Conditional Color Tests > renders success rate in re
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -97,7 +93,6 @@ exports[`<StatsDisplay /> > Conditional Color Tests > renders success rate in ye
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -118,10 +113,10 @@ exports[`<StatsDisplay /> > Conditional Rendering Tests > hides Efficiency secti
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │ │ │
│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ Model Usage
│ Model Reqs Input Tokens Cache Reads Output Tokens │
│ ──────────────────────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────────────────────── │
│ gemini-2.5-pro 1 100 0 100 │ │ gemini-2.5-pro 1 100 0 100 │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -141,7 +136,36 @@ exports[`<StatsDisplay /> > Conditional Rendering Tests > hides User Agreement w
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 123ms (100.0%) │ │ » Tool Time: 123ms (100.0%) │
│ │ │ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`;
exports[`<StatsDisplay /> > Quota Display > renders pooled quota information for auto mode 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │ │ │
│ Session Stats │
│ │
│ Interaction Summary │
│ Session ID: test-session-id │
│ Tool Calls: 0 ( ✓ 0 x 0 ) │
│ Success Rate: 0.0% │
│ │
│ Performance │
│ Wall Time: 1s │
│ Agent Active: 0s │
│ » API Time: 0s (0.0%) │
│ » Tool Time: 0s (0.0%) │
│ │
│ │
│ auto Usage │
│ 65% usage remaining │
│ Usage limit: 1,100 │
│ Usage limits span all sessions and reset daily. │
│ For a full token breakdown, run \`/stats model\`. │
│ │
│ Model Reqs Usage remaining │
│ ──────────────────────────────────────────────────────────── │
│ gemini-2.5-pro - │
│ gemini-2.5-flash - │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -162,16 +186,10 @@ exports[`<StatsDisplay /> > Quota Display > renders quota information for unused
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │ │ │
│ Model Usage Reqs Usage left │ Model Usage
│ Model Reqs Usage remaining │
│ ──────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────── │
│ gemini-2.5-flash - 50.0% (Resets in 2h) │ gemini-2.5-flash - 50.0% resets in 2h │
│ │
│ Usage limits span all sessions and reset daily. │
│ /auth to upgrade or switch to API key. │
│ │
│ │
│ » Tip: For a full token breakdown, run \`/stats model\`. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -192,16 +210,10 @@ exports[`<StatsDisplay /> > Quota Display > renders quota information when quota
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │ │ │
│ Model Usage Reqs Usage left │ Model Usage
│ Model Reqs Usage remaining │
│ ──────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────── │
│ gemini-2.5-pro 1 75.0% (Resets in 1h 30m) │ gemini-2.5-pro 1 75.0% resets in 1h 30m │
│ │
│ Usage limits span all sessions and reset daily. │
│ /auth to upgrade or switch to API key. │
│ │
│ │
│ » Tip: For a full token breakdown, run \`/stats model\`. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -221,7 +233,6 @@ exports[`<StatsDisplay /> > Title Rendering > renders the custom title when a ti
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -241,7 +252,6 @@ exports[`<StatsDisplay /> > Title Rendering > renders the default title when no
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -262,13 +272,13 @@ exports[`<StatsDisplay /> > renders a table with two models correctly 1`] = `
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │ │ │
│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ Model Usage
│ Model Reqs Input Tokens Cache Reads Output Tokens │
│ ──────────────────────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────────────────────── │
│ gemini-2.5-pro 3 500 500 2,000 │ │ gemini-2.5-pro 3 500 500 2,000 │
│ gemini-2.5-flash 5 15,000 10,000 15,000 │ │ gemini-2.5-flash 5 15,000 10,000 15,000 │
│ │ │ │
│ Savings Highlight: 10,500 (40.4%) of input tokens were served from the cache, reducing costs. │ │ Savings Highlight: 10,500 (40.4%) of input tokens were served from the cache, reducing costs. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -290,12 +300,12 @@ exports[`<StatsDisplay /> > renders all sections when all data is present 1`] =
│ » Tool Time: 123ms (55.2%) │ │ » Tool Time: 123ms (55.2%) │
│ │ │ │
│ │ │ │
│ Model Usage Reqs Input Tokens Cache Reads Output Tokens │ Model Usage
│ Model Reqs Input Tokens Cache Reads Output Tokens │
│ ──────────────────────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────────────────────── │
│ gemini-2.5-pro 1 50 50 100 │ │ gemini-2.5-pro 1 50 50 100 │
│ │ │ │
│ Savings Highlight: 50 (50.0%) of input tokens were served from the cache, reducing costs. │ │ Savings Highlight: 50 (50.0%) of input tokens were served from the cache, reducing costs. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -315,6 +325,5 @@ exports[`<StatsDisplay /> > renders only the Performance section in its zero sta
│ » API Time: 0s (0.0%) │ │ » API Time: 0s (0.0%) │
│ » Tool Time: 0s (0.0%) │ │ » Tool Time: 0s (0.0%) │
│ │ │ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -16,7 +16,6 @@ exports[`<ToolStatsDisplay /> > should display stats for a single tool correctly
│ » Modified: 0 │ │ » Modified: 0 │
│ ──────────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────────── │
│ Overall Agreement Rate: 100.0% │ │ Overall Agreement Rate: 100.0% │
│ │
╰────────────────────────────────────────────────────────────────────╯" ╰────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -37,7 +36,6 @@ exports[`<ToolStatsDisplay /> > should display stats for multiple tools correctl
│ » Modified: 1 │ │ » Modified: 1 │
│ ──────────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────────── │
│ Overall Agreement Rate: 33.3% │ │ Overall Agreement Rate: 33.3% │
│ │
╰────────────────────────────────────────────────────────────────────╯" ╰────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -58,7 +56,6 @@ exports[`<ToolStatsDisplay /> > should handle large values without wrapping or o
│ » Modified: 12345 │ │ » Modified: 12345 │
│ ──────────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────────── │
│ Overall Agreement Rate: 55.6% │ │ Overall Agreement Rate: 55.6% │
│ │
╰────────────────────────────────────────────────────────────────────╯" ╰────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -78,7 +75,6 @@ exports[`<ToolStatsDisplay /> > should handle zero decisions gracefully 1`] = `
│ » Modified: 0 │ │ » Modified: 0 │
│ ──────────────────────────────────────────────────────────────── │ │ ──────────────────────────────────────────────────────────────── │
│ Overall Agreement Rate: -- │ │ Overall Agreement Rate: -- │
│ │
╰────────────────────────────────────────────────────────────────────╯" ╰────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -86,6 +82,5 @@ exports[`<ToolStatsDisplay /> > should render "no tool calls" message when there
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ "╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │ │ │
│ No tool calls have been made in this session. │ │ No tool calls have been made in this session. │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯"
`; `;
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -10,6 +10,7 @@ import type {
ThoughtSummary, ThoughtSummary,
ConsoleMessageItem, ConsoleMessageItem,
ConfirmationRequest, ConfirmationRequest,
QuotaStats,
LoopDetectionConfirmationRequest, LoopDetectionConfirmationRequest,
HistoryItemWithoutId, HistoryItemWithoutId,
StreamingState, StreamingState,
@@ -54,6 +55,13 @@ import { type RestartReason } from '../hooks/useIdeTrustListener.js';
import type { TerminalBackgroundColor } from '../utils/terminalCapabilityManager.js'; import type { TerminalBackgroundColor } from '../utils/terminalCapabilityManager.js';
import type { BackgroundShell } from '../hooks/shellCommandProcessor.js'; import type { BackgroundShell } from '../hooks/shellCommandProcessor.js';
export interface QuotaState {
userTier: UserTierId | undefined;
stats: QuotaStats | undefined;
proQuotaRequest: ProQuotaDialogRequest | null;
validationRequest: ValidationDialogRequest | null;
}
export interface UIState { export interface UIState {
history: HistoryItem[]; history: HistoryItem[];
historyManager: UseHistoryManagerReturn; historyManager: UseHistoryManagerReturn;
@@ -120,9 +128,7 @@ export interface UIState {
queueErrorMessage: string | null; queueErrorMessage: string | null;
showApprovalModeIndicator: ApprovalMode; showApprovalModeIndicator: ApprovalMode;
// Quota-related state // Quota-related state
userTier: UserTierId | undefined; quota: QuotaState;
proQuotaRequest: ProQuotaDialogRequest | null;
validationRequest: ValidationDialogRequest | null;
currentModel: string; currentModel: string;
contextFileNames: string[]; contextFileNames: string[];
errorCount: number; errorCount: number;
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -157,7 +157,7 @@ describe('useQuotaAndFallback', () => {
const message = request!.message; const message = request!.message;
expect(message).toContain('Usage limit reached for gemini-pro.'); expect(message).toContain('Usage limit reached for gemini-pro.');
expect(message).toContain('Access resets at'); // From getResetTimeMessage expect(message).toContain('Access resets at'); // From getResetTimeMessage
expect(message).toContain('/stats for usage details'); expect(message).toContain('/stats model for usage details');
expect(message).toContain('/auth to switch to API key.'); expect(message).toContain('/auth to switch to API key.');
expect(mockHistoryManager.addItem).not.toHaveBeenCalled(); expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -78,7 +78,7 @@ export function useQuotaAndFallback({
const messageLines = [ const messageLines = [
`Usage limit reached for ${usageLimitReachedModel}.`, `Usage limit reached for ${usageLimitReachedModel}.`,
error.retryDelayMs ? getResetTimeMessage(error.retryDelayMs) : null, error.retryDelayMs ? getResetTimeMessage(error.retryDelayMs) : null,
`/stats for usage details`, `/stats model for usage details`,
`/model to switch models.`, `/model to switch models.`,
`/auth to switch to API key.`, `/auth to switch to API key.`,
].filter(Boolean); ].filter(Boolean);
+19 -9
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -153,20 +153,30 @@ export type HistoryItemHelp = HistoryItemBase & {
timestamp: Date; timestamp: Date;
}; };
export type HistoryItemStats = HistoryItemBase & { export interface HistoryItemQuotaBase extends HistoryItemBase {
selectedAuthType?: string;
userEmail?: string;
tier?: string;
currentModel?: string;
pooledRemaining?: number;
pooledLimit?: number;
pooledResetTime?: string;
}
export interface QuotaStats {
remaining: number | undefined;
limit: number | undefined;
resetTime?: string;
}
export type HistoryItemStats = HistoryItemQuotaBase & {
type: 'stats'; type: 'stats';
duration: string; duration: string;
quotas?: RetrieveUserQuotaResponse; quotas?: RetrieveUserQuotaResponse;
selectedAuthType?: string;
userEmail?: string;
tier?: string;
}; };
export type HistoryItemModelStats = HistoryItemBase & { export type HistoryItemModelStats = HistoryItemQuotaBase & {
type: 'model_stats'; type: 'model_stats';
selectedAuthType?: string;
userEmail?: string;
tier?: string;
}; };
export type HistoryItemToolStats = HistoryItemBase & { export type HistoryItemToolStats = HistoryItemBase & {
+4 -1
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -16,6 +16,9 @@ export const USER_AGREEMENT_RATE_MEDIUM = 45;
export const CACHE_EFFICIENCY_HIGH = 40; export const CACHE_EFFICIENCY_HIGH = 40;
export const CACHE_EFFICIENCY_MEDIUM = 15; export const CACHE_EFFICIENCY_MEDIUM = 15;
export const QUOTA_THRESHOLD_HIGH = 20;
export const QUOTA_THRESHOLD_MEDIUM = 5;
// --- Color Logic --- // --- Color Logic ---
export const getStatusColor = ( export const getStatusColor = (
value: number, value: number,
+25 -1
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -97,3 +97,27 @@ export function stripReferenceContent(text: string): string {
return text.replace(pattern, '').trim(); return text.replace(pattern, '').trim();
} }
export const formatResetTime = (resetTime: string): string => {
const diff = new Date(resetTime).getTime() - Date.now();
if (diff <= 0) return '';
const totalMinutes = Math.ceil(diff / (1000 * 60));
const hours = Math.floor(totalMinutes / 60);
const minutes = totalMinutes % 60;
const fmt = (val: number, unit: 'hour' | 'minute') =>
new Intl.NumberFormat('en', {
style: 'unit',
unit,
unitDisplay: 'narrow',
}).format(val);
if (hours > 0 && minutes > 0) {
return `resets in ${fmt(hours, 'hour')} ${fmt(minutes, 'minute')}`;
} else if (hours > 0) {
return `resets in ${fmt(hours, 'hour')}`;
}
return `resets in ${fmt(minutes, 'minute')}`;
};
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -43,6 +43,7 @@ describe('codeAssist', () => {
const mockUserData = { const mockUserData = {
projectId: 'test-project', projectId: 'test-project',
userTier: UserTierId.FREE, userTier: UserTierId.FREE,
userTierName: 'free-tier-name',
}; };
it('should create a server for LOGIN_WITH_GOOGLE', async () => { it('should create a server for LOGIN_WITH_GOOGLE', async () => {
@@ -70,7 +71,7 @@ describe('codeAssist', () => {
httpOptions, httpOptions,
'session-123', 'session-123',
'free-tier', 'free-tier',
undefined, 'free-tier-name',
); );
expect(generator).toBeInstanceOf(MockedCodeAssistServer); expect(generator).toBeInstanceOf(MockedCodeAssistServer);
}); });
@@ -99,7 +100,7 @@ describe('codeAssist', () => {
httpOptions, httpOptions,
undefined, // No session ID undefined, // No session ID
'free-tier', 'free-tier',
undefined, 'free-tier-name',
); );
expect(generator).toBeInstanceOf(MockedCodeAssistServer); expect(generator).toBeInstanceOf(MockedCodeAssistServer);
}); });
+224 -48
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -38,8 +38,9 @@ import { RipgrepFallbackEvent } from '../telemetry/types.js';
import { ToolRegistry } from '../tools/tool-registry.js'; import { ToolRegistry } from '../tools/tool-registry.js';
import { ACTIVATE_SKILL_TOOL_NAME } from '../tools/tool-names.js'; import { ACTIVATE_SKILL_TOOL_NAME } from '../tools/tool-names.js';
import type { SkillDefinition } from '../skills/skillLoader.js'; import type { SkillDefinition } from '../skills/skillLoader.js';
import type { McpClientManager } from '../tools/mcp-client-manager.js';
import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js'; import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js';
import { DEFAULT_GEMINI_MODEL, PREVIEW_GEMINI_MODEL } from './models.js'; import { DEFAULT_GEMINI_MODEL } from './models.js';
vi.mock('fs', async (importOriginal) => { vi.mock('fs', async (importOriginal) => {
const actual = await importOriginal<typeof import('fs')>(); const actual = await importOriginal<typeof import('fs')>();
@@ -169,6 +170,7 @@ const mockCoreEvents = vi.hoisted(() => ({
emitFeedback: vi.fn(), emitFeedback: vi.fn(),
emitModelChanged: vi.fn(), emitModelChanged: vi.fn(),
emitConsoleLog: vi.fn(), emitConsoleLog: vi.fn(),
emitQuotaChanged: vi.fn(),
on: vi.fn(), on: vi.fn(),
})); }));
@@ -203,7 +205,9 @@ import { getCodeAssistServer } from '../code_assist/codeAssist.js';
import { getExperiments } from '../code_assist/experiments/experiments.js'; import { getExperiments } from '../code_assist/experiments/experiments.js';
import type { CodeAssistServer } from '../code_assist/server.js'; import type { CodeAssistServer } from '../code_assist/server.js';
import { ContextManager } from '../services/contextManager.js'; import { ContextManager } from '../services/contextManager.js';
import { UserTierId } from 'src/code_assist/types.js'; import { UserTierId } from '../code_assist/types.js';
import type { ModelConfigService } from '../services/modelConfigService.js';
import type { ModelConfigServiceConfig } from '../services/modelConfigService.js';
import { ExitPlanModeTool } from '../tools/exit-plan-mode.js'; import { ExitPlanModeTool } from '../tools/exit-plan-mode.js';
import { EnterPlanModeTool } from '../tools/enter-plan-mode.js'; import { EnterPlanModeTool } from '../tools/enter-plan-mode.js';
@@ -253,7 +257,7 @@ describe('Server Config (config.ts)', () => {
describe('initialize', () => { describe('initialize', () => {
it('should throw an error if checkpointing is enabled and GitService fails', async () => { it('should throw an error if checkpointing is enabled and GitService fails', async () => {
const gitError = new Error('Git is not installed'); const gitError = new Error('Git is not installed');
(GitService.prototype.initialize as Mock).mockRejectedValue(gitError); vi.mocked(GitService.prototype.initialize).mockRejectedValue(gitError);
const config = new Config({ const config = new Config({
...baseParams, ...baseParams,
@@ -265,7 +269,7 @@ describe('Server Config (config.ts)', () => {
it('should not throw an error if checkpointing is disabled and GitService fails', async () => { it('should not throw an error if checkpointing is disabled and GitService fails', async () => {
const gitError = new Error('Git is not installed'); const gitError = new Error('Git is not installed');
(GitService.prototype.initialize as Mock).mockRejectedValue(gitError); vi.mocked(GitService.prototype.initialize).mockRejectedValue(gitError);
const config = new Config({ const config = new Config({
...baseParams, ...baseParams,
@@ -299,13 +303,16 @@ describe('Server Config (config.ts)', () => {
); );
let mcpStarted = false; let mcpStarted = false;
(McpClientManager as unknown as Mock).mockImplementation(() => ({ vi.mocked(McpClientManager).mockImplementation(
startConfiguredMcpServers: vi.fn().mockImplementation(async () => { () =>
await new Promise((resolve) => setTimeout(resolve, 50)); ({
mcpStarted = true; startConfiguredMcpServers: vi.fn().mockImplementation(async () => {
}), await new Promise((resolve) => setTimeout(resolve, 50));
getMcpInstructions: vi.fn(), mcpStarted = true;
})); }),
getMcpInstructions: vi.fn(),
}) as Partial<McpClientManager> as McpClientManager,
);
await config.initialize(); await config.initialize();
@@ -329,13 +336,16 @@ describe('Server Config (config.ts)', () => {
resolveMcp = resolve; resolveMcp = resolve;
}); });
(McpClientManager as unknown as Mock).mockImplementation(() => ({ (McpClientManager as unknown as Mock).mockImplementation(
startConfiguredMcpServers: vi.fn().mockImplementation(async () => { () =>
await mcpPromise; ({
mcpStarted = true; startConfiguredMcpServers: vi.fn().mockImplementation(async () => {
}), await mcpPromise;
getMcpInstructions: vi.fn(), mcpStarted = true;
})); }),
getMcpInstructions: vi.fn(),
}) as Partial<McpClientManager> as McpClientManager,
);
await config.initialize(); await config.initialize();
@@ -459,7 +469,9 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGeneratorConfig).mockImplementation( vi.mocked(createContentGeneratorConfig).mockImplementation(
async (_: Config, authType: AuthType | undefined) => async (_: Config, authType: AuthType | undefined) =>
({ authType }) as unknown as ContentGeneratorConfig, ({
authType,
}) as Partial<ContentGeneratorConfig> as ContentGeneratorConfig,
); );
await config.refreshAuth(AuthType.USE_GEMINI); await config.refreshAuth(AuthType.USE_GEMINI);
@@ -472,7 +484,9 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGeneratorConfig).mockImplementation( vi.mocked(createContentGeneratorConfig).mockImplementation(
async (_: Config, authType: AuthType | undefined) => async (_: Config, authType: AuthType | undefined) =>
({ authType }) as unknown as ContentGeneratorConfig, ({
authType,
}) as Partial<ContentGeneratorConfig> as ContentGeneratorConfig,
); );
await config.refreshAuth(AuthType.USE_GEMINI); await config.refreshAuth(AuthType.USE_GEMINI);
@@ -489,7 +503,9 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGeneratorConfig).mockImplementation( vi.mocked(createContentGeneratorConfig).mockImplementation(
async (_: Config, authType: AuthType | undefined) => async (_: Config, authType: AuthType | undefined) =>
({ authType }) as unknown as ContentGeneratorConfig, ({
authType,
}) as Partial<ContentGeneratorConfig> as ContentGeneratorConfig,
); );
await config.refreshAuth(AuthType.USE_GEMINI); await config.refreshAuth(AuthType.USE_GEMINI);
@@ -506,7 +522,9 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGeneratorConfig).mockImplementation( vi.mocked(createContentGeneratorConfig).mockImplementation(
async (_: Config, authType: AuthType | undefined) => async (_: Config, authType: AuthType | undefined) =>
({ authType }) as unknown as ContentGeneratorConfig, ({
authType,
}) as Partial<ContentGeneratorConfig> as ContentGeneratorConfig,
); );
await config.refreshAuth(AuthType.USE_VERTEX_AI); await config.refreshAuth(AuthType.USE_VERTEX_AI);
@@ -1268,7 +1286,7 @@ describe('setApprovalMode with folder trust', () => {
getTool: vi.fn().mockReturnValue(undefined), getTool: vi.fn().mockReturnValue(undefined),
unregisterTool: vi.fn(), unregisterTool: vi.fn(),
registerTool: vi.fn(), registerTool: vi.fn(),
} as unknown as ReturnType<Config['getToolRegistry']>); } as Partial<ToolRegistry> as ToolRegistry);
const updateSpy = vi.spyOn(config, 'updateSystemInstructionIfInitialized'); const updateSpy = vi.spyOn(config, 'updateSystemInstructionIfInitialized');
config.setApprovalMode(ApprovalMode.PLAN); config.setApprovalMode(ApprovalMode.PLAN);
@@ -1286,7 +1304,7 @@ describe('setApprovalMode with folder trust', () => {
getTool: vi.fn().mockReturnValue(undefined), getTool: vi.fn().mockReturnValue(undefined),
unregisterTool: vi.fn(), unregisterTool: vi.fn(),
registerTool: vi.fn(), registerTool: vi.fn(),
} as unknown as ReturnType<Config['getToolRegistry']>); } as Partial<ToolRegistry> as ToolRegistry);
const updateSpy = vi.spyOn(config, 'updateSystemInstructionIfInitialized'); const updateSpy = vi.spyOn(config, 'updateSystemInstructionIfInitialized');
config.setApprovalMode(ApprovalMode.DEFAULT); config.setApprovalMode(ApprovalMode.DEFAULT);
@@ -1310,11 +1328,11 @@ describe('setApprovalMode with folder trust', () => {
}); });
it('should register RipGrepTool when useRipgrep is true and it is available', async () => { it('should register RipGrepTool when useRipgrep is true and it is available', async () => {
(canUseRipgrep as Mock).mockResolvedValue(true); vi.mocked(canUseRipgrep).mockResolvedValue(true);
const config = new Config({ ...baseParams, useRipgrep: true }); const config = new Config({ ...baseParams, useRipgrep: true });
await config.initialize(); await config.initialize();
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls; const calls = vi.mocked(ToolRegistry.prototype.registerTool).mock.calls;
const wasRipGrepRegistered = calls.some( const wasRipGrepRegistered = calls.some(
(call) => call[0] instanceof vi.mocked(RipGrepTool), (call) => call[0] instanceof vi.mocked(RipGrepTool),
); );
@@ -1328,11 +1346,11 @@ describe('setApprovalMode with folder trust', () => {
}); });
it('should register GrepTool as a fallback when useRipgrep is true but it is not available', async () => { it('should register GrepTool as a fallback when useRipgrep is true but it is not available', async () => {
(canUseRipgrep as Mock).mockResolvedValue(false); vi.mocked(canUseRipgrep).mockResolvedValue(false);
const config = new Config({ ...baseParams, useRipgrep: true }); const config = new Config({ ...baseParams, useRipgrep: true });
await config.initialize(); await config.initialize();
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls; const calls = vi.mocked(ToolRegistry.prototype.registerTool).mock.calls;
const wasRipGrepRegistered = calls.some( const wasRipGrepRegistered = calls.some(
(call) => call[0] instanceof vi.mocked(RipGrepTool), (call) => call[0] instanceof vi.mocked(RipGrepTool),
); );
@@ -1346,17 +1364,17 @@ describe('setApprovalMode with folder trust', () => {
config, config,
expect.any(RipgrepFallbackEvent), expect.any(RipgrepFallbackEvent),
); );
const event = (logRipgrepFallback as Mock).mock.calls[0][1]; const event = vi.mocked(logRipgrepFallback).mock.calls[0][1];
expect(event.error).toBeUndefined(); expect(event.error).toBeUndefined();
}); });
it('should register GrepTool as a fallback when canUseRipgrep throws an error', async () => { it('should register GrepTool as a fallback when canUseRipgrep throws an error', async () => {
const error = new Error('ripGrep check failed'); const error = new Error('ripGrep check failed');
(canUseRipgrep as Mock).mockRejectedValue(error); vi.mocked(canUseRipgrep).mockRejectedValue(error);
const config = new Config({ ...baseParams, useRipgrep: true }); const config = new Config({ ...baseParams, useRipgrep: true });
await config.initialize(); await config.initialize();
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls; const calls = vi.mocked(ToolRegistry.prototype.registerTool).mock.calls;
const wasRipGrepRegistered = calls.some( const wasRipGrepRegistered = calls.some(
(call) => call[0] instanceof vi.mocked(RipGrepTool), (call) => call[0] instanceof vi.mocked(RipGrepTool),
); );
@@ -1370,7 +1388,7 @@ describe('setApprovalMode with folder trust', () => {
config, config,
expect.any(RipgrepFallbackEvent), expect.any(RipgrepFallbackEvent),
); );
const event = (logRipgrepFallback as Mock).mock.calls[0][1]; const event = vi.mocked(logRipgrepFallback).mock.calls[0][1];
expect(event.error).toBe(String(error)); expect(event.error).toBe(String(error));
}); });
@@ -1378,7 +1396,7 @@ describe('setApprovalMode with folder trust', () => {
const config = new Config({ ...baseParams, useRipgrep: false }); const config = new Config({ ...baseParams, useRipgrep: false });
await config.initialize(); await config.initialize();
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls; const calls = vi.mocked(ToolRegistry.prototype.registerTool).mock.calls;
const wasRipGrepRegistered = calls.some( const wasRipGrepRegistered = calls.some(
(call) => call[0] instanceof vi.mocked(RipGrepTool), (call) => call[0] instanceof vi.mocked(RipGrepTool),
); );
@@ -1526,8 +1544,11 @@ describe('Generation Config Merging (HACK)', () => {
}; };
const config = new Config(params); const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any const serviceConfig = (
const serviceConfig = (config.modelConfigService as any).config; config.modelConfigService as Partial<ModelConfigService> as {
config: ModelConfigServiceConfig;
}
).config;
// Assert that the default aliases are present // Assert that the default aliases are present
expect(serviceConfig.aliases).toEqual(DEFAULT_MODEL_CONFIGS.aliases); expect(serviceConfig.aliases).toEqual(DEFAULT_MODEL_CONFIGS.aliases);
@@ -1550,8 +1571,11 @@ describe('Generation Config Merging (HACK)', () => {
}; };
const config = new Config(params); const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any const serviceConfig = (
const serviceConfig = (config.modelConfigService as any).config; config.modelConfigService as Partial<ModelConfigService> as {
config: ModelConfigServiceConfig;
}
).config;
// Assert that the user's aliases are present // Assert that the user's aliases are present
expect(serviceConfig.aliases).toEqual(userAliases); expect(serviceConfig.aliases).toEqual(userAliases);
@@ -1574,8 +1598,11 @@ describe('Generation Config Merging (HACK)', () => {
}; };
const config = new Config(params); const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any const serviceConfig = (
const serviceConfig = (config.modelConfigService as any).config; config.modelConfigService as Partial<ModelConfigService> as {
config: ModelConfigServiceConfig;
}
).config;
// Assert that the user's aliases are used, not the defaults // Assert that the user's aliases are used, not the defaults
expect(serviceConfig.aliases).toEqual(userAliases); expect(serviceConfig.aliases).toEqual(userAliases);
@@ -1585,8 +1612,11 @@ describe('Generation Config Merging (HACK)', () => {
const params: ConfigParameters = { ...baseParams }; const params: ConfigParameters = { ...baseParams };
const config = new Config(params); const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any const serviceConfig = (
const serviceConfig = (config.modelConfigService as any).config; config.modelConfigService as Partial<ModelConfigService> as {
config: ModelConfigServiceConfig;
}
).config;
// Assert that the full default config is used // Assert that the full default config is used
expect(serviceConfig).toEqual(DEFAULT_MODEL_CONFIGS); expect(serviceConfig).toEqual(DEFAULT_MODEL_CONFIGS);
@@ -1942,8 +1972,10 @@ describe('Hooks configuration', () => {
describe('Config Quota & Preview Model Access', () => { describe('Config Quota & Preview Model Access', () => {
let config: Config; let config: Config;
// eslint-disable-next-line @typescript-eslint/no-explicit-any let mockCodeAssistServer: {
let mockCodeAssistServer: any; projectId: string;
retrieveUserQuota: Mock;
};
const baseParams: ConfigParameters = { const baseParams: ConfigParameters = {
cwd: '/tmp', cwd: '/tmp',
@@ -1965,14 +1997,22 @@ describe('Config Quota & Preview Model Access', () => {
projectId: 'test-project', projectId: 'test-project',
retrieveUserQuota: vi.fn(), retrieveUserQuota: vi.fn(),
}; };
vi.mocked(getCodeAssistServer).mockReturnValue(mockCodeAssistServer); vi.mocked(getCodeAssistServer).mockReturnValue(
mockCodeAssistServer as Partial<CodeAssistServer> as CodeAssistServer,
);
config = new Config(baseParams); config = new Config(baseParams);
}); });
describe('refreshUserQuota', () => { describe('refreshUserQuota', () => {
it('should update hasAccessToPreviewModel to true if quota includes preview model', async () => { it('should update hasAccessToPreviewModel to true if quota includes preview model', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({ mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [{ modelId: PREVIEW_GEMINI_MODEL }], buckets: [
{
modelId: 'gemini-3-pro-preview',
remainingAmount: '100',
remainingFraction: 1.0,
},
],
}); });
await config.refreshUserQuota(); await config.refreshUserQuota();
@@ -1981,13 +2021,82 @@ describe('Config Quota & Preview Model Access', () => {
it('should update hasAccessToPreviewModel to false if quota does not include preview model', async () => { it('should update hasAccessToPreviewModel to false if quota does not include preview model', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({ mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [{ modelId: 'some-other-model' }], buckets: [
{
modelId: 'some-other-model',
remainingAmount: '10',
remainingFraction: 0.1,
},
],
}); });
await config.refreshUserQuota(); await config.refreshUserQuota();
expect(config.getHasAccessToPreviewModel()).toBe(false); expect(config.getHasAccessToPreviewModel()).toBe(false);
}); });
it('should calculate pooled quota correctly for auto models', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [
{
modelId: 'gemini-2.5-pro',
remainingAmount: '10',
remainingFraction: 0.2,
},
{
modelId: 'gemini-2.5-flash',
remainingAmount: '80',
remainingFraction: 0.8,
},
],
});
config.setModel('auto-gemini-2.5');
await config.refreshUserQuota();
const pooled = (
config as Partial<Config> as {
getPooledQuota: () => {
remaining?: number;
limit?: number;
resetTime?: string;
};
}
).getPooledQuota();
// Pro: 10 / 0.2 = 50 total.
// Flash: 80 / 0.8 = 100 total.
// Pooled: (10 + 80) / (50 + 100) = 90 / 150 = 0.6
expect(pooled?.remaining).toBe(90);
expect(pooled?.limit).toBe(150);
expect((pooled?.remaining ?? 0) / (pooled?.limit ?? 1)).toBeCloseTo(0.6);
});
it('should return undefined pooled quota for non-auto models', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [
{
modelId: 'gemini-2.5-pro',
remainingAmount: '10',
remainingFraction: 0.2,
},
],
});
config.setModel('gemini-2.5-pro');
await config.refreshUserQuota();
expect(
(
config as Partial<Config> as {
getPooledQuota: () => {
remaining?: number;
limit?: number;
resetTime?: string;
};
}
).getPooledQuota(),
).toEqual({});
});
it('should update hasAccessToPreviewModel to false if buckets are undefined', async () => { it('should update hasAccessToPreviewModel to false if buckets are undefined', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({}); mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({});
@@ -2013,6 +2122,73 @@ describe('Config Quota & Preview Model Access', () => {
}); });
}); });
describe('refreshUserQuotaIfStale', () => {
beforeEach(() => {
vi.useFakeTimers();
vi.setSystemTime(new Date('2025-01-01T00:00:00Z'));
});
afterEach(() => {
vi.useRealTimers();
});
it('should refresh quota if stale', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [],
});
// First call to initialize lastQuotaFetchTime
await config.refreshUserQuota();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(1);
// Advance time by 31 seconds (default TTL is 30s)
vi.setSystemTime(Date.now() + 31_000);
await config.refreshUserQuotaIfStale();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(2);
});
it('should not refresh quota if fresh', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [],
});
// First call
await config.refreshUserQuota();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(1);
// Advance time by only 10 seconds
vi.setSystemTime(Date.now() + 10_000);
await config.refreshUserQuotaIfStale();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(1);
});
it('should respect custom staleMs', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [],
});
// First call
await config.refreshUserQuota();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(1);
// Advance time by 5 seconds
vi.setSystemTime(Date.now() + 5_000);
// Refresh with 2s staleMs -> should refresh
await config.refreshUserQuotaIfStale(2_000);
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(2);
// Advance by another 5 seconds
vi.setSystemTime(Date.now() + 5_000);
// Refresh with 10s staleMs -> should NOT refresh
await config.refreshUserQuotaIfStale(10_000);
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(2);
});
});
describe('getUserTier and getUserTierName', () => { describe('getUserTier and getUserTierName', () => {
it('should return undefined if contentGenerator is not initialized', () => { it('should return undefined if contentGenerator is not initialized', () => {
const config = new Config(baseParams); const config = new Config(baseParams);
@@ -2032,7 +2208,7 @@ describe('Config Quota & Preview Model Access', () => {
vi.mocked(createContentGenerator).mockResolvedValue({ vi.mocked(createContentGenerator).mockResolvedValue({
userTier: mockTier, userTier: mockTier,
userTierName: mockTierName, userTierName: mockTierName,
} as unknown as CodeAssistServer); } as Partial<CodeAssistServer> as CodeAssistServer);
await config.refreshAuth(AuthType.USE_GEMINI); await config.refreshAuth(AuthType.USE_GEMINI);
+240 -29
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -53,9 +53,14 @@ import { tokenLimit } from '../core/tokenLimits.js';
import { import {
DEFAULT_GEMINI_EMBEDDING_MODEL, DEFAULT_GEMINI_EMBEDDING_MODEL,
DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_MODEL_AUTO, DEFAULT_GEMINI_MODEL_AUTO,
isAutoModel,
isPreviewModel, isPreviewModel,
PREVIEW_GEMINI_FLASH_MODEL,
PREVIEW_GEMINI_MODEL, PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL_AUTO,
resolveModel,
} from './models.js'; } from './models.js';
import { shouldAttemptBrowserLaunch } from '../utils/browser.js'; import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
import type { MCPOAuthConfig } from '../mcp/oauth-provider.js'; import type { MCPOAuthConfig } from '../mcp/oauth-provider.js';
@@ -87,14 +92,14 @@ import { ContextManager } from '../services/contextManager.js';
import type { GenerateContentParameters } from '@google/genai'; import type { GenerateContentParameters } from '@google/genai';
// Re-export OAuth config type // Re-export OAuth config type
export type { MCPOAuthConfig, AnyToolInvocation }; export type { MCPOAuthConfig, AnyToolInvocation, AnyDeclarativeTool };
import type { AnyToolInvocation } from '../tools/tools.js'; import type { AnyToolInvocation, AnyDeclarativeTool } from '../tools/tools.js';
import { WorkspaceContext } from '../utils/workspaceContext.js'; import { WorkspaceContext } from '../utils/workspaceContext.js';
import { Storage } from './storage.js'; import { Storage } from './storage.js';
import type { ShellExecutionConfig } from '../services/shellExecutionService.js'; import type { ShellExecutionConfig } from '../services/shellExecutionService.js';
import { FileExclusions } from '../utils/ignorePatterns.js'; import { FileExclusions } from '../utils/ignorePatterns.js';
import type { EventEmitter } from 'node:events';
import { MessageBus } from '../confirmation-bus/message-bus.js'; import { MessageBus } from '../confirmation-bus/message-bus.js';
import type { EventEmitter } from 'node:events';
import { PolicyEngine } from '../policy/policy-engine.js'; import { PolicyEngine } from '../policy/policy-engine.js';
import { ApprovalMode, type PolicyEngineConfig } from '../policy/types.js'; import { ApprovalMode, type PolicyEngineConfig } from '../policy/types.js';
import { HookSystem } from '../hooks/index.js'; import { HookSystem } from '../hooks/index.js';
@@ -564,6 +569,31 @@ export class Config {
fallbackModelHandler?: FallbackModelHandler; fallbackModelHandler?: FallbackModelHandler;
validationHandler?: ValidationHandler; validationHandler?: ValidationHandler;
private quotaErrorOccurred: boolean = false; private quotaErrorOccurred: boolean = false;
private modelQuotas: Map<
string,
{ remaining: number; limit: number; resetTime?: string }
> = new Map();
private lastRetrievedQuota?: RetrieveUserQuotaResponse;
private lastQuotaFetchTime = 0;
private lastEmittedQuotaRemaining: number | undefined;
private lastEmittedQuotaLimit: number | undefined;
private emitQuotaChangedEvent(): void {
const pooled = this.getPooledQuota();
if (
this.lastEmittedQuotaRemaining !== pooled.remaining ||
this.lastEmittedQuotaLimit !== pooled.limit
) {
this.lastEmittedQuotaRemaining = pooled.remaining;
this.lastEmittedQuotaLimit = pooled.limit;
coreEvents.emitQuotaChanged(
pooled.remaining,
pooled.limit,
pooled.resetTime,
);
}
}
private readonly summarizeToolOutput: private readonly summarizeToolOutput:
| Record<string, SummarizeToolOutputSettings> | Record<string, SummarizeToolOutputSettings>
| undefined; | undefined;
@@ -1206,6 +1236,90 @@ export class Config {
return this.quotaErrorOccurred; return this.quotaErrorOccurred;
} }
setQuota(
remaining: number | undefined,
limit: number | undefined,
modelId?: string,
): void {
const activeModel = modelId ?? this.getActiveModel();
if (remaining !== undefined && limit !== undefined) {
const current = this.modelQuotas.get(activeModel);
if (
!current ||
current.remaining !== remaining ||
current.limit !== limit
) {
this.modelQuotas.set(activeModel, { remaining, limit });
this.emitQuotaChangedEvent();
}
}
}
private getPooledQuota(): {
remaining?: number;
limit?: number;
resetTime?: string;
} {
const model = this.getModel();
if (!isAutoModel(model)) {
return {};
}
const isPreview =
model === PREVIEW_GEMINI_MODEL_AUTO ||
isPreviewModel(this.getActiveModel());
const proModel = isPreview ? PREVIEW_GEMINI_MODEL : DEFAULT_GEMINI_MODEL;
const flashModel = isPreview
? PREVIEW_GEMINI_FLASH_MODEL
: DEFAULT_GEMINI_FLASH_MODEL;
const proQuota = this.modelQuotas.get(proModel);
const flashQuota = this.modelQuotas.get(flashModel);
if (proQuota || flashQuota) {
// For reset time, take the one that is furthest in the future (most conservative)
const resetTime = [proQuota?.resetTime, flashQuota?.resetTime]
.filter((t): t is string => !!t)
.sort()
.reverse()[0];
return {
remaining: (proQuota?.remaining ?? 0) + (flashQuota?.remaining ?? 0),
limit: (proQuota?.limit ?? 0) + (flashQuota?.limit ?? 0),
resetTime,
};
}
return {};
}
getQuotaRemaining(): number | undefined {
const pooled = this.getPooledQuota();
if (pooled.remaining !== undefined) {
return pooled.remaining;
}
const primaryModel = resolveModel(this.getModel());
return this.modelQuotas.get(primaryModel)?.remaining;
}
getQuotaLimit(): number | undefined {
const pooled = this.getPooledQuota();
if (pooled.limit !== undefined) {
return pooled.limit;
}
const primaryModel = resolveModel(this.getModel());
return this.modelQuotas.get(primaryModel)?.limit;
}
getQuotaResetTime(): string | undefined {
const pooled = this.getPooledQuota();
if (pooled.resetTime !== undefined) {
return pooled.resetTime;
}
const primaryModel = resolveModel(this.getModel());
return this.modelQuotas.get(primaryModel)?.resetTime;
}
getEmbeddingModel(): string { getEmbeddingModel(): string {
return this.embeddingModel; return this.embeddingModel;
} }
@@ -1285,6 +1399,35 @@ export class Config {
const quota = await codeAssistServer.retrieveUserQuota({ const quota = await codeAssistServer.retrieveUserQuota({
project: codeAssistServer.projectId, project: codeAssistServer.projectId,
}); });
if (quota.buckets) {
this.lastRetrievedQuota = quota;
this.lastQuotaFetchTime = Date.now();
for (const bucket of quota.buckets) {
if (
bucket.modelId &&
bucket.remainingAmount &&
bucket.remainingFraction != null
) {
const remaining = parseInt(bucket.remainingAmount, 10);
const limit =
bucket.remainingFraction > 0
? Math.round(remaining / bucket.remainingFraction)
: (this.modelQuotas.get(bucket.modelId)?.limit ?? 0);
if (!isNaN(remaining) && Number.isFinite(limit) && limit > 0) {
this.modelQuotas.set(bucket.modelId, {
remaining,
limit,
resetTime: bucket.resetTime,
});
}
}
}
this.emitQuotaChangedEvent();
}
const hasAccess = const hasAccess =
quota.buckets?.some((b) => b.modelId === PREVIEW_GEMINI_MODEL) ?? false; quota.buckets?.some((b) => b.modelId === PREVIEW_GEMINI_MODEL) ?? false;
this.setHasAccessToPreviewModel(hasAccess); this.setHasAccessToPreviewModel(hasAccess);
@@ -1295,6 +1438,41 @@ export class Config {
} }
} }
async refreshUserQuotaIfStale(
staleMs = 30_000,
): Promise<RetrieveUserQuotaResponse | undefined> {
const now = Date.now();
if (now - this.lastQuotaFetchTime > staleMs) {
return this.refreshUserQuota();
}
return this.lastRetrievedQuota;
}
getLastRetrievedQuota(): RetrieveUserQuotaResponse | undefined {
return this.lastRetrievedQuota;
}
getRemainingQuotaForModel(modelId: string):
| {
remainingAmount?: number;
remainingFraction?: number;
resetTime?: string;
}
| undefined {
const bucket = this.lastRetrievedQuota?.buckets?.find(
(b) => b.modelId === modelId,
);
if (!bucket) return undefined;
return {
remainingAmount: bucket.remainingAmount
? parseInt(bucket.remainingAmount, 10)
: undefined,
remainingFraction: bucket.remainingFraction,
resetTime: bucket.resetTime,
};
}
getCoreTools(): string[] | undefined { getCoreTools(): string[] | undefined {
return this.coreTools; return this.coreTools;
} }
@@ -2160,10 +2338,12 @@ export class Config {
const registry = new ToolRegistry(this, this.messageBus); const registry = new ToolRegistry(this, this.messageBus);
// helper to create & register core tools that are enabled // helper to create & register core tools that are enabled
// eslint-disable-next-line @typescript-eslint/no-explicit-any const maybeRegister = (
const registerCoreTool = (ToolClass: any, ...args: unknown[]) => { toolClass: { name: string; Name?: string },
const className = ToolClass.name; registerFn: () => void,
const toolName = ToolClass.Name || className; ) => {
const className = toolClass.name;
const toolName = toolClass.Name || className;
const coreTools = this.getCoreTools(); const coreTools = this.getCoreTools();
// On some platforms, the className can be minified to _ClassName. // On some platforms, the className can be minified to _ClassName.
const normalizedClassName = className.replace(/^_+/, ''); const normalizedClassName = className.replace(/^_+/, '');
@@ -2180,15 +2360,16 @@ export class Config {
} }
if (isEnabled) { if (isEnabled) {
// Pass message bus to tools (required for policy engine integration) registerFn();
const toolArgs = [...args, this.getMessageBus()];
registry.registerTool(new ToolClass(...toolArgs));
} }
}; };
registerCoreTool(LSTool, this); maybeRegister(LSTool, () =>
registerCoreTool(ReadFileTool, this); registry.registerTool(new LSTool(this, this.messageBus)),
);
maybeRegister(ReadFileTool, () =>
registry.registerTool(new ReadFileTool(this, this.messageBus)),
);
if (this.getUseRipgrep()) { if (this.getUseRipgrep()) {
let useRipgrep = false; let useRipgrep = false;
@@ -2199,30 +2380,60 @@ export class Config {
errorString = String(error); errorString = String(error);
} }
if (useRipgrep) { if (useRipgrep) {
registerCoreTool(RipGrepTool, this); maybeRegister(RipGrepTool, () =>
registry.registerTool(new RipGrepTool(this, this.messageBus)),
);
} else { } else {
logRipgrepFallback(this, new RipgrepFallbackEvent(errorString)); logRipgrepFallback(this, new RipgrepFallbackEvent(errorString));
registerCoreTool(GrepTool, this); maybeRegister(GrepTool, () =>
registry.registerTool(new GrepTool(this, this.messageBus)),
);
} }
} else { } else {
registerCoreTool(GrepTool, this); maybeRegister(GrepTool, () =>
registry.registerTool(new GrepTool(this, this.messageBus)),
);
} }
registerCoreTool(GlobTool, this); maybeRegister(GlobTool, () =>
registerCoreTool(ActivateSkillTool, this); registry.registerTool(new GlobTool(this, this.messageBus)),
registerCoreTool(EditTool, this); );
registerCoreTool(WriteFileTool, this); maybeRegister(ActivateSkillTool, () =>
registerCoreTool(WebFetchTool, this); registry.registerTool(new ActivateSkillTool(this, this.messageBus)),
registerCoreTool(ShellTool, this); );
registerCoreTool(MemoryTool); maybeRegister(EditTool, () =>
registerCoreTool(WebSearchTool, this); registry.registerTool(new EditTool(this, this.messageBus)),
registerCoreTool(AskUserTool); );
maybeRegister(WriteFileTool, () =>
registry.registerTool(new WriteFileTool(this, this.messageBus)),
);
maybeRegister(WebFetchTool, () =>
registry.registerTool(new WebFetchTool(this, this.messageBus)),
);
maybeRegister(ShellTool, () =>
registry.registerTool(new ShellTool(this, this.messageBus)),
);
maybeRegister(MemoryTool, () =>
registry.registerTool(new MemoryTool(this.messageBus)),
);
maybeRegister(WebSearchTool, () =>
registry.registerTool(new WebSearchTool(this, this.messageBus)),
);
maybeRegister(AskUserTool, () =>
registry.registerTool(new AskUserTool(this.messageBus)),
);
if (this.getUseWriteTodos()) { if (this.getUseWriteTodos()) {
registerCoreTool(WriteTodosTool); maybeRegister(WriteTodosTool, () =>
registry.registerTool(new WriteTodosTool(this.messageBus)),
);
} }
if (this.isPlanEnabled()) { if (this.isPlanEnabled()) {
registerCoreTool(ExitPlanModeTool, this); maybeRegister(ExitPlanModeTool, () =>
registerCoreTool(EnterPlanModeTool, this); registry.registerTool(new ExitPlanModeTool(this, this.messageBus)),
);
maybeRegister(EnterPlanModeTool, () =>
registry.registerTool(new EnterPlanModeTool(this, this.messageBus)),
);
} }
// Register Subagents as Tools // Register Subagents as Tools
+25 -26
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -139,12 +139,11 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: undefined, vertexai: undefined,
httpOptions: { httpOptions: expect.objectContaining({
headers: { headers: expect.objectContaining({
'User-Agent': expect.stringContaining('GeminiCLI/1.2.3/gemini-pro'), 'User-Agent': expect.stringContaining('GeminiCLI/1.2.3/gemini-pro'),
'x-gemini-api-privileged-user-id': expect.any(String), }),
}, }),
},
}); });
expect(generator).toEqual( expect(generator).toEqual(
new LoggingContentGenerator(mockGenerator.models, mockConfig), new LoggingContentGenerator(mockGenerator.models, mockConfig),
@@ -209,21 +208,21 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: undefined, vertexai: undefined,
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
'User-Agent': expect.any(String), 'User-Agent': expect.any(String),
'X-Test-Header': 'test', 'X-Test-Header': 'test',
Another: 'value', Another: 'value',
}), }),
}, }),
}); });
expect(GoogleGenAI).toHaveBeenCalledWith( expect(GoogleGenAI).toHaveBeenCalledWith(
expect.not.objectContaining({ expect.not.objectContaining({
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
Authorization: expect.any(String), Authorization: expect.any(String),
}), }),
}, }),
}), }),
); );
}); });
@@ -252,12 +251,12 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: undefined, vertexai: undefined,
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
'User-Agent': expect.any(String), 'User-Agent': expect.any(String),
Authorization: 'Bearer test-api-key', Authorization: 'Bearer test-api-key',
}), }),
}, }),
}); });
}); });
@@ -285,20 +284,20 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: undefined, vertexai: undefined,
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
'User-Agent': expect.any(String), 'User-Agent': expect.any(String),
}), }),
}, }),
}); });
// Explicitly assert that Authorization header is NOT present // Explicitly assert that Authorization header is NOT present
expect(GoogleGenAI).toHaveBeenCalledWith( expect(GoogleGenAI).toHaveBeenCalledWith(
expect.not.objectContaining({ expect.not.objectContaining({
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
Authorization: expect.any(String), Authorization: expect.any(String),
}), }),
}, }),
}), }),
); );
}); });
@@ -322,11 +321,11 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: undefined, vertexai: undefined,
httpOptions: { httpOptions: expect.objectContaining({
headers: { headers: {
'User-Agent': expect.any(String), 'User-Agent': expect.any(String),
}, },
}, }),
}); });
expect(generator).toEqual( expect(generator).toEqual(
new LoggingContentGenerator(mockGenerator.models, mockConfig), new LoggingContentGenerator(mockGenerator.models, mockConfig),
@@ -357,11 +356,11 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: undefined, vertexai: undefined,
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
'User-Agent': expect.any(String), 'User-Agent': expect.any(String),
}), }),
}, }),
apiVersion: 'v1', apiVersion: 'v1',
}); });
}); });
@@ -389,11 +388,11 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: undefined, vertexai: undefined,
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
'User-Agent': expect.any(String), 'User-Agent': expect.any(String),
}), }),
}, }),
}); });
expect(GoogleGenAI).toHaveBeenCalledWith( expect(GoogleGenAI).toHaveBeenCalledWith(
@@ -427,11 +426,11 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: undefined, vertexai: undefined,
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
'User-Agent': expect.any(String), 'User-Agent': expect.any(String),
}), }),
}, }),
}); });
expect(GoogleGenAI).toHaveBeenCalledWith( expect(GoogleGenAI).toHaveBeenCalledWith(
@@ -466,11 +465,11 @@ describe('createContentGenerator', () => {
expect(GoogleGenAI).toHaveBeenCalledWith({ expect(GoogleGenAI).toHaveBeenCalledWith({
apiKey: 'test-api-key', apiKey: 'test-api-key',
vertexai: true, vertexai: true,
httpOptions: { httpOptions: expect.objectContaining({
headers: expect.objectContaining({ headers: expect.objectContaining({
'User-Agent': expect.any(String), 'User-Agent': expect.any(String),
}), }),
}, }),
apiVersion: 'v1alpha', apiVersion: 'v1alpha',
}); });
}); });
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -51,6 +51,7 @@ describe('LoggingContentGenerator', () => {
getContentGeneratorConfig: vi.fn().mockReturnValue({ getContentGeneratorConfig: vi.fn().mockReturnValue({
authType: 'API_KEY', authType: 'API_KEY',
}), }),
refreshUserQuotaIfStale: vi.fn().mockResolvedValue(undefined),
} as unknown as Config; } as unknown as Config;
loggingContentGenerator = new LoggingContentGenerator(wrapped, config); loggingContentGenerator = new LoggingContentGenerator(wrapped, config);
vi.useFakeTimers(); vi.useFakeTimers();
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -34,6 +34,7 @@ import { CodeAssistServer } from '../code_assist/server.js';
import { toContents } from '../code_assist/converter.js'; import { toContents } from '../code_assist/converter.js';
import { isStructuredError } from '../utils/quotaErrorDetection.js'; import { isStructuredError } from '../utils/quotaErrorDetection.js';
import { runInDevTraceSpan, type SpanMetadata } from '../telemetry/trace.js'; import { runInDevTraceSpan, type SpanMetadata } from '../telemetry/trace.js';
import { debugLogger } from '../utils/debugLogger.js';
interface StructuredError { interface StructuredError {
status: number; status: number;
@@ -234,6 +235,9 @@ export class LoggingContentGenerator implements ContentGenerator {
req.config, req.config,
serverDetails, serverDetails,
); );
this.config
.refreshUserQuotaIfStale()
.catch((e) => debugLogger.debug('quota refresh failed', e));
return response; return response;
} catch (error) { } catch (error) {
const durationMs = Date.now() - startTime; const durationMs = Date.now() - startTime;
@@ -355,6 +359,9 @@ export class LoggingContentGenerator implements ContentGenerator {
req.config, req.config,
serverDetails, serverDetails,
); );
this.config
.refreshUserQuotaIfStale()
.catch((e) => debugLogger.debug('quota refresh failed', e));
spanMetadata.output = { spanMetadata.output = {
streamChunks: responses.map((r) => ({ streamChunks: responses.map((r) => ({
content: r.candidates?.[0]?.content ?? null, content: r.candidates?.[0]?.content ?? null,
+24 -1
View File
@@ -1,6 +1,6 @@
/** /**
* @license * @license
* Copyright 2025 Google LLC * Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
@@ -127,6 +127,15 @@ export interface AgentsDiscoveredPayload {
agents: AgentDefinition[]; agents: AgentDefinition[];
} }
/**
* Payload for the 'quota-changed' event.
*/
export interface QuotaChangedPayload {
remaining: number | undefined;
limit: number | undefined;
resetTime?: string;
}
export enum CoreEvent { export enum CoreEvent {
UserFeedback = 'user-feedback', UserFeedback = 'user-feedback',
ModelChanged = 'model-changed', ModelChanged = 'model-changed',
@@ -146,6 +155,7 @@ export enum CoreEvent {
AgentsDiscovered = 'agents-discovered', AgentsDiscovered = 'agents-discovered',
RequestEditorSelection = 'request-editor-selection', RequestEditorSelection = 'request-editor-selection',
EditorSelected = 'editor-selected', EditorSelected = 'editor-selected',
QuotaChanged = 'quota-changed',
} }
/** /**
@@ -161,6 +171,7 @@ export interface CoreEvents extends ExtensionEvents {
[CoreEvent.ConsoleLog]: [ConsoleLogPayload]; [CoreEvent.ConsoleLog]: [ConsoleLogPayload];
[CoreEvent.Output]: [OutputPayload]; [CoreEvent.Output]: [OutputPayload];
[CoreEvent.MemoryChanged]: [MemoryChangedPayload]; [CoreEvent.MemoryChanged]: [MemoryChangedPayload];
[CoreEvent.QuotaChanged]: [QuotaChangedPayload];
[CoreEvent.ExternalEditorClosed]: never[]; [CoreEvent.ExternalEditorClosed]: never[];
[CoreEvent.McpClientUpdate]: Array<Map<string, McpClient> | never>; [CoreEvent.McpClientUpdate]: Array<Map<string, McpClient> | never>;
[CoreEvent.OauthDisplayMessage]: string[]; [CoreEvent.OauthDisplayMessage]: string[];
@@ -311,6 +322,18 @@ export class CoreEventEmitter extends EventEmitter<CoreEvents> {
this._emitOrQueue(CoreEvent.AgentsDiscovered, payload); this._emitOrQueue(CoreEvent.AgentsDiscovered, payload);
} }
/**
* Notifies subscribers that the quota has changed.
*/
emitQuotaChanged(
remaining: number | undefined,
limit: number | undefined,
resetTime?: string,
): void {
const payload: QuotaChangedPayload = { remaining, limit, resetTime };
this.emit(CoreEvent.QuotaChanged, payload);
}
/** /**
* Flushes buffered messages. Call this immediately after primary UI listener * Flushes buffered messages. Call this immediately after primary UI listener
* subscribes. * subscribes.