Update models menu dialog to support manual model selection and group… (#80)

* Update models menu dialog to support manual model selection and group by model family

* fix tests

* update codebase investigator model setting

* fix test

* fix test

* update generated golden file

* regenerate scheme and doc for settings

* use Preview Auto if previewFeatures is set to true
This commit is contained in:
Sehoon Shon
2025-12-11 09:57:27 -05:00
committed by Tommaso Sciortino
parent c3f6e7132b
commit 17bf02b901
26 changed files with 461 additions and 535 deletions

View File

@@ -1283,7 +1283,7 @@ describe('loadCliConfig model selection', () => {
argv,
);
expect(config.getModel()).toBe('auto');
expect(config.getModel()).toBe('auto-gemini-2.5');
});
it('always prefers model from argv', async () => {

View File

@@ -31,6 +31,7 @@ import {
debugLogger,
loadServerHierarchicalMemory,
WEB_FETCH_TOOL_NAME,
PREVIEW_GEMINI_MODEL_AUTO,
} from '@google/gemini-cli-core';
import type { Settings } from './settings.js';
@@ -569,7 +570,9 @@ export async function loadCliConfig(
extraExcludes.length > 0 ? extraExcludes : undefined,
);
const defaultModel = DEFAULT_GEMINI_MODEL_AUTO;
const defaultModel = settings.general?.previewFeatures
? PREVIEW_GEMINI_MODEL_AUTO
: DEFAULT_GEMINI_MODEL_AUTO;
const resolvedModel: string =
argv.model ||
process.env['GEMINI_MODEL'] ||

View File

@@ -19,7 +19,7 @@ import {
DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES,
DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD,
DEFAULT_MODEL_CONFIGS,
GEMINI_MODEL_ALIAS_PRO,
DEFAULT_GEMINI_MODEL,
} from '@google/gemini-cli-core';
import type { CustomTheme } from '../ui/themes/theme.js';
import type { SessionRetentionSettings } from './settings.js';
@@ -1384,7 +1384,7 @@ const SETTINGS_SCHEMA = {
label: 'Model',
category: 'Experimental',
requiresRestart: true,
default: GEMINI_MODEL_ALIAS_PRO,
default: DEFAULT_GEMINI_MODEL,
description:
'The model to use for the Codebase Investigator agent.',
showInDialog: false,

View File

@@ -4,239 +4,186 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { render } from '../../test-utils/render.js';
import { cleanup } from 'ink-testing-library';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import {
GEMINI_MODEL_ALIAS_FLASH_LITE,
GEMINI_MODEL_ALIAS_FLASH,
GEMINI_MODEL_ALIAS_PRO,
DEFAULT_GEMINI_MODEL_AUTO,
} from '@google/gemini-cli-core';
import { render } from 'ink-testing-library';
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { ModelDialog } from './ModelDialog.js';
import { useKeypress } from '../hooks/useKeypress.js';
import { DescriptiveRadioButtonSelect } from './shared/DescriptiveRadioButtonSelect.js';
import { ConfigContext } from '../contexts/ConfigContext.js';
import type { Config } from '@google/gemini-cli-core';
import { KeypressProvider } from '../contexts/KeypressContext.js';
import {
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_MODEL_AUTO,
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
PREVIEW_GEMINI_MODEL,
} from '@google/gemini-cli-core';
import type { Config, ModelSlashCommandEvent } from '@google/gemini-cli-core';
vi.mock('../hooks/useKeypress.js', () => ({
useKeypress: vi.fn(),
}));
const mockedUseKeypress = vi.mocked(useKeypress);
vi.mock('./shared/DescriptiveRadioButtonSelect.js', () => ({
DescriptiveRadioButtonSelect: vi.fn(() => null),
}));
const mockedSelect = vi.mocked(DescriptiveRadioButtonSelect);
const renderComponent = (
props: Partial<React.ComponentProps<typeof ModelDialog>> = {},
contextValue: Partial<Config> | undefined = undefined,
) => {
const defaultProps = {
onClose: vi.fn(),
};
const combinedProps = { ...defaultProps, ...props };
const mockConfig = contextValue
? ({
// --- Functions used by ModelDialog ---
getModel: vi.fn(() => DEFAULT_GEMINI_MODEL_AUTO),
setModel: vi.fn(),
getPreviewFeatures: vi.fn(() => false),
// --- Functions used by ClearcutLogger ---
getUsageStatisticsEnabled: vi.fn(() => true),
getSessionId: vi.fn(() => 'mock-session-id'),
getDebugMode: vi.fn(() => false),
getContentGeneratorConfig: vi.fn(() => ({ authType: 'mock' })),
getUseSmartEdit: vi.fn(() => false),
getProxy: vi.fn(() => undefined),
isInteractive: vi.fn(() => false),
getExperiments: () => {},
// --- Spread test-specific overrides ---
...contextValue,
} as Config)
: undefined;
const renderResult = render(
<ConfigContext.Provider value={mockConfig}>
<ModelDialog {...combinedProps} />
</ConfigContext.Provider>,
);
// Mock dependencies
const mockGetDisplayString = vi.fn();
const mockLogModelSlashCommand = vi.fn();
const mockModelSlashCommandEvent = vi.fn();
vi.mock('@google/gemini-cli-core', async () => {
const actual = await vi.importActual('@google/gemini-cli-core');
return {
...renderResult,
props: combinedProps,
mockConfig,
...actual,
getDisplayString: (val: string) => mockGetDisplayString(val),
logModelSlashCommand: (config: Config, event: ModelSlashCommandEvent) =>
mockLogModelSlashCommand(config, event),
ModelSlashCommandEvent: class {
constructor(model: string) {
mockModelSlashCommandEvent(model);
}
},
};
};
});
describe('<ModelDialog />', () => {
const mockSetModel = vi.fn();
const mockGetModel = vi.fn();
const mockGetPreviewFeatures = vi.fn();
const mockOnClose = vi.fn();
interface MockConfig extends Partial<Config> {
setModel: (model: string) => void;
getModel: () => string;
getPreviewFeatures: () => boolean;
}
const mockConfig: MockConfig = {
setModel: mockSetModel,
getModel: mockGetModel,
getPreviewFeatures: mockGetPreviewFeatures,
};
beforeEach(() => {
vi.clearAllMocks();
vi.resetAllMocks();
mockGetModel.mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
mockGetPreviewFeatures.mockReturnValue(false);
// Default implementation for getDisplayString
mockGetDisplayString.mockImplementation((val: string) => {
if (val === 'auto-gemini-2.5') return 'Auto (Gemini 2.5)';
if (val === 'auto-gemini-3') return 'Auto (Preview)';
return val;
});
});
afterEach(() => {
cleanup();
});
const renderComponent = (contextValue = mockConfig as Config) =>
render(
<KeypressProvider>
<ConfigContext.Provider value={contextValue}>
<ModelDialog onClose={mockOnClose} />
</ConfigContext.Provider>
</KeypressProvider>,
);
it('renders the title and help text', () => {
const { lastFrame, unmount } = renderComponent();
const waitForUpdate = () =>
new Promise((resolve) => setTimeout(resolve, 150));
it('renders the initial "main" view correctly', () => {
const { lastFrame } = renderComponent();
expect(lastFrame()).toContain('Select Model');
expect(lastFrame()).toContain('(Press Esc to close)');
expect(lastFrame()).toContain(
'To use a specific Gemini model on startup, use the --model flag.',
);
unmount();
expect(lastFrame()).toContain('Auto');
expect(lastFrame()).toContain('Manual');
});
it('passes all model options to DescriptiveRadioButtonSelect', () => {
const { unmount } = renderComponent();
expect(mockedSelect).toHaveBeenCalledTimes(1);
const props = mockedSelect.mock.calls[0][0];
expect(props.items).toHaveLength(4);
expect(props.items[0].value).toBe(DEFAULT_GEMINI_MODEL_AUTO);
expect(props.items[1].value).toBe(GEMINI_MODEL_ALIAS_PRO);
expect(props.items[2].value).toBe(GEMINI_MODEL_ALIAS_FLASH);
expect(props.items[3].value).toBe(GEMINI_MODEL_ALIAS_FLASH_LITE);
expect(props.showNumbers).toBe(true);
unmount();
it('renders "main" view with preview options when preview features are enabled', () => {
mockGetPreviewFeatures.mockReturnValue(true);
const { lastFrame } = renderComponent();
expect(lastFrame()).toContain('Auto (Preview)');
});
it('initializes with the model from ConfigContext', () => {
const mockGetModel = vi.fn(() => GEMINI_MODEL_ALIAS_FLASH);
const { unmount } = renderComponent({}, { getModel: mockGetModel });
it('switches to "manual" view when "Manual" is selected', async () => {
const { lastFrame, stdin } = renderComponent();
expect(mockGetModel).toHaveBeenCalled();
expect(mockedSelect).toHaveBeenCalledWith(
expect.objectContaining({
initialIndex: 2,
}),
undefined,
);
unmount();
// Select "Manual" (index 1)
// Press down arrow to move to "Manual"
stdin.write('\u001B[B'); // Arrow Down
await waitForUpdate();
// Press enter to select
stdin.write('\r');
await waitForUpdate();
// Should now show manual options
expect(lastFrame()).toContain(DEFAULT_GEMINI_MODEL);
expect(lastFrame()).toContain(DEFAULT_GEMINI_FLASH_MODEL);
expect(lastFrame()).toContain(DEFAULT_GEMINI_FLASH_LITE_MODEL);
});
it('initializes with "auto" model if context is not provided', () => {
const { unmount } = renderComponent({}, undefined);
it('renders "manual" view with preview options when preview features are enabled', async () => {
mockGetPreviewFeatures.mockReturnValue(true);
mockGetModel.mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
const { lastFrame, stdin } = renderComponent();
expect(mockedSelect).toHaveBeenCalledWith(
expect.objectContaining({
initialIndex: 0,
}),
undefined,
);
unmount();
// Select "Manual" (index 2 because Preview Auto is first, then Auto (Gemini 2.5))
stdin.write('\u001B[B'); // Arrow Down (to Auto (Gemini 2.5))
await waitForUpdate();
stdin.write('\u001B[B'); // Arrow Down (to Manual)
await waitForUpdate();
// Press enter to select Manual
stdin.write('\r');
await waitForUpdate();
expect(lastFrame()).toContain(PREVIEW_GEMINI_MODEL);
});
it('initializes with "auto" model if getModel returns undefined', () => {
const mockGetModel = vi.fn(() => undefined);
// @ts-expect-error This test validates component robustness when getModel
// returns an unexpected undefined value.
const { unmount } = renderComponent({}, { getModel: mockGetModel });
it('sets model and closes when a model is selected in "main" view', async () => {
const { stdin } = renderComponent();
expect(mockGetModel).toHaveBeenCalled();
// Select "Auto" (index 0)
stdin.write('\r');
await waitForUpdate();
// When getModel returns undefined, preferredModel falls back to DEFAULT_GEMINI_MODEL_AUTO
// which has index 0, so initialIndex should be 0
expect(mockedSelect).toHaveBeenCalledWith(
expect.objectContaining({
initialIndex: 0,
}),
undefined,
);
expect(mockedSelect).toHaveBeenCalledTimes(1);
unmount();
expect(mockSetModel).toHaveBeenCalledWith(DEFAULT_GEMINI_MODEL_AUTO);
expect(mockOnClose).toHaveBeenCalled();
});
it('calls config.setModel and onClose when DescriptiveRadioButtonSelect.onSelect is triggered', () => {
const { props, mockConfig, unmount } = renderComponent({}, {}); // Pass empty object for contextValue
it('sets model and closes when a model is selected in "manual" view', async () => {
const { stdin } = renderComponent();
const childOnSelect = mockedSelect.mock.calls[0][0].onSelect;
expect(childOnSelect).toBeDefined();
// Navigate to Manual (index 1) and select
stdin.write('\u001B[B');
await waitForUpdate();
stdin.write('\r');
await waitForUpdate();
childOnSelect(GEMINI_MODEL_ALIAS_PRO);
// Now in manual view. Default selection is first item (DEFAULT_GEMINI_MODEL)
stdin.write('\r');
await waitForUpdate();
// Assert against the default mock provided by renderComponent
expect(mockConfig?.setModel).toHaveBeenCalledWith(GEMINI_MODEL_ALIAS_PRO);
expect(props.onClose).toHaveBeenCalledTimes(1);
unmount();
expect(mockSetModel).toHaveBeenCalledWith(DEFAULT_GEMINI_MODEL);
expect(mockOnClose).toHaveBeenCalled();
});
it('does not pass onHighlight to DescriptiveRadioButtonSelect', () => {
const { unmount } = renderComponent();
it('closes dialog on escape in "main" view', async () => {
const { stdin } = renderComponent();
const childOnHighlight = mockedSelect.mock.calls[0][0].onHighlight;
expect(childOnHighlight).toBeUndefined();
unmount();
stdin.write('\u001B'); // Escape
await waitForUpdate();
expect(mockOnClose).toHaveBeenCalled();
});
it('calls onClose prop when "escape" key is pressed', () => {
const { props, unmount } = renderComponent();
it('goes back to "main" view on escape in "manual" view', async () => {
const { lastFrame, stdin } = renderComponent();
expect(mockedUseKeypress).toHaveBeenCalled();
// Go to manual view
stdin.write('\u001B[B');
await waitForUpdate();
stdin.write('\r');
await waitForUpdate();
const keyPressHandler = mockedUseKeypress.mock.calls[0][0];
const options = mockedUseKeypress.mock.calls[0][1];
expect(lastFrame()).toContain(DEFAULT_GEMINI_MODEL);
expect(options).toEqual({ isActive: true });
// Press Escape
stdin.write('\u001B');
await waitForUpdate();
keyPressHandler({
name: 'escape',
ctrl: false,
meta: false,
shift: false,
paste: false,
insertable: false,
sequence: '',
});
expect(props.onClose).toHaveBeenCalledTimes(1);
keyPressHandler({
name: 'a',
ctrl: false,
meta: false,
shift: false,
paste: false,
insertable: true,
sequence: '',
});
expect(props.onClose).toHaveBeenCalledTimes(1);
unmount();
});
it('updates initialIndex when config context changes', () => {
const mockGetModel = vi.fn(() => DEFAULT_GEMINI_MODEL_AUTO);
const oldMockConfig = {
getModel: mockGetModel,
getPreviewFeatures: vi.fn(() => false),
} as unknown as Config;
const { rerender, unmount } = render(
<ConfigContext.Provider value={oldMockConfig}>
<ModelDialog onClose={vi.fn()} />
</ConfigContext.Provider>,
);
expect(mockedSelect.mock.calls[0][0].initialIndex).toBe(0);
mockGetModel.mockReturnValue(GEMINI_MODEL_ALIAS_FLASH_LITE);
const newMockConfig = {
getModel: mockGetModel,
getPreviewFeatures: vi.fn(() => false),
} as unknown as Config;
rerender(
<ConfigContext.Provider value={newMockConfig}>
<ModelDialog onClose={vi.fn()} />
</ConfigContext.Provider>,
);
// Should be called at least twice: initial render + re-render after context change
expect(mockedSelect).toHaveBeenCalledTimes(2);
expect(mockedSelect.mock.calls[1][0].initialIndex).toBe(3);
unmount();
expect(mockOnClose).not.toHaveBeenCalled();
// Should be back to main view (Manual option visible)
expect(lastFrame()).toContain('Manual');
});
});

View File

@@ -5,19 +5,19 @@
*/
import type React from 'react';
import { useCallback, useContext, useMemo } from 'react';
import { useCallback, useContext, useMemo, useState } from 'react';
import { Box, Text } from 'ink';
import {
PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
PREVIEW_GEMINI_MODEL_AUTO,
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
DEFAULT_GEMINI_MODEL_AUTO,
GEMINI_MODEL_ALIAS_FLASH,
GEMINI_MODEL_ALIAS_FLASH_LITE,
GEMINI_MODEL_ALIAS_PRO,
ModelSlashCommandEvent,
logModelSlashCommand,
getDisplayString,
} from '@google/gemini-cli-core';
import { useKeypress } from '../hooks/useKeypress.js';
import { theme } from '../semantic-colors.js';
@@ -31,61 +31,128 @@ interface ModelDialogProps {
export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
const config = useContext(ConfigContext);
const [view, setView] = useState<'main' | 'manual'>('main');
// Determine the Preferred Model (read once when the dialog opens).
const preferredModel = config?.getModel() || DEFAULT_GEMINI_MODEL_AUTO;
const manualModelSelected = useMemo(() => {
const manualModels = [
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_FLASH_LITE_MODEL,
PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
];
if (manualModels.includes(preferredModel)) {
return preferredModel;
}
return '';
}, [preferredModel]);
useKeypress(
(key) => {
if (key.name === 'escape') {
onClose();
if (view === 'manual') {
setView('main');
} else {
onClose();
}
}
},
{ isActive: true },
);
const options = useMemo(
() => [
const mainOptions = useMemo(() => {
const list = [
{
value: DEFAULT_GEMINI_MODEL_AUTO,
title: 'Auto',
description: 'Let the system choose the best model for your task.',
title: getDisplayString(DEFAULT_GEMINI_MODEL_AUTO),
description:
'Let Gemini CLI decide the best model for the task: gemini-2.5-pro, gemini-2.5-flash',
key: DEFAULT_GEMINI_MODEL_AUTO,
},
{
value: GEMINI_MODEL_ALIAS_PRO,
title: config?.getPreviewFeatures()
? `Pro (${PREVIEW_GEMINI_MODEL}, ${DEFAULT_GEMINI_MODEL})`
: `Pro (${DEFAULT_GEMINI_MODEL})`,
value: 'Manual',
title: manualModelSelected
? `Manual (${manualModelSelected})`
: 'Manual',
description: 'Manually select a model',
key: 'Manual',
},
];
if (config?.getPreviewFeatures()) {
list.unshift({
value: PREVIEW_GEMINI_MODEL_AUTO,
title: getDisplayString(PREVIEW_GEMINI_MODEL_AUTO),
description:
'For complex tasks that require deep reasoning and creativity',
key: GEMINI_MODEL_ALIAS_PRO,
'Let Gemini CLI decide the best model for the task: gemini-3-pro, gemini-3-flash',
key: PREVIEW_GEMINI_MODEL_AUTO,
});
}
return list;
}, [config, manualModelSelected]);
const manualOptions = useMemo(() => {
const list = [
{
value: DEFAULT_GEMINI_MODEL,
title: DEFAULT_GEMINI_MODEL,
key: DEFAULT_GEMINI_MODEL,
},
{
value: GEMINI_MODEL_ALIAS_FLASH,
title: `Flash (${DEFAULT_GEMINI_FLASH_MODEL})`,
description: 'For tasks that need a balance of speed and reasoning',
key: GEMINI_MODEL_ALIAS_FLASH,
value: DEFAULT_GEMINI_FLASH_MODEL,
title: DEFAULT_GEMINI_FLASH_MODEL,
key: DEFAULT_GEMINI_FLASH_MODEL,
},
{
value: GEMINI_MODEL_ALIAS_FLASH_LITE,
title: `Flash-Lite (${DEFAULT_GEMINI_FLASH_LITE_MODEL})`,
description: 'For simple tasks that need to be done quickly',
key: GEMINI_MODEL_ALIAS_FLASH_LITE,
value: DEFAULT_GEMINI_FLASH_LITE_MODEL,
title: DEFAULT_GEMINI_FLASH_LITE_MODEL,
key: DEFAULT_GEMINI_FLASH_LITE_MODEL,
},
],
[config],
);
];
if (config?.getPreviewFeatures()) {
list.unshift(
{
value: PREVIEW_GEMINI_MODEL,
title: PREVIEW_GEMINI_MODEL,
key: PREVIEW_GEMINI_MODEL,
},
{
value: PREVIEW_GEMINI_FLASH_MODEL,
title: PREVIEW_GEMINI_FLASH_MODEL,
key: PREVIEW_GEMINI_FLASH_MODEL,
},
);
}
return list;
}, [config]);
const options = view === 'main' ? mainOptions : manualOptions;
// Calculate the initial index based on the preferred model.
const initialIndex = useMemo(
() => options.findIndex((option) => option.value === preferredModel),
[preferredModel, options],
);
const initialIndex = useMemo(() => {
const idx = options.findIndex((option) => option.value === preferredModel);
if (idx !== -1) {
return idx;
}
if (view === 'main') {
const manualIdx = options.findIndex((o) => o.value === 'Manual');
return manualIdx !== -1 ? manualIdx : 0;
}
return 0;
}, [preferredModel, options, view]);
// Handle selection internally (Autonomous Dialog).
const handleSelect = useCallback(
(model: string) => {
if (model === 'Manual') {
setView('manual');
return;
}
if (config) {
config.setModel(model);
const event = new ModelSlashCommandEvent(model);

View File

@@ -12,7 +12,7 @@ import type { SelectionListItem } from '../../hooks/useSelectionList.js';
export interface DescriptiveRadioSelectItem<T> extends SelectionListItem<T> {
title: string;
description: string;
description?: string;
}
export interface DescriptiveRadioButtonSelectProps<T> {
@@ -62,7 +62,9 @@ export function DescriptiveRadioButtonSelect<T>({
renderItem={(item, { titleColor }) => (
<Box flexDirection="column" key={item.key}>
<Text color={titleColor}>{item.title}</Text>
<Text color={theme.text.secondary}>{item.description}</Text>
{item.description && (
<Text color={theme.text.secondary}>{item.description}</Text>
)}
</Box>
)}
/>

View File

@@ -256,9 +256,8 @@ export class Session {
try {
const model = getEffectiveModel(
this.config.isInFallbackMode(),
this.config.getModel(),
this.config.getPreviewFeatures(),
this.config.isInFallbackMode(),
);
const responseStream = await chat.sendMessageStream(
{ model },