From 2aa3667d0a78c5a842dc59fd3d9466e82926cd78 Mon Sep 17 00:00:00 2001 From: "mingholy.lmh" Date: Tue, 23 Sep 2025 23:56:02 +0800 Subject: [PATCH] fix: setModel failure --- packages/cli/src/ui/App.tsx | 37 +++++++++++----- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 44 ++++++++++++++++++- packages/cli/src/ui/hooks/useGeminiStream.ts | 12 +++-- .../src/ui/hooks/useVisionAutoSwitch.test.ts | 14 +++--- .../cli/src/ui/hooks/useVisionAutoSwitch.ts | 14 +++--- packages/core/src/config/config.test.ts | 6 +-- packages/core/src/config/config.ts | 13 +++--- .../core/src/config/flashFallback.test.ts | 12 ++--- packages/core/src/core/client.ts | 2 +- packages/core/src/core/geminiChat.ts | 2 +- packages/core/src/subagents/subagent.test.ts | 13 ++++++ packages/core/src/subagents/subagent.ts | 2 +- 12 files changed, 124 insertions(+), 47 deletions(-) diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index c813e71a..26090018 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -566,7 +566,9 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { } // Switch model for future use but return false to stop current retry - config.setModel(fallbackModel); + config.setModel(fallbackModel).catch((error) => { + console.error('Failed to switch to fallback model:', error); + }); config.setFallbackMode(true); logFlashFallback( config, @@ -650,17 +652,28 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { }, []); const handleModelSelect = useCallback( - (modelId: string) => { - config.setModel(modelId); - setCurrentModel(modelId); - setIsModelSelectionDialogOpen(false); - addItem( - { - type: MessageType.INFO, - text: `Switched model to \`${modelId}\` for this session.`, - }, - Date.now(), - ); + async (modelId: string) => { + try { + await config.setModel(modelId); + setCurrentModel(modelId); + setIsModelSelectionDialogOpen(false); + addItem( + { + type: MessageType.INFO, + text: `Switched model to \`${modelId}\` for this session.`, + }, + Date.now(), + ); + } catch (error) { + console.error('Failed to switch model:', error); + addItem( + { + type: MessageType.ERROR, + text: `Failed to switch to model \`${modelId}\`. Please try again.`, + }, + Date.now(), + ); + } }, [config, setCurrentModel, addItem], ); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 125620cf..57da20c1 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -60,7 +60,9 @@ const mockParseAndFormatApiError = vi.hoisted(() => vi.fn()); const mockHandleVisionSwitch = vi.hoisted(() => vi.fn().mockResolvedValue({ shouldProceed: true }), ); -const mockRestoreOriginalModel = vi.hoisted(() => vi.fn()); +const mockRestoreOriginalModel = vi.hoisted(() => + vi.fn().mockResolvedValue(undefined), +); vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { const actualCoreModule = (await importOriginal()) as any; @@ -301,6 +303,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ); }, { @@ -462,6 +466,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -541,6 +547,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -649,6 +657,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -758,6 +768,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -887,6 +899,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, cancelSubmitSpy, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1198,6 +1212,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1251,6 +1267,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1301,6 +1319,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1349,6 +1369,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1398,6 +1420,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1487,6 +1511,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1537,6 +1563,8 @@ describe('useGeminiStream', () => { vi.fn(), // setModelSwitched vi.fn(), // onEditorClose vi.fn(), // onCancelSubmit + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1602,6 +1630,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1680,6 +1710,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1734,6 +1766,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1943,6 +1977,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -1975,6 +2011,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -2028,6 +2066,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); @@ -2065,6 +2105,8 @@ describe('useGeminiStream', () => { () => {}, () => {}, () => {}, + false, // visionModelPreviewEnabled + undefined, // onVisionSwitchRequired (optional) ), ); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 2a1d4371..5bac2c41 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -765,7 +765,9 @@ export const useGeminiStream = ( if (processingStatus === StreamProcessingStatus.UserCancelled) { // Restore original model if it was temporarily overridden - restoreOriginalModel(); + restoreOriginalModel().catch((error) => { + console.error('Failed to restore original model:', error); + }); isSubmittingQueryRef.current = false; return; } @@ -780,10 +782,14 @@ export const useGeminiStream = ( } // Restore original model if it was temporarily overridden - restoreOriginalModel(); + restoreOriginalModel().catch((error) => { + console.error('Failed to restore original model:', error); + }); } catch (error: unknown) { // Restore original model if it was temporarily overridden - restoreOriginalModel(); + restoreOriginalModel().catch((error) => { + console.error('Failed to restore original model:', error); + }); if (error instanceof UnauthorizedError) { onAuthError(); diff --git a/packages/cli/src/ui/hooks/useVisionAutoSwitch.test.ts b/packages/cli/src/ui/hooks/useVisionAutoSwitch.test.ts index e4322b83..c04a2404 100644 --- a/packages/cli/src/ui/hooks/useVisionAutoSwitch.test.ts +++ b/packages/cli/src/ui/hooks/useVisionAutoSwitch.test.ts @@ -210,7 +210,7 @@ describe('useVisionAutoSwitch hook', () => { let currentModel = initialModel; const mockConfig: Partial = { getModel: vi.fn(() => currentModel), - setModel: vi.fn((m: string) => { + setModel: vi.fn(async (m: string) => { currentModel = m; }), getApprovalMode: vi.fn(() => approvalMode), @@ -335,8 +335,8 @@ describe('useVisionAutoSwitch hook', () => { }); // Now restore - act(() => { - result.current.restoreOriginalModel(); + await act(async () => { + await result.current.restoreOriginalModel(); }); expect(config.setModel).toHaveBeenLastCalledWith(initialModel, { reason: 'vision_auto_switch', @@ -369,8 +369,8 @@ describe('useVisionAutoSwitch hook', () => { }); // Restore should be a no-op since no one-time override was used - act(() => { - result.current.restoreOriginalModel(); + await act(async () => { + await result.current.restoreOriginalModel(); }); // Last call should still be the persisted model set expect((config.setModel as any).mock.calls.pop()?.[0]).toBe('coder-model'); @@ -565,8 +565,8 @@ describe('useVisionAutoSwitch hook', () => { }); // Now restore the original model - act(() => { - result.current.restoreOriginalModel(); + await act(async () => { + await result.current.restoreOriginalModel(); }); // Verify model was restored diff --git a/packages/cli/src/ui/hooks/useVisionAutoSwitch.ts b/packages/cli/src/ui/hooks/useVisionAutoSwitch.ts index 6da989d0..f489c843 100644 --- a/packages/cli/src/ui/hooks/useVisionAutoSwitch.ts +++ b/packages/cli/src/ui/hooks/useVisionAutoSwitch.ts @@ -256,7 +256,7 @@ export function useVisionAutoSwitch( if (config.getApprovalMode() === ApprovalMode.YOLO) { const vlModelId = getDefaultVisionModel(); originalModelRef.current = config.getModel(); - config.setModel(vlModelId, { + await config.setModel(vlModelId, { reason: 'vision_auto_switch', context: 'YOLO mode auto-switch for image content', }); @@ -292,7 +292,7 @@ export function useVisionAutoSwitch( if (visionSwitchResult.modelOverride) { // One-time model override originalModelRef.current = config.getModel(); - config.setModel(visionSwitchResult.modelOverride, { + await config.setModel(visionSwitchResult.modelOverride, { reason: 'vision_auto_switch', context: `Default VLM switch mode: ${defaultVlmSwitchMode} (one-time override)`, }); @@ -302,7 +302,7 @@ export function useVisionAutoSwitch( }; } else if (visionSwitchResult.persistSessionModel) { // Persistent session model change - config.setModel(visionSwitchResult.persistSessionModel, { + await config.setModel(visionSwitchResult.persistSessionModel, { reason: 'vision_auto_switch', context: `Default VLM switch mode: ${defaultVlmSwitchMode} (session persistent)`, }); @@ -319,7 +319,7 @@ export function useVisionAutoSwitch( if (visionSwitchResult.modelOverride) { // One-time model override originalModelRef.current = config.getModel(); - config.setModel(visionSwitchResult.modelOverride, { + await config.setModel(visionSwitchResult.modelOverride, { reason: 'vision_auto_switch', context: 'User-prompted vision switch (one-time override)', }); @@ -329,7 +329,7 @@ export function useVisionAutoSwitch( }; } else if (visionSwitchResult.persistSessionModel) { // Persistent session model change - config.setModel(visionSwitchResult.persistSessionModel, { + await config.setModel(visionSwitchResult.persistSessionModel, { reason: 'vision_auto_switch', context: 'User-prompted vision switch (session persistent)', }); @@ -346,9 +346,9 @@ export function useVisionAutoSwitch( [config, addItem, visionModelPreviewEnabled, onVisionSwitchRequired], ); - const restoreOriginalModel = useCallback(() => { + const restoreOriginalModel = useCallback(async () => { if (originalModelRef.current) { - config.setModel(originalModelRef.current, { + await config.setModel(originalModelRef.current, { reason: 'vision_auto_switch', context: 'Restoring original model after vision switch', }); diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index e4e1cd05..5d83ce20 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -755,7 +755,7 @@ describe('setApprovalMode with folder trust', () => { const logModelSwitchSpy = vi.spyOn(config['logger']!, 'logModelSwitch'); // Change the model - config.setModel('qwen-vl-max-latest', { + await config.setModel('qwen-vl-max-latest', { reason: 'vision_auto_switch', context: 'Test model switch', }); @@ -785,7 +785,7 @@ describe('setApprovalMode with folder trust', () => { const logModelSwitchSpy = vi.spyOn(config['logger']!, 'logModelSwitch'); // Set the same model - config.setModel('qwen3-coder-plus'); + await config.setModel('qwen3-coder-plus'); // Verify that logModelSwitch was not called expect(logModelSwitchSpy).not.toHaveBeenCalled(); @@ -807,7 +807,7 @@ describe('setApprovalMode with folder trust', () => { const logModelSwitchSpy = vi.spyOn(config['logger']!, 'logModelSwitch'); // Change the model without options - config.setModel('qwen-vl-max-latest'); + await config.setModel('qwen-vl-max-latest'); // Verify that logModelSwitch was called with default reason expect(logModelSwitchSpy).toHaveBeenCalledWith({ diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 83d0bce0..9ff19919 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -528,13 +528,13 @@ export class Config { return this.contentGeneratorConfig?.model || this.model; } - setModel( + async setModel( newModel: string, options?: { reason?: ModelSwitchEvent['reason']; context?: string; }, - ): void { + ): Promise { const oldModel = this.getModel(); if (this.contentGeneratorConfig) { @@ -559,13 +559,16 @@ export class Config { // Reinitialize chat with updated configuration while preserving history const geminiClient = this.getGeminiClient(); if (geminiClient && geminiClient.isInitialized()) { - // Use async operation but don't await to avoid blocking - geminiClient.reinitialize().catch((error) => { + // Now await the reinitialize operation to ensure completion + try { + await geminiClient.reinitialize(); + } catch (error) { console.error( 'Failed to reinitialize chat with updated config:', error, ); - }); + throw error; // Re-throw to let callers handle the error + } } } diff --git a/packages/core/src/config/flashFallback.test.ts b/packages/core/src/config/flashFallback.test.ts index a0034ea1..4173786c 100644 --- a/packages/core/src/config/flashFallback.test.ts +++ b/packages/core/src/config/flashFallback.test.ts @@ -41,7 +41,7 @@ describe('Flash Model Fallback Configuration', () => { // with the fallback mechanism. This will be necessary we introduce more // intelligent model routing. describe('setModel', () => { - it('should only mark as switched if contentGeneratorConfig exists', () => { + it('should only mark as switched if contentGeneratorConfig exists', async () => { // Create config without initializing contentGeneratorConfig const newConfig = new Config({ sessionId: 'test-session-2', @@ -52,15 +52,15 @@ describe('Flash Model Fallback Configuration', () => { }); // Should not crash when contentGeneratorConfig is undefined - newConfig.setModel(DEFAULT_GEMINI_FLASH_MODEL); + await newConfig.setModel(DEFAULT_GEMINI_FLASH_MODEL); expect(newConfig.isInFallbackMode()).toBe(false); }); }); describe('getModel', () => { - it('should return contentGeneratorConfig model if available', () => { + it('should return contentGeneratorConfig model if available', async () => { // Simulate initialized content generator config - config.setModel(DEFAULT_GEMINI_FLASH_MODEL); + await config.setModel(DEFAULT_GEMINI_FLASH_MODEL); expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL); }); @@ -88,8 +88,8 @@ describe('Flash Model Fallback Configuration', () => { expect(config.isInFallbackMode()).toBe(false); }); - it('should persist switched state throughout session', () => { - config.setModel(DEFAULT_GEMINI_FLASH_MODEL); + it('should persist switched state throughout session', async () => { + await config.setModel(DEFAULT_GEMINI_FLASH_MODEL); // Setting state for fallback mode as is expected of clients config.setFallbackMode(true); expect(config.isInFallbackMode()).toBe(true); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index ae0c4205..8b965001 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -1053,7 +1053,7 @@ export class GeminiClient { error, ); if (accepted !== false && accepted !== null) { - this.config.setModel(fallbackModel); + await this.config.setModel(fallbackModel); this.config.setFallbackMode(true); return fallbackModel; } diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index bf8aa804..9f541601 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -224,7 +224,7 @@ export class GeminiChat { error, ); if (accepted !== false && accepted !== null) { - this.config.setModel(fallbackModel); + await this.config.setModel(fallbackModel); this.config.setFallbackMode(true); return fallbackModel; } diff --git a/packages/core/src/subagents/subagent.test.ts b/packages/core/src/subagents/subagent.test.ts index 0388f3e5..eabd0a9d 100644 --- a/packages/core/src/subagents/subagent.test.ts +++ b/packages/core/src/subagents/subagent.test.ts @@ -72,6 +72,19 @@ async function createMockConfig( } as unknown as ToolRegistry; vi.spyOn(config, 'getToolRegistry').mockReturnValue(mockToolRegistry); + + // Mock getContentGeneratorConfig to return a valid config + vi.spyOn(config, 'getContentGeneratorConfig').mockReturnValue({ + model: DEFAULT_GEMINI_MODEL, + authType: AuthType.USE_GEMINI, + }); + + // Mock setModel method + vi.spyOn(config, 'setModel').mockResolvedValue(); + + // Mock getSessionId method + vi.spyOn(config, 'getSessionId').mockReturnValue('test-session'); + return { config, toolRegistry: mockToolRegistry }; } diff --git a/packages/core/src/subagents/subagent.ts b/packages/core/src/subagents/subagent.ts index 02cf0e73..19636b3c 100644 --- a/packages/core/src/subagents/subagent.ts +++ b/packages/core/src/subagents/subagent.ts @@ -826,7 +826,7 @@ export class SubAgentScope { ); if (this.modelConfig.model) { - this.runtimeContext.setModel(this.modelConfig.model); + await this.runtimeContext.setModel(this.modelConfig.model); } return new GeminiChat(