From a72e9d8a65fa207b84f6883836fa8f0aab5b1ac0 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 12:59:12 -0400 Subject: [PATCH 1/9] fix(core): dynamic fallback routing for exhausted quota models Fixes #23397 This fixes an infinite UI dialog loop caused when hardcoded background utility models (like llm-edit-fixer) run out of quota. Instead of repeatedly asking the user for a fallback and only applying that fallback to the active chat session, activateFallbackMode now registers a global routing override. When the background utility retries, ModelConfigService intercepts its request for the exhausted model and transparently swaps it with the user's chosen fallback model. Additionally, if the chosen fallback model matches the user's current session model, handleFallback will now silently apply the route override without triggering the UI dialog at all. --- packages/core/src/config/config.ts | 14 +++++++++++++- packages/core/src/config/flashFallback.test.ts | 18 ++++++++++++++++++ packages/core/src/fallback/handler.test.ts | 2 ++ packages/core/src/fallback/handler.ts | 16 ++++++++++++---- 4 files changed, 45 insertions(+), 5 deletions(-) diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 4429ce5de1e..4c37f4fc960 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -833,6 +833,7 @@ export class Config implements McpContext, AgentLoopContext { private ideMode: boolean; private _activeModel: string; + private fallbackOverrides = new Map(); private readonly maxSessionTurns: number; private readonly listSessions: boolean; private readonly deleteSession: string | undefined; @@ -1923,14 +1924,25 @@ export class Config implements McpContext, AgentLoopContext { this.modelAvailabilityService.reset(); } - activateFallbackMode(model: string): void { + activateFallbackMode(model: string, failedModel?: string): void { this.setModel(model, true); + if (failedModel) { + this.fallbackOverrides.set(failedModel, model); + this.modelConfigService.registerRuntimeModelOverride({ + match: { model: failedModel }, + modelConfig: { model }, + }); + } const authType = this.getContentGeneratorConfig()?.authType; if (authType) { logFlashFallback(this, new FlashFallbackEvent(authType)); } } + getFallbackOverride(model: string): string | undefined { + return this.fallbackOverrides.get(model); + } + getActiveModel(): string { return this._activeModel ?? this.model; } diff --git a/packages/core/src/config/flashFallback.test.ts b/packages/core/src/config/flashFallback.test.ts index 320d69c5657..1d9fede86a8 100644 --- a/packages/core/src/config/flashFallback.test.ts +++ b/packages/core/src/config/flashFallback.test.ts @@ -73,5 +73,23 @@ describe('Flash Model Fallback Configuration', () => { expect.any(FlashFallbackEvent), ); }); + + it('should set fallback override when failedModel is provided and register runtime override', () => { + config.activateFallbackMode( + DEFAULT_GEMINI_FLASH_MODEL, + DEFAULT_GEMINI_MODEL, + ); + expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL); + expect(config.getFallbackOverride(DEFAULT_GEMINI_MODEL)).toBe( + DEFAULT_GEMINI_FLASH_MODEL, + ); + + // Verify it registers the runtime model override with ModelConfigService + expect( + config + .getModelConfigService() + .getResolvedConfig({ model: DEFAULT_GEMINI_MODEL }).model, + ).toBe(DEFAULT_GEMINI_FLASH_MODEL); + }); }); }); diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index 0bc3096f70f..935a6246ead 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -191,6 +191,7 @@ describe('handleFallback', () => { expect(policyConfig.getFallbackModelHandler).not.toHaveBeenCalled(); expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith( DEFAULT_GEMINI_FLASH_MODEL, + MOCK_PRO_MODEL, ); } finally { chainSpy.mockRestore(); @@ -383,6 +384,7 @@ describe('handleFallback', () => { expect(result).toBe(true); expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith( FALLBACK_MODEL, + MOCK_PRO_MODEL, ); // TODO: add logging expect statement }); diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index 5c4fbe91ff0..9f60dc02263 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -42,8 +42,14 @@ export async function handleFallback( return { service: availability, policy: failedPolicy }; }; + const activeModel = config.getActiveModel(); let fallbackModel: string; + if (!candidates.length) { + if (failedModel !== activeModel) { + applyAvailabilityTransition(getAvailabilityContext, failureKind); + return processIntent(config, 'retry_always', activeModel, failedModel); + } fallbackModel = failedModel; } else { const selection = availability.selectFirstAvailable( @@ -69,10 +75,11 @@ export async function handleFallback( // failureKind is already declared and calculated above const action = resolvePolicyAction(failureKind, selectedPolicy); + const activeModel = config.getActiveModel(); - if (action === 'silent') { + if (action === 'silent' || fallbackModel === activeModel) { applyAvailabilityTransition(getAvailabilityContext, failureKind); - return processIntent(config, 'retry_always', fallbackModel); + return processIntent(config, 'retry_always', fallbackModel, failedModel); } // This will be used in the future when FallbackRecommendation is passed through UI @@ -103,7 +110,7 @@ export async function handleFallback( applyAvailabilityTransition(getAvailabilityContext, failureKind); } - return await processIntent(config, intent, fallbackModel); + return await processIntent(config, intent, fallbackModel, failedModel); } catch (handlerError) { debugLogger.error('Fallback handler failed:', handlerError); return null; @@ -131,12 +138,13 @@ async function processIntent( config: Config, intent: FallbackIntent | null, fallbackModel: string, + failedModel?: string, ): Promise { switch (intent) { case 'retry_always': // TODO(telemetry): Implement generic fallback event logging. Existing // logFlashFallback is specific to a single Model. - config.activateFallbackMode(fallbackModel); + config.activateFallbackMode(fallbackModel, failedModel); return true; case 'retry_once': From 6d0f601b6b8fd627141c43624f7f1c4a2a277de2 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 13:10:05 -0400 Subject: [PATCH 2/9] fix(core): clear fallback overrides on session, auth, and model changes This addresses CR feedback to ensure that runtime model fallback overrides do not persist across session boundaries. They are now actively cleared in setSessionId, refreshAuth, and setModel to match the lifecycle behavior of modelAvailabilityService. --- packages/core/src/config/config.test.ts | 35 +++++++++++++++++++ packages/core/src/config/config.ts | 6 ++++ .../core/src/services/modelConfigService.ts | 4 +++ 3 files changed, 45 insertions(+) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 15eeee82b0c..aa27ba09fc9 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -863,6 +863,16 @@ describe('Server Config (config.ts)', () => { expect(GeminiClient).toHaveBeenCalledWith(config); }); + it('should clear fallback overrides when refreshing auth', async () => { + const config = new Config(baseParams); + config.activateFallbackMode('fallback-model', 'failed-model'); + expect(config.getFallbackOverride('failed-model')).toBe('fallback-model'); + + await config.refreshAuth(AuthType.USE_GEMINI); + + expect(config.getFallbackOverride('failed-model')).toBeUndefined(); + }); + it('should pass Vertex AI routing settings when refreshing auth', async () => { const vertexAiRouting = { requestType: 'shared' as const, @@ -1902,6 +1912,21 @@ describe('Server Config (config.ts)', () => { ); }); + it('clears fallback overrides when session changes', async () => { + const config = new Config({ + ...baseParams, + sessionId: 'session-one', + }); + await config.initialize(); + + config.activateFallbackMode('fallback-model', 'failed-model'); + expect(config.getFallbackOverride('failed-model')).toBe('fallback-model'); + + config.setSessionId('session-two'); + + expect(config.getFallbackOverride('failed-model')).toBeUndefined(); + }); + it('does not throw when changing sessions before the previous plans dir exists', async () => { const config = new Config({ ...baseParams, @@ -2715,6 +2740,16 @@ describe('Config getHooks', () => { expect(spy).toHaveBeenCalled(); }); + it('should clear fallback overrides when setting a model', () => { + const config = new Config(baseParams); + config.activateFallbackMode('fallback-model', 'failed-model'); + expect(config.getFallbackOverride('failed-model')).toBe('fallback-model'); + + config.setModel('new-model'); + + expect(config.getFallbackOverride('failed-model')).toBeUndefined(); + }); + it('should allow setting auto model from auto model and reset availability', () => { const config = new Config({ cwd: '/tmp', diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 4c37f4fc960..0cd66689ce9 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1568,6 +1568,8 @@ export class Config implements McpContext, AgentLoopContext { ) { // Reset availability service when switching auth this.modelAvailabilityService.reset(); + this.fallbackOverrides.clear(); + this.modelConfigService.clearRuntimeOverrides(); // Vertex and Genai have incompatible encryption and sending history with // thoughtSignature from Genai to Vertex will fail, we need to strip them @@ -1829,6 +1831,8 @@ export class Config implements McpContext, AgentLoopContext { this._sessionId = sessionId; this.storage.setSessionId(sessionId); this.trackerService = undefined; + this.fallbackOverrides.clear(); + this.modelConfigService.clearRuntimeOverrides(); this.approvedPlanPath = undefined; this.topicState.reset(); this.skillManager.reset(); @@ -1922,6 +1926,8 @@ export class Config implements McpContext, AgentLoopContext { this.onModelChange(newModel); } this.modelAvailabilityService.reset(); + this.fallbackOverrides.clear(); + this.modelConfigService.clearRuntimeOverrides(); } activateFallbackMode(model: string, failedModel?: string): void { diff --git a/packages/core/src/services/modelConfigService.ts b/packages/core/src/services/modelConfigService.ts index 70ff2f0f66f..21a5e936dbe 100644 --- a/packages/core/src/services/modelConfigService.ts +++ b/packages/core/src/services/modelConfigService.ts @@ -344,6 +344,10 @@ export class ModelConfigService { this.runtimeOverrides.push(override); } + clearRuntimeOverrides(): void { + this.runtimeOverrides.length = 0; + } + /** * Resolves a model configuration by merging settings from aliases and applying overrides. * From 23bedc41f87c44371f4de5d3e784bafeaf98bc04 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 13:12:47 -0400 Subject: [PATCH 3/9] test(core): add unit test for clearRuntimeOverrides in modelConfigService --- .../src/services/modelConfigService.test.ts | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/packages/core/src/services/modelConfigService.test.ts b/packages/core/src/services/modelConfigService.test.ts index 70df1aa7b03..858ed81d5e2 100644 --- a/packages/core/src/services/modelConfigService.test.ts +++ b/packages/core/src/services/modelConfigService.test.ts @@ -668,6 +668,31 @@ describe('ModelConfigService', () => { // Specificity should win over order expect(resolved.generateContentConfig.temperature).toBe(0.1); }); + + it('should clear runtime overrides', () => { + const config: ModelConfigServiceConfig = { + aliases: {}, + overrides: [], + }; + const service = new ModelConfigService(config); + + service.registerRuntimeModelOverride({ + match: { model: 'gemini-pro' }, + modelConfig: { generateContentConfig: { temperature: 0.99 } }, + }); + + expect( + service.getResolvedConfig({ model: 'gemini-pro' }).generateContentConfig + .temperature, + ).toBe(0.99); + + service.clearRuntimeOverrides(); + + expect( + service.getResolvedConfig({ model: 'gemini-pro' }).generateContentConfig + .temperature, + ).toBeUndefined(); + }); }); describe('custom aliases', () => { From 700d0128149069d826c7799bfc7d880b7a17ec19 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 13:29:52 -0400 Subject: [PATCH 4/9] fix(core): ensure auto-routing fallbacks do not register permanent alias overrides This fixes an integration test failure caused by the recent fallback loop fix. While UI-driven fallbacks for hardcoded aliases (like utility models) should register a global runtime override, automated silent fallbacks (auto-routing) should only update the active session model to preserve normal retry and evaluation semantics on subsequent turns. --- .../autoRoutingFallback.integration.test.ts | 14 ++++++++----- packages/core/src/fallback/handler.test.ts | 2 +- packages/core/src/fallback/handler.ts | 20 ++++++++++++------- 3 files changed, 23 insertions(+), 13 deletions(-) diff --git a/packages/core/src/availability/autoRoutingFallback.integration.test.ts b/packages/core/src/availability/autoRoutingFallback.integration.test.ts index 9ea062e1ab0..2ac7e553a58 100644 --- a/packages/core/src/availability/autoRoutingFallback.integration.test.ts +++ b/packages/core/src/availability/autoRoutingFallback.integration.test.ts @@ -386,14 +386,18 @@ describe('Auto Routing Fallback Integration', () => { // Simulate start of next turn config.getModelAvailabilityService().resetTurn(); - // Turn 2: Pro should be attempted again! - // Let's make it succeed this time to verify it works! + // Turn 2: Since fallback was 'retry_always', it is permanent. Flash should be used! vi.spyOn(fakeGenerator, 'generateContent').mockImplementation( async (params) => { - if (params.model === PREVIEW_GEMINI_MODEL) { + if (params.model === PREVIEW_GEMINI_FLASH_MODEL) { return { candidates: [ - { content: { role: 'model', parts: [{ text: 'Pro success' }] } }, + { + content: { + role: 'model', + parts: [{ text: 'Flash Turn 2 success' }], + }, + }, ], } as unknown as GenerateContentResponse; } @@ -411,7 +415,7 @@ describe('Auto Routing Fallback Integration', () => { const result2 = await promise2; expect(result2.candidates?.[0]?.content?.parts?.[0]?.text).toBe( - 'Pro success', + 'Flash Turn 2 success', ); }); }); diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index 935a6246ead..d5bd418a828 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -191,7 +191,7 @@ describe('handleFallback', () => { expect(policyConfig.getFallbackModelHandler).not.toHaveBeenCalled(); expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith( DEFAULT_GEMINI_FLASH_MODEL, - MOCK_PRO_MODEL, + undefined, ); } finally { chainSpy.mockRestore(); diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index 9f60dc02263..7956cd1553c 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -46,10 +46,6 @@ export async function handleFallback( let fallbackModel: string; if (!candidates.length) { - if (failedModel !== activeModel) { - applyAvailabilityTransition(getAvailabilityContext, failureKind); - return processIntent(config, 'retry_always', activeModel, failedModel); - } fallbackModel = failedModel; } else { const selection = availability.selectFirstAvailable( @@ -75,11 +71,21 @@ export async function handleFallback( // failureKind is already declared and calculated above const action = resolvePolicyAction(failureKind, selectedPolicy); - const activeModel = config.getActiveModel(); - if (action === 'silent' || fallbackModel === activeModel) { + if ( + action === 'silent' || + (fallbackModel === activeModel && failedModel !== activeModel) + ) { applyAvailabilityTransition(getAvailabilityContext, failureKind); - return processIntent(config, 'retry_always', fallbackModel, failedModel); + // For standard auto-routing (silent), we only update the active model, so don't pass failedModel. + // For utility bypass, we want a hard runtime override, so pass failedModel. + const overrideFailedModel = action === 'silent' ? undefined : failedModel; + return processIntent( + config, + 'retry_always', + fallbackModel, + overrideFailedModel, + ); } // This will be used in the future when FallbackRecommendation is passed through UI From b4527cc98c48a9bc551ffafa83eba8baeaf4bbac Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 13:35:28 -0400 Subject: [PATCH 5/9] fix(core): preserve fallback overrides on setModel and check availability on utility bypass This addresses CR feedback to ensure runtime model overrides persist across model changes within the same session, preventing utility models from reverting to exhausted configurations. Additionally, the fallback handler now correctly validates the availability of the active model before initiating a silent utility bypass heuristic. --- packages/core/src/availability/testUtils.ts | 2 +- packages/core/src/config/config.test.ts | 4 ++-- packages/core/src/config/config.ts | 2 -- packages/core/src/fallback/handler.test.ts | 5 +++++ packages/core/src/fallback/handler.ts | 7 +++++++ 5 files changed, 15 insertions(+), 5 deletions(-) diff --git a/packages/core/src/availability/testUtils.ts b/packages/core/src/availability/testUtils.ts index d27cfc7ee9b..671c0d4c4e7 100644 --- a/packages/core/src/availability/testUtils.ts +++ b/packages/core/src/availability/testUtils.ts @@ -21,7 +21,7 @@ export function createAvailabilityServiceMock( markHealthy: vi.fn(), markRetryOncePerTurn: vi.fn(), consumeStickyAttempt: vi.fn(), - snapshot: vi.fn(), + snapshot: vi.fn().mockReturnValue({ available: true }), resetTurn: vi.fn(), selectFirstAvailable: vi.fn().mockReturnValue(selection), }; diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index aa27ba09fc9..e0db66712a7 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2740,14 +2740,14 @@ describe('Config getHooks', () => { expect(spy).toHaveBeenCalled(); }); - it('should clear fallback overrides when setting a model', () => { + it('should preserve fallback overrides when setting a new model', () => { const config = new Config(baseParams); config.activateFallbackMode('fallback-model', 'failed-model'); expect(config.getFallbackOverride('failed-model')).toBe('fallback-model'); config.setModel('new-model'); - expect(config.getFallbackOverride('failed-model')).toBeUndefined(); + expect(config.getFallbackOverride('failed-model')).toBe('fallback-model'); }); it('should allow setting auto model from auto model and reset availability', () => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 0cd66689ce9..bac3d053561 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1926,8 +1926,6 @@ export class Config implements McpContext, AgentLoopContext { this.onModelChange(newModel); } this.modelAvailabilityService.reset(); - this.fallbackOverrides.clear(); - this.modelConfigService.clearRuntimeOverrides(); } activateFallbackMode(model: string, failedModel?: string): void { diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index d5bd418a828..d08e4a0ff33 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -208,6 +208,9 @@ describe('handleFallback', () => { selectedModel: MOCK_PRO_MODEL, skipped: [], }); + // Mock activeModel to be unavailable so the utility bypass heuristic is skipped + vi.mocked(availability.snapshot).mockReturnValue({ available: false }); + policyHandler.mockResolvedValue('retry_once'); await handleFallback( @@ -352,6 +355,8 @@ describe('handleFallback', () => { vi.mocked(policyConfig.getModel).mockReturnValue( DEFAULT_GEMINI_MODEL_AUTO, ); + // Mock activeModel to be unavailable so the utility bypass heuristic is skipped + vi.mocked(availability.snapshot).mockReturnValue({ available: false }); const result = await handleFallback( policyConfig, diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index 7956cd1553c..b9bd84ce27d 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -46,6 +46,13 @@ export async function handleFallback( let fallbackModel: string; if (!candidates.length) { + if ( + failedModel !== activeModel && + availability.snapshot(activeModel).available + ) { + applyAvailabilityTransition(getAvailabilityContext, failureKind); + return processIntent(config, 'retry_always', activeModel, failedModel); + } fallbackModel = failedModel; } else { const selection = availability.selectFirstAvailable( From fce1b0129f705937294e93992e37b0dc225bef62 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 13:51:07 -0400 Subject: [PATCH 6/9] fix(core): accurately determine utility models for runtime overrides This addresses CR feedback to ensure that runtime model fallback overrides are correctly applied. The logic to identify utility models now correctly checks if the failed model differs from the active session model, regardless of whether the policy dictates a silent fallback or a user prompt. --- packages/core/src/fallback/handler.ts | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index b9bd84ce27d..e5a3de3d7ba 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -86,7 +86,8 @@ export async function handleFallback( applyAvailabilityTransition(getAvailabilityContext, failureKind); // For standard auto-routing (silent), we only update the active model, so don't pass failedModel. // For utility bypass, we want a hard runtime override, so pass failedModel. - const overrideFailedModel = action === 'silent' ? undefined : failedModel; + const overrideFailedModel = + failedModel !== activeModel ? failedModel : undefined; return processIntent( config, 'retry_always', @@ -109,6 +110,7 @@ export async function handleFallback( const handler = config.getFallbackModelHandler(); if (typeof handler !== 'function') { + throw new Error('STOP 2'); return null; } @@ -123,7 +125,12 @@ export async function handleFallback( applyAvailabilityTransition(getAvailabilityContext, failureKind); } - return await processIntent(config, intent, fallbackModel, failedModel); + return await processIntent( + config, + intent, + fallbackModel, + failedModel !== activeModel ? failedModel : undefined, + ); } catch (handlerError) { debugLogger.error('Fallback handler failed:', handlerError); return null; From 39316d010baaf129a3b3d9caa44f7b51bf74710c Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 13:57:50 -0400 Subject: [PATCH 7/9] chore(core): remove leftover debug code in fallback handler --- packages/core/src/fallback/handler.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index e5a3de3d7ba..2d26279cce8 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -110,7 +110,6 @@ export async function handleFallback( const handler = config.getFallbackModelHandler(); if (typeof handler !== 'function') { - throw new Error('STOP 2'); return null; } From 45bbb89e2b443b9c7b82f82ad9fa34c9c5687bc4 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 14:38:33 -0400 Subject: [PATCH 8/9] test(core): fix integration and unit test expectations for model routing This fixes uncommitted test regressions. The integration test was reverted to correctly expect a retry of the primary model on the next turn (since no hard override is registered for main chat models). The unit test was updated to expect 'undefined' for the failedModel argument when the failed model matches the active session model. --- .../autoRoutingFallback.integration.test.ts | 14 +++++--------- packages/core/src/fallback/handler.test.ts | 2 +- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/packages/core/src/availability/autoRoutingFallback.integration.test.ts b/packages/core/src/availability/autoRoutingFallback.integration.test.ts index 2ac7e553a58..9ea062e1ab0 100644 --- a/packages/core/src/availability/autoRoutingFallback.integration.test.ts +++ b/packages/core/src/availability/autoRoutingFallback.integration.test.ts @@ -386,18 +386,14 @@ describe('Auto Routing Fallback Integration', () => { // Simulate start of next turn config.getModelAvailabilityService().resetTurn(); - // Turn 2: Since fallback was 'retry_always', it is permanent. Flash should be used! + // Turn 2: Pro should be attempted again! + // Let's make it succeed this time to verify it works! vi.spyOn(fakeGenerator, 'generateContent').mockImplementation( async (params) => { - if (params.model === PREVIEW_GEMINI_FLASH_MODEL) { + if (params.model === PREVIEW_GEMINI_MODEL) { return { candidates: [ - { - content: { - role: 'model', - parts: [{ text: 'Flash Turn 2 success' }], - }, - }, + { content: { role: 'model', parts: [{ text: 'Pro success' }] } }, ], } as unknown as GenerateContentResponse; } @@ -415,7 +411,7 @@ describe('Auto Routing Fallback Integration', () => { const result2 = await promise2; expect(result2.candidates?.[0]?.content?.parts?.[0]?.text).toBe( - 'Flash Turn 2 success', + 'Pro success', ); }); }); diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index d08e4a0ff33..7931a230078 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -389,7 +389,7 @@ describe('handleFallback', () => { expect(result).toBe(true); expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith( FALLBACK_MODEL, - MOCK_PRO_MODEL, + undefined, ); // TODO: add logging expect statement }); From 7ffc7989fd4b9029e43e1d7db2a545eed5ec7d97 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Wed, 20 May 2026 14:46:08 -0400 Subject: [PATCH 9/9] fix(core): optimize fallback activation and flatten override chains This addresses CR feedback to: 1. Avoid unnecessary resets of the availability service by only calling setModel when the active model actually changes. 2. Explicitly flatten fallback override chains (e.g., A -> B, then B -> C now correctly resolves A -> C) to maintain consistency across multiple model failures. --- packages/core/src/config/config.ts | 17 ++++++++- .../core/src/config/flashFallback.test.ts | 35 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index bac3d053561..471957890ec 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1929,8 +1929,23 @@ export class Config implements McpContext, AgentLoopContext { } activateFallbackMode(model: string, failedModel?: string): void { - this.setModel(model, true); + if (this.getActiveModel() !== model) { + this.setModel(model, true); + } if (failedModel) { + // Chained fallback mitigation: If we already have overrides that point to the model + // that just failed, we need to update them to point to the new fallback model. + // e.g. A -> B, then B fails and we fallback to C. We must update A to point to C. + for (const [source, target] of this.fallbackOverrides.entries()) { + if (target === failedModel) { + this.fallbackOverrides.set(source, model); + this.modelConfigService.registerRuntimeModelOverride({ + match: { model: source }, + modelConfig: { model }, + }); + } + } + this.fallbackOverrides.set(failedModel, model); this.modelConfigService.registerRuntimeModelOverride({ match: { model: failedModel }, diff --git a/packages/core/src/config/flashFallback.test.ts b/packages/core/src/config/flashFallback.test.ts index 1d9fede86a8..96086413a0c 100644 --- a/packages/core/src/config/flashFallback.test.ts +++ b/packages/core/src/config/flashFallback.test.ts @@ -91,5 +91,40 @@ describe('Flash Model Fallback Configuration', () => { .getResolvedConfig({ model: DEFAULT_GEMINI_MODEL }).model, ).toBe(DEFAULT_GEMINI_FLASH_MODEL); }); + + it('should flatten override chains when a model that was previously a target fails', () => { + // 1. Initial fallback: A -> B + config.activateFallbackMode('model-B', 'model-A'); + expect(config.getFallbackOverride('model-A')).toBe('model-B'); + expect( + config.getModelConfigService().getResolvedConfig({ model: 'model-A' }) + .model, + ).toBe('model-B'); + + // 2. Chained fallback: B fails, fallback to C + // This should update A -> C as well. + config.activateFallbackMode('model-C', 'model-B'); + + expect(config.getFallbackOverride('model-A')).toBe('model-C'); + expect(config.getFallbackOverride('model-B')).toBe('model-C'); + + expect( + config.getModelConfigService().getResolvedConfig({ model: 'model-A' }) + .model, + ).toBe('model-C'); + expect( + config.getModelConfigService().getResolvedConfig({ model: 'model-B' }) + .model, + ).toBe('model-C'); + }); + + it('should not reset availability service if model has not changed', () => { + const resetSpy = vi.spyOn(config.getModelAvailabilityService(), 'reset'); + const currentModel = config.getActiveModel(); + + config.activateFallbackMode(currentModel); + + expect(resetSpy).not.toHaveBeenCalled(); + }); }); });