From ac04c388e0862697a6bf9481a0219b70c4a075e3 Mon Sep 17 00:00:00 2001 From: nityam Date: Mon, 23 Feb 2026 09:14:00 +0530 Subject: [PATCH] Fix: Persist manual model selection on restart #19864 (#19891) --- packages/core/src/config/config.test.ts | 100 ++++++++++++++++++++++-- packages/core/src/config/config.ts | 12 +-- 2 files changed, 102 insertions(+), 10 deletions(-) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index c78cfecf20..7a008f12a6 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -22,7 +22,10 @@ import { DEFAULT_OTLP_ENDPOINT, uiTelemetryService, } from '../telemetry/index.js'; -import type { ContentGeneratorConfig } from '../core/contentGenerator.js'; +import type { + ContentGeneratorConfig, + ContentGenerator, +} from '../core/contentGenerator.js'; import { AuthType, createContentGenerator, @@ -44,7 +47,11 @@ import { ACTIVATE_SKILL_TOOL_NAME } from '../tools/tool-names.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_GEMINI_MODEL } from './models.js'; +import { + DEFAULT_GEMINI_MODEL, + PREVIEW_GEMINI_3_1_MODEL, + DEFAULT_GEMINI_MODEL_AUTO, +} from './models.js'; import { Storage } from './storage.js'; vi.mock('fs', async (importOriginal) => { @@ -2302,7 +2309,8 @@ describe('Config Quota & Preview Model Access', () => { vi.mocked(getCodeAssistServer).mockReturnValue(undefined); const result = await config.refreshUserQuota(); expect(result).toBeUndefined(); - expect(config.getHasAccessToPreviewModel()).toBe(false); + // Never set => stays null (unknown); getter returns true so UI shows preview + expect(config.getHasAccessToPreviewModel()).toBe(true); }); it('should return undefined if retrieveUserQuota fails', async () => { @@ -2311,8 +2319,8 @@ describe('Config Quota & Preview Model Access', () => { ); const result = await config.refreshUserQuota(); expect(result).toBeUndefined(); - // Should remain default (false) - expect(config.getHasAccessToPreviewModel()).toBe(false); + // Never set => stays null (unknown); getter returns true so UI shows preview + expect(config.getHasAccessToPreviewModel()).toBe(true); }); }); @@ -2799,3 +2807,85 @@ describe('syncPlanModeTools', () => { expect(setToolsSpy).toHaveBeenCalled(); }); }); + +describe('Model Persistence Bug Fix (#19864)', () => { + const baseParams: ConfigParameters = { + sessionId: 'test-session', + cwd: '/tmp', + targetDir: '/path/to/target', + debugMode: false, + model: PREVIEW_GEMINI_3_1_MODEL, // User saved preview model + }; + + it('should NOT reset preview model for CodeAssist auth when refreshUserQuota is not called (no projectId)', async () => { + const mockContentConfig = { + authType: AuthType.LOGIN_WITH_GOOGLE, + } as Partial as ContentGeneratorConfig; + + const mockContentGenerator = { + generateContent: vi.fn(), + } as Partial as ContentGenerator; + + vi.mocked(createContentGeneratorConfig).mockResolvedValue( + mockContentConfig, + ); + vi.mocked(createContentGenerator).mockResolvedValue(mockContentGenerator); + // getCodeAssistServer returns undefined by default, so refreshUserQuota() isn't called; + // hasAccessToPreviewModel stays null; reset only when === false, so we don't reset. + const config = new Config(baseParams); + + // Verify initial model is the preview model + expect(config.getModel()).toBe(PREVIEW_GEMINI_3_1_MODEL); + + // Call refreshAuth to simulate restart (CodeAssist auth, no projectId) + await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE); + + // Verify the model was NOT reset (bug fix) + expect(config.getModel()).toBe(PREVIEW_GEMINI_3_1_MODEL); + expect(config.getModel()).not.toBe(DEFAULT_GEMINI_MODEL_AUTO); + }); + + it('should NOT reset preview model for USE_GEMINI (hasAccessToPreviewModel is set to true)', async () => { + const mockContentConfig = { + authType: AuthType.USE_GEMINI, + } as Partial as ContentGeneratorConfig; + + const mockContentGenerator = { + generateContent: vi.fn(), + } as Partial as ContentGenerator; + + vi.mocked(createContentGeneratorConfig).mockResolvedValue( + mockContentConfig, + ); + vi.mocked(createContentGenerator).mockResolvedValue(mockContentGenerator); + + const config = new Config(baseParams); + + // Verify initial model is the preview model + expect(config.getModel()).toBe(PREVIEW_GEMINI_3_1_MODEL); + + // Call refreshAuth + await config.refreshAuth(AuthType.USE_GEMINI); + + // For USE_GEMINI, hasAccessToPreviewModel should be set to true + // So the model should NOT be reset + expect(config.getModel()).toBe(PREVIEW_GEMINI_3_1_MODEL); + expect(config.getHasAccessToPreviewModel()).toBe(true); + }); + + it('should persist model when user selects it with persistMode=true', () => { + const onModelChange = vi.fn(); + const config = new Config({ + ...baseParams, + model: DEFAULT_GEMINI_MODEL_AUTO, // Initial model + onModelChange, + }); + + // User selects preview model with persist mode enabled + config.setModel(PREVIEW_GEMINI_3_1_MODEL, false); // isTemporary = false + + // Verify onModelChange was called to persist the model + expect(onModelChange).toHaveBeenCalledWith(PREVIEW_GEMINI_3_1_MODEL); + expect(config.getModel()).toBe(PREVIEW_GEMINI_3_1_MODEL); + }); +}); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index d885e5dc7f..42f8508697 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -578,7 +578,8 @@ export class Config { private readonly bugCommand: BugCommandSettings | undefined; private model: string; private readonly disableLoopDetection: boolean; - private hasAccessToPreviewModel: boolean = false; + // null = unknown (quota not fetched); true = has access; false = definitively no access + private hasAccessToPreviewModel: boolean | null = null; private readonly noBrowser: boolean; private readonly folderTrust: boolean; private ideMode: boolean; @@ -1113,8 +1114,9 @@ export class Config { this.setHasAccessToPreviewModel(true); } - // Update model if user no longer has access to the preview model - if (!this.hasAccessToPreviewModel && isPreviewModel(this.model)) { + // Only reset when we have explicit "no access" (hasAccessToPreviewModel === false). + // When null (quota not fetched) or true, we preserve the saved model. + if (isPreviewModel(this.model) && this.hasAccessToPreviewModel === false) { this.setModel(DEFAULT_GEMINI_MODEL_AUTO); } @@ -1451,10 +1453,10 @@ export class Config { } getHasAccessToPreviewModel(): boolean { - return this.hasAccessToPreviewModel; + return this.hasAccessToPreviewModel !== false; } - setHasAccessToPreviewModel(hasAccess: boolean): void { + setHasAccessToPreviewModel(hasAccess: boolean | null): void { this.hasAccessToPreviewModel = hasAccess; }