diff --git a/packages/cli/src/ui/components/ModelDialog.test.tsx b/packages/cli/src/ui/components/ModelDialog.test.tsx index 6f347faa1d..d5c89215b8 100644 --- a/packages/cli/src/ui/components/ModelDialog.test.tsx +++ b/packages/cli/src/ui/components/ModelDialog.test.tsx @@ -109,7 +109,15 @@ describe('', () => { unmount(); }); - it('switches to "manual" view when "Manual" is selected', async () => { + it('switches to "manual" view when "Manual" is selected and uses getDisplayString for models', async () => { + mockGetDisplayString.mockImplementation((val: string) => { + if (val === DEFAULT_GEMINI_MODEL) return 'Formatted Pro Model'; + if (val === DEFAULT_GEMINI_FLASH_MODEL) return 'Formatted Flash Model'; + if (val === DEFAULT_GEMINI_FLASH_LITE_MODEL) + return 'Formatted Lite Model'; + return val; + }); + const { lastFrame, stdin, waitUntilReady, unmount } = await renderComponent(); @@ -129,9 +137,9 @@ describe('', () => { // Should now show manual options await waitFor(() => { const output = lastFrame(); - expect(output).toContain(DEFAULT_GEMINI_MODEL); - expect(output).toContain(DEFAULT_GEMINI_FLASH_MODEL); - expect(output).toContain(DEFAULT_GEMINI_FLASH_LITE_MODEL); + expect(output).toContain('Formatted Pro Model'); + expect(output).toContain('Formatted Flash Model'); + expect(output).toContain('Formatted Lite Model'); }); unmount(); }); @@ -264,11 +272,16 @@ describe('', () => { unmount(); }); - it('shows the preferred manual model in the main view option', async () => { + it('shows the preferred manual model in the main view option using getDisplayString', async () => { mockGetModel.mockReturnValue(DEFAULT_GEMINI_MODEL); + mockGetDisplayString.mockImplementation((val: string) => { + if (val === DEFAULT_GEMINI_MODEL) return 'My Custom Model Display'; + if (val === 'auto-gemini-2.5') return 'Auto (Gemini 2.5)'; + return val; + }); const { lastFrame, unmount } = await renderComponent(); - expect(lastFrame()).toContain(`Manual (${DEFAULT_GEMINI_MODEL})`); + expect(lastFrame()).toContain('Manual (My Custom Model Display)'); unmount(); }); diff --git a/packages/cli/src/ui/components/ModelDialog.tsx b/packages/cli/src/ui/components/ModelDialog.tsx index d50e9b7153..7d7fea4d86 100644 --- a/packages/cli/src/ui/components/ModelDialog.tsx +++ b/packages/cli/src/ui/components/ModelDialog.tsx @@ -94,7 +94,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element { { value: 'Manual', title: manualModelSelected - ? `Manual (${manualModelSelected})` + ? `Manual (${getDisplayString(manualModelSelected)})` : 'Manual', description: 'Manually select a model', key: 'Manual', @@ -118,17 +118,17 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element { const list = [ { value: DEFAULT_GEMINI_MODEL, - title: DEFAULT_GEMINI_MODEL, + title: getDisplayString(DEFAULT_GEMINI_MODEL), key: DEFAULT_GEMINI_MODEL, }, { value: DEFAULT_GEMINI_FLASH_MODEL, - title: DEFAULT_GEMINI_FLASH_MODEL, + title: getDisplayString(DEFAULT_GEMINI_FLASH_MODEL), key: DEFAULT_GEMINI_FLASH_MODEL, }, { value: DEFAULT_GEMINI_FLASH_LITE_MODEL, - title: DEFAULT_GEMINI_FLASH_LITE_MODEL, + title: getDisplayString(DEFAULT_GEMINI_FLASH_LITE_MODEL), key: DEFAULT_GEMINI_FLASH_LITE_MODEL, }, ]; @@ -145,12 +145,12 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element { list.unshift( { value: previewProValue, - title: previewProModel, + title: getDisplayString(previewProModel), key: previewProModel, }, { value: PREVIEW_GEMINI_FLASH_MODEL, - title: PREVIEW_GEMINI_FLASH_MODEL, + title: getDisplayString(PREVIEW_GEMINI_FLASH_MODEL), key: PREVIEW_GEMINI_FLASH_MODEL, }, );