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,
},
);