From b4df7e351bd87aaaa9fd8feec84ae7cc68527e14 Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Mon, 1 Dec 2025 12:41:06 -0800 Subject: [PATCH] feat(core): enhance availability routing with wrapped fallback and single-model policies (#13874) --- .../core/src/availability/policyCatalog.ts | 7 ++- .../src/availability/policyHelpers.test.ts | 3 +- .../core/src/availability/policyHelpers.ts | 8 ++- packages/core/src/fallback/handler.test.ts | 52 ++++++++++++++++--- packages/core/src/fallback/handler.ts | 19 +++---- 5 files changed, 65 insertions(+), 24 deletions(-) diff --git a/packages/core/src/availability/policyCatalog.ts b/packages/core/src/availability/policyCatalog.ts index 74b810d7de..257eb7e072 100644 --- a/packages/core/src/availability/policyCatalog.ts +++ b/packages/core/src/availability/policyCatalog.ts @@ -72,8 +72,11 @@ export function getModelPolicyChain( /** * Provides a default policy scaffold for models not present in the catalog. */ -export function createDefaultPolicy(model: string): ModelPolicy { - return definePolicy({ model }); +export function createDefaultPolicy( + model: string, + options?: { isLastResort?: boolean }, +): ModelPolicy { + return definePolicy({ model, isLastResort: options?.isLastResort }); } export function validateModelPolicyChain(chain: ModelPolicyChain): void { diff --git a/packages/core/src/availability/policyHelpers.test.ts b/packages/core/src/availability/policyHelpers.test.ts index dfa86750a1..d1b7f1fd73 100644 --- a/packages/core/src/availability/policyHelpers.test.ts +++ b/packages/core/src/availability/policyHelpers.test.ts @@ -22,6 +22,7 @@ describe('policyHelpers', () => { isInFallbackMode: () => false, } as unknown as Config; const chain = resolvePolicyChain(config); + expect(chain).toHaveLength(1); expect(chain[0]?.model).toBe('custom-model'); }); @@ -46,7 +47,7 @@ describe('policyHelpers', () => { ]; const context = buildFallbackPolicyContext(chain, 'b'); expect(context.failedPolicy?.model).toBe('b'); - expect(context.candidates.map((p) => p.model)).toEqual(['c']); + expect(context.candidates.map((p) => p.model)).toEqual(['c', 'a']); }); it('returns full chain when model is not in policy list', () => { diff --git a/packages/core/src/availability/policyHelpers.ts b/packages/core/src/availability/policyHelpers.ts index 542af4ce58..c01a9ab890 100644 --- a/packages/core/src/availability/policyHelpers.ts +++ b/packages/core/src/availability/policyHelpers.ts @@ -34,7 +34,9 @@ export function resolvePolicyChain(config: Config): ModelPolicyChain { return chain; } - return [createDefaultPolicy(activeModel), ...chain]; + // If the user specified a model not in the default chain, we assume they want + // *only* that model. We do not fallback to the default chain. + return [createDefaultPolicy(activeModel, { isLastResort: true })]; } /** @@ -52,9 +54,11 @@ export function buildFallbackPolicyContext( if (index === -1) { return { failedPolicy: undefined, candidates: chain }; } + // Return [candidates_after, candidates_before] to prioritize downgrades + // (continuing the chain) before wrapping around to upgrades. return { failedPolicy: chain[index], - candidates: chain.slice(index + 1), + candidates: [...chain.slice(index + 1), ...chain.slice(0, index)], }; } diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index 020b7085db..8abcffc4c4 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -539,7 +539,7 @@ describe('handleFallback', () => { beforeEach(() => { vi.clearAllMocks(); availability = createAvailabilityMock({ - selectedModel: 'gemini-1.5-flash', + selectedModel: DEFAULT_GEMINI_FLASH_MODEL, skipped: [], }); policyHandler = vi.fn().mockResolvedValue('retry_once'); @@ -556,9 +556,16 @@ describe('handleFallback', () => { ); }); - it('uses availability selection when enabled', async () => { - await handleFallback(policyConfig, MOCK_PRO_MODEL, AUTH_OAUTH); - expect(availability.selectFirstAvailable).toHaveBeenCalled(); + it('uses availability selection with correct candidates when enabled', async () => { + vi.spyOn(policyConfig, 'getPreviewFeatures').mockReturnValue(true); + vi.spyOn(policyConfig, 'getModel').mockReturnValue(DEFAULT_GEMINI_MODEL); + + await handleFallback(policyConfig, DEFAULT_GEMINI_MODEL, AUTH_OAUTH); + + expect(availability.selectFirstAvailable).toHaveBeenCalledWith([ + DEFAULT_GEMINI_FLASH_MODEL, + PREVIEW_GEMINI_MODEL, + ]); }); it('falls back to last resort when availability returns null', async () => { @@ -611,6 +618,33 @@ describe('handleFallback', () => { } }); + it('wraps around to upgrade candidates if the current model was selected mid-chain (e.g. by router)', async () => { + // Last-resort failure (Flash) in [Preview, Pro, Flash] checks Preview then Pro (all upstream). + vi.spyOn(policyConfig, 'getPreviewFeatures').mockReturnValue(true); + + availability.selectFirstAvailable = vi.fn().mockReturnValue({ + selectedModel: MOCK_PRO_MODEL, + skipped: [], + }); + policyHandler.mockResolvedValue('retry_once'); + + await handleFallback( + policyConfig, + DEFAULT_GEMINI_FLASH_MODEL, + AUTH_OAUTH, + ); + + expect(availability.selectFirstAvailable).toHaveBeenCalledWith([ + PREVIEW_GEMINI_MODEL, + MOCK_PRO_MODEL, + ]); + expect(policyHandler).toHaveBeenCalledWith( + DEFAULT_GEMINI_FLASH_MODEL, + MOCK_PRO_MODEL, + undefined, + ); + }); + it('logs and returns null when handler resolves to null', async () => { policyHandler.mockResolvedValue(null); const debugLoggerErrorSpy = vi.spyOn(debugLogger, 'error'); @@ -656,7 +690,12 @@ describe('handleFallback', () => { ); }); - it('short-circuits when the failed model is already the last-resort policy', async () => { + it('short-circuits when the failed model is the last-resort policy AND candidates are unavailable', async () => { + // Ensure short-circuit when wrapping to an unavailable upstream model. + availability.selectFirstAvailable = vi + .fn() + .mockReturnValue({ selectedModel: null, skipped: [] }); + const result = await handleFallback( policyConfig, DEFAULT_GEMINI_FLASH_MODEL, @@ -664,7 +703,8 @@ describe('handleFallback', () => { ); expect(result).toBeNull(); - expect(policyConfig.getModelAvailabilityService).not.toHaveBeenCalled(); + // Service called to check upstream; no UI handler since nothing selected. + expect(policyConfig.getModelAvailabilityService).toHaveBeenCalled(); expect(policyConfig.getFallbackModelHandler).not.toHaveBeenCalled(); }); }); diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index c9837cdb1a..a6406d095b 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -135,20 +135,13 @@ async function handlePolicyDrivenFallback( candidates.map((policy) => policy.model), ); - let lastResortPolicy = candidates.find((policy) => policy.isLastResort); - if (!lastResortPolicy) { - debugLogger.warn( - 'No isLastResort policy found in candidates, using last candidate as fallback.', - ); - lastResortPolicy = candidates[candidates.length - 1]; - } + const lastResortPolicy = candidates.find((policy) => policy.isLastResort); + const fallbackModel = selection.selectedModel ?? lastResortPolicy?.model; + const selectedPolicy = candidates.find( + (policy) => policy.model === fallbackModel, + ); - const fallbackModel = selection.selectedModel ?? lastResortPolicy.model; - const selectedPolicy = - candidates.find((policy) => policy.model === fallbackModel) ?? - lastResortPolicy; - - if (!fallbackModel || fallbackModel === failedModel) { + if (!fallbackModel || fallbackModel === failedModel || !selectedPolicy) { return null; }