From 217f27758050afb3307017ae76007bd13cfd06d6 Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Tue, 20 Jan 2026 01:25:15 -0500 Subject: [PATCH] fix: update currentSequenceModel when modelChanged (#17051) --- .../src/ui/hooks/useQuotaAndFallback.test.ts | 20 -------- .../cli/src/ui/hooks/useQuotaAndFallback.ts | 4 -- packages/core/src/config/config.ts | 5 +- packages/core/src/core/client.test.ts | 51 +++++++++++++++++++ packages/core/src/core/client.ts | 11 ++++ packages/core/src/fallback/handler.test.ts | 20 +++++--- packages/core/src/fallback/handler.ts | 2 +- 7 files changed, 77 insertions(+), 36 deletions(-) diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts index ebcbe99752..a0066becd8 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts @@ -166,11 +166,6 @@ describe('useQuotaAndFallback', () => { const intent = await promise!; expect(intent).toBe('retry_always'); - // Verify activateFallbackMode was called - expect(mockConfig.activateFallbackMode).toHaveBeenCalledWith( - 'gemini-flash', - ); - // The pending request should be cleared from the state expect(result.current.proQuotaRequest).toBeNull(); expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(1); @@ -282,11 +277,6 @@ describe('useQuotaAndFallback', () => { const intent = await promise!; expect(intent).toBe('retry_always'); - // Verify activateFallbackMode was called - expect(mockConfig.activateFallbackMode).toHaveBeenCalledWith( - 'model-B', - ); - // The pending request should be cleared from the state expect(result.current.proQuotaRequest).toBeNull(); expect(mockConfig.setQuotaErrorOccurred).toHaveBeenCalledWith(true); @@ -342,11 +332,6 @@ To disable gemini-3-pro-preview, disable "Preview features" in /settings.`, const intent = await promise!; expect(intent).toBe('retry_always'); - // Verify activateFallbackMode was called - expect(mockConfig.activateFallbackMode).toHaveBeenCalledWith( - 'gemini-2.5-pro', - ); - expect(result.current.proQuotaRequest).toBeNull(); }); }); @@ -430,11 +415,6 @@ To disable gemini-3-pro-preview, disable "Preview features" in /settings.`, expect(intent).toBe('retry_always'); expect(result.current.proQuotaRequest).toBeNull(); - // Verify activateFallbackMode was called - expect(mockConfig.activateFallbackMode).toHaveBeenCalledWith( - 'gemini-flash', - ); - // Verify quota error flags are reset expect(mockSetModelSwitchedFromQuotaError).toHaveBeenCalledWith(false); expect(mockConfig.setQuotaErrorOccurred).toHaveBeenCalledWith(false); diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts index d83b082950..4922323567 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts @@ -135,10 +135,6 @@ export function useQuotaAndFallback({ config.setQuotaErrorOccurred(false); if (choice === 'retry_always') { - // Set the model to the fallback model for the current session. - // This ensures the Footer updates and future turns use this model. - // The change is not persisted, so the original model is restored on restart. - config.activateFallbackMode(proQuotaRequest.fallbackModel); historyManager.addItem( { type: MessageType.INFO, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 5db98732b1..0a541d8b92 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1955,9 +1955,8 @@ export class Config { */ async dispose(): Promise { coreEvents.off(CoreEvent.AgentsRefreshed, this.onAgentsRefreshed); - if (this.agentRegistry) { - this.agentRegistry.dispose(); - } + this.agentRegistry?.dispose(); + this.geminiClient?.dispose(); if (this.mcpClientManager) { await this.mcpClientManager.stop(); } diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 1c86d5175d..6be93a92d8 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -48,6 +48,7 @@ import type { import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'; import * as policyCatalog from '../availability/policyCatalog.js'; import { partToString } from '../utils/partUtils.js'; +import { coreEvents } from '../utils/events.js'; vi.mock('../services/chatCompressionService.js'); @@ -290,6 +291,7 @@ describe('Gemini Client (client.ts)', () => { }); afterEach(() => { + client.dispose(); vi.restoreAllMocks(); }); @@ -1579,6 +1581,55 @@ ${JSON.stringify( expect.any(AbortSignal), ); }); + + it('should re-route within the same prompt when the configured model changes', async () => { + mockTurnRunFn.mockClear(); + mockTurnRunFn.mockImplementation(async function* () { + yield { type: 'content', value: 'Hello' }; + }); + + mockRouterService.route.mockResolvedValueOnce({ + model: 'original-model', + reason: 'test', + }); + + let stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-1', + ); + await fromAsync(stream); + + expect(mockRouterService.route).toHaveBeenCalledTimes(1); + expect(mockTurnRunFn).toHaveBeenNthCalledWith( + 1, + { model: 'original-model' }, + [{ text: 'Hi' }], + expect.any(AbortSignal), + ); + + mockRouterService.route.mockResolvedValue({ + model: 'fallback-model', + reason: 'test', + }); + vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-flash'); + coreEvents.emitModelChanged('gemini-2.5-flash'); + + stream = client.sendMessageStream( + [{ text: 'Continue' }], + new AbortController().signal, + 'prompt-1', + ); + await fromAsync(stream); + + expect(mockRouterService.route).toHaveBeenCalledTimes(2); + expect(mockTurnRunFn).toHaveBeenNthCalledWith( + 2, + { model: 'fallback-model' }, + [{ text: 'Continue' }], + expect.any(AbortSignal), + ); + }); }); it('should use getGlobalMemory for system instruction when JIT is enabled', async () => { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 165aa7c00f..0fc489a627 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -58,6 +58,7 @@ import { import { resolveModel } from '../config/models.js'; import type { RetryAvailabilityContext } from '../utils/retry.js'; import { partToString } from '../utils/partUtils.js'; +import { coreEvents, CoreEvent } from '../utils/events.js'; const MAX_TURNS = 100; @@ -94,8 +95,14 @@ export class GeminiClient { this.loopDetector = new LoopDetectionService(config); this.compressionService = new ChatCompressionService(); this.lastPromptId = this.config.getSessionId(); + + coreEvents.on(CoreEvent.ModelChanged, this.handleModelChanged); } + private handleModelChanged = () => { + this.currentSequenceModel = null; + }; + // Hook state to deduplicate BeforeAgent calls and track response for // AfterAgent private hookStateMap = new Map< @@ -253,6 +260,10 @@ export class GeminiClient { this.updateTelemetryTokenCount(); } + dispose() { + coreEvents.off(CoreEvent.ModelChanged, this.handleModelChanged); + } + async resumeChat( history: Content[], resumedSessionData?: ResumedSessionData, diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index 8ae6348384..c6b0997737 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -65,6 +65,8 @@ const createMockConfig = (overrides: Partial = {}): Config => fallbackHandler: undefined, getFallbackModelHandler: vi.fn(), setActiveModel: vi.fn(), + setModel: vi.fn(), + activateFallbackMode: vi.fn(), getModelAvailabilityService: vi.fn(() => createAvailabilityServiceMock({ selectedModel: FALLBACK_MODEL, @@ -198,7 +200,7 @@ describe('handleFallback', () => { expect(result).toBe(true); expect(policyConfig.getFallbackModelHandler).not.toHaveBeenCalled(); - expect(policyConfig.setActiveModel).toHaveBeenCalledWith( + expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith( DEFAULT_GEMINI_FLASH_MODEL, ); } finally { @@ -273,7 +275,7 @@ describe('handleFallback', () => { expect(openBrowserSecurely).toHaveBeenCalledWith( 'https://goo.gle/set-up-gemini-code-assist', ); - expect(policyConfig.setActiveModel).not.toHaveBeenCalled(); + expect(policyConfig.activateFallbackMode).not.toHaveBeenCalled(); }); it('should catch errors from the handler, log an error, and return null', async () => { @@ -378,7 +380,7 @@ describe('handleFallback', () => { ); }); - it('calls setActiveModel and logs telemetry when handler returns "retry_always"', async () => { + it('calls activateFallbackMode when handler returns "retry_always"', async () => { policyHandler.mockResolvedValue('retry_always'); vi.mocked(policyConfig.getModel).mockReturnValue( DEFAULT_GEMINI_MODEL_AUTO, @@ -391,11 +393,13 @@ describe('handleFallback', () => { ); expect(result).toBe(true); - expect(policyConfig.setActiveModel).toHaveBeenCalledWith(FALLBACK_MODEL); + expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith( + FALLBACK_MODEL, + ); // TODO: add logging expect statement }); - it('does NOT call setActiveModel when handler returns "stop"', async () => { + it('does NOT call activateFallbackMode when handler returns "stop"', async () => { policyHandler.mockResolvedValue('stop'); const result = await handleFallback( @@ -405,11 +409,11 @@ describe('handleFallback', () => { ); expect(result).toBe(false); - expect(policyConfig.setActiveModel).not.toHaveBeenCalled(); + expect(policyConfig.activateFallbackMode).not.toHaveBeenCalled(); // TODO: add logging expect statement }); - it('does NOT call setActiveModel when handler returns "retry_once"', async () => { + it('does NOT call activateFallbackMode when handler returns "retry_once"', async () => { policyHandler.mockResolvedValue('retry_once'); const result = await handleFallback( @@ -419,7 +423,7 @@ describe('handleFallback', () => { ); expect(result).toBe(true); - expect(policyConfig.setActiveModel).not.toHaveBeenCalled(); + expect(policyConfig.activateFallbackMode).not.toHaveBeenCalled(); }); }); }); diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index ee3c0d1802..eb4bdf8f0b 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -131,7 +131,7 @@ async function processIntent( case 'retry_always': // TODO(telemetry): Implement generic fallback event logging. Existing // logFlashFallback is specific to a single Model. - config.setActiveModel(fallbackModel); + config.activateFallbackMode(fallbackModel); return true; case 'retry_once':